import collections
import math
import warnings

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import acovf
from scipy.optimize import curve_fit
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components

class Speciation:
    """
    Calculate the solvation shells of every solute.

    Speciation organizes the solvation data by the type of residue
    coordinated with the central solvent. It collects this information in a
    pandas.DataFrame indexed by the frame and solute number. Each column is
    one of the solvents in the res_name column of the solvation data. The
    column value is how many residue of that type are in the solvation shell.

    Speciation provides the speciation of each solute in the speciation
    attribute, it also calculates the percentage of each unique
    shell and makes it available in the speciation_percent attribute.

    Additionally, there are methods for finding solvation shells of
    interest and computing how common certain shell configurations are.

    Parameters
    ----------
    solvation_data : pandas.DataFrame
        The solvation data frame output by Solution.
    n_frames : int
        The number of frames in solvation_data.
    n_solutes : int
        The number of solutes in solvation_data.

    Attributes
    ----------
    speciation : pandas.DataFrame
        a dataframe containing the speciation of every li ion at
        every trajectory frame. Indexed by frame and solute numbers.
        Columns are the solvent molecules and values are the number
        of solvent in the shell.
    speciation_percent : pandas.DataFrame
        the percentage of shells of each type. Columns are the solvent
        molecules and and values are the number of solvent in the shell.
        The final column is the percentage of total shell of that
        particular composition.
    co_occurrence : pandas.DataFrame
        The actual co-occurrence of solvents divided by the expected co-occurrence.
        In other words, given one molecule of solvent i in the shell, what is the
        probability of finding a solvent j relative to choosing a solvent at random
        from the pool of all coordinated solvents. This matrix will
        likely not be symmetric.
    """

    def __init__(self, solvation_data, n_frames, n_solutes):
        self.solvation_data = solvation_data
        self.n_frames = n_frames
        self.n_solutes = n_solutes
        self.speciation_data, self.speciation_percent = self._compute_speciation()
        self.co_occurrence = self._solvent_co_occurrence()

    @staticmethod
    def from_solution(solution):
        """
        Generate a Speciation object from a solution.

        Parameters
        ----------
        solution : Solution

        Returns
        -------
        Pairing
        """
        assert solution.has_run, "The solution must be run before calling from_solution"
        return Speciation(
            solution.solvation_data,
            solution.n_frames,
            solution.n_solute,
        )

    def _compute_speciation(self):
        counts = self.solvation_data.groupby(["frame", "solvated_atom", "res_name"]).count()["res_ix"]
        counts_re = counts.reset_index(["res_name"])
        speciation_data = counts_re.pivot(columns=["res_name"]).fillna(0).astype(int)
        res_names = speciation_data.columns.levels[1]
        speciation_data.columns = res_names
        sum_series = speciation_data.groupby(speciation_data.columns.to_list()).size()
        sum_sorted = sum_series.sort_values(ascending=False)
        speciation_percent = sum_sorted.reset_index().rename(columns={0: 'count'})
        speciation_percent['count'] = speciation_percent['count'] / (self.n_frames * self.n_solutes)
        return speciation_data, speciation_percent

    @classmethod
    def _mean_speciation(cls, speciation_frames, solute_number, frame_number):
        means = speciation_frames.sum(axis=1) / (solute_number * frame_number)
        return means

    def shell_percent(self, shell_dict):
        """
        Calculate the percentage of shells matching shell_dict.

        This function computes the percent of solvation shells that exist with a particular
        composition. The composition is specified by the shell_dict. The percent
        will be of all shells that match that specification.

        Attributes
        ----------
        shell_dict : dict of {str: int}
            a specification for a shell composition. Keys are residue names (str)
            and values are the number of desired residues. e.g. if shell_dict =
            {'mol1': 4} then the function will return the percentage of shells
            that have 4 mol1. Note that this may include shells with 4 mol1 and
            any number of other solvents. To specify a shell with 4 mol1 and nothing
            else, enter a dict such as {'mol1': 4, 'mol2': 0, 'mol3': 0}.

        Returns
        -------
        float
            the percentage of shells

        Examples
        --------

         .. code-block:: python

            # first define Li, BN, and FEC AtomGroups
            >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
            >>> solution.run()
            >>> solution.speciation.shell_percent({'BN': 4, 'PF6': 1})
            0.0898
        """
        query_list = [f"{name} == {str(count)}" for name, count in shell_dict.items()]
        query = " and ".join(query_list)
        query_counts = self.speciation_percent.query(query)
        return query_counts['count'].sum()

    def find_shells(self, shell_dict):
        """
        Find all solvation shells that match shell_dict.

        This returns the frame, solute index, and composition of all solutes
        that match the composition given in shell_dict.

        Attributes
        ----------
        shell_dict : dict of {str: int}
            a specification for a shell composition. Keys are residue names (str)
            and values are the number of desired residues. e.g. if shell_dict =
            {'mol1': 4} then the function will return all shells
            that have 4 mol1. Note that this may include shells with 4 mol1 and
            any number of other solvents. To specify a shell with 4 mol1 and nothing
            else, enter a dict such as {'mol1': 4, 'mol2': 0, 'mol3': 0}.

        Returns
        -------
        pandas.DataFrame
            the index and composition of all shells that match shell_dict
        """
        query_list = [f"{name} == {str(count)}" for name, count in shell_dict.items()]
        query = " and ".join(query_list)
        query_counts = self.speciation_data.query(query)
        return query_counts

    def _solvent_co_occurrence(self):
        # calculate the co-occurrence of solvent molecules.
        expected_solvents_list = []
        actual_solvents_list = []
        for solvent in self.speciation_data.columns.values:
            # calculate number of available coordinating solvent slots
            shells_w_solvent = self.speciation_data.query(f'{solvent} > 0')
            n_solvents = shells_w_solvent.sum()
            # calculate expected number of coordinating solvents
            n_coordination_slots = n_solvents.sum() - len(shells_w_solvent)
            coordination_percentage = self.speciation_data.sum() / self.speciation_data.sum().sum()
            expected_solvents = coordination_percentage * n_coordination_slots
            # calculate actual number of coordinating solvents
            actual_solvents = n_solvents.copy()
            actual_solvents[solvent] = actual_solvents[solvent] - len(shells_w_solvent)
            # name series and append to list
            expected_solvents.name = solvent
            actual_solvents.name = solvent
            expected_solvents_list.append(expected_solvents)
            actual_solvents_list.append(actual_solvents)
        # make DataFrames
        actual_df = pd.concat(actual_solvents_list, axis=1)
        expected_df = pd.concat(expected_solvents_list, axis=1)
        # calculate correlation matrix
        correlation = actual_df / expected_df
        return correlation

    def plot_co_occurrence(self):
        """
        Plot the co-occurrence matrix of the solution.

        Co-occurrence as a heatmap with numerical values in addition to colors.

        Returns
        -------
        fig : matplotlib.Figure
        ax : matplotlib.Axes

        """
        solvent_names = self.speciation_data.columns.values
        fig, ax = plt.subplots()
        im = ax.imshow(self.co_occurrence)
        # We want to show all ticks...
        ax.set_xticks(np.arange(len(solvent_names)))
        ax.set_yticks(np.arange(len(solvent_names)))
        # ... and label them with the respective list entries
        ax.set_xticklabels(solvent_names, fontsize=14)
        ax.set_yticklabels(solvent_names, fontsize=14)
        # Let the horizontal axes labeling appear on top.
        ax.tick_params(top=True, bottom=False,
                       labeltop=True, labelbottom=False, )
        # Rotate the tick labels and set their alignment.
        plt.setp(ax.get_xticklabels(), rotation=-30, ha="right",
                 rotation_mode="anchor")
        # Loop over data dimensions and create text annotations.
        for i in range(len(solvent_names)):
            for j in range(len(solvent_names)):
                ax.text(j, i, round(self.co_occurrence.iloc[i, j], 2),
                        horizontalalignment="center",
                        verticalalignment="center",
                        color="black",
                        fontsize=14,
                        )
        fig.tight_layout()
        return fig, ax

class Coordination:
    """
    Calculate the coordination number for each solvent.

    Coordination calculates the coordination number by averaging the number of
    coordinated solvents in all of the solvation shells. This is equivalent to
    the typical method of integrating the solute-solvent RDF up to the solvation
    radius cutoff. As a result, Coordination calculates **species-species** coordination
    numbers, not the total coordination number of the solute. So if the coordination
    number of mol1 is 3.2, there are on average 3.2 mol1 residues within the solvation
    distance of each solute.

    The coordination numbers are made available as an mean over the whole
    simulation and by frame.

    Parameters
    ----------
    solvation_data : pandas.DataFrame
        The solvation data frame output by Solution.
    n_frames : int
        The number of frames in solvation_data.
    n_solutes : int
        The number of solutes in solvation_data.

    Attributes
    ----------
    cn_dict : dict of {str: float}
        a dictionary where keys are residue names (str) and values are the
        mean coordination number of that residue (float).
    cn_by_frame : pd.DataFrame
        a DataFrame of the mean coordination number of in each frame of the trajectory.
    coordinating_atoms : pd.DataFrame
        percent of each atom_type participating in solvation, calculated for each solvent.

    Examples
    --------

     .. code-block:: python

        # first define Li, BN, and FEC AtomGroups
        >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
        >>> solution.run()
        >>> solution.coordination.cn_dict
        {'BN': 4.328, 'FEC': 0.253, 'PF6': 0.128}

    """

    def __init__(self, solvation_data, n_frames, n_solutes, atom_group):
        self.solvation_data = solvation_data
        self.n_frames = n_frames
        self.n_solutes = n_solutes
        self.cn_dict, self.cn_by_frame = self._mean_cn()
        self.atom_group = atom_group
        self.coordinating_atoms = self._calculate_coordinating_atoms()

    @staticmethod
    def from_solution(solution):
        """
        Generate a Coordination object from a solution.

        Parameters
        ----------
        solution : Solution

        Returns
        -------
        Pairing
        """
        assert solution.has_run, "The solution must be run before calling from_solution"
        return Coordination(
            solution.solvation_data,
            solution.n_frames,
            solution.n_solute,
            solution.u.atoms,
        )

    def _mean_cn(self):
        counts = self.solvation_data.groupby(["frame", "solvated_atom", "res_name"]).count()["res_ix"]
        cn_series = counts.groupby(["res_name", "frame"]).sum() / (
                self.n_solutes * self.n_frames
        )
        cn_by_frame = cn_series.unstack()
        cn_dict = cn_series.groupby(["res_name"]).sum().to_dict()
        return cn_dict, cn_by_frame

    def _calculate_coordinating_atoms(self, tol=0.005):
        """
        Determine which atom types are actually coordinating
        return the types of those atoms
        """
        # lookup atom types
        atom_types = self.solvation_data.reset_index(['atom_ix'])
        atom_types['atom_type'] = self.atom_group[atom_types['atom_ix']].types
        # count atom types
        atoms_by_type = atom_types[['atom_type', 'res_name', 'atom_ix']]
        type_counts = atoms_by_type.groupby(['res_name', 'atom_type']).count()
        solvent_counts = type_counts.groupby(['res_name']).sum()['atom_ix']
        # calculate percent of each
        solvent_counts_list = [solvent_counts[solvent] for solvent in type_counts.index.get_level_values(0)]
        type_percents = type_counts['atom_ix'] / solvent_counts_list
        type_percents.name = 'percent'
        # change index type
        type_percents = (type_percents
                         .reset_index(level=1)
                         .astype({'atom_type': str})
                         .set_index('atom_type', append=True)
                         )
        return type_percents[type_percents.percent > tol]

