from __future__ import division

import os
from distutils.version import LooseVersion
from collections import OrderedDict, namedtuple

import numpy as np
import deepdish
import pandas as pd
import corner
import scipy.stats
import matplotlib
import matplotlib.pyplot as plt

from . import utils
from .utils import (logger, infer_parameters_from_function,
                    check_directory_exists_and_if_not_mkdir)
from .prior import Prior, PriorDict, DeltaFunction

def result_file_name(outdir, label):
    """ Returns the standard filename used for a result file

    Parameters
    ----------
    outdir: str
        Name of the output directory
    label: str
        Naming scheme of the output file

    Returns
    -------
    str: File name of the output file
    """
    return '{}/{}_result.h5'.format(outdir, label)

def read_in_result(filename=None, outdir=None, label=None):
    """ Read in a saved .h5 data file

    Parameters
    ----------
    filename: str
        If given, try to load from this filename
    outdir, label: str
        If given, use the default naming convention for saved results file

    Returns
    -------
    result: bilby.core.result.Result

    Raises
    -------
    ValueError: If no filename is given and either outdir or label is None
                If no bilby.core.result.Result is found in the path

    """
    if filename is None:
        if (outdir is None) and (label is None):
            raise ValueError("No information given to load file")
        else:
            filename = result_file_name(outdir, label)
    if os.path.isfile(filename):
        return Result(**deepdish.io.load(filename))
    else:
        raise IOError("No result '{}' found".format(filename))

