import numpy as np
import matplotlib.pyplot as plt
from allytools.logger import get_logger

log = get_logger(__name__)

def compare_scalar2d(
    s1,
    s2,
    *,
    normalize: str = "max",
    title1: str = "Field 1",
    title2: str = "Field 2",
):
    """
    Compare two Scalar2d fields on the same metric grid.

    This function performs:
      - optional normalization ("max", "sum", "none")
      - RMSE, max-abs-error, correlation, peak metrics
      - 2D comparison plots (field1, field2, difference)
      - 1D central cuts

    Returns
    -------
    dict with:
        rmse, max_abs_err, corr, peak1, peak2
    """

    log.debug("Starting Scalar2d comparison...")
    log.debug("Field1 shape=%s, extent=%s", s1.data.shape, s1.extent)
    log.debug("Field2 shape=%s, extent=%s", s2.data.shape, s2.extent)

    # -------------------------
    # Shape check
    # -------------------------
    if s1.data.shape != s2.data.shape:
        raise ValueError(
            f"Scalar2d shapes differ: {s1.data.shape} vs {s2.data.shape}. "
            "Both fields must have identical sampling."
        )

    ny, nx = s1.data.shape

    # -------------------------
    # Extent check
    # -------------------------
    def _resolve_extent(s):
        if s.extent is not None:
            return tuple(map(float, s.extent))
        return (0.0, float(nx), 0.0, float(ny))

    ext1 = _resolve_extent(s1)
    ext2 = _resolve_extent(s2)

    if not np.allclose(ext1, ext2):
        raise ValueError(
            f"Coordinate extents differ: {ext1} vs {ext2}. "
            "Comparison requires identical metric grids."
        )

    log.debug("Metric extent confirmed: %s", ext1)

    x_min, x_max, y_min, y_max = ext1
    x_mm = np.linspace(x_min, x_max, nx)
    y_mm = np.linspace(y_min, y_max, ny)

    # -------------------------
    # Copy and normalize
    # -------------------------
    p1 = s1.data.astype(float).copy()
    p2 = s2.data.astype(float).copy()

    log.debug("Applying normalization: %s", normalize)

    if normalize == "max":
        m1, m2 = p1.max(), p2.max()
        p1 /= m1 if m1 != 0 else 1.0
        p2 /= m2 if m2 != 0 else 1.0
        log.debug("Normalized by max: max1=%.6g, max2=%.6g", m1, m2)

    elif normalize == "sum":
        s_1, s_2 = p1.sum(), p2.sum()
        p1 /= s_1 if s_1 != 0 else 1.0
        p2 /= s_2 if s_2 != 0 else 1.0
        log.debug("Normalized by sum: sum1=%.6g, sum2=%.6g", s_1, s_2)

    elif normalize == "none":
        log.debug("No normalization applied.")

    else:
        raise ValueError(f"Unknown normalize='{normalize}'. Use 'max', 'sum', or 'none'.")

    # -------------------------
    # Metrics
    # -------------------------
    diff = p1 - p2
    mse = np.mean(diff**2)
    rmse = float(np.sqrt(mse))
    max_abs_err = float(np.max(np.abs(diff)))
    corr = float(np.corrcoef(p1.ravel(), p2.ravel())[0, 1])

    metrics = {
        "rmse": rmse,
        "max_abs_err": max_abs_err,
        "corr": corr,
        "peak1": float(p1.max()),
        "peak2": float(p2.max()),
    }

    log.info("Scalar2d Comparison Metrics:")
    log.info("  RMSE          = %.6g", rmse)
    log.info("  Max |Δ|       = %.6g", max_abs_err)
    log.info("  Correlation   = %.6g", corr)
    log.info("  Peak1, Peak2 = %.6g, %.6g",
             metrics["peak1"], metrics["peak2"])

    # -------------------------
    # 2D visualization
    # -------------------------
    log.debug("Rendering 2D comparison plots...")
    extent = ext1

    fig, axes = plt.subplots(1, 3, figsize=(14, 4), sharex=True, sharey=True)

    im0 = axes[0].imshow(
        p1, extent=extent, origin="lower",
        interpolation="nearest", cmap="jet"
    )
    axes[0].set_title(title1)
    plt.colorbar(im0, ax=axes[0], fraction=0.046)

    im1 = axes[1].imshow(
        p2, extent=extent, origin="lower",
        interpolation="nearest", cmap="jet"
    )
    axes[1].set_title(title2)
    plt.colorbar(im1, ax=axes[1], fraction=0.046)

    im2 = axes[2].imshow(
        diff, extent=extent, origin="lower",
        interpolation="nearest", cmap="bwr"
    )
    axes[2].set_title(f"Difference ({title1} - {title2})")
    plt.colorbar(im2, ax=axes[2], fraction=0.046)

    for ax in axes:
        ax.set_xlabel("x [mm]")
        ax.set_ylabel("y [mm]")

    plt.tight_layout()
    plt.show()
    log.debug("2D comparison plots rendered.")

    # -------------------------
    # 1D cross-sections
    # -------------------------
    log.debug("Rendering 1D center cuts...")
    H, W = p1.shape
    cy, cx = H // 2, W // 2

    fig2, axes2 = plt.subplots(1, 2, figsize=(10, 4))

    # Horizontal slice
    axes2[0].plot(x_mm, p1[cy, :], label=title1)
    axes2[0].plot(x_mm, p2[cy, :], "--", label=title2)
    axes2[0].set_title("X-cut (center line)")
    axes2[0].set_xlabel("x [mm]")
    axes2[0].set_ylabel("Value (normalized)")
    axes2[0].legend()
    axes2[0].grid(True)

    # Vertical slice
    axes2[1].plot(y_mm, p1[:, cx], label=title1)
    axes2[1].plot(y_mm, p2[:, cx], "--", label=title2)
    axes2[1].set_title("Y-cut (center line)")
    axes2[1].set_xlabel("y [mm]")
    axes2[1].set_ylabel("Value (normalized)")
    axes2[1].legend()
    axes2[1].grid(True)

    plt.tight_layout()
    plt.show()
    log.debug("1D center cuts rendered.")

    log.debug("Scalar2d comparison completed.")
    return metrics