class Pairing:
    """
    Calculate the percent of solutes that are coordinated with each solvent.

    The pairing percentage is the percent of solutes that are coordinated with
    ANY solvent with matching type. So if the pairing of mol1 is 0.5, then 50% of
    solutes are coordinated with at least 1 mol1.

    The pairing percentages are made available as an mean over the whole
    simulation and by frame.

    Parameters
    ----------
    solvation_data : pandas.DataFrame
        The solvation data frame output by Solution.
    n_frames : int
        The number of frames in solvation_data.
    n_solutes : int
        The number of solutes in solvation_data.
    n_solvents : dict of {str: int}
        The number of each kind of solvent.

    Attributes
    ----------
    pairing_dict : dict of {str: float}
        a dictionary where keys are residue names (str) and values are the
        percentage of solutes that contain that residue (float).
    pairing_by_frame : pd.DataFrame
        a dictionary tracking the mean percentage of each residue across frames.
    percent_free_solvents : dict of {str: float}
        a dictionary containing the percent of each solvent that is free. e.g.
        not coordinated to a solute.
    diluent_dict : dict of {str: float}
        the fraction of the diluent constituted by each solvent. The diluent is
        defined as everything that is not coordinated with the solute.
    diluent_by_frame : pd.DataFrame
        a DataFrame of the diluent composition in each frame of the trajectory.
    diluent_counts : pd.DataFrame
        a DataFrame of the raw solvent counts in the diluent in each frame of the trajectory.

    Examples
    --------

     .. code-block:: python

        # first define Li, BN, and FEC AtomGroups
        >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
        >>> solution.run()
        >>> solution.pairing.pairing_dict
        {'BN': 1.0, 'FEC': 0.210, 'PF6': 0.120}
    """

    def __init__(self, solvation_data, n_frames, n_solutes, n_solvents):
        self.solvation_data = solvation_data
        self.n_frames = n_frames
        self.n_solutes = n_solutes
        self.solvent_counts = n_solvents
        self.pairing_dict, self.pairing_by_frame = self._percent_coordinated()
        self.percent_free_solvents = self._percent_free_solvent()
        self.diluent_dict, self.diluent_by_frame, self.diluent_counts = self._diluent_composition()

    @staticmethod
    def from_solution(solution):
        """
        Generate a Pairing object from a solution.

        Parameters
        ----------
        solution : Solution

        Returns
        -------
        Pairing
        """
        assert solution.has_run, "The solution must be run before calling from_solution"
        return Pairing(
            solution.solvation_data,
            solution.n_frames,
            solution.n_solute,
            solution.solvent_counts
        )

    def _percent_coordinated(self):
        # calculate the percent of solute coordinated with each solvent
        counts = self.solvation_data.groupby(["frame", "solvated_atom", "res_name"]).count()["res_ix"]
        pairing_series = counts.astype(bool).groupby(["res_name", "frame"]).sum() / (
            self.n_solutes
        )  # mean coordinated overall
        pairing_by_frame = pairing_series.unstack()
        pairing_normalized = pairing_series / self.n_frames
        pairing_dict = pairing_normalized.groupby(["res_name"]).sum().to_dict()
        return pairing_dict, pairing_by_frame

    def _percent_free_solvent(self):
        # calculate the percent of each solvent NOT coordinated with the solute
        counts = self.solvation_data.groupby(["frame", "res_ix", "res_name"]).count()['dist']
        totals = counts.groupby(['res_name']).count() / self.n_frames
        n_solvents = np.array([self.solvent_counts[name] for name in totals.index.values])
        free_solvents = np.ones(len(totals)) - totals / n_solvents
        return free_solvents.to_dict()

    def _diluent_composition(self):
        coordinated_solvents = self.solvation_data.groupby(["frame", "res_name"]).nunique()["res_ix"]
        solvent_counts = pd.Series(self.solvent_counts)
        total_solvents = solvent_counts.reindex(coordinated_solvents.index, level=1)
        diluent_solvents = total_solvents - coordinated_solvents
        diluent_series = diluent_solvents / diluent_solvents.groupby(['frame']).sum()
        diluent_by_frame = diluent_series.unstack().T
        diluent_counts = diluent_solvents.unstack().T
        diluent_dict = diluent_by_frame.mean(axis=1).to_dict()
        return diluent_dict, diluent_by_frame, diluent_counts

class Residence:
    """
    Calculate the residence times of solvents.

    This class calculates the residence time of each solvent on the solute.
    The residence time is in units of Solution frames, so if the Solution object
    has 1000 frames over 1 nanosecond, then each frame will be 1 picosecond.
    Thus a residence time of 100 would translate to 100 picoseconds.

    Two residence time implementations are available. Both calculate the
    solute-solvent autocorrelation function for each solute-solvent pair,
    take and take the mean over the solvents of each type, this should yield
    an exponentially decaying autocorrelation function.

    The first implementation fits an exponential curve to the autocorrelation
    function and extract the time constant, which is inversely proportional to the
    residence time. This result is saved in the ``residence_times_fit`` attribute.
    Unfortunately, the fit often fails to converge (value is set to np.nan),
    making this method unreliable.

    Instead, the default implementation is to simply find point where the
    value of the autocorrelation function is 1/e, which is the time constant
    of an exact exponential. These values are saved in ``residence_times``.

    It is recommended that the user visually inspect the autocorrelation function
    with ``Residence.plot_autocorrelation_function`` to ensure an approximately
    exponential decay. The residence times are only valid insofar as the autocorrelation
    function resembles an exact exponential, it should decays to zero with a long tail.
    If the exponential does not decay to zero or its slope does not level off, increasing
    the simulation time may help. For this technique to be appropriate, the simulation time
    should exceed the residence time.

    A fuller description of the method can be found in
    `Self, Fong, and Persson <https://pubs-acs-org.libproxy.berkeley.edu/doi/full/10.1021/acsenergylett.9b02118>`_

    Parameters
    ----------
    solvation_data : pandas.DataFrame
        The solvation data frame output by Solution.
    step : int
        The spacing of frames in solvation_data. This should be equal
        to solution.step.

    Attributes
    ----------
    residence_times : dict of {str: float}
        a dictionary where keys are residue names and values are the
        residence times of the that residue on the solute, calculated
        with the 1/e cutoff method.
    residence_times_fit : dict of {str: float}
        a dictionary where keys are residue names and values are the
        residence times of the that residue on the solute, calculated
        with the exponential fit method.
    fit_parameters : pd.DataFrame
        a dictionary where keys are residue names and values are the
        arameters for the exponential fit to the autocorrelation function.

    Examples
    --------

     .. code-block:: python

        # first define Li, BN, and FEC AtomGroups
        >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
        >>> residence = Residence.from_solution(solution)
        >>> residence.residence_times
        {'BN': 4.02, 'FEC': 3.79, 'PF6': 1.15}
    """

    def __init__(self, solvation_data, step):
        self.solvation_data = solvation_data
        self.auto_covariances = self._calculate_auto_covariance_dict()
        self.residence_times = self._calculate_residence_times_with_cutoff(self.auto_covariances, step)
        self.residence_times_fit, self.fit_parameters = self._calculate_residence_times_with_fit(
            self.auto_covariances,
            step
        )

    @staticmethod
    def from_solution(solution):
        """
        Generate a Residence object from a solution.

        Parameters
        ----------
        solution : Solution

        Returns
        -------
        Residence
        """
        assert solution.has_run, "The solution must be run before calling from_solution"
        return Residence(
            solution.solvation_data,
            solution.step
        )

    def _calculate_auto_covariance_dict(self):
        frame_solute_index = np.unique(self.solvation_data.index.droplevel(2))
        auto_covariance_dict = {}
        for res_name, res_solvation_data in self.solvation_data.groupby(['res_name']):
            adjacency_mini = Residence.calculate_adjacency_dataframe(res_solvation_data)
            adjacency_df = adjacency_mini.reindex(frame_solute_index, fill_value=0)
            auto_covariance = Residence._calculate_auto_covariance(adjacency_df)
            # normalize
            auto_covariance = auto_covariance / np.max(auto_covariance)
            auto_covariance_dict[res_name] = auto_covariance
        return auto_covariance_dict

    @staticmethod
    def _calculate_residence_times_with_cutoff(auto_covariances, step, convergence_cutoff=0.1):
        residence_times = {}
        for res_name, auto_covariance in auto_covariances.items():
            if np.min(auto_covariance) > convergence_cutoff:
                residence_times[res_name] = np.nan
                warnings.warn(f'the autocovariance for {res_name} does not converge to zero '
                              'so a residence time cannot be calculated. A longer simulation '
                              'is required to get a valid estimate of the residence time.')
            unassigned = True
            for frame, val in enumerate(auto_covariance):
                if val < 1 / math.e:
                    residence_times[res_name] = frame * step
                    unassigned = False
                    break
            if unassigned:
                residence_times[res_name] = np.nan
        return residence_times

    @staticmethod
    def _calculate_residence_times_with_fit(auto_covariances, step):
        # calculate the residence times
        residence_times = {}
        fit_parameters = {}
        for res_name, auto_covariance in auto_covariances.items():
            res_time, params = Residence._fit_exponential(auto_covariance, res_name)
            residence_times[res_name], fit_parameters[res_name] = res_time * step, params
        return residence_times, fit_parameters

    def plot_auto_covariance(self, res_name):
        """
        Plot the autocovariance of a solvent on the solute.

        See the discussion in the class docstring for more information.

        Parameters
        ----------
        res_name : str
            the name of a solvent in the solution.

        Returns
        -------
        fig : matplotlib.Figure
        ax : matplotlib.Axes
        """
        auto_covariance = self.auto_covariances[res_name]
        frames = np.arange(len(auto_covariance))
        params = self.fit_parameters[res_name]
        exp_func = lambda x: self._exponential_decay(x, *params)
        exp_fit = np.array(map(exp_func, frames))
        fig, ax = plt.subplots()
        ax.plot(frames, auto_covariance, "b-", label="auto covariance")
        try:
            ax.scatter(frames, exp_fit, label="exponential fit")
        except:
            warnings.warn(f'The fit for {res_name} failed so the exponential '
                          f'fit will not be plotted.')
        ax.hlines(y=1/math.e, xmin=frames[0], xmax=frames[-1], label='1/e cutoff')
        ax.set_xlabel("Timestep (frames)")
        ax.set_ylabel("Normalized Autocovariance")
        ax.set_ylim(0, 1)
        ax.legend()
        return fig, ax

    @staticmethod
    def _exponential_decay(x, a, b, c):
        """
        An exponential decay function.

        Args:
            x: Independent variable.
            a: Initial quantity.
            b: Exponential decay constant.
            c: Constant.

        Returns:
            The acf
        """
        return a * np.exp(-b * x) + c

    @staticmethod
    def _fit_exponential(auto_covariance, res_name):
        auto_covariance_norm = auto_covariance / auto_covariance[0]
        try:
            params, param_covariance = curve_fit(
                Residence._exponential_decay,
                np.arange(len(auto_covariance_norm)),
                auto_covariance_norm,
                p0=(1, 0.1, 0.01),
            )
            tau = 1 / params[1]  # p
        except RuntimeError:
            warnings.warn(f'The fit for {res_name} failed so its values in'
                          f'residence_time_fits and fit_parameters will be'
                          f'set to np.nan.')
            tau, params = np.nan, (np.nan, np.nan, np.nan)
        return tau, params

    @staticmethod
    def _calculate_auto_covariance(adjacency_matrix):
        auto_covariances = []
        for solute_ix, df in adjacency_matrix.groupby(['solvated_atom']):
            non_zero_cols = df.loc[:, (df != 0).any(axis=0)]
            auto_covariance_df = non_zero_cols.apply(
                acovf,
                axis=0,
                result_type='expand',
                demean=False,
                unbiased=True,
                fft=True
            )
            auto_covariances.append(auto_covariance_df.values)
        auto_covariance = np.mean(np.concatenate(auto_covariances, axis=1), axis=1)
        return auto_covariance

    @staticmethod
    def calculate_adjacency_dataframe(solvation_data):
        """
        Calculate a frame-by-frame adjacency matrix from the solvation data.

        This will calculate the adjacency matrix of the solute and all possible
        solvents. It will maintain an index of ['frame', 'solvated_atom', 'res_ix']
        where each 'frame' is a sparse adjacency matrix between solvated atom ix
        and residue ix.

        Parameters
        ----------
        solvation_data : pd.DataFrame
            the solvation_data from a Solution.

        Returns
        -------
        adjacency_df : pandas.DataFrame
        """
        # generate an adjacency matrix from the solvation data
        adjacency_group = solvation_data.groupby(['frame', 'solvated_atom', 'res_ix'])
        adjacency_df = adjacency_group['dist'].count().unstack(fill_value=0)
        return adjacency_df

