"""
Atomic VSA Interactive Demo
Gradio frontend calling Julia VSA backend
"""

import gradio as gr
import requests
import time
import csv
import os

# Julia backend URL
JULIA_BACKEND = "http://localhost:8080"

# ============================================================================
# Backend Communication
# ============================================================================

def wait_for_backend(max_retries=30, delay=2):
    """Wait for Julia backend to be ready."""
    for i in range(max_retries):
        try:
            resp = requests.get(f"{JULIA_BACKEND}/api/health", timeout=5)
            if resp.status_code == 200:
                return True
        except:
            pass
        time.sleep(delay)
    return False

def get_backend_info():
    """Get info from Julia backend."""
    try:
        resp = requests.get(f"{JULIA_BACKEND}/api/info", timeout=10)
        if resp.status_code == 200:
            return resp.json()
    except:
        pass
    return None

def classify_symptoms(symptoms: list) -> dict:
    """Send symptoms to Julia backend for classification."""
    try:
        resp = requests.post(
            f"{JULIA_BACKEND}/api/classify",
            json={"symptoms": symptoms},
            timeout=30
        )
        if resp.status_code == 200:
            return resp.json()
    except Exception as e:
        return {"error": str(e)}
    return {"error": "Backend unavailable"}

# ============================================================================
# Symptoms List
# ============================================================================

SYMPTOMS = [
    "chest_pain", "shortness_of_breath", "fever", "cough",
    "headache", "fatigue", "nausea", "dizziness",
    "abdominal_pain", "back_pain", "joint_pain", "rash",
    "sore_throat", "runny_nose", "muscle_ache", "chills"
]

SYMPTOM_LABELS = {
    "chest_pain": "Chest Pain",
    "shortness_of_breath": "Shortness of Breath",
    "fever": "Fever",
    "cough": "Cough",
    "headache": "Headache",
    "fatigue": "Fatigue",
    "nausea": "Nausea",
    "dizziness": "Dizziness",
    "abdominal_pain": "Abdominal Pain",
    "back_pain": "Back Pain",
    "joint_pain": "Joint Pain",
    "rash": "Rash",
    "sore_throat": "Sore Throat",
    "runny_nose": "Runny Nose",
    "muscle_ache": "Muscle Ache",
    "chills": "Chills"
}

# ============================================================================
# Sample Test Cases - Loaded from CSV
# ============================================================================