class Result(object):
    def __init__(self, label='no_label', outdir='.', sampler=None,
                 search_parameter_keys=None, fixed_parameter_keys=None,
                 priors=None, sampler_kwargs=None, injection_parameters=None,
                 meta_data=None, posterior=None, samples=None,
                 nested_samples=None, log_evidence=np.nan,
                 log_evidence_err=np.nan, log_noise_evidence=np.nan,
                 log_bayes_factor=np.nan, log_likelihood_evaluations=None,
                 sampling_time=None, nburn=None, walkers=None,
                 max_autocorrelation_time=None, parameter_labels=None,
                 parameter_labels_with_unit=None, version=None):
        """ A class to store the results of the sampling run

        Parameters
        ----------
        label, outdir, sampler: str
            The label, output directory, and sampler used
        search_parameter_keys, fixed_parameter_keys: list
            Lists of the search and fixed parameter keys. Elemenents of the
            list should be of type `str` and matchs the keys of the `prior`
        priors: dict, bilby.core.prior.PriorDict
            A dictionary of the priors used in the run
        sampler_kwargs: dict
            Key word arguments passed to the sampler
        injection_parameters: dict
            A dictionary of the injection parameters
        meta_data: dict
            A dictionary of meta data to store about the run
        posterior: pandas.DataFrame
            A pandas data frame of the posterior
        samples, nested_samples: array_like
            An array of the output posterior samples and the unweighted samples
        log_evidence, log_evidence_err, log_noise_evidence, log_bayes_factor: float
            Natural log evidences
        log_likelihood_evaluations: array_like
            The evaluations of the likelihood for each sample point
        sampling_time: float
            The time taken to complete the sampling
        nburn: int
            The number of burn-in steps discarded for MCMC samplers
        walkers: array_like
            The samplers taken by a ensemble MCMC samplers
        max_autocorrelation_time: float
            The estimated maximum autocorrelation time for MCMC samplers
        parameter_labels, parameter_labels_with_unit: list
            Lists of the latex-formatted parameter labels
        version: str,
            Version information for software used to generate the result. Note,
            this information is generated when the result object is initialized

        Note:
            All sampling output parameters, e.g. the samples themselves are
            typically not given at initialisation, but set at a later stage.

        """

        self.label = label
        self.outdir = os.path.abspath(outdir)
        self.sampler = sampler
        self.search_parameter_keys = search_parameter_keys
        self.fixed_parameter_keys = fixed_parameter_keys
        self.parameter_labels = parameter_labels
        self.parameter_labels_with_unit = parameter_labels_with_unit
        self.priors = priors
        self.sampler_kwargs = sampler_kwargs
        self.meta_data = meta_data
        self.injection_parameters = injection_parameters
        self.posterior = posterior
        self.samples = samples
        self.nested_samples = nested_samples
        self.walkers = walkers
        self.nburn = nburn
        self.log_evidence = log_evidence
        self.log_evidence_err = log_evidence_err
        self.log_noise_evidence = log_noise_evidence
        self.log_bayes_factor = log_bayes_factor
        self.log_likelihood_evaluations = log_likelihood_evaluations
        self.sampling_time = sampling_time
        self.version = version
        self.max_autocorrelation_time = max_autocorrelation_time

    def __str__(self):
        """Print a summary """
        if getattr(self, 'posterior', None) is not None:
            if getattr(self, 'log_noise_evidence', None) is not None:
                return ("nsamples: {:d}\n"
                        "log_noise_evidence: {:6.3f}\n"
                        "log_evidence: {:6.3f} +/- {:6.3f}\n"
                        "log_bayes_factor: {:6.3f} +/- {:6.3f}\n"
                        .format(len(self.posterior), self.log_noise_evidence, self.log_evidence,
                                self.log_evidence_err, self.log_bayes_factor,
                                self.log_evidence_err))
            else:
                return ("nsamples: {:d}\n"
                        "log_evidence: {:6.3f} +/- {:6.3f}\n"
                        .format(len(self.posterior), self.log_evidence, self.log_evidence_err))
        else:
            return ''

    @property
    def priors(self):
        if self._priors is not None:
            return self._priors
        else:
            raise ValueError('Result object has no priors')

    @priors.setter
    def priors(self, priors):
        if isinstance(priors, dict):
            self._priors = PriorDict(priors)
            if self.parameter_labels is None:
                self.parameter_labels = [self.priors[k].latex_label for k in
                                         self.search_parameter_keys]
            if self.parameter_labels_with_unit is None:
                self.parameter_labels_with_unit = [
                    self.priors[k].latex_label_with_unit for k in
                    self.search_parameter_keys]

        elif priors is None:
            self._priors = priors
            self.parameter_labels = self.search_parameter_keys
            self.parameter_labels_with_unit = self.search_parameter_keys
        else:
            raise ValueError("Input priors not understood")

    @property
    def samples(self):
        """ An array of samples """
        if self._samples is not None:
            return self._samples
        else:
            raise ValueError("Result object has no stored samples")

    @samples.setter
    def samples(self, samples):
        self._samples = samples

    @property
    def nested_samples(self):
        """" An array of unweighted samples """
        if self._nested_samples is not None:
            return self._nested_samples
        else:
            raise ValueError("Result object has no stored nested samples")

    @nested_samples.setter
    def nested_samples(self, nested_samples):
        self._nested_samples = nested_samples

    @property
    def walkers(self):
        """" An array of the ensemble walkers """
        if self._walkers is not None:
            return self._walkers
        else:
            raise ValueError("Result object has no stored walkers")

    @walkers.setter
    def walkers(self, walkers):
        self._walkers = walkers

    @property
    def nburn(self):
        """" An array of the ensemble walkers """
        if self._nburn is not None:
            return self._nburn
        else:
            raise ValueError("Result object has no stored nburn")

    @nburn.setter
    def nburn(self, nburn):
        self._nburn = nburn

    @property
    def posterior(self):
        """ A pandas data frame of the posterior """
        if self._posterior is not None:
            return self._posterior
        else:
            raise ValueError("Result object has no stored posterior")

    @posterior.setter
    def posterior(self, posterior):
        self._posterior = posterior

    @property
    def version(self):
        return self._version

    @version.setter
    def version(self, version):
        if version is None:
            self._version = 'bilby={}'.format(utils.get_version_information())
        else:
            self._version = version

    def _get_save_data_dictionary(self):
        # This list defines all the parameters saved in the result object
        save_attrs = [
            'label', 'outdir', 'sampler', 'log_evidence', 'log_evidence_err',
            'log_noise_evidence', 'log_bayes_factor', 'priors', 'posterior',
            'injection_parameters', 'meta_data', 'search_parameter_keys',
            'fixed_parameter_keys', 'sampling_time', 'sampler_kwargs',
            'log_likelihood_evaluations', 'samples', 'nested_samples',
            'walkers', 'nburn', 'parameter_labels',
            'parameter_labels_with_unit', 'version']
        dictionary = OrderedDict()
        for attr in save_attrs:
            try:
                dictionary[attr] = getattr(self, attr)
            except ValueError as e:
                logger.debug("Unable to save {}, message: {}".format(attr, e))
                pass
        return dictionary

    def save_to_file(self, overwrite=False):
        """
        Writes the Result to a deepdish h5 file

        Parameters
        ----------
        overwrite: bool, optional
            Whether or not to overwrite an existing result file.
            default=False
        """
        file_name = result_file_name(self.outdir, self.label)
        utils.check_directory_exists_and_if_not_mkdir(self.outdir)
        if os.path.isfile(file_name):
            if overwrite:
                logger.debug('Removing existing file {}'.format(file_name))
                os.remove(file_name)
            else:
                logger.debug(
                    'Renaming existing file {} to {}.old'.format(file_name,
                                                                 file_name))
                os.rename(file_name, file_name + '.old')

        logger.debug("Saving result to {}".format(file_name))

        # Convert the prior to a string representation for saving on disk
        dictionary = self._get_save_data_dictionary()
        if dictionary.get('priors', False):
            dictionary['priors'] = {key: str(self.priors[key]) for key in self.priors}

        # Convert callable sampler_kwargs to strings to avoid pickling issues
        if dictionary.get('sampler_kwargs', None) is not None:
            for key in dictionary['sampler_kwargs']:
                if hasattr(dictionary['sampler_kwargs'][key], '__call__'):
                    dictionary['sampler_kwargs'][key] = str(dictionary['sampler_kwargs'])

        try:
            deepdish.io.save(file_name, dictionary)
        except Exception as e:
            logger.error("\n\n Saving the data has failed with the "
                         "following message:\n {} \n\n".format(e))

    def save_posterior_samples(self):
        """Saves posterior samples to a file"""
        filename = '{}/{}_posterior_samples.txt'.format(self.outdir, self.label)
        utils.check_directory_exists_and_if_not_mkdir(self.outdir)
        self.posterior.to_csv(filename, index=False, header=True)

    def get_latex_labels_from_parameter_keys(self, keys):
        """ Returns a list of latex_labels corresponding to the given keys

        Parameters
        ----------
        keys: list
            List of strings corresponding to the desired latex_labels

        Returns
        -------
        list: The desired latex_labels

        """
        latex_labels = []
        for k in keys:
            if k in self.search_parameter_keys:
                idx = self.search_parameter_keys.index(k)
                latex_labels.append(self.parameter_labels_with_unit[idx])
            elif k in self.parameter_labels:
                latex_labels.append(k)
            else:
                logger.debug(
                    'key {} not a parameter label or latex label'.format(k))
                latex_labels.append(' '.join(k.split('_')))
        return latex_labels

    @property
    def covariance_matrix(self):
        """ The covariance matrix of the samples the posterior """
        samples = self.posterior[self.search_parameter_keys].values
        return np.cov(samples.T)

    @property
    def posterior_volume(self):
        """ The posterior volume """
        if self.covariance_matrix.ndim == 0:
            return np.sqrt(self.covariance_matrix)
        else:
            return 1 / np.sqrt(np.abs(np.linalg.det(
                1 / self.covariance_matrix)))

    @staticmethod
    def prior_volume(priors):
        """ The prior volume, given a set of priors """
        return np.prod([priors[k].maximum - priors[k].minimum for k in priors])

    def occam_factor(self, priors):
        """ The Occam factor,

        See Chapter 28, `Mackay "Information Theory, Inference, and Learning
        Algorithms" <http://www.inference.org.uk/itprnn/book.html>`_ Cambridge
        University Press (2003).

        """
        return self.posterior_volume / self.prior_volume(priors)

    def get_one_dimensional_median_and_error_bar(self, key, fmt='.2f',
                                                 quantiles=[0.16, 0.84]):
        """ Calculate the median and error bar for a given key

        Parameters
        ----------
        key: str
            The parameter key for which to calculate the median and error bar
        fmt: str, ('.2f')
            A format string
        quantiles: list
            A length-2 list of the lower and upper-quantiles to calculate
            the errors bars for.

        Returns
        -------
        summary: namedtuple
            An object with attributes, median, lower, upper and string

        """
        summary = namedtuple('summary', ['median', 'lower', 'upper', 'string'])

        if len(quantiles) != 2:
            raise ValueError("quantiles must be of length 2")

        quants_to_compute = np.array([quantiles[0], 0.5, quantiles[1]])
        quants = np.percentile(self.posterior[key], quants_to_compute * 100)
        summary.median = quants[1]
        summary.plus = quants[2] - summary.median
        summary.minus = summary.median - quants[0]

        fmt = "{{0:{0}}}".format(fmt).format
        string_template = r"${{{0}}}_{{-{1}}}^{{+{2}}}$"
        summary.string = string_template.format(
            fmt(summary.median), fmt(summary.minus), fmt(summary.plus))
        return summary

    def plot_single_density(self, key, prior=None, cumulative=False,
                            title=None, truth=None, save=True,
                            file_base_name=None, bins=50, label_fontsize=16,
                            title_fontsize=16, quantiles=[0.16, 0.84], dpi=300):
        """ Plot a 1D marginal density, either probablility or cumulative.

        Parameters
        ----------
        key: str
            Name of the parameter to plot
        prior: {bool (True), bilby.core.prior.Prior}
            If true, add the stored prior probability density function to the
            one-dimensional marginal distributions. If instead a Prior
            is provided, this will be plotted.
        cumulative: bool
            If true plot the CDF
        title: bool
            If true, add 1D title of the median and (by default 1-sigma)
            error bars. To change the error bars, pass in the quantiles kwarg.
            See method `get_one_dimensional_median_and_error_bar` for further
            details). If `quantiles=None` is passed in, no title is added.
        truth: {bool, float}
            If true, plot self.injection_parameters[parameter].
            If float, plot this value.
        save: bool:
            If true, save plot to disk.
        file_base_name: str, optional
            If given, the base file name to use (by default `outdir/label_` is
            used)
        bins: int
            The number of histogram bins
        label_fontsize, title_fontsize: int
            The fontsizes for the labels and titles
        quantiles: list
            A length-2 list of the lower and upper-quantiles to calculate
            the errors bars for.
        dpi: int
            Dots per inch resolution of the plot

        Returns
        -------
        figure: matplotlib.pyplot.figure
            A matplotlib figure object
        """
        logger.info('Plotting {} marginal distribution'.format(key))
        label = self.get_latex_labels_from_parameter_keys([key])[0]
        fig, ax = plt.subplots()
        ax.hist(self.posterior[key].values, bins=bins, density=True,
                histtype='step', cumulative=cumulative)
        ax.set_xlabel(label, fontsize=label_fontsize)
        if truth is not None:
            ax.axvline(truth, ls='-', color='orange')

        summary = self.get_one_dimensional_median_and_error_bar(
            key, quantiles=quantiles)
        ax.axvline(summary.median - summary.minus, ls='--', color='C0')
        ax.axvline(summary.median + summary.plus, ls='--', color='C0')
        if title:
            ax.set_title(summary.string, fontsize=title_fontsize)

        if isinstance(prior, Prior):
            theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300)
            ax.plot(theta, Prior.prob(theta), color='C2')

        if save:
            fig.tight_layout()
            if cumulative:
                file_name = file_base_name + key + '_cdf'
            else:
                file_name = file_base_name + key + '_pdf'
            fig.savefig(file_name, dpi=dpi)

        return fig

    def plot_marginals(self, parameters=None, priors=None, titles=True,
                       file_base_name=None, bins=50, label_fontsize=16,
                       title_fontsize=16, quantiles=[0.16, 0.84], dpi=300):
        """ Plot 1D marginal distributions

        Parameters
        ----------
        parameters: (list, dict), optional
            If given, either a list of the parameter names to include, or a
            dictionary of parameter names and their "true" values to plot.
        priors: {bool (False), bilby.core.prior.PriorDict}
            If true, add the stored prior probability density functions to the
            one-dimensional marginal distributions. If instead a PriorDict
            is provided, this will be plotted.
        titles: bool
            If true, add 1D titles of the median and (by default 1-sigma)
            error bars. To change the error bars, pass in the quantiles kwarg.
            See method `get_one_dimensional_median_and_error_bar` for further
            details). If `quantiles=None` is passed in, no title is added.
        file_base_name: str, optional
            If given, the base file name to use (by default `outdir/label_` is
            used)
        bins: int
            The number of histogram bins
        label_fontsize, title_fontsize: int
            The fontsizes for the labels and titles
        quantiles: list
            A length-2 list of the lower and upper-quantiles to calculate
            the errors bars for.
        dpi: int
            Dots per inch resolution of the plot

        Returns
        -------
        """
        if isinstance(parameters, dict):
            plot_parameter_keys = list(parameters.keys())
            truths = parameters
        elif parameters is None:
            plot_parameter_keys = self.posterior.keys()
            if self.injection_parameters is None:
                truths = dict()
            else:
                truths = self.injection_parameters
        else:
            plot_parameter_keys = list(parameters)
            if self.injection_parameters is None:
                truths = dict()
            else:
                truths = self.injection_parameters

        if file_base_name is None:
            file_base_name = '{}/{}_1d/'.format(self.outdir, self.label)
            check_directory_exists_and_if_not_mkdir(file_base_name)

        if priors is True:
            priors = getattr(self, 'priors', dict())
        elif isinstance(priors, dict):
            pass
        elif priors in [False, None]:
            priors = dict()
        else:
            raise ValueError('Input priors={} not understood'.format(priors))

        for i, key in enumerate(plot_parameter_keys):
            if not isinstance(self.posterior[key].values[0], float):
                continue
            prior = priors.get(key, None)
            truth = truths.get(key, None)
            for cumulative in [False, True]:
                fig = self.plot_single_density(
                    key, prior=prior, cumulative=cumulative, title=titles,
                    truth=truth, save=True, file_base_name=file_base_name,
                    bins=bins, label_fontsize=label_fontsize, dpi=dpi,
                    title_fontsize=title_fontsize, quantiles=quantiles)
                plt.close(fig)

    def plot_corner(self, parameters=None, priors=None, titles=True, save=True,
                    filename=None, dpi=300, **kwargs):
        """ Plot a corner-plot

        Parameters
        ----------
        parameters: (list, dict), optional
            If given, either a list of the parameter names to include, or a
            dictionary of parameter names and their "true" values to plot.
        priors: {bool (False), bilby.core.prior.PriorDict}
            If true, add the stored prior probability density functions to the
            one-dimensional marginal distributions. If instead a PriorDict
            is provided, this will be plotted.
        titles: bool
            If true, add 1D titles of the median and (by default 1-sigma)
            error bars. To change the error bars, pass in the quantiles kwarg.
            See method `get_one_dimensional_median_and_error_bar` for further
            details). If `quantiles=None` is passed in, no title is added.
        save: bool, optional
            If true, save the image using the given label and outdir
        filename: str, optional
            If given, overwrite the default filename
        dpi: int, optional
            Dots per inch resolution of the plot
        **kwargs:
            Other keyword arguments are passed to `corner.corner`. We set some
            defaults to improve the basic look and feel, but these can all be
            overridden.

        Notes
        -----
            The generation of the corner plot themselves is done by the corner
            python module, see https://corner.readthedocs.io for more
            information.

        Returns
        -------
        fig:
            A matplotlib figure instance

        """

        # If in testing mode, not corner plots are generated
        if utils.command_line_args.test:
            return

        # bilby default corner kwargs. Overwritten by anything passed to kwargs
        defaults_kwargs = dict(
            bins=50, smooth=0.9, label_kwargs=dict(fontsize=16),
            title_kwargs=dict(fontsize=16), color='#0072C1',
            truth_color='tab:orange', quantiles=[0.16, 0.84],
            levels=(1 - np.exp(-0.5), 1 - np.exp(-2), 1 - np.exp(-9 / 2.)),
            plot_density=False, plot_datapoints=True, fill_contours=True,
            max_n_ticks=3)

        if LooseVersion(matplotlib.__version__) < "2.1":
            defaults_kwargs['hist_kwargs'] = dict(normed=True)
        else:
            defaults_kwargs['hist_kwargs'] = dict(density=True)

        if 'lionize' in kwargs and kwargs['lionize'] is True:
            defaults_kwargs['truth_color'] = 'tab:blue'
            defaults_kwargs['color'] = '#FF8C00'

        defaults_kwargs.update(kwargs)
        kwargs = defaults_kwargs

        # Handle if truths was passed in
        if 'truth' in kwargs:
            kwargs['truths'] = kwargs.pop('truth')
        if kwargs.get('truths'):
            truths = kwargs.get('truths')
            if isinstance(parameters, list) and isinstance(truths, list):
                if len(parameters) != len(truths):
                    raise ValueError(
                        "Length of parameters and truths don't match")
            elif isinstance(truths, dict) and parameters is None:
                parameters = kwargs.pop('truths')
            else:
                raise ValueError(
                    "Combination of parameters and truths not understood")

        # If injection parameters where stored, use these as parameter values
        # but do not overwrite input parameters (or truths)
        cond1 = getattr(self, 'injection_parameters', None) is not None
        cond2 = parameters is None
        if cond1 and cond2:
            parameters = {key: self.injection_parameters[key] for key in
                          self.search_parameter_keys}

        # If parameters is a dictionary, use the keys to determine which
        # parameters to plot and the values as truths.
        if isinstance(parameters, dict):
            plot_parameter_keys = list(parameters.keys())
            kwargs['truths'] = list(parameters.values())
        elif parameters is None:
            plot_parameter_keys = self.search_parameter_keys
        else:
            plot_parameter_keys = list(parameters)

        # Get latex formatted strings for the plot labels
        kwargs['labels'] = kwargs.get(
            'labels', self.get_latex_labels_from_parameter_keys(
                plot_parameter_keys))

        # Create the data array to plot and pass everything to corner
        xs = self.posterior[plot_parameter_keys].values
        fig = corner.corner(xs, **kwargs)
        axes = fig.get_axes()

        #  Add the titles
        if titles and kwargs.get('quantiles', None) is not None:
            for i, par in enumerate(plot_parameter_keys):
                ax = axes[i + i * len(plot_parameter_keys)]
                if ax.title.get_text() == '':
                    ax.set_title(self.get_one_dimensional_median_and_error_bar(
                        par, quantiles=kwargs['quantiles']).string,
                        **kwargs['title_kwargs'])

        #  Add priors to the 1D plots
        if priors is True:
            priors = getattr(self, 'priors', False)
        if isinstance(priors, dict):
            for i, par in enumerate(plot_parameter_keys):
                ax = axes[i + i * len(plot_parameter_keys)]
                theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300)
                ax.plot(theta, priors[par].prob(theta), color='C2')
        elif priors in [False, None]:
            pass
        else:
            raise ValueError('Input priors={} not understood'.format(priors))

        if save:
            if filename is None:
                utils.check_directory_exists_and_if_not_mkdir(self.outdir)
                filename = '{}/{}_corner.png'.format(self.outdir, self.label)
            logger.debug('Saving corner plot to {}'.format(filename))
            fig.savefig(filename, dpi=dpi)

        return fig

    def plot_walkers(self, **kwargs):
        """ Method to plot the trace of the walkers in an ensemble MCMC plot """
        if hasattr(self, 'walkers') is False:
            logger.warning("Cannot plot_walkers as no walkers are saved")
            return

        if utils.command_line_args.test:
            return

        nwalkers, nsteps, ndim = self.walkers.shape
        idxs = np.arange(nsteps)
        fig, axes = plt.subplots(nrows=ndim, figsize=(6, 3 * ndim))
        walkers = self.walkers[:, :, :]
        for i, ax in enumerate(axes):
            ax.plot(idxs[:self.nburn + 1], walkers[:, :self.nburn + 1, i].T,
                    lw=0.1, color='r')
            ax.set_ylabel(self.parameter_labels[i])

        for i, ax in enumerate(axes):
            ax.plot(idxs[self.nburn:], walkers[:, self.nburn:, i].T, lw=0.1,
                    color='k')
            ax.set_ylabel(self.parameter_labels[i])

        fig.tight_layout()
        filename = '{}/{}_walkers.png'.format(self.outdir, self.label)
        logger.debug('Saving walkers plot to {}'.format('filename'))
        utils.check_directory_exists_and_if_not_mkdir(self.outdir)
        fig.savefig(filename)

    def plot_with_data(self, model, x, y, ndraws=1000, npoints=1000,
                       xlabel=None, ylabel=None, data_label='data',
                       data_fmt='o', draws_label=None, filename=None,
                       maxl_label='max likelihood', dpi=300):
        """ Generate a figure showing the data and fits to the data

        Parameters
        ----------
        model: function
            A python function which when called as `model(x, **kwargs)` returns
            the model prediction (here `kwargs` is a dictionary of key-value
            pairs of the model parameters.
        x, y: np.ndarray
            The independent and dependent data to plot
        ndraws: int
            Number of draws from the posterior to plot
        npoints: int
            Number of points used to plot the smoothed fit to the data
        xlabel, ylabel: str
            Labels for the axes
        data_label, draws_label, maxl_label: str
            Label for the data, draws, and max likelihood legend
        data_fmt: str
            Matpltolib fmt code, defaults to `'-o'`
        dpi: int
            Passed to `plt.savefig`
        filename: str
            If given, the filename to use. Otherwise, the filename is generated
            from the outdir and label attributes.

        """

        # Determine model_posterior, the subset of the full posterior which
        # should be passed into the model
        model_keys = infer_parameters_from_function(model)
        model_posterior = self.posterior[model_keys]

        xsmooth = np.linspace(np.min(x), np.max(x), npoints)
        fig, ax = plt.subplots()
        logger.info('Plotting {} draws'.format(ndraws))
        for _ in range(ndraws):
            s = model_posterior.sample().to_dict('records')[0]
            ax.plot(xsmooth, model(xsmooth, **s), alpha=0.25, lw=0.1, color='r',
                    label=draws_label)
        try:
            if all(~np.isnan(self.posterior.log_likelihood)):
                logger.info('Plotting maximum likelihood')
                s = model_posterior.iloc[self.posterior.log_likelihood.idxmax()]
                ax.plot(xsmooth, model(xsmooth, **s), lw=1, color='k',
                        label=maxl_label)
        except AttributeError:
            logger.debug(
                "No log likelihood values stored, unable to plot max")

        ax.plot(x, y, data_fmt, markersize=2, label=data_label)

        if xlabel is not None:
            ax.set_xlabel(xlabel)
        if ylabel is not None:
            ax.set_ylabel(ylabel)

        handles, labels = plt.gca().get_legend_handles_labels()
        by_label = OrderedDict(zip(labels, handles))
        plt.legend(by_label.values(), by_label.keys())
        ax.legend(numpoints=3)
        fig.tight_layout()
        if filename is None:
            utils.check_directory_exists_and_if_not_mkdir(self.outdir)
            filename = '{}/{}_plot_with_data'.format(self.outdir, self.label)
        fig.savefig(filename, dpi=dpi)

    def samples_to_posterior(self, likelihood=None, priors=None,
                             conversion_function=None):
        """
        Convert array of samples to posterior (a Pandas data frame)

        Also applies the conversion function to any stored posterior

        Parameters
        ----------
        likelihood: bilby.likelihood.GravitationalWaveTransient, optional
            GravitationalWaveTransient likelihood used for sampling.
        priors: dict, optional
            Dictionary of prior object, used to fill in delta function priors.
        conversion_function: function, optional
            Function which adds in extra parameters to the data frame,
            should take the data_frame, likelihood and prior as arguments.
        """
        try:
            data_frame = self.posterior
        except ValueError:
            data_frame = pd.DataFrame(
                self.samples, columns=self.search_parameter_keys)
            for key in priors:
                if isinstance(priors[key], DeltaFunction):
                    data_frame[key] = priors[key].peak
                elif isinstance(priors[key], float):
                    data_frame[key] = priors[key]
            data_frame['log_likelihood'] = getattr(
                self, 'log_likelihood_evaluations', np.nan)
        if conversion_function is not None:
            data_frame = conversion_function(data_frame, likelihood, priors)
        self.posterior = data_frame

    def calculate_prior_values(self, priors):
        """
        Evaluate prior probability for each parameter for each sample.

        Parameters
        ----------
        priors: dict, PriorDict
            Prior distributions
        """
        self.prior_values = pd.DataFrame()
        for key in priors:
            if key in self.posterior.keys():
                if isinstance(priors[key], DeltaFunction):
                    continue
                else:
                    self.prior_values[key]\
                        = priors[key].prob(self.posterior[key].values)

    def get_all_injection_credible_levels(self):
        """
        Get credible levels for all parameters in self.injection_parameters

        Returns
        -------
        credible_levels: dict
            The credible levels at which the injected parameters are found.
        """
        if self.injection_parameters is None:
            raise(TypeError, "Result object has no 'injection_parameters'. "
                             "Cannot copmute credible levels.")
        credible_levels = {key: self.get_injection_credible_level(key)
                           for key in self.search_parameter_keys
                           if isinstance(self.injection_parameters[key], float)}
        return credible_levels

    def get_injection_credible_level(self, parameter):
        """
        Get the credible level of the injected parameter

        Calculated as CDF(injection value)

        Parameters
        ----------
        parameter: str
            Parameter to get credible level for
        Returns
        -------
        float: credible level
        """
        if self.injection_parameters is None:
            raise(TypeError, "Result object has no 'injection_parameters'. "
                             "Cannot copmute credible levels.")
        if parameter in self.posterior and\
                parameter in self.injection_parameters:
            credible_level =\
                sum(self.posterior[parameter].values <
                    self.injection_parameters[parameter]) / len(self.posterior)
            return credible_level
        else:
            return np.nan

    def _check_attribute_match_to_other_object(self, name, other_object):
        """ Check attribute name exists in other_object and is the same

        Parameters
        ----------
        name: str
            Name of the attribute in this instance
        other_object: object
            Other object with attributes to compare with

        Returns
        -------
        bool: True if attribute name matches with an attribute of other_object, False otherwise

        """
        A = getattr(self, name, False)
        B = getattr(other_object, name, False)
        logger.debug('Checking {} value: {}=={}'.format(name, A, B))
        if (A is not False) and (B is not False):
            typeA = type(A)
            typeB = type(B)
            if typeA == typeB:
                if typeA in [str, float, int, dict, list]:
                    try:
                        return A == B
                    except ValueError:
                        return False
                elif typeA in [np.ndarray]:
                    return np.all(A == B)
        return False

    @property
    def kde(self):
        """ Kernel density estimate built from the stored posterior

        Uses `scipy.stats.gaussian_kde` to generate the kernel density
        """
        try:
            return self._kde
        except AttributeError:
            self._kde = scipy.stats.gaussian_kde(
                self.posterior[self.search_parameter_keys].values.T)
            return self._kde

    def posterior_probability(self, sample):
        """ Calculate the posterior probabily for a new sample

        This queries a Kernel Density Estimate of the posterior to calculate
        the posterior probability density for the new sample.

        Parameters
        ----------
        sample: dict, or list of dictionaries
            A dictionary containing all the keys from
            self.search_parameter_keys and corresponding values at which to
            calculate the posterior probability

        Returns
        -------
        p: array-like,
            The posterior probability of the sample

        """
        if isinstance(sample, dict):
            sample = [sample]
        ordered_sample = [[s[key] for key in self.search_parameter_keys]
                          for s in sample]
        return self.kde(ordered_sample)

