import json, sys, sqlite3
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rcParams['figure.dpi'] = 110
matplotlib.rcParams['figure.facecolor'] = 'white'

# Find worktree root by looking for AGENTS.md marker
_cwd = Path('.').resolve()
if (_cwd / 'AGENTS.md').exists():
    REPO = _cwd
else:
    REPO = _cwd
    for _p in [_cwd] + list(_cwd.parents):
        if (_p / 'AGENTS.md').exists():
            REPO = _p
            break

CACHE = REPO / 'data' / 'forge_cache' / 'seaad'

# TARGET_GENES: 11-gene AD vulnerability set (from forge/seaad_analysis.py)
TARGET_GENES = [
    "TREM2", "GFAP", "SLC17A7", "PDGFRA", "PDGFRB",
    "APOE", "MAPT", "APP", "PSEN1", "TYROBP", "CLU"
]

def load(name, _cache=CACHE):
    """Load cached JSON, return empty dict if file missing."""
    _path = _cache / f'{name}.json'
    if _path.exists():
        return json.loads(_path.read_text())
    return {}

# Forge provenance: tool calls this session invoked
try:
    db = sqlite3.connect(str(REPO / 'scidex.db'))
    prov = pd.read_sql_query("""
        SELECT skill_id, status, COUNT(*) AS n_calls,
               ROUND(AVG(duration_ms),0) AS mean_ms,
               MIN(created_at) AS first_call
        FROM tool_calls
        WHERE created_at >= date('now','-1 day')
        GROUP BY skill_id, status
        ORDER BY n_calls DESC
    """, db)
    db.close()
    prov.rename(columns={'skill_id':'tool'}, inplace=True)
    prov['tool'] = prov['tool'].str.replace('tool_', '', regex=False)
    print(f'{len(prov)} tool-call aggregates from the last 24h of Forge provenance:')
    prov.head(20)
except Exception as e:
    print(f"Provenance query skipped: {e}")
    print("(tool_calls table may not exist in this worktree)")

anno = load('mygene_TREM2')  # probe one
ann_rows = []
for g in ['TREM2','GFAP','SLC17A7','PDGFRA','PDGFRB','APOE','MAPT','APP','PSEN1','TYROBP','CLU']:
    mg = load(f'mygene_{g}')
    hpa = load(f'hpa_{g}')
    ann_rows.append({
        'gene': g,
        'name': (mg.get('name') or '')[:55],
        'protein_class': ', '.join((hpa.get('protein_class') or [])[:2])[:55],
        'disease_involvement': ', '.join((hpa.get('disease_involvement') or [])[:2])[:55] if isinstance(hpa.get('disease_involvement'), list) else str(hpa.get('disease_involvement') or '')[:55],
        'ensembl_id': hpa.get('ensembl_id') or '',
    })
anno_df = pd.DataFrame(ann_rows)
anno_df

go_bp = load('enrichr_GO_Biological_Process')[:10]
go_df = pd.DataFrame(go_bp)[['term','p_value','odds_ratio','genes']]
go_df['p_value'] = go_df['p_value'].apply(lambda p: f'{p:.2e}')
go_df['odds_ratio'] = go_df['odds_ratio'].round(1)
go_df['term'] = go_df['term'].str[:60]
go_df['n_hits'] = go_df['genes'].apply(len)
go_df['genes'] = go_df['genes'].apply(lambda g: ', '.join(g))
go_df[['term','n_hits','p_value','odds_ratio','genes']]

# Visualize top GO BP enrichment (−log10 p-value bar chart)
import numpy as np
go_bp = load('enrichr_GO_Biological_Process')[:8]
terms = [t['term'][:45] for t in go_bp][::-1]
neglogp = [-np.log10(t['p_value']) for t in go_bp][::-1]
fig, ax = plt.subplots(figsize=(9, 4.5))
ax.barh(terms, neglogp, color='#4fc3f7')
ax.set_xlabel('-log10(p-value)')
ax.set_title('Top GO:BP enrichment for SEA-AD vulnerability gene set (Enrichr)')
ax.grid(axis='x', alpha=0.3)
plt.tight_layout()
plt.show()

