"""Evaluate Scale-MAE reconstruction, kNN transfer and GSD sensitivity."""

import json
import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import yaml

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


def knn_predict(train_features, train_labels, test_features, k=3):
    train = train_features / np.maximum(np.linalg.norm(train_features, axis=1, keepdims=True), 1e-8)
    test = test_features / np.maximum(np.linalg.norm(test_features, axis=1, keepdims=True), 1e-8)
    nearest = np.argsort(-(test @ train.T), axis=1)[:, :min(k, len(train))]
    return np.asarray([np.bincount(train_labels[row]).argmax() for row in nearest])


def rgb(image):
    return np.clip(image[:3].transpose(1, 2, 0), 0, 1)


def main():
    parser = argparse.ArgumentParser(description="Evaluate Scale-MAE outputs")
    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
    args = parser.parse_args()
    cfg = yaml.safe_load(Path(args.config).read_text())
    source = ROOT / cfg["paths"]["inference_dir"] / "reconstruction.npz"
    if not source.exists():
        raise FileNotFoundError("Run inference before evaluation")
    data = np.load(source)
    for key in ("prediction", "target", "test_features", "train_features", "labels", "gsd"):
        if key not in data: raise ValueError(f"inference archive missing {key}")
    if not np.isfinite(data["prediction"]).all(): raise FloatingPointError("non-finite inference output")
    prediction = knn_predict(data["train_features"], data["train_labels"], data["test_features"])
    accuracy = float(np.mean(prediction == data["labels"]))
    mse = np.mean((data["prediction"] - data["target"]) ** 2, axis=(1, 2, 3))
    low_mse = float(np.mean((data["low_prediction"] - data["low_target"]) ** 2))
    high_mse = float(np.mean((data["high_prediction"] - data["high_target"]) ** 2))
    gsd_values = sorted(np.unique(data["gsd"]).tolist())
    gsd_mse = {str(value): float(np.mean(mse[data["gsd"] == value])) for value in gsd_values}
    gsd_accuracy = {str(value): float(np.mean(prediction[data["gsd"] == value] == data["labels"][data["gsd"] == value]))
                    for value in gsd_values}
    result = {"reconstruction_mse": float(np.mean(mse)), "low_frequency_mse": low_mse,
              "high_frequency_mse": high_mse, "knn_accuracy": accuracy,
              "gsd_reconstruction_mse": gsd_mse, "gsd_knn_accuracy": gsd_accuracy,
              "data_source": "synthetic", "protocol": cfg["data"]["protocol"]}
    output = ROOT / cfg["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")

    figure, axes = plt.subplots(2, 3, figsize=(9, 6))
    images = ((data["target"][0], "Original"), (data["low_target"][0], "Low target"),
              (data["high_target"][0] + 0.5, "High target"), (data["prediction"][0], "Reconstruction"),
              (data["low_prediction"][0], "Low prediction"), (data["high_prediction"][0] + 0.5, "High prediction"))
    for axis, (image, title) in zip(axes.flat, images):
        axis.imshow(rgb(image)); axis.set_title(title); axis.axis("off")
    figure.tight_layout(); figure.savefig(output / "bandpass_reconstruction.png", dpi=160); plt.close(figure)

    figure, axis = plt.subplots(figsize=(6, 3.5))
    axis.plot(gsd_values, [gsd_mse[str(x)] for x in gsd_values], marker="o", color="#1f77b4")
    axis.set(xlabel="GSD (m/pixel)", ylabel="Reconstruction MSE", title="Scale-aware Reconstruction")
    axis.grid(alpha=0.25)
    figure.tight_layout(); figure.savefig(output / "gsd_reconstruction_error.png", dpi=160); plt.close(figure)

    figure, axis = plt.subplots(figsize=(6, 3.5))
    axis.plot(gsd_values, [gsd_accuracy[str(x)] for x in gsd_values], marker="s", color="#d62728")
    axis.set(xlabel="GSD (m/pixel)", ylabel="kNN accuracy", title="Scale-aware Feature Transfer")
    axis.set_ylim(0, 1.05); axis.grid(alpha=0.25)
    figure.tight_layout(); figure.savefig(output / "gsd_knn_accuracy.png", dpi=160); plt.close(figure)

    figure, axis = plt.subplots(figsize=(5, 3))
    axis.bar(["Low frequency", "High frequency"], [low_mse, high_mse], color=["#2a9d8f", "#e76f51"])
    axis.set(ylabel="MSE", title="Bandpass Reconstruction Error")
    figure.tight_layout(); figure.savefig(output / "frequency_error.png", dpi=160); plt.close(figure)

    np.save(output / "features.npy", data["test_features"])
    print(json.dumps(result, indent=2)); print("evaluation=", output)


if __name__ == "__main__":
    main()