# -*- coding: utf-8 -*-
"""Mapas de calor de extremos de corriente de celda (MicrotrackingFer).

Para cada uno de los N logs mas nuevos de logs_nuevos/ genera una figura con:

  Panel A  Heatmap X-Y construido EXCLUSIVAMENTE con los puntos cuyo valor de
           corriente de celda C supera el percentil 80 de la sesion (solo dia).
           El color es la Isc en mA. Las celdas sin punto extremo quedan vacias,
           de modo que el mapa muestra unicamente la region de maxima corriente.
           Encima se dibuja el campo de irradiancia G_POA medido por el
           piranometro (isolineas) y su propia region > percentil 80, para
           correlacionar espacialmente corriente extrema vs. intensidad.

  Panel B  Correlacion C (Isc) vs G_POA de ese mismo subconjunto extremo,
           con regresion, Pearson r, Spearman rho y la linea de corte P80.

Mapeo adc_manager.cpp: C=canal1(fotocorriente), V=canal3(Voc), P=canal2(piranometro,
Apogee SP-212 on-board, AIN2). Calibracion ADS122C04 24-bit, Vref 2.048 V
(identica a metrica1.py y parse_log.py). El canal P se registra en complemento a 2
de 24 bits sin signo: se extiende el signo y los valores negativos (noche) se apodan a 0.

Uso:  py -3.12 heatmap_extremos_p80.py [n_logs]   (por defecto n_logs = 5)
"""

import glob
import os
import re
import sys

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize

ROOT = os.path.dirname(os.path.abspath(__file__))
LOGS_DIR = os.path.join(ROOT, "logs_nuevos")
OUT_DIR = os.path.join(ROOT, "figures", "heatmap_extremos_p80")

ADC_MAX = 8388607.0
GAIN_SHUNT = 20.0
R_SHUNT = 0.075
DEG_PER_USTEP = 80.0 / 53000.0

PCT = 80.0
DAY_G_MIN = 20.0
PAD_PCT = 0.5

RE = re.compile(
    r"LOG:\[(?P<t>[^\]]+)\]\s*"
    r"X:(?P<x>-?\d+),Y:(?P<y>-?\d+),C:(?P<c>-?\d+),V:(?P<v>-?\d+),"
    r"P:(?P<p>-?\d+),S:(?P<s>-?\d+),er:(?P<er>-?\d+)"
)


def sign24(raw):
    raw &= 0xFFFFFF
    return raw - 0x1000000 if raw & 0x800000 else raw


def isc_ma(c_raw):
    vsh_mV = c_raw / ADC_MAX * 2048.0 / GAIN_SHUNT
    return (vsh_mV / 1000.0) / R_SHUNT * 1000.0