def plot_multiple(results, filename=None, labels=None, colours=None,
                  save=True, evidences=False, **kwargs):
    """ Generate a corner plot overlaying two sets of results

    Parameters
    ----------
    results: list
        A list of `bilby.core.result.Result` objects containing the samples to
        plot.
    filename: str
        File name to save the figure to. If None (default), a filename is
        constructed from the outdir of the first element of results and then
        the labels for all the result files.
    labels: list
        List of strings to use when generating a legend. If None (default), the
        `label` attribute of each result in `results` is used.
    colours: list
        The colours for each result. If None, default styles are applied.
    save: bool
        If true, save the figure
    kwargs: dict
        All other keyword arguments are passed to `result.plot_corner`.
        However, `show_titles` and `truths` are ignored since they would be
        ambiguous on such a plot.
    evidences: bool, optional
        Add the log-evidence calculations to the legend. If available, the
        Bayes factor will be used instead.

    Returns
    -------
    fig:
        A matplotlib figure instance

    """

    kwargs['show_titles'] = False
    kwargs['truths'] = None

    fig = results[0].plot_corner(save=False, **kwargs)
    default_filename = '{}/{}'.format(results[0].outdir, 'combined')
    lines = []
    default_labels = []
    for i, result in enumerate(results):
        if colours:
            c = colours[i]
        else:
            c = 'C{}'.format(i)
        hist_kwargs = kwargs.get('hist_kwargs', dict())
        hist_kwargs['color'] = c
        fig = result.plot_corner(fig=fig, save=False, color=c, **kwargs)
        default_filename += '_{}'.format(result.label)
        lines.append(matplotlib.lines.Line2D([0], [0], color=c))
        default_labels.append(result.label)

    # Rescale the axes
    for i, ax in enumerate(fig.axes):
        ax.autoscale()
    plt.draw()

    if labels is None:
        labels = default_labels

    if evidences:
        if np.isnan(results[0].log_bayes_factor):
            template = ' $\mathrm{{ln}}(Z)={lnz:1.3g}$'
        else:
            template = ' $\mathrm{{ln}}(B)={lnbf:1.3g}$'
        labels = [template.format(lnz=result.log_evidence,
                                  lnbf=result.log_bayes_factor)
                  for ii, result in enumerate(results)]

    axes = fig.get_axes()
    ndim = int(np.sqrt(len(axes)))
    axes[ndim - 1].legend(lines, labels)

    if filename is None:
        filename = default_filename

    if save:
        fig.savefig(filename)
    return fig

