#!/usr/bin/env python3
"""
NQR-SNN Detector — Gradio UI
==============================
Drag-and-drop any file format → ¹⁴N compound detection.

Supports: CSV, NPY, NPZ, MAT, HDF5, WAV, JSON, Parquet, Excel,
          binary IQ (cf32/cs16/cu8), SigMF, JCAMP-DX, Bruker FID,
          plain text, and more.

Usage:
    python app.py                    # Launch locally
    python app.py --share            # Public URL
    python app.py --port 7861        # Custom port
"""

import os
import sys
import tempfile
import time
import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from nqr_snn import config
from nqr_snn.data.ingest import load_signals, detect_format, PARSERS
from nqr_snn.data.dataset import extract_features_batch
from nqr_snn.snn.model import SpikingClassifier
from nqr_snn.snn.ensemble import SNNEnsemble

# -- Global detector (loaded once) --
DETECTOR = None


def get_detector():
    """Lazy-load the detector."""
    global DETECTOR
    if DETECTOR is None:
        from nqr_snn.inference import NQRDetector
        DETECTOR = NQRDetector(ensemble_size=config.ENSEMBLE_SIZE)
    return DETECTOR


# -- Supported file extensions for Gradio file picker --
NQR_FILE_EXTENSIONS = [
    # Scientific
    ".npy", ".npz", ".mat", ".h5", ".hdf5",
    # Tabular
    ".csv", ".tsv", ".json", ".parquet", ".xlsx", ".xls",
    ".txt", ".dat", ".asc",
    # RF/SDR
    ".bin", ".raw", ".iq", ".cf32", ".cf64", ".cs16", ".cu8",
    # Audio
    ".wav",
    # NQR-specific
    ".sigmf", ".jdx", ".dx", ".fid",
]


def process_file(filepath, threshold, format_override):
    """Main processing function for Gradio UI.
    
    Args:
        filepath: Uploaded file path (from gr.File).
        threshold: Detection threshold slider value.
        format_override: Format override dropdown value.
    
    Returns:
        Tuple of (verdict_text, results_table, metadata_json, plot_or_none)
    """
    if filepath is None:
        return "⏳ Upload a file to begin detection.", None, "{}", None
    
    t0 = time.time()
    
    # Parse format override
    fmt_override = None
    if format_override and format_override != "Auto-detect":
        fmt_override = format_override.lower().replace(" ", "_")
    
    try:
        # Load signals
        signals, metadata = load_signals(
            filepath, 
            format_override=fmt_override,
        )
    except Exception as e:
        return f"❌ **Parse Error:** {str(e)}", None, "{}", None
    
    t_parse = time.time() - t0
    
    # Feature extraction
    t1 = time.time()
    features = extract_features_batch(signals)
    t_feat = time.time() - t1
    
    # Run ensemble via detector (uses trained encoder)
    t2 = time.time()
    detector = get_detector()
    
    import torch
    feat_tensor = torch.from_numpy(features)
    # v3.2: Use the detector's encoder (LearnableTemporalEncoder with trained weights)
    x_seq = detector.encoder.encode(feat_tensor).to(detector.device)
    
    mean_p, std_p = detector.ensemble.predict(x_seq)
    probs = mean_p.cpu().numpy()
    stds = std_p.cpu().numpy()
    t_pred = time.time() - t2
    
    total_time = time.time() - t0
    
    # Build results
    detections = (probs >= threshold).astype(int)
    n_total = len(signals)
    n_detected = int(detections.sum())
    
    # Verdict
    if n_detected > 0:
        pct = n_detected / n_total * 100
        verdict = (
            f"## ⚠️ ¹⁴N COMPOUND DETECTED\n\n"
            f"**{n_detected}/{n_total}** signals ({pct:.1f}%) contain nitrogen-based compound signatures.\n\n"
            f"Mean confidence: **{probs[detections == 1].mean():.1%}** | "
            f"Mean uncertainty: **{stds[detections == 1].mean():.4f}**"
        )
    else:
        verdict = (
            f"## ✅ NO ¹⁴N COMPOUND DETECTED\n\n"
            f"All **{n_total}** signals classified as noise-only.\n\n"
            f"Mean probability: **{probs.mean():.4f}** | "
            f"Mean uncertainty: **{stds.mean():.4f}**"
        )
    
    # Results table
    results_df = pd.DataFrame({
        'Signal #': range(1, n_total + 1),
        'Detection': ['⚠️ DETECTED' if d else '✅ Clear' for d in detections],
        'Probability': [f"{p:.4f}" for p in probs],
        'Uncertainty': [f"{s:.4f}" for s in stds],
    })
    
    # Metadata
    meta_display = {
        "File": metadata.get('filename', 'unknown'),
        "Format detected": metadata.get('format', 'unknown'),
        "File size": f"{metadata.get('file_size_bytes', 0):,} bytes",
        "Original samples": f"{metadata.get('original_samples', 0):,}",
        "Signals extracted": n_total,
        "Signal length": config.SIGNAL_LENGTH,
        "Threshold": threshold,
        "Parse time": f"{t_parse:.3f}s",
        "Feature extraction": f"{t_feat:.3f}s",
        "Prediction time": f"{t_pred:.3f}s",
        "Total time": f"{total_time:.3f}s",
        "Throughput": f"{n_total / max(total_time, 0.001):.1f} signals/sec",
        "Model": f"CNN+SNN ensemble ({detector.ensemble_size} members)",
        "Device": detector.device,
    }
    
    # Generate signal visualization
    plot = generate_signal_plot(signals, probs, detections, metadata)
    
    return verdict, results_df, meta_display, plot


