#!/usr/bin/env python3
"""Reload best checkpoint; evaluate exact DNS and DNS-interpolated 200x200 grids."""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from scipy.interpolate import RectBivariateSpline

from scripts.common import (ROOT, gather, load_checkpoint, load_dns, predict, read_config,
                            relative_l2, resolve, save_json, setup, tensor)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(ROOT / "config/config.yaml"))
    parser.add_argument("--checkpoint")
    parser.add_argument("--device")
    args = parser.parse_args()
    cfg = read_config(args.config)
    device, dtype = setup(cfg, args.device)
    model, checkpoint = load_checkpoint(cfg, device, args.checkpoint)
    dns = load_dns(cfg)
    if checkpoint["data_sha256"] != dns["sha256"]:
        raise ValueError("Checkpoint was trained on different DNS bytes")
    out = resolve(cfg["outputs"]["results"])
    sampling = np.load(out / "sampling_indices.npz")
    xy, times = dns["xy"], dns["t"]
    nx, ny = cfg["evaluation"]["grid_size"]
    gx = np.linspace(xy[:, 0].min(), xy[:, 0].max(), nx)
    gy = np.linspace(xy[:, 1].min(), xy[:, 1].max(), ny)
    xx, yy = np.meshgrid(gx, gy)
    grid_xy = np.column_stack((xx.ravel(), yy.ravel()))
    native_x, native_y = np.unique(xy[:, 0]), np.unique(xy[:, 1])
    ix, iy = np.searchsorted(native_x, xy[:, 0]), np.searchsorted(native_y, xy[:, 1])
    dns_grid_coincidence = np.any(np.isclose(gx[:, None], native_x[None, :], atol=1e-10, rtol=0), 1)[None, :] & np.any(np.isclose(gy[:, None], native_y[None, :], atol=1e-10, rtol=0), 1)[:, None]
    grid_unseen = ~dns_grid_coincidence.ravel()
    native_pred = np.empty((len(times), len(xy), 3), dtype=np.float64)
    per_time, grid_sse, grid_denom = [], np.zeros(2), np.zeros(2)
    unseen_sse, unseen_denom = np.zeros(2), np.zeros(2)
    held_sse, held_denom = np.zeros(2), np.zeros(2)
    snapshot = int(np.argmin(abs(times - cfg["evaluation"]["snapshot_time"])))
    batch_size = cfg["evaluation"]["batch_size"]
    native_heldout_count = 0
    for k, t in enumerate(times):
        q = np.column_stack((xy, np.full(len(xy), t)))
        pred = predict(model, q, device, dtype, batch_size)
        native_pred[k] = pred
        errors = relative_l2(pred[:, :2], dns["uv"][k])
        unused = np.ones(len(xy), dtype=bool)
        unused[sampling["boundary"][k]] = False
        unused[sampling["residual"][k]] = False
        native_heldout_count += int(unused.sum())
        held_sse += np.sum((pred[unused, :2] - dns["uv"][k, unused])**2, 0)
        held_denom += np.sum(dns["uv"][k, unused]**2, 0)
        gq = np.column_stack((grid_xy, np.full(len(grid_xy), t)))
        gp = predict(model, gq, device, dtype, batch_size)[:, :2]
        reference = np.empty((ny, nx, 2), dtype=np.float64)
        for channel in range(2):
            field = np.empty((len(native_y), len(native_x)))
            field[iy, ix] = dns["uv"][k, :, channel]
            reference[:, :, channel] = RectBivariateSpline(native_y, native_x, field, kx=3, ky=3)(gy, gx)
        ref = reference.reshape(-1, 2)
        ge = relative_l2(gp, ref)
        grid_sse += np.sum((gp - ref)**2, 0)
        grid_denom += np.sum(ref**2, 0)
        unseen_sse += np.sum((gp[grid_unseen] - ref[grid_unseen])**2, 0)
        unseen_denom += np.sum(ref[grid_unseen]**2, 0)
        row = {"time_index": k, "time": float(t), "dns_u": errors[0], "dns_v": errors[1], "grid_u": ge[0], "grid_v": ge[1]}
        per_time.append(row)
        if k == snapshot:
            np.savez_compressed(out / "snapshot_grid.npz", x=gx, y=gy, reference=reference,
                                prediction=gp.reshape(ny, nx, 2), time=t,
                                unseen_mask=grid_unseen.reshape(ny, nx))
        if k % 10 == 0 or k == len(times) - 1:
            print(f"INFERENCE time_index={k:3d}/199 t={t:.2f} DNS_relL2(u,v)=({errors[0]:.8e},{errors[1]:.8e}) grid_relL2(u,v)=({ge[0]:.8e},{ge[1]:.8e})", flush=True)
    np.savez_compressed(out / "dns_predictions.npz", X_star=xy, t=times, uvp_prediction=native_pred,
                        uv_reference=dns["uv"], p_reference=dns["p"])
    with (out / "per_time_metrics.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(per_time[0]))
        writer.writeheader()
        writer.writerows(per_time)
    primary = relative_l2(native_pred[:, :, :2].reshape(-1, 2), dns["uv"].reshape(-1, 2))
    snap_error = np.array([per_time[snapshot]["dns_u"], per_time[snapshot]["dns_v"]])
    grid_aggregate = np.sqrt(grid_sse / grid_denom)
    grid_snapshot = np.array([per_time[snapshot]["grid_u"], per_time[snapshot]["grid_v"]])
    table = cfg["evaluation"]["table3"]
    table_values = np.array([table["tsa_u"], table["tsa_v"]])
    nearest = int(np.argmin(np.sum((xy - np.array(cfg["evaluation"]["probe"]))**2, 1)))
    probe_q = np.column_stack((np.tile(xy[nearest], (len(times), 1)), times))
    probe_prediction = predict(model, probe_q, device, dtype, batch_size)
    np.savez_compressed(out / "probe_series.npz", xy=xy[nearest], requested_xy=cfg["evaluation"]["probe"],
                        t=times, reference=dns["uv"][:, nearest], prediction=probe_prediction[:, :2])
    diagnostics_q, _ = gather(dns, sampling["validation"])
    diagnostic_count = min(len(diagnostics_q), 1000)
    diagnostics_q = diagnostics_q[np.linspace(0, len(diagnostics_q) - 1, diagnostic_count, dtype=int)]
    momentum_sse, div_sse = np.zeros(2), 0.0
    ad_max_abs = {"flow": 0.0, "momentum": 0.0}
    # Use the independent reverse-mode implementation for physical diagnostics.
    for start in range(0, diagnostic_count, 250):
        q = tensor(diagnostics_q[start:start + 250], device, dtype)
        flow_values, residual, div = model.residual_autograd(q)
        with torch.no_grad():
            taylor_values, taylor_residual, _ = model.residual_taylor(q)
        torch.testing.assert_close(taylor_values, flow_values, atol=1e-8, rtol=1e-7)
        torch.testing.assert_close(taylor_residual, residual, atol=1e-8, rtol=1e-7)
        ad_max_abs["flow"] = max(ad_max_abs["flow"], float((taylor_values-flow_values).abs().max().detach()))
        ad_max_abs["momentum"] = max(ad_max_abs["momentum"], float((taylor_residual-residual).abs().max().detach()))
        momentum_sse += (residual.detach().cpu().numpy()**2).sum(0)
        div_sse += float((div.detach()**2).sum())
    def pair(array):
        return {"u": float(array[0]), "v": float(array[1])}
    metrics = {"checkpoint_completed_updates": checkpoint["completed_updates"], "data_path": dns["path"],
               "data_sha256": dns["sha256"], "metric": "||prediction-reference||_2 / ||reference||_2 (dimensionless ratio)",
               "dns_all_spacetime": pair(primary), "dns_all_spacetime_percent": pair(primary * 100),
               "dns_unused_spacetime": pair(np.sqrt(held_sse / held_denom)), "dns_unused_count": native_heldout_count,
               "dns_snapshot": {"time": float(times[snapshot]), **pair(snap_error)},
               "grid200_all_spacetime": pair(grid_aggregate), "grid200_snapshot": {"time": float(times[snapshot]), **pair(grid_snapshot)},
               "grid200_strictly_unseen_spacetime": pair(np.sqrt(unseen_sse / unseen_denom)),
               "grid200_unseen_spatial_count": int(grid_unseen.sum()),
               "grid_reference": "Tensor-product cubic spline interpolation of this exact DNS; not additional DNS ground truth",
               "table3_tsa_pinn": pair(table_values), "table3_pinn": {"u": table["pinn_u"], "v": table["pinn_v"]},
               "snapshot_over_table3_ratio": pair(snap_error / table_values),
               "grid_snapshot_over_table3_ratio": pair(grid_snapshot / table_values),
               "all_spacetime_over_table3_ratio": pair(primary / table_values),
               "probe": {"requested": cfg["evaluation"]["probe"], "actual_dns_point": xy[nearest].tolist(),
                         "relative_l2": pair(relative_l2(probe_prediction[:, :2], dns["uv"][:, nearest]))},
               "physics_diagnostics": {"count": diagnostic_count, "sampling": "uniformly spaced across held-out locations over all 200 frames", "momentum_rmse": pair(np.sqrt(momentum_sse / diagnostic_count)),
                                       "divergence_rmse": float(np.sqrt(div_sse / diagnostic_count)),
                                       "trained_checkpoint_ad_equivalence": {"status": "passed", "max_absolute_difference": ad_max_abs}},
               "comparison_limits": ["Table3 exact seed/checkpoint/evaluation aggregation not specified in paper", "Paper Eq42 includes x100, but Table3 numerical entries are reported without a percent marker; both ratio and percent provided", "Native DNS covers t=0..19.9 at dt=0.1; no synthetic temporal resampling", "PyTorch port of paper TensorFlow method; coordinate AD checked against reverse-mode PyTorch"]}
    save_json(out / "metrics.json", metrics)
    print("FINAL METRICS " + json.dumps(metrics, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()