cm = load('enrichr_CellMarker_Cell_Types')[:10]
cm_df = pd.DataFrame(cm)[['term','p_value','odds_ratio','genes']]
cm_df['genes'] = cm_df['genes'].apply(lambda g: ', '.join(g))
cm_df['p_value'] = cm_df['p_value'].apply(lambda p: f'{p:.2e}')
cm_df['odds_ratio'] = cm_df['odds_ratio'].round(1)
cm_df

ppi = load('string_network')
ppi_df = pd.DataFrame(ppi)
if not ppi_df.empty:
    ppi_df = ppi_df.sort_values('score', ascending=False)
    display_cols = [c for c in ['protein1','protein2','score','escore','tscore'] if c in ppi_df.columns]
    print(f'{len(ppi_df)} STRING edges among {len(set(list(ppi_df.protein1)+list(ppi_df.protein2)))} proteins')
    ppi_df[display_cols].head(20)
else:
    print('No STRING edges returned (API may be rate-limited)')

# Simple network figure using matplotlib (no networkx dep)
ppi = load('string_network')
if ppi:
    import math
    nodes = sorted({p for e in ppi for p in (e['protein1'], e['protein2'])})
    n = len(nodes)
    pos = {n_: (math.cos(2*math.pi*i/n), math.sin(2*math.pi*i/n)) for i, n_ in enumerate(nodes)}
    fig, ax = plt.subplots(figsize=(7, 7))
    for e in ppi:
        x1,y1 = pos[e['protein1']]; x2,y2 = pos[e['protein2']]
        ax.plot([x1,x2],[y1,y2], color='#888', alpha=0.3+0.5*e['score'], linewidth=0.5+2*e['score'])
    for name,(x,y) in pos.items():
        ax.scatter([x],[y], s=450, color='#ffd54f', edgecolors='#333', zorder=3)
        ax.annotate(name, (x,y), ha='center', va='center', fontsize=9, fontweight='bold', zorder=4)
    ax.set_aspect('equal'); ax.axis('off')
    ax.set_title(f'STRING physical PPI network ({len(ppi)} edges, score ≥ 0.4)')
    plt.tight_layout(); plt.show()

pw_rows = []
for g in ['TREM2','GFAP','SLC17A7','PDGFRA','PDGFRB','APOE','MAPT','APP','PSEN1','TYROBP','CLU']:
    pws = load(f'reactome_{g}')
    pw_rows.append({'gene': g, 'n_pathways': len(pws),
                    'top_pathway': (pws[0]['name'] if pws else '—')[:70]})
pw_df = pd.DataFrame(pw_rows).sort_values('n_pathways', ascending=False)
pw_df

from collections import Counter
ac = load('allen_celltypes_TREM2')  # same for any gene (not gene-filtered at API level)
ct = pd.DataFrame(ac.get('cell_types', []))
if not ct.empty:
    ct_display = ct.head(15)
else:
    ct_display = pd.DataFrame()
ct_display

ish_rows = []
for g in TARGET_GENES:
    ish = load(f'allen_ish_{g}')
    regions = ish.get('regions') or []
    ish_rows.append({
        'gene': g,
        'n_ish_regions': len(regions),
        'top_region': (regions[0].get('structure','') if regions else '—')[:45],
        'top_energy': round(regions[0].get('expression_energy',0), 2) if regions else None,
        'note': (ish.get('note') or '')[:60],
    })
ish_df = pd.DataFrame(ish_rows)
ish_df

hid = 'h-seaad-5b3cb8ea'
papers = load(f'pubmed_{hid}')
if papers:
    lit = pd.DataFrame(papers)[['year','journal','title','pmid']]
    lit['title'] = lit['title'].str[:80]
    lit['journal'] = lit['journal'].str[:30]
    lit.sort_values('year', ascending=False, inplace=True)
    display_df = lit
else:
    display_df = pd.DataFrame([{'note':'no PubMed results for this hypothesis query'}])
display_df

hid = 'h-seaad-51323624'
papers = load(f'pubmed_{hid}')
if papers:
    lit = pd.DataFrame(papers)[['year','journal','title','pmid']]
    lit['title'] = lit['title'].str[:80]
    lit['journal'] = lit['journal'].str[:30]
    lit.sort_values('year', ascending=False, inplace=True)
    display_df = lit
