import os
import numpy as np
from dynasor.trajectory import Trajectory
from dynasor.qpoints import get_spherical_qpoints
from dynasor.correlation_functions import compute_static_structure_factors


def test_regression():
    # parameters
    q_max = 5.0

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

    # 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)
    sample = compute_static_structure_factors(traj, q_points)

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

    # compare against old results
    assert len(old_results) == len(sample.dimensions + sample.available_correlation_functions)
    for key, array in old_results.items():
        assert np.allclose(array, sample[key])