class Networking:
    """
    Calculate the number and size of solute-solvent networks.

    A network is defined as a bipartite graph of solutes and solvents, where edges
    are defined by coordination in the solvation_data DataFrame. A single solvent
    or multiple solvents can be selected, but coordination between solvents will
    not be included, only coordination between solutes and solvents.

    Networking uses the solvation_data to construct an adjacency matrix and then
    extracts the connected subgraphs within it. These connected subgraphs are stored
    in a DataFrame in Networking.network_df.

    Several other representations of the networking data are included in the attributes.

    Parameters
    ----------
    solvents : str or list[str]
        the solvents to include in the solute-solvent network.
    solvation_data : pandas.DataFrame
        a dataframe of solvation data with columns "frame", "solvated_atom", "atom_ix",
        "dist", "res_name", and "res_ix".
    solute_res_ix : np.ndarray
        the residue indices of the solutes in solvation_data
    res_name_map : pd.Series
        a mapping between residue indices and the solute & solvent names in a Solution.

    Attributes
    ----------
    network_df : pd.DataFrame
        the dataframe containing all networking data. the indices are the frame and
        network index, respectively. the columns are the res_name and res_ix.
    network_sizes : pd.DataFrame
        a dataframe of network sizes. the index is the frame. the column headers
        are network sizes, or the number of solutes + solvents in the network, so
        the columns might be [2, 3, 4, ...]. the values in each column are the
        number of networks with that size in each frame.
    solute_status : dict of {str: float}
        a dictionary where the keys are the "status" of the solute and the values
        are the fraction of solute with that status, averaged over all frames.
        "alone" means that the solute not coordinated with any of the networking
        solvents, network size is 1.
        "paired" means the solute and is coordinated with a single networking
        solvent and that solvent is not coordinated to any other solutes, network
        size is 2.
        "in_network" means that the solute is coordinated to more than one solvent
        or its solvent is coordinated to more than one solute, network size >= 3.
    solute_status_by_frame : pd.DataFrame
        as described above, except organized into a dataframe where each
        row is a unique frame and the columns are "alone", "paired", and "in_network".

    Examples
    --------
     .. code-block:: python

        # first define Li, BN, and FEC AtomGroups
        >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
        >>> networking = Networking.from_solution(solution, 'PF6')
    """

    def __init__(self, solvents, solvation_data, solute_res_ix, res_name_map):
        self.solvents = solvents
        self.solvation_data = solvation_data
        solvent_present = np.isin(self.solvents, self.solvation_data['res_name'].unique())
        if not solvent_present.all():
            raise Exception(f"Solvent(s) {np.array(self.solvents)[~solvent_present]} not found in solvation data.")
        self.solute_res_ix = solute_res_ix
        self.res_name_map = res_name_map
        self.n_solute = len(solute_res_ix)
        self.network_df = self._generate_networks()
        self.network_sizes = self._calculate_network_sizes()
        self.solute_status, self.solute_status_by_frame = self._calculate_solute_status()
        self.solute_status = self.solute_status.to_dict()

    @staticmethod
    def from_solution(solution, solvents):
        """
        Generate a Networking object from a solution and solvent names.

        Parameters
        ----------
        solution : Solution
        solvents : str or list of str
            the strings should be the name of solvents in the Solution. The
            strings must match exactly for Networking to work properly. The
            selected solvents will be used to construct the networking graph
            that is described in documentation for the Networking class.

        Returns
        -------
        Networking
        """
        return Networking(
            solvents,
            solution.solvation_data,
            solution.solute_res_ix,
            solution.res_name_map,
        )

    @staticmethod
    def _unwrap_adjacency_dataframe(df):
        # this class will transform the biadjacency matrix into a proper adjacency matrix
        connections = df.reset_index(level=0).drop(columns='frame')
        idx = connections.columns.append(connections.index)
        directed = connections.reindex(index=idx, columns=idx, fill_value=0)
        undirected = directed.values + directed.values.T
        adjacency_matrix = csr_matrix(undirected)
        return adjacency_matrix

    def _generate_networks(self):
        """
        This function generates a dataframe containing all the solute-solvent networks
        in every frame of the simulation. The rough approach is as follows:

        1. transform the solvation_data DataFrame into an adjacency matrix
        2. determine the connected subgraphs in the adjacency matrix
        3. tabulate the res_ix involved in each network and store in a DataFrame
        """
        solvents = [self.solvents] if isinstance(self.solvents, str) else self.solvents
        solvation_subset = self.solvation_data[np.isin(self.solvation_data.res_name, solvents)]
        # reindex solvated_atom to residue indexes
        reindexed_subset = solvation_subset.reset_index(level=1)
        reindexed_subset.solvated_atom = self.solute_res_ix[reindexed_subset.solvated_atom]
        dropped_reindexed = reindexed_subset.set_index(['solvated_atom'], append=True)
        reindexed_subset = dropped_reindexed.reorder_levels(['frame', 'solvated_atom', 'atom_ix'])
        # create adjacency matrix from reindexed df
        graph = Residence.calculate_adjacency_dataframe(reindexed_subset)
        network_arrays = []
        # loop through each time step / frame
        for frame, df in graph.groupby('frame'):
            # drop empty columns
            df = df.loc[:, (df != 0).any(axis=0)]
            # save map from local index to residue index
            solute_map = df.index.get_level_values(1).values
            solvent_map = df.columns.values
            ix_to_res_ix = np.concatenate([solvent_map, solute_map])
            adjacency_df = Networking._unwrap_adjacency_dataframe(df)
            _, network = connected_components(
                csgraph=adjacency_df,
                directed=False,
                return_labels=True
            )
            network_array = np.vstack([
                np.full(len(network), frame),  # frame
                network,  # network
                self.res_name_map[ix_to_res_ix],  # res_names
                ix_to_res_ix,  # res index
            ]).T
            network_arrays.append(network_array)
        # create and return network dataframe
        all_clusters = np.concatenate(network_arrays)
        cluster_df = (
            pd.DataFrame(all_clusters, columns=['frame', 'network', 'res_name', 'res_ix'])
                .set_index(['frame', 'network'])
                .sort_values(['frame', 'network'])
        )
        return cluster_df

    def _calculate_network_sizes(self):
        # This utility calculates the network sizes and returns a convenient dataframe.
        cluster_df = self.network_df
        cluster_sizes = cluster_df.groupby(['frame', 'network']).count()
        size_counts = cluster_sizes.groupby(['frame', 'res_name']).count().unstack(fill_value=0)
        size_counts.columns = size_counts.columns.droplevel()
        return size_counts

    def _calculate_solute_status(self):
        """
        This utility calculates the percentage of each solute with a given "status".
        Namely, whether the solvent is "alone", "paired" (with a single solvent), or
        "in_network" of > 2 species.
        """
        status = self.network_sizes.rename(columns={2: 'paired'})
        status['in_network'] = status.iloc[:, 1:].sum(axis=1).astype(int)
        status['alone'] = self.n_solute - status.loc[:, ['paired', 'in_network']].sum(axis=1)
        status = status.loc[:, ['alone', 'paired', 'in_network']]
        solute_status_by_frame = status / self.n_solute
        solute_status = solute_status_by_frame.mean()
        return solute_status, solute_status_by_frame

    def get_network_res_ix(self, network_index, frame):
        """
        Return the indexes of all residues in a selected network.

        The network_index and frame must be provided to fully specify the network.
        Once the indexes are returned, they can be used to select an AtomGroup with
        the species of interest, see Examples.

        Parameters
        ----------
        network_index : int
            The index of the network of interest
        frame : int
            the frame in the trajectory to perform selection at. Defaults to the
            current trajectory frame.
        Returns
        -------
        res_ix : np.ndarray

        Examples
        --------
         .. code-block:: python

            # first define Li, BN, and FEC AtomGroups
            >>> solution = Solution(Li, {'BN': BN, 'FEC': FEC, 'PF6': PF6})
            >>> networking = Networking.from_solution(solution, 'PF6')
            >>> res_ix = networking.get_network_res_ix(1, 5)
            >>> solution.u.residues[res_ix].atoms
            <AtomGroup with 126 Atoms>

        """
        res_ix = self.network_df.loc[pd.IndexSlice[frame, network_index], 'res_ix'].values
        return res_ix.astype(int)


from .biencoder import AddLabelFixedCandsTRA
from .modules import TransformerEncoder
from .modules import get_n_positions_from_options
from parlai.core.torch_ranker_agent import TorchRankerAgent
from .transformer import TransformerRankerAgent
from .modules import BasicAttention, MultiHeadAttention
import torch