def load_log(path):
    ts, xs, ys, cs, ps = [], [], [], [], []
    with open(path, "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            m = RE.search(line)
            if m:
                ts.append(m.group("t"))
                xs.append(int(m.group("x")))
                ys.append(int(m.group("y")))
                cs.append(int(m.group("c")))
                ps.append(sign24(int(m.group("p"))))
    return (np.array(ts), np.array(xs, dtype=float), np.array(ys, dtype=float),
            np.array(cs, dtype=float), np.array(ps, dtype=float))


def bin_stat(x, y, vals, xbins, ybins, stat="mean"):
    xi = np.clip(np.searchsorted(xbins, x, side="right") - 1, 0, len(xbins) - 2)
    yi = np.clip(np.searchsorted(ybins, y, side="right") - 1, 0, len(ybins) - 2)
    shape = (len(ybins) - 1, len(xbins) - 1)
    flat_sum = np.zeros(shape[0] * shape[1])
    flat_cnt = np.zeros(shape[0] * shape[1])
    for a, b, v in zip(yi, xi, vals):
        k = a * shape[1] + b
        flat_sum[k] += v
        flat_cnt[k] += 1.0
    with np.errstate(invalid="ignore"):
        flat = flat_sum / np.where(flat_cnt == 0, np.nan, flat_cnt)
    return flat.reshape(shape)


def robust_limits(arr, lo=PAD_PCT, hi=100.0 - PAD_PCT):
    """Limites con clipping de outliers (unos pocos puntos con Y saltando a 0 o
    a 116000 estiran los ejes y dejan el mapa vacio)."""
    a = np.asarray(arr, dtype=float)
    return float(np.percentile(a, lo)), float(np.percentile(a, hi))


def bin_max(x, y, vals, xbins, ybins):
    xi = np.clip(np.searchsorted(xbins, x, side="right") - 1, 0, len(xbins) - 2)
    yi = np.clip(np.searchsorted(ybins, y, side="right") - 1, 0, len(ybins) - 2)
    grid = np.full((len(ybins) - 1, len(xbins) - 1), -np.inf)
    np.maximum.at(grid, (yi, xi), vals)
    grid[grid == -np.inf] = np.nan
    return grid


def nan_smooth(grid, k=3):
    """Media movil nan-aware k x k sobre una rejilla con celdas vacias."""
    out = np.full_like(grid, np.nan)
    r = k // 2
    for i in range(grid.shape[0]):
        for j in range(grid.shape[1]):
            blk = grid[max(0, i - r):i + r + 1, max(0, j - r):j + r + 1]
            if np.isfinite(blk).any():
                out[i, j] = np.nanmean(blk)
    return out


def deg_axis(ax):
    """Ejes secundarios en grados equivalentes (80 deg = 53000 micro-pasos)."""
    tx = ax.secondary_xaxis("top", functions=(lambda s: s * DEG_PER_USTEP,
                                              lambda d: d / DEG_PER_USTEP))
    tx.set_xlabel("X [deg]", fontsize=8, labelpad=2)
    tx.tick_params(labelsize=7)
    ty = ax.secondary_yaxis("right", functions=(lambda s: s * DEG_PER_USTEP,
                                               lambda d: d / DEG_PER_USTEP))
    ty.set_ylabel("Y [deg]", fontsize=8, labelpad=2, color="dimgray")
    ty.tick_params(labelsize=7, colors="dimgray")


def process(path):
    """Devuelve dict con datos y estadisticos, o None si no hay datos de dia."""
    ts, xs, ys, cs, ps = load_log(path)
    g = np.maximum(0.0, ps) / ADC_MAX * 2048.0 / 2.0
    isc = isc_ma(cs)

    day = g > DAY_G_MIN
    if day.sum() < 50:
        return None

    x0d, x1d = robust_limits(xs[day])
    y0d, y1d = robust_limits(ys[day])
    inside = day & (xs >= x0d) & (xs <= x1d) & (ys >= y0d) & (ys <= y1d)
    n_drop = int(day.sum() - inside.sum())
    if inside.sum() < 50:
        return None

    x, y, c, gg = xs[inside], ys[inside], isc[inside], g[inside]
    t_str = ts[inside]

    thr_c = np.percentile(c, PCT)
    thr_g = np.percentile(gg, PCT)
    ext = c > thr_c

    n_ext = int(ext.sum())
    if n_ext < 10:
        return None

    xe, ye, ce, ge, te = x[ext], y[ext], c[ext], gg[ext], t_str[ext]

    r = float(np.corrcoef(ce, ge)[0, 1]) if ce.std() > 0 and ge.std() > 0 else np.nan
    ce_rank = np.argsort(np.argsort(ce))
    ge_rank = np.argsort(np.argsort(ge))
    rho = float(np.corrcoef(ce_rank, ge_rank)[0, 1]) if n_ext > 2 else np.nan
    if ce.std() > 0 and ge.std() > 0:
        slope, intercept = np.polyfit(ge, ce, 1)
    else:
        slope, intercept = np.nan, np.nan

    x0e, x1e = robust_limits(xe, 1.0, 99.0)
    y0e, y1e = robust_limits(ye, 1.0, 99.0)
    m_x = 0.15 * max(1.0, x1e - x0e)
    m_y = 0.15 * max(1.0, y1e - y0e)
    in_win = (x >= x0e - m_x) & (x <= x1e + m_x) & (y >= y0e - m_y) & (y <= y1e + m_y)
    if in_win.sum() < 20:
        return None

    xbins = np.linspace(x0e - m_x, x1e + m_x, 61)
    ybins = np.linspace(y0e - m_y, y1e + m_y, 46)
    xc = (xbins[:-1] + xbins[1:]) / 2
    yc = (ybins[:-1] + ybins[1:]) / 2

    grid_c = bin_max(xe, ye, ce, xbins, ybins)
    grid_g = bin_stat(x[in_win], y[in_win], gg[in_win], xbins, ybins)
    grid_g_s = nan_smooth(grid_g, 3)

    cell_ext_c = ~np.isnan(grid_c)
    cell_ext_g = ~np.isnan(grid_g) & (grid_g > thr_g)
    inter = cell_ext_c & cell_ext_g
    union = cell_ext_c | cell_ext_g
    iou = float(inter.sum()) / float(union.sum()) if union.sum() else np.nan

    xi = np.clip(np.searchsorted(xbins, xe, side="right") - 1, 0, len(xbins) - 2)
    yi = np.clip(np.searchsorted(ybins, ye, side="right") - 1, 0, len(ybins) - 2)
    contained = float(cell_ext_g[yi, xi].mean())

    # pico del campo G suavizado, restringido a celdas con dato crudo
    valid = np.isfinite(grid_g) & np.isfinite(grid_g_s)
    gi, gj = np.unravel_index(np.argmax(np.where(valid, grid_g_s, -np.inf)),
                              grid_g_s.shape)
    g_peak = float(grid_g[gi, gj])

    wx = ce / ce.sum()
    com_x = float((wx * xe).sum())
    com_y = float((wx * ye).sum())

    i_cmax = int(np.argmax(ce))

    return dict(
        name=os.path.splitext(os.path.basename(path))[0],
        n_day=int(day.sum()), n_used=int(inside.sum()), n_drop=n_drop, n_ext=n_ext,
        x0d=x0d, x1d=x1d, y0d=y0d, y1d=y1d,
        t0=t_str[0], t1=t_str[-1],
        xe=xe, ye=ye, ce=ce, ge=ge, te=te,
        thr_c=thr_c, thr_g=thr_g,
        xbins=xbins, ybins=ybins, xc=xc, yc=yc,
        grid_c=grid_c, grid_g=grid_g, grid_g_s=grid_g_s,
        iou=iou, contained=contained,
        com_x=com_x, com_y=com_y,
        c_max=float(ce[i_cmax]), c_max_x=float(xe[i_cmax]), c_max_y=float(ye[i_cmax]),
        g_max=float(ge.max()), g_peak=g_peak,
        g_max_x=float(xc[gj]), g_max_y=float(yc[gi]),
        r=r, rho=rho, slope=float(slope), intercept=float(intercept),
    )


def make_figure(d):
    fig, axes = plt.subplots(1, 2, figsize=(15.5, 6.6))
    axA, axB = axes

    im = axA.imshow(
        np.ma.masked_invalid(d["grid_c"]),
        extent=[d["xbins"][0], d["xbins"][-1], d["ybins"][0], d["ybins"][-1]],
        origin="lower", aspect="auto", cmap="inferno",
        norm=Normalize(vmin=np.nanmin(d["grid_c"]), vmax=np.nanmax(d["grid_c"])),
        interpolation="nearest",
    )
    cb = fig.colorbar(im, ax=axA, pad=0.16)
    cb.set_label("Isc maxima por celda [mA]  (solo puntos > P%d)" % int(PCT),
                 fontsize=9)
    cb.ax.tick_params(labelsize=8)

    cs = axA.contour(d["xc"], d["yc"], d["grid_g"], levels=5, colors="white",
                     linewidths=0.8, alpha=0.7)
    axA.clabel(cs, inline=True, fontsize=6.5, fmt="%.0f")
    axA.contour(d["xc"], d["yc"], d["grid_g"], levels=[d["thr_g"]],
                colors="deepskyblue", linewidths=2.0, linestyles="--")

    axA.plot(d["xe"], d["ye"], ".", ms=1.6, color="cyan", alpha=0.35, zorder=3,
             label="puntos > P%d de Isc" % int(PCT))
    axA.plot(d["g_max_x"], d["g_max_y"], "H", ms=14, mfc="none", mec="deepskyblue",
             mew=2.4, zorder=6,
             label="piranometro: pico de G_POA (%.0f W/m2)" % d["g_peak"])
    axA.plot(d["c_max_x"], d["c_max_y"], "*", ms=19, color="lime", mec="black",
             mew=0.8, zorder=7,
             label="celda: max Isc (%.1f mA)" % d["c_max"])
    axA.plot(d["com_x"], d["com_y"], "x", ms=13, color="yellow", mew=2.4, zorder=7,
             label="centroide ponderado por Isc")
    axA.annotate(
        "", xy=(d["c_max_x"], d["c_max_y"]), xytext=(d["g_max_x"], d["g_max_y"]),
        arrowprops=dict(arrowstyle="-|>", color="lime", lw=1.8, alpha=0.9), zorder=6)

    dx = d["c_max_x"] - d["g_max_x"]
    dy = d["c_max_y"] - d["g_max_y"]
    axA.set_title("A) Zonas de corriente extrema (Isc > P%d) y campo del piranometro"
                  % int(PCT), fontsize=11)
    axA.set_xlabel("X (micro-pasos)")
    axA.set_ylabel("Y (micro-pasos)")
    axA.legend(loc="upper center", bbox_to_anchor=(0.5, -0.13), ncol=2,
               fontsize=7.5, framealpha=0.95, borderpad=0.4, handletextpad=0.5,
               labelspacing=0.4, columnspacing=1.0)
    deg_axis(axA)

    sc = axB.scatter(d["ge"], d["ce"], c=np.arange(len(d["ce"])), cmap="viridis",
                     s=9, alpha=0.75, edgecolors="none")
    if np.isfinite(d["slope"]):
        gx = np.linspace(d["ge"].min(), d["ge"].max(), 50)
        axB.plot(gx, d["slope"] * gx + d["intercept"], "-", color="crimson", lw=1.8,
                 label="ajuste lineal: Isc = %.3f*G + %.1f" % (d["slope"], d["intercept"]))
    axB.axvline(d["thr_g"], color="deepskyblue", ls="--", lw=1.6,
                label="P%d de G_POA = %.0f W/m2" % (int(PCT), d["thr_g"]))
    axB.axhline(d["thr_c"], color="black", ls=":", lw=1.6,
                label="P%d de Isc = %.1f mA" % (int(PCT), d["thr_c"]))
    cb2 = fig.colorbar(sc, ax=axB, pad=0.02)
    cb2.set_label("indice temporal (temprano -> tarde)")

    axB.set_title("B) Correlacion espacial Isc vs G_POA en el subconjunto extremo",
                  fontsize=11)
    axB.set_xlabel("G_POA del piranometro [W/m2]")
    axB.set_ylabel("Isc de la celda [mA]")
    axB.grid(True, ls="--", alpha=0.35)
    axB.legend(loc="upper left", fontsize=7.5, framealpha=0.9)
    axB.text(0.985, 0.02,
             "n = %d   Pearson r = %.3f   Spearman rho = %.3f\n"
             "extremos dentro del area G > P%d: %.0f%%   IoU = %.3f"
             % (d["n_ext"], d["r"], d["rho"], int(PCT),
                100.0 * d["contained"], d["iou"]),
             transform=axB.transAxes, ha="right", va="bottom", fontsize=8,
             bbox=dict(fc="white", alpha=0.9, ec="0.6"))

    fig.suptitle(
        "%s  -  %d muestras de dia, %d extremas (%.0f%%)  -  %s a %s\n"
        "recorrido de la sesion: X %d..%d  Y %d..%d upaso  "
        "(el marco del panel A se recorta a la zona de los puntos extremos)\n"
        "desfase pico piranometro -> pico celda: dX=%+.0f dY=%+.0f upaso "
        "(%.1f deg, %.1f deg)"
        % (d["name"], d["n_day"], d["n_ext"], 100.0 - PCT, d["t0"][11:], d["t1"][11:],
           d["x0d"], d["x1d"], d["y0d"], d["y1d"],
           dx, dy, dx * DEG_PER_USTEP, dy * DEG_PER_USTEP),
        fontsize=12, y=1.06)
    fig.tight_layout()
    out = os.path.join(OUT_DIR, "extremos_p80_%s.png" % d["name"])
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return out


def make_summary(ds):
    fig, axes = plt.subplots(2, 3, figsize=(20, 11))
    axs = axes.ravel()
    for ax, d in zip(axs, ds):
        im = ax.imshow(
            np.ma.masked_invalid(d["grid_c"]),
            extent=[d["xbins"][0], d["xbins"][-1], d["ybins"][0], d["ybins"][-1]],
            origin="lower", aspect="auto", cmap="inferno",
            norm=Normalize(vmin=np.nanmin(d["grid_c"]), vmax=np.nanmax(d["grid_c"])),
            interpolation="nearest")
        fig.colorbar(im, ax=ax, pad=0.02).set_label("Isc max [mA]", fontsize=8)
        ax.contour(d["xc"], d["yc"], d["grid_g"],
                   levels=[d["thr_g"]], colors="deepskyblue", linewidths=1.8,
                   linestyles="--")
        ax.contour(d["xc"], d["yc"], d["grid_g"], levels=6, colors="white",
                   linewidths=0.6, alpha=0.6)
        ax.plot(d["g_max_x"], d["g_max_y"], "H", ms=12, mfc="none", mec="deepskyblue",
                mew=2.2)
        ax.plot(d["c_max_x"], d["c_max_y"], "*", ms=17, color="lime", mec="black",
                mew=0.8)
        ax.set_title("%s\nn_ext=%d  r=%.2f  rho=%.2f  cont=%.0f%%  IoU=%.2f"
                     % (d["name"], d["n_ext"], d["r"], d["rho"],
                        100.0 * d["contained"], d["iou"]), fontsize=10)
        ax.set_xlabel("X (micro-pasos)", fontsize=8)
        ax.set_ylabel("Y (micro-pasos)", fontsize=8)
        ax.tick_params(labelsize=7)

    ax = axs[len(ds)]
    ax.axis("off")
    lines = ["Resumen: extremos de Isc (percentil %d) y piranometro\n" % int(PCT),
             "regla: solo puntos con Isc > P%d (dia G>%.0f W/m2, ventana P%.1f-P%.1f\n"
             "       en X e Y para descartar outliers de la trayectoria)\n"
             % (int(PCT), DAY_G_MIN, PAD_PCT, 100.0 - PAD_PCT),
             "%-12s %6s %6s %6s %6s %6s %6s" %
             ("log", "n_dia", "n_ext", "r", "rho", "cont", "IoU"),
             "-" * 56]
    for d in ds:
        lines.append("%-12s %6d %6d %6.3f %6.3f %5.0f%% %6.3f"
                     % (d["name"].replace("log_", ""), d["n_used"], d["n_ext"],
                        d["r"], d["rho"], 100.0 * d["contained"], d["iou"]))
    lines.append("")
    lines.append("%-12s %10s %10s %10s %8s" % ("log", "Isc_max_mA", "G_max", "P80_mA", "atipic"))
    lines.append("-" * 56)
    for d in ds:
        lines.append("%-12s %10.1f %10.0f %10.1f %8d"
                     % (d["name"].replace("log_", ""), d["c_max"], d["g_max"],
                        d["thr_c"], d["n_drop"]))
    lines.append("")
    lines.append("G_max = maximo instantaneo de G_POA; el hexagono marca el pico")
    lines.append("        del campo G suavizado (3x3 celdas), mas robusto")
    lines.append("cont = % de puntos extremos de Isc que caen dentro del area")
    lines.append("       > P%d medida por el piranometro (G_POA)")
    lines.append("IoU  = solapamiento de areas en la rejilla (celda a celda)")
    lines.append("atipic = muestras de dia descartadas por la ventana robusta")
    lines.append("")
    lines.append("marca hexagono = piranometro (pico del campo G_POA)")
    lines.append("marca estrella  = celda (maximo de corriente)")
    lines.append("linea cian ---  = frontera del area > P%d de G_POA" % int(PCT))
    lines.append("(aqui se destacan las zonas de mayor intensidad medida")
    lines.append(" por el piranometro, sobre la nube de puntos extremos de Isc)")
    ax.text(0.0, 0.98, "\n".join(lines), va="top", ha="left", fontsize=9,
            family="monospace", transform=ax.transAxes)

    fig.suptitle("Extremos de corriente de celda (percentil %d) vs piranometro - "
                 "5 logs mas nuevos" % int(PCT), fontsize=14, y=0.995)
    fig.tight_layout()
    out = os.path.join(OUT_DIR, "extremos_p80_resumen.png")
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return out


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 5
    os.makedirs(OUT_DIR, exist_ok=True)

    paths = sorted(glob.glob(os.path.join(LOGS_DIR, "log_*.txt")),
                   key=lambda p: re.search(r"(\d{4}-\d{2}-\d{2})", p).group(1))
    paths = paths[-n:]
    print("Logs seleccionados (%d mas nuevos):" % len(paths))
    for p in paths:
        print("   ", os.path.basename(p))

    ds = []
    for p in paths:
        d = process(p)
        if d is None:
            print("!! %s: sin datos de dia suficientes, se omite" % os.path.basename(p))
            continue
        ds.append(d)
        out = make_figure(d)
        print("%s -> n_dia=%d (usados=%d, atipicos=%d) n_ext=%d r=%.3f rho=%.3f "
              "cont=%.2f IoU=%.3f | %s"
              % (d["name"], d["n_day"], d["n_used"], d["n_drop"], d["n_ext"],
                 d["r"], d["rho"], d["contained"], d["iou"], os.path.basename(out)))

    if not ds:
        print("Sin figuras generadas.")
        return
    print("Resumen ->", os.path.basename(make_summary(ds)))


if __name__ == "__main__":
    main()