import sys, json, sqlite3, warnings, textwrap
import pandas as pd
%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
from pathlib import Path
from datetime import datetime

warnings.filterwarnings('ignore')
pd.set_option('display.max_colwidth', 100)
pd.set_option('display.max_rows', 30)

REPO = Path('/home/ubuntu/scidex')
sys.path.insert(0, str(REPO))

ANALYSIS_ID = 'SDA-2026-04-03-gap-crispr-neurodegeneration-20260402'
DB = REPO / 'scidex.db'
KEY_GENES = ["APOE", "SOD1", "HTT", "MSH3", "TARDBP"]

# Consistent plot style
plt.rcParams.update({
    'figure.figsize': (10, 5),
    'axes.titlesize': 13,
    'axes.labelsize': 11,
    'xtick.labelsize': 9,
    'ytick.labelsize': 9,
    'figure.dpi': 100,
})

print(f"Analysis : {ANALYSIS_ID}")
print(f"Key genes: {', '.join(KEY_GENES)}")
print(f"Generated: {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}")

db = sqlite3.connect(str(DB))

hyps = pd.read_sql_query('''
    SELECT title, target_gene, composite_score,
           confidence_score, novelty_score, feasibility_score, impact_score
    FROM hypotheses
    WHERE analysis_id = ?
    ORDER BY composite_score DESC
''', db, params=[ANALYSIS_ID])

print(f"Hypotheses found   : {len(hyps)}")
print(f"Mean composite     : {hyps['composite_score'].mean():.3f}")
print(f"Top composite      : {hyps['composite_score'].max():.3f}")
print(f"Bottom composite   : {hyps['composite_score'].min():.3f}")
print()
hyps[['title', 'target_gene', 'composite_score']].head(14)

fig, ax = plt.subplots(figsize=(12, 6))

# Truncate titles for readability
labels = [t[:55] + '...' if len(t) > 55 else t for t in hyps['title']]
colors = plt.cm.RdYlGn(hyps['composite_score'] / hyps['composite_score'].max())

bars = ax.barh(range(len(hyps)), hyps['composite_score'], color=colors, edgecolor='#333', linewidth=0.5)
ax.set_yticks(range(len(hyps)))
ax.set_yticklabels(labels, fontsize=8)
ax.set_xlabel('Composite Score')
ax.set_title('CRISPR Neurodegeneration — Hypothesis Composite Scores')
ax.invert_yaxis()
ax.axvline(hyps['composite_score'].mean(), color='navy', linestyle='--', alpha=0.7, label=f"Mean = {hyps['composite_score'].mean():.3f}")
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()

dims = ['confidence_score', 'novelty_score', 'feasibility_score', 'impact_score']
dim_labels = ['Confidence', 'Novelty', 'Feasibility', 'Impact']

heatmap_data = hyps[dims].values
row_labels = [t[:50] + '...' if len(t) > 50 else t for t in hyps['title']]

fig, ax = plt.subplots(figsize=(8, 7))
im = ax.imshow(heatmap_data, cmap='YlOrRd', aspect='auto', vmin=0, vmax=1)

ax.set_xticks(range(len(dim_labels)))
ax.set_xticklabels(dim_labels, fontsize=10)
ax.set_yticks(range(len(row_labels)))
ax.set_yticklabels(row_labels, fontsize=7)

# Annotate cells
for i in range(len(row_labels)):
    for j in range(len(dim_labels)):
        val = heatmap_data[i, j]
        color = 'white' if val > 0.6 else 'black'
        ax.text(j, i, f'{val:.2f}', ha='center', va='center', fontsize=7, color=color)

plt.colorbar(im, ax=ax, shrink=0.8, label='Score')
ax.set_title('Hypothesis Scoring Dimensions')
plt.tight_layout()
plt.show()

edges = pd.read_sql_query('''
    SELECT source_type, target_type, relation, COUNT(*) as count,
           ROUND(AVG(evidence_strength), 3) as avg_strength
    FROM knowledge_edges
    WHERE analysis_id = ?
    GROUP BY source_type, target_type, relation
    ORDER BY count DESC
''', db, params=[ANALYSIS_ID])

