#!/usr/bin/env python3
"""SarcoMyokinEngine (사르코마이오킨엔진)
Sarcopenic obesity myokine composite biomarker hypothesis engine.

Stdlib only: sqlite3, json, csv, argparse, math, itertools, statistics.
"""
import argparse
import csv
import json
import math
import os
import random
import sqlite3
import sys
from itertools import combinations

DISCLAIMER = (
    "WARNING: 본 도구는 연구·참고용입니다. 임상 의사결정에 사용 금지. "
    "Recommended panels require pilot validation before grant submission. "
    "Cost estimates are illustrative."
)

ROOT = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(ROOT, "data")
DB_PATH = os.path.join(ROOT, "myokines.db")


def _print_disclaimer():
    print("=" * 78)
    print(DISCLAIMER)
    print("=" * 78)


# ---------- INIT / DB BUILD ----------

def _load_json(name):
    with open(os.path.join(DATA, name), "r", encoding="utf-8") as f:
        return json.load(f)


def _load_csv(name):
    rows = []
    with open(os.path.join(DATA, name), "r", encoding="utf-8") as f:
        rdr = csv.DictReader(f)
        for r in rdr:
            rows.append(r)
    return rows


def cmd_init(args):
    _print_disclaimer()
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()

    c.execute("""CREATE TABLE myokines (
        symbol TEXT PRIMARY KEY,
        aliases TEXT,
        primary_action TEXT,
        serum_measurable INTEGER,
        elisa_cost_usd REAL,
        uniprot TEXT
    )""")
    c.execute("""CREATE TABLE pubmed_evidence (
        symbol TEXT PRIMARY KEY,
        so_total INTEGER,
        muscle_loss INTEGER,
        adiposity INTEGER,
        aging INTEGER
    )""")
    c.execute("""CREATE TABLE gtex_expression (
        symbol TEXT PRIMARY KEY,
        muscle_skeletal_tpm REAL,
        adipose_subcutaneous_tpm REAL,
        adipose_visceral_tpm REAL,
        liver_tpm REAL
    )""")
    c.execute("""CREATE TABLE hpa_expression (
        symbol TEXT PRIMARY KEY,
        muscle_protein_level TEXT,
        adipose_protein_level TEXT,
        detection_method TEXT
    )""")
    c.execute("""CREATE TABLE pathways (
        symbol TEXT,
        pathway TEXT
    )""")
    c.execute("""CREATE TABLE cooccurrence (
        symbol_a TEXT,
        symbol_b TEXT,
        n_papers INTEGER
    )""")

    myokines = _load_json("myokines.json")["myokines"]
    for m in myokines:
        c.execute(
            "INSERT INTO myokines VALUES (?,?,?,?,?,?)",
            (m["symbol"], ",".join(m["aliases"]), m["primary_action"],
             1 if m["serum_measurable"] else 0, m["elisa_cost_usd"], m["uniprot"])
        )

    ev = _load_json("pubmed_evidence.json")
    for sym, vals in ev["evidence_counts"].items():
        c.execute(
            "INSERT INTO pubmed_evidence VALUES (?,?,?,?,?)",
            (sym, vals["so_total"], vals["muscle_loss"], vals["adiposity"], vals["aging"])
        )

    co = ev["co_occurrence_matrix"]
    for a, partners in co.items():
        if a == "_note":
            continue
        for b, n in partners.items():
            c.execute("INSERT INTO cooccurrence VALUES (?,?,?)", (a, b, n))

    for r in _load_csv("gtex_expression.csv"):
        c.execute(
            "INSERT INTO gtex_expression VALUES (?,?,?,?,?)",
            (r["myokine"], float(r["muscle_skeletal_tpm"]),
             float(r["adipose_subcutaneous_tpm"]), float(r["adipose_visceral_tpm"]),
             float(r["liver_tpm"]))
        )

    for r in _load_csv("hpa_expression.csv"):
        c.execute(
            "INSERT INTO hpa_expression VALUES (?,?,?,?)",
            (r["myokine"], r["muscle_protein_level"],
             r["adipose_protein_level"], r["detection_method"])
        )

    pw = _load_json("reactome_pathways.json")["myokine_to_pathways"]
    for sym, paths in pw.items():
        for p in paths:
            c.execute("INSERT INTO pathways VALUES (?,?)", (sym, p))

    conn.commit()

    # Stats
    c.execute("SELECT COUNT(*) FROM myokines")
    n_myo = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM pathways")
    n_pw = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM cooccurrence")
    n_co = c.fetchone()[0]
    conn.close()

    print(f"[init] DB built at {DB_PATH}")
    print(f"[init] {n_myo} myokines, {n_pw} pathway links, {n_co} co-occurrence pairs.")


