from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_RESULTS_DIR = PROJECT_ROOT / "outputs" / "inference"


def _load_array(path: Path, name: str) -> np.ndarray:
    value = np.load(path)[name]
    if not np.all(np.isfinite(value)):
        raise ValueError(f"{path.name}:{name} contains NaN or Inf values.")
    return value


def validate_sample(path: Path, expected_channels: int = 20) -> dict[str, Any]:
    with np.load(path) as result:
        required = {"input", "reconstruction", "error", "mask", "step_idx", "time_index"}
        missing = required - set(result.files)
        if missing:
            raise ValueError(f"{path} is missing fields: {sorted(missing)}")
        inputs = _load_array(path, "input")
        reconstruction = _load_array(path, "reconstruction")
        error = _load_array(path, "error")
        mask = _load_array(path, "mask")

    expected_shape = (expected_channels, 720, 1440)
    if inputs.shape != expected_shape:
        raise ValueError(f"input shape must be {expected_shape}, got {inputs.shape}.")
    if reconstruction.shape != inputs.shape or error.shape != inputs.shape:
        raise ValueError("reconstruction and error must have the same shape as input.")
    if mask.shape != (90, 180):
        raise ValueError(f"mask must be a patch grid with shape (90, 180), got {mask.shape}.")
    if not np.array_equal(mask, mask.astype(bool)):
        raise ValueError("mask must contain only binary values.")
    if not np.allclose(error, reconstruction - inputs):
        raise ValueError("error does not equal reconstruction - input.")

    spatial_mask = np.repeat(np.repeat(mask.astype(bool), 8, axis=0), 8, axis=1)
    visible_error = np.abs(error[:, ~spatial_mask])
    masked_error = np.abs(error[:, spatial_mask])
    mask_fraction = float(mask.mean())
    return {
        "file": str(path),
        "shape": list(inputs.shape),
        "mask_fraction": mask_fraction,
        "masked_pixels": int(spatial_mask.sum()),
        "visible_mae": np.mean(visible_error, axis=1).tolist() if visible_error.size else [0.0] * expected_channels,
        "masked_mae": np.mean(masked_error, axis=1).tolist() if masked_error.size else [0.0] * expected_channels,
        "mae": np.mean(np.abs(error), axis=(1, 2)).tolist(),
        "rmse": np.sqrt(np.mean(error**2, axis=(1, 2))).tolist(),
    }


