"""
make_lite_db.py
Creates a deployment-ready stripped version of drugbank.db under 500MB.
Keeps ALL drugs, targets, enzymes, food_interactions, pathways.
Samples 500k interactions — priority drugs first, then sequential fill.

Run from the Meta hackathon folder:
    python make_lite_db.py
"""

import sqlite3
import os

SRC = "drugbank.db"
DST = "drugbank_lite.db"

PRIORITY_DRUGS = [
    "Amlodipine", "Lisinopril", "Atenolol", "Hydrochlorothiazide",
    "Ramipril", "Metoprolol", "Valsartan", "Losartan", "Nifedipine",
    "Hydroflumethiazide", "Furosemide", "Spironolactone",
    "Metformin", "Glibenclamide", "Sitagliptin", "Empagliflozin",
    "Dapagliflozin", "Pioglitazone", "Glipizide",
    "Digoxin", "Carvedilol", "Bisoprolol", "Enalapril",
    "Sacubitril", "Ivabradine", "Torsemide", "Bumetanide",
    "Methotrexate", "Hydroxychloroquine", "Sulfasalazine",
    "Leflunomide", "Prednisolone",
    "Salbutamol", "Budesonide", "Fluticasone", "Salmeterol",
    "Montelukast", "Ipratropium", "Theophylline",
    "Valproate", "Carbamazepine", "Lamotrigine", "Levetiracetam",
    "Phenytoin", "Topiramate", "Oxcarbazepine", "Clonazepam",
    "Levothyroxine", "Liothyronine",
    "Sertraline", "Fluoxetine", "Escitalopram", "Venlafaxine",
    "Amitriptyline", "Mirtazapine", "Duloxetine", "Paroxetine",
    "Omeprazole", "Lansoprazole", "Pantoprazole",
    "Clarithromycin", "Amoxicillin", "Metronidazole",
    "Warfarin", "Apixaban", "Rivaroxaban", "Dabigatran",
    "Amiodarone", "Flecainide", "Sotalol",
    "Aspirin", "Ibuprofen", "Atorvastatin", "Simvastatin",
    "Clopidogrel", "Heparin", "Codeine", "Tramadol",
]

TARGET_INTERACTIONS = 500000

if os.path.exists(DST):
    os.remove(DST)
    print(f"Removed existing {DST}")

src = sqlite3.connect(SRC)
dst = sqlite3.connect(DST)
dst.execute("PRAGMA journal_mode=OFF")
dst.execute("PRAGMA synchronous=OFF")
dst.execute("PRAGMA cache_size=100000")
src_cursor = src.cursor()

def get_schema(table):
    src_cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name=?", (table,))
    row = src_cursor.fetchone()
    return row[0] if row else None

def copy_full_table(table):
    print(f"Copying ALL from {table}...", end=" ", flush=True)
    schema = get_schema(table)
    if not schema:
        print("not found.")
        return
    dst.execute(schema)
    src_cursor.execute(f"SELECT * FROM {table}")
    rows = src_cursor.fetchall()
    if rows:
        ph = ",".join(["?"] * len(rows[0]))
        dst.executemany(f"INSERT INTO {table} VALUES ({ph})", rows)
        dst.commit()
    print(f"{len(rows)} rows.")

def copy_interactions():
    print("Copying interactions...")
    schema = get_schema("interactions")
    if not schema:
        print("  Not found.")
        return
    dst.execute(schema)

    seen = set()
    total = 0

    # Step 1: priority drugs
    print("  Step 1: priority drugs...", flush=True)
    for drug in PRIORITY_DRUGS:
        src_cursor.execute(
            "SELECT * FROM interactions WHERE drug1_name LIKE ? OR drug2_name LIKE ? LIMIT 5000",
            (drug, drug)
        )
        rows = src_cursor.fetchall()
        new_rows = [r for r in rows if (r[0], r[2]) not in seen]
        for r in new_rows:
            seen.add((r[0], r[2]))
        if new_rows:
            ph = ",".join(["?"] * len(new_rows[0]))
            dst.executemany(f"INSERT INTO interactions VALUES ({ph})", new_rows)
            total += len(new_rows)
    dst.commit()
    print(f"  Priority done: {total} rows")

    # Step 2: sequential fill to TARGET
    remaining = TARGET_INTERACTIONS - total
    if remaining > 0:
        print(f"  Step 2: filling {remaining} more rows sequentially...", flush=True)
        BATCH = 5000
        offset = 0
        added = 0
        while added < remaining:
            src_cursor.execute("SELECT * FROM interactions LIMIT ? OFFSET ?", (BATCH, offset))
            rows = src_cursor.fetchall()
            if not rows:
                break
            new_rows = []
            for r in rows:
                if added >= remaining:
                    break
                key = (r[0], r[2])
                if key not in seen:
                    seen.add(key)
                    new_rows.append(r)
                    added += 1
            if new_rows:
                ph = ",".join(["?"] * len(new_rows[0]))
                dst.executemany(f"INSERT INTO interactions VALUES ({ph})", new_rows)
            offset += BATCH
            if added % 50000 == 0 and added > 0:
                dst.commit()
                print(f"    ...{total + added} so far", flush=True)
        dst.commit()
        total += added
        print(f"  Step 2 done: added {added} rows")

    print(f"  Total interactions: {total}")

for table in ["drugs", "targets", "enzymes", "food_interactions", "pathways"]:
    copy_full_table(table)

copy_interactions()
src.close()

print("Vacuuming...", end=" ", flush=True)
dst.execute("VACUUM")
dst.close()
print("done.")

size_mb = os.path.getsize(DST) / 1e6
print(f"\nResult: drugbank_lite.db = {size_mb:.1f} MB")
if size_mb > 900:
    print("Still large — HF LFS should handle it fine up to 5GB.")
else:
    print("Good to push to HuggingFace.")