# ---------- HELPERS ----------

def _ensure_db():
    if not os.path.exists(DB_PATH):
        print("[warn] DB not found — running init() automatically.")
        cmd_init(None)


def _fetch_all_myokines(conn):
    c = conn.cursor()
    c.execute("SELECT symbol FROM myokines")
    return [r[0] for r in c.fetchall()]


def _myokine_expression(conn, sym):
    c = conn.cursor()
    c.execute(
        "SELECT muscle_skeletal_tpm, adipose_subcutaneous_tpm, adipose_visceral_tpm, liver_tpm "
        "FROM gtex_expression WHERE symbol=?", (sym,))
    return c.fetchone()


def _myokine_pathways(conn, sym):
    c = conn.cursor()
    c.execute("SELECT pathway FROM pathways WHERE symbol=?", (sym,))
    return set(r[0] for r in c.fetchall())


def _myokine_meta(conn, sym):
    c = conn.cursor()
    c.execute("SELECT serum_measurable, elisa_cost_usd, primary_action FROM myokines WHERE symbol=?", (sym,))
    return c.fetchone()


def _myokine_evidence(conn, sym):
    c = conn.cursor()
    c.execute("SELECT so_total FROM pubmed_evidence WHERE symbol=?", (sym,))
    r = c.fetchone()
    return r[0] if r else 0


def _pair_cooccur(conn, a, b):
    c = conn.cursor()
    c.execute(
        "SELECT n_papers FROM cooccurrence WHERE (symbol_a=? AND symbol_b=?) OR (symbol_a=? AND symbol_b=?)",
        (a, b, b, a))
    r = c.fetchone()
    return r[0] if r else 0


# ---------- SCORING ----------

def _tissue_diversity_score(triplet_expr):
    """Higher = more tissue-diverse expression across the triplet.
    triplet_expr: list of (muscle, adipose_sc, adipose_v, liver) tuples.
    Use the per-tissue total, normalized by sum across tissues."""
    totals = [0.0, 0.0, 0.0, 0.0]
    for tup in triplet_expr:
        for i, v in enumerate(tup):
            totals[i] += v
    s = sum(totals)
    if s == 0:
        return 0.0
    p = [v / s for v in totals]
    # Shannon entropy normalized to [0,1] over 4 tissues
    h = 0.0
    for x in p:
        if x > 0:
            h -= x * math.log(x)
    return h / math.log(4)


def _mechanism_nonoverlap(pathway_sets):
    """Jaccard-style non-overlap for k sets. Higher = less mechanistic overlap."""
    if not pathway_sets:
        return 0.0
    union = set().union(*pathway_sets)
    if not union:
        return 0.0
    inter = set.intersection(*pathway_sets) if len(pathway_sets) > 1 else pathway_sets[0]
    overlap = len(inter) / len(union)
    return 1.0 - overlap


def _measurability(metas):
    """metas: list of (serum_measurable, cost, action). Penalize biopsy-only."""
    score = 0.0
    for m in metas:
        score += 1.0 if m[0] == 1 else 0.3
    return score / len(metas)


def _unexplored_bonus(conn, triplet):
    """Lower SO co-occurrence among the triplet -> larger unexplored bonus."""
    pairs = list(combinations(triplet, 2))
    co_total = sum(_pair_cooccur(conn, a, b) for a, b in pairs)
    # Map: 0 papers -> 1.0, ~30 papers -> ~0.5, 100+ -> ~0.0
    return 1.0 / (1.0 + co_total / 15.0)