class PolyencoderAgent(TorchRankerAgent):
    """
    Poly-encoder Agent.

    Equivalent of bert_ranker/polyencoder and biencoder_multiple_output but does not
    rely on an external library (hugging face).
    """

    @classmethod
    def add_cmdline_args(cls, argparser):
        """
        Add command-line arguments specifically for this agent.
        """
        TransformerRankerAgent.add_cmdline_args(argparser)
        agent = argparser.add_argument_group('Polyencoder Arguments')
        agent.add_argument(
            '--polyencoder-type',
            type=str,
            default='codes',
            choices=['codes', 'n_first'],
            help='Type of polyencoder, either we compute'
            'vectors using codes + attention, or we '
            'simply take the first N vectors.',
            recommended='codes',
        )
        agent.add_argument(
            '--poly-n-codes',
            type=int,
            default=64,
            help='number of vectors used to represent the context'
            'in the case of n_first, those are the number'
            'of vectors that are considered.',
            recommended=64,
        )
        agent.add_argument(
            '--poly-attention-type',
            type=str,
            default='basic',
            choices=['basic', 'sqrt', 'multihead'],
            help='Type of the top aggregation layer of the poly-'
            'encoder (where the candidate representation is'
            'the key)',
            recommended='basic',
        )
        agent.add_argument(
            '--polyencoder-attention-keys',
            type=str,
            default='context',
            choices=['context', 'position'],
            help='Input emb vectors for the first level of attention. '
            'Context refers to the context outputs; position refers to the '
            'computed position embeddings.',
            recommended='context',
        )
        agent.add_argument(
            '--poly-attention-num-heads',
            type=int,
            default=4,
            help='In case poly-attention-type is multihead, '
            'specify the number of heads',
        )

        # Those arguments are here in case where polyencoder type is 'code'
        agent.add_argument(
            '--codes-attention-type',
            type=str,
            default='basic',
            choices=['basic', 'sqrt', 'multihead'],
            help='Type ',
            recommended='basic',
        )
        agent.add_argument(
            '--codes-attention-num-heads',
            type=int,
            default=4,
            help='In case codes-attention-type is multihead, '
            'specify the number of heads',
        )
        return agent

    def __init__(self, opt, shared=None):
        super().__init__(opt, shared)
        self.rank_loss = torch.nn.CrossEntropyLoss(reduce=True, size_average=True)
        if self.use_cuda:
            self.rank_loss.cuda()
        self.data_parallel = opt.get('data_parallel') and self.use_cuda
        if self.data_parallel:
            from parlai.utils.distributed import is_distributed

            if is_distributed():
                raise ValueError('Cannot combine --data-parallel and distributed mode')
            if shared is None:
                self.model = torch.nn.DataParallel(self.model)

    def build_model(self, states=None):
        """
        Return built model.
        """
        return PolyEncoderModule(self.opt, self.dict, self.NULL_IDX)

    def vectorize(self, *args, **kwargs):
        """
        Add the start and end token to the labels.
        """
        kwargs['add_start'] = True
        kwargs['add_end'] = True
        obs = super().vectorize(*args, **kwargs)
        return obs

    def _set_text_vec(self, *args, **kwargs):
        """
        Add the start and end token to the text.
        """
        obs = super()._set_text_vec(*args, **kwargs)
        if 'text_vec' in obs and 'added_start_end_tokens' not in obs:
            obs.force_set(
                'text_vec', self._add_start_end_tokens(obs['text_vec'], True, True)
            )
            obs['added_start_end_tokens'] = True
        return obs

    def vectorize_fixed_candidates(self, *args, **kwargs):
        """
        Vectorize fixed candidates.

        Override to add start and end token when computing the candidate encodings in
        interactive mode.
        """
        kwargs['add_start'] = True
        kwargs['add_end'] = True
        return super().vectorize_fixed_candidates(*args, **kwargs)

    def _make_candidate_encs(self, vecs):
        """
        Make candidate encs.

        The polyencoder module expects cand vecs to be 3D while torch_ranker_agent
        expects it to be 2D. This requires a little adjustment (used in interactive mode
        only)
        """
        rep = super()._make_candidate_encs(vecs)
        return rep.transpose(0, 1).contiguous()

    def encode_candidates(self, padded_cands):
        """
        Encode candidates.
        """
        padded_cands = padded_cands.unsqueeze(1)
        _, _, _, cand_rep = self.model(cand_tokens=padded_cands)
        return cand_rep

    def score_candidates(self, batch, cand_vecs, cand_encs=None):
        """
        Score candidates.

        The Poly-encoder encodes the candidate and context independently. Then, the
        model applies additional attention before ultimately scoring a candidate.
        """
        bsz = batch.text_vec.size(0)
        ctxt_rep, ctxt_rep_mask, ctxt_pos, _ = self.model(ctxt_tokens=batch.text_vec)

        if cand_encs is not None:
            if bsz == 1:
                cand_rep = cand_encs
            else:
                cand_rep = cand_encs.expand(bsz, cand_encs.size(1), -1)
        # bsz x num cands x seq len
        elif len(cand_vecs.shape) == 3:
            _, _, _, cand_rep = self.model(cand_tokens=cand_vecs)
        # bsz x seq len (if batch cands) or num_cands x seq len (if fixed cands)
        elif len(cand_vecs.shape) == 2:
            _, _, _, cand_rep = self.model(cand_tokens=cand_vecs.unsqueeze(1))
            num_cands = cand_rep.size(0)  # will be bsz if using batch cands
            cand_rep = cand_rep.expand(num_cands, bsz, -1).transpose(0, 1).contiguous()

        scores = self.model(
            ctxt_rep=ctxt_rep,
            ctxt_rep_mask=ctxt_rep_mask,
            cand_rep=cand_rep,
            ctxt_pos=ctxt_pos,
        )
        return scores

    def load_state_dict(self, state_dict):
        """
        Override to account for codes.
        """
        if self.model.type == 'codes' and 'codes' not in state_dict:
            state_dict['codes'] = self.model.codes
        super().load_state_dict(state_dict)

class PolyEncoderModule(torch.nn.Module):
    """
    Poly-encoder model.

    See https://arxiv.org/abs/1905.01969 for more details
    """

    def __init__(self, opt, dict, null_idx):
        super(PolyEncoderModule, self).__init__()
        self.null_idx = null_idx
        self.encoder_ctxt = self.get_encoder(opt, dict, null_idx, 'none_with_pos_embs')
        self.encoder_cand = self.get_encoder(opt, dict, null_idx, opt['reduction_type'])

        self.type = opt['polyencoder_type']
        self.n_codes = opt['poly_n_codes']
        self.attention_type = opt['poly_attention_type']
        self.attention_keys = opt.get('polyencoder_attention_keys', 'context')
        self.attention_num_heads = opt['poly_attention_num_heads']
        self.codes_attention_type = opt['codes_attention_type']
        self.codes_attention_num_heads = opt['codes_attention_num_heads']
        embed_dim = opt['embedding_size']

        # In case it's a polyencoder with code.
        if self.type == 'codes':
            # experimentally it seems that random with size = 1 was good.
            codes = torch.empty(self.n_codes, embed_dim)
            codes = torch.nn.init.uniform_(codes)
            self.codes = torch.nn.Parameter(codes)

            # The attention for the codes.
            if self.codes_attention_type == 'multihead':
                self.code_attention = MultiHeadAttention(
                    self.codes_attention_num_heads, embed_dim, opt['dropout']
                )
            elif self.codes_attention_type == 'sqrt':
                self.code_attention = PolyBasicAttention(
                    self.type, self.n_codes, dim=2, attn='sqrt', get_weights=False
                )
            elif self.codes_attention_type == 'basic':
                self.code_attention = PolyBasicAttention(
                    self.type, self.n_codes, dim=2, attn='basic', get_weights=False
                )

        # The final attention (the one that takes the candidate as key)
        if self.attention_type == 'multihead':
            self.attention = MultiHeadAttention(
                self.attention_num_heads, opt['embedding_size'], opt['dropout']
            )
        else:
            self.attention = PolyBasicAttention(
                self.type,
                self.n_codes,
                dim=2,
                attn=self.attention_type,
                get_weights=False,
            )

    def get_encoder(self, opt, dict, null_idx, reduction_type):
        """
        Return encoder, given options.

        :param opt:
            opt dict
        :param dict:
            dictionary agent
        :param null_idx:
            null/pad index into dict
        :reduction_type:
            reduction type for the encoder

        :return:
            a TransformerEncoder, initialized correctly
        """
        n_positions = get_n_positions_from_options(opt)
        embeddings = torch.nn.Embedding(
            len(dict), opt['embedding_size'], padding_idx=null_idx
        )
        torch.nn.init.normal_(embeddings.weight, 0, opt['embedding_size'] ** -0.5)
        return TransformerEncoder(
            n_heads=opt['n_heads'],
            n_layers=opt['n_layers'],
            embedding_size=opt['embedding_size'],
            ffn_size=opt['ffn_size'],
            vocabulary_size=len(dict),
            embedding=embeddings,
            dropout=opt['dropout'],
            attention_dropout=opt['attention_dropout'],
            relu_dropout=opt['relu_dropout'],
            padding_idx=null_idx,
            learn_positional_embeddings=opt['learn_positional_embeddings'],
            embeddings_scale=opt['embeddings_scale'],
            reduction_type=reduction_type,
            n_positions=n_positions,
            n_segments=2,
            activation=opt['activation'],
            variant=opt['variant'],
            output_scaling=opt['output_scaling'],
        )

    def attend(self, attention_layer, queries, keys, values, mask):
        """
        Apply attention.

        :param attention_layer:
            nn.Module attention layer to use for the attention
        :param queries:
            the queries for attention
        :param keys:
            the keys for attention
        :param values:
            the values for attention
        :param mask:
            mask for the attention keys

        :return:
            the result of applying attention to the values, with weights computed
            wrt to the queries and keys.
        """
        if keys is None:
            keys = values
        if isinstance(attention_layer, PolyBasicAttention):
            return attention_layer(queries, keys, mask_ys=mask, values=values)
        elif isinstance(attention_layer, MultiHeadAttention):
            return attention_layer(queries, keys, values, mask)
        else:
            raise Exception('Unrecognized type of attention')

    def encode(self, ctxt_tokens, cand_tokens):
        """
        Encode a text sequence.

        :param ctxt_tokens:
            2D long tensor, batchsize x sent_len
        :param cand_tokens:
            3D long tensor, batchsize x num_cands x sent_len
            Note this will actually view it as a 2D tensor
        :return:
            (ctxt_rep, ctxt_mask, ctxt_pos, cand_rep)
            - ctxt_rep 3D float tensor, batchsize x n_codes x dim
            - ctxt_mask byte:  batchsize x n_codes (all 1 in case
            of polyencoder with code. Which are the vectors to use
            in the ctxt_rep)
            - ctxt_pos 3D float tensor, batchsize x sent_len x dim
            - cand_rep (3D float tensor) batchsize x num_cands x dim
        """
        cand_embed = None
        ctxt_rep = None
        ctxt_rep_mask = None
        ctxt_pos = None
        if cand_tokens is not None:
            assert len(cand_tokens.shape) == 3
            bsz = cand_tokens.size(0)
            num_cands = cand_tokens.size(1)
            cand_embed = self.encoder_cand(cand_tokens.view(bsz * num_cands, -1))
            cand_embed = cand_embed.view(bsz, num_cands, -1)

        if ctxt_tokens is not None:
            assert len(ctxt_tokens.shape) == 2
            bsz = ctxt_tokens.size(0)
            # get context_representation. Now that depends on the cases.
            ctxt_out, ctxt_mask, ctxt_pos = self.encoder_ctxt(ctxt_tokens)
            att_keys = ctxt_out if self.attention_keys == 'context' else ctxt_pos
            dim = ctxt_out.size(2)

            if self.type == 'codes':
                ctxt_rep = self.attend(
                    self.code_attention,
                    queries=self.codes.repeat(bsz, 1, 1),
                    keys=att_keys,
                    values=ctxt_out,
                    mask=ctxt_mask,
                )
                ctxt_pos = None  # we don't need this anymore
                ctxt_rep_mask = ctxt_rep.new_ones(bsz, self.n_codes).byte()

            elif self.type == 'n_first':
                # Expand the output if it is not long enough
                if ctxt_out.size(1) < self.n_codes:
                    difference = self.n_codes - ctxt_out.size(1)
                    extra_rep = ctxt_out.new_zeros(bsz, difference, dim)
                    ctxt_rep = torch.cat([ctxt_out, extra_rep], dim=1)
                    ctxt_pos = torch.cat([ctxt_pos, extra_rep], dim=1)
                    extra_mask = ctxt_mask.new_zeros(bsz, difference)
                    ctxt_rep_mask = torch.cat([ctxt_mask, extra_mask], dim=1)
                else:
                    ctxt_rep = ctxt_out[:, 0 : self.n_codes, :]
                    ctxt_pos = ctxt_pos[:, 0 : self.n_codes, :]
                    ctxt_rep_mask = ctxt_mask[:, 0 : self.n_codes]

        return ctxt_rep, ctxt_rep_mask, ctxt_pos, cand_embed

    def score(self, ctxt_rep, ctxt_rep_mask, ctxt_pos, cand_embed):
        """
        Score the candidates.

        :param ctxt_rep:
            3D float tensor, bsz x ctxt_len x dim
        :param ctxt_rep_mask:
            2D byte tensor, bsz x ctxt_len, in case there are some elements
            of the ctxt that we should not take into account.
        :param ctx_pos: 3D float tensor, bsz x sent_len x dim
        :param cand_embed: 3D float tensor, bsz x num_cands x dim

        :return: scores, 2D float tensor: bsz x num_cands
        """
        # Attention keys determined by self.attention_keys
        # 'context' == use context final rep; otherwise use context position embs
        keys = ctxt_rep if self.attention_keys == 'context' else ctxt_pos
        # reduces the context representation to a 3D tensor bsz x num_cands x dim
        ctxt_final_rep = self.attend(
            self.attention, cand_embed, keys, ctxt_rep, ctxt_rep_mask
        )
        scores = torch.sum(ctxt_final_rep * cand_embed, 2)
        return scores

    def forward(
        self,
        ctxt_tokens=None,
        cand_tokens=None,
        ctxt_rep=None,
        ctxt_rep_mask=None,
        ctxt_pos=None,
        cand_rep=None,
    ):
        """
        Forward pass of the model.

        Due to a limitation of parlai, we have to have one single model
        in the agent. And because we want to be able to use data-parallel,
        we need to have one single forward() method.
        Therefore the operation_type can be either 'encode' or 'score'.

        :param ctxt_tokens:
            tokenized contexts
        :param cand_tokens:
            tokenized candidates
        :param ctxt_rep:
            (bsz x num_codes x hsz)
            encoded representation of the context. If self.type == 'codes', these
            are the context codes. Otherwise, they are the outputs from the
            encoder
        :param ctxt_rep_mask:
            mask for ctxt rep
        :param ctxt_pos:
            position embeddings for the ctxt_rep. If self.type == 'codes', these
            are None, as their use is earlier in the pipeline.
        :param cand_rep:
            encoded representation of the candidates
        """
        if ctxt_tokens is not None or cand_tokens is not None:
            return self.encode(ctxt_tokens, cand_tokens)
        elif (
            ctxt_rep is not None and ctxt_rep_mask is not None and cand_rep is not None
        ):
            # ctxt_pos can be none, if we are using codes (not first M)
            return self.score(ctxt_rep, ctxt_rep_mask, ctxt_pos, cand_rep)
        raise Exception('Unsupported operation')

