from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

COORDS = ["trust_obs", "shock_obs", "policy_proxy"]
OVERLAP_ORDER = ["pre--shock", "shock--transition", "transition--post"]


@dataclass(frozen=True)
class LocalFit:
    coefficients: np.ndarray
    centre: float
    scale: float
    degree: int


def expected_quarters() -> list[str]:
    return [f"{year}Q{quarter}" for year in range(2000, 2020) for quarter in range(1, 5)]


def load_panel(root: Path) -> pd.DataFrame:
    path = root / "data/input/tunisia_aligned_quarterly_panel.csv"
    panel = pd.read_csv(path)
    required = {"quarter", "year", "quarter_index", *COORDS}
    missing = required.difference(panel.columns)
    if missing:
        raise ValueError(f"Analysis input is missing columns: {sorted(missing)}")
    if len(panel) != 80:
        raise ValueError(f"Expected 80 quarterly rows, found {len(panel)}")
    if panel["quarter"].tolist() != expected_quarters():
        raise ValueError("Quarter sequence is not exactly 2000Q1--2019Q4")
    if panel["quarter"].duplicated().any():
        raise ValueError("Duplicate quarter labels found")
    expected_index = np.arange(80)
    if not np.array_equal(panel["quarter_index"].to_numpy(int), expected_index):
        raise ValueError("quarter_index must be the zero-based row index 0--79")
    if panel[COORDS].isna().any().any():
        raise ValueError("Missing core coordinate values found")
    values = panel[COORDS].to_numpy(float)
    if np.any((values < 0.0) | (values > 1.0)):
        raise ValueError("Core coordinates must lie on [0,1]")
    return panel


def _chart(indices: range | list[int], centre: float, scale: float) -> dict[str, object]:
    return {"indices": list(indices), "centre": float(centre), "scale": float(scale)}


def baseline_charts(extended_shock: bool = False) -> dict[str, dict[str, object]]:
    shock = _chart(range(42, 46), 43.5, float(np.std(np.arange(42, 46)))) if extended_shock else _chart(range(43, 45), 43.5, 1.0)
    return {
        "pre": _chart(range(0, 43), 21.0, 10.744),
        "shock": shock,
        "transition": _chart(range(45, 60), 52.0, 4.320),
        "post": _chart(range(60, 80), 69.5, 5.000),
    }


def audit_edges() -> list[tuple[str, str, str, np.ndarray]]:
    return [
        ("pre--shock", "pre", "shock", np.arange(40, 46)),
        ("shock--transition", "shock", "transition", np.arange(43, 50)),
        ("transition--post", "transition", "post", np.arange(56, 64)),
    ]


def fit_local(values: np.ndarray, chart: dict[str, object], degree: int = 2, ridge: float = 1.0) -> LocalFit:
    if degree not in (1, 2):
        raise ValueError("degree must be 1 or 2")
    if ridge <= 0:
        raise ValueError("ridge must be positive")
    indices = np.asarray(chart["indices"], dtype=int)
    centre = float(chart["centre"])
    scale = float(chart["scale"])
    if len(indices) == 0 or scale <= 0:
        raise ValueError("chart must have non-empty indices and positive scale")
    x = (indices - centre) / scale
    design = np.column_stack([x ** power for power in range(degree + 1)])
    penalty = np.diag([0.0] + [1.0] * degree)
    coefficients = np.linalg.solve(design.T @ design + ridge * penalty, design.T @ values[indices])
    return LocalFit(coefficients=coefficients, centre=centre, scale=scale, degree=degree)


def predict(fit: LocalFit, indices: np.ndarray | list[int]) -> np.ndarray:
    q = np.asarray(indices, dtype=float)
    x = (q - fit.centre) / fit.scale
    design = np.column_stack([x ** power for power in range(fit.degree + 1)])
    return design @ fit.coefficients


