"""Signal correlation and causality analysis.

Cross-correlation between signals, and automated root cause analysis
("what signals predict this error going high?").

Usage:
    corr = ewal.correlate(sig_a, sig_b, clock=clk, max_lag=20)
    predictors = ewal.find_predictors(target=error, candidates=t.signals, ...)
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ewal.trace import Trace
    from ewal.signal import SignalRef


@dataclass
class CorrelationResult:
    peak_lag: int
    peak_r: float
    lags: list[int] = field(default_factory=list)
    correlations: list[float] = field(default_factory=list)


@dataclass
class Predictor:
    signal: str
    lag: int
    r: float


def correlate(sig_a, sig_b, clock=None, max_lag: int = 20) -> CorrelationResult:
    """Cross-correlate two signals.

    Computes correlation at each lag from -max_lag to +max_lag.
    Returns the lag and correlation coefficient at peak.
    """
    trace = None
    for s in [sig_a, sig_b]:
        if hasattr(s, "_trace"):
            trace = s._trace
            break

    if trace is None:
        return CorrelationResult(peak_lag=0, peak_r=0.0)

    # Get signal values at clock edges (or all indices)
    if clock is not None and hasattr(clock, "_trace"):
        from ewal.signal import _compile_expr
        compiled = _compile_expr(clock.rising(), trace)
        indices = [i for i in range(trace.max_index + 1)
                   if trace._core.eval_expr_bool(compiled, i)]
    else:
        indices = list(range(trace.max_index + 1))

    n = len(indices)
    if n < 2:
        return CorrelationResult(peak_lag=0, peak_r=0.0)

    # Extract values
    path_a = sig_a._path if hasattr(sig_a, "_path") else str(sig_a)
    path_b = sig_b._path if hasattr(sig_b, "_path") else str(sig_b)

    vals_a = [float(trace._get_signal_value(path_a, i) or 0) for i in indices]
    vals_b = [float(trace._get_signal_value(path_b, i) or 0) for i in indices]

    mean_a = sum(vals_a) / n
    mean_b = sum(vals_b) / n
    std_a = (sum((v - mean_a) ** 2 for v in vals_a) / n) ** 0.5
    std_b = (sum((v - mean_b) ** 2 for v in vals_b) / n) ** 0.5

    if std_a == 0 or std_b == 0:
        return CorrelationResult(peak_lag=0, peak_r=0.0)

    lags = []
    correlations = []
    best_lag = 0
    best_r = 0.0

    for lag in range(-max_lag, max_lag + 1):
        # Cross-correlation at this lag
        s = 0.0
        count = 0
        for i in range(n):
            j = i + lag
            if 0 <= j < n:
                s += (vals_a[i] - mean_a) * (vals_b[j] - mean_b)
                count += 1

        if count > 0:
            r = s / (count * std_a * std_b)
        else:
            r = 0.0

        lags.append(lag)
        correlations.append(r)

        if abs(r) > abs(best_r):
            best_r = r
            best_lag = lag

    return CorrelationResult(
        peak_lag=best_lag,
        peak_r=best_r,
        lags=lags,
        correlations=correlations,
    )


def find_predictors(target, candidates: list, clock=None,
                    window: int = 50) -> list[Predictor]:
    """Find signals that predict the target signal going high.

    Tests each candidate signal's cross-correlation with the target
    and returns the best predictors sorted by correlation strength.
    """
    trace = None
    if hasattr(target, "_trace"):
        trace = target._trace

    if trace is None:
        return []

    target_path = target._path if hasattr(target, "_path") else str(target)
    results = []

    for cand in candidates:
        if isinstance(cand, str):
            # Build a SignalRef from name
            from ewal.signal import SignalRef
            cand_ref = SignalRef(trace, cand)
        else:
            cand_ref = cand

        cand_path = cand_ref._path if hasattr(cand_ref, "_path") else str(cand_ref)

        # Skip self-correlation
        if cand_path == target_path:
            continue

        corr = correlate(target, cand_ref, clock=clock, max_lag=window)

        if abs(corr.peak_r) > 0.1:  # Threshold
            results.append(Predictor(
                signal=cand_path,
                lag=corr.peak_lag,
                r=corr.peak_r,
            ))

    results.sort(key=lambda p: -abs(p.r))
    return results