total_edges = edges['count'].sum()
print(f"Total knowledge edges  : {total_edges}")
print(f"Distinct relation types: {len(edges)}")
print()

# Bar chart of top edge types
top_edges = edges.head(10).copy()
top_edges['label'] = top_edges['source_type'] + ' → ' + top_edges['target_type'] + ' (' + top_edges['relation'].str[:25] + ')'

fig, ax = plt.subplots(figsize=(10, 5))
ax.barh(range(len(top_edges)), top_edges['count'],
        color=plt.cm.Blues(np.linspace(0.4, 0.9, len(top_edges))),
        edgecolor='#333', linewidth=0.5)
ax.set_yticks(range(len(top_edges)))
ax.set_yticklabels(top_edges['label'], fontsize=8)
ax.set_xlabel('Edge Count')
ax.set_title(f'Knowledge Graph — Top Edge Types ({total_edges} total)')
ax.invert_yaxis()
plt.tight_layout()
plt.show()

# Debate session metadata
sess = pd.read_sql_query('''
    SELECT id, num_rounds, num_hypotheses_generated, num_hypotheses_surviving,
           quality_score, debate_type, created_at
    FROM debate_sessions
    WHERE analysis_id = ?
''', db, params=[ANALYSIS_ID])

if len(sess) > 0:
    s = sess.iloc[0]
    print(f"Debate session      : {s['id']}")
    print(f"Type                : {s['debate_type']}")
    print(f"Rounds              : {s['num_rounds']}")
    print(f"Hypotheses proposed : {s['num_hypotheses_generated']}")
    print(f"Hypotheses surviving: {s['num_hypotheses_surviving']}")
    print(f"Quality score       : {s['quality_score']}")
    print(f"Created             : {s['created_at']}")
else:
    print("No debate session found")

# Debate rounds
rounds = pd.read_sql_query('''
    SELECT round_number, agent_persona, action, substr(content, 1, 300) as excerpt
    FROM debate_rounds
    WHERE session_id = ?
    ORDER BY round_number
''', db, params=[sess.iloc[0]['id'] if len(sess) > 0 else ''])

print(f"\nDebate Rounds: {len(rounds)}")
print("=" * 70)
for _, r in rounds.iterrows():
    persona = r['agent_persona'].upper()
    action = r['action']
    excerpt = r['excerpt'][:250]
    print(f"\nRound {r['round_number']}: {persona} ({action})")
    print(f"  {excerpt}...")
    print("-" * 70)

from tools import get_gene_info

gene_data = {}
for gene in KEY_GENES:
    try:
        info = get_gene_info(gene)
        if info and not info.get('error'):
            gene_data[gene] = info
            print(f"\n{'='*60}")
            print(f"  {gene} — {info.get('name', 'N/A')}")
            print(f"{'='*60}")
            summary = (info.get('summary', '') or '')[:300]
            print(f"  Summary : {summary}")
            aliases = info.get('aliases', [])
            if aliases:
                print(f"  Aliases : {', '.join(aliases[:6])}")
        else:
            print(f"  {gene}: no data returned")
    except Exception as e:
        print(f"  {gene}: error — {e}")

print(f"\nAnnotated {len(gene_data)}/{len(KEY_GENES)} genes successfully")

from tools import clinvar_variants

all_variants = []
for gene in KEY_GENES:
    try:
        variants = clinvar_variants(gene, max_results=8)
        if variants and isinstance(variants, list):
            for v in variants:
                v['query_gene'] = gene
            all_variants.extend(variants)
            print(f"{gene}: {len(variants)} ClinVar variants")
        else:
            print(f"{gene}: {variants}")
    except Exception as e:
        print(f"{gene}: error — {e}")

print(f"\nTotal ClinVar variants retrieved: {len(all_variants)}")

if all_variants:
    cv_df = pd.DataFrame(all_variants)
    display_cols = [c for c in ['query_gene', 'title', 'clinical_significance', 'conditions', 'variation_type']
                   if c in cv_df.columns]
    if not display_cols:
        display_cols = cv_df.columns.tolist()[:5]
    cv_df[display_cols].head(15)

