from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor
from dataclasses import asdict, dataclass
from hashlib import sha256
from pathlib import Path
import argparse
import json
import os
import platform
import sys

import numpy as np
import pandas as pd

from common_models import (
    acceptable_indices,
    causal_ewma_adjust,
    ensemble_quantile_es,
    fit_models,
    model_quantile_es,
    observation_log_losses,
    score_weights,
)


HERE = Path(__file__).resolve().parent
PROJECT = HERE.parent
DATA_DIR = (
    PROJECT.parent
    / "distribution_uncertainty_traceback"
    / "output_crossasset"
    / "parsed"
)
DEFAULT_OUTPUT = PROJECT / "results" / "study2"
PROTOCOL = PROJECT / "08_FROZEN_CONFIRMATORY_PROTOCOL.md"

ASSET_FILES = {
    "SPY": "SPY.csv",
    "QQQ": "QQQ.csv",
    "IWM": "IWM.csv",
    "EFA": "EFA.csv",
    "TLT": "TLT.csv",
    "IEF": "IEF.csv",
    "LQD": "LQD.csv",
    "HYG": "HYG.csv",
    "GLD": "GLD.csv",
    "EURUSD": "EURUSD_X.csv",
    "BTC": "BTC-USD.csv",
}
POLICIES = ("gaussian", "selected", "ensemble", "acceptable_set")


@dataclass(frozen=True)
class Config:
    seed: int = 2026072603
    training_n: int = 1000
    validation_n: int = 120
    test_n: int = 20
    ewma_decay: float = 0.94
    alpha: float = 0.05
    negligible_margin: float = 0.02
    bootstrap_reps: int = 2000
    capital_margin: float = 0.10


def load_returns(path: Path) -> pd.DataFrame:
    data = pd.read_csv(path)
    price_column = (
        "Adj Close"
        if "Adj Close" in data.columns
        else "AdjClose"
        if "AdjClose" in data.columns
        else None
    )
    if "Date" not in data.columns or price_column is None:
        raise ValueError(f"{path.name} lacks Date or adjusted close")
    data["Date"] = pd.to_datetime(data["Date"], errors="coerce")
    data["adjusted_close"] = pd.to_numeric(
        data[price_column], errors="coerce"
    )
    data = (
        data.dropna(subset=["Date", "adjusted_close"])
        .loc[lambda frame: frame["adjusted_close"] > 0]
        .sort_values("Date")
        .drop_duplicates("Date", keep="last")
        .reset_index(drop=True)
    )
    data["return"] = np.log(data["adjusted_close"]).diff()
    return data.dropna(subset=["return"]).reset_index(drop=True)


def evaluate_asset(
    task: tuple[str, str, Config, int | None]
) -> list[dict]:
    asset, filename, config, development_blocks = task
    data = load_returns(DATA_DIR / filename)
    residuals, scales = causal_ewma_adjust(
        data["return"].to_numpy(), decay=config.ewma_decay
    )
    first_test = config.training_n + config.validation_n
    starts = list(range(first_test, len(data) - config.test_n + 1, config.test_n))
    if development_blocks is not None:
        starts = starts[:development_blocks]
    rows: list[dict] = []
    prior_winner: str | None = None
    for block_index, test_start in enumerate(starts):
        training_start = test_start - config.validation_n - config.training_n
        training_end = test_start - config.validation_n
        validation_end = test_start
        training = residuals[training_start:training_end]
        validation = residuals[training_end:validation_end]
        test = residuals[test_start:test_start + config.test_n]
        try:
            models = fit_models(
                training,
                config.seed
                + list(ASSET_FILES).index(asset) * 1_000_000
                + block_index,
            )
            losses = observation_log_losses(models, validation)
            accepted, winner, _, _ = acceptable_indices(
                losses,
                alpha=config.alpha,
                negligible_margin=config.negligible_margin,
            )
            weights = score_weights(losses)
            model_decisions = [
                model_quantile_es(model, training, grid_n=2001)
                for model in models
            ]
            ensemble_decision = ensemble_quantile_es(
                models, weights, training, grid_n=2001
            )
            quantiles = np.asarray([item[0] for item in model_decisions])
            shortfalls = np.asarray([item[1] for item in model_decisions])
            policy_decisions = {
                "gaussian": (quantiles[0], shortfalls[0]),
                "selected": (quantiles[winner], shortfalls[winner]),
                "ensemble": ensemble_decision,
                "acceptable_set": (
                    float(np.min(quantiles[accepted])),
                    float(np.min(shortfalls[accepted])),
                ),
            }
            selected_name = models[winner].name
            turnover = (
                np.nan
                if prior_winner is None
                else float(selected_name != prior_winner)
            )
            prior_winner = selected_name
            for policy, (quantile, expected_shortfall) in policy_decisions.items():
                violations = test < quantile
                severity = (
                    float(np.mean(quantile - test[violations]))
                    if np.any(violations)
                    else 0.0
                )
                es_error = (
                    float(abs(np.mean(test[violations]) - expected_shortfall))
                    if np.any(violations)
                    else np.nan
                )
                rows.append(
                    {
                        "asset": asset,
                        "block": block_index,
                        "test_start": data.loc[test_start, "Date"],
                        "test_end": data.loc[
                            test_start + config.test_n - 1, "Date"
                        ],
                        "policy": policy,
                        "fit_failed": False,
                        "test_n": config.test_n,
                        "violation_count": int(np.sum(violations)),
                        "violation_rate": float(np.mean(violations)),
                        "standardized_capital": float(-quantile),
                        "expected_shortfall_capital": float(-expected_shortfall),
                        "exceedance_severity": severity,
                        "es_error": es_error,
                        "acceptable_count": len(accepted),
                        "selected_model": selected_name,
                        "selected_turnover": turnover,
                        "mean_raw_scale": float(
                            np.mean(scales[test_start:test_start + config.test_n])
                        ),
                    }
                )
        except Exception as error:
            for policy in POLICIES:
                rows.append(
                    {
                        "asset": asset,
                        "block": block_index,
                        "test_start": data.loc[test_start, "Date"],
                        "test_end": data.loc[
                            test_start + config.test_n - 1, "Date"
                        ],
                        "policy": policy,
                        "fit_failed": True,
                        "failure_type": type(error).__name__,
                        "failure_message": str(error)[:240],
                    }
                )
    return rows


