#!/usr/bin/env python3
"""Audit temporal integrity and reproducibility of the frozen studies.

This audit leaves the frozen analysis code and outputs unchanged.  It checks
the released snapshots, reconstructs every Study 2 time boundary, verifies
prefix invariance of the EWMA transformation, and independently recomputes
selected Study 1 and Study 2 cells from the preserved inputs.
"""

from __future__ import annotations

from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
import json
import sys

import numpy as np
import pandas as pd


HERE = Path(__file__).resolve().parent
PROJECT = HERE.parent
ANALYSIS_CODE = (
    HERE if (HERE / "common_models.py").exists() else HERE / "frozen"
)
sys.path.insert(0, str(ANALYSIS_CODE))

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


PACKAGE = (
    PROJECT / "replication_package"
    if (PROJECT / "replication_package").exists()
    else PROJECT
)
RAW_DATA = PACKAGE / "data" / "raw"
FROZEN_CODE = PACKAGE / "code" / "frozen"
STUDY1_CELLS = PROJECT / "results" / "study1" / "study1_cells.csv"
STUDY2_BLOCKS = PROJECT / "results" / "study2" / "study2_blocks.csv"
STUDY1_MANIFEST = PROJECT / "results" / "study1" / "study1_manifest.json"
STUDY2_MANIFEST = PROJECT / "results" / "study2" / "study2_manifest.json"
OUTPUT_DIR = PROJECT / "results" / "temporal_integrity_audit"
REPORT = (
    PROJECT / "paper_rewriting_output" / "temporal_integrity_audit.md"
    if (PROJECT / "paper_rewriting_output").exists()
    else OUTPUT_DIR / "temporal_integrity_audit.md"
)


@dataclass
class Check:
    check_id: str
    scope: str
    status: str
    evidence: str


CHECKS: list[Check] = []


def record(check_id: str, scope: str, passed: bool, evidence: str) -> None:
    CHECKS.append(
        Check(check_id, scope, "PASS" if passed else "FAIL", evidence)
    )


def file_hash(path: Path) -> str:
    return sha256(path.read_bytes()).hexdigest()


def check_release_identity() -> None:
    study1_manifest = json.loads(STUDY1_MANIFEST.read_text(encoding="utf-8"))
    study2_manifest = json.loads(STUDY2_MANIFEST.read_text(encoding="utf-8"))
    for name, manifest in (
        ("common_models.py", study2_manifest),
        ("run_study2.py", study2_manifest),
        ("run_study1.py", study1_manifest),
    ):
        local = ANALYSIS_CODE / name
        frozen = FROZEN_CODE / name
        expected_key = (
            "common_models_sha256"
            if name == "common_models.py"
            else "script_sha256"
        )
        hashes = (file_hash(local), file_hash(frozen), manifest[expected_key])
        record(
            f"release_identity_{name}",
            "release",
            len(set(hashes)) == 1,
            f"working=frozen=manifest={hashes[0]}",
        )

    inventory = pd.read_csv(
        PACKAGE / "data" / "asset_inventory.csv"
    ).set_index("asset")
    for asset, filename in run_study2.ASSET_FILES.items():
        snapshot = RAW_DATA / filename
        expected = study2_manifest["source_sha256"][asset]
        actual = file_hash(snapshot)
        inventory_hash = str(inventory.loc[asset, "sha256"])
        record(
            f"source_snapshot_{asset}",
            "Study 2 source data",
            actual == expected == inventory_hash,
            f"{filename}: {actual}",
        )