class PolyBasicAttention(BasicAttention):
    """
    Override basic attention to account for edge case for polyencoder.
    """

    def __init__(self, poly_type, n_codes, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.poly_type = poly_type
        self.n_codes = n_codes

    def forward(self, *args, **kwargs):
        """
        Forward pass.

        Account for accidental dimensionality reduction when num_codes is 1 and the
        polyencoder type is 'codes'
        """
        lhs_emb = super().forward(*args, **kwargs)
        if self.poly_type == 'codes' and self.n_codes == 1 and len(lhs_emb.shape) == 2:
            lhs_emb = lhs_emb.unsqueeze(self.dim - 1)
        return lhs_emb

class IRFriendlyPolyencoderAgent(AddLabelFixedCandsTRA, PolyencoderAgent):
    """
    Poly-encoder agent that allows for adding label to fixed cands.
    """

    @classmethod
    def add_cmdline_args(cls, argparser):
        """
        Add cmd line args.
        """
        super(AddLabelFixedCandsTRA, cls).add_cmdline_args(argparser)
        super(PolyencoderAgent, cls).add_cmdline_args(argparser)


from django.conf.urls import url, include
from django.contrib import admin
from short_app.views import BookmarkCreateView, ShortView, UserCreateView, BookmarkUpdateView

urlpatterns = [
    url(r'^admin/', admin.site.urls),
    url('^', include('django.contrib.auth.urls')),
    url(r'^create_user/$', UserCreateView.as_view(), name='user_create_view'),
    url(r'^$', BookmarkCreateView.as_view(), name='bookmark_create_view'),
    url(r'^(?P<pk>\d+)$', BookmarkUpdateView.as_view(), name='bookmark_update_view'),
    url(r'^(?P<short_url>\w+)/$', ShortView.as_view(), name='short_view'),

]


from collections import namedtuple

class Obstacle(object):
    def __init__(self, obs_id, obs_type):
        self.id = obs_id
        self.type = obs_type
        self.data_pt = namedtuple(
            'data_pt', 'timestamp, x, y, global_x, global_y, vx, vy, length, width, height')
        self.column_names = list(self.data_pt._fields)
        self.data = []
        return

    def append(self, timestamp, x, y, global_x, global_y, vx, vy, length, width, height):
        d = self.data_pt(timestamp, x, y, global_x, global_y,
                         vx, vy, length, width, height)
        self.data.append(d)
        return

    def get(self):
        return self.data

    def concatenate(self, new_obs):
        assert self.__class__ == new_obs.__class__, "Different class cannot concatenate"
        assert self.data.shape[1] == new_obs.get(
        ).shape[1], "Different shape cannot concatenate"
        return self.data.append(new_obs.get())

    def __str__(self):
        return '%s(%s, %s)' % (self.__class__.__name__, self.id, self.type)

class LaneMarker(object):
    def __init__(self, lane_pos):
        assert lane_pos in ['left', 'right', 'leftleft',
                            'rightright'], "Unrecognized LaneType."
        self.lane_pos = lane_pos
        self.data_pt = namedtuple(
            'data_pt', 'timestamp, c0, c1, c2, c3, longitude_start, longitude_end')
        self.column_names = list(self.data_pt._fields)
        self.data = []
        return

    def append(self, timestamp, c0, c1, c2, c3, longitude_start, longitude_end):
        d = self.data_pt(timestamp, c0, c1, c2, c3,
                         longitude_start, longitude_end)
        self.data.append(d)
        return

    def get(self):
        return self.data

    def __str__(self):
        return '%s(%s)' % (self.__class__.__name__, self.lane_pos)

# Testing
if __name__ == "__main__":
    test_obs = Obstacle(1, 'VEHICLE')
    test_obs.append(timestamp=123, x=1, y=2, vx=0.1, vy=0.0)
    print(test_obs)
    print(test_obs.get())

    test_lane = LaneMarker('left')
    test_lane.append(timestamp=123, c0=1, c1=1, c2=1, c3=1,
                     longitudinal_start=0, longitudinal_end=0)
    print(test_lane)
    print(test_lane.get())


import sqlite3
connection = sqlite3.connect("game.db")

cursor = connection.cursor()

def add(name, score, difficulty):
    sql_command = """INSERT INTO data (id ,name, score, difficulty)
        VALUES (NULL, "{}",{},'{}');""".format(name.lower(),score, difficulty.lower())
    cursor.execute(sql_command)

    connection.commit()
    connection.close()


"""Module for annotating Raw objects and other data."""
import mne

def annotate_trials(
    raw: mne.io.BaseRaw,
    keyword: str = "EMG",
    inplace: bool = False,
    keep_original_annotations: bool = True,
) -> mne.io.BaseRaw:
    """Create squared data (0s and 1s) from events and add to Raw object.

    Parameters
    ----------
    raw : MNE Raw object
        The MNE Raw object for this function to modify.
    keyword : str
        Keyword for the events to be added.
    inplace : bool. Default: False
        Set to True if Raw object should be modified in place.
    keep_original_annotations : bool. Default: True
        Set to False if existing annotations should be discarded.

    Returns
    -------
    raw : MNE Raw object
        The Raw object containing the added squared channel.
    """
    if not inplace:
        raw = raw.copy()

    annotations_orig = raw.annotations.copy()
    raw.plot(
        # scalings="auto",
        block=True,
        title=(
            f"Please annotate {keyword} onset and {keyword} end with the label"
            f" `{keyword}`."
        ),
    )

    events_emg, _ = mne.events_from_annotations(raw, event_id={keyword: 1})
    events_emg[1::2, 2] = -1
    # Set event durations to 0
    events_emg[..., 1] = 0

    annotations_new = mne.annotations_from_events(
        events_emg,
        raw.info["sfreq"],
        event_desc={1: f"{keyword}_onset", -1: f"{keyword}_end"},
        orig_time=annotations_orig.orig_time,
    )
    if keep_original_annotations:
        annotations_new += annotations_orig
    raw.set_annotations(annotations_new)

    raw.plot(
        block=True,
        title=f"{keyword} onset and end annotated. Please check annotations.",
    )
    return raw


from collections import OrderedDict
from functools import update_wrapper

def _get_key(*a, **kw):
    if not kw:
        return hash(a)
    # tuple(sorted(kw.items())) must have a great cost
    return hash(a + tuple(sorted(kw.items())))

def lru_cache(max_size=100):

    # an OrderedDict already maintains the order of inserted items, so it
    # feels more likely to work with a LRU cache then a plain dict (and we
    # dont need timestamps)
    cache = OrderedDict()

    def inner(user_function):
        def wrapper(*a, **kw):
            key = _get_key(*a, **kw)

            if key in cache:
                # key already in cache, let's just pop and insert again to
                # maintain the order (kind of inverse, last in cache are used
                # most)
                value = cache.pop(key)
                cache[key] = value
                return value
            else:
                # key not in cache, first we need to get the call value
                res = user_function(*a, **kw)
                try:
                    # let's try to remove it
                    cache.pop(key)
                except KeyError:
                    # the key doesn't exists, but we need to check if there's
                    # room for a new key (by NOT popping the last item)
                    if len(cache) >= max_size:
                        cache.popitem(False)
                # and the result is cached in the end (well, it was the last
                # used value afaik)
                cache[key] = res
                return res

        wrapper.cache = cache
        update_wrapper(wrapper, user_function)
        return wrapper

    return inner

__all__ = "lru_cache"


# -*- coding: utf-8 -*-
"""
Created on Wed Jul 11 14:39:11 2018

@author: mayn
"""

import matplotlib.pyplot as plt

import DNA_Geometry
import DNA_Object
import DNA_File
import DrawShape
import Util
from shapely.geometry import LineString

def show_room_layout(dna, vec_room2bed, vec_room2window, trans_vec, room_name):
    vec_window2bed = {}
    vec_window2bed['x'] = []
    vec_window2bed['y'] = []
    room = DNA_Object.get_room_by_usagename(dna, room_name)
    #print(dna['solutionId'], room['roomUsageName'])

    if 'doors' in dna:
        door_num, door_pos = DNA_Object.get_door_from_dna(dna)
    else:
        door_center, door_num, door_pos = DNA_Object.get_door_from_room(room)

    if 'walls' in dna:
        wall_num, wall_pos = DNA_Object.get_wall_from_dna(dna)
    else:
        wall_num, wall_pos = DNA_Object.get_wall_from_room(room)

    if 'windows' in dna:
        window_num, window_pos = DNA_Object.get_window_info_from_dna(dna)
    else:
        window_center,window_num,window_pos = DNA_Object.get_window_info_from_room(room)

    '''
    if door_num > 0:
        room_polygon = DNA_Geometry.generate_room_polygon(room)
        door_polygon_list = DNA_Geometry.generate_door_geometry(door_num, door_pos)

        for w in range(len(door_polygon_list)):
            #print('hi',room_polygon.boundary.coords, len(room_polygon.boundary.coords), list(room_polygon.boundary.coords))
            if DNA_Geometry.polygon_intersect_polygon_loosely(room_polygon, door_polygon_list[w]):
                #print('hi',room_polygon.boundary.coords, len(room_polygon.boundary.coords), list(room_polygon.boundary.coords))
                #print('[', w, ']', wall_pos['x'][w], wall_pos['y'][w])
                Util.trans_by_vec(door_pos, trans_vec)
                plt.plot(door_pos['x'][w], door_pos['y'][w], linewidth='1.5')
                '''

    bed_num, bed_center, bed_size, bed_point= DNA_Object.get_obj_info_from_room(room, 318)
    area_center, area = DNA_Object.get_room_area(room)

    if bed_num > 0:
        if ((bed_size[0] > 1000 and bed_size[1] > 1100) or (bed_size[0] > 1100 and bed_size[1] > 1000)):
            print('DNA',dna['solutionId'],'room', room_name, 'bed size maybe incorrect! bed_size is:',bed_size[0]*2,'x',bed_size[1]*2)
            return
        vec_room2bed['x'].append(bed_center[0]-area_center[0])
        vec_room2bed['y'].append(bed_center[1]-area_center[1])

        Util.trans_by_vec(bed_point, trans_vec)
        Util.trans_by_vec(area_center, trans_vec)
        Util.trans_by_vec(bed_center, trans_vec)
        plt.plot(bed_point['x'], bed_point['y'], linewidth='0.5')
        #print(bed_center)
        plt.plot(bed_center[0], bed_center[1], 'o')

    wardrobe_num, wardobe_center, wardrobe_size, wardrobe_point = DNA_Object.get_obj_info_from_room(room, 120)
    if wardrobe_num > 0:
        Util.trans_by_vec(wardrobe_point, trans_vec)
        Util.trans_by_vec(wardobe_center, trans_vec)
        plt.plot(wardrobe_point['x'], wardrobe_point['y'], linewidth='0.5')
        #print(bed_center)
        plt.plot(wardobe_center[0], wardobe_center[1], 'o')
        if (wardobe_center[0] > 600):
            print(dna['solutionId'])

#------------------------------------------------------------------------------
if __name__ == '__main__':

    path = "E:/simulation_tools/simulation_py/dnafiles";
    dna_file_list = DNA_File.find_json_files(path)
    count = 0
    #dna_list = ['suiyueruge.txt', 'lanseduonaohe.txt', 'xizhaomisuli.txt', 'qinqishuhua.txt', 'bimozhiyan.txt']
    #dna_list = ['suiyueruge.txt', 'lanseduonaohe.txt', 'xizhaomisuli.txt']
    vec_room2bed = {}
    vec_room2bed['x'] = []
    vec_room2bed['y'] = []

    vec_door2bed = {}
    vec_door2bed['x'] = []
    vec_door2bed['y'] = []

    flag = [False, False, False, True]
    for i in range(len(dna_file_list)):

        dna = DNA_File.load_dna_by_file_name(dna_file_list[i])

        if 'windows' in dna:
            if flag[0] == False and DrawShape.draw_house_area(dna):
                flag[0] = True

            if flag[1] == False and DrawShape.draw_house_door(dna):
                flag[1] = True

            if flag[2] == False and DrawShape.draw_house_window(dna):
                flag[2] = True

            if flag[3] == False and DrawShape.draw_house_wall(dna):
                flag[3] = True

        #处理户型平移（同一户型的位置坐标不一样，由于户型只画了一次，家具信息每次都画，因此家具相对于户型的位置不准确，需做平移处理）
        trans_vec1 = [0, 0]
        trans_vec2 = [0, 0]

        if (i == 0):
            base_room_center1, base_area1 = DNA_Object.get_room_area(DNA_Object.get_room_by_usagename(dna, '主卧'))
            base_room_center2, base_area2 = DNA_Object.get_room_area(DNA_Object.get_room_by_usagename(dna, '次卧'))
        else:
            cur_room_center1, cur_area1 = DNA_Object.get_room_area(DNA_Object.get_room_by_usagename(dna, '主卧'))
            cur_room_center2, cur_area2 = DNA_Object.get_room_area(DNA_Object.get_room_by_usagename(dna, '次卧'))
            trans_vec1 = Util.get_trans_vector(base_room_center1, cur_room_center1)
            trans_vec2 = Util.get_trans_vector(base_room_center2, cur_room_center2)

        #show_room_layout(dna, vec_room2bed, trans_vec1, '主卧')
        #show_room_layout(dna, vec_room2bed, trans_vec2, '次卧')
        count += 1
        #print("-------------------------------------------------------------\n")
    DrawShape.draw_relative_info(vec_room2bed)
    print(count, 'files processed.\n')
    print(vec_room2bed)


import torch.utils.data as data
import torch
import torchvision.transforms as transforms
from torchvision.utils import save_image
from PIL import Image, ImageFilter
import os
import numpy as np
import pandas as pd
import random
from scipy.stats import norm
import json
import glob
import csv
import argparse

from torch.utils.data import DataLoader
from tqdm import tqdm
from config import Config

cfg = Config()

OldFrame_dir = os.path.join(cfg.data_path, 'Frame')
Frame_dir = os.path.join(cfg.data_path, 'Frame_person2')

_ = input('Press enter to confirm old frame path: ' + OldFrame_dir)
_ = input('Press enter to confirm new frame path: ' + Frame_dir)

class ttmData(data.Dataset):
    def __init__(self, root):
        self.root = root
        self.transform = None
        self.data = []

        seg_filenames = glob.glob(os.path.join(root, 'seg', '*.csv'))
        seg_filenames = [x for x in sorted(seg_filenames)]

        box_filenames = glob.glob(os.path.join(root, 'bbox', '*.csv'))
        box_filenames = [x for x in sorted(box_filenames)]

        # seg_filenames = ['./student_data/train/seg/b2735994-59ae-4478-9c5f-93f5ccf18c8b_seg.csv']
        # box_filenames = ['./student_data/train/bbox/b2735994-59ae-4478-9c5f-93f5ccf18c8b_bbox.csv']

        for seg_file, box_file in tqdm(zip(seg_filenames, box_filenames), total=len(seg_filenames), desc='Init the dataloader'):

            assert seg_file.split('/')[-1].split('_')[0] == box_file.split('/')[-1].split('_')[0]

            # print(seg_file.split('/')[-1].split('_')[0])

            seg = pd.read_csv(seg_file)
            box = pd.read_csv(box_file)

            start = seg['start_frame'].to_list()
            end = seg['end_frame'].to_list()

            sampled_frames = [list(range(s, e+1)) for s, e in zip(start, end)]

            # start = seg['start_frame'].to_numpy()[:, np.newaxis]
            # end = seg['end_frame'].to_numpy()[:, np.newaxis]
            # sampled_frames = start.copy()
            # iters = 4*cfg.samples
            # for i in range(1, iters):
            #     sampled_frames = np.hstack((sampled_frames, start+(end-start)*(i/iters)))
            # sampled_frames = sampled_frames.astype(int).tolist()
            last_pid = -1
            for i, (pid, frames) in enumerate(zip(seg['person_id'], sampled_frames)):
                if last_pid != pid:
                    memory = box[box['person_id'] == pid]
                last_pid = pid
                frames = [frame for frame in frames if memory[(memory['frame_id'] == frame)].values.tolist()[
                    0][-1] != -1]
                sampled_frames[i] = frames

            seg['sampled_frames'] = sampled_frames
            seg.drop(columns=['start_frame', 'end_frame'], inplace=True)
            bbox = []
            last_pid = -1
            for pid, frames in zip(seg['person_id'], sampled_frames):
                if last_pid != pid:
                    memory = box[box['person_id'] == pid]
                last_pid = pid
                # print(memory)
                # print(memory[(memory['frame_id'] == 1)].values.tolist()[0][-4:])
                bbox.append([memory[(memory['frame_id'] == frame)].values.tolist()[0][-4:] for frame in frames])

            seg['bbox'] = bbox
            seg['hash'] = seg_file.split('/')[-1].split('_')[0]
            # seg.drop(columns=['person_id'], inplace=True)
            # print(seg.columns)
            self.data.extend(seg.values.tolist())
            # print(self.data[0])
            # exit()

    def __getitem__(self, index):

        if 'train' in self.root:
            pid, label, frames, bboxs, fname = self.data[index]
        else:
            pid, frames, bboxs, fname = self.data[index]

        if len(frames) == 0:
            return 0

        for i, frame in enumerate(frames):
            if len(glob.glob(os.path.join(Frame_dir, f'{fname}_{frame}_{pid}.jpg'))) > 0:
                continue
            img = Image.open(os.path.join(OldFrame_dir, f'{fname}_{frame}.jpg'))

            # bbox = bboxs[i]
            x1, y1, x2, y2 = bboxs[i]
            x1, x2 = min(x1, x2), max(x1, x2)
            y1, y2 = min(y1, y2), max(y1, y2)
            avg_x, avg_y = (x1 + x2) / 2., (y1 + y2) / 2.
            x1 = (x1 - avg_x) * 1.3 + avg_x
            x2 = (x2 - avg_x) * 1.3 + avg_x
            y1 = (y1 - avg_y) * 1.4 + avg_y
            y2 = (y2 - avg_y) * 1.4 + avg_y

            top, height = int(y1), int(y2 - y1)
            left, width = int(x1), int(x2 - x1)

            img = transforms.functional.crop(img, top, left, height, width)
            # img = transforms.functional.resized_crop(
            #     img, top, left, height, width, size=(int(shape[0]*scale), cfg.input_size))

            # if (index % 10) == 0:
            #     img.save("demo.jpg", "JPEG")
            img.save(os.path.join(Frame_dir, f'{fname}_{frame}_{pid}.jpg'), "JPEG")
            # exit()

            # print(imgs.shape, img.shape)
            # imgs = torch.cat((imgs, img), 1)

        return 0

    def __len__(self):
        return len(self.data)

if __name__ == '__main__':
    train_data = ttmData(root=os.path.join(cfg.data_path, 'train'))
    train_loader = DataLoader(train_data, batch_size=1, shuffle=False,
                              num_workers=cfg.num_workers, pin_memory=False, drop_last=False)
    test_data = ttmData(root=os.path.join(cfg.data_path, 'test'))
    test_loader = DataLoader(test_data, batch_size=1, shuffle=False,
                             num_workers=cfg.num_workers, pin_memory=False, drop_last=False)
    for epoch in range(1):
        iters = len(train_loader)
        with tqdm(train_loader, unit="batch", desc='Preprocessing_train') as tepoch:
            for i, (trivial_data) in enumerate(tepoch):
                tepoch.set_description(f"Preprocessing")
                trivial_data = trivial_data.to(cfg.device)

    for epoch in range(1):
        iters = len(test_loader)
        with tqdm(test_loader, unit="batch", desc='Preprocessing_test') as tepoch:
            for i, (trivial_data) in enumerate(tepoch):
                tepoch.set_description(f"Preprocessing")
                trivial_data = trivial_data.to(cfg.device)


import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F

import time, random, numpy as np, argparse
import pandas as pd
import numpy as np
from CovostDataset import CovostDataset
from transfomer import PytorchTransformer

def generate_square_subsequent_mask(sz, device):
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

def create_mask(src, tgt, pad_idx, device):

    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device)
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

    src_padding_mask = (src == pad_idx).transpose(0, 1)
    tgt_padding_mask = (tgt == pad_idx).transpose(0, 1)
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tsv_root = '/datasets/CS678/es/train_st_es_en.tsv'
zip_fbank_path = '/datasets/CS678/es/fbank80.zip'
corpus_path = '/datasets/CS678/es/corpus_train_es_en.txt'

parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", type=int, default=8)
parser.add_argument("--src_lan", type=str, default='es', help='source language (can be en, es, fr, fa, ca)')
parser.add_argument("--tgt_lan", type=str, default='en', help='target language (can be en, ca, fa)')
parser.add_argument("--vocab_size", type=int, default=2000, help='size of vocab')
args = parser.parse_args()

train_dataset = CovostDataset(tsv_root=tsv_root, zip_fbank_path=zip_fbank_path, corpus_path=corpus_path, args=args)

train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, collate_fn=train_dataset.collate_fn)

# import ipdb
# ipdb.set_trace()
#  vocabs = [train_dataset.tokenizer.sp.id_to_piece(id) for id in range(train_dataset.tokenizer.sp.get_piece_size())]

embed_size = 80
heads = 20
inner_layer_size = 512
vocab_size = args.vocab_size
n_blocks = 5
padding_indx = 0

batch_size = 128
n_epochs = 50

model = PytorchTransformer(n_blocks, n_blocks, embed_size,
                                 heads, vocab_size, vocab_size, inner_layer_size)

model.to(DEVICE)

optimizer = torch.optim.AdamW(model.parameters(), lr=0.00001, betas=(0.9, 0.98), eps=1e-9)
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=padding_indx)

# import ipdb
# ipdb.set_trace()

