"""
Sundew Live Monitor - Enhanced "Wow" Demo
Production-quality interface showcasing neurosymbolic ECG monitoring
"""
import io
import json
import math
import os
import sys
from typing import Any, Dict, List

import gradio as gr
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
import pandas as pd

ROOT = os.path.dirname(os.path.abspath(__file__))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from app.ml.gating import gate_signal
from app.ml.inference import infer_ecg, load_model
from app.rules.engine import evaluate_ecg_rules

load_model()

SCENARIOS = {
    "healthy": {
        "name": "Healthy Adult (60yo)",
        "age": 60,
        "has_prior_stroke": False,
        "signal_type": "normal",
        "description": "Normal sinus rhythm, no risk factors, routine monitoring",
        "icon": "✓"
    },
    "afib_high_risk": {
        "name": "AFib Suspect (85yo, Prior Stroke)",
        "age": 85,
        "has_prior_stroke": True,
        "signal_type": "afib",
        "description": "Irregular rhythm detected, high-risk patient requiring immediate review",
        "icon": "⚠"
    },
    "tachycardia": {
        "name": "Tachycardia Episode (45yo)",
        "age": 45,
        "has_prior_stroke": False,
        "signal_type": "tachy",
        "description": "Elevated heart rate (120+ bpm), otherwise healthy patient",
        "icon": "↑"
    },
    "elderly_normal": {
        "name": "Elderly Patient (78yo, Normal ECG)",
        "age": 78,
        "has_prior_stroke": True,
        "signal_type": "normal",
        "description": "High-risk profile but currently stable rhythm",
        "icon": "👤"
    },
    "noisy": {
        "name": "Poor Signal Quality",
        "age": 60,
        "has_prior_stroke": False,
        "signal_type": "noise",
        "description": "Motion artifacts, low-quality signal requiring gating",
        "icon": "~"
    }
}


def generate_signal(signal_type: str, length: int = 512) -> List[float]:
    if signal_type == "normal":
        return [0.05 * math.sin(2 * math.pi * 2 * (i / length)) +
                0.02 * math.sin(2 * math.pi * 0.5 * (i / length)) for i in range(length)]
    elif signal_type == "afib":
        return [
            0.25 * math.sin(2 * math.pi * 6 * (i / length)) +
            0.05 * math.sin(2 * math.pi * 15 * (i / length)) +
            (0.15 if i % 40 == 0 else 0.0) +
            0.03 * (hash(i) % 100 - 50) / 500
            for i in range(length)
        ]
    elif signal_type == "tachy":
        return [0.08 * math.sin(2 * math.pi * 4.5 * (i / length)) +
                0.03 * math.sin(2 * math.pi * 1 * (i / length)) for i in range(length)]
    elif signal_type == "noise":
        return [0.02 * math.sin(2 * math.pi * 1 * (i / length)) +
                (0.01 if i % 13 == 0 else 0.0) +
                0.005 * (hash(i) % 100 - 50) / 50 for i in range(length)]
    return [0.0] * length


