from __future__ import annotations

import warnings

import numpy as np
import LabelLib as ll
import IMP.core
import IMP.atom

import scikit_fluorescence as skf
import scikit_fluorescence.io.structure
import scikit_fluorescence.math.datatools
import scikit_fluorescence.decay.rate_spectra

from scikit_fluorescence.modeling.dye.av import functions
from scikit_fluorescence import typing


class BasicAV(object):
    """Simulates the accessible volume of a dye

    Attributes
    ----------
    _av : LabelLib.Grid3D
        The X coordinate
    _density : numpy.array
        The Y coordinate
    attachment_coordinates : numpy.array
        The Y coordinate
    verbose : bool
        The Y coordinate
    _atoms : numpy.array
        The Y coordinate

    Examples
    --------
    Create AV from coordinate files

    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.io
    >>> import scikit_fluorescence.modeling
    >>> atoms = skf.io.structure.read_coordinates('./test/data/hgbp1/protein.pdb')
    >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, chain='A', atom_name='CB')
    >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, chain='B', atom_name='CB')
    >>> av1.dRDA(av2)
    53.66615067748922

    Create AV from IMP Hierarchy
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.modeling
    >>> import IMP.atom
    >>> model = IMP.Model()
    >>> filename = './test/data/hgbp1/protein.pdb'
    >>> mp = IMP.atom.read_pdb(filename, model, IMP.atom.NonWaterPDBSelector())
    >>> av1 = skf.modeling.dye.av.BasicAV(atoms=None, hierarchy=mp, residue_seq_number=577, atom_name='CB')
    >>> av1.dRDA(av2)
    53.66615067748922

    """

    _av: ll.Grid3D = None
    _density: np.ndarray = None
    attachment_coordinates: np.ndarray = None
    verbose: bool = True
    _atoms: np.ndarray = None

    def __init__(
            self,
            atoms: np.ndarray,
            simulation_grid_resolution: float = 0.5,
            allowed_sphere_radius: float = 1.5,
            radius1: float = 3.5,
            radius2: float = 4.5,
            radius3: float = 1.5,
            linker_width: float = 0.5,
            linker_length: float = 20.5,
            simulation_type: str = "AV1",
            hierarchy: IMP.core.Hierarchy = None,
            chain_identifier: str = None,
            residue_name: str = None,
            residue_seq_number: int = None,
            atom_name: str = None,
            position_name: str = None,
            contact_volume_thickness: float = None,
            contact_volume_trapped_fraction: float = None,
            min_sphere_volume_fraction: float = None,
            strip_mask: str = None,
            chain_weighting: bool = None,
            vdw_max: float = 2.5,
            **kwargs
    ):
        # Raise warnings for optional options that are not implemented
        if isinstance(contact_volume_thickness, float):
            warnings.warn("Contact_volume_thickness is not implemented")
        if isinstance(contact_volume_trapped_fraction, float):
            warnings.warn("contact_volume_trapped_fraction is not implemented")
        if isinstance(contact_volume_trapped_fraction, float):
            warnings.warn("contact_volume_trapped_fraction is not implemented")
        if isinstance(min_sphere_volume_fraction, float):
            warnings.warn("min_sphere_volume_fraction is not implemented")
        if isinstance(strip_mask, str):
            warnings.warn("strip_mask is not implemented")
        if isinstance(chain_weighting, bool):
            warnings.warn("chain_weighting is not implemented")

        if isinstance(hierarchy, IMP.core.Hierarchy):
            atoms = scikit_fluorescence.io.structure.convert_atoms(
                IMP.atom.get_leaves(hierarchy)
            )
        self.dg = simulation_grid_resolution
        self.allowed_sphere_radius = allowed_sphere_radius

        self.position_name = position_name
        self.residue_name = residue_name
        self.attachment_residue = residue_seq_number
        self.attachment_atom = atom_name

        self.radius1 = radius1
        self.radius2 = radius2
        self.radius3 = radius3
        self.linker_width = linker_width
        self.linker_length = linker_length
        self.simulation_type = simulation_type
        attachment_atom_index = kwargs.get(
            'attachment_atom_index',
            skf.io.structure.find_atom_index(
                atoms,
                chain_identifier,
                self.attachment_residue,
                self.attachment_atom,
                self.residue_name
            )
        )
        # set attachment atom radius to zero otherwise the
        # search cannot start and yields an empty AV.
        r = np.copy(atoms['radius'])
        r[attachment_atom_index] = 0.0
        # set radii of all obstacles in allow sphere radius to zero to
        # make a bit more room for the dye.
        if allowed_sphere_radius > 0:
            xyz_a, vdw_a, indices = functions.atoms_in_reach(
                xyz=atoms['xyz'],
                vdw=atoms['radius'],
                dmaxsq=(allowed_sphere_radius + vdw_max) ** 2,
                atom_i=attachment_atom_index
            )
            r[indices] = 0.0
        x = atoms['xyz'][:, 0]
        y = atoms['xyz'][:, 1]
        z = atoms['xyz'][:, 2]
        xyzr = np.vstack([x, y, z, r])
        dye_attachment_point = np.array(
            [
                x[attachment_atom_index],
                y[attachment_atom_index],
                z[attachment_atom_index]
            ]
        )
        if self.verbose:
            print("Computing AV")
            print("-- allowed_sphere_radius: %s" % allowed_sphere_radius)
            print("-- simulation_type: %s" % simulation_type)
            print("-- dye_attachment_point: %s" % dye_attachment_point)
            print("-- linker_length: %s" % linker_length)
            print("-- linker_width: %s" % linker_width)
            print("-- radius1: %s" % radius1)
            print("-- radius2: %s" % radius2)
            print("-- radius3: %s" % radius3)
            print("-- simulation_grid_resolution: %s" % simulation_grid_resolution)
        if self.simulation_type == 'AV3':

            av = ll.dyeDensityAV3(
                xyzr,
                dye_attachment_point,
                linker_length,
                linker_width,
                [radius1, radius2, radius3],
                simulation_grid_resolution
            )
        else:
            av = ll.dyeDensityAV1(
                xyzr,
                dye_attachment_point,
                linker_length,
                linker_width,
                radius1,
                simulation_grid_resolution
            )
        self._atoms = atoms
        self._density = np.array(av.grid, dtype=np.float64).reshape(av.shape)
        self.attachment_coordinates = dye_attachment_point
        self._av = av

    @property
    def ng(self):
        return self.density.shape[0]

    @property
    def density(self):
        return self._density

    @property
    def points(self):
        return self._av.points().T

    def save(
            self,
            filename: str,
            mode: str = 'xyz',
            **kwargs
    ):
        """Saves the accessible volume as xyz-file or open-dx density file

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=1, atom_name='N')
        >>> av.save('test', reading_routine='xyz')
        >>> av.save('test', reading_routine='dx')
        >>> av.save('test', reading_routine='mrc')
        """
        density = kwargs.get('density', self.density)
        if mode == 'dx':
            ng = self.ng
            dg = self.dg
            offset = (ng - 1) / 2 * dg
            skf.io.density.write_open_dx(
                filename,
                density,
                self.attachment_coordinates - offset,
                ng, ng, ng,
                dg, dg, dg
            )
        elif mode == 'mrc':
            ng = self.ng
            dg = self.dg
            offset = (ng - 1) / 2 * dg
            skf.io.density.write_mrc(
                filename,
                density,
                self.attachment_coordinates - offset,
                ng, ng, ng,
                dg
            )
        else:
            p = kwargs.get('points', self.points)
            d = p[:, [3]].flatten()
            d /= max(d) * 50.0
            xyz = p[:, [0, 1, 2]]
            scikit_fluorescence.io.structure.write_points(
                filename=filename + '.' + mode,
                points=xyz,
                mode=mode,
                verbose=self.verbose,
                density=d
            )

    def get_fret_rate_constant_spectrum(
            self,
            av: scikit_fluorescence.modeling.dye.av.BasicAV,
            forster_radius: float,
            tau0: float,
            kappa2: float = 2. / 3.,
            rda_axis: np.ndarray = None,
            rda_min: float = 10,
            rda_max: float = 200,
            n_rda_bins = 100,
            use_log: bool = True,
            n_samples: int = 100000
    ):
        p, rda_axis = scikit_fluorescence.modeling.dye.av.functions.histogram_rda(
            av1=self,
            av2=av,
            n_samples=n_samples,
            rda_axis=rda_axis,
            rda_min=rda_min,
            rda_max=rda_max,
            use_log=use_log,
            n_rda_bins=n_rda_bins
        )
        # compute FRET rate constants that belong to the distances
        return skf.math.datatools.two_column_to_interleaved(
            p, skf.decay.rate_spectra.distance_to_fret_rate_constant(
                r=rda_axis[1:],
                forster_radius=forster_radius,
                tau0=tau0,
                kappa2=kappa2
            )
        )

    def get_donor_fret_lifetime_spectrum(
            self,
            donor_lifetime_spectrum: np.ndarray,
            **kwargs
    ):
        """
        :param donor_lifetime_spectrum:
        :param kwargs: parameters passed to get_fret_rate_constant_spectrum
        :return: interleaved spectrum of donor lifetimes in the presence of FRET
        """
        if not isinstance(donor_lifetime_spectrum, np.ndarray):
            donor_lifetime_spectrum = np.array(donor_lifetime_spectrum, dtype=np.float64)
        fret_rate_constant_spectrum = self.get_fret_rate_constant_spectrum(**kwargs)
        return skf.math.datatools.invert_interleaved(
            skf.math.datatools.ere2(
                fret_rate_constant_spectrum,
                skf.math.datatools.invert_interleaved(
                    donor_lifetime_spectrum
                )
            )
        )

    def dRmp(
            self,
            av: BasicAV,
    ):
        """
        Calculate the distance between the mean positions of the AVs

        :param av: accessible volume object
        :return:

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=1, atom_name='N')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, atom_name='CB')
        >>> av1.dRmp(av2)
        """
        return functions.dRmp(self, av)

    def dRDA(
            self,
            av: BasicAV,
            **kwargs
    ):
        """Calculate the average distance to the second accessible volume

        :param av:
        :return:

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, atom_name='CB')
        >>> av1.dRDA(av2)
        """
        return functions.RDAMean(self, av, **kwargs)

    def widthRDA(
            self,
            av: BasicAV,
    ):
        """Calculates the width of a DA-distance distribution

        :param av:
        :return:

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, atom_name='CB')
        >>> av1.widthRDA(av2)

        """
        return functions.widthRDA(self, av)

    def dRDAE(
            self,
            av: BasicAV,
            forster_radius: float
    ):
        """Calculate the FRET-averaged mean distance to the second accessible volume

        :param av: Accessible volume
        :return:

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, atom_name='CB')
        >>> av1.dRDAE(av2)
        """
        return functions.RDAMeanE(self, av, forster_radius)

    def pRDA(
            self,
            av: BasicAV,
            normalize: bool = False,
            **kwargs
    ) -> typing.Tuple[
        np.ndarray,
        np.ndarray
    ]:
        """Calculate a histogram of the distance distribution

        This method computes a histogram of the inter-dye distance distribution
        between two accessible volumes. This function returns the distance / histogram
        axis and corresponding histogram.

        :param av: Accessible volume
        :param kwargs: these parameters are passed to
        :py:meth:`scikit_fluorescence.modeling.dye.av.functions.histogram_rda`
         rda_axis (axis used to compute the histogram), rda_min (minimum distance
         value), rda_max (maximum distance value), n_rda_bins (number of histogram
         bins), use_log (if True computes a logarithmic axis), n_samples (number
         of samples taken to compute the histogram).
        :return: Tuple (histogram, histogram axis)

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        >>> av2 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=577, atom_name='CB')
        >>> y, x = av1.pRDA(av2)

        """
        return functions.histogram_rda(
            self, av, normalize=normalize, **kwargs
        )

    @property
    def Rmp(self) -> np.ndarray:
        """
        The mean position of the accessible volume (average x, y, z coordinate)

        Examples
        --------
        >>> import scikit_fluorescence as skf
        >>> import scikit_fluorescence.io
        >>> import scikit_fluorescence.modeling
        >>> atoms = skf.io.structure.read_coordinates('./test/data/1fat.cif')
        >>> av1 = skf.modeling.dye.av.BasicAV(atoms, residue_seq_number=18, atom_name='CB')
        """
        weights = self.points[:, 3]
        weights /= weights.sum()
        xyz = self.points[:, [0, 1, 2]]
        return np.average(xyz, weights=weights, axis=0)
