#!/usr/bin/env python3
"""Evaluate forecasts and render meteorological diagnostic figures."""

from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Mapping, Sequence, Tuple

import h5py
import numpy as np
import yaml

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.colors import Normalize, TwoSlopeNorm
except ImportError as exc:  # pragma: no cover
    raise RuntimeError(
        "scripts/result.py requires matplotlib for scientific figures"
    ) from exc

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from scripts.data_loader import read_metadata, resolve_data_dir


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
    parser.add_argument("--data-dir")
    parser.add_argument("--output-dir", default=str(PROJECT_ROOT / "result/output"))
    parser.add_argument("--plot-count", type=int, default=3, help="Initialization times to visualize")
    parser.add_argument("--plot-leads", type=int, default=3, help="Lead times per initialization to visualize")
    parser.add_argument("--dpi", type=int, default=160)
    return parser.parse_args()


def _pretty_name(name: str) -> str:
    replacements = {
        "geopotential_": "Geopotential ",
        "temperature_": "Temperature ",
        "specific_humidity_": "Specific humidity ",
        "2m_temperature": "2 m temperature",
        "total_precipitation": "Total precipitation",
        "mean_sea_level_pressure": "Mean sea-level pressure",
        "10m_u_component_of_wind": "10 m U wind",
        "10m_v_component_of_wind": "10 m V wind",
        "100m_u_component_of_wind": "100 m U wind",
        "100m_v_component_of_wind": "100 m V wind",
        "u_component_of_wind_": "U wind ",
        "v_component_of_wind_": "V wind ",
        "outgoing_longwave_radiation": "Outgoing longwave radiation",
        "sea_surface_temperature": "Sea-surface temperature",
        "significant_wave_height": "Significant wave height",
        "mean_wave_direction": "Mean wave direction",
        "mean_wave_period": "Mean wave period",
        "soil_moisture_0_7cm": "Soil moisture 0-7 cm",
    }
    for prefix, replacement in replacements.items():
        if name.startswith(prefix):
            suffix = name[len(prefix) :]
            return f"{replacement}{suffix}" if suffix.isdigit() else replacement
    return name.replace("_", " ").title()


def _safe_percentile(values: np.ndarray, percentile: float, default: float = 1.0) -> float:
    finite = np.asarray(values, dtype=np.float64)
    finite = finite[np.isfinite(finite)]
    if finite.size == 0:
        return float(default)
    result = float(np.percentile(finite, percentile))
    return result if np.isfinite(result) else float(default)


def _absolute_limits(values: np.ndarray) -> Tuple[float, float]:
    lo, hi = _safe_percentile(values, 2.0, -1.0), _safe_percentile(values, 98.0, 1.0)
    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
        center = float(np.nanmean(values)) if np.isfinite(values).any() else 0.0
        span = max(abs(center) * 0.05, 1.0)
        return center - span, center + span
    return lo, hi


def _error_limit(values: np.ndarray) -> float:
    limit = _safe_percentile(np.abs(values), 98.0, 1.0)
    if not np.isfinite(limit) or limit <= 1.0e-8:
        limit = max(_safe_percentile(np.abs(values), 100.0, 1.0), 1.0)
    return float(limit)


def _field(field: np.ndarray, lookup: Mapping[str, int], name: str) -> np.ndarray | None:
    index = lookup.get(name)
    return None if index is None else np.asarray(field[index], dtype=np.float32)


