import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import pandas as pd

# EEG + tACS simulation functions
def generate_eeg_signal(duration=5, sampling_rate=250):
    """Generate simulated EEG signal with alpha waves"""
    t = np.linspace(0, duration, int(duration * sampling_rate), endpoint=False)
    alpha_wave = 0.5 * np.sin(2 * np.pi * 10 * t)  # 10 Hz alpha rhythm
    noise = 0.2 * np.random.normal(size=len(t))
    return t, alpha_wave + noise

def apply_tacs(eeg_signal, tacs_freq, tacs_amp, tacs_phase, sampling_rate=250):
    """Apply tACS to EEG signal"""
    t = np.linspace(0, len(eeg_signal)/sampling_rate, len(eeg_signal))
    tacs = tacs_amp * np.sin(2 * np.pi * tacs_freq * t + np.deg2rad(tacs_phase))
    return eeg_signal + tacs, tacs

def filter_signal(raw_signal, sampling_rate=250, lowcut=1, highcut=40):
    """Apply bandpass filter to EEG signal"""
    nyquist = 0.5 * sampling_rate
    low = lowcut / nyquist
    high = highcut / nyquist
    b, a = signal.butter(4, [low, high], btype='band')
    return signal.filtfilt(b, a, raw_signal)

def plot_signals(t, eeg, tacs, combined, title):
    """Create matplotlib plot of signals"""
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(t, eeg, label='EEG Original', alpha=0.7)
    ax.plot(t, tacs, label='tACS Stimulation', alpha=0.7)
    ax.plot(t, combined, label='EEG + tACS', linewidth=1.5)
    ax.set_title(title, fontsize=14)
    ax.set_xlabel('Time (s)', fontsize=12)
    ax.set_ylabel('Amplitude (μV)', fontsize=12)
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.legend()
    ax.set_xlim(0, 2)  # Show first 2 seconds for clarity
    plt.tight_layout()
    return fig

def simulate_tacs(tacs_frequency, tacs_amplitude, tacs_phase, show_fft):
    # Generate EEG signal
    t, eeg = generate_eeg_signal()
    
    # Apply tACS stimulation
    combined, tacs_wave = apply_tacs(eeg, tacs_frequency, tacs_amplitude, tacs_phase)
    
    # Filter the combined signal
    filtered = filter_signal(combined)
    
    # Create plots
    time_plot = plot_signals(t, eeg, tacs_wave, combined, 
                            f"EEG + tACS Simulation | {tacs_frequency}Hz, {tacs_amplitude}μA")
    
    # Optional FFT analysis
    if show_fft:
        fft_plot = plot_fft(combined, tacs_frequency)
        return time_plot, fft_plot, filtered[:250]  # Return first second of filtered signal
    
    return time_plot, None, filtered[:250]

def plot_fft(signal, stim_freq, fs=250):
    """Plot frequency spectrum of signal"""
    n = len(signal)
    freq = np.fft.rfftfreq(n, d=1/fs)
    mag = np.abs(np.fft.rfft(signal)) / n
    
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(freq, mag)
    ax.axvline(x=stim_freq, color='r', linestyle='--', label=f'Stim Freq: {stim_freq}Hz')
    ax.set_title('Frequency Spectrum (FFT)', fontsize=14)
    ax.set_xlabel('Frequency (Hz)', fontsize=12)
    ax.set_ylabel('Magnitude', fontsize=12)
    ax.set_xlim(0, 50)
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.legend()
    plt.tight_layout()
    return fig

# Gradio application
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 EEG + tACS Simulator")
    gr.Markdown("Simulieren Sie die Wirkung von transkranieller Wechselstromstimulation (tACS) auf EEG-Signale")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## ⚡ Stimulationsparameter")
            tacs_freq = gr.Slider(1, 100, value=10, step=0.5, 
                                 label="tACS Frequenz (Hz)", info="1-100 Hz")
            tacs_amp = gr.Slider(0.1, 5, value=1, step=0.1, 
                                label="Stromstärke (μA)", info="0.1-5 μA")
            tacs_phase = gr.Slider(0, 360, value=0, step=15, 
                                  label="Phase (Grad)", info="0-360°")
            show_fft = gr.Checkbox(label="Frequenzspektrum anzeigen (FFT)")
            btn = gr.Button("Simulation starten", variant="primary")
            
            gr.Markdown("### ℹ️ Über tACS")
            gr.Markdown("""
            **Transkranielle Wechselstromstimulation (tACS)** ist eine nicht-invasive Methode zur Modulation neuronaler Aktivität:
            - Beeinflusst die Gehirnrhythmen durch externe Wechselströme
            - Typische Frequenzen: Alpha (8-12 Hz), Beta (12-30 Hz), Gamma (30-100 Hz)
            - Anwendungen: Kognitionsverbesserung, Stimmungsregulation, Schmerztherapie
            """)
            
        with gr.Column(scale=2):
            gr.Markdown("## 📊 Simulationsergebnisse")
            time_plot = gr.Plot(label="Zeitbereich")
            fft_plot = gr.Plot(label="Frequenzbereich", visible=False)
            signal_output = gr.Dataframe(
                label="Gefiltertes Signal (erste Sekunde)",
                headers=["Zeitindex", "Amplitude (μV)"],
                datatype=["number", "number"],
                max_rows=10
            )
            
            gr.Markdown("### 💡 Interpretation")
            gr.Markdown("""
            - **Rote Linie**: tACS Stimulationsfrequenz
            - **Blaue Linie**: Natürliche EEG-Aktivität (Alpha-Wellen ~10Hz)
            - **Grüne Linie**: Kombinierter Effekt von EEG + tACS
            """)
    
    # Built with anycoder
    gr.Markdown(
        "[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)",
        elem_classes=["footer"]
    )
    
    # Event handling
    btn.click(
        simulate_tacs,
        inputs=[tacs_freq, tacs_amp, tacs_phase, show_fft],
        outputs=[time_plot, fft_plot, signal_output]
    )
    
    show_fft.change(
        lambda x: gr.update(visible=x),
        inputs=show_fft,
        outputs=fft_plot
    )

# Launch with modern theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="indigo",
        font=[gr.themes.GoogleFont("Montserrat"), "sans-serif"]
    ),
    css=".footer {text-align: center; margin-top: 1em; font-size: 0.8em}",
    footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
)