"""LT1 · 30 · SEGUNDA IMPLEMENTACIÓN (DEFINICIONES_PREVIAS §6): pandas + pyarrow + numpy, sin DuckDB, desde los parquet.
Celdas: L05 × M_1183 y M_REP (SS); L11 × M_REP (SS frente a SSI sin salud); L12 × NM_EMER (SS y O).
Escribe interno/impl2_celdas.json y salidas/30_impl2.json (puntuales)."""
import json, math, time
import numpy as np, pandas as pd, pyarrow as pa, pyarrow.dataset as ds, pyarrow.compute as pc
T0 = time.time()
def log(*a): print(f"[{time.time()-T0:6.1f}s]", *a, flush=True)
D = '/d/unified_20261018'
COLS = ['id_canonical', 'unidad', 'vigente', 'es_contrato', 'copia_de', 'correccion_imp_adj', 'tipo_contrato', 'fecha_adj', 'es_menor',
        'imp_adj_eur_comparable', 'iva_base', 'adjudicatario_nif', 'adjudicatario_nif_valid', 'organo_nif', 'organo_nombre', 'dir3', 'tramitacion']
dset = ds.dataset(D, format='parquet', partitioning='hive')
NUL = '\x00'
of = pd.read_parquet(D + '.identidad/organos_filas.parquet', columns=['organo_nif', 'organo_nombre', 'dir3', 'organo_clave'])
for k in ('organo_nif', 'organo_nombre', 'dir3'): of[k] = of[k].astype(object).where(of[k].notna(), NUL)
assert not of.duplicated(['organo_nif', 'organo_nombre', 'dir3']).any()
ot = pd.read_parquet(D + '.identidad/organos_tipo.parquet', columns=['organo_clave', 'tipo', 'agente_secti'])
vf = pd.read_parquet(D + '.vinculos/filas.parquet', columns=['source', 'id_canonical', 'cuenta'])
vf_no = vf[vf['cuenta'] == False][['source', 'id_canonical']].copy(); vf_no['fuera'] = True; del vf
log('identidad y vínculos', len(of), len(ot), len(vf_no))

def leer(filtro, cols):
    t = dset.to_table(columns=cols + ['source'], filter=filtro)
    df = t.to_pandas(date_as_object=False)
    log('leídas', len(df))
    df = df[(df['unidad'] == 'adjudicacion') & (df['vigente'] == True) & (df['es_contrato'] == True) & df['copia_de'].isna()
            & (df['correccion_imp_adj'].fillna('') != 'copia_otra_fuente') & df['tipo_contrato'].isin(['Servicios', 'Suministros', 'Obras'])]
    df = df.merge(vf_no, on=['source', 'id_canonical'], how='left'); df = df[df['fuera'].isna()].drop(columns='fuera')
    for k in ('organo_nif', 'organo_nombre', 'dir3'): df[k] = df[k].astype(object).where(df[k].notna(), NUL)
    df = df.merge(of, on=['organo_nif', 'organo_nombre', 'dir3'], how='left'); df = df[df['organo_clave'].notna()]
    df = df.merge(ot, on='organo_clave', how='left')
    df['secti'] = df['agente_secti'].fillna(False).astype(bool); df['salud'] = (df['tipo'] == 'salud')
    df['grupo'] = np.where(df['tipo_contrato'] == 'Obras', 'O', np.where(df['secti'], 'SSI', 'SS'))
    df['mi'] = (df['fecha_adj'].dt.year - 2012) * 12 + df['fecha_adj'].dt.month - 1
    return df

# ---------- menores 2016-2021, fuente a fuente (la repetición y el acumulado son dentro de la fuente)
import os
FUENTES = sorted(p.split('=', 1)[1] for p in os.listdir(D) if p.startswith('source='))
partes = []
for fu in FUENTES:
    f = (ds.field('source') == fu) & (ds.field('es_menor') == True) & (ds.field('fecha_adj') >= pa.scalar(pd.Timestamp('2016-01-01').date())) & (ds.field('fecha_adj') <= pa.scalar(pd.Timestamp('2021-12-31').date()))
    m = leer(f, [c for c in COLS if c != 'tramitacion'])
    if len(m) == 0: continue
    ok = (m['imp_adj_eur_comparable'] > 0) & m['iva_base'].isin(['sin', 'sin_medido']) & ~((m['source'] == 'ARAGON') & (m['iva_base'] == 'sin_medido'))
    m['imp'] = np.where(ok, np.round(m['imp_adj_eur_comparable'].astype(float), 2), np.nan)
    m['nif'] = m['adjudicatario_nif'].where(m['adjudicatario_nif_valid'] == True)
    m['anio'] = m['fecha_adj'].dt.year
    m = m[['id_canonical', 'source', 'organo_clave', 'grupo', 'salud', 'fecha_adj', 'mi', 'anio', 'imp', 'nif']]
    m = m.sort_values(['source', 'organo_clave', 'grupo', 'nif', 'anio', 'fecha_adj', 'id_canonical'], kind='mergesort')
    conn = m['nif'].notna()
    g = m[conn].groupby(['source', 'organo_clave', 'grupo', 'nif', 'anio'], sort=False)
    m['rep'] = np.nan; m.loc[conn, 'rep'] = (g.cumcount() > 0).astype(float)
    m['cent'] = np.where(m['imp'].notna(), np.round(m['imp'].fillna(0) * 100), 0).astype(np.int64)
    m['cum'] = np.nan; m.loc[conn, 'cum'] = m[conn].groupby(['source', 'organo_clave', 'grupo', 'nif', 'anio'], sort=False)['cent'].cumsum()
    tope = np.where(m['grupo'] == 'O', 4000000, 1500000)
    m['sobre'] = np.where(conn & m['imp'].notna(), (m['cum'] > tope).astype(float), np.nan)
    partes.append(m.groupby(['organo_clave', 'grupo', 'source', 'mi', 'salud'], sort=True).agg(n=('id_canonical', 'size'), rep_num=('rep', 'sum'), rep_den=('rep', 'count'),
                                                                                            t_num=('sobre', 'sum'), t_den=('sobre', 'count')).reset_index())
    log(fu, len(m)); del m, g
