from __future__ import annotations
from scikit_fluorescence import typing

import math
import numpy as np
import numba as nb
import scikit_fluorescence as skf


distance_samples: int = 10000


def histogram_rda(
        av1: skf.modeling.dye.av.BasicAV,
        av2: skf.modeling.dye.av.BasicAV,
        rda_axis: np.ndarray = None,
        rda_min: float = 1,
        rda_max: float = 200,
        n_rda_bins = 100,
        use_log: bool = False,
        n_samples: int = distance_samples,
        normalize: bool = False
) -> typing.Tuple[
    np.ndarray,
    np.ndarray
]:
    """Calculate a histogram of the distance distribution

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

    :param av1: Accessible volume
    :param av2: Accessible volume
    :param rda_axis: if provided this axis will be used to compute a histogram
    :param rda_min: minimum value of the axis used to compute the histogram
    :param rda_max: maximum value of the axis used to compute the histogram
    :param n_rda_bins: number of axis bins
    :param use_log: if True (default is True) the axis is logarithmic
    :param n_samples: number of samples take to compute the histogram
    :param normalize: if True (default is False) the area is noramlized to one.
    :return: Tuple (histogram, histogram axis)

    Examples
    --------
    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.modeling
    >>> import scikit_fluorescence.io
    >>> atoms = skf.io.structure.read_coordinates('./test/data/atomic_coordinates/pdb_files/hGBP1_closed.pdb')
    >>> 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')
    >>> p_rda, rda = av1.pRDA(av2)

    """
    if rda_axis is None:
        if use_log:
            rda_axis = np.logspace(
                start=np.log10(rda_min),
                stop=np.log10(rda_max),
                num=n_rda_bins,
                dtype=np.float64
            )
        else:
            rda_axis = np.linspace(
                start=rda_min,
                stop=rda_max,
                num=n_rda_bins,
                dtype=np.float64
            )
    ds = random_distances(av1.points, av2.points, n_samples)
    r = ds[:, 0]
    w = ds[:, 1]
    p = np.histogram(r, bins=rda_axis, weights=w)[0]
    if normalize:
        p /= np.sum(p)
    return p, rda_axis


def RDAMean(
        av1: skf.modeling.dye.av.BasicAV,
        av2: skf.modeling.dye.av.BasicAV,
        n_samples: int = distance_samples
) -> float:
    """Calculate the mean distance between two accessible volumes

    >>> import scikit_fluorescence as skf
    >>> import scikit_fluorescence.modeling
    >>> pdb_filename = './test/data/structure/T4L_Topology.pdb'
    >>> structure = skf.structure(pdb_filename)
    >>> av1 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=72, atom_name='CB')
    >>> av2 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=134, atom_name='CB')
    >>> skf.modeling.dye.av.functions.RDAMean(av1, av2)
    52.93390285282142
    """
    d = random_distances(av1.points, av2.points, n_samples)
    return np.dot(d[:, 0], d[:, 1]) / d[:, 1].sum()


def widthRDA(
        av1: skf.modeling.dye.av.BasicAV,
        av2: skf.modeling.dye.av.BasicAV,
        n_samples: int = distance_samples
):
    """Calculate the width of the distance distribution between two accessible volumes

    >>> import scikit_fluorescence as skl
    >>> import scikit_fluorescence.modeling
    >>> pdb_filename = './test/data/structure/T4L_Topology.pdb'
    >>> structure = mfm.structure.Structure(pdb_filename)
    >>> av1 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=72, atom_name='CB')
    >>> av2 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=134, atom_name='CB')
    >>> skf.modeling.dye.av.functions.widthRDA(av1, av2)
    52.93390285282142
    """
    d = random_distances(av1.points, av2.points, n_samples)
    s = np.dot(d[:, 0]**2.0, d[:, 1])
    f = np.dot(d[:, 0], d[:, 1])**2.0
    v = s - f
    return np.sqrt(v)


def RDAMeanE(
        av1: skf.modeling.dye.av.BasicAV,
        av2: skf.modeling.dye.av.BasicAV,
        R0: float = 52.0,
        n_samples: int = distance_samples
) -> float:
    """Calculate the FRET-averaged (PDA/Intensity) distance between two accessible volumes

    >>> import scikit_fluorescence as skl
    >>> import scikit_fluorescence.modeling
    >>> import scikit_fluorescence.io
    >>> pdb_filename = './test/data/structure/T4L_Topology.pdb'
    >>> structure = skf.io.structure.read_coordinates(pdb_filename)
    >>> av1 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=72, atom_name='CB')
    >>> av2 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=134, atom_name='CB')
    >>> skf.modeling.dye.av.RDAMeanE(av1, av2)
    52.602731299544686
    """
    d = random_distances(av1.points, av2.points, n_samples)
    r = d[:, 0]
    w = d[:, 1]
    e = (1./(1.+(r/R0)**6.0))
    mean_fret = np.dot(w, e) / w.sum()
    return (1./mean_fret - 1.)**(1./6.) * R0


