import logging
import os

import numpy as np
import pytest

from dynasor.correlation_functions import (
    compute_dynamic_structure_factors,
    compute_static_structure_factors,
)
from dynasor.logging_tools import logger as dynasor_logger
from dynasor.sample import Sample
from dynasor.trajectory import Trajectory


@pytest.fixture
def dynasor_caplog(caplog):
    """caplog fixture that works with dynasor's non-propagating logger.

    Adds caplog's handler directly to the dynasor logger so records are captured
    regardless of the logger's propagate setting.
    """
    dynasor_logger.addHandler(caplog.handler)
    yield caplog
    dynasor_logger.removeHandler(caplog.handler)


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


@pytest.fixture
def traj_fname_lammps():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir,
        'trajectory_reader/trajectory_files/positions.lammpstrj')
    return traj_fname


@pytest.fixture
def q_points():
    q_points = np.array([i * np.array([1/2, 2/3, 1.235]) for i in np.linspace(0, 10, 11)])
    return q_points


def test_API_dynamic_structure_factor(traj_fname_xyz, q_points):
    traj = Trajectory(traj_fname_xyz, 'extxyz')
    res = compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=3)
    assert isinstance(res, Sample)


def test_API_dynamic_structure_factor_integer_q_points(traj_fname_xyz):
    traj = Trajectory(traj_fname_xyz, 'extxyz')
    q_points = np.array([[1, 1, 0]])
    res = compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=3)
    assert isinstance(res, Sample)


def test_API_dynamic_structure_factor_selfpart(traj_fname_xyz, q_points):
    traj = Trajectory(traj_fname_xyz, 'extxyz')
    res = compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=3,
                                            calculate_incoherent=True)
    assert isinstance(res, Sample)


def test_API_static_structure_factor(traj_fname_xyz_long, q_points):
    # setup
    n_atoms = 320
    n_A = 50
    atomic_indices = dict()
    atomic_indices['A'] = np.arange(0, n_A, 1)
    atomic_indices['B'] = np.arange(n_A, n_atoms, 1)

    # calculate static
    traj = Trajectory(
        traj_fname_xyz_long, trajectory_format='extxyz', atomic_indices=atomic_indices)
    res1 = compute_static_structure_factors(traj, q_points=q_points)
    assert res1.simulation_data['number_of_frames'] == 30

    # calculate S(q) via compute_dynamic_structure_factors (use dummy dt and window_size)
    traj = Trajectory(
        traj_fname_xyz_long, trajectory_format='extxyz', atomic_indices=atomic_indices)
    res2 = compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=3)
    assert res2.simulation_data['number_of_frames'] == 30

    # compare results
    pairs = ['A_A', 'A_B', 'B_B']
    for pair in pairs:
        Sq1 = getattr(res1, f'Sq_{pair}')
        Sq2 = getattr(res2, f'Fqt_coh_{pair}')[:, 0].reshape(-1, 1)
        assert np.allclose(Sq1, Sq2)
    assert np.allclose(res1.Sq, res2.Fqt_coh[:, 0].reshape(-1, 1))


# Comparing the mdanalysis lammps reader to the internal lammps reader
def test_API_calling_mdanalysis_reader(traj_fname_lammps, q_points):
    window = 3
    dt = 1.0

    traj1 = Trajectory(traj_fname_lammps, 'lammps_internal')
    traj2 = Trajectory(traj_fname_lammps, 'lammps_mdanalysis',
                       length_unit='Angstrom', time_unit='fs')

    res1 = compute_dynamic_structure_factors(traj1, q_points=q_points, dt=dt, window_size=window)
    res2 = compute_dynamic_structure_factors(traj2, q_points=q_points, dt=dt, window_size=window)
    assert isinstance(res1, Sample)
    assert isinstance(res2, Sample)

    # check that results are the same
    assert res1.available_correlation_functions == res2.available_correlation_functions
    for key in res1.available_correlation_functions:
        np.testing.assert_almost_equal(res1[key], res2[key], decimal=5)


# test of time-sampling parameters
# --------------------------------
def test_API_dynamic_step(traj_fname_xyz_long, q_points):
    """
    Ensure that setting step > 1 affects the resulting time
    """
    dt = 1.0
    time_window = 4
    step = 1

    for step in [1, 2, 3]:

        traj = Trajectory(traj_fname_xyz_long, 'extxyz', frame_step=step)
        res = compute_dynamic_structure_factors(
            traj, q_points=q_points, dt=dt, window_size=time_window)
        assert isinstance(res, Sample)
        time = res.time
        assert np.allclose(np.diff(time), step * dt), (np.diff(time), step * dt)


# test ValueErrors from bad input arguments
# -----------------------------------------

def test_raises_with_bad_input_args(traj_fname_xyz_long, q_points):
    time_window = 4
    dt = 1.0

    # q-points wrong shape
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'q-points array has the wrong shape'):
        compute_dynamic_structure_factors(traj, q_points=q_points[:, 1:], dt=dt,
                                          window_size=time_window)

    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'q-points array has the wrong shape'):
        compute_static_structure_factors(traj, q_points=q_points[:, 1:])

    # dt not positive
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'dt must be positive'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=-1.0, window_size=time_window)

    # bad window_size
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'window_size must be larger than 2'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=dt, window_size=1)

    # bad window_step
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'window_step must be positive'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=dt, window_size=time_window,
                                          window_step=-1)


