"""Causal, inspectable lag-correlation estimation for aligned time series.

The convention is deliberately one-way: lag ``L`` means a source observation
at grid position ``i - L`` is compared with the target observation at ``i``.
Thus a positive L says the source leads the target by L grid intervals.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Mapping

import numpy as np
import pandas as pd


class Transform(str, Enum):
    LEVELS = "levels"
    DIFFERENCE = "difference"
    PCT_CHANGE = "pct_change"
    LOG_RETURN = "log_return"


class SelectionMode(str, Enum):
    STRONGEST_POSITIVE = "strongest_positive"
    STRONGEST_NEGATIVE = "strongest_negative"
    STRONGEST_ABSOLUTE = "strongest_absolute"


@dataclass(frozen=True)
class LagCorrelationConfig:
    """Estimator configuration; all lag and window values are grid intervals."""

    min_lag: int = 30
    max_lag: int = 120
    lag_step: int = 1
    window: int = 180
    min_observations: int = 120
    transform: Transform = Transform.LEVELS
    selection_mode: SelectionMode = SelectionMode.STRONGEST_POSITIVE
    tie_tolerance: float = 1e-12
    distinct_peak_exclusion_radius: int = 7

    def __post_init__(self) -> None:
        if self.min_lag < 0:
            raise ValueError("min_lag must be >= 0")
        if self.max_lag < self.min_lag:
            raise ValueError("max_lag must be >= min_lag")
        if self.lag_step <= 0:
            raise ValueError("lag_step must be > 0")
        if self.window <= 1:
            raise ValueError("window must be > 1")
        if not 2 <= self.min_observations <= self.window:
            raise ValueError("min_observations must be between 2 and window")
        if self.tie_tolerance < 0:
            raise ValueError("tie_tolerance must be >= 0")
        if self.distinct_peak_exclusion_radius < 0:
            raise ValueError("distinct_peak_exclusion_radius must be >= 0")

    @property
    def lags(self) -> range:
        return range(self.min_lag, self.max_lag + 1, self.lag_step)


@dataclass(frozen=True)
class CandidateCorrelation:
    lag: int
    correlation: float
    observations: int


@dataclass(frozen=True)
class LagConfidence:
    """Peak separation, not a statistical probability of the economic thesis."""

    runner_up_lag: int | None
    runner_up_correlation: float | None
    margin: float | None
    relative_peak_separation: float | None
    distinct_peak_lag: int | None
    distinct_peak_correlation: float | None
    distinct_peak_margin: float | None
    distinct_peak_relative_separation: float | None


@dataclass(frozen=True)
class LagEstimate:
    evaluated_at: pd.Timestamp
    selection_mode: SelectionMode
    best_lag: int | None
    best_correlation: float | None
    observations: int | None
    confidence: LagConfidence
    spectrum: Mapping[int, CandidateCorrelation]

    def top_candidates(self, limit: int = 5) -> list[CandidateCorrelation]:
        """Candidates ordered by the configured criterion then shortest lag."""
        if limit < 1:
            return []
        return sorted(self.spectrum.values(), key=lambda item: (-_score(item.correlation, self.selection_mode), item.lag))[:limit]


def align_to_regular_grid(
    source: pd.Series,
    target: pd.Series,
    *,
    frequency: str = "1D",
    source_forward_fill: bool = True,
    source_fill_limit: int | None = None,
) -> tuple[pd.Series, pd.Series]:
    """Resample two timestamped series deterministically onto a common grid.

    Each grid bucket receives its last observed value. Source values may be
    carried forward only from their availability timestamp; target values are
    never forward-filled. Callers must pass a source indexed by publication /
    availability time rather than a later-revised economic-period timestamp.
    """
    _validate_series(source, "source")
    _validate_series(target, "target")
    source_grid = source.astype(float).resample(frequency).last()
    if source_forward_fill:
        source_grid = source_grid.ffill(limit=source_fill_limit)
    target_grid = target.astype(float).resample(frequency).last()
    index = source_grid.index.union(target_grid.index).sort_values()
    return source_grid.reindex(index), target_grid.reindex(index)


class AdaptiveLagCorrelationEngine:
    def __init__(self, config: LagCorrelationConfig) -> None:
        self.config = config

    def estimate(self, source: pd.Series, target: pd.Series) -> list[LagEstimate]:
        """Return the causal estimate at every aligned target timestamp."""
        self._validate_aligned(source, target)
        return [self.estimate_at(source, target, timestamp) for timestamp in target.index]

    def estimate_at(
        self, source: pd.Series, target: pd.Series, evaluated_at: pd.Timestamp | str
    ) -> LagEstimate:
        """Estimate using observations at or before ``evaluated_at`` only."""
        self._validate_aligned(source, target)
        timestamp = pd.Timestamp(evaluated_at)
        if timestamp not in target.index:
            raise ValueError("evaluated_at must be an index value in target")
        end = target.index.get_loc(timestamp)
        if not isinstance(end, (int, np.integer)):
            raise ValueError("target index must be unique")

        source_values = transform_series(source, self.config.transform).to_numpy(dtype=float)
        target_values = transform_series(target, self.config.transform).to_numpy(dtype=float)
        start = max(0, int(end) - self.config.window + 1)
        spectrum: dict[int, CandidateCorrelation] = {}
        for lag in self.config.lags:
            target_positions = np.arange(start, int(end) + 1)
            source_positions = target_positions - lag
            usable = source_positions >= 0
            x = source_values[source_positions[usable]]
            y = target_values[target_positions[usable]]
            valid = np.isfinite(x) & np.isfinite(y)
            if valid.sum() < self.config.min_observations:
                continue
            x, y = x[valid], y[valid]
            if np.std(x) == 0 or np.std(y) == 0:
                continue
            spectrum[lag] = CandidateCorrelation(
                lag=lag,
                correlation=float(np.corrcoef(x, y)[0, 1]),
                observations=int(valid.sum()),
            )

        if not spectrum:
            return LagEstimate(timestamp, self.config.selection_mode, None, None, None, _empty_confidence(), spectrum)

        ranked = sorted(spectrum.values(), key=self._rank_key)
        best = ranked[0]
        # Deterministic tie policy: scores within tolerance are tied; use shortest lag.
        tied = [candidate for candidate in ranked if self._scores_tied(candidate, best)]
        best = min(tied, key=lambda candidate: candidate.lag)
        runner = next((candidate for candidate in ranked if candidate.lag != best.lag), None)
        confidence = self._confidence(best, runner, spectrum)
        return LagEstimate(timestamp, self.config.selection_mode, best.lag, best.correlation, best.observations, confidence, spectrum)

    def _rank_key(self, candidate: CandidateCorrelation) -> tuple[float, int]:
        score = self._selection_score(candidate.correlation)
        return (-score, candidate.lag)

    def _selection_score(self, correlation: float) -> float:
        return _score(correlation, self.config.selection_mode)

    def _scores_tied(self, left: CandidateCorrelation, right: CandidateCorrelation) -> bool:
        return abs(self._selection_score(left.correlation) - self._selection_score(right.correlation)) <= self.config.tie_tolerance

    def _confidence(
        self,
        best: CandidateCorrelation,
        runner: CandidateCorrelation | None,
        spectrum: Mapping[int, CandidateCorrelation],
    ) -> LagConfidence:
        margin = self._margin(best, runner) if runner is not None else None
        distinct_peak = self._best_distinct_local_peak(best, spectrum)
        distinct_margin = self._margin(best, distinct_peak) if distinct_peak is not None else None
        denominator = max(abs(best.correlation), np.finfo(float).eps)
        return LagConfidence(
            runner_up_lag=runner.lag if runner else None,
            runner_up_correlation=runner.correlation if runner else None,
            margin=margin,
            relative_peak_separation=margin / denominator if margin is not None else None,
            distinct_peak_lag=distinct_peak.lag if distinct_peak else None,
            distinct_peak_correlation=distinct_peak.correlation if distinct_peak else None,
            distinct_peak_margin=distinct_margin,
            distinct_peak_relative_separation=(
                distinct_margin / denominator if distinct_margin is not None else None
            ),
        )

    def _best_distinct_local_peak(
        self, best: CandidateCorrelation, spectrum: Mapping[int, CandidateCorrelation]
    ) -> CandidateCorrelation | None:
        ordered = sorted(spectrum.values(), key=lambda candidate: candidate.lag)
        local_peaks = []
        for position, candidate in enumerate(ordered):
            left = ordered[position - 1] if position else None
            right = ordered[position + 1] if position + 1 < len(ordered) else None
            if left and self._selection_score(candidate.correlation) < self._selection_score(left.correlation) - self.config.tie_tolerance:
                continue
            if right and self._selection_score(candidate.correlation) < self._selection_score(right.correlation) - self.config.tie_tolerance:
                continue
            # Represent a flat local maximum once, at its shortest lag.
            if left and self._scores_tied(candidate, left):
                continue
            if abs(candidate.lag - best.lag) <= self.config.distinct_peak_exclusion_radius:
                continue
            local_peaks.append(candidate)
        return min(local_peaks, key=self._rank_key) if local_peaks else None

    def _margin(self, best: CandidateCorrelation, competitor: CandidateCorrelation) -> float:
        if self._scores_tied(best, competitor):
            return 0.0
        return max(0.0, self._selection_score(best.correlation) - self._selection_score(competitor.correlation))

    @staticmethod
    def _validate_aligned(source: pd.Series, target: pd.Series) -> None:
        _validate_series(source, "source")
        _validate_series(target, "target")
        if not source.index.equals(target.index):
            raise ValueError("source and target must share the same aligned index")


def _empty_confidence() -> LagConfidence:
    return LagConfidence(None, None, None, None, None, None, None, None)


def _score(correlation: float, mode: SelectionMode) -> float:
    if mode is SelectionMode.STRONGEST_POSITIVE:
        return correlation
    if mode is SelectionMode.STRONGEST_NEGATIVE:
        return -correlation
    return abs(correlation)


def _validate_series(series: pd.Series, name: str) -> None:
    if not isinstance(series, pd.Series):
        raise TypeError(f"{name} must be a pandas Series")
    if not isinstance(series.index, pd.DatetimeIndex):
        raise ValueError(f"{name} index must be a DatetimeIndex")
    if not series.index.is_monotonic_increasing:
        raise ValueError(f"{name} index must be monotonic increasing")
    if not series.index.is_unique:
        raise ValueError(f"{name} index must be unique")


def transform_series(series: pd.Series, transform: Transform) -> pd.Series:
    numeric = series.astype(float)
    if transform is Transform.LEVELS:
        return numeric
    if transform is Transform.DIFFERENCE:
        return numeric.diff()
    if transform is Transform.PCT_CHANGE:
        return numeric.pct_change(fill_method=None)
    positive = numeric.where(numeric > 0)
    return np.log(positive).diff()