"""Reload a trained checkpoint; stream evaluation from the supplied MATLAB HDF5."""
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import h5py
import numpy as np
import torch
from common import (ROOT, atomic_json, build_model, load_checkpoint, load_config,
                    predict, relative_metrics, resolved, seed_everything, sha256,
                    validation_indices)


def infer(config, checkpoint_path, output_dir, metrics_only=False):
    seed_everything(0, config["training"]["num_threads"])
    model = build_model(config)
    checkpoint = load_checkpoint(checkpoint_path, model)
    if checkpoint["config"]["model"] != config["model"]:
        raise ValueError("Checkpoint and requested model configuration differ")
    model.eval()
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    source_path = resolved(config, "reference")
    report = {"checkpoint": str(Path(checkpoint_path).resolve()), "checkpoint_sha256": sha256(checkpoint_path),
              "source": str(source_path), "source_sha256": sha256(source_path),
              "seed": checkpoint["seed"], "epoch": checkpoint["epoch"], "selection_score": checkpoint["selection_score"],
              "validation_metrics": checkpoint["validation_metrics"],
              "metric": "raw ||pred-ref||2/||ref||2; no pressure gauge alignment",
              "groups": {}, "reference_regenerated": False, "checkpoint_reloaded": True}
    size = config["inference"]["batch_size"]
    global_num, global_den = np.zeros(3), np.zeros(3)
    output = None if metrics_only else h5py.File(out/"predictions.h5.tmp", "w")
    try:
        with h5py.File(source_path, "r") as source:
            for key in config["inference"]["reference_keys"]:
                data = source[key]
                if data.ndim != 2 or data.shape[0] != 6: raise ValueError(f"Invalid {key} shape: {data.shape}")
                count = data.shape[1]
                numer, denom, hold_num, hold_den = [np.zeros(3) for _ in range(4)]
                if key == config["validation"]["reference_key"]:
                    valids = validation_indices(config, count)
                group = None
                if output is not None:
                    group = output.create_group(key)
                    group.create_dataset("coordinates", (count, 3), dtype="f8", chunks=True)
                    group.create_dataset("prediction", (count, 3), dtype="f4", chunks=True)
                for start in range(0, count, size):
                    end = min(start+size, count)
                    rows = data[:, start:end].T
                    if not np.isfinite(rows).all(): raise ValueError(f"Nonfinite input in {key} at {start}")
                    preds = predict(model, rows[:, :3], size)
                    if not np.isfinite(preds).all(): raise FloatingPointError("Nonfinite predictions")
                    error2 = (preds.astype(np.float64)-rows[:, 3:])**2
                    ref2 = rows[:, 3:]**2
                    numer += error2.sum(0)
                    denom += ref2.sum(0)
                    if key == config["validation"]["reference_key"]:
                        keep = ~np.isin(np.arange(start, end), valids, assume_unique=True)
                        hold_num += error2[keep].sum(0)
                        hold_den += ref2[keep].sum(0)
                    if group is not None:
                        group["coordinates"][start:end] = rows[:, :3]
                        group["prediction"][start:end] = preds
                metrics = relative_metrics(numer, denom)
                report["groups"][key] = {"count": count, "metrics": metrics}
                global_num += numer
                global_den += denom
                if key == config["validation"]["reference_key"]:
                    report["internal_holdout"] = {"count": count-len(valids), "metrics": relative_metrics(hold_num, hold_den)}
                print(f"INFERENCE key={key} count={count} metrics={json.dumps(metrics)}", flush=True)
        report["all_reference"] = {"metrics": relative_metrics(global_num, global_den),
                                    "count": sum(v["count"] for v in report["groups"].values())}
        if output is not None:
            output.attrs["provenance"] = json.dumps({k: v for k, v in report.items() if k not in ["groups", "internal_holdout", "all_reference"]})
            output.close()
            output = None
            (out/"predictions.h5.tmp").replace(out/"predictions.h5")
        atomic_json(out/"metrics.json", report)
    finally:
        if output is not None: output.close()
    print("INFERENCE_COMPLETED", flush=True)
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(ROOT/"config/config.yaml"))
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--seed", type=int)
    parser.add_argument("--kind", choices=["best", "last"], default="best")
    parser.add_argument("--metrics-only", action="store_true")
    args = parser.parse_args()
    cfg = load_config(args.config)
    checkpoint = args.checkpoint or (resolved(cfg, "weights")/f"seed_{args.seed}"/("best_model.pt" if args.kind == "best" else "last_model.pt") if args.seed is not None else resolved(cfg, "weights")/"best_model.pt")
    output = args.output_dir or (resolved(cfg, "results")/f"seed_{args.seed}"/("best_eval" if args.kind == "best" else "final_eval") if args.seed is not None else resolved(cfg, "results")/"best_model")
    infer(cfg, checkpoint, output, args.metrics_only)