"""SOB1 · sellar.py — sobres sellados de predicción para expedientes CNMC abiertos (protocolo: PROTOCOLO_SOBRES.md).

Tres órdenes (Python 3 + pyarrow; sin red, sin base de datos):

  sellar  --exp S0004-25 --config cfg.json [--seco]
      Lee salidas/universos/<exp>.parquet, filtra el universo de la config, puntúa cada lote con los hechos elegidos,
      aplica la regla de decisión fijada y escribe sobre_<exp>.json (JSON canónico) + .md5 + copia .original.json (444).
      Con --seco escribe en tmp/seco/ y no pone 444: sirve para probar sin sellar nada.
  verificar --sobre ruta/sobre_<exp>.json
      Recalcula el md5 y lo compara con el .md5 de al lado.
  abrir   --sobre ruta/sobre_<exp>.json --verdad verdad.tsv [--seco]
      Verdad = TSV con una columna id_lic (md5 de source|licitacion_id) o id_lote de los lotes que cite la resolución;
      o la palabra ARCHIVO (sin infracción). Mide precisión, recall, F1, tasa de señalados en R y fuera, AUC (Mann-Whitney
      con empates) con IC bootstrap por licitación, y aplica el criterio de acierto escrito en el sobre.

La config (lo decide el orquestador ANTES de sellar) tiene la forma:
{
 "universo": "en_mercado AND en_periodo",            # expresión sobre columnas booleanas del parquet (and/or/not, paréntesis)
 "hechos": [ {"col": "d1_pct_s3", "sentido": "+", "tipo": "percentil"},   # ya en [0,1]
             {"col": "ou1_unica", "sentido": "+", "tipo": "binario", "esperado": "ou1_p"},  # residuo observado - esperado
             {"col": "rs1_hermana", "sentido": "+", "tipo": "binario"} ],
 "combinar": "media_disponibles",                    # o "max" o "primero" (solo el primer hecho)
 "min_hechos": 1,                                     # lotes con menos hechos disponibles: puntuación nula, no señalados
 "regla_lote": {"umbral": 0.8},                       # señalado = puntuación >= umbral
 "regla_expediente": {"tasa_referencia": 0.203, "fuente_referencia": "...", "margen": 0.25},
                                                       # predicción del expediente = tasa_señalados >= ref*(1+margen)
 "criterio_acierto": {"recall_min": 0.6, "auc_min": 0.65, "auc_lb_min": 0.5},
 "sesgo_etiqueta": "texto declarado"
}
"""
import argparse, hashlib, json, os, random, re, stat, sys, time

import pyarrow.parquet as pq

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
VERSION = 'sobres-v1'


# ----------------------------------------------------------------------------------------------- utilidades
def md5_bytes(b):
    return hashlib.md5(b).hexdigest()


def md5_fichero(p):
    return md5_bytes(open(p, 'rb').read())


def canon(obj):
    """JSON canónico: claves ordenadas, sin espacios, UTF-8, floats redondeados a 6 decimales."""
    def r(x):
        if isinstance(x, float):
            return None if x != x else round(x, 6)
        if isinstance(x, dict):
            return {k: r(v) for k, v in x.items()}
        if isinstance(x, list):
            return [r(v) for v in x]
        return x
    return json.dumps(r(obj), sort_keys=True, ensure_ascii=False, separators=(',', ':')).encode('utf-8')


_TOK = re.compile(r'\s*(\(|\)|and\b|or\b|not\b|[A-Za-z_][A-Za-z0-9_]*)\s*', re.I)


def compila_expr(expr, columnas):
    """Expresión booleana segura sobre columnas: solo nombres de columna, and/or/not y paréntesis."""
    pos, out = 0, []
    while pos < len(expr):
        m = _TOK.match(expr, pos)
        if not m:
            raise ValueError(f'expresión no válida cerca de: {expr[pos:pos+20]!r}')
        t = m.group(1)
        tl = t.lower()
        if tl in ('and', 'or', 'not', '(', ')'):
            out.append(f' {tl} ')
        else:
            if t not in columnas:
                raise ValueError(f'columna desconocida en el universo: {t}')
            out.append(f'bool(f[{t!r}])')
        pos = m.end()
    codigo = ''.join(out)
    return lambda f: eval(codigo, {'__builtins__': {'bool': bool}}, {'f': f})