def run_pipeline(scenario_key: str):
    scenario = SCENARIOS[scenario_key]
    signal = generate_signal(scenario["signal_type"], length=512)

    gated, gating_meta = gate_signal(signal, return_windows=True)
    model_output = infer_ecg(gated, original_len=len(signal), gating_meta=gating_meta)

    patient_context = {
        "patient_id": scenario_key,
        "age": scenario["age"],
        "has_prior_stroke": scenario["has_prior_stroke"],
    }
    rules_result = evaluate_ecg_rules(patient_context, model_output)

    # Build comprehensive results
    energy_saved = (1 - gating_meta.get("ratio", 1.0)) * 100

    # Summary card
    summary_html = f"""
    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 25px; border-radius: 15px; margin: 10px 0;">
        <h2 style="margin: 0 0 15px 0;">Patient: {scenario['name']}</h2>
        <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px;">
            <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px;">
                <h3 style="margin: 0; font-size: 14px; opacity: 0.9;">Diagnosis</h3>
                <p style="margin: 5px 0 0 0; font-size: 24px; font-weight: bold;">{model_output.get('label', 'Unknown').upper()}</p>
                <p style="margin: 5px 0 0 0; opacity: 0.8;">Confidence: {model_output.get('score', 0.0):.1%}</p>
            </div>
            <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px;">
                <h3 style="margin: 0; font-size: 14px; opacity: 0.9;">Alert Level</h3>
                <p style="margin: 5px 0 0 0; font-size: 24px; font-weight: bold;">{rules_result.get('alert_level', 'NONE').upper()}</p>
                <p style="margin: 5px 0 0 0; opacity: 0.8;">HR: {model_output.get('hr')} bpm</p>
            </div>
        </div>
        <div style="margin-top: 15px; background: rgba(46,213,115,0.2); padding: 12px; border-radius: 8px; border-left: 4px solid #2ed573;">
            <strong>Energy Savings: {energy_saved:.1f}%</strong> | Windows: {gating_meta.get('selected_windows', 0)}/{gating_meta.get('total_windows', 0)}
        </div>
    </div>
    """

    # Signal visualization
    fig1, axes = plt.subplots(2, 1, figsize=(12, 6))
    axes[0].plot(signal, color='#3498db', linewidth=1.5, alpha=0.8)
    axes[0].set_title('Original ECG Signal', fontsize=13, fontweight='bold')
    axes[0].set_ylabel('Amplitude')
    axes[0].grid(alpha=0.3)
    axes[0].set_xlim(0, len(signal))

    axes[1].plot(gated, color='#e74c3c', linewidth=1.5, alpha=0.8)
    axes[1].set_title(f'Gated Signal (Compression: {gating_meta.get("ratio", 1.0):.1%})', fontsize=13, fontweight='bold')
    axes[1].set_xlabel('Sample Index')
    axes[1].set_ylabel('Amplitude')
    axes[1].grid(alpha=0.3)

    fig1.tight_layout()
    buf1 = io.BytesIO()
    fig1.savefig(buf1, format='png', dpi=150, bbox_inches='tight')
    plt.close(fig1)
    buf1.seek(0)
    signal_img = mpimg.imread(buf1)

    # Energy bar chart - Enhanced version
    fig2 = plt.figure(figsize=(12, 7))
    gs = fig2.add_gridspec(2, 2, height_ratios=[2, 1], hspace=0.4, wspace=0.3)

    # Main comparison chart
    ax1 = fig2.add_subplot(gs[0, :])
    categories = ['Baseline\n(Traditional ML)', 'Sundew\n(Neurosymbolic)']
    compute = [100, gating_meta.get("ratio", 1.0) * 100]
    colors = ['#e74c3c', '#2ecc71']

    bars = ax1.barh(categories, compute, color=colors, edgecolor='white', linewidth=2, height=0.6)

    # Add gradient-like effect with alpha
    for bar, color in zip(bars, colors):
        bar.set_alpha(0.85)

    ax1.set_xlabel('Computational Load (%)', fontsize=13, fontweight='bold', color='#2c3e50')
    ax1.set_xlim(0, 115)
    ax1.set_facecolor('#f8f9fa')

    # Value labels on bars
    for bar, val in zip(bars, compute):
        ax1.text(val + 1.5, bar.get_y() + bar.get_height()/2,
                f'{val:.1f}%', va='center', fontsize=14, fontweight='bold', color='#2c3e50')

    # Dramatic savings callout
    ax1.text(57, 1.2, f'{energy_saved:.1f}%',
            ha='center', fontsize=32, fontweight='bold', color='#27ae60')
    ax1.text(57, 0.95, 'ENERGY SAVED',
            ha='center', fontsize=11, fontweight='bold', color='#27ae60', alpha=0.8)

    # Add arrow showing reduction
    ax1.annotate('', xy=(compute[1], 0.5), xytext=(compute[0], 0.5),
                arrowprops=dict(arrowstyle='->', lw=3, color='#f39c12', alpha=0.7))

    ax1.set_title('Computational Efficiency Comparison', fontsize=15, fontweight='bold',
                 pad=15, color='#2c3e50')
    ax1.spines['top'].set_visible(False)
    ax1.spines['right'].set_visible(False)
    ax1.spines['left'].set_linewidth(1.5)
    ax1.spines['bottom'].set_linewidth(1.5)

    # Bottom left - Battery life impact
    ax2 = fig2.add_subplot(gs[1, 0])
    battery_baseline = 24  # hours
    battery_sundew = battery_baseline / gating_meta.get("ratio", 1.0)
    battery_data = [battery_baseline, battery_sundew]
    battery_colors = ['#e74c3c', '#2ecc71']

    bars2 = ax2.bar(['Baseline', 'Sundew'], battery_data, color=battery_colors,
                    edgecolor='white', linewidth=2, alpha=0.85)
    ax2.set_ylabel('Battery Life (hours)', fontsize=11, fontweight='bold')
    ax2.set_ylim(0, max(battery_data) * 1.2)
    ax2.set_title('Battery Life Extension', fontsize=12, fontweight='bold', color='#2c3e50')

    for bar, val in zip(bars2, battery_data):
        height = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width()/2, height + 5,
                f'{val:.0f}h', ha='center', fontsize=11, fontweight='bold')

    ax2.set_facecolor('#f8f9fa')
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)

    # Bottom right - Cost & environmental impact
    ax3 = fig2.add_subplot(gs[1, 1])
    ax3.axis('off')
    ax3.set_facecolor('#f8f9fa')

    # Impact metrics
    daily_savings_kwh = 0.024 * (energy_saved / 100)  # Assume 24Wh baseline daily consumption
    annual_co2_kg = daily_savings_kwh * 365 * 0.5  # ~0.5 kg CO2 per kWh

    impact_text = f"""
Impact per Device (Annual):

Energy Saved: {daily_savings_kwh * 365:.1f} kWh
CO2 Reduced: {annual_co2_kg:.1f} kg
Cost Savings: ${daily_savings_kwh * 365 * 0.12:.2f}

Scaling to 10,000 devices:
Energy: {daily_savings_kwh * 365 * 10000 / 1000:.1f} MWh/year
CO2: {annual_co2_kg * 10000 / 1000:.1f} tonnes/year
    """

    ax3.text(0.1, 0.95, impact_text.strip(),
            fontsize=10, verticalalignment='top', fontfamily='monospace',
            bbox=dict(boxstyle='round,pad=0.8', facecolor='#fff3cd',
                     edgecolor='#f39c12', linewidth=2, alpha=0.9))

    fig2.patch.set_facecolor('white')
    fig2.suptitle('Sundew Energy Efficiency Analysis', fontsize=16, fontweight='bold',
                 y=0.98, color='#2c3e50')

    buf2 = io.BytesIO()
    fig2.savefig(buf2, format='png', dpi=150, bbox_inches='tight')
    plt.close(fig2)
    buf2.seek(0)
    energy_img = mpimg.imread(buf2)

    # Rule chain
    rule_md = f"""### Rule Chain Trace

**Neural Network Output:**
- Label: `{model_output.get('label')}` (Confidence: {model_output.get('score', 0.0):.3f})
- Estimated HR: `{model_output.get('hr')} bpm`

**Patient Context:**
- Age: {scenario['age']} years
- Prior Stroke: {'Yes' if scenario['has_prior_stroke'] else 'No'}

**Rules Evaluated:**
"""
    for exp in rules_result.get('explanations', []):
        rule_md += f"\n- {exp}"

    rule_md += f"\n\n**Final Alert:** `{rules_result.get('alert_level', 'NONE').upper()}`"

    return summary_html, signal_img, energy_img, rule_md