model.train()

for epoch in range(n_epochs):

    losses = 0
    for i_batch, sample_batched in enumerate(train_dataloader):
        b_waves, b_frames, b_texts, b_speakers = sample_batched[0]['fbank_waves'], sample_batched[0]['n_frames'], \
            sample_batched[0]['tgt_texts'], sample_batched[0]['speakers']

        tgt = torch.tensor(b_texts)

        src = b_waves.permute(1, 0, 2).to(DEVICE)
        tgt = tgt.permute(1, 0).to(DEVICE)

        # import ipdb
        # ipdb.set_trace()

        tgt_input = tgt[:-1, :]

        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src[:, :, 0], tgt_input, padding_indx, DEVICE)

        logits = model(src, tgt_input, src_mask, tgt_mask,src_padding_mask, tgt_padding_mask, src_padding_mask)

        optimizer.zero_grad()

        tgt_out = tgt[1:, :]

        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        loss.backward()

        optimizer.step()
        losses += loss.item()

        # if i_batch % 500 == 0:
        print(f"Epoch {epoch}, batch {i_batch}, loss {loss.item()}")

    print("---0---------------0--------------------0---------------0-------")
    print(f"Epoch {epoch}, loss {losses/len(train_dataloader)}\n")

    file_path = "/home/jmbuya/nlp/cs678_project/src/models"
    torch.save(model.state_dict(), f"{file_path}/model_lr_0.00001_e{epoch}_3.pt")

    # import ipdb
    # ipdb.set_trace()


from config import config
from state import getOrCreateUser, findUserByNick

def parseAmount(text, min=0, max=config["max-amount"]):
	try:
		if text.find(",") != -1:
			text = text.replace(",", ".")
		val = float(text)
		if val < 0:
			return None, "less than zero"
		elif val == 0:
			return None, "zero"
		elif val > max:
			return None, "larger than the maximum allowed amount"
		elif int(val * 100) / float(100) != val:
			return None, "too precise ({} would turn into {})" \
				.format(val, int(val * 100) / float(100))
		else:
			return int(val * 100), None
	except:
		return None, "not a number"

ARG_TEXT = 0
ARG_INT = 1
ARG_AMOUNT = 2
ARG_USER = 3
ARG_REST = 4
def parseArgs(msg, argDef, defaults, usage=""):
	split = msg.text.split(" ")
	result = []
	error = None

	offset = len(split[0]) + 1
	split = split[1 : ]

	for i, expected in enumerate(argDef):
		if i < len(split):
			arg = split[i]
		elif i < len(defaults) and defaults[i] is not None:
			result.append(defaults[i])
			continue
		else:
			error = "Argument {} not specified".format(i)
			break

		if expected == ARG_TEXT:
			result.append(arg)
		elif expected == ARG_INT:
			try:
				val = int(arg)
				result.append(val)
			except:
				error = "Argument {} should be an int but is '{}'".format(i, arg)
				break
		elif expected == ARG_AMOUNT:
			val, error = parseAmount(arg)
			if val is None:
				error = "Argument {} should be an amount but is {}".format(i, error)
				break
			result.append(val)
		elif expected == ARG_USER:
			user = None
			for entity in msg.entities:
				if entity.offset == offset:
					if entity.type == "mention":
						user = findUserByNick(arg[1 : ])
						break
					elif entity.type == "text_mention":
						user = getOrCreateUser(entity.user)
						break

			if user is None:
				error = "Argument {} should be an user but is '{}'".format(i, arg)
				break

			result.append(user)
		elif expected == ARG_REST:
			result.append(" ".join(split[i : ]))
			break

		offset = offset + len(arg) + 1

	if error is None:
		return result
	else:
		error = error + usage
		msg.reply_text(error)
		raise Exception(error)


#  Author:   Niels Nuyttens  <niels@nannyml.com>
#
#  License: Apache Software License 2.0

"""Unit tests for drift ranking."""

import pandas as pd
import pytest

from nannyml import BinaryClassificationMetadata, Feature, FeatureType
from nannyml.drift import UnivariateDriftResult
from nannyml.drift.ranking import AlertCountRanking
from nannyml.exceptions import InvalidArgumentsException

@pytest.fixture
def sample_drift_result() -> UnivariateDriftResult:  # noqa: D103
    return UnivariateDriftResult(
        analysis_data=[],
        model_metadata=BinaryClassificationMetadata(),
        drift_data=pd.DataFrame(
            {
                'f1_alert': [0, 0, 0, 0, 1, 1],
                'f2_alert': [0, 0, 0, 1, 1, 1],
                'f3_alert': [0, 0, 0, 1, 0, 1],
                'f4_alert': [0, 0, 0, 0, 0, 0],
                'f1': [0, 0, 0, 0, 0, 0],
                'f2': [1, 1, 1, 1, 1, 1],
                'f3': [0, 0, 0, 0, 0, 0],
                'f4': [1, 1, 1, 1, 1, 1],
            }
        ),
    )

@pytest.fixture
def sample_metadata(sample_drift_result):  # noqa: D103
    md = BinaryClassificationMetadata(
        features=[
            Feature(column_name='f1', label='f1', feature_type=FeatureType.CONTINUOUS),
            Feature(column_name='f2', label='f2', feature_type=FeatureType.CONTINUOUS),
            Feature(column_name='f3', label='f3', feature_type=FeatureType.CONTINUOUS),
            Feature(column_name='f4', label='f4', feature_type=FeatureType.CONTINUOUS),
        ]
    )
    md.predicted_probability_column_name = 'y_pred_proba'
    return md

def test_alert_count_ranking_raises_invalid_arguments_exception_when_drift_result_is_empty(  # noqa: D103
    sample_metadata,
):
    ranking = AlertCountRanking()
    with pytest.raises(InvalidArgumentsException, match='drift results contain no data to use for ranking'):
        ranking.rank(
            UnivariateDriftResult(
                analysis_data=[], drift_data=pd.DataFrame(columns=['f1', 'f1_alert']), model_metadata=sample_metadata
            ),
            sample_metadata,
        )

def test_alert_count_ranking_contains_rank_column(sample_drift_result, sample_metadata):  # noqa: D103
    ranking = AlertCountRanking()
    sut = ranking.rank(sample_drift_result, sample_metadata)
    assert 'rank' in sut.columns

def test_alert_count_ranks_by_sum_of_alerts_per_feature(sample_drift_result, sample_metadata):  # noqa: D103
    ranking = AlertCountRanking()
    sut = ranking.rank(sample_drift_result, sample_metadata)
    assert sut.loc[sut['rank'] == 1, 'feature'].values[0] == 'f2'
    assert sut.loc[sut['rank'] == 2, 'feature'].values[0] == 'f1'
    assert sut.loc[sut['rank'] == 3, 'feature'].values[0] == 'f3'
    assert sut.loc[sut['rank'] == 4, 'feature'].values[0] == 'f4'

