#!/usr/bin/env python3
"""Banco sintético de metricas_impl2.py con respuestas conocidas.

1) Casos pequeños calculados a mano (AUC con empates, AP con empates, AUC ponderada).
2) Contraste con algoritmos independientes: AUC por fuerza bruta de pares, sklearn roc_auc_score /
   average_precision_score / log_loss / brier_score_loss, statsmodels GLM binomial (con freq_weights).
3) Pendiente ≈ 1 si y ~ Bernoulli(p) y ≈ 0,7 si y ~ Bernoulli(expit(0,3 + 0,7·logit p)).
4) Extremo a extremo: ficheros CSV sintéticos (órganos con «|», comas y nulos; p con empates; p_b2 con huecos;
   celdas con cobertura conocida) pasados por main() y comparados con los valores esperados.
Se ejecuta en la imagen bquantagents-worker con /w = impl2/. Escribe solo en /w/sint/.
"""
import json
import math
import os
import sys

import numpy as np
import pandas as pd
from scipy.special import expit, logit
from sklearn.metrics import average_precision_score, brier_score_loss, log_loss, roc_auc_score
import statsmodels.api as sm

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import metricas_impl2 as M  # noqa: E402

FALLOS = []
N_OK = [0]


def chk(nombre, obtenido, esperado, tol):
    ok = (obtenido is None and esperado is None) or (
        obtenido is not None and esperado is not None and abs(obtenido - esperado) <= tol)
    if ok:
        N_OK[0] += 1
    else:
        FALLOS.append('%s: obtenido=%r esperado=%r tol=%g' % (nombre, obtenido, esperado, tol))
    print('%-58s %s  obt=%.15g  esp=%.15g' % (nombre, 'OK ' if ok else 'MAL', obtenido if obtenido is not None
                                               else float('nan'), esperado if esperado is not None else float('nan')))


def chk_bool(nombre, cond):
    if cond:
        N_OK[0] += 1
    else:
        FALLOS.append(nombre)
    print('%-58s %s' % (nombre, 'OK ' if cond else 'MAL'))


def auc_fuerza_bruta(s, y, w=None):
    s = np.asarray(s, float)
    y = np.asarray(y, int)
    w = np.ones(len(s)) if w is None else np.asarray(w, float)
    sp, sn = s[y == 1], s[y == 0]
    wp, wn = w[y == 1], w[y == 0]
    gana = (sp[:, None] > sn[None, :]).astype(float) + 0.5 * (sp[:, None] == sn[None, :])
    return float((wp[:, None] * wn[None, :] * gana).sum() / (wp.sum() * wn.sum()))


def pendiente_sm(p, y, w=None):
    x = logit(np.clip(p, 1e-6, 1 - 1e-6))
    X = sm.add_constant(x)
    r = sm.GLM(y, X, family=sm.families.Binomial(), freq_weights=w).fit(tol=1e-14, maxiter=200)
    return float(r.params[1])


# ------------------------------------------------------------------ 1) casos a mano
print('== 1) casos a mano')
# pos = {0,35; 0,8; 0,4}, neg = {0,1; 0,4; 0,4}: 0,35→1 par; 0,8→3; 0,4→1+½+½ = 2  ⇒ 6/9
s = [0.1, 0.4, 0.35, 0.8, 0.4, 0.4]
y = [0, 0, 1, 1, 1, 0]
chk('AUC a mano con empates (6/9)', M.auc_mw(s, y), 6 / 9, 1e-15)
chk('AUC todo empatado (½)', M.auc_mw([0.3] * 6, y), 0.5, 1e-15)
chk('AUC perfecta', M.auc_mw([0, 0, 1, 1], [0, 0, 1, 1]), 1.0, 1e-15)
chk('AUC invertida', M.auc_mw([1, 1, 0, 0], [0, 0, 1, 1]), 0.0, 1e-15)
chk_bool('AUC sin positivos → None', M.auc_mw([0.1, 0.2], [0, 0]) is None)
chk('AUC ponderada w=1 = no ponderada', M.auc_ponderada(s, y, [1] * 6), 6 / 9, 1e-15)
# pesos: pos 0,35 (w=2), 0,8 (w=1), 0,4 (w=3); neg 0,1 (w=1), 0,4 (w=1), 0,4 (w=2)
# num = 2·(1·1) + 1·(1+1+2) + 3·(1 + ½·1 + ½·2) = 2 + 4 + 7,5 = 13,5 ; den = 6·4 = 24
chk('AUC ponderada a mano (13,5/24)', M.auc_ponderada(s, y, [1, 1, 2, 1, 3, 2]), 13.5 / 24, 1e-15)
# AP: s = [0,9; 0,8; 0,8; 0,1], y = [1,0,1,0] → umbral 0,9: P=1, ΔR=½; 0,8: P=2/3, ΔR=½ ⇒ ½ + ⅓
chk('AP a mano con empates (5/6)', M.precision_media([0.9, 0.8, 0.8, 0.1], [1, 0, 1, 0]), 5 / 6, 1e-15)
# log-loss y Brier a mano
chk('Brier a mano', M.brier(np.array([0.2, 0.9]), np.array([0, 1])), (0.04 + 0.01) / 2, 1e-15)
chk('log-loss a mano con recorte (p=0 con y=1)', M.log_loss(np.array([0.0, 0.5]), np.array([1, 0])),
    (-math.log(1e-6) - math.log(0.5)) / 2, 1e-12)