def make_pp_plot(results, filename=None, save=True, **kwargs):
    """
    Make a P-P plot for a set of runs with injected signals.

    Parameters
    ----------
    results: list
        A list of Result objects, each of these should have injected_parameters
    filename: str, optional
        The name of the file to save, the default is "outdir/pp.png"
    save: bool, optional
        Whether to save the file, default=True
    kwargs:
        Additional kwargs to pass to matplotlib.pyplot.plot

    Returns
    -------
    fig:
        Matplotlib figure
    """
    fig = plt.figure()
    credible_levels = pd.DataFrame()
    for result in results:
        credible_levels = credible_levels.append(
            result.get_all_injection_credible_levels(), ignore_index=True)
    n_parameters = len(credible_levels.keys())
    x_values = np.linspace(0, 1, 101)
    for key in credible_levels:
        plt.plot(x_values, [sum(credible_levels[key].values < xx) /
                            len(credible_levels) for xx in x_values],
                 color='k', alpha=min([1, 4 / n_parameters]), **kwargs)
    plt.plot([0, 1], [0, 1], linestyle='--', color='r')
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.tight_layout()
    if save:
        if filename is None:
            filename = 'outdir/pp.png'
        plt.savefig(filename)
    return fig


# Copyright (c) 2021-present, Facebook, Inc.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#

import logging
import math
from copy import deepcopy
from types import SimpleNamespace
from typing import Any, Dict, List, Optional, Tuple

import gym
import hydra
import numpy as np
import torch as th
from omegaconf import DictConfig
from torch import nn
from torch.nn import functional as F

from hucc import ReplayBuffer
from hucc.agents import Agent

log = logging.getLogger(__name__)

