"""Rebuild averages, per-design scores, details and explanations from Summary JSON only."""
from __future__ import annotations

import argparse
import html
import os
import sys
from collections import Counter, defaultdict
from pathlib import Path

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from common import ROOT, dump, read, stage_dir, model_size_key
from Postprocess.main import METRICS
from Summary.model_checks import weak_model_tables, weak_result


def averages(records):
    return {key: sum(values) / len(values) if (values := [r["metrics"][key] for r in records if r["metrics"].get(key) is not None]) else None
            for key in METRICS}


def rebuild(directory):
    directory = Path(directory)
    provenance = read(directory / "data_provenance.json")
    records = [read(path) for path in sorted((directory / "results").rglob("*.json"))]
    grouped, designs = defaultdict(list), defaultdict(list)
    for record in records:
        grouped[(record["model_id"], record["protocol"])].append(record)
        designs[(record["model_id"], record["protocol"], record["design_id"])].append(record)
    avg = []
    for (model, protocol), group in sorted(grouped.items()):
        scored = [r for r in group if r["score_status"] == "scored"]
        avg.append({"model_id": model, "protocol": protocol, "macro_average_scores": averages(scored),
                    "scored_datasets": len(scored), "total_result_records": len(group),
                    "status_counts": dict(Counter(r["status"] for r in group)),
                    "aggregation": "unweighted mean across scored dataset records; null metrics excluded",
                    "evaluation": group[0]["evaluation"]})
    individual = [{"model_id": model, "protocol": protocol, "design_id": did,
                   "scores": averages([r for r in group if r["score_status"] == "scored"]),
                   "datasets": [{k: r.get(k) for k in ("dataset_id", "fold_id", "status", "score_status", "metrics", "detail")} for r in group]}
                  for (model, protocol, did), group in sorted(designs.items())]
    explanations = {"evaluation_protocols": {
        "leave_one_design_out": "All records sharing a design-file SHA256 are held out together. Fit and preprocessing statistics use training records only.",
        "per_design": "Fit each design's reviewed records and evaluate those same records. This is training fit only.",
        "train_all": "Fit all reviewed records and evaluate those same records. This is training fit only.",
        "diagnostic": "Untrained actual-input forward passes. No accuracy is calculated."},
        "data_selection": provenance.get('data', {}).get('selection', 'CAD + Excel + PLY; CSV required for NG and optional for explicit OK.'),
        "targets": "Externally reviewed common-CAD node risk; raw CSV pixel coordinates are not automatically mapped to CAD.",
        "scores": "MAE/RMSE/Brier use reviewed probability targets; AP/AUC and confusion scores threshold reviewed targets. Null means unavailable/undefined.",
        "averaging": "Macro means are per dataset, not per node. Unscored records never enter the average.",
        "uncertainty": "GP posterior standard deviations are available when fitted. Other models have no calibrated uncertainty claim.",
        "deformation": "Physics GNN deformation/type heads are not supervised by this node-risk contract and are not scored.",
        "limitations": "Six distinct designs are the current holdout groups. All complete records are NG; OK without CSV is excluded. No pixel-to-CAD calibration or validated deformation target was supplied.",
        "statuses": dict(Counter(r["status"] for r in records)),
        "failures": [{k: r.get(k) for k in ("model_id", "protocol", "design_id", "dataset_id", "status", "reason")}
                     for r in records if r["status"] not in ("completed", "untrained_diagnostic")]}
    if provenance['run'].get('config', {}).get('allow_weak_labels'):
        explanations['targets'] = 'Automatic scan_projected_weak CAD-node targets from CSV pixels, measured height rasters and proper rigid scan-to-CAD fits. Unannotated covered nodes outside defect buffers are assumed negative.'
        explanations['scores'] = 'Independently reviewed accuracy is unavailable for weak supervision. Real optimizer epoch losses and checkpoints are saved; weak-label predictions never enter reviewed score averages.'
        explanations['evaluation_protocols']['train_all'] = 'Experimental fit on eligible automatic scan-projected weak labels. Predictions are training-fit diagnostics, not independently reviewed accuracy.'
        explanations['limitations'] = 'Raster calibration and registration are inferred and unreviewed. CSV source conflicts and poor-overlap cases are recorded. Datasets with failed projection gates are excluded from weak training. All source records are NG; unannotated regions are not reviewed healthy.'
        if provenance['run']['config'].get('user_selected_fit_datasets'):
            explanations['targets'] = 'User-selected scan/CAD pairs with recorded mirrored frames and CAD-supported scan regions. Node targets remain automatic weak CSV/raster projections; unannotated covered nodes outside guarded regions are assumed negative.'
            explanations['limitations'] = 'User confirmation applies to pair inclusion. CAD-node ground truth remains unreviewed. Rejected CSV rows and their nearby assumed negatives are masked; source choices, reflection, raw residuals and partial scan coverage remain recorded.'
    weak_groups = defaultdict(list)
    for record in records:
        if weak_result(record) and not record.get('training_excluded'): weak_groups[(record['model_id'], record['protocol'])].append(weak_result(record))
    weak_checks = []
    for (model, protocol), group in sorted(weak_groups.items()):
        entry = {'model_id': model, 'protocol': protocol, 'datasets': len(group), 'independently_reviewed': False}
        if provenance['run'].get('config', {}).get('size_named_outputs'):
            entry['input_size'] = model_size_key(provenance['run']['config'], model)
        for metric in ('bce', 'average_precision', 'iou'):
            values = [r[metric] for r in group if r.get(metric) is not None]
            entry[metric] = sum(values)/len(values) if values else None
        weak_checks.append(entry)
    dump(directory / "avg_scores.json", avg)
    dump(directory / "individual_scores.json", individual)
    dump(directory / "detailed_scores.json", records)
    dump(directory / "explanations.json", explanations)
    summary = {"run": provenance["run"], "dataset_count": provenance["data"]["dataset_count"],
               "design_count": provenance["data"]["design_count"], "result_record_count": len(records),
               "scored_record_count": sum(r["score_status"] == "scored" for r in records),
               "average_scores": avg, "individual_scores": individual, "explanations": explanations,
               "self_contained_files": ["data_provenance.json", "avg_scores.json", "individual_scores.json",
                                        "detailed_scores.json", "explanations.json", "results/**/*.json"]}
    from Plots.main import training_diagnosis
    summary["training_diagnosis"] = training_diagnosis(records, provenance.get("prepared"))
    summary['weak_model_checks'] = weak_checks
    summary['excluded_train_dataset_ids'] = provenance['run'].get('config', {}).get('excluded_train_dataset_ids', [])
    explanations['training_exclusions'] = {'dataset_ids': summary['excluded_train_dataset_ids'],
        'policy': 'Excluded from every model fit and model-check cohort; retained for inspection and separately marked held-out predictions.'}
    if provenance['run'].get('config', {}).get('include_ok_weak_labels'):
        explanations['targets'] += ' Explicit OK production records provide weak healthy CAD-node targets without requiring a defect CSV.'
        explanations['limitations'] += ' Production OK status is a weak part-level healthy assumption; it is not independently reviewed CAD-node ground truth.'
    if provenance['run'].get('config', {}).get('evaluation_reference_prepared'):
        explanations['averaging'] += ' Size comparisons use frozen reference CAD-node positions, targets and masks, shared across every size.'
    dump(directory / 'explanations.json', explanations)
    dump(directory / "summary.json", summary)
    # The HTML is derived entirely from the JSON snapshots in this directory.
    rows = []
    for record in avg:
        metrics = record["macro_average_scores"]
        cells = [record["model_id"], record["protocol"], str(record["scored_datasets"])]
        cells += ["—" if metrics[key] is None else f"{metrics[key]:.5f}" for key in ("mae", "rmse", "average_precision", "iou", "dice")]
        rows.append("<tr>" + "".join(f"<td>{html.escape(value)}</td>" for value in cells) + "</tr>")
    run = summary["run"]
    plots_link = os.path.relpath(ROOT / "Plots" / "outputs" / run["date"] / run["run_id"] / "index.html", directory)
    viewer_link = os.path.relpath(ROOT / "Visualizer" / "index.html", directory) + f'?run={run["date"]}/{run["run_id"]}'
    weak_section = ''
    if weak_checks:
        weak_section = ('<h2>Model checks against weak labels</h2><p>Training fit and held-out-design agreement are separate; '
                        'these are not independently reviewed accuracy.</p>' + weak_model_tables(weak_checks))
    content = f'''<!doctype html><html lang="en"><meta charset="utf-8"><title>Moltex v2 summary</title>
<style>body{{font:15px system-ui;margin:40px;max-width:1400px}}table{{border-collapse:collapse}}td,th{{padding:8px;border:1px solid #ddd}}pre{{white-space:pre-wrap}}</style>
<h1>Moltex v2 — {summary['run']['date']} / {summary['run']['run_id']}</h1>
<p>{summary['dataset_count']} datasets · {summary['design_count']} designs · {summary['scored_record_count']} scored records</p>
<p><strong>Training status: {summary['training_diagnosis']['training_status']}</strong> · {summary['training_diagnosis']['trained_prediction_records']} trained prediction records · {summary['training_diagnosis']['folds_with_epoch_losses']} folds with epoch losses</p>
<p><a href="{html.escape(plots_link)}">PNG results, losses &amp; training diagnosis</a> · <a href="{html.escape(viewer_link)}">CAD / scan / label / model-input visualizer</a></p>
<p>— = no valid score. Diagnostic and training-fit outputs are separate from unseen-design evaluation.</p>
{weak_section}
<h2>Average scores</h2><table><tr><th>Model</th><th>Protocol</th><th>Scored datasets</th><th>MAE</th><th>RMSE</th><th>AP</th><th>IoU</th><th>Dice</th></tr>{''.join(rows)}</table>
<h2>Individual and detailed scores</h2><p><a href="individual_scores.json">Individual scores</a> · <a href="detailed_scores.json">Detailed predictions/targets</a> · <a href="data_provenance.json">Data and code provenance</a></p>
<h2>Explanations</h2><pre>{html.escape(str(explanations))}</pre></html>'''
    (directory / "index.html").write_text(content, encoding="utf-8")
    return summary


