"""Evaluate dense PINN-TC predictions and create sparse/target/prediction/error panels."""

import importlib.util
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import yaml


ROOT = Path(__file__).resolve().parents[1]


def load_model_module():
    spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    module = load_model_module()
    archive = np.load(ROOT / config["paths"]["prediction"])
    if str(archive["format_version"]) != config["data"]["format_version"]:
        raise ValueError("prediction format_version mismatch")
    prediction = archive["predictions"]
    expected = (251, 251, 44, 3, 4)
    if prediction.shape != expected or prediction.dtype != np.float32:
        raise ValueError(f"prediction must be float32 with shape {expected}")
    y_axis, x_axis, pressures, times = archive["y"], archive["x"], archive["pressure"], archive["time"]
    sum_speed_error = sum_pred = sum_true = sum_pred2 = sum_true2 = sum_cross = 0.0
    count = 0
    radial_edges = np.linspace(0.0, np.sqrt(2.0) * config["data"]["horizontal_extent_m"], 13)
    radial_sse = np.zeros(len(radial_edges) - 1); radial_count = np.zeros(len(radial_edges) - 1, dtype=np.int64)
    yy, xx = np.meshgrid(y_axis, x_axis, indexing="ij")
    radius = np.sqrt(xx * xx + yy * yy)
    for pressure_index, pressure in enumerate(pressures):
        for time_index, time in enumerate(times):
            coordinates = np.column_stack((yy.ravel(), xx.ravel(), np.full(yy.size, time), np.full(yy.size, pressure))).astype(np.float32)
            target = module.analytic_vortex_numpy(coordinates).reshape(251, 251, 4)
            predicted_speed = np.linalg.norm(prediction[:, :, pressure_index, time_index, :2].astype(np.float64), axis=-1)
            true_speed = np.linalg.norm(target[:, :, :2].astype(np.float64), axis=-1)
            difference = predicted_speed - true_speed
            sum_speed_error += np.square(difference).sum(); count += difference.size
            sum_pred += predicted_speed.sum(); sum_true += true_speed.sum()
            sum_pred2 += np.square(predicted_speed).sum(); sum_true2 += np.square(true_speed).sum()
            sum_cross += (predicted_speed * true_speed).sum()
            bins = np.clip(np.digitize(radius.ravel(), radial_edges) - 1, 0, len(radial_sse) - 1)
            radial_sse += np.bincount(bins, weights=np.square(difference).ravel(), minlength=len(radial_sse))
            radial_count += np.bincount(bins, minlength=len(radial_count))
    covariance = sum_cross - sum_pred * sum_true / count
    variance_pred = sum_pred2 - sum_pred * sum_pred / count
    variance_true = sum_true2 - sum_true * sum_true / count
    pearson = covariance / np.sqrt(max(variance_pred * variance_true, 1e-30))

    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location="cpu", weights_only=False)
    if checkpoint.get("format_version") != config["data"]["format_version"]:
        raise ValueError("checkpoint format_version mismatch")
    model = module.PINNTC(**checkpoint["model_config"]); model.load_state_dict(checkpoint["model"]); model.eval()
    rng = np.random.default_rng(int(config["seed"]) + 99)
    residual_count = int(config["evaluation"]["pde_points"])
    residual_coordinates = np.column_stack((rng.uniform(-400_000, 400_000, (residual_count, 2)),
                                            rng.uniform(-21_600, 21_600, residual_count),
                                            rng.uniform(15_000, 90_000, residual_count))).astype(np.float32)
    tensor = torch.from_numpy(residual_coordinates).requires_grad_(True)
    ru, rv, rc = module.pde_residuals(model, tensor, **config["physics"])
    report = {"dense_shape": list(expected), "dense_points": int(np.prod(expected[:-1])),
              "wind_speed_rmse_m_s": float(np.sqrt(sum_speed_error / count)), "wind_speed_pearson": float(pearson),
              "pde_residual_rms": {"ru_m_s2": float(torch.sqrt(torch.mean(ru.square())).detach()),
                                   "rv_m_s2": float(torch.sqrt(torch.mean(rv.square())).detach()),
                                   "rc_s-1": float(torch.sqrt(torch.mean(rc.square())).detach())},
              "radial_rmse": [{"radius_mid_km": float((left + right) / 2000.0), "rmse_m_s": float(np.sqrt(error / max(samples, 1))), "count": int(samples)}
                              for left, right, error, samples in zip(radial_edges[:-1], radial_edges[1:], radial_sse, radial_count)],
              "scope": "Structured analytic engineering validation, not a paper-performance claim."}
    values = [report["wind_speed_rmse_m_s"], report["wind_speed_pearson"], *report["pde_residual_rms"].values()]
    if not np.isfinite(values).all():
        raise FloatingPointError("evaluation produced a non-finite metric")
    output = ROOT / config["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")

    pressure_index, time_index = int(np.argmin(abs(pressures - 50_000))), int(np.argmin(abs(times)))
    coordinates = np.column_stack((yy.ravel(), xx.ravel(), np.full(yy.size, times[time_index]),
                                   np.full(yy.size, pressures[pressure_index]))).astype(np.float32)
    target_speed = np.linalg.norm(module.analytic_vortex_numpy(coordinates)[:, :2], axis=1).reshape(251, 251)
    predicted_speed = np.linalg.norm(prediction[:, :, pressure_index, time_index, :2], axis=-1)
    sparse = np.load(ROOT / config["data"]["root"] / "training_points.npz")["observation_coordinates"]
    selected = sparse[(abs(sparse[:, 2] - times[time_index]) < 1) & (abs(sparse[:, 3] - pressures[pressure_index]) < 1000)]
    figure, axes_plot = plt.subplots(1, 4, figsize=(16, 4), constrained_layout=True)
    axes_plot[0].scatter(selected[:, 1] / 1000, selected[:, 0] / 1000, s=5, c="black"); axes_plot[0].set_title("Sparse observations")
    panels = ((target_speed, "Target wind speed", "viridis"), (predicted_speed, "Predicted wind speed", "viridis"),
              (predicted_speed - target_speed, "Prediction error", "coolwarm"))
    for axis, (field, title, color_map) in zip(axes_plot[1:], panels):
        image = axis.imshow(field, origin="lower", extent=(-400, 400, -400, 400), cmap=color_map)
        axis.set_title(title); figure.colorbar(image, ax=axis, shrink=0.75)
    for axis in axes_plot:
        axis.set_xlabel("x (km)"); axis.set_ylabel("y (km)")
    figure.savefig(output / "comparison.png", dpi=140); plt.close(figure)
    print(f"evaluation={output.relative_to(ROOT)} wind_rmse={report['wind_speed_rmse_m_s']:.6g} pearson={pearson:.6g}")


if __name__ == "__main__":
    main()