class DIAYNAgent(Agent):
    '''
    Diversity is All You Need with Soft-Actor Critic.
    '''

    def __init__(
        self,
        env: gym.Env,
        model: nn.Module,
        optim: SimpleNamespace,
        cfg: DictConfig,
    ):
        super().__init__(cfg)
        if not hasattr(model, 'pi'):
            raise ValueError('Model needs "pi" module')
        if not hasattr(model, 'q'):
            raise ValueError('Model needs "q" module')
        if not hasattr(model, 'phi'):
            raise ValueError('Model needs "phi" module')

        self._model = model
        self._optim = optim
        self._n_skills = int(cfg.n_skills)
        self._bsz = int(cfg.batch_size)
        self._gamma = float(cfg.gamma)
        self._polyak = float(cfg.polyak)
        self._rpbuf_size = int(cfg.rpbuf_size)
        self._samples_per_update = int(cfg.samples_per_update)
        self._num_updates = int(cfg.num_updates)
        self._warmup_samples = int(cfg.warmup_samples)
        self._phi_obs = str(cfg.phi_obs)
        self._phi_obs_feats: Optional[List[int]] = None
        if cfg.phi_obs_feats is not None:
            self._phi_obs_feats = list(
                map(int, str(cfg.phi_obs_feats).split('#'))
            )
        elif cfg.phi_obs_n > 0:
            self._phi_obs_feats = list(range(cfg.phi_obs_n))
        self._add_p_z = bool(cfg.add_p_z)

        self._target_entropy = -np.prod(env.action_space.shape)
        # Optimize log(alpha) so that we'll always have a positive factor
        log_alpha = np.log(cfg.alpha)
        if cfg.optim_alpha is None:
            self._log_alpha = th.tensor(log_alpha)
            self._optim_alpha = None
        else:
            self._log_alpha = th.tensor(log_alpha, requires_grad=True)
            self._optim_alpha = hydra.utils.instantiate(
                cfg.optim_alpha, [self._log_alpha]
            )

        self._log_p_z = th.tensor(1 / self._n_skills).log().to(env.device)

        self._buffer = ReplayBuffer(
            size=self._rpbuf_size, interleave=env.num_envs
        )
        self._n_samples_since_update = 0

        self._target = deepcopy(model)
        # We'll never need gradients for the target network
        for param in self._target.parameters():
            param.requires_grad_(False)

        self._action_space = env.action_space
        self._action_factor = env.action_space.high[0]

        self.set_checkpoint_attr(
            '_model',
            '_target',
            '_optim',
            '_log_alpha',
            '_optim_alpha',
        )

        assert isinstance(
            env.action_space, gym.spaces.Box
        ), f'DIAYNAgent requires a continuous (Box) action space (but got {type(env.action_space)})'

    @staticmethod
    def effective_observation_space(env: gym.Env, cfg: DictConfig) -> gym.Space:
        obs_space = env.observation_space['observation']
        n_skills = int(cfg.n_skills)
        # Add one-hot encoding of current skill to all observations
        return gym.spaces.Box(
            low=np.concatenate([obs_space.low, np.zeros(n_skills)]),
            high=np.concatenate([obs_space.high, np.ones(n_skills)]),
        )

    def action(self, env, obs) -> Tuple[th.Tensor, Any]:
        with th.no_grad():
            prev_z = env.ctx.get(
                'z', th.zeros((env.num_envs,), dtype=th.long, device=env.device)
            )
            # Sample a new skill at the beginning of an episode
            sample_z = (obs['time'] == 0).squeeze(1)
            if sample_z.any():
                if self.training:
                    sample = th.randint(
                        0,
                        self._n_skills,
                        size=(env.num_envs,),
                        device=env.device,
                    )
                else:
                    # For evaluation, just enumerate all skills
                    sample = th.cat(
                        [th.arange(0, self._n_skills)]
                        * math.ceil(env.num_envs / self._n_skills)
                    )[: env.num_envs].to(env.device)
                z = (th.logical_not(sample_z)) * prev_z + (sample_z) * sample
            else:
                z = prev_z
            env.ctx['z'] = z

            aug_obs = th.cat(
                [obs['observation'], F.one_hot(z, self._n_skills).float()],
                dim=-1,
            )
            dist = self._model.pi(aug_obs)
            assert (
                dist.has_rsample
            ), f'rsample() required for policy distribution'
            if self.training:
                action = dist.sample() * self._action_factor
            else:
                action = dist.mean * self._action_factor
        return action, {'z': z}

    def step(
        self,
        env,
        obs,
        action: th.Tensor,
        extra: Any,
        result: Tuple[th.Tensor, th.Tensor, th.Tensor, List[Dict]],
    ) -> None:
        oobs = obs['observation']
        z = extra['z']
        next_obs, reward, done, info = result
        next_oobs = next_obs['observation']

        # Ignore terminal state if we have a timeout
        for i in range(len(info)):
            if 'TimeLimit.truncated' in info[i]:
                done[i] = False

        batch = dict(
            obs=oobs,
            z=z,
            next_obs=next_oobs,
            action=action.squeeze(-1),
            terminal=done,
        )
        if self._phi_obs != 'observation':
            batch['phi_obs'] = obs[self._phi_obs]
        self._buffer.put_row(batch)

        self._n_steps += 1
        self._n_samples += done.nelement()
        self._n_samples_since_update += done.nelement()
        if self._buffer.size < self._warmup_samples:
            return
        if self._n_samples_since_update >= self._samples_per_update:
            self.update()
            self._n_samples_since_update = 0

    def _update(self):
        def act_logp(obs):
            dist = self._model.pi(obs)
            action = dist.rsample()
            log_prob = dist.log_prob(action).sum(dim=-1)
            action = action * self._action_factor
            return action, log_prob

        rewards = []
        for _ in range(self._num_updates):
            batch = self._buffer.get_batch(self._bsz)
            # Ensure that action has a batch dimension
            action = batch['action'].view(batch['obs'].shape[0], -1)
            z = batch['z']
            z_one_hot = F.one_hot(z, self._n_skills).float()
            phi_obs = batch.get('phi_obs', batch['obs'])
            if self._phi_obs_feats is not None:
                phi_obs = phi_obs[:, self._phi_obs_feats]
            not_done = th.logical_not(batch['terminal'])

            # Compute pseudo-reward with discriminator
            with th.no_grad():
                reward = -F.cross_entropy(
                    self._model.phi(phi_obs), z, reduction='none'
                )
                # Subtract baseline
                if self._add_p_z:
                    reward = reward - self._log_p_z
                rewards.append(reward.mean().item())

            # Backup for Q-Function
            with th.no_grad():
                obs_p = th.cat([batch['next_obs'], z_one_hot], dim=1)
                a_p, log_prob_p = act_logp(obs_p)

                q_in = th.cat([obs_p, a_p], dim=1)
                q_tgt = th.min(self._target.q(q_in), dim=-1).values
                backup = reward + self._gamma * not_done * (
                    q_tgt - self._log_alpha.detach().exp() * log_prob_p
                )

            # Q-Function update
            obs = th.cat([batch['obs'], z_one_hot], dim=1)
            q_in = th.cat([obs, action], dim=1)
            q = self._model.q(q_in)
            q1 = q[:, 0]
            q2 = q[:, 1]
            q1_loss = F.mse_loss(q1, backup)
            q2_loss = F.mse_loss(q2, backup)
            q_loss = q1_loss + q2_loss
            self._optim.q.zero_grad()
            q_loss.backward()
            self._optim.q.step()

            # Policy update
            for param in self._model.q.parameters():
                param.requires_grad_(False)

            a, log_prob = act_logp(obs)
            q_in = th.cat([obs, a], dim=1)
            q = th.min(self._model.q(q_in), dim=-1).values
            pi_loss = (self._log_alpha.detach().exp() * log_prob - q).mean()
            self._optim.pi.zero_grad()
            pi_loss.backward()
            self._optim.pi.step()

            for param in self._model.q.parameters():
                param.requires_grad_(True)

            # Optional temperature update
            if self._optim_alpha:
                # This is slight reording of the formulation in
                # https://github.com/rail-berkeley/softlearning, mostly so we
                # don't need to create temporary tensors. log_prob is the only
                # non-scalar tensor, so we can compute its mean first.
                alpha_loss = -(
                    self._log_alpha.exp()
                    * (log_prob.mean().cpu() + self._target_entropy).detach()
                )
                self._optim_alpha.zero_grad()
                alpha_loss.backward()
                self._optim_alpha.step()

            # Update discriminator
            self._optim.phi.zero_grad()
            phi_loss = F.cross_entropy(self._model.phi(phi_obs), z)
            phi_loss.backward()
            self._optim.phi.step()

            # Update target network
            with th.no_grad():
                for tp, p in zip(
                    self._target.q.parameters(), self._model.q.parameters()
                ):
                    tp.data.lerp_(p.data, 1.0 - self._polyak)

        self.tbw_add_scalar('Loss/Policy', pi_loss.item())
        self.tbw_add_scalar('Loss/QValue', q_loss.item())
        self.tbw_add_scalar('Loss/Discriminator', phi_loss.item())
        self.tbw_add_scalar('Avg Reward', np.mean(rewards))
        self.tbw_add_scalar('Health/Entropy', log_prob.mean().item())
        if self._optim_alpha:
            self.tbw_add_scalar('Health/Alpha', self._log_alpha.exp().item())

        msg = log.debug
        if (self._n_updates * self._num_updates) % 50 == 0:
            msg = log.info
        msg(
            f'Sample {self._n_samples}, up {self._n_updates*self._num_updates}, pi loss {pi_loss.item():+.03f}, q loss {q_loss.item():+.03f}, phi loss {phi_loss.item():+.03f}, avg reward {np.mean(rewards):+.03}, alpha {self._log_alpha.exp().item():.03f}'
        )


#!/usr/bin/env python
from astropy.io import fits
import argparse
import numpy as np
from astropy.table import Table

parser = argparse.ArgumentParser(description='Verify foreground/background quasars and measure absorbing gas around the foreground in the background quasar spectrum')
parser.add_argument('-m', metavar='mask name', type=str, help='name of the mask to be redshifted', required=True)
args = parser.parse_args()
objects = Table.read('{0}_spec1D/{0}_objects.fits'.format(args.m))
objects=objects[~objects['alignbox']]

ntot=len(objects)
n2=sum(objects['quality']==2)
n2gal=sum((objects['quality']==2) & ((objects['class']=='galaxy') | (objects['class']=='quasar')))
n2star=sum((objects['quality']==2) & (objects['class']=='star'))

n1=sum(objects['quality']==1)
n0=sum(objects['quality']==0)
nf=sum(objects['quality']==-1)

print '{:d} targets\n{:d} Q=2, {:d} galaxies, {:d} stars\n{:d} single line\n{:d} No redshift\n{:d} Bad spectrum'.format(ntot,n2,n2gal,n2star,n1,n0,nf)


# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import logging
import numpy as np
from typing import Dict, List
import torch

from cvpods.layers import ShapeSpec
from cvpods.structures import Boxes, Instances, pairwise_iou, pairwise_ioa
from cvpods.utils import get_event_storage

from cvpods.modeling.matcher import MatcherIgnore
from cvpods.modeling.roi_heads import StandardROIHeads

from cvpods.modeling.proposal_generator.proposal_utils import add_ground_truth_to_proposals

logger = logging.getLogger(__name__)