def _composite_score(tissue_div, mech_nonoverlap, measurability, unexplored):
    return (0.30 * tissue_div +
            0.30 * mech_nonoverlap +
            0.20 * measurability +
            0.20 * unexplored)


def _sample_triplets(myokines, k=3, n_samples=300, seed=42):
    """Sample without replacement from C(n,k) avoiding full enumeration."""
    rng = random.Random(seed)
    seen = set()
    out = []
    n = len(myokines)
    max_possible = math.comb(n, k)
    target = min(n_samples, max_possible)
    while len(out) < target:
        chosen = tuple(sorted(rng.sample(range(n), k)))
        if chosen in seen:
            continue
        seen.add(chosen)
        out.append(tuple(myokines[i] for i in chosen))
    return out


# ---------- RANK TRIPLETS ----------

def cmd_rank_triplets(args):
    _print_disclaimer()
    _ensure_db()
    conn = sqlite3.connect(DB_PATH)

    myos = _fetch_all_myokines(conn)
    samples = _sample_triplets(myos, k=3, n_samples=300, seed=42)

    scored = []
    for trip in samples:
        exprs = [_myokine_expression(conn, s) for s in trip]
        if any(e is None for e in exprs):
            continue
        pws = [_myokine_pathways(conn, s) for s in trip]
        metas = [_myokine_meta(conn, s) for s in trip]

        td = _tissue_diversity_score(exprs)
        mn = _mechanism_nonoverlap(pws)
        mz = _measurability(metas)
        ub = _unexplored_bonus(conn, trip)
        comp = _composite_score(td, mn, mz, ub)

        scored.append({
            "triplet": trip,
            "tissue_diversity": round(td, 4),
            "mechanism_nonoverlap": round(mn, 4),
            "measurability": round(mz, 4),
            "unexplored_bonus": round(ub, 4),
            "composite": round(comp, 4),
        })

    scored.sort(key=lambda x: x["composite"], reverse=True)
    top = scored[: args.top]

    print(f"\n[rank-triplets] Top {len(top)} of {len(scored)} sampled triplets")
    print("-" * 100)
    print(f"{'rank':<5}{'triplet':<48}{'tissue':<10}{'mech':<10}{'meas':<10}{'unexp':<10}{'COMP':<8}")
    print("-" * 100)
    for i, s in enumerate(top, 1):
        trip_str = ", ".join(s["triplet"])
        if len(trip_str) > 47:
            trip_str = trip_str[:44] + "..."
        print(f"{i:<5}{trip_str:<48}{s['tissue_diversity']:<10}{s['mechanism_nonoverlap']:<10}"
              f"{s['measurability']:<10}{s['unexplored_bonus']:<10}{s['composite']:<8}")

    conn.close()


# ---------- PANEL ----------

def _power_two_sample_t(n_per_group, effect_d, alpha=0.05):
    """Closed-form approximation of two-sample t-test power.
    Uses normal approximation: power = Phi(d * sqrt(n/2) - z_{1-alpha/2})."""
    # Standard normal CDF
    def Phi(x):
        return 0.5 * (1 + math.erf(x / math.sqrt(2)))

    z_alpha = 1.959963984540054  # alpha=0.05 two-sided
    if alpha != 0.05:
        # Approximate for other alpha via symmetric two-sided
        # Use Newton-like inversion of Phi; here we keep alpha=0.05 by default
        pass
    ncp = effect_d * math.sqrt(n_per_group / 2.0)
    return Phi(ncp - z_alpha)


