#!/usr/bin/env python3
"""MP1 · segunda implementación independiente de las métricas PUNTUALES (DEFINICIONES_MP1 §7, §12(7)).

Escrita solo con la especificación sellada (md5 3da2be07…), sin leer la primera implementación.

Uso:  python3 metricas_impl2.py <h>     con h ∈ {ou, des, mod, ser, rec}
Entradas en $MP1_IN (por defecto /in) / mp1_<h>/…; salida en $MP1_OUT (por defecto /w/salidas)/impl2_<h>.json.

Algoritmos (elegidos a propósito distintos del camino obvio sklearn/statsmodels):
- AUC: Mann-Whitney con rangos medios (scipy.stats.rankdata 'average'; empates ½).
- AUC ponderada (rec): agrupación por valor único de la puntuación y peso negativo acumulado por debajo;
  cada par positivo-negativo pesa w_pos·w_neg, empates ½.
- Pendiente: IRLS propio en numpy (mínimos cuadrados ponderados con lstsq + reducción de paso),
  y ~ a + b·logit(p), pesos de frecuencia opcionales.
- Log-loss: p recortada a [1e-6, 1 − 1e-6], logaritmo natural, media por fila.
- AP: a mano, definición de sklearn average_precision_score: Σ_umbrales (R_n − R_{n−1})·P_n, con los empates
  de puntuación entrando juntos en el mismo umbral.
- Deciles: pandas.qcut(p, 10, duplicates='drop') tal cual (definición sellada).
- Los CSV se leen como texto y cada número se convierte con float() de Python (ida y vuelta exacta).
"""
import json
import math
import os
import sys

import numpy as np
import pandas as pd
import scipy
from scipy.stats import rankdata

EPS_LL = 1e-6        # recorte de log-loss (§7)
EPS_LOGIT = 1e-6     # recorte para logit(p) en la pendiente (ambigüedad: la spec no lo fija; se usa el de log-loss)
DESENLACES = ('ou', 'des', 'mod', 'ser', 'rec')


# ----------------------------------------------------------------------------------------------- lectura
def leer_csv(ruta):
    """Todo como texto, sin convertir 'NA'/'nan'/'' en nulos: los nulos los decide cada columna."""
    return pd.read_csv(ruta, dtype=str, keep_default_na=False, na_filter=False)


def a_float(col):
    """Texto → float64 con float() de Python (exacto). '' / 'nan' / 'NaN' / 'None' / 'NA' → NaN."""
    nulos = {'', 'nan', 'NaN', 'NAN', 'None', 'NA', 'null', 'NULL'}
    return np.array([np.nan if v.strip() in nulos else float(v) for v in col], dtype=np.float64)


def a_y(col):
    y = a_float(col)
    if np.isnan(y).any():
        raise ValueError('y tiene nulos')
    if not np.isin(y, (0.0, 1.0)).all():
        raise ValueError('y no es 0/1: valores %r' % sorted(set(y.tolist()))[:10])
    return y.astype(np.int64)


# ----------------------------------------------------------------------------------------------- métricas
def auc_mw(score, y):
    """AUC de Mann-Whitney con rangos medios (empates ½). Filas con score nulo fuera."""
    s = np.asarray(score, dtype=np.float64)
    y = np.asarray(y, dtype=np.int64)
    ok = ~np.isnan(s)
    s, y = s[ok], y[ok]
    n1 = int(y.sum())
    n0 = int(len(y) - n1)
    if n1 == 0 or n0 == 0:
        return None
    r = rankdata(s, method='average')
    suma_r1 = math.fsum(r[y == 1].tolist())
    u1 = suma_r1 - n1 * (n1 + 1) / 2.0
    return u1 / (float(n1) * float(n0))


def auc_ponderada(score, y, w):
    """Σ_{i∈pos, j∈neg} w_i·w_j·[s_i > s_j + ½·(s_i = s_j)] / (Σ w_pos · Σ w_neg)."""
    s = np.asarray(score, dtype=np.float64)
    y = np.asarray(y, dtype=np.int64)
    w = np.asarray(w, dtype=np.float64)
    ok = ~np.isnan(s)
    s, y, w = s[ok], y[ok], w[ok]
    valores, inv = np.unique(s, return_inverse=True)       # orden ascendente
    wp = np.bincount(inv, weights=w * (y == 1), minlength=len(valores))
    wn = np.bincount(inv, weights=w * (y == 0), minlength=len(valores))
    Wp, Wn = math.fsum(wp.tolist()), math.fsum(wn.tolist())
    if Wp <= 0 or Wn <= 0:
        return None
    neg_debajo = np.cumsum(wn) - wn                         # peso negativo estrictamente por debajo
    num = math.fsum((wp * (neg_debajo + 0.5 * wn)).tolist())
    return num / (Wp * Wn)