def check_study2_boundaries() -> dict[str, pd.DataFrame]:
    config = run_study2.Config()
    saved = pd.read_csv(
        STUDY2_BLOCKS, parse_dates=["test_start", "test_end"]
    )
    loaded: dict[str, pd.DataFrame] = {}
    representatives = (
        saved[saved["policy"] == "gaussian"]
        .sort_values(["asset", "block"])
        .copy()
    )
    boundary_ok = True
    nonoverlap_ok = True
    row_shape_ok = True
    date_ok = True
    counts = []

    for asset, filename in run_study2.ASSET_FILES.items():
        data = run_study2.load_returns(RAW_DATA / filename)
        loaded[asset] = data
        starts = list(
            range(
                config.training_n + config.validation_n,
                len(data) - config.test_n + 1,
                config.test_n,
            )
        )
        observed = representatives[representatives.asset == asset]
        counts.append((asset, len(starts), len(observed)))
        row_shape_ok &= list(observed.block) == list(range(len(starts)))
        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_start = training_end
            validation_end = test_start
            test_end = test_start + config.test_n
            boundary_ok &= (
                training_start >= 0
                and training_end - training_start == config.training_n
                and validation_end - validation_start == config.validation_n
                and test_end - test_start == config.test_n
                and training_end == validation_start
                and validation_end == test_start
            )
            row = observed.iloc[block_index]
            date_ok &= (
                row.test_start == data.loc[test_start, "Date"]
                and row.test_end == data.loc[test_end - 1, "Date"]
            )
        if len(starts) > 1:
            nonoverlap_ok &= bool(
                np.all(np.diff(np.asarray(starts)) >= config.test_n)
            )

    policies_per_block = (
        saved.groupby(["asset", "block"])["policy"].nunique()
    )
    record(
        "study2_window_order",
        "all 2,710 Study 2 blocks",
        boundary_ok,
        "training[1000] < validation[120] < sealed test[20] for every block",
    )
    record(
        "study2_test_nonoverlap",
        "all assets",
        nonoverlap_ok,
        "successive test starts are separated by exactly 20 observations",
    )
    record(
        "study2_saved_dates",
        "all 2,710 Study 2 blocks",
        date_ok,
        "saved test_start/test_end equal dates reconstructed from raw snapshots",
    )
    record(
        "study2_block_inventory",
        "all assets",
        row_shape_ok and all(a == b for _, a, b in counts),
        "; ".join(f"{asset}={expected}" for asset, expected, _ in counts),
    )
    record(
        "study2_policy_rows",
        "all 2,710 Study 2 blocks",
        bool((policies_per_block == 4).all()),
        "four frozen policies are present for every asset-block",
    )
    return loaded


def check_ewma_prefix_invariance(
    loaded: dict[str, pd.DataFrame],
) -> None:
    config = run_study2.Config()
    all_equal = True
    max_difference = 0.0
    comparisons = 0
    for data in loaded.values():
        values = data["return"].to_numpy()
        full_adjusted, full_scales = causal_ewma_adjust(
            values, config.ewma_decay
        )
        starts = list(
            range(
                config.training_n + config.validation_n,
                len(values) - config.test_n + 1,
                config.test_n,
            )
        )
        cutoffs = sorted(
            {
                starts[0],
                starts[len(starts) // 2],
                starts[-1] + config.test_n,
            }
        )
        for cutoff in cutoffs:
            prefix_adjusted, prefix_scales = causal_ewma_adjust(
                values[:cutoff], config.ewma_decay
            )
            difference = max(
                float(
                    np.max(
                        np.abs(prefix_adjusted - full_adjusted[:cutoff])
                    )
                ),
                float(np.max(np.abs(prefix_scales - full_scales[:cutoff]))),
            )
            comparisons += 1
            max_difference = max(max_difference, difference)
            all_equal &= difference == 0.0
    record(
        "ewma_prefix_invariance",
        f"{comparisons} asset-cutoff comparisons",
        all_equal,
        f"maximum absolute change after truncating future data={max_difference:.3g}",
    )


def recompute_study2_block(
    asset: str,
    block_index: int,
    data: pd.DataFrame,
) -> list[dict]:
    config = run_study2.Config()
    values = data["return"].to_numpy()
    residuals, scales = causal_ewma_adjust(values, config.ewma_decay)
    test_start = (
        config.training_n + config.validation_n + block_index * config.test_n
    )
    training_start = (
        test_start - config.validation_n - config.training_n
    )
    training_end = test_start - config.validation_n
    training = residuals[training_start:training_end]
    validation = residuals[training_end:test_start]
    test = residuals[test_start : test_start + config.test_n]
    models = fit_models(
        training,
        config.seed
        + list(run_study2.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)
    decisions = [
        model_quantile_es(model, training, grid_n=2001) for model in models
    ]
    ensemble = ensemble_quantile_es(
        models, weights, training, grid_n=2001
    )
    quantiles = np.asarray([item[0] for item in decisions])
    shortfalls = np.asarray([item[1] for item in decisions])
    policies = {
        "gaussian": (quantiles[0], shortfalls[0]),
        "selected": (quantiles[winner], shortfalls[winner]),
        "ensemble": ensemble,
        "acceptable_set": (
            float(np.min(quantiles[accepted])),
            float(np.min(shortfalls[accepted])),
        ),
    }
    rows = []
    for policy, (quantile, expected_shortfall) in policies.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,
                "policy": policy,
                "violation_count": int(np.sum(violations)),
                "standardized_capital": float(-quantile),
                "expected_shortfall_capital": float(-expected_shortfall),
                "exceedance_severity": severity,
                "es_error": es_error,
                "acceptable_count": len(accepted),
                "selected_model": models[winner].name,
                "mean_raw_scale": float(
                    np.mean(scales[test_start : test_start + config.test_n])
                ),
            }
        )
    return rows


