"""Evaluate reconstruction, spherical embeddings, quantization and low-shot transfer."""

import json
from pathlib import Path

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


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


def balanced_accuracy(target, prediction):
    scores = [(prediction[target == label] == label).mean() for label in np.unique(target)]
    return float(np.mean(scores))


def transfer_metrics(embedding, class_target, regression_target, seed):
    rng = np.random.default_rng(seed)
    features = embedding.transpose(0, 2, 3, 1).reshape(-1, embedding.shape[1])
    classes = class_target.reshape(-1)
    regression = regression_target.reshape(-1)
    train_indices, test_indices = [], []
    for label in np.unique(classes):
        indices = np.flatnonzero(classes == label)
        rng.shuffle(indices)
        split = min(10, max(1, len(indices) // 3))
        train_indices.extend(indices[:split])
        test_indices.extend(indices[split:])
    train_indices, test_indices = np.asarray(train_indices), np.asarray(test_indices)
    x_train, x_test = features[train_indices], features[test_indices]
    y_train, y_test = classes[train_indices], classes[test_indices]
    distances = ((x_test[:, None] - x_train[None]) ** 2).sum(axis=-1)
    transfer = {}
    for k in (1, 3):
        neighbors = np.argpartition(distances, min(k, len(x_train)) - 1, axis=1)[:, :k]
        votes = y_train[neighbors]
        prediction = np.asarray([np.bincount(row).argmax() for row in votes])
        transfer[f"knn_k{k}_balanced_accuracy"] = balanced_accuracy(y_test, prediction)
    labels = np.unique(classes)
    one_hot = np.stack([np.where(y_train == label, 1.0, -1.0) for label in labels], axis=1)
    design = np.column_stack([x_train, np.ones(len(x_train))])
    coefficients = np.linalg.lstsq(design, one_hot, rcond=None)[0]
    class_prediction = labels[np.argmax(np.column_stack([x_test, np.ones(len(x_test))]) @ coefficients, axis=1)]
    transfer["linear_balanced_accuracy"] = balanced_accuracy(y_test, class_prediction)
    regression_coefficients = np.linalg.lstsq(design, regression[train_indices], rcond=None)[0]
    regression_prediction = np.column_stack([x_test, np.ones(len(x_test))]) @ regression_coefficients
    residual = ((regression[test_indices] - regression_prediction) ** 2).sum()
    total = ((regression[test_indices] - regression[test_indices].mean()) ** 2).sum()
    transfer["linear_regression_r2"] = float(1.0 - residual / max(total, 1e-12))
    transfer["train_pixels"] = int(len(train_indices))
    transfer["test_pixels"] = int(len(test_indices))
    return transfer


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    predictions = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
    embedding, restored = predictions["embedding"], predictions["embedding_dequantized"]
    metrics = {
        "samples": int(len(embedding)),
        "mean_embedding_norm": float(np.linalg.norm(embedding, axis=1).mean()),
        "s8_power2_quantization_mae": float(np.abs(embedding - restored).mean()),
        "reconstruction_mae": {},
        "low_shot_transfer": transfer_metrics(
            embedding, predictions["target_nlcd"], predictions["target_sentinel2"][:, 0], config["seed"]
        ),
    }
    for name, spec in config["data"]["target_sources"].items():
        prediction = predictions[f"reconstruction_{name}"]
        target = predictions[f"target_{name}"]
        mask = predictions[f"mask_{name}"]
        if spec["type"] == "categorical":
            metrics["reconstruction_mae"][name] = float(
                (((prediction.argmax(axis=1) != target) * mask[:, 0]).sum()) / max(mask[:, 0].sum(), 1)
            )
        else:
            metrics["reconstruction_mae"][name] = float((np.abs(prediction - target) * mask).sum() / max(mask.sum(), 1))
    evaluation_dir = ROOT / config["paths"]["evaluation_dir"]
    evaluation_dir.mkdir(parents=True, exist_ok=True)
    (evaluation_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    rgb = embedding[0, [1, 16, 9]].transpose(1, 2, 0)
    rgb = np.clip((rgb + 0.3) / 0.6, 0, 1)
    figure, axes = plt.subplots(1, 2, figsize=(8, 4))
    axes[0].imshow(rgb)
    axes[0].set_title("AEF axes A01/A16/A09")
    axes[1].imshow(predictions["target_nlcd"][0], cmap="tab20", vmin=0, vmax=15)
    axes[1].set_title("Synthetic NLCD target")
    for axis in axes:
        axis.axis("off")
    figure.tight_layout()
    figure.savefig(evaluation_dir / "comparison.png", dpi=160)
    plt.close(figure)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()