def run_assets(
    config: Config,
    workers: int,
    assets: tuple[str, ...],
    development_blocks: int | None,
) -> pd.DataFrame:
    tasks = [
        (asset, ASSET_FILES[asset], config, development_blocks)
        for asset in assets
    ]
    rows: list[dict] = []
    if workers == 1:
        for task in tasks:
            rows.extend(evaluate_asset(task))
    else:
        with ProcessPoolExecutor(max_workers=workers) as executor:
            for result in executor.map(evaluate_asset, tasks):
                rows.extend(result)
    return pd.DataFrame(rows)


def asset_policy_summary(blocks: pd.DataFrame) -> pd.DataFrame:
    valid = blocks[~blocks.fit_failed].copy()
    summary = (
        valid.groupby(["asset", "policy"])
        .agg(
            blocks=("block", "nunique"),
            observations=("test_n", "sum"),
            violations=("violation_count", "sum"),
            mean_capital=("standardized_capital", "mean"),
            mean_es_capital=("expected_shortfall_capital", "mean"),
            mean_exceedance_severity=("exceedance_severity", "mean"),
            mean_es_error=("es_error", "mean"),
            mean_acceptable_count=("acceptable_count", "mean"),
            selected_turnover=("selected_turnover", "mean"),
        )
        .reset_index()
    )
    summary["violation_rate"] = summary.violations / summary.observations
    summary["coverage_error"] = abs(summary.violation_rate - 0.05)
    return summary


def holm_adjust(pvalues: dict[str, float]) -> dict[str, float]:
    ordered = sorted(pvalues, key=pvalues.get)
    adjusted: dict[str, float] = {}
    running = 0.0
    count = len(ordered)
    for rank, key in enumerate(ordered):
        value = min((count - rank) * pvalues[key], 1.0)
        running = max(running, value)
        adjusted[key] = running
    return adjusted


