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}")