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


def test_Sqw_calculation_with_time_window_longer_than_traj():
    # traj index files
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir, 'trajectory_reader/trajectory_files/dump_long_with_velocities.xyz')

    # input parameters
    time_window = 40  # Trajectory is 30 long
    dt = 2
    q_max = 1.0

    # run dynasor
    traj = Trajectory(traj_fname, trajectory_format='extxyz', atomic_indices='read_from_trajectory')
    q_points = get_spherical_qpoints(traj.cell, q_max)
    sample = compute_dynamic_structure_factors(traj, q_points, dt=dt, window_size=time_window,
                                               calculate_currents=True, calculate_incoherent=True)

    time_reference = dt * np.arange(0, 30)
    assert np.allclose(sample.time, time_reference)