def collect(run, manifest, prepared, index):
    directory = stage_dir("Summary", run)
    provenance = {"run": run, "data": manifest, "prepared": prepared,
                  "code": read(Path(run["manifest_dir"]) / "code_provenance.json"),
                  "dependencies": read(Path(run["manifest_dir"]) / "dependencies.json")}
    dump(directory / "data_provenance.json", provenance)
    for source in index["metric_files"]:
        record = read(source)
        path = (directory / "results" / record["model_id"] / record["protocol"] / record["fold_id"]
                / record["design_id"] / f"{record['dataset_id']}.json")
        dump(path, record)
    summary = rebuild(directory)
    report_dir = ROOT / "reports" / run["date"] / run["run_id"]
    report_dir.mkdir(parents=True, exist_ok=True)
    content = (directory / "index.html").read_text(encoding="utf-8")
    for stage in ("Plots", "Visualizer"):
        content = content.replace(os.path.relpath(ROOT / stage, directory) + "/",
                                  os.path.relpath(ROOT / stage, report_dir) + "/")
    (report_dir / "results.html").write_text(content, encoding="utf-8")
    dump(report_dir / "summary.json", summary)
    # Colocated snapshots make the report's JSON links resolve without NAS access.
    for name in ("individual_scores.json", "detailed_scores.json", "data_provenance.json"):
        dump(report_dir / name, read(directory / name))
    return summary, directory


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--summary-dir", type=Path, required=True)
    args = parser.parse_args()
    summary = rebuild(args.summary_dir)
    print(f"Rebuilt {summary['result_record_count']} records from Summary JSON only: {args.summary_dir}")


if __name__ == "__main__": main()