def generate_signal_plot(signals, probs, detections, metadata):
    """Generate a matplotlib figure showing the first few signals."""
    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        
        n_show = min(4, len(signals))
        fig, axes = plt.subplots(n_show, 2, figsize=(12, 3 * n_show))
        if n_show == 1:
            axes = axes.reshape(1, 2)
        
        for i in range(n_show):
            sig = signals[i]
            
            # Time domain
            ax = axes[i, 0]
            ax.plot(np.abs(sig), linewidth=0.5, color='steelblue', alpha=0.8)
            ax.set_ylabel(f'Signal {i+1}')
            if i == 0:
                ax.set_title('Magnitude (time domain)')
            if i == n_show - 1:
                ax.set_xlabel('Sample')
            
            det_str = '⚠️ DETECTED' if detections[i] else '✅ Clear'
            ax.text(0.98, 0.95, f'{det_str} (p={probs[i]:.3f})',
                    transform=ax.transAxes, ha='right', va='top',
                    fontsize=9, fontweight='bold',
                    color='red' if detections[i] else 'green',
                    bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
            
            # Frequency domain
            ax = axes[i, 1]
            fft_mag = np.abs(np.fft.fft(sig))[:len(sig)//2]
            ax.plot(fft_mag, linewidth=0.5, color='coral', alpha=0.8)
            if i == 0:
                ax.set_title('FFT Magnitude (frequency domain)')
            if i == n_show - 1:
                ax.set_xlabel('Frequency bin')
        
        fig.suptitle(f"NQR Signal Analysis — {metadata.get('filename', 'uploaded file')}", 
                     fontsize=13, fontweight='bold')
        plt.tight_layout()
        return fig
    except Exception:
        return None


def create_demo():
    """Build the Gradio interface."""
    import gradio as gr
    
    with gr.Blocks(
        title="NQR-SNN Detector — ¹⁴N Compound Detection",
        theme=gr.themes.Soft(),
    ) as demo:
        
        gr.Markdown("""
        # 🔬 NQR-SNN Detector v3.2
        ### Nuclear Quadrupole Resonance — ¹⁴N Compound Detection
        
        **Drag & drop any signal file** to detect nitrogen-based explosives/narcotics 
        (RDX, TNT, HMX, PETN, cocaine, heroin) from NQR receiver RF signals.
        
        Supports: **CSV, NPY, NPZ, MAT, HDF5, WAV, JSON, Parquet, Excel, 
        binary IQ (cf32/cs16/cu8), SigMF, JCAMP-DX, Bruker FID, plain text**
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                file_input = gr.File(
                    label="📡 Upload NQR/RF Signal File",
                    file_count="single",
                    file_types=NQR_FILE_EXTENSIONS,
                    type="filepath",
                    height=120,
                )
            
            with gr.Column(scale=1):
                threshold_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.5, step=0.05,
                    label="Detection Threshold",
                    info="Lower = more sensitive (more detections), Higher = more specific"
                )
                
                format_dropdown = gr.Dropdown(
                    choices=["Auto-detect"] + sorted([
                        "CSV", "JSON", "NPY", "NPZ", "HDF5", "MATLAB",
                        "WAV", "Parquet", "Excel", "Text",
                        "Raw CF32", "Raw CS16", "Raw CU8",
                        "SigMF", "JCAMP", "Bruker FID",
                    ]),
                    value="Auto-detect",
                    label="Format Override",
                    info="Usually auto-detect works. Override if needed."
                )
                
                detect_btn = gr.Button("🔍 Detect", variant="primary", size="lg")
        
        # Results
        verdict_output = gr.Markdown(
            value="⏳ Upload a file to begin detection.",
            label="Verdict"
        )
        
        with gr.Row():
            with gr.Column(scale=2):
                results_table = gr.DataFrame(
                    label="Signal-by-Signal Results",
                    interactive=False,
                )
            with gr.Column(scale=1):
                metadata_output = gr.JSON(label="Analysis Metadata")
        
        signal_plot = gr.Plot(label="Signal Visualization")
        
        # Wire up
        detect_btn.click(
            fn=process_file,
            inputs=[file_input, threshold_slider, format_dropdown],
            outputs=[verdict_output, results_table, metadata_output, signal_plot],
        )
        
        file_input.upload(
            fn=process_file,
            inputs=[file_input, threshold_slider, format_dropdown],
            outputs=[verdict_output, results_table, metadata_output, signal_plot],
        )
        
        # Examples
        gr.Markdown("""
        ---
        ### Supported Data Formats
        
        | Category | Formats | Notes |
        |----------|---------|-------|
        | **Scientific** | `.npy`, `.npz`, `.mat`, `.h5`/`.hdf5` | Auto-detects I/Q keys |
        | **Tabular** | `.csv`, `.tsv`, `.json`, `.parquet`, `.xlsx` | Columns: `I`+`Q`, `real`+`imag`, or single array |
        | **RF/SDR** | `.bin`, `.cf32`, `.cs16`, `.cu8`, `.iq` | Interleaved binary (GNU Radio, RTL-SDR, HackRF) |
        | **Audio** | `.wav` | Stereo: L=I, R=Q |
        | **NQR** | `.sigmf`, `.jdx`/`.dx`, `.fid` | SigMF archive, JCAMP-DX, Bruker FID |
        | **Text** | `.txt`, `.dat`, `.asc` | Whitespace/comma-separated numbers |
        
        **Signal handling:** Files with >1024 samples are split into segments. 
        Files with <1024 samples are zero-padded. Multiple signals per file supported.
        """)
    
    return demo


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()
    
    demo = create_demo()
    demo.launch(server_port=args.port, share=args.share)