class StandardROIHeadsIgnore(StandardROIHeads):
    """
    It's "standard" in a sense that there is no ROI transform sharing
    or feature sharing between tasks.
    The cropped rois go to separate branches (boxes and masks) directly.
    This way, it is easier to make separate abstractions for different branches.

    This class is used by most models, such as FPN and C5.
    To implement more models, you can subclass it and implement a different
    :meth:`forward()` or a head.
    """

    def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
        super(StandardROIHeadsIgnore, self).__init__(cfg, input_shape)

        # Matcher to assign box proposals to gt boxes
        self.proposal_matcher = MatcherIgnore(
            cfg.MODEL.ROI_HEADS.IOU_THRESHOLDS,
            cfg.MODEL.ROI_HEADS.IOU_LABELS,
            allow_low_quality_matches=False,
        )

    @torch.no_grad()
    def label_and_sample_proposals(
        self, proposals: List[Instances], targets: List[Instances]
    ) -> List[Instances]:
        """
        Prepare some proposals to be used to train the ROI heads.
        It performs box matching between `proposals` and `targets`, and assigns
        training labels to the proposals.
        It returns ``self.batch_size_per_image`` random samples from proposals and groundtruth
        boxes, with a fraction of positives that is no larger than
        ``self.positive_sample_fraction``.

        Args:
            See :meth:`ROIHeads.forward`

        Returns:
            list[Instances]:
                length `N` list of `Instances`s containing the proposals
                sampled for training. Each `Instances` has the following fields:

                - proposal_boxes: the proposal boxes
                - gt_boxes: the ground-truth box that the proposal is assigned to
                  (this is only meaningful if the proposal has a label > 0; if label = 0
                  then the ground-truth box is random)

                Other fields such as "gt_classes", "gt_masks", that's included in `targets`.
        """
        gt_boxes = [x.gt_boxes for x in targets]
        # Augment proposals with ground-truth boxes.
        # In the case of learned proposals (e.g., RPN), when training starts
        # the proposals will be low quality due to random initialization.
        # It's possible that none of these initial
        # proposals have high enough overlap with the gt objects to be used
        # as positive examples for the second stage components (box head,
        # cls head, mask head). Adding the gt boxes to the set of proposals
        # ensures that the second stage components will have some positive
        # examples from the start of training. For RPN, this augmentation improves
        # convergence and empirically improves box AP on COCO by about 0.5
        # points (under one tested configuration).
        if self.proposal_append_gt:
            proposals = add_ground_truth_to_proposals(gt_boxes, proposals)

        proposals_with_gt = []

        num_fg_samples = []
        num_bg_samples = []
        for proposals_per_image, targets_per_image in zip(proposals, targets):
            has_gt = len(targets_per_image) > 0
            match_quality_matrix = pairwise_iou(
                targets_per_image.gt_boxes, proposals_per_image.proposal_boxes
            )

            match_quality_ignore = pairwise_ioa(
                targets_per_image.gt_boxes, proposals_per_image.proposal_boxes,
                targets_per_image.gt_classes
            )

            matched_idxs, matched_labels = self.proposal_matcher(
                match_quality_matrix, match_quality_ignore)
            sampled_idxs, gt_classes = self._sample_proposals(
                matched_idxs, matched_labels, targets_per_image.gt_classes
            )

            # Set target attributes of the sampled proposals:
            proposals_per_image = proposals_per_image[sampled_idxs]
            proposals_per_image.gt_classes = gt_classes

            # We index all the attributes of targets that start with "gt_"
            # and have not been added to proposals yet (="gt_classes").
            if has_gt:
                sampled_targets = matched_idxs[sampled_idxs]
                # NOTE: here the indexing waste some compute, because heads
                # like masks, keypoints, etc, will filter the proposals again,
                # (by foreground/background, or number of keypoints in the image, etc)
                # so we essentially index the data twice.
                for (trg_name, trg_value) in targets_per_image.get_fields().items():
                    if trg_name.startswith("gt_") and not proposals_per_image.has(trg_name):
                        proposals_per_image.set(trg_name, trg_value[sampled_targets])
            else:
                gt_boxes = Boxes(
                    targets_per_image.gt_boxes.tensor.new_zeros((len(sampled_idxs), 4))
                )
                proposals_per_image.gt_boxes = gt_boxes

            num_bg_samples.append((gt_classes == self.num_classes).sum().item())
            num_fg_samples.append(gt_classes.numel() - num_bg_samples[-1])
            proposals_with_gt.append(proposals_per_image)

        # Log the number of fg/bg samples that are selected for training ROI heads
        storage = get_event_storage()
        storage.put_scalar("roi_head/num_fg_samples", np.mean(num_fg_samples))
        storage.put_scalar("roi_head/num_bg_samples", np.mean(num_bg_samples))

        return proposals_with_gt


# OpenCV
import cv2 as cv
import numpy as np

import os

from screenshot import *
from functools import cmp_to_key
from typing import Tuple, List

from utils import waitForCycles, waitForNextCycle
import pyautogui as gui
from config import GAME_LOCATION, GAME_SCREEN_SIZE_16_9, GAME_SCREEN_SIZE_2_1, GAME_SCREEN_SIZE_3_2, \
    SCREEN_SCALE, SORT_OFFSET, ONE_CYCLE

left, top, end_left, end_top = GAME_LOCATION
width = (end_left - left) * SCREEN_SCALE
height = (end_top - top) * SCREEN_SCALE
print("=> Game Size is ({}, {})".format(width, height))
monitor = {"top": top * SCREEN_SCALE, "left": left * SCREEN_SCALE, "width": width, "height": height}

# determine best input size
ratio = height / width
INPUT_SIZE = 0
ACCURACY = 0.8
if ratio < 1.6:
    INPUT_SIZE = GAME_SCREEN_SIZE_3_2
elif ratio < 1.85:
    INPUT_SIZE = GAME_SCREEN_SIZE_16_9
else:
    # NOTE: all templates are based on 2 : 1 game ratio
    # ACCURACY = 0.9
    INPUT_SIZE = GAME_SCREEN_SIZE_2_1
print("=> Ratio is {:.3f}. Resize to {}. Accuracy is set to {}".format(ratio, INPUT_SIZE, ACCURACY))

# track original size to go back to extract point
width_scale = width / INPUT_SIZE[0]
height_scale = height / INPUT_SIZE[1]
print("=> Scales: Width - {:.3f}, Height - {:.3f}".format(width_scale, height_scale))

def find(template: str, img) -> Tuple[bool, Tuple[int, int]]:
    """
    Find template in img and tap on it if needed.
    img doesn't need to be in gray scale because it will be converted in this function.
    """
    result = (False, ())
    if os.path.exists(template):
        template_img = cv.imread(template, 0)
        # print("👀 Looing for '{}'".format(template))
        # showImage(template_img, "template")

        # NOTE: resize is very important for the detection to be consistent
        img = cv.resize(img, INPUT_SIZE)
        # showImage(img, "Resize")

        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        matches = cv.matchTemplate(gray, template_img, cv.TM_CCOEFF_NORMED)
        # showImage(matches, "Raw")

        # Check if anything matches with the template
        w, h = template_img.shape[::-1]
        locations = __sorted_matches(matches, ACCURACY, offset=SORT_OFFSET)
        for pt in locations:
            # draw rectangles around all locations
            result = (True, ())
            cv.rectangle(img, pt, (pt[0] + w, pt[1] + h), (255, 0, 255), 3)

        if result[0]:
            # only consider the first one
            x, y = locations[0]
            # print("=> Original ({}, {})".format(x, y))
            # NOTE: consider the scale change when we resize the image
            x *= width_scale
            y *= height_scale
            # print("=> Scaled ({}, {})".format(x, y))
            # print("=> w {}, h {}".format(w, h))

            # need to scale x, y to original and get center point
            # NOTE: consider template scale as well
            # TODO: make sure this still works on other machines
            x += GAME_LOCATION[0] +  w * width_scale / 2
            y += GAME_LOCATION[1] + h * height_scale / 2

            result = (True, (x, y))
            print("=> ({}, {})".format(x, y))
            # resized_img = cv.resize(img, (300, 570))
            # showImage(resized_img, "Matches")
    else:
        exit("Cannot find template at {}".format(template))
    return result

def tap(template: str, img, delay=1) -> bool:
    """
    Find template from img and tap on it
    """

    # location can be () here so only get x, y if found
    found, location = find(template, img)
    if found:
        x, y = location
        gui.moveTo(x, y)
        gui.mouseDown()
        gui.mouseUp()

        waitForCycles(delay)
    return found

def swipe_up(up=True):
    # find center point
    if up:
        offset = height / 3.3
    else:
        offset = - height / 3.3

    x, y = width / 2 + left, height / 2 + top
    gui.moveTo(x, y - offset)
    gui.mouseDown()
    gui.moveTo(x, y + offset, duration=0.5)
    gui.mouseUp()

def swipe_down():
    swipe_up(False)

def touch():
    """
    Tap the center point of Game Location
    """
    x = width * 0.5 + left
    y = height * 0.75 + top
    gui.moveTo(x, y)
    gui.mouseDown()
    gui.mouseUp()

def __sorted_matches(matches, threshold: int, offset=25, allow_duplicates=False) -> list:
    """
    Filter out matches base on the threshold.
    Return a list of sorted and filtered matches
    """
    output = []
    curr_pt = None

    filtered_matches = np.where(matches >= threshold)
    for pt in sorted(zip(*filtered_matches[::-1]), key=cmp_to_key(lambda a, b : __sort_element(a, b, offset))):
        if not allow_duplicates:
            # ignore points which are close to each other
            if curr_pt is None:
                curr_pt = pt
            elif abs(curr_pt[0] - pt[0]) < offset and abs(curr_pt[1] - pt[1]) < offset:
                continue
            else:
                curr_pt = pt

        output.append(pt)
    return output

def __sort_element(a, b, offset: int) -> bool:
    """
    Sort elements based on the offset
    """
    # 140, 141, 142 are treated as one point here
    if abs(a[0] - b[0]) < offset:
        return a[1] - b[1]
    return a[0] - b[0]

def showImage(img, name="image"):
    """
    Show an image using opencv with title
    """
    cv.imshow(name, img)
    cv.waitKey()
    cv.destroyWindow(name)

def __testFind():
    """
    Debug only
    """
    left, top, end_left, end_top = GAME_LOCATION
    width = (end_left - left) * SCREEN_SCALE
    height = (end_top - top) * SCREEN_SCALE
    print("=> Game Size is ({}, {})".format(width, height))

    monitor = {"top": top * SCREEN_SCALE, "left": left * SCREEN_SCALE, "width": width, "height": height}
    game_img = np.array(take_screenshot(monitor))

    # print(find(u"game/dungeons/tamadora.png", game_img))
    # print(tap(u"game/dungeons/mugen-kairou.png", game_img))
    # print(find(u"game/sell.png", game_img))
    # print(tap(u"game/buttons/special.png", game_img))
    # print(tap(u"game/home.png", game_img))

    # go to the required dungeon
    tap(u"game/dungeons/kairou/sub1.png", game_img)

def tapInOrder(instructions: List[str]) -> bool:
    """
    Find templates in order with delay
    """

    result = False

    # at least, try all of them and see if one if successful
    for template in instructions:
        # NOTE: move the cursor outside the game screen so that it doesn't cover up the screen
        gui.moveTo(end_left + 10, height / 2)

        game_img = np.array(take_screenshot(monitor))
        success = tap(template, game_img)
        if success:
            result = True

    # for template in instructions:
    #     while not success:

        # if not success:
        #     print("❌ Couldn't find '{}'".format(template))
        #     return False

    return True

def __mugen_loop():
    """
    Go to mugen kairou, quit and repeat
    """
    while True:
        # test join and quit
        tapInOrder([
            # "game/dungeons/kairou/main.png",
            "game/dungeons/kairou/sub1.png",
            "game/friends/you.png",
            "game/buttons/challenge.png",
        ])

        waitForCycles(16)

        # Quit any battles
        tapInOrder([
            "game/battle/menu.png",
            "game/buttons/quit_battle.png",
            "game/buttons/yes.png",
        ])

        waitForCycles(6)