# deciles: p = 1..20 /21 → 10 tramos de 2
pp = np.arange(1, 21) / 21.0
yy = (np.arange(20) % 2).astype(int)
dd = M.deciles(pp, yy)
chk_bool('deciles 1..20: 10 tramos de 2', len(dd) == 10 and all(d['n'] == 2 for d in dd))
chk('deciles 1..20: predicha del decil 1', dd[0]['predicha'], 1.5 / 21, 1e-15)
chk('deciles 1..20: observada del decil 1', dd[0]['observada'], 0.5, 1e-15)
# empates fuertes: cuantiles lineales q.5 = pos 49,5 → 0,15; q.8 = pos 79,2 → 0,26; q.9 = pos 89,1 → 0,5
# aristas únicas [0,1; 0,15; 0,2; 0,26; 0,5; 0,9] → 5 tramos, el (0,2; 0,26] VACÍO (n=0, medias null)
pt = np.array([0.1] * 50 + [0.2] * 30 + [0.5] * 15 + [0.9] * 5)
yt = (np.arange(100) % 4 == 0).astype(int)
dt = M.deciles(pt, yt)
chk_bool('deciles con empates: n = [50, 30, 0, 15, 5], vacío con null',
         [d['n'] for d in dt] == [50, 30, 0, 15, 5] and dt[2]['predicha'] is None and dt[2]['observada'] is None)
chk('deciles con empates: predicha del tramo 4', dt[3]['predicha'], 0.5, 1e-15)
chk('deciles con empates: observada del tramo 1', dt[0]['observada'], 13 / 50, 1e-15)