def lee_universo(exp):
    p = os.path.join(BASE, 'salidas', 'universos', f'{exp}.parquet')
    t = pq.read_table(p)
    filas = t.to_pylist()
    return p, filas, t.column_names


def puntua(f, hechos, combinar, min_hechos):
    vals = []
    for h in hechos:
        v = f.get(h['col'])
        if v is None:
            vals.append(None)
            continue
        v = float(v)
        if h.get('tipo') == 'binario' and h.get('esperado'):
            e = f.get(h['esperado'])
            if e is None:
                vals.append(None)
                continue
            v = v - float(e)            # residuo en [-1, 1]
            v = (v + 1) / 2             # a [0, 1]
        if h.get('sentido', '+') == '-':
            v = 1 - v
        vals.append(v)
    disp = [v for v in vals if v is not None]
    if len(disp) < min_hechos:
        return None, vals
    if combinar == 'max':
        return max(disp), vals
    if combinar == 'primero':
        return vals[0], vals
    return sum(disp) / len(disp), vals


def escribe(p, b, solo_lectura=False):
    if os.path.exists(p):
        raise SystemExit(f'ya existe {p}: un sobre no se sobrescribe')
    with open(p, 'wb') as fh:
        fh.write(b)
    if solo_lectura:
        os.chmod(p, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)


# ----------------------------------------------------------------------------------------------- sellar
def sellar(a):
    cfg_bytes = open(a.config, 'rb').read()
    cfg = json.loads(cfg_bytes)
    p_univ, filas, cols = lee_universo(a.exp)
    filtro = compila_expr(cfg['universo'], set(cols))
    hechos = cfg['hechos']
    for c in cfg.get('secundarios', []):
        if c not in cols:
            raise SystemExit(f'secundario desconocido: {c}')
    for h in hechos:
        if h['col'] not in cols or (h.get('esperado') and h['esperado'] not in cols):
            raise SystemExit(f"hecho desconocido: {h}")
    umbral = float(cfg['regla_lote']['umbral'])
    lotes = []
    for f in filas:
        if not filtro(f):
            continue
        s, vals = puntua(f, hechos, cfg.get('combinar', 'media_disponibles'), int(cfg.get('min_hechos', 1)))
        lote = {'id_lote': f['id_lote'], 'id_lic': f['id_lic'], 'anio': f.get('anio'),
                'hechos': vals, 'puntuacion': s, 'senalado': bool(s is not None and s >= umbral)}
        if cfg.get('secundarios'):   # puntuaciones registradas, NO decisorias (valor bruto, sin sentido)
            lote['secundarios'] = [None if f.get(c) is None else (float(f[c]) if not isinstance(f[c], bool) else int(f[c]))
                                   for c in cfg['secundarios']]
        lotes.append(lote)
    lotes.sort(key=lambda x: x['id_lote'])
    n = len(lotes)
    n_p = sum(1 for x in lotes if x['puntuacion'] is not None)
    n_s = sum(1 for x in lotes if x['senalado'])
    re_ = cfg['regla_expediente']
    tasa = n_s / n_p if n_p else None
    pred_exp = bool(tasa is not None and tasa >= float(re_['tasa_referencia']) * (1 + float(re_.get('margen', 0))))
    ahora = time.strftime('%Y-%m-%d %H:%M:%S %Z')
    sobre = {
        'protocolo': VERSION, 'expediente': a.exp, 'sellado': ahora,
        'universo': {'fichero': os.path.basename(p_univ), 'md5': md5_fichero(p_univ), 'expresion': cfg['universo'],
                     'lotes': n, 'lotes_puntuados': n_p, 'licitaciones': len({x['id_lic'] for x in lotes})},
        'config': cfg, 'config_md5': md5_bytes(cfg_bytes), 'codigo_md5': md5_fichero(os.path.abspath(__file__)),
        'prediccion_expediente': {'tasa_senalados': tasa, 'regla': re_, 'prediccion': pred_exp,
                                  'texto': 'la resolución acreditará manipulación en este universo' if pred_exp
                                  else 'los hechos no distinguen este universo de su referencia'},
        'senalados': n_s, 'lotes': lotes,
        'se_abre': 'cuando la CNMC publique la resolución del expediente (o su archivo)',
    }
    b = canon(sobre)
    h = md5_bytes(b)
    dest = os.path.join(BASE, 'tmp', 'seco') if a.seco else os.path.join(BASE, 'sellados')
    os.makedirs(dest, exist_ok=True)
    pj = os.path.join(dest, f'sobre_{a.exp}.json')
    if a.seco:   # en seco se puede repetir: se borran las pruebas anteriores (nunca en sellados/)
        for q in (pj, pj + '.md5', pj.replace('.json', '.original.json')):
            if os.path.exists(q):
                os.chmod(q, 0o600); os.remove(q)
    escribe(pj, b)
    escribe(pj.replace('.json', '.original.json'), b, solo_lectura=not a.seco)
    escribe(pj + '.md5', f'{h}  {os.path.basename(pj)}  sellado {ahora}\n'.encode(), solo_lectura=not a.seco)
    print(json.dumps({'sobre': pj, 'md5': h, 'sellado': ahora, 'lotes': n, 'puntuados': n_p, 'senalados': n_s,
                      'tasa': tasa, 'prediccion_expediente': pred_exp, 'seco': a.seco}, ensure_ascii=False))