def brier(p, y):
    return float(np.mean((p - y) ** 2))


def log_loss(p, y):
    pc = np.clip(p, EPS_LL, 1.0 - EPS_LL)
    return float(-np.mean(y * np.log(pc) + (1 - y) * np.log1p(-pc)))


def precision_media(score, y):
    """AP como sklearn: umbrales = valores únicos de score en orden descendente."""
    s = np.asarray(score, dtype=np.float64)
    y = np.asarray(y, dtype=np.int64)
    P = int(y.sum())
    if P == 0:
        return None
    valores, inv = np.unique(s, return_inverse=True)
    tp_g = np.bincount(inv, weights=y.astype(np.float64), minlength=len(valores))[::-1]
    n_g = np.bincount(inv, minlength=len(valores)).astype(np.float64)[::-1]
    tp = np.cumsum(tp_g)
    tot = np.cumsum(n_g)
    prec = tp / tot
    d_recall = tp_g / P
    return math.fsum((d_recall * prec).tolist())


def _expit(z):
    return np.where(z >= 0, 1.0 / (1.0 + np.exp(-np.abs(z))), np.exp(-np.abs(z)) / (1.0 + np.exp(-np.abs(z))))


def _loglik(beta, X, y, w):
    eta = X @ beta
    return math.fsum((w * (y * eta - np.logaddexp(0.0, eta))).tolist())


