import numpy as np
import pandas as pd
import pytest
from pathlib import Path

from strategies.adaptive_lag_correlation import (
    AdaptiveLagCorrelationEngine,
    LagCorrelationConfig,
    SelectionMode,
    Transform,
    align_to_regular_grid,
)


def _series(values, start="2020-01-01"):
    return pd.Series(values, index=pd.date_range(start, periods=len(values), freq="D"))


def _leading_pair(n=500, lag=45, noise=0.02, seed=7):
    rng = np.random.default_rng(seed)
    source = rng.normal(size=n).cumsum()
    target = np.full(n, np.nan)
    target[lag:] = source[:-lag] + rng.normal(scale=noise, size=n - lag)
    return _series(source), _series(target)


def _engine(**overrides):
    settings = dict(
        min_lag=30,
        max_lag=100,
        lag_step=1,
        window=160,
        min_observations=100,
        transform=Transform.LEVELS,
    )
    settings.update(overrides)
    config = LagCorrelationConfig(**settings)
    return AdaptiveLagCorrelationEngine(config)


def test_recovers_known_positive_source_lead_and_spectrum():
    source, target = _leading_pair(lag=45)

    result = _engine().estimate_at(source, target, target.index[-1])

    assert result.best_lag == 45
    assert result.best_correlation > 0.99
    assert result.spectrum[45].correlation == pytest.approx(result.best_correlation)
    assert result.confidence.relative_peak_separation > 0
    assert result.observations == 160


def test_selection_mode_is_explicit_for_negative_relationships():
    source, target = _leading_pair(lag=45, noise=0.0)
    target = -target

    positive = _engine().estimate_at(source, target, target.index[-1])
    negative = _engine(selection_mode=SelectionMode.STRONGEST_NEGATIVE).estimate_at(
        source, target, target.index[-1]
    )
    absolute = _engine(selection_mode=SelectionMode.STRONGEST_ABSOLUTE).estimate_at(
        source, target, target.index[-1]
    )

    assert positive.best_lag != 45
    assert negative.best_lag == 45
    assert negative.best_correlation < -0.99
    assert absolute.best_lag == 45


def test_equal_scores_choose_smallest_lag_and_report_zero_separation():
    source = _series(np.arange(300, dtype=float))
    target = source.copy()

    result = AdaptiveLagCorrelationEngine(
        LagCorrelationConfig(min_lag=1, max_lag=3, window=100, min_observations=100)
    ).estimate_at(source, target, target.index[-1])

    assert result.best_lag == 1
    assert result.confidence.margin == pytest.approx(0.0)
    assert result.confidence.relative_peak_separation == pytest.approx(0.0)


def test_scores_within_tie_tolerance_report_zero_separation():
    source = _series(np.arange(300, dtype=float))
    target = source.copy()
    target.iloc[-1] += 1e-12

    result = AdaptiveLagCorrelationEngine(
        LagCorrelationConfig(min_lag=1, max_lag=3, window=100, min_observations=100, tie_tolerance=1e-6)
    ).estimate_at(source, target, target.index[-1])

    assert result.best_lag == 1
    assert result.confidence.margin == pytest.approx(0.0)


def test_reports_a_distant_local_peak_separately_from_adjacent_runner_up():
    rng = np.random.default_rng(22)
    n = 700
    source_values = rng.normal(size=n)
    target_values = np.full(n, np.nan)
    for position in range(72, n):
        target_values[position] = (
            source_values[position - 45]
            + 0.95 * source_values[position - 46]
            + 0.60 * source_values[position - 72]
            + rng.normal(scale=0.03)
        )
    source, target = _series(source_values), _series(target_values)
    engine = AdaptiveLagCorrelationEngine(
        LagCorrelationConfig(
            min_lag=30,
            max_lag=100,
            window=300,
            min_observations=250,
            distinct_peak_exclusion_radius=7,
        )
    )

    result = engine.estimate_at(source, target, target.index[-1])

    assert result.best_lag == 46
    assert result.confidence.runner_up_lag == 45
    assert result.confidence.distinct_peak_lag == 72
    assert result.confidence.distinct_peak_correlation > 0.2
    assert result.confidence.distinct_peak_margin > 0