def test_raises_with_dt_zero(traj_fname_xyz_long, q_points):
    """dt=0 is invalid (division by zero in frequency axis); must raise ValueError."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with pytest.raises(ValueError, match=r'dt must be positive'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=0.0, window_size=4)


def test_raises_with_currents_but_no_velocities(traj_fname_lammps, q_points):
    """calculate_currents=True requires velocities; must raise a clear ValueError
    instead of an opaque TypeError from indexing None."""
    traj = Trajectory(traj_fname_lammps, 'lammps_internal')
    with pytest.raises(ValueError, match=r'[Vv]elocities'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=3,
                                          calculate_currents=True)


def test_empty_q_points_returns_empty_sample(traj_fname_xyz_long):
    """An empty (0, 3) q_points array produces a sample with zero q-points."""
    q_empty = np.zeros((0, 3))

    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    res_static = compute_static_structure_factors(traj, q_points=q_empty)
    assert res_static.Sq.shape[0] == 0

    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    res_dynamic = compute_dynamic_structure_factors(
        traj, q_points=q_empty, dt=1.0, window_size=4)
    assert res_dynamic.Fqt_coh.shape[0] == 0


def test_window_size_exceeding_n_frames_completes_without_nan(traj_fname_xyz_long, q_points):
    """window_size >= n_frames is handled gracefully: NaNs are truncated internally.

    The long trajectory has 30 frames. When window_size equals or exceeds n_frames the
    code emits a warning and truncates NaN time lags before returning, so the result
    contains no NaN values.
    """
    n_frames = 30  # known frame count of the long trajectory

    # Largest valid even window_size: produces a result with no NaN
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    res_valid = compute_dynamic_structure_factors(
        traj, q_points=q_points, dt=1.0, window_size=n_frames - 2)
    assert not np.any(np.isnan(res_valid.Fqt_coh))

    # window_size equal to n_frames: code truncates NaNs and returns a valid result
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    res_exceed = compute_dynamic_structure_factors(
        traj, q_points=q_points, dt=1.0, window_size=n_frames)
    assert not np.any(np.isnan(res_exceed.Fqt_coh))


# --- logging_interval tests ---

def test_dynamic_logging_interval_default_emits_info_at_frame_0(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With default logging_interval=1000, frame 0 is logged at INFO."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.INFO, logger='dynasor'):
        compute_dynamic_structure_factors(traj, q_points=q_points, dt=1.0, window_size=4)
    info_msgs = [r.message for r in dynasor_caplog.records if r.levelno == logging.INFO]
    assert any('Processing window 0' in m for m in info_msgs)


def test_dynamic_logging_interval_1_emits_info_for_every_window(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With logging_interval=1, every window start is logged at INFO."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.DEBUG, logger='dynasor'):
        compute_dynamic_structure_factors(
            traj, q_points=q_points, dt=1.0, window_size=4, logging_interval=1)
    info_window_msgs = [
        r.message for r in dynasor_caplog.records
        if r.levelno == logging.INFO and 'Processing window' in r.message
    ]
    # With 30 frames and window_size=4 (N_tc=5), window_step=1,
    # we get windows starting at frames 0..25 — all should be logged at INFO.
    assert len(info_window_msgs) > 1


def test_dynamic_logging_interval_0_emits_no_info_progress(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With logging_interval=0, no 'Processing window' lines appear at INFO."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.INFO, logger='dynasor'):
        compute_dynamic_structure_factors(
            traj, q_points=q_points, dt=1.0, window_size=4, logging_interval=0)
    info_window_msgs = [
        r.message for r in dynasor_caplog.records
        if r.levelno == logging.INFO and 'Processing window' in r.message
    ]
    assert len(info_window_msgs) == 0


def test_dynamic_logging_no_duplicate_at_info_level(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """At INFO level, windows not on the interval are NOT logged (no duplicates)."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.INFO, logger='dynasor'):
        compute_dynamic_structure_factors(
            traj, q_points=q_points, dt=1.0, window_size=4, logging_interval=1000)
    # At INFO, debug records are suppressed; only interval-matched windows appear.
    window_msgs = [
        r for r in dynasor_caplog.records if 'Processing window' in r.message
    ]
    # All captured window messages must be at INFO (no DEBUG leaking through).
    assert all(r.levelno == logging.INFO for r in window_msgs)


def test_static_logging_interval_1_emits_info_for_every_frame(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With logging_interval=1, every frame is logged at INFO for static SF."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.DEBUG, logger='dynasor'):
        compute_static_structure_factors(traj, q_points=q_points, logging_interval=1)
    info_frame_msgs = [
        r.message for r in dynasor_caplog.records
        if r.levelno == logging.INFO and 'Processing frame' in r.message
    ]
    # Trajectory has 30 frames; all should be logged at INFO.
    assert len(info_frame_msgs) == 30


def test_static_logging_interval_0_emits_no_info_progress(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With logging_interval=0, no 'Processing frame' lines appear at INFO."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.INFO, logger='dynasor'):
        compute_static_structure_factors(traj, q_points=q_points, logging_interval=0)
    info_frame_msgs = [
        r.message for r in dynasor_caplog.records
        if r.levelno == logging.INFO and 'Processing frame' in r.message
    ]
    assert len(info_frame_msgs) == 0


def test_static_logging_interval_default_emits_info_at_frame_0(
        traj_fname_xyz_long, q_points, dynasor_caplog):
    """With default logging_interval=1000, frame 0 is logged at INFO."""
    traj = Trajectory(traj_fname_xyz_long, 'extxyz')
    with dynasor_caplog.at_level(logging.INFO, logger='dynasor'):
        compute_static_structure_factors(traj, q_points=q_points)
    info_msgs = [r.message for r in dynasor_caplog.records if r.levelno == logging.INFO]
    assert any('Processing frame 0' in m for m in info_msgs)