else:
    display_df = pd.DataFrame([{'note':'no PubMed results for this hypothesis query'}])
display_df

hid = 'h-seaad-7f15df4c'
papers = load(f'pubmed_{hid}')
if papers:
    lit = pd.DataFrame(papers)[['year','journal','title','pmid']]
    lit['title'] = lit['title'].str[:80]
    lit['journal'] = lit['journal'].str[:30]
    lit.sort_values('year', ascending=False, inplace=True)
    display_df = lit
else:
    display_df = pd.DataFrame([{'note':'no PubMed results for this hypothesis query'}])
display_df

hid = 'h-seaad-fa5ea82d'
papers = load(f'pubmed_{hid}')
if papers:
    lit = pd.DataFrame(papers)[['year','journal','title','pmid']]
    lit['title'] = lit['title'].str[:80]
    lit['journal'] = lit['journal'].str[:30]
    lit.sort_values('year', ascending=False, inplace=True)
    display_df = lit
else:
    display_df = pd.DataFrame([{'note':'no PubMed results for this hypothesis query'}])
display_df

hid = 'h-seaad-56fa6428'
papers = load(f'pubmed_{hid}')
if papers:
    lit = pd.DataFrame(papers)[['year','journal','title','pmid']]
    lit['title'] = lit['title'].str[:80]
    lit['journal'] = lit['journal'].str[:30]
    lit.sort_values('year', ascending=False, inplace=True)
    display_df = lit
else:
    display_df = pd.DataFrame([{'note':'no PubMed results for this hypothesis query'}])
display_df

from pathlib import Path
bundle_path = REPO / 'data/analysis_outputs/analysis-SEAAD-20260402/mechanistic_de/bundle.json'
if bundle_path.exists():
    mech_bundle = json.loads(bundle_path.read_text())
    print("Mechanistic highlights:")
    for item in mech_bundle.get('mechanistic_highlights', []):
        print(f"- {item}")
    mech_df = pd.DataFrame([
        {
            'gene': gene,
            'dx_hits': len((payload.get('differential_expression') or {}).get('experiments', [])),
            'top_pathway': ((payload.get('reactome_pathways') or [{}])[0].get('name', '')),
            'top_paper': ((payload.get('literature') or [{}])[0].get('title', '')),
        }
        for gene, payload in mech_bundle.get('per_gene', {}).items()
    ])
    mech_df
else:
    print(f"Missing mechanistic evidence bundle: {bundle_path}")

benchmark_path = REPO / 'data/analysis_outputs/analysis-SEAAD-20260402/cell_type_benchmarks_v1.json'
benchmark = json.loads(benchmark_path.read_text())

def leave_one_out_interval(values):
    vals = list(values)
    loo = [(sum(vals) - v) / (len(vals) - 1) for v in vals]
    return min(loo), max(loo)

summary_rows = []
for item in benchmark['cell_type_scores']:
    values = list(item['study_scores'].values())
    ci_low, ci_high = leave_one_out_interval(values)
    summary_rows.append({
        'cell_type': item['display_name'],
        'mean_vulnerability': round(sum(values) / len(values), 3),
        'ci_low': round(ci_low, 3),
        'ci_high': round(ci_high, 3),
        'n_benchmarks': len(values),
        'top_markers': ', '.join(item['primary_markers'][:4]),
        'rationale': item['rationale'],
    })
summary = pd.DataFrame(summary_rows).sort_values('mean_vulnerability', ascending=False)

score_lookup = {item['cell_type']: list(item['study_scores'].values()) for item in benchmark['cell_type_scores']}
display_lookup = {item['cell_type']: item['display_name'] for item in benchmark['cell_type_scores']}
driver_rows = []
for driver in benchmark['candidate_drivers']:
    values = score_lookup[driver['cell_type']]
    ci_low, ci_high = leave_one_out_interval(values)
    mean_support = sum(values) / len(values)
    driver_score = 0.55 * driver['prior_confidence'] + 0.45 * mean_support
    driver_rows.append({
        'candidate_driver': driver['driver'],
        'cell_type': display_lookup[driver['cell_type']],
        'driver_score': round(driver_score, 3),
        'ci_low': round(0.55 * driver['prior_confidence'] + 0.45 * ci_low, 3),
        'ci_high': round(0.55 * driver['prior_confidence'] + 0.45 * ci_high, 3),
        'genes': ', '.join(driver['genes']),
        'testable_direction': driver['expected_direction'],
    })