def _diagnostics(channels: Sequence[str], prediction: np.ndarray, truth: np.ndarray) -> List[Dict[str, Any]]:
    """Create fields with meteorological meaning from a channel tensor."""
    lookup = {name: index for index, name in enumerate(channels)}
    diagnostics: List[Dict[str, Any]] = []

    temperature_name = "2m_temperature" if "2m_temperature" in lookup else "temperature_500"
    pred, true = _field(prediction, lookup, temperature_name), _field(truth, lookup, temperature_name)
    if pred is not None and true is not None:
        diagnostics.append({"key": "temperature", "title": _pretty_name(temperature_name), "unit": "data units", "prediction": pred, "truth": true, "cmap": "coolwarm"})

    pred, true = _field(prediction, lookup, "geopotential_500"), _field(truth, lookup, "geopotential_500")
    if pred is not None and true is not None:
        diagnostics.append({"key": "geopotential", "title": "500 hPa geopotential", "unit": "data units", "prediction": pred, "truth": true, "cmap": "viridis"})

    wind_pairs = (
        ("10m_u_component_of_wind", "10m_v_component_of_wind", "10 m wind speed"),
        ("u_component_of_wind_850", "v_component_of_wind_850", "850 hPa wind speed"),
    )
    for u_name, v_name, title in wind_pairs:
        up, vp = _field(prediction, lookup, u_name), _field(prediction, lookup, v_name)
        ut, vt = _field(truth, lookup, u_name), _field(truth, lookup, v_name)
        if all(value is not None for value in (up, vp, ut, vt)):
            diagnostics.append({"key": "wind_speed", "title": title, "unit": "data units", "prediction": np.hypot(up, vp), "truth": np.hypot(ut, vt), "cmap": "magma", "prediction_uv": (up, vp), "truth_uv": (ut, vt)})
            break

    pred, true = _field(prediction, lookup, "total_precipitation"), _field(truth, lookup, "total_precipitation")
    if pred is not None and true is not None:
        diagnostics.append({"key": "precipitation", "title": "Total precipitation", "unit": "data units", "prediction": pred, "truth": true, "cmap": "YlGnBu"})

    if len(diagnostics) < 2:
        used = {item["title"] for item in diagnostics}
        for index, name in enumerate(channels):
            if _pretty_name(name) in used:
                continue
            diagnostics.append({"key": f"channel_{index}", "title": _pretty_name(name), "unit": "data units", "prediction": prediction[index], "truth": truth[index], "cmap": "coolwarm"})
            if len(diagnostics) >= 3:
                break
    return diagnostics


