from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np


def plot_reconstruction_batch(
    path: str | Path,
    target_loglam,
    target_flux,
    pred_flux,
    loss_mask,
    valid_patch,
    z_true,
    z_pred,
    max_items: int = 4,
) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    bsz = min(max_items, len(target_flux))
    fig, axes = plt.subplots(bsz, 1, figsize=(13, 3.2 * bsz), squeeze=False)
    for i in range(bsz):
        ax = axes[i, 0]
        wave = np.exp(target_loglam[i])
        valid = valid_patch[i].astype(bool)
        masked = loss_mask[i].astype(bool)
        ax.plot(wave[valid], target_flux[i][valid], color="black", linewidth=1.0, label="target")
        ax.plot(wave[valid], pred_flux[i][valid], color="#d62728", linewidth=1.0, alpha=0.9, label="recon")
        if masked.any():
            ax.scatter(wave[masked], target_flux[i][masked], s=8, color="#1f77b4", alpha=0.55, label="masked target")
        ax.set_ylabel("norm flux")
        ax.set_title(f"z true={z_true[i]:.5f}  z pred={z_pred[i]:.5f}")
        ax.grid(alpha=0.2)
        if i == 0:
            ax.legend(loc="best", fontsize=8)
    axes[-1, 0].set_xlabel("wavelength Angstrom")
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)


def plot_redshift_scatter(path: str | Path, y_true: np.ndarray, y_pred: np.ndarray) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    z_true = np.expm1(y_true)
    z_pred = np.expm1(y_pred)
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    axes[0].scatter(z_true, z_pred, s=6, alpha=0.35)
    lim = [float(np.nanmin(z_true)), float(np.nanmax(z_true))]
    axes[0].plot(lim, lim, color="black", linewidth=1)
    axes[0].set_xlabel("z true")
    axes[0].set_ylabel("z pred")
    axes[0].grid(alpha=0.2)
    dz = (z_pred - z_true) / (1 + z_true)
    axes[1].scatter(z_true, dz, s=6, alpha=0.35)
    axes[1].axhline(0, color="black", linewidth=1)
    axes[1].axhline(0.01, color="#d62728", linewidth=1, linestyle="--")
    axes[1].axhline(-0.01, color="#d62728", linewidth=1, linestyle="--")
    axes[1].set_xlabel("z true")
    axes[1].set_ylabel("delta z / (1 + z)")
    axes[1].grid(alpha=0.2)
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)