# ------------------------------------------------------------------ 2) contraste con algoritmos independientes
print('== 2) contraste aleatorio con fuerza bruta / sklearn / statsmodels')
rng = np.random.default_rng(20261026)
for rep in range(5):
    n = 1500
    pr = np.round(rng.beta(2, 5, n), 2)            # muchos empates
    yr = rng.binomial(1, pr)
    ww = np.where(yr == 1, rng.choice([1.0, 1.25, 2.0, 1 / 0.7], n), 1.0)
    chk('rep%d AUC rangos vs fuerza bruta' % rep, M.auc_mw(pr, yr), auc_fuerza_bruta(pr, yr), 1e-12)
    chk('rep%d AUC rangos vs sklearn' % rep, M.auc_mw(pr, yr), roc_auc_score(yr, pr), 1e-12)
    chk('rep%d AUC ponderada vs fuerza bruta' % rep, M.auc_ponderada(pr, yr, ww), auc_fuerza_bruta(pr, yr, ww),
        1e-12)
    chk('rep%d AUC ponderada vs sklearn sample_weight' % rep, M.auc_ponderada(pr, yr, ww),
        roc_auc_score(yr, pr, sample_weight=ww), 1e-12)
    wi = rng.integers(1, 4, n)
    chk('rep%d AUC ponderada entera = filas replicadas' % rep, M.auc_ponderada(pr, yr, wi),
        M.auc_mw(np.repeat(pr, wi), np.repeat(yr, wi)), 1e-12)
    chk('rep%d AP vs sklearn' % rep, M.precision_media(pr, yr), average_precision_score(yr, pr), 1e-12)
    chk('rep%d log-loss vs sklearn' % rep, M.log_loss(pr, yr), log_loss(yr, np.clip(pr, 1e-6, 1 - 1e-6)), 1e-12)
    chk('rep%d Brier vs sklearn' % rep, M.brier(pr, yr), brier_score_loss(yr, pr), 1e-15)
    prc = np.clip(pr, 0.005, 0.995)
    chk('rep%d pendiente IRLS vs statsmodels' % rep, M.pendiente_irls(prc, yr)[0], pendiente_sm(prc, yr), 1e-8)
    chk('rep%d pendiente IRLS ponderada vs statsmodels freq_w' % rep, M.pendiente_irls(prc, yr, ww)[0],
        pendiente_sm(prc, yr, ww), 1e-8)
    chk('rep%d pendiente w entera = filas replicadas' % rep, M.pendiente_irls(prc, yr, wi)[0],
        M.pendiente_irls(np.repeat(prc, wi), np.repeat(yr, wi))[0], 1e-9)
    # deciles: contraste con groupby de pandas sobre la misma qcut
    cats = pd.qcut(pd.Series(pr), 10, duplicates='drop')
    g = pd.DataFrame({'c': cats, 'p': pr, 'y': yr}).groupby('c', observed=False).agg(
        n=('p', 'size'), pm=('p', 'mean'), ym=('y', 'mean'))
    dm = M.deciles(pr, yr)
    chk_bool('rep%d deciles: mismos n que groupby (%d tramos)' % (rep, len(dm)),
             [d['n'] for d in dm] == g['n'].tolist())
    chk('rep%d deciles: máx |Δ predicha| vs groupby' % rep,
        max(abs(d['predicha'] - b) for d, b in zip(dm, g['pm'])), 0.0, 1e-15)
    chk('rep%d deciles: máx |Δ observada| vs groupby' % rep,
        max(abs(d['observada'] - b) for d, b in zip(dm, g['ym'])), 0.0, 1e-15)

# ------------------------------------------------------------------ 3) pendiente con respuesta conocida
print('== 3) pendiente con respuesta conocida')
n = 400_000
pz = expit(rng.normal(-1.0, 1.0, n))
yz = rng.binomial(1, pz)
b1, a1, it1, _ = M.pendiente_irls(pz, yz)
chk('pendiente y~Bernoulli(p) ≈ 1 (tol 0,02)', b1, 1.0, 0.02)
chk('ordenada y~Bernoulli(p) ≈ 0 (tol 0,02)', a1, 0.0, 0.02)
chk('pendiente IRLS vs statsmodels n=400k', b1, pendiente_sm(pz, yz), 1e-8)
yq = rng.binomial(1, expit(0.3 + 0.7 * logit(pz)))
chk('pendiente y~Bernoulli(expit(0,3+0,7·logit p)) ≈ 0,7', M.pendiente_irls(pz, yq)[0], 0.7, 0.02)

# ------------------------------------------------------------------ 4) extremo a extremo con ficheros
print('== 4) extremo a extremo con CSV sintéticos')
BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'sint')
IN = os.path.join(BASE, 'in')
OUT = os.path.join(BASE, 'out')
for sub in ('mp1_ou', 'mp1_rec', 'mp1_ser'):
    os.makedirs(os.path.join(IN, sub), exist_ok=True)
os.makedirs(OUT, exist_ok=True)

organos = ['S79|CONSEJERIA DE HACIENDA, MELILLA', 'D:L01300367', 'D:EA0003089', 'NA', '', 'L|X|Y']
cpvs = ['45', '79', '??', '09']
niveles = np.round(np.linspace(0.05, 0.6, 12), 4)   # p con empates (como una isotónica)


def lotes(nn, semilla):
    r = np.random.default_rng(semilla)
    p = r.choice(niveles, nn)
    df = pd.DataFrame({
        'licitacion_id': ['PLACSP:x/%d' % i for i in range(nn)],
        'lote_u': ['PLACSP:x/%d' % i for i in range(nn)],
        'organo_clave': r.choice(organos, nn, p=[0.3, 0.25, 0.2, 0.1, 0.1, 0.05]),
        'cpv_cel': r.choice(cpvs, nn),
        'y': r.binomial(1, p),
        'p': p,
        'p_raw': np.clip(p + r.normal(0, 0.05, nn), 0.001, 0.999),
        'b1': r.uniform(0, 1, nn),
        'b3': np.round(r.uniform(0, 1, nn), 1),
    })
    pb2 = np.clip(p + r.normal(0, 0.1, nn), 0.001, 0.999)
    pb2[r.uniform(size=nn) < 0.3] = np.nan
    df['p_b2'] = pb2
    return df