def save_diagnostic_png(path: Path, output_path: Path, channel: int) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    with np.load(path) as result:
        input_field = result["input"][channel]
        reconstruction = result["reconstruction"][channel]
        error = result["error"][channel]
        mask = result["mask"].astype(bool)
        time_index = result["time_index"].item()

    input_low, input_high = np.nanpercentile(input_field, [1, 99])
    reconstruction_low, reconstruction_high = np.nanpercentile(reconstruction, [1, 99])
    field_low = min(input_low, reconstruction_low)
    field_high = max(input_high, reconstruction_high)
    if not np.isfinite(field_low) or not np.isfinite(field_high) or field_high <= field_low:
        field_low, field_high = float(np.nanmin(input_field)), float(np.nanmax(input_field))
    error_limit = float(np.nanpercentile(np.abs(error), 99))
    if not np.isfinite(error_limit) or error_limit <= 0:
        error_limit = max(float(np.nanmax(np.abs(error))), 1.0)

    latitude = np.linspace(90, -90, input_field.shape[0])
    longitude = np.linspace(0, 360, input_field.shape[1], endpoint=False)
    extent = [longitude[0], longitude[-1], latitude[-1], latitude[0]]
    masked_pixels = np.repeat(np.repeat(mask, 8, axis=0), 8, axis=1)
    masked_overlay = np.ma.masked_where(~masked_pixels, masked_pixels)

    figure, axes = plt.subplots(2, 3, figsize=(18, 8), constrained_layout=True)
    field_kwargs = {"cmap": "viridis", "vmin": field_low, "vmax": field_high, "extent": extent, "aspect": "auto"}
    error_kwargs = {"cmap": "RdBu_r", "vmin": -error_limit, "vmax": error_limit, "extent": extent, "aspect": "auto"}
    panels = [
        (axes[0, 0], input_field, "Input (normalized)", field_kwargs),
        (axes[0, 1], reconstruction, "Reconstruction (normalized)", field_kwargs),
        (axes[0, 2], error, "Signed error", error_kwargs),
        (axes[1, 0], np.abs(error), "Absolute error", {**error_kwargs, "cmap": "magma", "vmin": 0}),
    ]
    for axis, field, title, kwargs in panels:
        image = axis.imshow(field, **kwargs)
        axis.set_title(title)
        axis.set_xlabel("Longitude (degrees east)")
        axis.set_ylabel("Latitude (degrees)")
        axis.contour(masked_pixels, levels=[0.5], colors="white", linewidths=0.35, extent=extent)
        figure.colorbar(image, ax=axis, shrink=0.8)

    mask_image = axes[1, 1].imshow(mask.astype(float), cmap="gray_r", vmin=0, vmax=1, extent=extent, aspect="auto")
    axes[1, 1].set_title("Patch mask (white = masked)")
    axes[1, 1].set_xlabel("Longitude (degrees east)")
    axes[1, 1].set_ylabel("Latitude (degrees)")
    figure.colorbar(mask_image, ax=axes[1, 1], ticks=[0, 1], shrink=0.8)

    axes[1, 2].hist(error.ravel(), bins=80, color="#315f8c", alpha=0.85)
    axes[1, 2].axvline(0, color="black", linewidth=0.8)
    axes[1, 2].set_title(f"Error distribution\nMAE={np.mean(np.abs(error)):.4f}, RMSE={np.sqrt(np.mean(error**2)):.4f}")
    axes[1, 2].set_xlabel("Signed error (normalized)")
    axes[1, 2].set_ylabel("Pixel count")

    figure.suptitle(f"W-MAE reconstruction diagnostic | channel={channel} | time={time_index}")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    figure.savefig(output_path, dpi=150)
    plt.close(figure)


def main() -> None:
    parser = argparse.ArgumentParser(description="Validate W-MAE inference outputs and make diagnostics.")
    parser.add_argument(
        "results_dir",
        type=Path,
        nargs="?",
        default=DEFAULT_RESULTS_DIR,
        help="Inference output directory (default: outputs/inference).",
    )
    parser.add_argument("--expected-channels", type=int, default=20)
    parser.add_argument("--mask-ratio", type=float, default=None)
    parser.add_argument(
        "--channel",
        type=int,
        default=0,
        help="Channel for the diagnostic image (default: 0). Use --channel -1 to disable it.",
    )
    parser.add_argument("--diagnostic-dir", type=Path, default=None)
    parser.add_argument(
        "--visualization-samples",
        type=int,
        default=1,
        help="Number of samples to visualize; validation still covers all samples (default: 1).",
    )
    args = parser.parse_args()
    if args.visualization_samples <= 0:
        raise ValueError("--visualization-samples must be positive.")
    files = sorted(args.results_dir.glob("sample_*.npz"))
    if not files:
        raise FileNotFoundError(f"No sample_*.npz files found in {args.results_dir}.")
    reports = [validate_sample(path, args.expected_channels) for path in files]
    if args.mask_ratio is not None:
        for report in reports:
            if not np.isclose(report["mask_fraction"], args.mask_ratio):
                raise ValueError(
                    f"{report['file']} mask fraction {report['mask_fraction']} "
                    f"does not match expected {args.mask_ratio}."
                )
    if args.channel is not None and args.channel >= 0:
        if not 0 <= args.channel < args.expected_channels:
            raise ValueError("--channel is outside the configured channel range.")
        diagnostic_dir = args.diagnostic_dir or args.results_dir / "diagnostics"
        for path in files[:args.visualization_samples]:
            save_diagnostic_png(path, diagnostic_dir / f"{path.stem}_channel_{args.channel:02d}.png", args.channel)
    summary = {
        "samples": len(reports),
        "mean_mask_fraction": float(np.mean([report["mask_fraction"] for report in reports])),
        "mean_mae": np.mean([report["mae"] for report in reports], axis=0).tolist(),
        "mean_rmse": np.mean([report["rmse"] for report in reports], axis=0).tolist(),
        "reports": reports,
    }
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()