from tools import search_trials

trials = search_trials("CRISPR gene editing neurodegenerative disease therapy", max_results=10)

if trials and isinstance(trials, list):
    trials_df = pd.DataFrame(trials)
    print(f"Clinical trials found: {len(trials_df)}")
    if 'status' in trials_df.columns:
        print(f"Status breakdown: {dict(trials_df['status'].value_counts())}")
    if 'phase' in trials_df.columns:
        print(f"Phase breakdown : {dict(trials_df['phase'].value_counts())}")
    print()
    display_cols = [c for c in ['title', 'status', 'phase', 'conditions', 'enrollment']
                   if c in trials_df.columns]
    if not display_cols:
        display_cols = trials_df.columns.tolist()[:5]
    trials_df[display_cols].head(10)
elif trials and isinstance(trials, dict) and trials.get('error'):
    print(f"ClinicalTrials.gov: {trials['error']}")
    trials_df = pd.DataFrame()
else:
    print(f"ClinicalTrials.gov returned: {trials}")
    trials_df = pd.DataFrame()

from tools import pubmed_search

papers = pubmed_search(
    "CRISPR gene editing neurodegenerative disease therapy Alzheimer Parkinson Huntington",
    max_results=15
)

if papers and isinstance(papers, list):
    papers_df = pd.DataFrame(papers)
    print(f"PubMed papers found: {len(papers_df)}")
    if 'year' in papers_df.columns:
        print(f"Year range: {papers_df['year'].min()} – {papers_df['year'].max()}")
    print()
    display_cols = [c for c in ['title', 'journal', 'year', 'pmid'] if c in papers_df.columns]
    if display_cols:
        papers_df[display_cols].head(10)
    else:
        papers_df.head(10)
else:
    print(f"PubMed returned: {papers}")
    papers_df = pd.DataFrame()

from tools import string_protein_interactions

interactions = string_protein_interactions(KEY_GENES, score_threshold=400)

if interactions and isinstance(interactions, list):
    ppi_df = pd.DataFrame(interactions)
    print(f"Protein-protein interactions: {len(ppi_df)}")
    if len(ppi_df) > 0 and 'score' in ppi_df.columns:
        print(f"Score range: {ppi_df['score'].min():.3f} – {ppi_df['score'].max():.3f}")
    print()
    ppi_df.head(15)
else:
    print(f"STRING returned: {interactions}")
    ppi_df = pd.DataFrame()

from tools import reactome_pathways

all_pathways = []
for gene in KEY_GENES:
    try:
        pathways = reactome_pathways(gene, max_results=5)
        if pathways and isinstance(pathways, list):
            for p in pathways:
                p['query_gene'] = gene
            all_pathways.extend(pathways)
            print(f"{gene}: {len(pathways)} pathways")
        else:
            print(f"{gene}: {pathways}")
    except Exception as e:
        print(f"{gene}: error — {e}")

print(f"\nTotal pathways across all genes: {len(all_pathways)}")

if all_pathways:
    pw_df = pd.DataFrame(all_pathways)
    display_cols = [c for c in ['query_gene', 'pathway_name', 'pathway_id', 'name']
                   if c in pw_df.columns]
    if not display_cols:
        display_cols = pw_df.columns.tolist()[:4]
    pw_df[display_cols].head(20)

linked_papers = pd.read_sql_query('''
    SELECT DISTINCT hp.pmid, hp.evidence_direction, hp.claim
    FROM hypothesis_papers hp
    INNER JOIN hypotheses h ON hp.hypothesis_id = h.id
    WHERE h.analysis_id = ?
    LIMIT 25
''', db, params=[ANALYSIS_ID])

print(f"Papers linked to hypotheses: {len(linked_papers)}")
if len(linked_papers) > 0:
    print(f"Evidence breakdown: {dict(linked_papers['evidence_direction'].value_counts())}")
    print()
    linked_papers.head(12)
else:
    print("No linked papers found")