ev = lotes(2400, 1)
ev.to_csv(os.path.join(IN, 'mp1_ou', 'eval_prueba.csv'), index=False)
rc = lotes(1800, 2)
pred = pd.concat([rc.assign(ventana='recal'), ev.assign(ventana='test'), lotes(50, 3).assign(ventana='pred')])
pred.to_csv(os.path.join(IN, 'mp1_ou', 'predicciones.csv'), index=False)

# celdas esperadas, con bucle Python llano (otra vía que el groupby del script)
acc = {}
for o, c, yv in zip(ev['organo_clave'], ev['cpv_cel'], ev['y']):
    if o == '':
        continue
    k = o + '|' + c
    n0, s0 = acc.get(k, (0, 0))
    acc[k] = (n0 + 1, s0 + int(yv))
cel = sorted(k for k, v in acc.items() if v[0] >= 10)
filas = []
dentro_esp = 0
for i, k in enumerate(cel):
    n0, s0 = acc[k]
    if i % 3 == 0:          # fuera: lo por encima de lo observado
        lo, hi = s0 + 1, s0 + 5
    else:                   # dentro, con los bordes exactos en la mitad de los casos
        lo, hi = (s0, s0) if i % 2 else (max(s0 - 2, 0), s0 + 2)
        dentro_esp += 1
    filas.append({'celda': k, 'n': n0, 'observados': float(s0), 'lo': float(lo), 'hi': float(hi)})
pd.DataFrame(filas).to_csv(os.path.join(IN, 'mp1_ou', 'celdas_test.csv'), index=False)
cob_esp = dentro_esp / len(cel)

os.environ['MP1_IN'] = IN
os.environ['MP1_OUT'] = OUT
M.main(['x', 'ou'])
J = json.load(open(os.path.join(OUT, 'impl2_ou.json')))
P = J['prueba']
pe = ev['p'].to_numpy()
ye = ev['y'].to_numpy()
chk_bool('e2e: claves del JSON', set(J) == {'desenlace', 'versiones', 'prueba', 'recal'})
chk_bool('e2e: n y positivos', P['n'] == len(ev) and P['positivos'] == int(ye.sum()))
for c in ('p', 'p_raw', 'b1', 'b3'):
    chk('e2e: AUC %s vs fuerza bruta' % c, P['auc'][c], auc_fuerza_bruta(ev[c].to_numpy(), ye), 1e-12)
m = ev['p_b2'].notna().to_numpy()
chk_bool('e2e: n del subconjunto B2', P['auc_b2_subconjunto']['n'] == int(m.sum()))
chk('e2e: AUC p en subconjunto B2', P['auc_b2_subconjunto']['p'], roc_auc_score(ye[m], pe[m]), 1e-12)
chk('e2e: AUC p_b2 en subconjunto B2', P['auc_b2_subconjunto']['p_b2'],
    roc_auc_score(ye[m], ev['p_b2'].to_numpy()[m]), 1e-12)
chk('e2e: Brier', P['brier'], brier_score_loss(ye, pe), 1e-15)
chk('e2e: log-loss', P['log_loss'], log_loss(ye, pe), 1e-12)
chk('e2e: AP', P['ap'], average_precision_score(ye, pe), 1e-12)
chk('e2e: calibración global (puntos)', P['cal_global_puntos'], 100 * (pe.mean() - ye.mean()), 1e-12)
chk('e2e: O/E', P['o_e'], ye.sum() / pe.sum(), 1e-12)
chk('e2e: pendiente vs statsmodels', P['pendiente'], pendiente_sm(pe, ye), 1e-8)
g = pd.DataFrame({'c': pd.qcut(ev['p'], 10, duplicates='drop'), 'p': pe, 'y': ye}).groupby(
    'c', observed=False).agg(n=('p', 'size'), pm=('p', 'mean'), ym=('y', 'mean'))
