from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import matplotlib

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

from common import DEFAULT_CONFIG, load_config, resolve_path


def latitude_weights(height: int) -> np.ndarray:
    latitude = np.linspace(np.pi / 2, -np.pi / 2, height)
    weights = np.cos(latitude).clip(min=0)
    return weights / weights.mean()


def compute_metrics(prediction: np.ndarray, target: np.ndarray) -> dict[str, list[float]]:
    weights = latitude_weights(target.shape[-2])[None, None, None, :, None]
    error = prediction - target
    rmse = np.sqrt(np.mean(error**2 * weights, axis=(0, 1, 3, 4)))
    spatial_weight = weights / (
        weights.sum(axis=(-2, -1), keepdims=True) * target.shape[-1]
    )
    pred_mean = np.sum(prediction * spatial_weight, axis=(-2, -1), keepdims=True)
    target_mean = np.sum(target * spatial_weight, axis=(-2, -1), keepdims=True)
    pred_anomaly = prediction - pred_mean
    target_anomaly = target - target_mean
    numerator = np.sum(pred_anomaly * target_anomaly * weights, axis=(0, 1, 3, 4))
    denominator = np.sqrt(
        np.sum(pred_anomaly**2 * weights, axis=(0, 1, 3, 4))
        * np.sum(target_anomaly**2 * weights, axis=(0, 1, 3, 4))
    )
    acc = numerator / np.maximum(denominator, 1e-12)
    return {"rmse": rmse.tolist(), "acc": acc.tolist()}


def plot_sample(
    prediction: np.ndarray,
    target: np.ndarray,
    variable: str,
    channel_index: int,
    cmap: str,
    output_path: Path,
) -> None:
    predicted = prediction[0, 0, channel_index]
    expected = target[0, 0, channel_index]
    error = predicted - expected
    value_min = min(predicted.min(), expected.min())
    value_max = max(predicted.max(), expected.max())
    error_limit = max(abs(error.min()), abs(error.max()), 1e-12)
    extent = (0, 360, -90, 90)

    figure, axes = plt.subplots(3, 1, figsize=(12, 10), constrained_layout=True)
    image = axes[0].imshow(
        expected, origin="upper", extent=extent, aspect="auto", cmap=cmap,
        vmin=value_min, vmax=value_max,
    )
    axes[0].set_title(f"Target {variable}")
    figure.colorbar(image, ax=axes[0], orientation="vertical")
    image = axes[1].imshow(
        predicted, origin="upper", extent=extent, aspect="auto", cmap=cmap,
        vmin=value_min, vmax=value_max,
    )
    axes[1].set_title(f"Prediction {variable}")
    figure.colorbar(image, ax=axes[1], orientation="vertical")
    image = axes[2].imshow(
        error, origin="upper", extent=extent, aspect="auto", cmap="RdBu_r",
        vmin=-error_limit, vmax=error_limit,
    )
    axes[2].set_title(f"Error {variable}")
    figure.colorbar(image, ax=axes[2], orientation="vertical")
    for axis in axes:
        axis.set_xlabel("Longitude")
        axis.set_ylabel("Latitude")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    figure.savefig(output_path, dpi=160)
    plt.close(figure)


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate and plot FCNv2 output")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--input")
    args = parser.parse_args()

    config = load_config(args.config)
    inference_dir = resolve_path(config, config["inference"]["output_dir"])
    input_path = Path(args.input).expanduser().resolve() if args.input else None
    files = [input_path] if input_path else sorted(inference_dir.glob("sample_*.npz"))
    if not files:
        raise FileNotFoundError(f"No inference outputs found in {inference_dir}")

    predictions = []
    targets = []
    for path in files:
        with np.load(path) as data:
            predictions.append(data["prediction"])
            targets.append(data["target"])
    prediction = np.concatenate(predictions)
    target = np.concatenate(targets)
    metrics = compute_metrics(prediction, target)

    output_dir = resolve_path(config, config["visualization"]["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(
        json.dumps(metrics, indent=2), encoding="utf-8"
    )
    variable = config["visualization"]["variable"]
    channel_index = config["data"]["variables"].index(variable)
    sample_index = config["visualization"]["sample_index"]
    if not 0 <= sample_index < prediction.shape[0]:
        raise IndexError(
            f"visualization.sample_index={sample_index} is outside "
            f"the available range [0, {prediction.shape[0] - 1}]"
        )
    plot_sample(
        prediction[sample_index : sample_index + 1],
        target[sample_index : sample_index + 1],
        variable,
        channel_index,
        config["visualization"]["cmap"],
        output_dir / f"{variable}_forecast.png",
    )
    print(output_dir / "metrics.json")
    print(output_dir / f"{variable}_forecast.png")


if __name__ == "__main__":
    main()