cel = pd.concat(partes, ignore_index=True); del partes
log('celdas menores', len(cel))

def mi_(a, mth): return (a - 2012) * 12 + mth - 1
SEMS = [(0, 5), (6, 11), (12, 17), (19, 24), (25, 30), (31, 36)]
def wins():
    out = []
    for cen in (0, -6, 6):
        out.append((list(range(18 + cen - 12, 18 + cen)), list(range(18 + cen + 1, 18 + cen + 13))))
    return out

def brazo_men(d, grupo, num, den, sin_salud=False):
    c = cel[(cel['grupo'] == grupo) & (cel['mi'] >= d - 18) & (cel['mi'] <= d + 18)]
    if sin_salud: c = c[c['salud'] == False]
    tot = c.groupby(['organo_clave', 'source'])['n'].sum().reset_index()
    tot['pl'] = (tot['source'] != 'PLACSP').astype(int)
    tot = tot.sort_values(['organo_clave', 'n', 'pl', 'source'], ascending=[True, False, True, True], kind='mergesort').drop_duplicates('organo_clave')
    c = c.merge(tot[['organo_clave', 'source']], on=['organo_clave', 'source'])
    return matriz(c, d, num, den)

def matriz(c, d, num, den):
    A = {}
    for o, sub in c.groupby('organo_clave'):
        M = np.zeros((37, 2))
        for k, a, b in zip(sub['mi'] - (d - 18), sub[num], sub[den]): M[int(k), 0] += a; M[int(k), 1] += b
        if all(M[a:b + 1, 1].sum() > 0 for a, b in SEMS): A[o] = M
    return A

def puntual(A):
    out = []
    for b, a in wins():
        nb = sum(M[b, 0].sum() for M in A.values()); db = sum(M[b, 1].sum() for M in A.values())
        na = sum(M[a, 0].sum() for M in A.values()); da = sum(M[a, 1].sum() for M in A.values())
        out.append(math.log(na / da) - math.log(nb / db))
    return out

def volcar(A):
    return {o: [[float(M[w[0]].sum(0)[j]) for w in wins() for j in (0, 1)], [float(M[w[1]].sum(0)[j]) for w in wins() for j in (0, 1)]] for o, M in A.items()}

VOL, RES = {}, {}
def guarda(clave, AT, AC=None):
    VOL[clave] = {'T': volcar(AT)} | ({'C': volcar(AC)} if AC is not None else {})
    DT = puntual(AT)
    if AC is not None:
        DC = puntual(AC); Ds = [a - b for a, b in zip(DT, DC)]
    else: Ds = DT
    RES[clave] = dict(n_T=len(AT), n_C=None if AC is None else len(AC), delta=Ds[0], K=Ds[0] - 0.5 * (Ds[1] + Ds[2]), placebo_menos6=Ds[1], placebo_mas6=Ds[2])
    log(clave, RES[clave])

d = mi_(2018, 3)
guarda('L05|M_1183|SS', brazo_men(d, 'SS', 't_num', 't_den'))
guarda('L05|M_REP|SS', brazo_men(d, 'SS', 'rep_num', 'rep_den'))
d = mi_(2020, 2)
guarda('L11|M_REP|SS vs SSI sin salud', brazo_men(d, 'SS', 'rep_num', 'rep_den'), brazo_men(d, 'SSI', 'rep_num', 'rep_den', sin_salud=True))
del cel
# ---------- no menores PLACSP, emergencia
f = (ds.field('source') == 'PLACSP') & (ds.field('fecha_adj') >= pa.scalar(pd.Timestamp('2018-09-01').date())) & (ds.field('fecha_adj') <= pa.scalar(pd.Timestamp('2021-09-30').date()))
n = leer(f, COLS)
n = n[n['es_menor'].fillna(False) == False]
n['emer_num'] = (n['tramitacion'] == 'emergencia').astype(float); n['emer_den'] = n['tramitacion'].notna().astype(float)
d = mi_(2020, 3)
for gr in ('SS', 'O'):
    c = n[n['grupo'] == gr].groupby(['organo_clave', 'mi']).agg(emer_num=('emer_num', 'sum'), emer_den=('emer_den', 'sum')).reset_index()
    guarda(f'L12|NM_EMER|{gr}', matriz(c, d, 'emer_num', 'emer_den'))
json.dump(VOL, open('/w/interno/impl2_celdas.json', 'w'))
json.dump(RES, open('/w/salidas/30_impl2.json', 'w'), indent=1)
log('fin')