"""SQL provider example — extract from a simulated invitroDB and tensorize.

This shows the provider-aware path: SQL source tables are preserved as
Arrow tables, keys/lineage are tracked, and pushdown filters reduce what's
pulled before canonicalization.
"""

import sqlite3
import tempfile
import os
import warnings
warnings.filterwarnings("ignore")

from toxarrow.adapters import SqlProviderAdapter, SqlTableMapping
from toxarrow.compile.arrow_store import ArrowStudyStore
from toxarrow.schemas.tensor import DenseGridSpec
from toxarrow.materialize.materializer import TensorMaterializer

# -- 1. Build a synthetic SQL database (simulating invitroDB / ToxCast schema) --
db_path = os.path.join(tempfile.mkdtemp(), "invitro.sqlite")
conn = sqlite3.connect(db_path)
conn.executescript("""
    CREATE TABLE chemical (chid TEXT PRIMARY KEY, chnm TEXT);
    CREATE TABLE assay (aeid TEXT PRIMARY KEY, aenm TEXT, assay_family TEXT);
    CREATE TABLE mc5 (
        m5id INTEGER PRIMARY KEY,
        chid TEXT, aeid TEXT,
        modl_acc REAL, modl_ga REAL, hitc INTEGER,
        FOREIGN KEY (chid) REFERENCES chemical(chid),
        FOREIGN KEY (aeid) REFERENCES assay(aeid)
    );

    INSERT INTO chemical VALUES
        ('CHEM001','Bisphenol A'),
        ('CHEM002','Genistein'),
        ('CHEM003','17b-Estradiol'),
        ('CHEM004','Atrazine');

    INSERT INTO assay VALUES
        ('A01','ER_alpha_agonism','nuclear_receptor'),
        ('A02','AR_antagonism','nuclear_receptor'),
        ('A03','Mitochondrial_tox','cytotoxicity');

    -- Only active hits (hitc=1) with measurable AC50
    INSERT INTO mc5 VALUES
        (1,'CHEM001','A01',0.45,85.0,1),
        (2,'CHEM001','A03',12.0,40.0,1),
        (3,'CHEM002','A01',0.08,95.0,1),
        (4,'CHEM002','A02',2.4,60.0,1),
        (5,'CHEM003','A01',0.002,100.0,1),
        (6,'CHEM003','A02',0.15,75.0,1),
        (7,'CHEM004','A03',45.0,25.0,1);
    -- CHEM001 has NO hit for A02 (inactive in AR assay)
    -- CHEM004 has hits only in A03
""")
conn.commit()

# -- 2. Define provider-aware mappings --
mappings = [
    SqlTableMapping(
        table="chemical", entity="unit", primary_key="chid",
        column_map={"unit_id": "chid", "label": "chnm"},
        order_by="chid",
    ),
    SqlTableMapping(
        table="assay", entity="endpoint", primary_key="aeid",
        column_map={
            "endpoint_id": "aeid", "endpoint_name": "aenm",
            "endpoint_family": "assay_family",
        },
        order_by="aeid",
    ),
    # Pushdown: only pull hitc=1 rows (active hits)
    SqlTableMapping(
        table="mc5", entity="observation", primary_key="m5id",
        column_map={
            "unit_id": "chid", "endpoint_id": "aeid",
            "value": "modl_acc",
        },
        where="hitc = 1",
        order_by="chid, aeid",
    ),
]

adapter = SqlProviderAdapter(
    conn, mappings, study_id="invitro-example",
    provider_name="invitrodb-synthetic", organism_or_system="in_vitro",
)

# -- 3. Stage (source tables preserved) + canonicalize --
tables = adapter.stage()
print(f"Staged {len(tables)} source tables:")
for name, tbl in tables.items():
    print(f"  {name}: {tbl.num_rows} rows, columns={tbl.column_names}")

group = next(adapter.read())
print(f"\nCanonical: {len(group.units)} chemicals, {len(group.endpoints)} assays, "
      f"{len(group.observations)} active hits")

# -- 4. Compile to Arrow, build Chemical × Assay grid --
store = ArrowStudyStore.from_records(group)
import numpy as np

rows = store.observations.to_pylist()
chem_ids = sorted([u.unit_id for u in group.units])
assay_order = sorted([e.endpoint_id for e in group.endpoints])
assay_idx = {a: i for i, a in enumerate(assay_order)}

vals = np.full((len(chem_ids), len(assay_order)), np.nan, dtype=np.float32)
mask = np.ones((len(chem_ids), len(assay_order)), dtype=bool)
for r in rows:
    ci = chem_ids.index(r["unit_id"])
    ai = assay_idx[r["endpoint_id"]]
    v = r.get("value")
    if v is not None:
        vals[ci, ai] = float(v)
        mask[ci, ai] = False

print(f"\nChemical × Assay grid: shape={vals.shape}")
print(f"Missing cells (masked): {mask.sum()}")
for ci, cid in enumerate(chem_ids):
    hits = [f"{a}:{vals[ci,i]:.2f}" if not mask[ci,i] else f"{a}:---" for i,a in enumerate(assay_order)]
    print(f"  {cid}: {hits}")

# Verify known missing cells
assert mask[chem_ids.index("CHEM001"), assay_idx["A02"]] == True, "CHEM001 missing AR"
assert mask[chem_ids.index("CHEM004"), assay_idx["A01"]] == True, "CHEM004 missing ER"
assert mask[chem_ids.index("CHEM004"), assay_idx["A02"]] == True, "CHEM004 missing AR"

print("\nAll checks passed: missingness masks correct, source tables preserved")
conn.close()