"""
Atomic VSA Interactive Demo
A Gradio-based demonstration of Vector Symbolic Architecture for clinical triage.
"""

import gradio as gr
import numpy as np
from typing import Dict, List, Tuple

# ============================================================================
# VSA Core Implementation (Python port of Julia logic)
# ============================================================================

D = 2048  # Dimensionality for demo (production uses 10,048)
np.random.seed(42)

class AtomRegistry:
    """Registry of atomic vectors (symbols)."""
    def __init__(self, dim: int = D):
        self.dim = dim
        self.atoms: Dict[str, np.ndarray] = {}
    
    def get_or_create(self, name: str) -> np.ndarray:
        if name not in self.atoms:
            # Create bipolar random vector {-1, +1}^D
            self.atoms[name] = np.random.choice([-1, 1], size=self.dim).astype(np.float32)
        return self.atoms[name]
    
    def __getitem__(self, name: str) -> np.ndarray:
        return self.get_or_create(name)

# Global registry
registry = AtomRegistry()

def bind(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """BIND operation: Element-wise multiplication (⊗)."""
    return a * b

def bundle(vectors: List[np.ndarray]) -> np.ndarray:
    """BUNDLE operation: Element-wise addition (⊕)."""
    return np.sum(vectors, axis=0)

def similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between vectors."""
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a, b) / (norm_a * norm_b))

def thermometer_encode(value: float, min_val: float, max_val: float, levels: int = 20) -> np.ndarray:
    """Thermometer encoding for continuous values."""
    level_atoms = [registry.get_or_create(f"_level_{i}") for i in range(levels)]
    active_levels = int((value - min_val) / (max_val - min_val) * levels)
    active_levels = max(0, min(levels, active_levels))
    if active_levels == 0:
        return np.zeros(D, dtype=np.float32)
    return bundle(level_atoms[:active_levels])

# ============================================================================
# Clinical Triage System
# ============================================================================

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"
]

TRIAGE_CATEGORIES = {
    "Emergency - Cardiac": ["chest_pain", "shortness_of_breath", "dizziness"],
    "Urgent - Respiratory": ["shortness_of_breath", "cough", "fever", "chills"],
    "Urgent - Infection": ["fever", "chills", "fatigue", "muscle_ache"],
    "Standard - Flu-like": ["fever", "cough", "sore_throat", "runny_nose", "muscle_ache"],
    "Standard - GI": ["nausea", "abdominal_pain", "fever"],
    "Standard - Musculoskeletal": ["back_pain", "joint_pain", "muscle_ache"],
    "Low Priority - Minor": ["headache", "fatigue", "runny_nose"],
}

# Pre-build prototype vectors for each triage category
PROTOTYPES: Dict[str, np.ndarray] = {}
for category, symptom_list in TRIAGE_CATEGORIES.items():
    molecules = []
    for symptom in symptom_list:
        field_atom = registry[f"symptom_field"]
        value_atom = registry[symptom]
        molecules.append(bind(field_atom, value_atom))
    PROTOTYPES[category] = bundle(molecules)

def create_patient_vector(symptoms: List[str], vitals: Dict[str, float]) -> np.ndarray:
    """Create a patient record as a bundled molecule."""
    molecules = []
    
    # Encode symptoms
    for symptom in symptoms:
        if symptom in SYMPTOMS:
            field_atom = registry["symptom_field"]
            value_atom = registry[symptom]
            molecules.append(bind(field_atom, value_atom))
    
    # Encode vitals with thermometer encoding
    if "heart_rate" in vitals:
        hr_encoded = thermometer_encode(vitals["heart_rate"], 40, 180)
        molecules.append(bind(registry["heart_rate_field"], hr_encoded))
    
    if "temperature" in vitals:
        temp_encoded = thermometer_encode(vitals["temperature"], 35, 42)
        molecules.append(bind(registry["temperature_field"], temp_encoded))
    
    if "blood_pressure_sys" in vitals:
        bp_encoded = thermometer_encode(vitals["blood_pressure_sys"], 80, 200)
        molecules.append(bind(registry["bp_field"], bp_encoded))
    
    if not molecules:
        return np.zeros(D, dtype=np.float32)
    
    return bundle(molecules)

def triage_patient(patient_vector: np.ndarray) -> List[Tuple[str, float]]:
    """Classify patient against triage prototypes."""
    scores = []
    for category, prototype in PROTOTYPES.items():
        sim = similarity(patient_vector, prototype)
        scores.append((category, sim))
    
    # Sort by similarity descending
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores

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

def process_triage(
    chest_pain: bool, shortness_of_breath: bool, fever: bool, cough: bool,
    headache: bool, fatigue: bool, nausea: bool, dizziness: bool,
    abdominal_pain: bool, back_pain: bool, joint_pain: bool, rash: bool,
    sore_throat: bool, runny_nose: bool, muscle_ache: bool, chills: bool,
    heart_rate: float, temperature: float, blood_pressure: float
) -> Tuple[str, str, str]:
    """Process symptoms and return triage classification."""
    
    # Collect selected symptoms
    symptom_map = {
        "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
    }
    
    selected_symptoms = [s for s, v in symptom_map.items() if v]
    
    if not selected_symptoms:
        return "No symptoms selected", "", ""
    
    # Create patient vector
    vitals = {
        "heart_rate": heart_rate,
        "temperature": temperature,
        "blood_pressure_sys": blood_pressure
    }
    
    patient_vec = create_patient_vector(selected_symptoms, vitals)
    
    # Get triage scores
    scores = triage_patient(patient_vec)
    
    # Format results
    primary = scores[0]
    primary_result = f"**{primary[0]}**\nSimilarity: {primary[1]:.4f}"
    
    # All scores
    all_scores = "\n".join([f"{cat}: {sim:.4f}" for cat, sim in scores])
    
    # Explanation
    explanation = f"""
**VSA Explanation:**
- Patient encoded as {D}-dimensional bipolar vector
- Symptoms: {', '.join(selected_symptoms)}
- Heart Rate: {heart_rate} bpm → Thermometer encoded
- Temperature: {temperature}°C → Thermometer encoded  
- Blood Pressure: {blood_pressure} mmHg → Thermometer encoded

**How it works:**
1. Each symptom is BOUND (⊗) with a field identifier
2. All bound pairs are BUNDLED (⊕) into a superposition
3. Cosine similarity computed against {len(PROTOTYPES)} prototypes
4. Classification is deterministic and fully explainable
"""
    
    return primary_result, all_scores, explanation

def demo_bind_operation(concept_a: str, concept_b: str) -> str:
    """Demonstrate BIND operation."""
    if not concept_a or not concept_b:
        return "Enter two concepts"
    
    vec_a = registry[concept_a]
    vec_b = registry[concept_b]
    bound = bind(vec_a, vec_b)
    
    # Verify self-inverse property
    unbound = bind(bound, vec_b)
    recovery_sim = similarity(unbound, vec_a)
    
    return f"""
**BIND Operation: {concept_a} ⊗ {concept_b}**

Vector A (first 10 dims): {vec_a[:10].astype(int).tolist()}
Vector B (first 10 dims): {vec_b[:10].astype(int).tolist()}
Bound (first 10 dims): {bound[:10].astype(int).tolist()}

**Self-Inverse Property:**
(A ⊗ B) ⊗ B = A
Recovery similarity: {recovery_sim:.6f} (should be 1.0)

**Orthogonality:**
Sim(A, B): {similarity(vec_a, vec_b):.4f} (random vectors ≈ 0)
Sim(A, Bound): {similarity(vec_a, bound):.4f} (bound dissimilar to inputs)
"""

def demo_bundle_operation(concepts: str) -> str:
    """Demonstrate BUNDLE operation."""
    if not concepts:
        return "Enter comma-separated concepts"
    
    concept_list = [c.strip() for c in concepts.split(",") if c.strip()]
    if len(concept_list) < 2:
        return "Enter at least 2 concepts"
    
    vectors = [registry[c] for c in concept_list]
    bundled = bundle(vectors)
    
    # Show similarity to each component
    sims = [(c, similarity(bundled, registry[c])) for c in concept_list]
    
    sim_report = "\n".join([f"  Sim(Bundle, {c}): {s:.4f}" for c, s in sims])
    
    return f"""
**BUNDLE Operation: {' ⊕ '.join(concept_list)}**

Bundled {len(concept_list)} vectors into superposition.

**Holographic Property - Bundle is similar to ALL inputs:**
{sim_report}

Unlike classical storage, the bundle simultaneously represents 
all {len(concept_list)} concepts in the same {D}-dimensional space!
"""

# Build the Gradio app
with gr.Blocks(title="Atomic VSA Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # ⚛️ Atomic VSA: Interactive Demo
    
    **Physics-Inspired Hyperdimensional Computing for Explainable Clinical AI**
    
    [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.18650281.svg)](https://doi.org/10.5281/zenodo.18650281)
    [![GitHub](https://img.shields.io/badge/GitHub-atomic--vsa--research-blue)](https://github.com/muhammadarshad/atomic-vsa-research)
    
    This demo showcases the Atomic Vector Symbolic Architecture (Atomic VSA) — 
    a deterministic, interpretable AI framework achieving **92.5% F1** on clinical triage.
    """)
    
    with gr.Tabs():
        # Tab 1: Clinical Triage Demo
        with gr.TabItem("🏥 Clinical Triage"):
            gr.Markdown("### Enter Patient Symptoms & Vitals")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("**Symptoms:**")
                    chest_pain = gr.Checkbox(label="Chest Pain")
                    shortness_of_breath = gr.Checkbox(label="Shortness of Breath")
                    fever = gr.Checkbox(label="Fever")
                    cough = gr.Checkbox(label="Cough")
                    headache = gr.Checkbox(label="Headache")
                    fatigue = gr.Checkbox(label="Fatigue")
                    nausea = gr.Checkbox(label="Nausea")
                    dizziness = gr.Checkbox(label="Dizziness")
                
                with gr.Column():
                    gr.Markdown("**More Symptoms:**")
                    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")
                    muscle_ache = gr.Checkbox(label="Muscle Ache")
                    chills = gr.Checkbox(label="Chills")
                
                with gr.Column():
                    gr.Markdown("**Vitals:**")
                    heart_rate = gr.Slider(40, 180, value=75, label="Heart Rate (bpm)")
                    temperature = gr.Slider(35.0, 42.0, value=37.0, step=0.1, label="Temperature (°C)")
                    blood_pressure = gr.Slider(80, 200, value=120, label="Systolic BP (mmHg)")
            
            triage_btn = gr.Button("🔬 Run Triage", variant="primary")
            
            with gr.Row():
                primary_output = gr.Markdown(label="Primary Classification")
                scores_output = gr.Textbox(label="All Scores", lines=7)
            
            explanation_output = gr.Markdown(label="VSA Explanation")
            
            triage_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,
                    heart_rate, temperature, blood_pressure
                ],
                outputs=[primary_output, scores_output, explanation_output]
            )
        
        # Tab 2: VSA Operations
        with gr.TabItem("🧮 VSA Operations"):
            gr.Markdown("""
            ### Explore the Algebraic Operations
            
            Atomic VSA uses two fundamental operations:
            - **BIND (⊗)**: Creates relational structures (self-inverse)
            - **BUNDLE (⊕)**: Creates holographic superpositions
            """)
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### BIND Operation")
                    bind_a = gr.Textbox(label="Concept A", value="HeartRate")
                    bind_b = gr.Textbox(label="Concept B", value="115bpm")
                    bind_btn = gr.Button("Compute BIND")
                    bind_output = gr.Markdown()
                    
                    bind_btn.click(demo_bind_operation, [bind_a, bind_b], bind_output)
                
                with gr.Column():
                    gr.Markdown("#### BUNDLE Operation")
                    bundle_input = gr.Textbox(
                        label="Concepts (comma-separated)", 
                        value="fever, cough, fatigue"
                    )
                    bundle_btn = gr.Button("Compute BUNDLE")
                    bundle_output = gr.Markdown()
                    
                    bundle_btn.click(demo_bundle_operation, [bundle_input], bundle_output)
        
        # 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
            
            ### Links
            
            - **Paper (DOI)**: [10.5281/zenodo.18650281](https://doi.org/10.5281/zenodo.18650281)
            - **Code**: [GitHub](https://github.com/muhammadarshad/atomic-vsa-research)
            - **Author**: Muhammad Arshad (marshad.dev@gmail.com)
            """)

if __name__ == "__main__":
    demo.launch()