def dRmp(
        av1: skf.modeling.dye.av.BasicAV,
        av2: skf.modeling.dye.av.BasicAV
) -> float:
    """Calculate the distance between the mean position of two accessible volumes

    >>> import scikit_fluorescence as skl
    >>> import scikit_fluorescence.modeling
    >>> import scikit_fluorescence.io
    >>> pdb_filename = './test/data/structure/T4L_Topology.pdb'
    >>> structure = skf.io.structure.read_coordinates(pdb_filename)
    >>> av1 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=72, atom_name='CB')
    >>> av2 = skf.modeling.dye.av.BasicAV(structure, residue_seq_number=134, atom_name='CB')
    >>> skf.modeling.dye.av.dRmp(av1, av2)
    49.724995634807691
    """
    return np.sqrt(((av1.Rmp-av2.Rmp)**2).sum())


@nb.jit
def random_distances(
        p1: np.ndarray,
        p2: np.ndarray,
        n_samples: int
):
    """

    :param xyzw: a 4-dim vector xyz and the weight of the coordinate
    :param nSamples:
    :return:
    """
    n_p1 = p1.shape[0]
    n_p2 = p2.shape[0]
    distances = np.empty((n_samples, 2), dtype=np.float64)
    for i in range(n_samples):
        i1 = np.random.randint(0, n_p1)
        i2 = np.random.randint(0, n_p2)
        distances[i, 0] = math.sqrt(
            (p1[i1, 0] - p2[i2, 0]) ** 2.0 +
            (p1[i1, 1] - p2[i2, 1]) ** 2.0 +
            (p1[i1, 2] - p2[i2, 2]) ** 2.0
        )
        distances[i, 1] = p1[i1, 3] * p2[i2, 3]
    return distances


@nb.jit
def split_av_acv(density, dg, radius, rs, r0):
    """

    :param density: numpy-array
        density of the accessible volume dimension ng, ng, ng uint8 numpy array as obtained of fps
    :param ng: int
        number of grid points
    :param dg: float
        grid resolution of the density grid
    :param radius: float
        radius around the list of points. all points within a radius of r0 around the points in rs are part
        of the subav. each slow point is assiciated to one slow-radius
    :param rs: list
        list of points (x,y,z) defining the subav
    :param r0:
        is the position of the accessible volume


    """
    ng = density.shape[0]
    n_radii = rs.shape[0]

    radius = np.array(radius, dtype=np.float64)
    if len(radius) != n_radii:
        radius = np.zeros(n_radii, dtpye=np.float64) + radius[0]

    d1 = np.zeros_like(density)
    d2 = np.zeros_like(density)

    # iterate through all possible grid points in the density
    n1 = int(0)
    n2 = int(0)
    for ix in range(ng):
        for iy in range(ng):
            for iz in range(ng):
                if density[ix, iy, iz] <= 0:
                    continue
                # count the overlaps with slow-centers
                overlapped = 0

                for isa in range(n_radii):
                    ix0 = int(((rs[isa, 0]-r0[0])/dg) + (ng - 1)/2)
                    iy0 = int(((rs[isa, 1]-r0[1])/dg) + (ng - 1)/2)
                    iz0 = int(((rs[isa, 2]-r0[2])/dg) + (ng - 1)/2)
                    radius_idx = int((radius[isa] / dg))
                    if ((ix - ix0)**2 + (iy - iy0)**2 + (iz - iz0)**2) < radius_idx**2:
                        overlapped = 1
                        break

                if overlapped > 0:
                    d1[ix, iy, iz] = 1
                    n1 += 1
                else:
                    d2[ix, iy, iz] = 1
                    n2 += 1
    return n1, n2, d1, d2


@nb.jit(nopython=True)
def atoms_in_reach(
        xyz: np.ndarray,
        vdw: np.ndarray,
        dmaxsq: float,
        atom_i: int,
        return_subset: bool = True
) -> (np.ndarray, np.ndarray, typing.List[int]):
    """Return the xyz coordinates of atoms within reach (defined by dmaxsq) of a list of atoms

    :param xyz:
    :param vdw:
    :param dmaxsq:
    :param atom_i:
    :param return_subset: if set to True the coordinates and vdw radii are returned
    :return:

    Example
    -------
    >>> import scikit_fluorescence.io.structure
    >>> filename = './test/data/modelling/pdb_files/hGBP1_open.pdb'
    >>> atoms = scikit_fluorescence.io.structure.read_coordinates(filename)
    >>> xs, vs = atoms_in_reach(atoms['xyz'], atoms['radius'], 10.0, 1)

    """
    # copy all atoms in proximity to the dye into a smaller array and move coordinate frame to attachment point
    n_atoms = xyz.shape[0]
    atomindex = np.empty(n_atoms, dtype=np.uint32)
    r0 = xyz[atom_i]
    natomsgrid = 0
    for i in range(0, n_atoms):
        dsq = ((xyz[i] - r0)**2.0).sum()
        if (dsq < dmaxsq) and (i != atom_i):
            atomindex[natomsgrid] = i
            natomsgrid += 1
    if not return_subset:
        natomsgrid = 0
    re_idx = np.empty(natomsgrid, dtype=np.uint32)
    ra = np.empty((natomsgrid, 3), dtype=np.float64)
    vdwr = np.empty(natomsgrid, dtype=np.float64)
    for i in range(natomsgrid):
        n = atomindex[i]
        re_idx[i] = n
        ra[i] = xyz[n]
        vdwr[i] = vdw[n]
    return ra, vdwr, re_idx

