#!/usr/bin/env python3
"""Recalculate saved-chain summaries without rerunning the expensive sampler."""
from pathlib import Path
import argparse
import csv
import gzip
import json
import math
from timing_posterior import summarize


def load_runs(path):
    filename = path if path.exists() else path.with_suffix(path.suffix + ".gz")
    opener = gzip.open if filename.suffix == ".gz" else open
    with opener(filename, "rt") as file:
        rows = list(csv.DictReader(file))
    columns = [key for key in rows[0] if key not in {"ensemble", "draw"}]
    groups = {}
    for row in rows:
        groups.setdefault(int(row["ensemble"]), []).append(
            [float(row[key]) for key in columns]
        )
    return [{"samples": groups[index]} for index in sorted(groups)], columns


def analyze(directory):
    results = {}
    for experiment in ["posterior", "prior-sensitivity", "injection"]:
        base = directory / experiment
        frozen = json.loads((base / "timing-results.json").read_text())
        results[experiment] = {"data_kind": frozen["data_kind"], "models": {}}
        for model in ["slow", "chopping"]:
            runs, columns = load_runs(base / (model + "-samples.csv"))
            summary = summarize(runs)
            for name, values in summary.items():
                for quantile in ["q16", "median", "q84"]:
                    original = frozen["models"][model]["summary"][name][quantile]
                    assert abs(values[quantile] - original) < 1e-7 * max(
                        1, abs(original)
                    ), (experiment, model, name)
            all_samples = [row for run in runs for row in run["samples"]]
            largest_rhat = max(
                value["ensemble_mean_rhat"] for value in summary.values()
            )
            minimum_effective_means = min(
                value["effective_ensemble_means"] for value in summary.values()
            )
            results[experiment]["models"][model] = {
                "summary": summary,
                "maximum_ensemble_mean_rhat": largest_rhat,
                "minimum_effective_ensemble_means": minimum_effective_means,
                "diagnostic_threshold_pass": largest_rhat < 1.05
                and minimum_effective_means > 100,
                "free_vector_norm_above_0_1_fraction": sum(
                    math.hypot(row[2], row[3]) > 0.1 for row in all_samples
                )
                / len(all_samples),
                "middle_mass_above_55_fraction": sum(row[0] > 55 for row in all_samples)
                / len(all_samples),
                "retained_correlated_draws": len(all_samples),
            }
            print(
                experiment,
                model,
                "draws",
                len(all_samples),
                "max ensemble-mean Rhat",
                round(largest_rhat, 4),
                "minimum effective ensemble means",
                round(minimum_effective_means),
                flush=True,
            )
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--directory", type=Path, default=Path("."))
    parser.add_argument("--out", type=Path, default=Path("analysis-summary.json"))
    arguments = parser.parse_args()
    arguments.out.write_text(json.dumps(analyze(arguments.directory), indent=2) + "\n")
    print(
        "PASS: saved quantiles reproduced; diagnostic failures, if any, remain explicit"
    )