def compare_recomputed_study2(
    loaded: dict[str, pd.DataFrame],
) -> None:
    saved = pd.read_csv(STUDY2_BLOCKS)
    recomputed: list[dict] = []
    sampled_blocks = []
    for asset, data in loaded.items():
        block_count = int(saved.loc[saved.asset == asset, "block"].max()) + 1
        for block_index in sorted({0, block_count // 2, block_count - 1}):
            sampled_blocks.append((asset, block_index))
            recomputed.extend(
                recompute_study2_block(asset, block_index, data)
            )
    new = pd.DataFrame(recomputed).sort_values(
        ["asset", "block", "policy"]
    )
    old = saved.merge(
        pd.DataFrame(sampled_blocks, columns=["asset", "block"]),
        on=["asset", "block"],
        how="inner",
    ).sort_values(["asset", "block", "policy"])
    keys = ["asset", "block", "policy"]
    new = new.set_index(keys)
    old = old.set_index(keys).loc[new.index]
    exact_fields = [
        "violation_count",
        "acceptable_count",
        "selected_model",
    ]
    numeric_fields = [
        "standardized_capital",
        "expected_shortfall_capital",
        "exceedance_severity",
        "es_error",
        "mean_raw_scale",
    ]
    exact_ok = all(
        new[field].astype(str).equals(old[field].astype(str))
        for field in exact_fields
    )
    numeric_ok = all(
        np.allclose(
            new[field].to_numpy(float),
            old[field].to_numpy(float),
            rtol=2e-12,
            atol=2e-12,
            equal_nan=True,
        )
        for field in numeric_fields
    )
    max_difference = max(
        float(
            np.nanmax(
                np.abs(
                    new[field].to_numpy(float)
                    - old[field].to_numpy(float)
                )
            )
        )
        for field in numeric_fields
    )
    record(
        "study2_independent_recomputation",
        f"{len(sampled_blocks)} blocks / {len(new)} policy rows",
        exact_ok and numeric_ok,
        f"first/middle/last block per asset; max numeric difference={max_difference:.3g}",
    )


def compare_recomputed_study1() -> None:
    config = run_study1.Config()
    recomputed = []
    for dgp in run_study1.DGP_NAMES:
        recomputed.extend(run_study1.evaluate_replication((dgp, 0, config)))
    new = pd.DataFrame(recomputed).sort_values(
        ["dgp", "universe", "evidence_n"]
    )
    old = pd.read_csv(STUDY1_CELLS)
    old = old[old.rep == 0].sort_values(
        ["dgp", "universe", "evidence_n"]
    )
    keys = ["dgp", "rep", "universe", "evidence_n"]
    new = new.set_index(keys)
    old = old.set_index(keys).loc[new.index]
    exact_fields = [
        "model_count",
        "fit_failed",
        "winner",
        "winner_reversal",
        "acceptable_count",
    ]
    numeric_fields = [
        "best_validation_score",
        "effective_count",
        "var_diameter",
        "es_diameter",
        "sealed_test_regret",
        "minimum_test_score_gap",
        "score_difference_variance",
        "maximum_acceptable_lower_bound",
        "maximum_acceptable_se",
    ]
    exact_ok = all(
        new[field].astype(str).equals(old[field].astype(str))
        for field in exact_fields
    )
    numeric_ok = all(
        np.allclose(
            new[field].to_numpy(float),
            old[field].to_numpy(float),
            rtol=2e-12,
            atol=2e-12,
            equal_nan=True,
        )
        for field in numeric_fields
    )
    max_difference = max(
        float(
            np.nanmax(
                np.abs(
                    new[field].to_numpy(float)
                    - old[field].to_numpy(float)
                )
            )
        )
        for field in numeric_fields
    )
    record(
        "study1_independent_recomputation",
        f"8 DGP replications / {len(new)} cells",
        exact_ok and numeric_ok,
        f"models fit on training only; validation/test recomputed; max numeric difference={max_difference:.3g}",
    )


def write_outputs() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    checks = pd.DataFrame([item.__dict__ for item in CHECKS])
    checks.to_csv(OUTPUT_DIR / "temporal_integrity_checks.csv", index=False)
    failed = checks[checks.status == "FAIL"]
    overall = "PASS" if failed.empty else "FAIL"
    summary = {
        "overall": overall,
        "checks": len(checks),
        "passed": int((checks.status == "PASS").sum()),
        "failed": int((checks.status == "FAIL").sum()),
        "frozen_code_modified": False,
    }
    (OUTPUT_DIR / "summary.json").write_text(
        json.dumps(summary, indent=2), encoding="utf-8"
    )

    lines = [
        "# Temporal Integrity and Future-Data Audit",
        "",
        f"**Overall executable audit: {overall}.** "
        f"{summary['passed']} of {summary['checks']} checks passed.",
        "",
        "## Scope",
        "",
        "The audit verifies release identity, source-snapshot hashes, every "
        "Study 2 training-validation-test boundary, non-overlap of sealed "
        "test blocks, prefix invariance of the EWMA transformation, and "
        "independent recomputation of selected frozen outputs. The frozen "
        "analysis scripts and results were not altered.",
        "",
        "## Executable checks",
        "",
        checks.to_markdown(index=False),
        "",
        "## Manual information-flow review",
        "",
        "- **Study 1:** each model is fitted only on the first 1,000 simulated "
        "observations. Validation stream A constructs the empirical acceptable "
        "set and selects the winner; disjoint validation stream B measures "
        "winner reversal; the final 250 observations are used only for sealed-"
        "test regret. No sealed-test value enters fitting, screening, or tuning.",
        "- **Study 2:** models are fitted on 1,000 past standardized returns. "
        "The following 120 observations determine validation loss, the winner, "
        "weights, and the acceptable set. Quantiles are fixed before the next "
        "20 observations are read. Test observations are used only for "
        "violations, severity, and ES error.",
        "- **EWMA:** after an initial 60-observation variance estimate, the "
        "scale at time t is recorded before return t updates the recursion. "
        "All first decisions occur after 1,120 observations. The executable "
        "prefix tests show that deleting every later observation changes no "
        "earlier adjusted return or scale.",
        "- **Downstream inference:** the bootstrap reads only already sealed "
        "block outcomes. It does not refit models or feed test outcomes back "
        "into later policy choices. Prior selected winners are used only to "
        "report turnover.",
        "",
        "## Limitations that are not code leakage",
        "",
        "1. The first 60 EWMA-standardized observations share a batch variance "
        "initializer. This is not fully one-step online inside that historical "
        "initialization segment, but the segment lies entirely within the "
        "1,000-observation estimation history and cannot see any validation or "
        "test outcome. A strict-online initializer would be a new sensitivity "
        "specification, not a correction to the frozen results.",
        "2. The preserved files contain adjusted-close histories from a current "
        "data snapshot, not a point-in-time vendor database. Corporate-action "
        "revisions and current-vintage conventions are data-provenance risks. "
        "They are not caused by the train/validation/test code and cannot be "
        "eliminated without vintage-stamped source data.",
        "3. The eleven assets are a fixed surviving panel. Asset selection and "
        "survivorship limit external validity. Overlapping estimation histories "
        "and cross-asset dependence affect uncertainty quantification, not the "
        "direction of the information flow.",
        "",
        "## Conclusion",
        "",
        "No validation-to-training or test-to-model information path was found. "
        "The frozen numerical outputs sampled by this audit reproduce from the "
        "preserved inputs. The strongest defensible statement is therefore: "
        "**the released code is free of decision-time future-data leakage under "
        "the preserved data snapshot, while point-in-time data-vintage and "
        "survivorship limitations remain.**",
    ]
    REPORT.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    check_release_identity()
    loaded = check_study2_boundaries()
    check_ewma_prefix_invariance(loaded)
    compare_recomputed_study2(loaded)
    compare_recomputed_study1()
    write_outputs()
    failed = [item for item in CHECKS if item.status == "FAIL"]
    print(
        f"Temporal integrity audit: {len(CHECKS) - len(failed)}/"
        f"{len(CHECKS)} checks passed"
    )
    print(REPORT)
    if failed:
        for item in failed:
            print(f"FAIL {item.check_id}: {item.evidence}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()
