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


def to_numpy_img(t):
    """Convert torch CHW / numpy HWC (or HW) to an array matplotlib can show."""
    if isinstance(t, torch.Tensor):
        t = t.detach().cpu().numpy()
        if t.ndim == 4:  # (B,C,H,W) → first sample
            t = t[0]
        if t.ndim == 3:  # CHW → HWC
            t = np.transpose(t, (1, 2, 0))
    elif isinstance(t, np.ndarray) and t.ndim == 3:
        # Only transpose if channels-first (C,H,W), not already HWC
        if t.shape[0] in (1, 3, 4) and t.shape[-1] not in (1, 3, 4):
            t = np.transpose(t, (1, 2, 0))

    if isinstance(t, np.ndarray) and t.ndim == 3 and t.shape[-1] == 1:
        t = t.squeeze(-1)
    return t


def show_synthesis_result(normal_img, jnd_map, anomaly_img, severity_mask, panel_name: str = ""):

    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    axes[0].imshow(to_numpy_img(normal_img))
    axes[0].set_title("Normal Image")
    axes[1].imshow(to_numpy_img(jnd_map), cmap="viridis")
    axes[1].set_title("JND Map")
    axes[2].imshow(to_numpy_img(anomaly_img))
    axes[2].set_title(f"Synthesized Anomaly (panel={panel_name})")
    axes[3].imshow(to_numpy_img(severity_mask), cmap="jet", vmin=0, vmax=3)
    axes[3].set_title("Severity Mask (0=bg,1=easy,2=med,3=hard)")
    for ax in axes:
        ax.axis("off")
    plt.tight_layout()
    plt.show()


def show_model_output(anomaly_img, rec_img, rec_jnd, seg_map, gt_mask=None):

    n_cols = 5 if gt_mask is not None else 4
    fig, axes = plt.subplots(1, n_cols, figsize=(4 * n_cols, 4))
    axes[0].imshow(to_numpy_img(anomaly_img))
    axes[0].set_title("Input (anomaly)")
    axes[1].imshow(to_numpy_img(rec_img))
    axes[1].set_title("Reconstruction")
    axes[2].imshow(to_numpy_img(rec_jnd), cmap="viridis")
    axes[2].set_title("Reconstructed JND")
    axes[3].imshow(to_numpy_img(seg_map), cmap="jet")
    axes[3].set_title("Predicted anomaly map")
    if gt_mask is not None:
        axes[4].imshow(to_numpy_img(gt_mask), cmap="gray")
        axes[4].set_title("Ground Truth mask")
    for ax in axes:
        ax.axis("off")
    plt.tight_layout()
    plt.show()


def plot_loss_curves(history: dict):

    plt.figure(figsize=(8, 5))
    for name, values in history.items():
        plt.plot(values, label=name)
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.title("Training Loss Curves")
    plt.grid(alpha=0.3)
    plt.show()