def __battle_counter():
    """
    Count battles, boss battle and also when exiting the dungeon
    """
    battle = 0
    boss = False
    while True:
        game_img = np.array(take_screenshot(monitor))
        # if find(u"game/battle/empty1.png", game_img)[0] or find(u"game/battle/empty2.png", game_img)[0]:
        #     print("=> Falling")
        if not boss:
            if find(u"game/battle/battle_number.png", game_img)[0]:
                battle += 1
                print("=> Battle {}".format(battle))
                # wait for one extra cycle to prevent counting twice
                waitForCycles(2)
            elif find(u"game/battle/boss_alert.png", game_img)[0]:
                battle += 1
                boss = True
                print("=> Boss ({} / {})".format(battle, battle))

        if tap(u"game/buttons/ok.png", game_img):
            print("=> End")
        waitForNextCycle()

def __get_resized_screenshot():
    screenshot_img = np.array(take_screenshot(monitor))
    showImage(cv.resize(screenshot_img, INPUT_SIZE), "Screenshot")

if __name__ == "__main__":
    # __get_resized_screenshot()
    # swipe_up()
    # swipe_down()
    # __testFind()

    # __mugen_loop()
    __get_resized_screenshot()
    # __battle_counter()


import cv2
import numpy as np
import dlib
from math import hypot

# Loading Camera and Nose image and Creating mask
cap = cv2.VideoCapture(0)
nose_image = cv2.imread("mask/bat.png")
_, frame = cap.read()
rows, cols, _ = frame.shape
nose_mask = np.zeros((rows, cols), np.uint8)

# Loading Face detector
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("resources/shape_predictor_68_face_landmarks.dat")

while True:
    _, frame = cap.read()
    nose_mask.fill(0)
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    faces = detector(frame)
    for face in faces:
        landmarks = predictor(gray_frame, face)

        # Nose coordinates
        top_nose = (landmarks.part(29).x, landmarks.part(29).y)
        center_nose = (landmarks.part(30).x, landmarks.part(30).y-85)
        left_nose = (landmarks.part(31).x, landmarks.part(31).y)
        right_nose = (landmarks.part(35).x, landmarks.part(35).y)

        nose_width = int(hypot(left_nose[0] - right_nose[0],
                           left_nose[1] - right_nose[1]) * 7)
        nose_height = int(nose_width * 1.6)

        # New nose position
        top_left = (int(center_nose[0] - nose_width / 2),
                              int(center_nose[1] - nose_height / 2))
        bottom_right = (int(center_nose[0] + nose_width / 2),
                       int(center_nose[1] + nose_height / 2))

        # Adding the new nose
        nose_pig = cv2.resize(nose_image, (nose_width, nose_height))
        nose_pig_gray = cv2.cvtColor(nose_pig, cv2.COLOR_BGR2GRAY)
        _, nose_mask = cv2.threshold(nose_pig_gray, 25, 255, cv2.THRESH_BINARY_INV)

        nose_area = frame[top_left[1]: top_left[1] + nose_height,
                    top_left[0]: top_left[0] + nose_width]
        nose_area_no_nose = cv2.bitwise_and(nose_area, nose_area, mask=nose_mask)
        final_nose = cv2.add(nose_area_no_nose, nose_pig)

        frame[top_left[1]: top_left[1] + nose_height,
                    top_left[0]: top_left[0] + nose_width] = final_nose

    cv2.imshow("Frame", frame)

    key = cv2.waitKey(1)
    if key == 27:
        break


import argparse
import os
import re
from random import choice

import cv2 as cv
import numpy as np
from PIL import Image, ImageOps
from tensorflow.keras import models

from utils import decoder

parser = argparse.ArgumentParser(description='Predict number-plate using the CNN model')
parser.add_argument('-i', '--image', dest='image', action='store', default='assets/test.png', help='input image path')
parser.add_argument('-m', '--model', dest='model', action='store', default='models/anpr-300k.h5', help='trained model path')
parser.add_argument('-r', '--random', dest='random', action='store', default=None, help='use random examples from given dataset')

arguments = parser.parse_args()

if os.path.isfile(arguments.model):
  model = models.load_model(arguments.model)

  if arguments.random:
    if not os.path.isdir(arguments.random):
      print(f'Dataset was not found at {arguments.random}')
      exit()
    files = list(os.scandir(arguments.random))
    pattern = re.compile("\d+_(.+)\.png")

    while True:
      file = choice(files)

      mat = cv.imread(file.path, cv.IMREAD_GRAYSCALE)
      image = np.array(Image.open(file.path), dtype=np.uint8).reshape((128, 64, 1)) / 255

      label = pattern.search(file.name).group(1)
      predictions = model.predict(np.array([image]))

      cv.imshow('Input', mat)
      print('Ground truth vs prediction')
      print(label)
      print(decoder(predictions[0]))

      if cv.waitKey() == 27:
        break
  else:
    if os.path.isfile(arguments.image):
      image_pil = ImageOps.grayscale(Image.open(arguments.image).resize((128, 64)))
      image = np.array(image_pil, dtype=np.uint8).reshape((128, 64, 1)) / 255

      predictions = model.predict(np.array([image]))

      print(decoder(predictions[0]))
    else:
      print(f'Input image was not found at {arguments.image}')
else:
  print(f'Model was not found at {arguments.model}')


import pyaudio
import librosa
import numpy as np
import requests

# ring buffer will keep the last 2 seconds worth of audio
ringBuffer = RingBuffer(2 * 22050)

def callback(in_data, frame_count, time_info, flag):
    audio_data = np.fromstring(in_data, dtype=np.float32)

    # we trained on audio with a sample rate of 22050 so we need to convert it
    audio_data = librosa.resample(audio_data, 44100, 22050)
    ringBuffer.extend(audio_data)

    # machine learning model takes wavform as input and
    # decides if the last 2 seconds of audio contains a goal
    if model.is_goal(ringBuffer.get()):
        # GOAL!! Trigger light show
        requests.get("http://127.0.0.1:8082/goal")

    return (in_data, pyaudio.paContinue)

# function that finds the index of the Soundflower
# input device and HDMI output device
dev_indexes = findAudioDevices()

stream = pa.open(format = pyaudio.paFloat32,
                 channels = 1,
                 rate = 44100,
                 output = True,
                 input = True,
                 input_device_index = dev_indexes['input'],
                 output_device_index = dev_indexes['output'],
                 stream_callback = callback)

# start the stream
stream.start_stream()

while stream.is_active():
    sleep(0.25)

stream.close()
pa.terminate()


# 摄像头实时人脸识别

import dlib          # 人脸处理的库 Dlib
import numpy as np   # 数据处理的库 numpy
import cv2           # 图像处理的库 OpenCv
import pandas as pd  # 数据处理的库 Pandas
from PIL import ImageFont
NAME_LIST = []

with open('labels.txt', 'r', encoding='utf-8') as fw:
    for i in NAME_LIST:
        fw.write(i + '\n')

NAME_LIST = [line.strip() for line in open('labels.txt', 'r', encoding='utf-8').readlines()]
# face recognition model, the object maps human faces into 128D vectors
# Refer this tutorial: http://dlib.net/python/index.html#dlib.face_recognition_model_v1
facerec = dlib.face_recognition_model_v1("data/data_dlib/dlib_face_recognition_resnet_model_v1.dat")

# 计算两个128D向量间的欧式距离
# compute the e-distance between two 128D features
def return_euclidean_distance(feature_1, feature_2):
    feature_1 = np.array(feature_1)
    feature_2 = np.array(feature_2)
    dist = np.sqrt(np.sum(np.square(feature_1 - feature_2)))
    return dist

# 处理存放所有人脸特征的 csv
path_features_known_csv = "data/features_all.csv"
csv_rd = pd.read_csv(path_features_known_csv, header=None)

# 用来存放所有录入人脸特征的数组
# the array to save the features of faces in the database
features_known_arr = []

# 读取已知人脸数据
# print known faces
for i in range(csv_rd.shape[0]):
    features_someone_arr = []
    for j in range(0, len(csv_rd.ix[i, :])):
        features_someone_arr.append(csv_rd.ix[i, :][j])
    features_known_arr.append(features_someone_arr)
print("Faces in Database：", len(features_known_arr))

# Dlib 检测器和预测器
# The detector and predictor will be used
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('data/data_dlib/shape_predictor_68_face_landmarks.dat')

# 创建 cv2 摄像头对象
# cv2.VideoCapture(0) to use the default camera of PC,
# and you can use local video name by use cv2.VideoCapture(filename)
cap = cv2.VideoCapture(0)

# cap.set(propId, value)
# 设置视频参数，propId 设置的视频参数，value 设置的参数值
cap.set(3, 480)