def _draw_grid(
    ax: Any,
    field: np.ndarray,
    title: str,
    cmap: str,
    unit: str,
    error: bool = False,
    overlay: Tuple[np.ndarray, np.ndarray] | None = None,
    limits: Tuple[float, float] | None = None,
) -> None:
    """Draw a regular global latitude/longitude panel without cartopy."""
    height, width = field.shape
    longitude = np.linspace(0.0, 360.0, width, endpoint=False)
    latitude = np.linspace(-90.0, 90.0, height)
    if error:
        limit = _error_limit(field) if limits is None else max(abs(float(limits[0])), abs(float(limits[1])))
        norm = TwoSlopeNorm(vmin=-limit, vcenter=0.0, vmax=limit)
        image = ax.imshow(field, extent=(0, 360, -90, 90), origin="lower", aspect="auto", cmap="RdBu_r", norm=norm, interpolation="nearest")
    else:
        lo, hi = _absolute_limits(field) if limits is None else limits
        image = ax.imshow(field, extent=(0, 360, -90, 90), origin="lower", aspect="auto", cmap=cmap, norm=Normalize(vmin=lo, vmax=hi), interpolation="nearest")
    ax.set_title(title, fontsize=10, pad=5)
    ax.set_xlim(0, 360)
    ax.set_ylim(-90, 90)
    ax.set_xticks((0, 60, 120, 180, 240, 300, 360))
    ax.set_yticks((-60, -30, 0, 30, 60))
    ax.set_xlabel("Longitude (deg)", fontsize=8)
    ax.set_ylabel("Latitude (deg)", fontsize=8)
    ax.tick_params(labelsize=8)
    ax.grid(color="white", alpha=0.35, linewidth=0.45)
    ax.axhline(0, color="black", alpha=0.28, linewidth=0.55)
    ax.axvline(180, color="black", alpha=0.28, linewidth=0.55)
    if overlay is not None and not error:
        u, v = overlay
        step_y, step_x = max(1, height // 12), max(1, width // 24)
        speed = np.hypot(u, v)
        scale = max(_safe_percentile(speed, 95, 1.0) * 12, 1.0)
        quiver = ax.quiver(longitude[::step_x], latitude[::step_y], u[::step_y, ::step_x], v[::step_y, ::step_x], color="black", alpha=0.68, scale=scale, width=0.0022, headwidth=3.5, minlength=0.1)
        ax.quiverkey(quiver, 0.86, -0.16, _safe_percentile(speed, 75, 1.0), "75th pct", labelpos="E", coordinates="axes", fontproperties={"size": 7})
    colorbar = ax.figure.colorbar(image, ax=ax, fraction=0.045, pad=0.025)
    colorbar.ax.tick_params(labelsize=7)
    colorbar.set_label(("Error (" + unit + ")") if error else unit, fontsize=8)


def _plot_overview(path: Path, source: str, record: Mapping[str, Any], channels: Sequence[str], prediction: np.ndarray, truth: np.ndarray, dpi: int) -> None:
    diagnostics = _diagnostics(channels, prediction, truth)
    lead = int(record["lead"])
    time_step = int(record.get("time_step", 6))
    lead_hours = (lead + 1) * time_step
    fig, axes = plt.subplots(len(diagnostics), 3, figsize=(16.5, max(4.0 * len(diagnostics), 8.0)), squeeze=False, facecolor="#f5f7fa")
    for row, diagnostic in enumerate(diagnostics):
        pred, true = diagnostic["prediction"], diagnostic["truth"]
        error = pred - true
        limits = _absolute_limits(np.concatenate((pred.reshape(-1), true.reshape(-1))))
        _draw_grid(axes[row, 0], pred, f"{diagnostic['title']} | forecast", diagnostic["cmap"], diagnostic["unit"], overlay=diagnostic.get("prediction_uv"), limits=limits)
        _draw_grid(axes[row, 1], true, f"{diagnostic['title']} | verifying field", diagnostic["cmap"], diagnostic["unit"], overlay=diagnostic.get("truth_uv"), limits=limits)
        _draw_grid(axes[row, 2], error, f"{diagnostic['title']} | forecast error", diagnostic["cmap"], diagnostic["unit"], error=True)
        axes[row, 2].text(0.02, 0.03, f"RMSE {np.sqrt(np.nanmean(error * error)):.3g}   bias {np.nanmean(error):+.3g}", transform=axes[row, 2].transAxes, fontsize=8, color="black", bbox={"facecolor": "white", "alpha": 0.82, "edgecolor": "none", "pad": 2.5})
    valid = str(record["valid_timestamp"])
    analysis_time = datetime.strptime(source, "%Y%m%d%H") - timedelta(hours=time_step)
    analysis = analysis_time.strftime("%Y%m%d%H")
    fig.suptitle(f"FengWu-W2S global forecast diagnostics\nInit {analysis[:8]} {analysis[8:]} UTC  |  Valid {valid[:8]} {valid[8:]} UTC  |  F{lead_hours:03d}", fontsize=15, fontweight="bold", y=0.995)
    fig.subplots_adjust(top=0.93, left=0.04, right=0.98, bottom=0.045, hspace=0.56, wspace=0.20)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
    plt.close(fig)


def _plot_lead_summary(path: Path, lead_metrics: Sequence[Mapping[str, Any]], dpi: int) -> None:
    if not lead_metrics:
        return
    leads = np.asarray([item["lead_hours"] for item in lead_metrics], dtype=float)
    nrmse = np.asarray([item["mean_normalized_rmse"] for item in lead_metrics], dtype=float)
    acc = np.asarray([item["mean_acc"] for item in lead_metrics], dtype=float)
    fig, axes = plt.subplots(1, 2, figsize=(14, 4.8), facecolor="#f5f7fa")
    for ax in axes:
        ax.set_facecolor("white")
        ax.grid(alpha=0.25, linewidth=0.7)
        ax.set_xlabel("Forecast lead (hours)")
    axes[0].plot(leads, nrmse, color="#b43f3f", marker="o", linewidth=2.2, markersize=5)
    axes[0].fill_between(leads, 0, nrmse, color="#e8a1a1", alpha=0.28)
    axes[0].set_title("Normalized error growth", fontweight="bold")
    axes[0].set_ylabel("Mean RMSE / channel standard deviation")
    axes[0].set_ylim(bottom=0)
    axes[1].plot(leads, acc, color="#20639b", marker="o", linewidth=2.2, markersize=5)
    axes[1].axhline(0, color="#555", linewidth=0.8)
    axes[1].set_title("Anomaly correlation by lead", fontweight="bold")
    axes[1].set_ylabel("Mean ACC")
    axes[1].set_ylim(-1, 1)
    fig.suptitle("FengWu-W2S forecast skill across available leads", fontsize=14, fontweight="bold")
    fig.tight_layout(rect=(0, 0, 1, 0.93))
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
    plt.close(fig)


def _plot_channel_summary(path: Path, channels: Sequence[str], rmse: np.ndarray, normalized_rmse: np.ndarray, acc: np.ndarray, dpi: int) -> None:
    count = min(12, len(channels))
    ranking = np.argsort(normalized_rmse)[::-1][:count]
    labels = [_pretty_name(channels[index]) for index in ranking][::-1]
    values, acc_values = normalized_rmse[ranking][::-1], acc[ranking][::-1]
    fig, axes = plt.subplots(1, 2, figsize=(15, max(5.5, count * 0.46)), facecolor="#f5f7fa")
    for ax in axes:
        ax.set_facecolor("white")
        ax.grid(axis="x", alpha=0.25)
    axes[0].barh(labels, values, color="#d46a6a", alpha=0.9)
    axes[0].set_title("Largest normalized RMSE", fontweight="bold")
    axes[0].set_xlabel("RMSE / channel standard deviation")
    axes[1].barh(labels, acc_values, color="#4b88b8", alpha=0.9)
    axes[1].set_title("ACC for the same channels", fontweight="bold")
    axes[1].set_xlabel("Anomaly correlation coefficient")
    axes[1].set_xlim(-1, 1)
    fig.suptitle("Channel-level forecast diagnostics", fontsize=14, fontweight="bold")
    fig.tight_layout(rect=(0, 0, 1, 0.94))
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
    plt.close(fig)


def _plot_training_history(path: Path, train: np.ndarray, valid: np.ndarray, dpi: int) -> None:
    if train.size == 0 and valid.size == 0:
        return
    fig, ax = plt.subplots(figsize=(9.5, 4.8), facecolor="#f5f7fa")
    ax.set_facecolor("white")
    if train.size:
        ax.plot(np.arange(1, train.size + 1), train, marker="o", color="#2166ac", linewidth=2, label="Training")
    if valid.size:
        ax.plot(np.arange(1, valid.size + 1), valid, marker="o", color="#b2182b", linewidth=2, label="Validation")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Reported non-negative probability loss")
    ax.set_title("FengWu-W2S optimization history", fontweight="bold")
    ax.grid(alpha=0.28)
    ax.legend(frameon=False)
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
    plt.close(fig)


def _truth_for_timestamp(data_dir: Path, timestamp: str, indices: np.ndarray, time_step: int) -> np.ndarray:
    dt = datetime.strptime(timestamp, "%Y%m%d%H")
    index = int((dt - datetime(dt.year, 1, 1)).total_seconds() // (3600 * time_step))
    path = data_dir / "data" / f"{dt.year}.h5"
    with h5py.File(path, "r") as source:
        fields = source["fields"]
        if index < 0 or index >= fields.shape[0]:
            raise IndexError(f"Timestamp {timestamp} maps to index {index}, outside {path}")
        return np.asarray(fields[index, indices, :, :], dtype=np.float32)


def _load_prediction(output_dir: Path, record: Mapping[str, Any]) -> np.ndarray:
    path = Path(record["path"])
    if not path.is_absolute():
        path = output_dir / path
    if not path.exists():
        raise FileNotFoundError(f"Forecast field listed in index.json does not exist: {path}")
    field = np.load(path).astype(np.float32)
    if field.ndim != 3:
        raise ValueError(f"Forecast field must have [channel, lat, lon] shape, got {field.shape}")
    return field


def _selected_records(records: Sequence[Mapping[str, Any]], plot_count: int, plot_leads: int) -> List[Mapping[str, Any]]:
    grouped: Dict[str, List[Mapping[str, Any]]] = defaultdict(list)
    for record in records:
        grouped[str(record["source_timestamp"])].append(record)
    selected: List[Mapping[str, Any]] = []
    for source in sorted(grouped)[: max(0, int(plot_count))]:
        source_records = sorted(grouped[source], key=lambda item: int(item["lead"]))
        number = max(1, min(int(plot_leads), len(source_records)))
        positions = np.linspace(0, len(source_records) - 1, number, dtype=int)
        selected.extend(source_records[int(position)] for position in np.unique(positions))
    return selected


def main() -> None:
    args = parse_args()
    with Path(args.config).open(encoding="utf-8") as source:
        config = yaml.safe_load(source)
    data_cfg = config["data"]
    data_dir = resolve_data_dir(args.data_dir or data_cfg["data_dir"], PROJECT_ROOT)
    output_dir = Path(args.output_dir).expanduser()
    if not output_dir.is_absolute():
        output_dir = (PROJECT_ROOT / output_dir).resolve()
    index_path = output_dir / "index.json"
    if not index_path.exists():
        raise FileNotFoundError(f"Forecast index not found: {index_path}; run scripts/inference.py first")
    records = json.loads(index_path.read_text(encoding="utf-8"))
    if not isinstance(records, list):
        raise ValueError("Forecast index must contain a list")
    channels = list(data_cfg["channels"])
    metadata = read_metadata(data_dir, channels)
    indices, time_step = metadata["indices"], int(metadata["time_step"])
    means = np.asarray(metadata["means"], dtype=np.float64).reshape(-1)
    stds = np.asarray(metadata["stds"], dtype=np.float64).reshape(-1)
    sum_squared = np.zeros(len(channels), dtype=np.float64)
    sum_pred_anom_sq = np.zeros(len(channels), dtype=np.float64)
    sum_truth_anom_sq = np.zeros(len(channels), dtype=np.float64)
    sum_cross = np.zeros(len(channels), dtype=np.float64)
    lead_sum_squared = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
    lead_sum_pred_anom_sq = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
    lead_sum_truth_anom_sq = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
    lead_sum_cross = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
    lead_counts: Dict[int, int] = defaultdict(int)
    count = 0
    for record in records:
        prediction = _load_prediction(output_dir, record)
        truth = _truth_for_timestamp(data_dir, str(record["valid_timestamp"]), indices, time_step)
        if prediction.shape != truth.shape:
            raise ValueError(f"Prediction/truth shape mismatch: {prediction.shape} vs {truth.shape}")
        if not np.isfinite(prediction).all() or not np.isfinite(truth).all():
            raise ValueError(f"Non-finite forecast or truth field for {record}")
        error = prediction - truth
        pred_anom, truth_anom = prediction - means[:, None, None], truth - means[:, None, None]
        pixels = int(error.shape[1] * error.shape[2])
        sum_squared += np.sum(error * error, axis=(1, 2))
        sum_pred_anom_sq += np.sum(pred_anom * pred_anom, axis=(1, 2))
        sum_truth_anom_sq += np.sum(truth_anom * truth_anom, axis=(1, 2))
        sum_cross += np.sum(pred_anom * truth_anom, axis=(1, 2))
        lead = int(record["lead"])
        lead_sum_squared[lead] += np.sum(error * error, axis=(1, 2))
        lead_sum_pred_anom_sq[lead] += np.sum(pred_anom * pred_anom, axis=(1, 2))
        lead_sum_truth_anom_sq[lead] += np.sum(truth_anom * truth_anom, axis=(1, 2))
        lead_sum_cross[lead] += np.sum(pred_anom * truth_anom, axis=(1, 2))
        lead_counts[lead] += pixels
        count += pixels
    rmse = np.sqrt(sum_squared / max(1, count))
    acc = sum_cross / (np.sqrt(sum_pred_anom_sq * sum_truth_anom_sq) + 1.0e-8)
    normalized_rmse = rmse / np.maximum(stds, 1.0e-8)
    lead_metrics: List[Dict[str, Any]] = []
    pixels_per_field = int(metadata["shape"][-2] * metadata["shape"][-1])
    for lead in sorted(lead_counts):
        denominator = max(1, lead_counts[lead])
        lead_rmse = np.sqrt(lead_sum_squared[lead] / denominator)
        lead_acc = lead_sum_cross[lead] / (np.sqrt(lead_sum_pred_anom_sq[lead] * lead_sum_truth_anom_sq[lead]) + 1.0e-8)
        lead_metrics.append({"lead": lead, "lead_hours": (lead + 1) * time_step, "mean_rmse": float(np.mean(lead_rmse)), "mean_normalized_rmse": float(np.mean(lead_rmse / np.maximum(stds, 1.0e-8))), "mean_acc": float(np.mean(lead_acc)), "records": int(lead_counts[lead] // max(1, pixels_per_field))})
    result_root, plots_dir = output_dir.parent, output_dir.parent / "plots"
    plots_dir.mkdir(parents=True, exist_ok=True)
    for old_plot in plots_dir.glob("*.png"):
        old_plot.unlink()
    metrics = {"count": len(records), "channels": channels, "rmse": rmse.tolist(), "normalized_rmse": normalized_rmse.tolist(), "acc": acc.tolist(), "mean_rmse": float(np.mean(rmse)), "mean_normalized_rmse": float(np.mean(normalized_rmse)), "mean_acc": float(np.mean(acc)), "lead_metrics": lead_metrics, "visualization": {"domain": "meteorology", "coordinate_convention": "regular grid, longitude 0-360 degrees, latitude -90 to 90 degrees", "figures": ["plots/*_overview.png", "plots/lead_skill.png", "plots/channel_skill.png", "plots/training_history.png"]}}
    (result_root / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    np.save(result_root / "rmse.npy", rmse.astype(np.float32))
    np.save(result_root / "acc.npy", acc.astype(np.float32))
    np.save(result_root / "normalized_rmse.npy", normalized_rmse.astype(np.float32))
    for record in _selected_records(records, args.plot_count, args.plot_leads):
        prediction = _load_prediction(output_dir, record)
        truth = _truth_for_timestamp(data_dir, str(record["valid_timestamp"]), indices, time_step)
        source, lead = str(record["source_timestamp"]), int(record["lead"])
        _plot_overview(plots_dir / f"{source}_lead{lead:03d}_overview.png", source, {**record, "time_step": time_step}, channels, prediction, truth, max(80, int(args.dpi)))
    _plot_lead_summary(plots_dir / "lead_skill.png", lead_metrics, max(80, int(args.dpi)))
    _plot_channel_summary(plots_dir / "channel_skill.png", channels, rmse, normalized_rmse, acc, max(80, int(args.dpi)))
    checkpoint_dir = PROJECT_ROOT / "data" / "checkpoints"
    train_path, valid_path = checkpoint_dir / "trloss.npy", checkpoint_dir / "valoss.npy"
    if train_path.exists() and valid_path.exists():
        _plot_training_history(plots_dir / "training_history.png", np.load(train_path), np.load(valid_path), max(80, int(args.dpi)))
    print(json.dumps(metrics, indent=2))
    print(f"Saved scientific diagnostics under {plots_dir}")


if __name__ == "__main__":
    main()