import gradio as gr
import mne
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load open-source LLM (no training needed)
model_name = "tiiuae/falcon-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")

def process_eeg(file):
    # Load EEG data using MNE
    raw = mne.io.read_raw_fif(file.name, preload=True)
    # Compute some features (e.g., average band powers)
    psd, freqs = mne.time_frequency.psd_welch(raw, fmin=1, fmax=40)
    alpha_power = compute_band_power(psd, freqs, 8, 12)
    beta_power = compute_band_power(psd, freqs, 13, 30)
    # Create a human-readable summary of features
    data_summary = f"Alpha power: {alpha_power}, Beta power: {beta_power}. The data shows stable alpha rhythms and slightly elevated beta."
    
    # Prompt the LLM
    prompt = f"""You are a neuroscientist analyzing EEG features. 
    Data Summary: {data_summary}

    Provide a concise, user-friendly interpretation of these findings."""
    
    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(inputs, max_length=200, do_sample=True, top_k=50, top_p=0.95)
    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return summary

iface = gr.Interface(
    fn=process_eeg,
    inputs=gr.File(label="Upload your EEG data (FIF format)"),
    outputs="text",
    title="NeuroNarrative-Lite: EEG Summary",
    description="Upload EEG data to receive a text-based summary from an open-source LLM. No training required!"
)

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