"""Evaluate reconstruction and cross-sensor representation consistency."""

import json
from pathlib import Path

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


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


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    predictions = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
    sensors = list(config["data"]["sensors"])
    metrics = {"samples": int(len(predictions["class_target"])), "sensors": {}}
    projected = []
    for name in sensors:
        embedding = predictions[f"embedding_{name}"]
        projection = predictions[f"projected_embedding_{name}"]
        projected.append(projection)
        metrics["sensors"][name] = {
            "reconstruction_mae": float(np.abs(predictions[f"reconstruction_{name}"] - predictions[f"pixels_{name}"]).mean()),
            "embedding_norm": float(np.linalg.norm(embedding, axis=1).mean()),
            "projected_embedding_norm": float(np.linalg.norm(projection, axis=1).mean()),
            "representation_loss": float(predictions[f"representation_loss_{name}"]),
        }
    similarities = []
    for left in range(len(projected)):
        for right in range(left + 1, len(projected)):
            similarities.append((projected[left] * projected[right]).sum(axis=1))
    metrics["mean_cross_sensor_cosine_similarity"] = float(np.concatenate(similarities).mean())
    output_dir = ROOT / config["paths"]["evaluation_dir"]
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")

    figure, axes = plt.subplots(len(sensors), 3, figsize=(9, 3 * len(sensors)), squeeze=False)
    for row, name in enumerate(sensors):
        rgb = config["data"]["sensors"][name]["rgb_indices"]
        source = predictions[f"pixels_{name}"][0, rgb].transpose(1, 2, 0)
        reconstruction = predictions[f"reconstruction_{name}"][0, rgb].transpose(1, 2, 0)
        for image in (source, reconstruction):
            image -= image.min()
            image /= image.max() + 1e-6
        axes[row, 0].imshow(source)
        axes[row, 0].set_title(f"{name} input")
        axes[row, 1].imshow(reconstruction)
        axes[row, 1].set_title("MAE reconstruction")
        axes[row, 2].imshow(np.abs(source - reconstruction).mean(axis=2), cmap="magma")
        axes[row, 2].set_title("absolute error")
        for axis in axes[row]:
            axis.axis("off")
    figure.tight_layout()
    figure.savefig(output_dir / "comparison.png", dpi=150)
    plt.close(figure)
    print(f"metrics={output_dir.relative_to(ROOT) / 'metrics.json'}")


if __name__ == "__main__":
    main()