chk_bool('e2e: deciles (n) = groupby (%d tramos)' % len(g), [d['n'] for d in P['deciles']] == g['n'].tolist())
chk('e2e: deciles máx |Δ observada|', max(abs(d['observada'] - b) for d, b in zip(P['deciles'], g['ym'])), 0, 1e-15)
chk('e2e: cobertura conocida (%d/%d)' % (dentro_esp, len(cel)), P['cobertura'], cob_esp, 1e-15)
chk_bool('e2e: celdas_coinciden = true', P['celdas_coinciden'] is True)
yrc = rc['y'].to_numpy()
chk('e2e: AUC recal p', J['recal']['auc']['p'], roc_auc_score(yrc, rc['p']), 1e-12)
chk('e2e: AUC recal b1', J['recal']['auc']['b1'], roc_auc_score(yrc, rc['b1']), 1e-12)

# celdas con un recuento alterado → celdas_coinciden = false (la cobertura sigue con mis observados)
alt = pd.DataFrame(filas)
alt.loc[0, 'observados'] = alt.loc[0, 'observados'] + 1
alt.to_csv(os.path.join(IN, 'mp1_ou', 'celdas_test.csv'), index=False)
M.main(['x', 'ou'])
P2 = json.load(open(os.path.join(OUT, 'impl2_ou.json')))['prueba']
chk_bool('e2e: observados alterados → celdas_coinciden = false', P2['celdas_coinciden'] is False)
chk('e2e: cobertura con fichero alterado (mis observados)', P2['cobertura'], cob_esp, 1e-15)
pd.DataFrame(filas).to_csv(os.path.join(IN, 'mp1_ou', 'celdas_test.csv'), index=False)

# ser: sin celdas ni p_b2
ser = lotes(900, 4).drop(columns=['cpv_cel', 'p_b2'])
ser.to_csv(os.path.join(IN, 'mp1_ser', 'eval_prueba.csv'), index=False)
pd.concat([lotes(700, 5).drop(columns=['cpv_cel', 'p_b2']).assign(ventana='recal'),
           ser.assign(ventana='test')]).to_csv(os.path.join(IN, 'mp1_ser', 'predicciones.csv'), index=False)
M.main(['x', 'ser'])
PS = json.load(open(os.path.join(OUT, 'impl2_ser.json')))['prueba']
chk_bool('e2e ser: cobertura null, sin auc_b2_subconjunto',
         PS['cobertura'] is None and 'auc_b2_subconjunto' not in PS)
chk('e2e ser: AUC b3', PS['auc']['b3'], roc_auc_score(ser['y'], ser['b3']), 1e-12)

# rec: pesos 1/s en positivos, ventana train/test
r = np.random.default_rng(6)
nn = 3000
p = np.round(r.beta(1, 12, nn), 3)
yv = r.binomial(1, np.minimum(p * 1.1, 1))
s_frac = r.choice([1.0, 0.8, 0.5], nn)
w = np.where(yv == 1, 1 / s_frac, 1.0)
recdf = pd.DataFrame({'ventana': np.where(np.arange(nn) < 1000, 'train', 'test'), 'licitacion_id': np.arange(nn),
                      'y': yv, 'w': w, 'p': p, 'p_raw': np.clip(p + r.normal(0, 0.01, nn), 1e-4, 1),
                      'b3': r.normal(0, 1, nn)})
recdf.to_csv(os.path.join(IN, 'mp1_rec', 'predicciones.csv'), index=False)
M.main(['x', 'rec'])
JR = json.load(open(os.path.join(OUT, 'impl2_rec.json')))
t = recdf[recdf['ventana'] == 'test']
for c in ('p', 'p_raw', 'b3'):
    chk('e2e rec: AUC ponderada %s vs fuerza bruta' % c, JR['prueba']['auc'][c],
        auc_fuerza_bruta(t[c].to_numpy(), t['y'].to_numpy(), t['w'].to_numpy()), 1e-12)
chk('e2e rec: O/E = Σw·y/Σp', JR['prueba']['o_e'], (t['w'] * t['y']).sum() / t['p'].sum(), 1e-12)
chk('e2e rec: pendiente vs statsmodels freq_weights', JR['prueba']['pendiente'],
    pendiente_sm(t['p'].to_numpy(), t['y'].to_numpy(), t['w'].to_numpy()), 1e-8)

print('\n== RESUMEN: %d comprobaciones OK, %d fallos' % (N_OK[0], len(FALLOS)))
for f in FALLOS:
    print('  FALLO', f)
sys.exit(1 if FALLOS else 0)