"""Render heatmaps (concept P by layer x position) and the layer profile."""

import os, sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch

RES = Path(__file__).resolve().parent.parent / os.environ.get("JLENS_RESULTS", "results")


def plot_heat(path):
    d = torch.load(path)
    m, words, toks = d["map"].detach(), d["words"], d["tokens"]
    fig, ax = plt.subplots(figsize=(min(16, 0.28 * m.shape[1] + 2), 6))
    im = ax.imshow(m.numpy(), aspect="auto", origin="lower", cmap="magma",
                   vmin=0)
    ax.set_ylabel("layer")
    ax.set_xlabel("position")
    if toks and len(toks) <= 60:
        ax.set_xticks(range(len(toks)))
        ax.set_xticklabels([t.replace("\n", "\\n") for t in toks],
                           rotation=90, fontsize=6)
    ax.set_title(f"J-lens P({'/'.join(words)})")
    fig.colorbar(im)
    out = path.with_suffix(".png")
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    plt.close(fig)
    print("wrote", out)


def plot_profile(path):
    d = torch.load(path)
    fig, ax1 = plt.subplots(figsize=(8, 4.5))
    L = len(d["entropy"])
    ax1.plot(range(L), d["entropy"], "k-", label="lens entropy")
    ax1.set_xlabel("layer")
    ax1.set_ylabel("entropy (nats)")
    ax2 = ax1.twinx()
    ax2.plot(range(L), d["agree_in"], "b--", label="top tok == current input")
    ax2.plot(range(L), d["agree_next"], "r--", label="top tok == next output")
    ax2.set_ylabel("agreement")
    fig.legend(loc="upper center", ncol=3, fontsize=8)
    fig.tight_layout()
    out = path.with_suffix(".png")
    fig.savefig(out, dpi=130)
    plt.close(fig)
    print("wrote", out)


if __name__ == "__main__":
    for p in sorted(RES.glob("heat_*.pt")):
        plot_heat(p)
    lp = RES / "layer_profile.pt"
    if lp.exists():
        plot_profile(lp)