import os
from itertools import combinations_with_replacement

import numpy as np
import pytest

from dynasor.sample import Sample, DynamicSample, StaticSample


# parameters
N_qpoints = 20
window_size = 100


@pytest.fixture
def q_points():
    return np.array([np.array([1, 0, 0]) * amp for amp in np.linspace(0, 1, N_qpoints)])


@pytest.fixture
def time():
    return np.linspace(0, 10, window_size)


@pytest.fixture
def omega():
    return np.linspace(0, 5, window_size)


@pytest.fixture
def traj_fname_xyz_long():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(this_dir,
                              'trajectory_reader/trajectory_files/dump_long_with_velocities.xyz')
    return traj_fname


@pytest.fixture
def simple_data_dict(q_points, time, omega):
    # setup data dict
    data_dict = dict()
    data_dict['q_points'] = q_points
    data_dict['time'] = time
    data_dict['omega'] = omega

    size = (N_qpoints, window_size)
    data_dict['Fqt_coh_A_A'] = np.linspace(-1, 1, N_qpoints * window_size).reshape(size)
    data_dict['Fqt_coh_B_B'] = np.linspace(-0.4, 0.4, N_qpoints * window_size).reshape(size)
    data_dict['Fqt_coh_A_B'] = np.linspace(-0.3, 0.3, N_qpoints * window_size).reshape(size)
    data_dict['Sqw_coh_A_A'] = np.linspace(-10, 10, N_qpoints * window_size).reshape(size)
    data_dict['Sqw_coh_B_B'] = np.linspace(-4, 4, N_qpoints * window_size).reshape(size)
    data_dict['Sqw_coh_A_B'] = np.linspace(-3, 3, N_qpoints * window_size).reshape(size)
    return data_dict


@pytest.fixture
def data_dict_with_incoh(q_points, time, omega):
    # setup data dict
    data_dict = dict()
    data_dict['q_points'] = q_points
    data_dict['time'] = time
    data_dict['omega'] = omega

    # coherent and currents
    names = ['Fqt_coh', 'Sqw_coh', 'Clqt', 'Clqw', 'Ctqt', 'Ctqw']
    pairs = ['A_A', 'B_B', 'A_B']
    size = (N_qpoints, window_size)
    for it, name in enumerate(names):
        total = np.zeros(size)
        for jt, pair in enumerate(pairs):
            # some determinstic values
            minval = 0.5*it**2 - 4.8*jt
            maxval = 1.5 * it + 14.8*jt + 25.0
            F = np.linspace(minval, maxval, N_qpoints * window_size).reshape(size)
            total += F
            data_dict[f'{name}_{pair}'] = F
        data_dict[name] = total

    # incoherent
    names = ['Fqt_incoh', 'Sqw_incoh']
    atom_types = ['A', 'B']
    for it, name in enumerate(names):
        total = np.zeros(size)
        for jt, atom_type in enumerate(atom_types):
            # some determinstic values
            minval = 0.331*it**2.2 - 51.8*jt
            maxval = 3.1 * it + 84.8*jt + 121.2
            F = np.linspace(minval, maxval, N_qpoints * window_size).reshape(size)
            total += F
            data_dict[f'{name}_{atom_type}'] = F
        data_dict[name] = total

    # total
    data_dict['Fqt'] = data_dict['Fqt_coh'] + data_dict['Fqt_incoh']
    data_dict['Sqw'] = data_dict['Sqw_coh'] + data_dict['Sqw_incoh']
    return data_dict


@pytest.fixture
def data_dict_without_incoh(q_points, time, omega):
    # setup data dict
    data_dict = dict()
    data_dict['q_points'] = q_points
    data_dict['time'] = time
    data_dict['omega'] = omega

    # coherent and currents
    names = ['Fqt_coh', 'Sqw_coh', 'Clqt', 'Clqw', 'Ctqt', 'Ctqw']
    pairs = ['A_A', 'B_B', 'A_B']
    size = (N_qpoints, window_size)
    for it, name in enumerate(names):
        total = np.zeros(size)
        for jt, pair in enumerate(pairs):
            # some determinstic values
            minval = -0.225*it**2 - 14.8*jt
            maxval = 15.7 * it + 2.38*jt + 225.2
            F = np.linspace(minval, maxval, N_qpoints * window_size).reshape(size)
            total += F
            data_dict[f'{name}_{pair}'] = F
        data_dict[name] = total
    data_dict['Fqt'] = data_dict['Fqt_coh']
    data_dict['Sqw'] = data_dict['Sqw_coh']
    return data_dict