def test_distinct_peak_respects_the_configured_exclusion_radius():
    source, target = _leading_pair(lag=45)
    result = AdaptiveLagCorrelationEngine(
        LagCorrelationConfig(
            min_lag=30,
            max_lag=60,
            window=160,
            min_observations=100,
            distinct_peak_exclusion_radius=30,
        )
    ).estimate_at(source, target, target.index[-1])

    assert result.confidence.distinct_peak_lag is None


def test_estimates_lag_that_changes_over_time():
    rng = np.random.default_rng(4)
    n, change, first_lag, second_lag = 900, 450, 45, 72
    source_values = rng.normal(size=n).cumsum()
    target_values = np.full(n, np.nan)
    target_values[first_lag:change] = source_values[: change - first_lag]
    target_values[change + second_lag :] = source_values[change : n - second_lag]
    target_values += rng.normal(scale=0.01, size=n)
    source, target = _series(source_values), _series(target_values)
    engine = _engine(window=180)

    early = engine.estimate_at(source, target, target.index[change - 20])
    late = engine.estimate_at(source, target, target.index[-1])

    assert abs(early.best_lag - first_lag) <= 1
    assert abs(late.best_lag - second_lag) <= 1


def test_missing_observations_are_dropped_pairwise_and_minimum_is_enforced():
    source, target = _leading_pair(lag=45)
    source.iloc[-100:-80] = np.nan
    target.iloc[-15:-5] = np.nan

    result = _engine(min_observations=100).estimate_at(source, target, target.index[-1])
    insufficient = _engine(min_observations=151).estimate_at(source, target, target.index[-1])

    assert result.best_lag == 45
    assert result.observations == 130
    assert insufficient.best_lag is None
    assert not insufficient.spectrum


def test_estimates_are_invariant_to_future_mutation():
    source, target = _leading_pair(n=700, lag=45)
    evaluation_time = target.index[500]
    engine = _engine()

    before = engine.estimate_at(source, target, evaluation_time)
    source.loc[source.index > evaluation_time] = 1_000_000
    target.loc[target.index > evaluation_time] = -1_000_000
    after = engine.estimate_at(source, target, evaluation_time)

    assert after == before


def test_invalid_configuration_and_non_monotonic_index_are_rejected():
    with pytest.raises(ValueError, match="min_lag"):
        LagCorrelationConfig(min_lag=-1, max_lag=5)
    with pytest.raises(ValueError, match="lag_step"):
        LagCorrelationConfig(min_lag=1, max_lag=5, lag_step=0)

    source, target = _leading_pair()
    source.index = source.index[::-1]
    with pytest.raises(ValueError, match="monotonic"):
        _engine().estimate_at(source, target, target.index[-1])


def test_regular_grid_forward_fills_only_source_from_prior_observations():
    index = pd.to_datetime(["2020-01-01 12:00", "2020-01-03 12:00"], utc=True)
    source = pd.Series([10.0, 30.0], index=index)
    target = pd.Series([100.0, 300.0], index=index)

    source_grid, target_grid = align_to_regular_grid(source, target, frequency="1D")

    assert source_grid.loc["2020-01-02"] == 10.0
    assert pd.isna(target_grid.loc["2020-01-02"])


def test_pine_estimator_short_circuits_an_empty_candidate_spectrum():
    """Prevent a zero-length candidate array from reaching array.get()."""
    library = Path("libraries/LibraryAdaptiveLagCorrelation.pine").read_text()

    assert "candidateCount = array.size(lags)" in library
    assert "if candidateCount == 0\n        [na, na, na, na, na, na, na, na]" in library
    assert "lastIndex = candidateCount - 1" in library
    assert "for index = 0 to array.size(lags) - 1" not in library