"""A1 · base del piloto (sin modelos): contratos de carburantes y electricidad 2022-2025, todas las fuentes,
unidad = adjudicación (contrato o lote), universo NOT es_copia (copia_de OR copia_otra_fuente OR NOT cuenta en vínculos:
definición de la web), importe comparable sin IVA (imp_adj_eur_comparable: iva_base 'sin'/'sin_medido'), reglas v1 de
mercado + exclusión v2, cantidad/producto/capacidad por reglas v2. Salida: /p/priv/a1_base.parquet (objetos: carpeta 700)."""
import sys, time
sys.path.insert(0, '/p/scripts')
import duckdb, pyarrow as pa, pyarrow.parquet as pq
from reglas import MERCADOS
from reglas_v2 import cantidad_principal, producto2, excluido_v2, cantidades2
from comun_a1 import provincia_ine
t0 = time.time()
con = duckdb.connect(); con.execute("SET memory_limit='4500MB'; SET threads=4; SET temp_directory='/p/tmp/duck'; SET preserve_insertion_order=false")
D = '/d/unified_20261018'; I = D + '.identidad/'
con.execute(f"CREATE VIEW u AS SELECT * FROM read_parquet('{D}/*/*/*.parquet', hive_partitioning=1, union_by_name=1)")
con.execute(f"CREATE VIEW vf AS SELECT source, id_canonical, cuenta FROM read_parquet('{D}.vinculos/filas.parquet')")
partes = []
for m in ('carburantes', 'electricidad'):
    r = MERCADOS[m]
    partes.append(f"""SELECT '{m}' mercado, CASE WHEN {r['cpv']} THEN 'cpv' ELSE 'texto' END via, * FROM b
      WHERE {r['cpv']} OR (regexp_matches(o, '{r['txt']}') AND NOT regexp_matches(o, '{r['excl']}'))""")
con.execute("""CREATE TEMP TABLE b AS SELECT u.source, u.id_canonical, licitacion_id, lote_id, coalesce(es_menor, FALSE) menor, anio_ref,
  cpv_code cpv, objeto, lower(strip_accents(coalesce(objeto,''))) o, imp_adj_eur_comparable imp, imp_adj_eur, iva_base, rol_am, nuts,
  provincia, ccaa, organo_nif, organo_nombre, dir3, fecha_adj, fecha_formalizacion, fecha_pub, duracion_meses, procedimiento,
  num_ofertas, url_origen, copia_de, correccion_imp_adj, v.cuenta,
  (copia_de IS NOT NULL OR coalesce(correccion_imp_adj = 'copia_otra_fuente', FALSE) OR coalesce(NOT v.cuenta, FALSE)) es_copia
 FROM u LEFT JOIN vf v ON v.source = u.source AND v.id_canonical = u.id_canonical
 WHERE anio_ref BETWEEN 2022 AND 2025 AND unidad = 'adjudicacion'""")
con.execute("CREATE TEMP TABLE m AS " + " UNION ALL ".join(partes))
print('mercados (con copias)', con.execute("SELECT mercado, es_copia, count(*) FROM m GROUP BY ALL ORDER BY ALL").fetchall(), round(time.time()-t0), flush=True)
con.execute(f"""CREATE TEMP TABLE m2 AS SELECT m.* EXCLUDE (o), o.organo_clave, coalesce(t.tipo,'sin_tipo') tipo_organo
 FROM m LEFT JOIN (SELECT organo_nif, organo_nombre, dir3, any_value(organo_clave) organo_clave FROM read_parquet('{I}organos_filas.parquet') GROUP BY ALL) o
   ON o.organo_nif IS NOT DISTINCT FROM m.organo_nif AND o.organo_nombre IS NOT DISTINCT FROM m.organo_nombre AND o.dir3 IS NOT DISTINCT FROM m.dir3
 LEFT JOIN read_parquet('{I}organos_tipo.parquet') t ON t.organo_clave = o.organo_clave""")
con.execute(f"""CREATE TEMP TABLE dl AS SELECT licitacion_id, bool_or(tipo_documento='pliego_tecnico') ppt, bool_or(tipo_documento='pliego_administrativo') pcap
 FROM read_parquet('{D}.documentos/documentos.parquet') WHERE licitacion_id IN (SELECT DISTINCT licitacion_id FROM m2 WHERE licitacion_id IS NOT NULL) GROUP BY 1""")
rows = con.execute("SELECT m2.*, coalesce(dl.ppt, FALSE) ppt, coalesce(dl.pcap, FALSE) pcap FROM m2 LEFT JOIN dl USING (licitacion_id)").fetch_arrow_table()
print('filas', rows.num_rows, round(time.time()-t0), flush=True)
obj = rows.column('objeto').to_pylist(); mk = rows.column('mercado').to_pylist()
nuts = rows.column('nuts').to_pylist(); prov = rows.column('provincia').to_pylist()
cv, cl, cn, ccap, prod, excl, pine = [], [], [], [], [], [], []
for o, m, n, p in zip(obj, mk, nuts, prov):
    clase = 'litros' if m == 'carburantes' else 'energia'
    c = cantidad_principal(o, clase)
    cv.append(c[0] if c else None); cl.append(c[1] if c else None); cn.append(c[2] if c else 0)
    ccap.append(any(x[3] for x in cantidades2(o, {clase})))
    prod.append(producto2(o) if m == 'carburantes' else 'electricidad')
    excl.append(excluido_v2(o, m)); pine.append(provincia_ine(n, p))
for nombre, col in [('cant_v2', cv), ('literal_v2', cl), ('n_cant_v2', cn), ('capacidad_v2', ccap), ('producto_v2', prod),
                    ('excluido_v2', excl), ('prov_ine', pine)]:
    rows = rows.append_column(nombre, pa.array(col))
pq.write_table(rows, '/p/priv/a1_base.parquet')
con.execute("CREATE VIEW a AS SELECT * FROM read_parquet('/p/priv/a1_base.parquet')")
for q in ["SELECT mercado, count(*), sum(es_copia::int), sum(excluido_v2::int), sum((cant_v2 IS NOT NULL)::int), sum(capacidad_v2::int), sum((imp IS NOT NULL)::int), sum((prov_ine IS NULL)::int) FROM a GROUP BY 1",
          "SELECT producto_v2, count(*) FILTER (WHERE NOT es_copia), count(*) FILTER (WHERE NOT es_copia AND cant_v2 IS NOT NULL) FROM a WHERE mercado='carburantes' GROUP BY 1 ORDER BY 2 DESC",
          "SELECT iva_base, count(*) FROM a WHERE NOT es_copia GROUP BY 1"]:
    print(con.execute(q).fetchall())
print('fin', round(time.time()-t0))