def cmd_panel(args):
    _print_disclaimer()
    _ensure_db()
    conn = sqlite3.connect(DB_PATH)
    panel = [args.m1, args.m2, args.m3]

    # Validate
    valid = _fetch_all_myokines(conn)
    for m in panel:
        if m not in valid:
            print(f"[error] Unknown myokine: {m}")
            print(f"        Available: {', '.join(valid)}")
            conn.close()
            sys.exit(1)

    metas = [_myokine_meta(conn, s) for s in panel]
    exprs = [_myokine_expression(conn, s) for s in panel]
    pws = [_myokine_pathways(conn, s) for s in panel]

    print(f"\n[panel] {' + '.join(panel)}")
    print("=" * 78)

    # Per-myokine summary
    print("\n## Per-myokine summary")
    print(f"{'symbol':<14}{'serum?':<10}{'ELISA $':<10}{'muscle TPM':<12}{'adipose TPM':<14}{'#pathways':<10}")
    for s, meta, expr, pw in zip(panel, metas, exprs, pws):
        serum = "yes" if meta[0] == 1 else "biopsy"
        cost = meta[1]
        muscle_tpm = expr[0]
        adipose_tpm = (expr[1] + expr[2]) / 2.0
        print(f"{s:<14}{serum:<10}{cost:<10.0f}{muscle_tpm:<12.1f}{adipose_tpm:<14.1f}{len(pw):<10}")

    # Cost estimate
    total_elisa = sum(m[1] for m in metas)
    n_subjects = 60  # 30 SO + 30 control assumption
    n_replicates = 2
    # Per-subject cost: (kit_cost / 96 wells per kit) * n_replicates * 3 markers approximated
    per_subject_cost = sum((m[1] / 96.0) * n_replicates for m in metas)
    total_assay_cost = per_subject_cost * n_subjects
    print("\n## Cost estimate (illustrative; n=60 subjects, duplicate wells)")
    print(f"  ELISA kit list price (3 markers): ${total_elisa:,.0f}")
    print(f"  Per-subject reagent cost:         ${per_subject_cost:,.2f}")
    print(f"  Total assay reagent cost:         ${total_assay_cost:,.0f}")

    # Blood volume
    serum_per_assay_uL = 50.0  # typical sandwich ELISA
    serum_buffer_factor = 1.3  # losses, reruns
    total_serum_uL = serum_per_assay_uL * len(panel) * n_replicates * serum_buffer_factor
    whole_blood_factor = 2.5  # ~40% serum yield -> need 2.5x
    whole_blood_mL = (total_serum_uL * whole_blood_factor) / 1000.0
    print("\n## Blood volume per subject")
    print(f"  Serum needed (3 markers x dup x 1.3 buffer):  {total_serum_uL:.0f} uL")
    print(f"  Whole blood draw (2.5x for serum yield):      {whole_blood_mL:.1f} mL")
    print("  (Recommended: collect 1x SST tube ~5 mL to allow re-runs.)")

    # Statistical power table (closed-form two-sample t)
    print("\n## Statistical power table (two-sample t, alpha=0.05, two-sided)")
    print("  Detecting standardized mean difference (Cohen d) between SO vs control")
    print(f"  {'n/group':<10}{'d=0.3':<10}{'d=0.5':<10}{'d=0.8':<10}{'d=1.0':<10}")
    for n in [15, 20, 30, 40, 60, 80, 120]:
        row = f"  {n:<10}"
        for d in [0.3, 0.5, 0.8, 1.0]:
            pw = _power_two_sample_t(n, d)
            row += f"{pw:<10.3f}"
        print(row)

    # Mechanism overlap
    print("\n## Mechanism (pathway) coverage")
    union = set().union(*pws)
    inter = set.intersection(*pws) if pws else set()
    print(f"  Distinct pathways covered: {len(union)}")
    print(f"  Common across all 3:       {len(inter)}")
    if inter:
        print(f"  Shared pathways:           {', '.join(sorted(inter))}")
    nonov = _mechanism_nonoverlap(pws)
    print(f"  Mechanism non-overlap:     {nonov:.3f}  (1.0 = fully complementary)")

    # Co-occurrence note
    print("\n## SO literature co-occurrence (paper counts)")
    for a, b in combinations(panel, 2):
        n = _pair_cooccur(conn, a, b)
        flag = "(unexplored)" if n < 5 else ""
        print(f"  {a} + {b}: {n} {flag}")

    print("\n## Suggested experimental protocol")
    print("  1. Recruit n=30 SO patients (ASMM/BMI < 0.789 [F] / 1.027 [M]; BMI >= 25 kg/m^2)")
    print("     and n=30 BMI-matched non-sarcopenic obese controls.")
    print("  2. Fasting morning blood draw (>=5 mL SST, serum separation within 30 min).")
    print("  3. Aliquot serum at -80 C; assay all panel markers in single batch (avoid freeze-thaw).")
    print("  4. Phenotype: DXA body composition + handgrip strength + 4-m gait speed.")
    print("  5. Primary analysis: two-sample t (or Wilcoxon if non-normal) per marker;")
    print("     secondary: composite z-score, ROC AUC vs clinical SO definition.")
    print("  6. Multiple testing: Holm-Bonferroni across the 3 markers.")
    print("  7. Pilot validation REQUIRED before grant submission.")

    conn.close()