# cap.isOpened() 返回 true/false 检查初始化是否成功
# when the camera is open
while cap.isOpened():

    flag, img_rd = cap.read()
    kk = cv2.waitKey(1)

    # 取灰度
    img_gray = cv2.cvtColor(img_rd, cv2.COLOR_RGB2GRAY)

    # 人脸数 faces
    faces = detector(img_gray, 0)

    # 待会要写的字体 font to write later
    font = cv2.FONT_HERSHEY_COMPLEX

    # 存储当前摄像头中捕获到的所有人脸的坐标/名字
    # the list to save the positions and names of current faces captured
    pos_namelist = []
    name_namelist = []

    # 按下 q 键退出
    # press 'q' to exit
    if kk == ord('q'):
        break
    else:
        # 检测到人脸 when face detected
        if len(faces) != 0:
            # 获取当前捕获到的图像的所有人脸的特征，存储到 features_cap_arr
            # get the features captured and save into features_cap_arr
            features_cap_arr = []
            for i in range(len(faces)):
                shape = predictor(img_rd, faces[i])
                features_cap_arr.append(facerec.compute_face_descriptor(img_rd, shape))

            # 遍历捕获到的图像中所有的人脸
            # traversal all the faces in the database
            for k in range(len(faces)):
                print("##### camera person", k+1, "#####")
                # 让人名跟随在矩形框的下方
                # 确定人名的位置坐标
                # 先默认所有人不认识，是 unknown
                # set the default names of faces with "unknown"
                name_namelist.append("unknown")

                # 每个捕获人脸的名字坐标 the positions of faces captured
                pos_namelist.append(tuple([faces[k].left(), int(faces[k].bottom() + (faces[k].bottom() - faces[k].top())/4)]))

                # 对于某张人脸，遍历所有存储的人脸特征
                # for every faces detected, compare the faces in the database
                e_distance_list = []
                for i in range(len(features_known_arr)):
                    # 如果 person_X 数据不为空
                    if str(features_known_arr[i][0]) != '0.0':
                        print("with person", str(i + 1), "the e distance: ", end='')
                        e_distance_tmp = return_euclidean_distance(features_cap_arr[k], features_known_arr[i])
                        print(e_distance_tmp)
                        e_distance_list.append(e_distance_tmp)
                    else:
                        # 空数据 person_X
                        e_distance_list.append(999999999)
                # Find the one with minimum e distance
                similar_person_num = e_distance_list.index(min(e_distance_list))
                print("Minimum e distance with person", int(similar_person_num)+1)

                if min(e_distance_list) < 0.4:
                    # 在这里修改 person_1, person_2 ... 的名字
                    # 可以在这里改称 Jack, Tom and others
                    # Here you can modify the names shown on the camera
                    name_namelist[k] = NAME_LIST[int(similar_person_num)]

                    print("May be person "+str(int(similar_person_num)+1))
                else:
                    print("Unknown person")

                # 矩形框
                # draw rectangle
                for kk, d in enumerate(faces):
                    # 绘制矩形框
                    cv2.rectangle(img_rd, tuple([d.left(), d.top()]), tuple([d.right(), d.bottom()]), (0, 255, 255), 2)
                print('\n')

            # 在人脸框下面写人脸名字
            # write names under rectangle

            for i in range(len(faces)):
                from PIL import Image, ImageDraw, ImageFont
                cv2img = cv2.cvtColor(img_rd, cv2.COLOR_BGR2RGB)
                pilimg = Image.fromarray(cv2img)

                draw = ImageDraw.Draw(pilimg)
                font1 = ImageFont.truetype('simhei.ttf', 30, encoding='utf-8')
                draw.text( pos_namelist[i], str(name_namelist[i]), 'black', font=font1)

                img_rd = cv2.cvtColor(np.array(pilimg), cv2.COLOR_RGB2BGR)

                # cv2.putText(img_rd, str(name_namelist[i]), pos_namelist[i],font, 0.8, (0, 255, 255), 1, cv2.LINE_AA)

    print("Faces in camera now:", name_namelist, "\n")

    cv2.putText(img_rd, "Press 'q': Quit", (20, 450), font, 0.8, (84, 255, 159), 1, cv2.LINE_AA)
    cv2.putText(img_rd, "Face Recognition", (20, 40), font, 1, (0, 0, 0), 1, cv2.LINE_AA)
    cv2.putText(img_rd, "Faces: " + str(len(faces)), (20, 100), font, 1, (0, 0, 255), 1, cv2.LINE_AA)

    # 窗口显示 show with opencv
    cv2.imshow("camera", img_rd)

# 释放摄像头 release camera
cap.release()

# 删除建立的窗口 delete all the windows
cv2.destroyAllWindows()


import re
import numpy as np

class nbsegmenter(object):
    def __init__(self,
                 model='../Models/model_nbbow.out',
                 topics=['History','Labs','Medications','PhysicalExam','Courses'],
                 labels=['A', 'B', 'C', 'D', 'E'],
                 seged='../Output/segmented.txt'):
        self.__model=model
        self.__scores={}
        self.__topics=topics
        self.__topics_len=len(topics)
        self.__output=seged
        self.__labels=labels

    def __loadnbmodel(self):
        f = open(self.__model)
        for line in f:
            word, counts = line.strip().rsplit('\t', 1)
            self.__scores[word] = eval(counts)
        f.close()

    def __BoundaryFinder(self,text):
        values = []; noteline = ""; count = 0;
        for line in text:
            noteline += line + " " + "\n"
            words = re.split(' ', noteline)
            preValues = np.zeros(self.__topics_len)
            pre = np.zeros(self.__topics_len)
            for word in words:
                if word in self.__scores.keys():
                    preValues+=np.asarray(self.__scores[word])
            preValues=max(preValues)-preValues
            if (len(values) is not 0):
                pre = preValues - values[len(values)-1]
            if ((max(pre) + min(pre)) < 0):
                break
            values.append(preValues)
            count += 1
        return count, self.__topics[preValues.tolist().index(0)]

    def get_Boundary_Position(self,note):
        self.__loadnbmodel()
        BoundaryPostion=[]
        T=note
        while 1:
            index, label = self.__BoundaryFinder(T)
            T = T[index:]
            while index:
                BoundaryPostion.append(self.__labels[self.__topics.index(label)])
                index=index-1
            if (T == []):
                break
        return BoundaryPostion

    def get_Seg_index(self,notelines):
        count = 0
        Segindex = []
        BoundaryPositions=self.get_Boundary_Position(notelines)
        lastele = BoundaryPositions[0]
        for ele in BoundaryPositions:
            if ele == lastele: count = count + 1
            else:
                Segindex.append([lastele, count])
                count = 1
                lastele = ele
        return Segindex

    def print_Segs(self,notelines):
        lines = notelines
        Segindex=self.get_Seg_index(notelines)
        for ele in Segindex:
            print("========" + self.__topics [self.__labels.index(ele[0])] + "==========")
            for line in lines[:ele[1]]:
                print(line)
                print("\n")
            lines = lines[ele[1]:]


# Copyright (c) 2019, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: MIT
# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/MIT

import torch
import numpy as np
from torch.distributions import Normal
from base.actors.base import BaseActor
from base.learners.distance import BaseDistanceLearner, BaseSiblingRivalryLearner
from base.learners.her import BaseHERLearner
from agents.maze_agents.modules import Policy, Critic
from agents.maze_agents.ant_maze.env import Env

class Agent(BaseActor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        assert self.noise is not None
        assert self.epsilon is not None

    def _make_modules(self, policy):
        self.policy = policy

    @property
    def rollout(self):
        states = torch.stack([e['state'] for e in self.episode] + [self.episode[-1]['next_state']]).data.numpy()
        xs = states[:, 0]
        ys = states[:, 1]
        return [xs, ys]

    def step(self, do_eval=False):
        s = self.env.state
        g = self.env.goal
        a = self.policy(s.view(1, -1), g.view(1, -1)).view(-1)

        if not do_eval:
            if np.random.rand() < self.epsilon:
                a = np.random.uniform(
                    low=-self.env.action_range,
                    high=self.env.action_range,
                    size=(self.env.action_size,)
                )
                a = torch.from_numpy(a.astype(np.float32))
            else:
                z = Normal(torch.zeros_like(a), torch.ones_like(a) * self.noise * self.env.action_range)
                a = a + z.sample()
        a = torch.clamp(a, -self.env.action_range, self.env.action_range)

        self.env.step(a)
        complete = float(self.env.is_success) * torch.ones(1)
        terminal = float(self.env.is_done) * torch.ones(1)
        s_next = self.env.state
        r = -1 * torch.ones(1)

        self.episode.append({
            'state': s,
            'goal': g,
            'achieved': self.env.achieved,
            'action': a,
            'next_state': s_next,
            'terminal': terminal.view([]),
            'complete': complete.view([]),
            'reward': r.view([]),
        })

class DistanceLearner(BaseDistanceLearner):
    def create_env(self):
        return Env(**self.env_params)

    def _make_agent_modules(self):
        self.policy = Policy(self._dummy_env, 128)
        self.p_target = Policy(self._dummy_env, 128)
        self.p_target.load_state_dict(self.policy.state_dict())

        self.q_module = Critic(self._dummy_env, 128)
        self.q_target = Critic(self._dummy_env, 128)
        self.q_target.load_state_dict(self.q_module.state_dict())

    def _make_agent(self):
        return Agent(noise=self.noise, epsilon=self.epsilon, env=self.create_env(), policy=self.policy)

    def soft_update(self):
        module_pairs = [
            dict(source=self.q_module, target=self.q_target),
            dict(source=self.policy, target=self.p_target),
        ]
        for pair in module_pairs:
            for p, p_targ in zip(pair['source'].parameters(), pair['target'].parameters()):
                p_targ.data *= self.polyak
                p_targ.data += (1 - self.polyak) * p.data

    def get_icm_loss(self, batch):
        return self.icm(batch['state'], batch['next_state'], batch['action'])

    def get_next_qs(self, batch):
        next_policy_actions = self.p_target(batch['next_state'], batch['goal'])
        return self.q_target(batch['next_state'], next_policy_actions, batch['goal'], batch.get('antigoal', None))

    def get_action_qs(self, batch):
        return self.q_module(batch['state'], batch['action'], batch['goal'], batch.get('antigoal', None))

    def get_policy_loss_and_actions(self, batch):
        policy_actions = self.policy(batch['state'], batch['goal'])
        p_losses = -self.q_target.q_no_grad(batch['state'], policy_actions, batch['goal'], batch.get('antigoal', None))
        p_loss = p_losses.mean()
        return p_loss, policy_actions

class SiblingRivalryLearner(BaseSiblingRivalryLearner, DistanceLearner):
    def _make_agent_modules(self):
        self.policy = Policy(self._dummy_env, 128)
        self.p_target = Policy(self._dummy_env, 128)
        self.p_target.load_state_dict(self.policy.state_dict())

        self.q_module = Critic(self._dummy_env, 128, use_antigoal=self.use_antigoal)
        self.q_target = Critic(self._dummy_env, 128, use_antigoal=self.use_antigoal)
        self.q_target.load_state_dict(self.q_module.state_dict())

class HERLearner(BaseHERLearner, DistanceLearner):
    pass