cooccurrence = pd.read_sql_query('''
    SELECT source_id, target_id, relation, evidence_strength
    FROM knowledge_edges
    WHERE analysis_id = ?
      AND source_type = 'gene' AND target_type = 'gene'
    ORDER BY evidence_strength DESC
    LIMIT 30
''', db, params=[ANALYSIS_ID])

print(f"Gene-gene edges: {len(cooccurrence)}")
if len(cooccurrence) > 0:
    all_genes_in_net = set(cooccurrence['source_id'].tolist() + cooccurrence['target_id'].tolist())
    print(f"Unique genes: {len(all_genes_in_net)}")
    print(f"Relations: {dict(cooccurrence['relation'].value_counts())}")
    print()
    cooccurrence.head(15)
else:
    print("No gene-gene edges found")

fig, ax = plt.subplots(figsize=(10, 7))

x = hyps['novelty_score']
y = hyps['feasibility_score']
sizes = hyps['impact_score'] * 400
colors = hyps['confidence_score']

scatter = ax.scatter(x, y, s=sizes, c=colors, cmap='RdYlGn',
                     edgecolors='#333', linewidth=0.8, alpha=0.85, vmin=0, vmax=1)

# Label top 5
for i in range(min(5, len(hyps))):
    label = hyps.iloc[i]['title'][:35] + '...' if len(hyps.iloc[i]['title']) > 35 else hyps.iloc[i]['title']
    ax.annotate(label, (x.iloc[i], y.iloc[i]),
               fontsize=6.5, ha='left', va='bottom',
               xytext=(5, 5), textcoords='offset points',
               arrowprops=dict(arrowstyle='-', color='gray', lw=0.5))

plt.colorbar(scatter, ax=ax, shrink=0.8, label='Confidence Score')
ax.set_xlabel('Novelty Score')
ax.set_ylabel('Feasibility Score')
ax.set_title('CRISPR Hypothesis Landscape — Novelty vs Feasibility')
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# Quadrant lines at 0.5
ax.axhline(0.5, color='gray', linestyle=':', alpha=0.5)
ax.axvline(0.5, color='gray', linestyle=':', alpha=0.5)
ax.text(0.75, 0.75, 'HIGH PRIORITY\n(novel + feasible)', ha='center', fontsize=8, color='green', alpha=0.6)
ax.text(0.25, 0.25, 'LOW PRIORITY', ha='center', fontsize=8, color='red', alpha=0.6)

plt.tight_layout()
plt.show()

print("=" * 70)
print("CRISPR-Based Therapeutic Approaches for Neurodegenerative Diseases")
print("=" * 70)

print(f"""
Research Question:
  Evaluate the potential of CRISPR/Cas9 and related gene-editing technologies
  for treating neurodegenerative diseases.

Hypotheses         : {len(hyps)}
  Mean composite   : {hyps['composite_score'].mean():.3f}
  Score range      : {hyps['composite_score'].min():.3f} – {hyps['composite_score'].max():.3f}
  Top hypothesis   : {hyps.iloc[0]['title']}

Knowledge Graph    : {total_edges} edges across {len(edges)} relation types

Multi-Agent Debate : {len(rounds)} rounds (theorist → skeptic → domain expert → synthesizer)
  Proposed {sess.iloc[0]['num_hypotheses_generated']} hypotheses, {sess.iloc[0]['num_hypotheses_surviving']} survived critique

Gene Annotations   : {len(gene_data)} / {len(KEY_GENES)} genes annotated
  Genes: {', '.join(gene_data.keys())}

ClinVar Variants   : {len(all_variants)} clinical variants across target genes

Clinical Trials    : {len(trials_df)} trials on ClinicalTrials.gov

PubMed Papers      : {len(papers_df)} papers found

STRING Interactions: {len(ppi_df)} protein-protein interactions

Reactome Pathways  : {len(all_pathways)} pathways

Forge Tools Used   : Gene Info, ClinVar, Clinical Trials, PubMed, STRING PPI, Reactome (6 tools)
""")

db.close()
print(f"Notebook executed: {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}")