import os
import numpy as np
from ase.io import read
from dynasor.trajectory import Trajectory
from dynasor.qpoints import get_spherical_qpoints
from dynasor.correlation_functions import compute_spectral_energy_density


def test_regression():
    # parameters
    q_max = 1.8
    dt = 1.5

    # trajectory
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(this_dir, 'NaClI_size4x3x2_100frames_with_velocity.xyz')

    # setup prim and supercell
    prim = read(os.path.join(this_dir, 'NaCl_prim.xyz'))
    supercell = prim.repeat((4, 3, 2))

    # run dynasor
    traj = Trajectory(traj_fname, trajectory_format='extxyz', atomic_indices='read_from_trajectory')
    q_points = get_spherical_qpoints(traj.cell, q_max=q_max)
    w, SED = compute_spectral_energy_density(traj, supercell, prim, q_points, dt, partial=True)

    # read old results
    archive = np.load(os.path.join(this_dir, 'dynasor_regression_results_SED.npz'))
    old_results = {k: archive[k] for k in archive.files}

    # compare against old results
    assert len(old_results) == 2
    assert np.allclose(old_results['w'], w)
    assert np.allclose(old_results['sed'], SED)