def residual_from_difference(difference: np.ndarray, gauge_fraction: float = 0.15) -> dict[str, np.ndarray | float]:
    difference = np.asarray(difference, dtype=float)
    if difference.ndim != 2 or difference.shape[0] == 0:
        raise ValueError("difference must be a non-empty two-dimensional array")
    if gauge_fraction < 0:
        raise ValueError("gauge_fraction must be non-negative")
    raw_rms = float(np.sqrt(np.mean(difference ** 2)))
    translation = np.median(difference, axis=0)
    translated = difference - translation
    after_translation = float(np.sqrt(np.mean(translated ** 2)))
    component_rms = np.sqrt(np.mean(translated ** 2, axis=0))
    gauge_bounds = gauge_fraction * component_rms
    correction = np.clip(np.mean(translated, axis=0), -gauge_bounds, gauge_bounds)
    corrected = translated - correction
    after_gauge = float(np.sqrt(np.mean(corrected ** 2)))
    corrected_component_rms = np.sqrt(np.mean(corrected ** 2, axis=0))
    return {
        "difference": difference,
        "raw_rms": raw_rms,
        "translation": translation,
        "translated_difference": translated,
        "after_translation": after_translation,
        "gauge_bounds": gauge_bounds,
        "correction": correction,
        "corrected_difference": corrected,
        "after_gauge": after_gauge,
        "translation_norm": float(np.linalg.norm(translation)),
        "corrected_component_rms": corrected_component_rms,
    }


def run_specification(
    values: np.ndarray,
    degree: int = 2,
    ridge: float = 1.0,
    gauge_fraction: float = 0.15,
    extended_shock: bool = False,
) -> dict[str, dict[str, np.ndarray | float]]:
    charts = baseline_charts(extended_shock)
    fits = {name: fit_local(values, chart, degree=degree, ridge=ridge) for name, chart in charts.items()}
    results: dict[str, dict[str, np.ndarray | float]] = {}
    for overlap, left, right, window in audit_edges():
        difference = predict(fits[right], window) - predict(fits[left], window)
        results[overlap] = residual_from_difference(difference, gauge_fraction)
    return results


def circular_block_bootstrap(
    difference: np.ndarray,
    rng: np.random.Generator,
    replications: int,
    block_length: int,
    gauge_fraction: float,
) -> np.ndarray:
    difference = np.asarray(difference, dtype=float)
    n = len(difference)
    if n == 0 or replications <= 0 or block_length <= 0:
        raise ValueError("non-empty difference and positive replication/block counts are required")
    blocks_per_draw = int(np.ceil(n / block_length))
    draws = np.empty(replications, dtype=float)
    for replication in range(replications):
        starts = rng.integers(0, n, size=blocks_per_draw)
        indices: list[int] = []
        for start in starts:
            indices.extend((int(start) + offset) % n for offset in range(block_length))
        sampled = difference[np.asarray(indices[:n], dtype=int)]
        draws[replication] = float(residual_from_difference(sampled, gauge_fraction)["after_gauge"])
    return draws


def common_window_score(values: np.ndarray, boundary_index: int) -> float:
    if boundary_index < 8 or boundary_index + 8 > len(values):
        raise ValueError("boundary requires eight quarters on each side")
    left_indices = np.arange(boundary_index - 8, boundary_index)
    right_indices = np.arange(boundary_index, boundary_index + 8)
    left_chart = _chart(left_indices.tolist(), float(np.mean(left_indices)), float(np.std(left_indices)))
    right_chart = _chart(right_indices.tolist(), float(np.mean(right_indices)), float(np.std(right_indices)))
    left_fit = fit_local(values, left_chart, degree=2, ridge=1.0)
    right_fit = fit_local(values, right_chart, degree=2, ridge=1.0)
    audit_window = np.arange(boundary_index - 3, boundary_index + 3)
    difference = predict(right_fit, audit_window) - predict(left_fit, audit_window)
    return float(residual_from_difference(difference, 0.15)["after_gauge"])


def mechanical_class(score: float, tau_low: float, tau_high: float) -> str:
    if tau_high <= tau_low:
        raise ValueError("tau_high must exceed tau_low")
    if score < tau_low:
        return "Conditional pooling"
    if score <= tau_high:
        return "Review before pooling"
    return "Keep separate"