def pendiente_irls(p, y, w=None, tol=1e-14, max_iter=200):
    """Coeficiente b de y ~ a + b·logit(p) por IRLS propio (logística sin penalizar, pesos de frecuencia).

    Devuelve (b, a, iteraciones, filas_recortadas)."""
    p = np.asarray(p, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    w = np.ones_like(p) if w is None else np.asarray(w, dtype=np.float64)
    recortadas = int(((p < EPS_LOGIT) | (p > 1.0 - EPS_LOGIT)).sum())
    pc = np.clip(p, EPS_LOGIT, 1.0 - EPS_LOGIT)
    x = np.log(pc) - np.log1p(-pc)
    X = np.column_stack([np.ones_like(x), x])
    # arranque: a = logit de la tasa ponderada, b = 0
    tasa = math.fsum((w * y).tolist()) / math.fsum(w.tolist())
    beta = np.array([math.log(tasa / (1.0 - tasa)), 0.0])
    ll = _loglik(beta, X, y, w)
    for it in range(1, max_iter + 1):
        eta = X @ beta
        mu = _expit(eta)
        v = np.maximum(mu * (1.0 - mu), 1e-300)
        z = eta + (y - mu) / v                       # respuesta de trabajo
        sw = np.sqrt(w * v)
        nuevo, *_ = np.linalg.lstsq(X * sw[:, None], z * sw, rcond=None)
        paso = nuevo - beta
        t = 1.0
        while True:
            cand = beta + t * paso
            ll_c = _loglik(cand, X, y, w)
            if ll_c >= ll - 1e-12 * (1.0 + abs(ll)) or t < 1e-8:
                break
            t *= 0.5
        delta = np.max(np.abs(cand - beta))
        beta, ll = cand, ll_c
        if delta < tol * (1.0 + np.max(np.abs(beta))):
            return float(beta[1]), float(beta[0]), it, recortadas
    raise RuntimeError('IRLS no convergió en %d iteraciones (beta=%r)' % (max_iter, beta))


def deciles(p, y):
    cats = pd.qcut(pd.Series(p), 10, duplicates='drop')
    codigos = cats.cat.codes.to_numpy()
    out = []
    for k in range(len(cats.cat.categories)):
        m = codigos == k
        n = int(m.sum())
        out.append({'n': n,
                    'predicha': float(np.mean(p[m])) if n else None,
                    'observada': float(np.mean(y[m])) if n else None})
    if (codigos < 0).any():
        raise RuntimeError('qcut dejó filas sin decil')
    return out


# ----------------------------------------------------------------------------------------------- celdas
def celdas(df, y, ruta_celdas):
    """Observados por celda organo_clave|cpv_cel (≥ 10 lotes, órgano no nulo) y cobertura con lo/hi del fichero."""
    org = df['organo_clave']
    cpv = df['cpv_cel']
    con_org = org.str.len().to_numpy() > 0
    clave = (org + '|' + cpv)[con_org].to_numpy()
    yy = y[con_org]
    t = pd.DataFrame({'celda': clave, 'y': yy})
    g = t.groupby('celda', sort=True)['y'].agg(['size', 'sum'])
    g = g[g['size'] >= 10]
    mias = {c: (int(r['size']), int(r['sum'])) for c, r in g.iterrows()}

    f = leer_csv(ruta_celdas)
    fn, fobs, flo, fhi = a_float(f['n']), a_float(f['observados']), a_float(f['lo']), a_float(f['hi'])
    fich = {}
    duplicadas = 0
    for i, c in enumerate(f['celda'].tolist()):
        if c in fich:
            duplicadas += 1
        fich[c] = (fn[i], fobs[i], flo[i], fhi[i])

    solo_mias = sorted(set(mias) - set(fich))
    solo_fich = sorted(set(fich) - set(mias))
    distintas = [c for c in mias if c in fich and (fich[c][0] != mias[c][0] or fich[c][1] != mias[c][1])]
    coinciden = (not solo_mias) and (not solo_fich) and (not distintas) and duplicadas == 0

    evaluables = [c for c in mias if c in fich]
    dentro = sum(1 for c in evaluables if fich[c][2] <= mias[c][1] <= fich[c][3])
    cobertura = dentro / len(evaluables) if evaluables else None
    diag = {'celdas_mias': len(mias), 'celdas_fichero': len(fich), 'solo_mias': len(solo_mias),
            'solo_fichero': len(solo_fich), 'n_u_obs_distintos': len(distintas), 'duplicadas_fichero': duplicadas,
            'lotes_sin_organo': int((~con_org).sum()), 'dentro': dentro, 'evaluables': len(evaluables),
            'ejemplos_solo_mias': solo_mias[:3], 'ejemplos_solo_fichero': solo_fich[:3],
            'ejemplos_distintas': [(c, mias[c], fich[c][:2]) for c in distintas[:3]]}
    return cobertura, coinciden, diag


# ----------------------------------------------------------------------------------------------- principal
def versiones():
    return {'numpy': np.__version__, 'pandas': pd.__version__, 'scipy': scipy.__version__}


def limpio(o):
    """NaN/inf → None, numpy → Python, para json con allow_nan=False."""
    if isinstance(o, dict):
        return {k: limpio(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [limpio(v) for v in o]
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (float, np.floating)):
        o = float(o)
        return o if math.isfinite(o) else None
    return o


def metricas_lotes(h, base):
    d = os.path.join(base, 'mp1_%s' % h)
    ev = leer_csv(os.path.join(d, 'eval_prueba.csv'))
    y = a_y(ev['y'])
    p = a_float(ev['p'])
    if np.isnan(p).any():
        raise ValueError('p tiene %d nulos en eval_prueba' % int(np.isnan(p).sum()))
    cols = set(ev.columns)
    print('[%s] eval_prueba: %d filas, %d positivos, columnas=%s' % (h, len(y), int(y.sum()), list(ev.columns)))

    aucs = {'p': auc_mw(p, y)}
    for c in ('p_raw', 'b1', 'b3'):
        if c in cols:
            s = a_float(ev[c])
            if np.isnan(s).any():
                print('  AVISO: %s tiene %d nulos (fuera de su AUC)' % (c, int(np.isnan(s).sum())))
            aucs[c] = auc_mw(s, y)
        elif c != 'b3':
            aucs[c] = None
            print('  AVISO: falta la columna %s' % c)

    b2 = None
    if h in ('ou', 'des'):
        if 'p_b2' in cols:
            pb2 = a_float(ev['p_b2'])
            m = ~np.isnan(pb2)
            b2 = {'n': int(m.sum()), 'p': auc_mw(p[m], y[m]), 'p_b2': auc_mw(pb2[m], y[m])}
        else:
            b2 = {'n': 0, 'p': None, 'p_b2': None}
            print('  AVISO: falta la columna p_b2')

    b, a, it, rec = pendiente_irls(p, y)
    print('  pendiente IRLS: b=%.15g a=%.15g iter=%d filas_recortadas=%d' % (b, a, it, rec))

    prueba = {
        'n': int(len(y)),
        'positivos': int(y.sum()),
        'auc': aucs,
    }
    if b2 is not None:
        prueba['auc_b2_subconjunto'] = b2
    prueba.update({
        'brier': brier(p, y),
        'log_loss': log_loss(p, y),
        'ap': precision_media(p, y),
        'cal_global_puntos': 100.0 * (float(np.mean(p)) - float(np.mean(y))),
        'o_e': float(y.sum()) / math.fsum(p.tolist()),
        'pendiente': b,
        'deciles': deciles(p, y),
    })

    ruta_celdas = os.path.join(d, 'celdas_test.csv')
    if h != 'ser' and os.path.exists(ruta_celdas) and 'cpv_cel' in cols:
        cob, coinc, diag = celdas(ev, y, ruta_celdas)
        print('  celdas:', json.dumps(limpio(diag), ensure_ascii=False))
        prueba['cobertura'] = cob
        prueba['celdas_coinciden'] = coinc
    else:
        if h != 'ser':
            print('  AVISO: sin celdas_test.csv o sin cpv_cel: cobertura null')
        prueba['cobertura'] = None
        prueba['celdas_coinciden'] = None

    pr = leer_csv(os.path.join(d, 'predicciones.csv'))
    rc = pr[pr['ventana'] == 'recal']
    print('  predicciones.csv ventanas:', pr['ventana'].value_counts().to_dict())
    yr = a_y(rc['y'])
    recal = {'auc': {'p': auc_mw(a_float(rc['p']), yr), 'b1': auc_mw(a_float(rc['b1']), yr)}}
    nt = int((pr['ventana'] == 'test').sum())
    if nt != len(y):
        print('  AVISO: predicciones.csv tiene %d filas test y eval_prueba %d' % (nt, len(y)))
    return {'desenlace': h, 'versiones': versiones(), 'prueba': prueba, 'recal': recal}


def metricas_rec(base):
    d = os.path.join(base, 'mp1_rec')
    pr = leer_csv(os.path.join(d, 'predicciones.csv'))
    print('[rec] predicciones.csv: %d filas, columnas=%s' % (len(pr), list(pr.columns)))
    if 'ventana' in pr.columns:
        vc = pr['ventana'].value_counts().to_dict()
        print('  ventanas:', vc)
        etiqueta = 'test' if 'test' in vc else ('prueba' if 'prueba' in vc else None)
        if etiqueta is None:
            raise ValueError('predicciones.csv de rec tiene ventana pero ni test ni prueba: %r' % vc)
        pr = pr[pr['ventana'] == etiqueta]
    y = a_y(pr['y'])
    w = a_float(pr['w'])
    p = a_float(pr['p'])
    if np.isnan(w).any() or np.isnan(p).any() or (w < 0).any():
        raise ValueError('w o p con nulos o negativos')
    aucs = {'p': auc_ponderada(p, y, w)}
    for c in ('p_raw', 'b3'):
        aucs[c] = auc_ponderada(a_float(pr[c]), y, w) if c in pr.columns else None
    b, a, it, rec = pendiente_irls(p, y, w)
    o_e = math.fsum((w * y).tolist()) / math.fsum(p.tolist())
    print('  n=%d positivos=%d Σw·y=%.10g Σp=%.10g pendiente b=%.15g iter=%d recortadas=%d'
          % (len(y), int(y.sum()), math.fsum((w * y).tolist()), math.fsum(p.tolist()), b, it, rec))
    return {'desenlace': 'rec', 'versiones': versiones(),
            'prueba': {'auc': aucs, 'o_e': o_e, 'pendiente': b}}


def main(argv):
    if len(argv) != 2 or argv[1] not in DESENLACES:
        sys.exit('uso: metricas_impl2.py {%s}' % '|'.join(DESENLACES))
    h = argv[1]
    base = os.environ.get('MP1_IN', '/in')
    out_dir = os.environ.get('MP1_OUT', '/w/salidas')
    res = metricas_rec(base) if h == 'rec' else metricas_lotes(h, base)
    os.makedirs(out_dir, exist_ok=True)
    ruta = os.path.join(out_dir, 'impl2_%s.json' % h)
    with open(ruta, 'w', encoding='utf-8') as fh:
        json.dump(limpio(res), fh, ensure_ascii=False, indent=1, allow_nan=False)
        fh.write('\n')
    print('escrito', ruta)
    return res


if __name__ == '__main__':
    main(sys.argv)