import itertools as it

import numpy as np
import mdtraj as md
from progressbar import ProgressBar

from scattering.utils.utils import get_dt
from scattering.utils.constants import get_form_factor


def compute_van_hove(trj, chunk_length, water=False,
                     r_range=(0, 1.0), bin_width=0.005, n_bins=None,
                     self_correlation=True, periodic=True, opt=True, partial=False):
    """Compute the partial van Hove function of a trajectory

    Parameters
    ----------
    trj : mdtraj.Trajectory
        trajectory on which to compute the Van Hove function
    chunk_length : int
        length of time between restarting averaging
    water : bool
        use X-ray form factors for water that account for polarization
    r_range : array-like, shape=(2,), optional, default=(0.0, 1.0)
        Minimum and maximum radii.
    bin_width : float, optional, default=0.005
        Width of the bins in nanometers.
    n_bins : int, optional, default=None
        The number of bins. If specified, this will override the `bin_width`
         parameter.
    self_correlation : bool, default=True
        Whether or not to include the self-self correlations

    Returns
    -------
    r : numpy.ndarray
        r positions generated by histogram binning
    g_r_t : numpy.ndarray
        Van Hove function at each time and position
    """

    n_physical_atoms = len([a for a in trj.top.atoms if a.element.mass > 0])
    unique_elements = list(set([a.element for a in trj.top.atoms if a.element.mass > 0]))

    partial_dict = dict()

    for elem1, elem2 in it.combinations_with_replacement(unique_elements[::-1], 2):
        print('doing {0} and {1} ...'.format(elem1, elem2))
        r, g_r_t_partial = compute_partial_van_hove(trj=trj,
                                                    chunk_length=chunk_length,
                                                    selection1='element {}'.format(elem1.symbol),
                                                    selection2='element {}'.format(elem2.symbol),
                                                    r_range=r_range,
                                                    bin_width=bin_width,
                                                    n_bins=n_bins,
                                                    self_correlation=self_correlation,
                                                    periodic=periodic,
                                                    opt=opt)
        partial_dict[(elem1, elem2)] = g_r_t_partial

    if partial:
        return partial_dict

    norm = 0
    g_r_t = None

    for key, val in partial_dict.items():
        elem1, elem2 = key
        concentration1 = trj.atom_slice(trj.top.select('element {}'.format(elem1.symbol))).n_atoms / n_physical_atoms
        concentration2 = trj.atom_slice(trj.top.select('element {}'.format(elem2.symbol))).n_atoms / n_physical_atoms
        form_factor1 = get_form_factor(element_name=elem1.symbol, water=water)
        form_factor2 = get_form_factor(element_name=elem2.symbol, water=water)

        coeff = form_factor1 * concentration1 * form_factor2 * concentration2

        if g_r_t is None:
            g_r_t = np.zeros_like(val)
        g_r_t += val * coeff

        norm += coeff

    # Reshape g_r_t to better represent the discretization in both r and t
    g_r_t_final = np.empty(shape=(chunk_length, len(r)))
    for i in range(chunk_length):
        g_r_t_final[i, :] = np.mean(g_r_t[i::chunk_length], axis=0)

    g_r_t_final /= norm

    t = trj.time[:chunk_length]

    return r, t, g_r_t_final


def compute_partial_van_hove(trj, chunk_length=10, selection1=None, selection2=None,
                             r_range=(0, 1.0), bin_width=0.005, n_bins=200,
                             self_correlation=True, periodic=True, opt=True):
    """Compute the partial van Hove function of a trajectory

    Parameters
    ----------
    trj : mdtraj.Trajectory
        trajectory on which to compute the Van Hove function
    chunk_length : int
        length of time between restarting averaging
    selection1 : str
        selection to be considered, in the style of MDTraj atom selection
    selection2 : str
        selection to be considered, in the style of MDTraj atom selection
    r_range : array-like, shape=(2,), optional, default=(0.0, 1.0)
        Minimum and maximum radii.
    bin_width : float, optional, default=0.005
        Width of the bins in nanometers.
    n_bins : int, optional, default=None
        The number of bins. If specified, this will override the `bin_width`
         parameter.
    self_correlation : bool, default=True
        Whether or not to include the self-self correlations

    Returns
    -------
    r : numpy.ndarray
        r positions generated by histogram binning
    g_r_t : numpy.ndarray
        Van Hove function at each time and position
    """

    unique_elements = (
        set([a.element for a in trj.atom_slice(trj.top.select(selection1)).top.atoms]),
        set([a.element for a in trj.atom_slice(trj.top.select(selection2)).top.atoms]),
    )

    if any([len(val) > 1 for val in unique_elements]):
        raise UserWarning(
            'Multiple elements found in a selection(s). Results may not be '
            'direcitly comprable to scattering experiments.'
        )

    # Don't need to store it, but this serves to check that dt is constant
    dt = get_dt(trj)

    pairs = trj.top.select_pairs(selection1=selection1, selection2=selection2)

    n_chunks = int(trj.n_frames / chunk_length)

    g_r_t = None
    pbar = ProgressBar()

    for i in pbar(range(n_chunks)):
        times = list()
        for j in range(chunk_length):
            times.append([chunk_length*i, chunk_length*i+j])
        r, g_r_t_frame = md.compute_rdf_t(
            traj=trj,
            pairs=pairs,
            times=times,
            r_range=r_range,
            bin_width=bin_width,
            n_bins=n_bins,
            period_length=chunk_length,
            self_correlation=self_correlation,
            periodic=periodic,
            opt=opt,
        )

        if g_r_t is None:
            g_r_t = np.zeros_like(g_r_t_frame)
        g_r_t += g_r_t_frame

    return r, g_r_t

[(47, 'itertools.combinations_with_replacement', 'it.combinations_with_replacement', True, 'import itertools as it\n'), (139, 'scattering.utils.utils.get_dt', 'get_dt', False, 'from scattering.utils.utils import get_dt\n'), (146, 'progressbar.ProgressBar', 'ProgressBar', False, 'from progressbar import ProgressBar\n'), (71, 'scattering.utils.constants.get_form_factor', 'get_form_factor', False, 'from scattering.utils.constants import get_form_factor\n'), (72, 'scattering.utils.constants.get_form_factor', 'get_form_factor', False, 'from scattering.utils.constants import get_form_factor\n'), (85, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (152, 'mdtraj.compute_rdf_t', 'md.compute_rdf_t', True, 'import mdtraj as md\n'), (77, 'numpy.zeros_like', 'np.zeros_like', True, 'import numpy as np\n'), (166, 'numpy.zeros_like', 'np.zeros_like', True, 'import numpy as np\n')]