def load_sample_cases(csv_filename="sample_cases.csv"):
    """Load sample test cases from CSV file for transparency and trust."""
    cases = []
    
    # Try multiple paths for Docker and local environments
    possible_paths = [
        os.path.join("/app", csv_filename),  # Docker workdir
        os.path.join(os.path.dirname(os.path.abspath(__file__)), csv_filename),  # Same dir as script
        csv_filename,  # Current working directory
        os.path.join(os.getcwd(), csv_filename),  # Explicit cwd
    ]
    
    csv_path = None
    for path in possible_paths:
        if os.path.exists(path):
            csv_path = path
            break
    
    if csv_path is None:
        print(f"Warning: {csv_filename} not found in any of: {possible_paths}")
        return cases
    
    try:
        with open(csv_path, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                cases.append({
                    "name": f"{row['icon']} Case {row['case_id']}: {row['name']}",
                    "symptoms": [s.strip() for s in row['symptoms'].split(',')],
                    "expected": row['expected_category'],
                    "rationale": row['clinical_rationale'],
                    "confidence": f"~{row['confidence_range']}"
                })
        print(f"Loaded {len(cases)} sample cases from {csv_path}")
    except Exception as e:
        print(f"Error loading sample cases from {csv_path}: {e}")
    
    return cases

# Load cases from CSV at startup
SAMPLE_CASES = load_sample_cases()

def get_sample_cases_markdown():
    """Generate markdown table of sample cases."""
    md = """
## 📚 Sample Test Cases

> **Data Source:** Loaded from [`sample_cases.csv`](https://huggingface.co/spaces/marshad180/atomic-vsa/blob/main/sample_cases.csv) — fully auditable and version-controlled.

Use these cases to validate the system is working correctly.

| Case | Symptoms | Expected Result | Confidence |
|------|----------|-----------------|------------|
"""
    for case in SAMPLE_CASES:
        symptoms_str = ", ".join([SYMPTOM_LABELS.get(s, s) for s in case["symptoms"]])
        md += f"| {case['name']} | {symptoms_str} | **{case['expected']}** | {case['confidence']} |\n"
    
    md += f"""

**{len(SAMPLE_CASES)} cases loaded from CSV**

### How to Use
1. Go to the **Clinical Triage** tab
2. Select the symptoms listed for a case
3. Click **Classify Patient**
4. Verify the result matches the expected classification

### Understanding Results
- **Confidence** is the cosine similarity between patient vector and category prototype
- Higher values (>80%) indicate strong match
- Multiple categories may have similar scores for overlapping symptoms
"""
    return md

# ============================================================================
# UI Functions
# ============================================================================

def process_triage(*symptom_states) -> tuple:
    """Process symptoms and return triage classification."""
    
    # Collect selected symptoms
    selected = []
    for symptom, state in zip(SYMPTOMS, symptom_states):
        if state:
            selected.append(symptom)
    
    if not selected:
        return (
            "⚠️ Please select at least one symptom",
            "",
            "No symptoms selected"
        )
    
    # Call Julia backend
    result = classify_symptoms(selected)
    
    if "error" in result:
        return (
            f"❌ Error: {result['error']}",
            "",
            "Backend error"
        )
    
    # Format results
    top_result = result["results"][0]
    category = top_result["category"]
    similarity = top_result["similarity"]
    
    # Determine urgency styling
    if "Emergency" in category:
        urgency_icon = "🚨"
        urgency_color = "red"
    elif "Urgent" in category:
        urgency_icon = "⚠️"
        urgency_color = "orange"
    elif "Standard" in category:
        urgency_icon = "📋"
        urgency_color = "blue"
    else:
        urgency_icon = "✅"
        urgency_color = "green"
    
    # Main result
    main_result = f"""
## {urgency_icon} {category}

**Confidence:** {similarity:.1%}

**Symptoms Analyzed:** {len(selected)}
- {', '.join([SYMPTOM_LABELS[s] for s in selected])}
"""
    
    # All scores
    all_scores = "\n".join([
        f"| {r['category']} | {r['similarity']:.1%} |"
        for r in result["results"]
    ])
    
    scores_table = f"""
| Category | Similarity |
|----------|------------|
{all_scores}
"""
    
    # Technical details
    tech_details = f"""
**Engine:** Julia VSA (D={result.get('dimensionality', 2048)})
**Method:** Cosine similarity to pre-computed prototypes
**Decision:** Winner-Take-All (deterministic)
"""
    
    return main_result, scores_table, tech_details

# ============================================================================
# Gradio Interface
# ============================================================================

with gr.Blocks(title="Atomic VSA", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # ⚛️ Atomic VSA: Clinical Triage Demo
    
    **Physics-Inspired Hyperdimensional Computing for Explainable AI**
    
    Select symptoms below to receive a triage classification powered by the Julia VSA engine.
    """)
    
    with gr.Tabs():
        # Tab 1: Triage Demo
        with gr.TabItem("🏥 Clinical Triage"):
            gr.Markdown("### Select Patient Symptoms")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("**Cardiopulmonary**")
                    chest_pain = gr.Checkbox(label="Chest Pain")
                    shortness_of_breath = gr.Checkbox(label="Shortness of Breath")
                    cough = gr.Checkbox(label="Cough")
                    
                with gr.Column():
                    gr.Markdown("**Systemic**")
                    fever = gr.Checkbox(label="Fever")
                    chills = gr.Checkbox(label="Chills")
                    fatigue = gr.Checkbox(label="Fatigue")
                    muscle_ache = gr.Checkbox(label="Muscle Ache")
                    
                with gr.Column():
                    gr.Markdown("**Neurological**")
                    headache = gr.Checkbox(label="Headache")
                    dizziness = gr.Checkbox(label="Dizziness")
                    nausea = gr.Checkbox(label="Nausea")
                    
                with gr.Column():
                    gr.Markdown("**Other**")
                    abdominal_pain = gr.Checkbox(label="Abdominal Pain")
                    back_pain = gr.Checkbox(label="Back Pain")
                    joint_pain = gr.Checkbox(label="Joint Pain")
                    rash = gr.Checkbox(label="Rash")
                    sore_throat = gr.Checkbox(label="Sore Throat")
                    runny_nose = gr.Checkbox(label="Runny Nose")
            
            classify_btn = gr.Button("🔬 Classify Patient", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    result_output = gr.Markdown(label="Classification Result")
                with gr.Column():
                    scores_output = gr.Markdown(label="All Scores")
            
            tech_output = gr.Markdown(label="Technical Details")
            
            # Wire up the button
            classify_btn.click(
                process_triage,
                inputs=[
                    chest_pain, shortness_of_breath, fever, cough,
                    headache, fatigue, nausea, dizziness,
                    abdominal_pain, back_pain, joint_pain, rash,
                    sore_throat, runny_nose, muscle_ache, chills
                ],
                outputs=[result_output, scores_output, tech_output]
            )
        
        # Tab 2: Sample Cases (Knowledge Base)
        with gr.TabItem("📚 Sample Cases"):
            gr.Markdown(get_sample_cases_markdown())
            
            gr.Markdown("---\n### 🧪 Detailed Case Explanations")
            
            for case in SAMPLE_CASES:
                with gr.Accordion(case["name"], open=False):
                    gr.Markdown(f"""
**Symptoms to Select:**
{chr(10).join(['- ' + SYMPTOM_LABELS[s] for s in case['symptoms']])}

**Expected Classification:** `{case['expected']}`

**Clinical Rationale:** {case['rationale']}

**Expected Confidence Range:** {case['confidence']}

---
*This pattern is encoded in the VSA prototype vectors built from clinical knowledge.*
""")
        
        # Tab 3: About
        with gr.TabItem("📄 About"):
            gr.Markdown("""
            ## Atomic Vector Symbolic Architecture
            
            ### Key Innovation
            
            Atomic VSA applies **physics-inspired principles** to hyperdimensional computing:
            
            | Physics | Computing (AVSA) | Clinical Medicine |
            |---------|------------------|-------------------|
            | Atom | 10,048-dim vector | Semantic concept |
            | Proton (+) | Positive evidence | Finding FOR diagnosis |
            | Electron (−) | Negative evidence | Finding AGAINST diagnosis |
            | Molecule | BIND(a ⊗ b) | Clinical fact pair |
            | Superposition | BUNDLE(⊕) | Patient record |
            
            ### Performance
            
            | Metric | Result |
            |--------|--------|
            | F1 Score | 92.5% (25-category ICD-11) |
            | Label Recall | 91.9% (comorbidity) |
            | Latency | 11.97ms (p50) |
            | Power | 15W (edge) |
            
            ### Why VSA?
            
            - ✅ **Deterministic**: Same input → same output, always
            - ✅ **Interpretable**: Algebraic operations are transparent
            - ✅ **Efficient**: No training, no GPU required
            - ✅ **Green AI**: 160× lower power than neural networks
            
            ### Technical Stack
            
            - **Backend**: Julia VSA engine (authentic implementation)
            - **Frontend**: Gradio
            - **Architecture**: HTTP REST API
            
            ### Links
            
            - **Paper (DOI)**: [10.5281/zenodo.18650281](https://doi.org/10.5281/zenodo.18650281)
            - **Author**: Muhammad Arshad (marshad.dev@gmail.com)
            """)

if __name__ == "__main__":
    print("Waiting for Julia backend...")
    if wait_for_backend():
        print("Julia backend ready!")
        info = get_backend_info()
        if info:
            print(f"Backend: {info.get('name')} v{info.get('version')} ({info.get('engine')})")
    else:
        print("Warning: Julia backend not responding, some features may not work")
    
    demo.launch(server_name="0.0.0.0", server_port=7860)