driver_df = pd.DataFrame(driver_rows).sort_values('driver_score', ascending=False)

source_df = pd.DataFrame([
    {
        'benchmark': s['short_name'],
        'pmid': s.get('pmid', ''),
        'doi': s.get('doi', ''),
        'design': s['region_or_design'],
    }
    for s in benchmark['sources']
])

print('Benchmark artifact:', benchmark_path.relative_to(REPO))
print('Interval method: leave-one-benchmark-out sensitivity over source-study support scores')
display(summary)
display(driver_df)
display(source_df)

driver_plan_path = REPO / 'data/analysis_outputs/analysis-SEAAD-20260402/driver_candidate_tests_v2.json'
driver_plan = json.loads(driver_plan_path.read_text())

def sensitivity_interval_for_model(model):
    values = list(model['benchmark_support_by_source'].values())
    loo = [(sum(values) - v) / (len(values) - 1) for v in values]
    # Combine cross-dataset support, prior confidence, and assay readiness into a conservative
    # rank score. This is not a fitted causal effect; it prioritizes which test should be run first.
    center = (0.45 * (sum(values) / len(values))
              + 0.35 * model['prior_confidence']
              + 0.20 * model['assay_readiness'])
    lo = 0.45 * min(loo) + 0.35 * model['prior_confidence'] + 0.20 * model['assay_readiness']
    hi = 0.45 * max(loo) + 0.35 * model['prior_confidence'] + 0.20 * model['assay_readiness']
    return center, lo, hi

plan_rows = []
for model in driver_plan['candidate_driver_models']:
    center, lo, hi = sensitivity_interval_for_model(model)
    plan_rows.append({
        'candidate_driver': model['candidate_driver'],
        'cell_type': model['cell_type'].replace('_', ' '),
        'priority_score': round(center, 3),
        'interval_low': round(lo, 3),
        'interval_high': round(hi, 3),
        'assay_readiness': model['assay_readiness'],
        'module': ', '.join(model['marker_module']),
        'primary_endpoint': model['primary_endpoint'],
        'falsifying_pattern': model['falsifying_pattern'],
    })
plan_df = pd.DataFrame(plan_rows).sort_values('priority_score', ascending=False)

print('Driver-test artifact:', driver_plan_path.relative_to(REPO))
print('Shared covariates:', ', '.join(driver_plan['shared_covariates']))
display(plan_df[['candidate_driver', 'cell_type', 'priority_score', 'interval_low', 'interval_high', 'assay_readiness', 'module']])
display(plan_df[['candidate_driver', 'primary_endpoint', 'falsifying_pattern']])

fig, ax = plt.subplots(figsize=(8.5, 3.8))
plot_df = plan_df.sort_values('priority_score')
y = range(len(plot_df))
left_err = plot_df['priority_score'] - plot_df['interval_low']
right_err = plot_df['interval_high'] - plot_df['priority_score']
colors = ['#2d7dd2' if ct in {'astrocytes', 'excitatory neurons'} else '#7a7a7a' for ct in plot_df['cell_type']]
ax.errorbar(plot_df['priority_score'], y, xerr=[left_err, right_err], fmt='o', color='#1f2933', ecolor='#8aa0b8', capsize=4)
ax.scatter(plot_df['priority_score'], y, s=90, c=colors, zorder=3)
ax.set_yticks(list(y))
ax.set_yticklabels(plot_df['candidate_driver'])
ax.set_xlabel('Driver-test priority score with leave-one-source-out sensitivity')
ax.set_xlim(0.45, 0.90)
ax.axvline(0.70, color='#b23b3b', linestyle='--', linewidth=1, label='Exchange-ready threshold')
ax.grid(axis='x', alpha=0.25)
ax.legend(loc='lower right')
plt.tight_layout()
plt.show()