# Build Gradio Interface
with gr.Blocks(title="Sundew ECG Monitor") as demo:

    # Header
    gr.HTML("""
    <div style="text-align: center; padding: 30px 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; margin-bottom: 20px;">
        <h1 style="margin: 0; font-size: 42px; font-weight: 800;">Sundew ECG Monitor</h1>
        <p style="margin: 10px 0 0 0; font-size: 18px; opacity: 0.95;">Neurosymbolic AI for Energy-Efficient Medical Monitoring</p>
        <div style="margin-top: 15px; display: inline-flex; gap: 20px; flex-wrap: wrap; justify-content: center;">
            <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px;">⚡ 85% Energy Savings</span>
            <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px;">🧠 Explainable AI</span>
            <span style="background: rgba(255,255,255,0.2); padding: 8px 16px; border-radius: 20px;">🏥 Clinical-Grade Rules</span>
        </div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Select Patient Scenario")
            scenario_dropdown = gr.Radio(
                choices=list(SCENARIOS.keys()),
                value="afib_high_risk",
                label="",
                info="Choose a patient to analyze"
            )

            for key, val in SCENARIOS.items():
                gr.Markdown(f"**{val['icon']} {val['name']}**\n{val['description']}", visible=(key=="afib_high_risk"))

            run_btn = gr.Button("Run Analysis", variant="primary", size="lg")

            gr.Markdown("---")
            gr.Markdown("""
            **Architecture:**
            ```
            ECG Signal → Sundew Gating → ML Inference → Rule Engine
                         (50-90% reduction)  (PyTorch)    (Symbolic)
            ```
            """)

        with gr.Column(scale=2):
            summary_card = gr.HTML()

            with gr.Tabs():
                with gr.Tab("📊 Signal Analysis"):
                    signal_plot = gr.Image(label="ECG: Original vs Gated")

                with gr.Tab("⚡ Energy Efficiency"):
                    energy_plot = gr.Image(label="Compute Savings")

                with gr.Tab("🔗 Rule Chain"):
                    rule_trace = gr.Markdown()

    run_btn.click(
        run_pipeline,
        inputs=scenario_dropdown,
        outputs=[summary_card, signal_plot, energy_plot, rule_trace]
    )

    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 20px; margin-top: 30px; border-top: 1px solid #eee;">
        <p style="color: #666; font-size: 14px;">
            Built with Sundew Algorithm · FastAPI · PyTorch · Gradio
        </p>
    </div>
    """)

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