def bootstrap_inference(
    blocks: pd.DataFrame, config: Config
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    valid = blocks[~blocks.fit_failed].copy()
    assets = sorted(valid.asset.unique())
    policy_index = {policy: index for index, policy in enumerate(POLICIES)}
    arrays: dict[str, tuple[np.ndarray, np.ndarray, int]] = {}
    for asset in assets:
        subset = valid[valid.asset == asset]
        pivot_v = subset.pivot(
            index="block", columns="policy", values="violation_count"
        )[list(POLICIES)]
        pivot_c = subset.pivot(
            index="block", columns="policy", values="standardized_capital"
        )[list(POLICIES)]
        arrays[asset] = (
            pivot_v.to_numpy(),
            pivot_c.to_numpy(),
            config.test_n,
        )

    rng = np.random.default_rng(config.seed)
    comparison_names = ("gaussian", "selected")
    coverage_draws = {
        comparator: np.zeros(config.bootstrap_reps)
        for comparator in comparison_names
    }
    capital_draws = np.zeros(config.bootstrap_reps)
    for replication in range(config.bootstrap_reps):
        sampled_assets = rng.choice(assets, size=len(assets), replace=True)
        errors = []
        capitals = []
        for asset in sampled_assets:
            violations, capital, test_n = arrays[asset]
            indices = rng.integers(0, len(violations), size=len(violations))
            rates = violations[indices].sum(axis=0) / (len(indices) * test_n)
            errors.append(np.abs(rates - 0.05))
            capitals.append(capital[indices].mean(axis=0))
        mean_errors = np.mean(errors, axis=0)
        mean_capital = np.mean(capitals, axis=0)
        for comparator in comparison_names:
            coverage_draws[comparator][replication] = (
                mean_errors[policy_index["acceptable_set"]]
                - mean_errors[policy_index[comparator]]
            )
        capital_draws[replication] = (
            mean_capital[policy_index["acceptable_set"]]
            / mean_capital[policy_index["selected"]]
            - 1
        )

    summary = asset_policy_summary(valid)
    wide_error = summary.pivot(
        index="asset", columns="policy", values="coverage_error"
    )
    wide_capital = summary.pivot(
        index="asset", columns="policy", values="mean_capital"
    )
    raw_pvalues = {
        comparator: float(
            (1 + np.sum(coverage_draws[comparator] >= 0))
            / (config.bootstrap_reps + 1)
        )
        for comparator in comparison_names
    }
    adjusted = holm_adjust(raw_pvalues)
    rows = []
    for comparator in comparison_names:
        observed = float(
            (
                wide_error["acceptable_set"] - wide_error[comparator]
            ).mean()
        )
        draw = coverage_draws[comparator]
        rows.append(
            {
                "comparison": f"acceptable_set_minus_{comparator}",
                "outcome": "coverage_error",
                "estimate": observed,
                "ci_low": float(np.quantile(draw, 0.025)),
                "ci_high": float(np.quantile(draw, 0.975)),
                "p_one_sided": raw_pvalues[comparator],
                "p_holm": adjusted[comparator],
                "assets_improved": int(
                    (wide_error["acceptable_set"] < wide_error[comparator]).sum()
                ),
                "assets_total": len(assets),
            }
        )
    observed_capital = float(
        wide_capital["acceptable_set"].mean()
        / wide_capital["selected"].mean()
        - 1
    )
    rows.append(
        {
            "comparison": "acceptable_set_relative_to_selected",
            "outcome": "relative_capital",
            "estimate": observed_capital,
            "ci_low": float(np.quantile(capital_draws, 0.025)),
            "ci_high": float(np.quantile(capital_draws, 0.975)),
            "p_one_sided": np.nan,
            "p_holm": np.nan,
            "assets_improved": np.nan,
            "assets_total": len(assets),
        }
    )
    inference = pd.DataFrame(rows)

    loo_rows = []
    for omitted in assets:
        remaining = wide_error.drop(index=omitted)
        for comparator in comparison_names:
            loo_rows.append(
                {
                    "omitted_asset": omitted,
                    "comparator": comparator,
                    "coverage_error_difference": float(
                        (
                            remaining["acceptable_set"]
                            - remaining[comparator]
                        ).mean()
                    ),
                }
            )
    leave_one_out = pd.DataFrame(loo_rows)

    gaussian = inference[
        inference.comparison == "acceptable_set_minus_gaussian"
    ].iloc[0]
    selected = inference[
        inference.comparison == "acceptable_set_minus_selected"
    ].iloc[0]
    capital = inference[
        inference.outcome == "relative_capital"
    ].iloc[0]
    gates = pd.DataFrame(
        [
            {
                "gate": "coverage_better_than_gaussian",
                "observed": gaussian.estimate,
                "rule": "estimate < 0 and Holm p < .05",
                "passed": gaussian.estimate < 0 and gaussian.p_holm < 0.05,
            },
            {
                "gate": "coverage_better_than_selected",
                "observed": selected.estimate,
                "rule": "estimate < 0 and Holm p < .05",
                "passed": selected.estimate < 0 and selected.p_holm < 0.05,
            },
            {
                "gate": "capital_margin",
                "observed": capital.ci_high,
                "rule": "95% CI upper bound <= 0.10",
                "passed": capital.ci_high <= config.capital_margin,
            },
            {
                "gate": "majority_assets_gaussian",
                "observed": gaussian.assets_improved,
                "rule": f"> {len(assets) / 2}",
                "passed": gaussian.assets_improved > len(assets) / 2,
            },
            {
                "gate": "majority_assets_selected",
                "observed": selected.assets_improved,
                "rule": f"> {len(assets) / 2}",
                "passed": selected.assets_improved > len(assets) / 2,
            },
            {
                "gate": "leave_one_out_direction",
                "observed": float(
                    (leave_one_out.coverage_error_difference < 0).all()
                ),
                "rule": "all differences < 0",
                "passed": bool(
                    (leave_one_out.coverage_error_difference < 0).all()
                ),
            },
        ]
    )
    return inference, leave_one_out, gates


def write_summary(
    output: Path,
    blocks: pd.DataFrame,
    summary: pd.DataFrame,
    inference: pd.DataFrame,
    gates: pd.DataFrame,
) -> None:
    failure = float(
        blocks.groupby(["asset", "block"]).fit_failed.max().mean()
    )
    all_passed = bool(gates.passed.all() and failure < 0.05)
    lines = [
        "# Study 2 Financial Tail-Decision Results",
        "",
        f"- Assets: {summary.asset.nunique()}",
        f"- Non-overlapping asset-blocks: "
        f"{blocks[['asset', 'block']].drop_duplicates().shape[0]:,}",
        f"- Block-level fit failure rate: {failure:.3%}",
        "",
        "## Equal-weight asset inference",
        "",
        inference.to_markdown(index=False),
        "",
        "## Decision",
        "",
        "The frozen financial-value claim "
        + ("passes." if all_passed else "does not pass."),
        "",
        gates.to_markdown(index=False),
    ]
    (output / "study2_results.md").write_text(
        "\n".join(lines) + "\n", encoding="utf-8"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=min(4, os.cpu_count() or 1))
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument(
        "--assets",
        default=",".join(ASSET_FILES),
        help="Comma-separated asset names; confirmation uses the default.",
    )
    parser.add_argument(
        "--development-blocks",
        type=int,
        default=None,
        help="Limit blocks only for code development; never use for confirmation.",
    )
    parser.add_argument("--bootstrap-reps", type=int, default=2000)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    assets = tuple(item.strip() for item in args.assets.split(",") if item.strip())
    unknown = set(assets) - set(ASSET_FILES)
    if unknown:
        raise ValueError(f"Unknown assets: {sorted(unknown)}")
    config = Config(bootstrap_reps=args.bootstrap_reps)
    args.output.mkdir(parents=True, exist_ok=True)
    blocks = run_assets(
        config, args.workers, assets, args.development_blocks
    )
    summary = asset_policy_summary(blocks)
    inference, leave_one_out, gates = bootstrap_inference(blocks, config)
    failure_rate = float(
        blocks.groupby(["asset", "block"]).fit_failed.max().mean()
    )
    gates = pd.concat(
        [
            gates,
            pd.DataFrame(
                [
                    {
                        "gate": "fit_failure_rate",
                        "observed": failure_rate,
                        "rule": "< 0.05",
                        "passed": failure_rate < 0.05,
                    }
                ]
            ),
        ],
        ignore_index=True,
    )
    gates = pd.concat(
        [
            gates,
            pd.DataFrame(
                [
                    {
                        "gate": "all_study2_gates",
                        "observed": float(gates.passed.all()),
                        "rule": "equals 1",
                        "passed": bool(gates.passed.all()),
                    }
                ]
            ),
        ],
        ignore_index=True,
    )
    blocks.to_csv(args.output / "study2_blocks.csv", index=False)
    summary.to_csv(args.output / "study2_asset_policy.csv", index=False)
    inference.to_csv(args.output / "study2_inference.csv", index=False)
    leave_one_out.to_csv(args.output / "study2_leave_one_out.csv", index=False)
    gates.to_csv(args.output / "study2_gates.csv", index=False)
    write_summary(args.output, blocks, summary, inference, gates)
    source_hashes = {
        asset: sha256((DATA_DIR / ASSET_FILES[asset]).read_bytes()).hexdigest()
        for asset in assets
    }
    manifest = {
        "config": asdict(config),
        "assets": assets,
        "asset_files": {asset: ASSET_FILES[asset] for asset in assets},
        "source_sha256": source_hashes,
        "script_sha256": sha256(Path(__file__).read_bytes()).hexdigest(),
        "common_models_sha256": sha256(
            (HERE / "common_models.py").read_bytes()
        ).hexdigest(),
        "protocol_sha256": sha256(PROTOCOL.read_bytes()).hexdigest(),
        "python": sys.version,
        "platform": platform.platform(),
    }
    (args.output / "study2_manifest.json").write_text(
        json.dumps(manifest, indent=2), encoding="utf-8"
    )


if __name__ == "__main__":
    main()