@pytest.fixture
def simulation_parameters():
    simulation_parameters = dict()
    simulation_parameters['atom_types'] = ['A', 'B']
    simulation_parameters['pairs'] = [('A', 'A'), ('A', 'B'), ('B', 'B')]
    simulation_parameters['particle_counts'] = dict(A=2050, B=1020)
    simulation_parameters['cell'] = np.diag([1.5, 2.0, 3.0])
    return simulation_parameters


@pytest.fixture
def simple_sample(data_dict, simulation_parameters):
    meta_data = {'type': 'data_dict'}
    sample = Sample(data_dict, **simulation_parameters, **meta_data)
    return sample


@pytest.fixture
def dynamic_sample_with_incoh(data_dict_with_incoh, simulation_parameters):
    meta_data = {'type': 'data_dict_with_incoh'}
    sample = DynamicSample(data_dict_with_incoh, **simulation_parameters, **meta_data)
    return sample


@pytest.fixture
def dynamic_sample_without_incoh(data_dict_without_incoh, simulation_parameters):
    meta_data = {'type': 'data_dict_without_incoh'}
    sample = DynamicSample(data_dict_without_incoh, **simulation_parameters, **meta_data)
    return sample


@pytest.fixture
def static_sample(q_points, simulation_parameters):
    data_dict = dict()
    data_dict['q_points'] = q_points
    data_dict['Sq_A_A'] = np.linspace(-1.1, 1, N_qpoints).reshape((N_qpoints, 1))
    data_dict['Sq_A_B'] = np.linspace(0, 1.5, N_qpoints).reshape((N_qpoints, 1))
    data_dict['Sq_B_B'] = np.linspace(5.1, 9.0, N_qpoints).reshape((N_qpoints, 1))
    data_dict['Sq'] = data_dict['Sq_A_A'] + data_dict['Sq_A_B'] + data_dict['Sq_B_B']
    sample = StaticSample(data_dict, **simulation_parameters)
    return sample


@pytest.fixture(scope='function')
def sample_with_species(request, q_points, time, omega):
    """
    Returns a dynamic sample with coherent, incoherent and currents
    for the requested species.
    """
    species = request.param
    # Simulation simulation_parameters
    simulation_parameters = dict()
    simulation_parameters['atom_types'] = species
    pairs = list(combinations_with_replacement(species, 2))
    simulation_parameters['pairs'] = pairs
    simulation_parameters['particle_counts'] = {}
    for i, s in enumerate(species):
        simulation_parameters['particle_counts'][s] = i*1000 + 100
    simulation_parameters['cell'] = np.diag([1.5, 2.0, 3.0])

    # Correlation functions
    names = ['Fqt_coh', 'Sqw_coh', 'Clqt', 'Clqw', 'Ctqt', 'Ctqw']
    data_dict = dict()
    data_dict['q_points'] = q_points
    data_dict['time'] = time
    data_dict['omega'] = omega

    size = (N_qpoints, window_size)
    # list of the form [A_A, B_B, A_B, ...]
    pair_strings = ['_'.join(combination) for combination in pairs]
    for it, name in enumerate(names):
        total = np.zeros(size)
        for jt, pair in enumerate(pair_strings):
            # some determinstic values
            minval = 0.5*it**2 - 4.8*jt
            maxval = 1.5 * it + 14.8*jt + 25.0
            F = np.linspace(minval, maxval, N_qpoints * window_size).reshape(size)
            total += F
            data_dict[f'{name}_{pair}'] = F
        data_dict[name] = total

    # incoherent
    names = ['Fqt_incoh', 'Sqw_incoh']
    for it, name in enumerate(names):
        total = np.zeros(size)
        for jt, atom_type in enumerate(species):
            # some determinstic values
            minval = 0.331*it**2.2 - 51.8*jt
            maxval = 3.1 * it + 84.8*jt + 121.2
            F = np.linspace(minval, maxval, N_qpoints * window_size).reshape(size)
            total += F
            data_dict[f'{name}_{atom_type}'] = F
        data_dict[name] = total

    # total
    data_dict['Fqt'] = data_dict['Fqt_coh'] + data_dict['Fqt_incoh']
    data_dict['Sqw'] = data_dict['Sqw_coh'] + data_dict['Sqw_incoh']

    # Create the DynamicSample
    meta_data = {'type': 'data_dict'}
    sample = DynamicSample(data_dict, **simulation_parameters, **meta_data)
    return sample