def test_alert_count_ranking_should_exclude_zero_alert_features_when_exclude_option_set(  # noqa: D103
    sample_drift_result,
    sample_metadata,
):
    ranking = AlertCountRanking()
    sut = ranking.rank(sample_drift_result, sample_metadata, only_drifting=True)
    assert len(sut) == 3
    assert len(sut[sut['feature'] == 'f4']) == 0

def test_alert_count_ranking_should_raise_invalid_arguments_exception_when_given_wrong_drift_results(  # noqa: D103
    sample_metadata,
):
    with pytest.raises(InvalidArgumentsException, match="drift results are not statistical drift results."):
        _ = AlertCountRanking().rank(
            UnivariateDriftResult(
                analysis_data=[],
                drift_data=pd.DataFrame({'alert': [False, False, True]}),
                model_metadata=sample_metadata,
            ),
            sample_metadata,
        )


bl_info = {
    "name": "gltf_auto_export",
    "author": "kaosigh",
    "version": (0, 6, 0),
    "blender": (3, 4, 0),
    "location": "File > Import-Export",
    "description": "glTF/glb auto-export",
    "warning": "",
    "wiki_url": "https://github.com/kaosat-dev/Blender_bevy_components_workflow",
    "tracker_url": "https://github.com/kaosat-dev/Blender_bevy_components_workflow/issues/new",
    "category": "Import-Export"
}
import bpy
import os

from bpy.app.handlers import persistent
from bpy.props import (IntProperty)

from . import helpers
from .internals import (SceneLink,
                        SceneLinks,
                        CollectionToExport,
                        CollectionsToExport,
                        CUSTOM_PG_sceneName
                        )
from .auto_export import auto_export
from .preferences import (AutoExportGltfPreferenceNames,
                          AutoExportGltfAddonPreferences
                          )
from .ui.main import (GLTF_PT_auto_export_main,
                      GLTF_PT_auto_export_root,
                      GLTF_PT_auto_export_blueprints,
                      GLTF_PT_auto_export_collections_list,
                      GLTF_PT_auto_export_gltf,
                      SCENE_UL_GLTF_auto_export,
                      AutoExportGLTF
                      )
from .ui.various import (SCENES_LIST_OT_actions)
from .helpers_scenes import (is_scene_ok)

bpy.context.window_manager['changed_objects_per_scene'] = {}
bpy.context.window_manager['previous_params'] = {}
bpy.context.window_manager['__gltf_auto_export_initialized'] = False
bpy.context.window_manager['__gltf_auto_export_gltf_params_changed'] = False

######################################################
""" there are two places where we load settings for auto_export from:
- in ui/main AutoExportGLTF -> invoke
- in auto_export.py -> auto_export
This is a workaround needed because of the way the settings are stored , perhaps there is a better way to deal with it ? ie by calling the AutoExportGLTF operator from the auto_export function ?
"""

#see here for original gltf exporter infos https://github.com/KhronosGroup/glTF-Blender-IO/blob/main/addons/io_scene_gltf2/__init__.py
@persistent
def deps_update_handler(scene, depsgraph):
    if scene.name != "temp_scene": # actually do we care about anything else than the main scene(s) ?
        print("depsgraph_update_post", scene.name)
        print("-------------")
        changed = scene.name or ""

        # depsgraph = bpy.context.evaluated_depsgraph_get()
        if not 'changed_objects_per_scene' in bpy.context.window_manager:
            bpy.context.window_manager['changed_objects_per_scene'] = {}

        if not changed in bpy.context.window_manager['changed_objects_per_scene']:
            bpy.context.window_manager['changed_objects_per_scene'][changed] = {}

        for obj in depsgraph.updates:
            if isinstance(obj.id, bpy.types.Object):
                # get the actual object
                object = bpy.data.objects[obj.id.name]
                bpy.context.window_manager['changed_objects_per_scene'][scene.name][obj.id.name] = object

        bpy.context.window_manager.changedScene = changed

@persistent
def save_handler(dummy):
    print("-------------")
    print("saved", bpy.data.filepath)
    if not 'changed_objects_per_scene' in bpy.context.window_manager:
        bpy.context.window_manager['changed_objects_per_scene'] = {}
    changes_per_scene =  bpy.context.window_manager['changed_objects_per_scene']

    #determine changed parameters
    addon_prefs = bpy.context.preferences.addons["gltf_auto_export"].preferences

    prefs = {}
    for (k,v) in addon_prefs.items():
        if k not in AutoExportGltfPreferenceNames:
            prefs[k] = v

    previous_params = bpy.context.window_manager['previous_params'] if 'previous_params' in bpy.context.window_manager else {}
    set1 = set(previous_params.items())
    set2 = set(prefs.items())
    difference = dict(set1 ^ set2)

    changed_param_names = list(set(difference.keys())- set(AutoExportGltfPreferenceNames))
    changed_parameters = len(changed_param_names) > 0
    # do the export
    auto_export(changes_per_scene, changed_parameters)

    # save the parameters
    # todo add back
    for (k, v) in prefs.items():
        bpy.context.window_manager['previous_params'][k] = v

    # reset a few things after exporting
    # reset wether the gltf export paramters were changed since the last save
    bpy.context.window_manager['__gltf_auto_export_gltf_params_changed'] = False
    # reset whether there have been changed objects since the last save
    bpy.context.window_manager['changed_objects_per_scene'] = {}

def get_changedScene(self):
    return self["changedScene"]

def set_changedScene(self, value):
    self["changedScene"] = value

classes = [
    SceneLink,
    SceneLinks,
    CUSTOM_PG_sceneName,
    SCENE_UL_GLTF_auto_export,
    SCENES_LIST_OT_actions,

    AutoExportGLTF,
    AutoExportGltfAddonPreferences,

    CollectionToExport,
    CollectionsToExport,

    GLTF_PT_auto_export_main,
    GLTF_PT_auto_export_root,
    GLTF_PT_auto_export_blueprints,
    GLTF_PT_auto_export_collections_list,
    GLTF_PT_auto_export_gltf
]

def menu_func_import(self, context):
    self.layout.operator(AutoExportGLTF.bl_idname, text="glTF auto Export (.glb/gltf)")

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

    bpy.types.Scene.main_scene = bpy.props.PointerProperty(type=bpy.types.Scene, name="main scene", description="main_scene_chooser", poll=is_scene_ok)
    bpy.types.Scene.library_scene = bpy.props.PointerProperty(type=bpy.types.Scene, name="library scene", description="library_scene_picker", poll=is_scene_ok)

    # setup handlers for updates & saving
    bpy.app.handlers.depsgraph_update_post.append(deps_update_handler)
    bpy.app.handlers.save_post.append(save_handler)

    bpy.types.WindowManager.changedScene = bpy.props.StringProperty(get=get_changedScene, set=set_changedScene)
    bpy.types.WindowManager.exportedCollections = bpy.props.CollectionProperty(type=CollectionsToExport)

    # add our addon to the toolbar
    bpy.types.TOPBAR_MT_file_export.append(menu_func_import)

    ## just experiments
    bpy.types.Scene.main_scenes_list_index = IntProperty(name = "Index for main scenes list", default = 0)
    bpy.types.Scene.library_scenes_list_index = IntProperty(name = "Index for library scenes list", default = 0)

    """
    mock_main_scenes = []
    main_scenes = bpy.context.preferences.addons["gltf_auto_export"].preferences.main_scenes
    for item_name in mock_main_scenes:
        item = main_scenes.add()
        item.name = item_name

    mock_library_scenes = []
    library_scenes = bpy.context.preferences.addons["gltf_auto_export"].preferences.library_scenes
    for item_name in mock_library_scenes:
        item = library_scenes.add()
        item.name = item_name"""

    bpy.context.preferences.addons["gltf_auto_export"].preferences.main_scenes_index = 0
    bpy.context.preferences.addons["gltf_auto_export"].preferences.library_scenes_index = 0

def unregister():
    for cls in classes:
        bpy.utils.unregister_class(cls)

    bpy.types.TOPBAR_MT_file_export.remove(menu_func_import)

    # remove handlers & co
    bpy.app.handlers.depsgraph_update_post.remove(deps_update_handler)
    bpy.app.handlers.save_post.remove(save_handler)

    del bpy.types.WindowManager.changedScene
    del bpy.types.WindowManager.exportedCollections

    del bpy.types.Scene.main_scene
    del bpy.types.Scene.library_scene

    del bpy.types.Scene.main_scenes_list_index
    del bpy.types.Scene.library_scenes_list_index

if "gltf_auto_export" == "__main__":
    register()


from external_api.dbx import RussDropBox
from external_api.gcp_pubsub import ImgPathPubisher
from processing.image_publishing import ImgPublisher
from processing.image_labeller_pipeline import ImgLabelPipeLine
from processing.image_labeller_pipeline import ImgLabelPersitanceMock
import configparser
import os
import logging
import sys

def utils_read_dbx_cfg_from_gcs():
    from google.cloud import storage
    storage_client = storage.Client()
    bucket = storage_client.bucket('application-core-config')
    blob = bucket.blob('core-config.ini')
    return blob.download_as_string().decode('utf-8')

def setup_logger():
    logging.basicConfig(
    filename='process_pipeline.log',
    level=logging.INFO,
    format='[%(asctime)s]{%(pathname)s:%(lineno)d}%(levelname)s- %(message)s',
    datefmt='%H:%M:%S'
    )
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))

if __name__ == "__main__":

    # Setup logger
    setup_logger()

    # read app config
    config = configparser.ConfigParser()
    config.read('./config/main.ini')
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"]=config['DEFAULT']['CredentialsFile']
    logging.info('gcp credentials set to be read from file {}'.format(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]))
    config.read_string(utils_read_dbx_cfg_from_gcs())

    # Connect to dropbox, find all images and publish to subsub
    publisher = ImgPathPubisher(config['DEFAULT']['ProjectName'], config['DEFAULT']['ProcessEntryTopic'])
    dropb = RussDropBox(config['DROPBOX-SECRETS']['Token'], max_file_count = int(config['DROPBOX']['MaxSize']), batch_size = int(config['DROPBOX']['BatchSize']))

    processor = ImgPublisher(dropb, publisher)
    mockStorage  =ImgLabelPersitanceMock('mock','mock')
    lblPipeLine = ImgLabelPipeLine(dropb, mockStorage)

    for path in processor.publish_dbx_library(True):
        lblPipeLine.label(path)

    # test write to bigquery
    # import external_api.gcp_biquery
    # db = external_api.gcp_biquery.BQWriter('image_labels', 'labels')
    # db.write_image_data('russ test', 'default', 100)
    # print(db.image_data_exists('/camera uploads/2012-08-01 17.43.26.jpg'))

    # db.remove_duplicates()
    # db.delete_image_data('russ test')