# ---------- STATS ----------

def cmd_stats(args):
    _print_disclaimer()
    _ensure_db()
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()

    c.execute("SELECT COUNT(*) FROM myokines")
    n_myo = c.fetchone()[0]
    c.execute("SELECT SUM(so_total) FROM pubmed_evidence")
    so_total = c.fetchone()[0]
    c.execute("SELECT symbol, so_total FROM pubmed_evidence ORDER BY so_total DESC LIMIT 5")
    top5 = c.fetchall()
    c.execute("SELECT symbol, so_total FROM pubmed_evidence ORDER BY so_total ASC LIMIT 5")
    bot5 = c.fetchall()
    c.execute("SELECT COUNT(DISTINCT pathway) FROM pathways")
    n_path = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM cooccurrence")
    n_co = c.fetchone()[0]
    c.execute("SELECT AVG(elisa_cost_usd), MIN(elisa_cost_usd), MAX(elisa_cost_usd) FROM myokines")
    avg_c, min_c, max_c = c.fetchone()

    print(f"\n[stats] Corpus summary")
    print("-" * 60)
    print(f"  Total myokines:              {n_myo}")
    print(f"  Total SO papers (synthetic): {so_total}")
    print(f"  Distinct Reactome pathways:  {n_path}")
    print(f"  SO co-occurrence pairs:      {n_co}")
    print(f"  ELISA cost: avg ${avg_c:.0f} / min ${min_c:.0f} / max ${max_c:.0f}")
    print("\n  Top 5 most-studied in SO:")
    for s, n in top5:
        print(f"    {s:<14} {n}")
    print("\n  Bottom 5 (most underexplored — candidates for novel hypotheses):")
    for s, n in bot5:
        print(f"    {s:<14} {n}")

    conn.close()


# ---------- MAIN ----------

def main():
    parser = argparse.ArgumentParser(
        prog="sarcomyokin",
        description="SarcoMyokinEngine — sarcopenic obesity myokine biomarker hypothesis engine. "
                    "WARNING: 연구·참고용 only. Not for clinical decisions."
    )
    sub = parser.add_subparsers(dest="cmd", required=True)

    sub.add_parser("init", help="Build SQLite DB from data/ files")

    p_rank = sub.add_parser("rank-triplets", help="Score and rank candidate triplet panels")
    p_rank.add_argument("--top", type=int, default=10, help="Number of top triplets to display")

    p_panel = sub.add_parser("panel", help="Detailed panel report (cost, blood volume, power)")
    p_panel.add_argument("m1")
    p_panel.add_argument("m2")
    p_panel.add_argument("m3")

    sub.add_parser("stats", help="Corpus statistics")

    args = parser.parse_args()
    if args.cmd == "init":
        cmd_init(args)
    elif args.cmd == "rank-triplets":
        cmd_rank_triplets(args)
    elif args.cmd == "panel":
        cmd_panel(args)
    elif args.cmd == "stats":
        cmd_stats(args)


if __name__ == "__main__":
    main()