# ----------------------------------------------------------------------------------------------- verificar
def verificar(a):
    b = open(a.sobre, 'rb').read()
    h = md5_bytes(b)
    esperado = open(a.sobre + '.md5').read().split()[0]
    ok = h == esperado
    print(json.dumps({'sobre': a.sobre, 'md5': h, 'md5_sellado': esperado, 'coincide': ok}))
    return ok


# ----------------------------------------------------------------------------------------------- abrir
def auc(pos, neg):
    if not pos or not neg:
        return None
    neg = sorted(neg)
    import bisect
    s = 0.0
    for x in pos:
        lo = bisect.bisect_left(neg, x); hi = bisect.bisect_right(neg, x)
        s += lo + 0.5 * (hi - lo)
    return s / (len(pos) * len(neg))


def abrir(a):
    if not verificar(a):
        raise SystemExit('el md5 no coincide: el sobre no se abre')
    s = json.loads(open(a.sobre, 'rb').read())
    lotes = s['lotes']
    crit = s['config'].get('criterio_acierto', {})
    verdad_txt = open(a.verdad, encoding='utf-8').read().strip()
    out = {'expediente': s['expediente'], 'md5': md5_bytes(open(a.sobre, 'rb').read()), 'abierto': time.strftime('%Y-%m-%d %H:%M:%S %Z')}
    if verdad_txt.upper() == 'ARCHIVO':
        n_p = sum(1 for x in lotes if x['puntuacion'] is not None)
        out.update(resultado='archivo sin infracción', tasa_falsos_positivos=(s['senalados'] / n_p if n_p else None),
                   prediccion_expediente=s['prediccion_expediente']['prediccion'],
                   expediente_acierta=not s['prediccion_expediente']['prediccion'])
    else:
        ids = {l.strip().split('\t')[0] for l in verdad_txt.splitlines()[1:] if l.strip()}
        col = verdad_txt.splitlines()[0].strip().split('\t')[0]
        if col not in ('id_lic', 'id_lote'):
            raise SystemExit('la verdad debe tener cabecera id_lic o id_lote')
        # métricas sobre los lotes PUNTUADOS (los demás no tienen predicción); los citados sin puntuación se cuentan aparte
        R_todos = [x for x in lotes if x[col] in ids]
        R = [x for x in R_todos if x['puntuacion'] is not None]
        noR = [x for x in lotes if x[col] not in ids and x['puntuacion'] is not None]
        tp = sum(x['senalado'] for x in R); fp = sum(x['senalado'] for x in noR)
        fn = len(R) - tp
        prec = tp / (tp + fp) if tp + fp else None
        rec = tp / len(R) if R else None
        f1 = 2 * prec * rec / (prec + rec) if prec and rec else None
        pos = [x['puntuacion'] for x in R if x['puntuacion'] is not None]
        neg = [x['puntuacion'] for x in noR if x['puntuacion'] is not None]
        A = auc(pos, neg)
        # IC bootstrap por licitación (semilla fija)
        rnd = random.Random(20261022)
        por_lic = {}
        for x in lotes:
            if x['puntuacion'] is not None:
                por_lic.setdefault(x['id_lic'], []).append((x[col] in ids, x['puntuacion']))
        claves = sorted(por_lic)
        reps = []
        for _ in range(int(a.bootstrap)):
            mu = [por_lic[claves[rnd.randrange(len(claves))]] for _ in claves] if claves else []
            pp = [v for g in mu for (r, v) in g if r]; nn = [v for g in mu for (r, v) in g if not r]
            v = auc(pp, nn)
            if v is not None:
                reps.append(v)
        reps.sort()
        ic = [reps[int(0.025 * len(reps))], reps[int(0.975 * len(reps)) - 1]] if len(reps) >= 40 else None
        acierto = (rec is not None and rec >= crit.get('recall_min', 0.6))
        if 'controles_max' in crit:   # como D2: la tasa de señalados en los controles (referencia sellada) no pasa del tope
            acierto = acierto and float(s['config']['regla_expediente']['tasa_referencia']) <= float(crit['controles_max'])
        debil = (A is not None and ic is not None and A >= crit.get('auc_min', 0.65) and ic[0] > crit.get('auc_lb_min', 0.5))
        out.update(resultado='resolución con infracción', clave=col, verdad_citados=len(ids),
                   R_en_universo=len(R_todos), R_sin_puntuacion=len(R_todos) - len(R), R_puntuados=len(R), fuera_de_universo=len(ids - {x[col] for x in lotes}),
                   tp=tp, fp=fp, fn=fn, precision=prec, recall=rec, f1=f1,
                   tasa_senalados_R=(tp / len(R) if R else None), tasa_senalados_resto=(fp / len(noR) if noR else None),
                   auc=A, auc_ic95=ic, n_pos=len(pos), n_neg=len(neg),
                   veredicto='acierto' if acierto else ('acierto débil' if debil else 'fallo'),
                   prediccion_expediente=s['prediccion_expediente']['prediccion'],
                   expediente_acierta=bool(s['prediccion_expediente']['prediccion']),
                   sesgo_etiqueta=s['config'].get('sesgo_etiqueta'))
    dest = os.path.join(BASE, 'tmp', 'seco') if a.seco else os.path.join(BASE, 'abiertos')
    os.makedirs(dest, exist_ok=True)
    p = os.path.join(dest, f"apertura_{s['expediente']}.json")
    if os.path.exists(p) and not a.seco:
        raise SystemExit(f'ya existe {p}')
    open(p, 'w').write(json.dumps(out, ensure_ascii=False, indent=1))
    print(json.dumps(out, ensure_ascii=False))


def main():
    ap = argparse.ArgumentParser()
    sp = ap.add_subparsers(dest='orden', required=True)
    s1 = sp.add_parser('sellar'); s1.add_argument('--exp', required=True); s1.add_argument('--config', required=True)
    s1.add_argument('--seco', action='store_true')
    s2 = sp.add_parser('verificar'); s2.add_argument('--sobre', required=True)
    s3 = sp.add_parser('abrir'); s3.add_argument('--sobre', required=True); s3.add_argument('--verdad', required=True)
    s3.add_argument('--bootstrap', default=2000); s3.add_argument('--seco', action='store_true')
    a = ap.parse_args()
    for k in ('config', 'sobre', 'verdad'):   # rutas relativas: respecto a sobres/
        v = getattr(a, k, None)
        if v and not os.path.isabs(v):
            setattr(a, k, os.path.join(BASE, v))
    if a.orden == 'sellar':
        sellar(a)
    elif a.orden == 'verificar':
        sys.exit(0 if verificar(a) else 1)
    else:
        abrir(a)


if __name__ == '__main__':
    main()