# PSD Analysis
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from scipy import signal
import gradio as gr
import lcmv_stats as ls
import logging
from typing import Dict, List, Tuple

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# =============================================================================
# 1. CONFIGURATION & CONSTANTS
# =============================================================================

DATA_FILES = {
    "drug": "study_ocd_drug.npz",
    "neutral": "study_ocd_neutral.npz",
    "positive": "study_ocd_positive.npz",
}

PSD_WINDOW_SECONDS: float = 4.0
PSD_OVERLAP_FRACTION: float = 0.75
PSD_EPSILON: float = 1e-15
REFERENCE_BAND_HZ: Tuple[float, float] = (1.0, 4.0)
FREQ_MAX_PLOT_HZ: float = 50.0

PLOT_BANDS = [
    (1, 4, 'Delta', '#90B3F9'),
    (4, 8, 'Theta', '#FFF9B2'),
    (8, 13, 'Alpha', '#AAFCD2'),
    (13, 20, 'Low Beta', '#97C2F9'),
    (20, 30, 'High Beta', '#90BEF5'),
]

BAND_OPTIONS = {
    'Delta (1-4 Hz)': (1, 4),
    'Theta (4-8 Hz)': (4, 8),
    'Alpha (8-13 Hz)': (8, 13),
    'Low Beta (13-20 Hz)': (13, 20),
    'High Beta (20-30 Hz)': (20, 30),
    'Low Gamma (30-50 Hz)': (30, 50),
}

CONDITION_COLORS = {
    'Drug': '#D62728',
    'Neutral': '#1F77B4',
    'Positive': '#2CA02C',
}

# =============================================================================
# 2. DATA LOADING & ATLAS MANAGEMENT
# =============================================================================

def load_study_data() -> Tuple[dict, dict, dict, float]:
    """Load NPZ data files and extract sampling frequency."""
    drug = ls.load_tensor(DATA_FILES["drug"])
    neutral = ls.load_tensor(DATA_FILES["neutral"])
    positive = ls.load_tensor(DATA_FILES["positive"])

    sfreq = float(drug['sfreq'])
    assert sfreq > 0, "Sampling frequency must be positive"
    return drug, neutral, positive, sfreq


def build_cascading_roi_map(atlas_df: pd.DataFrame) -> Tuple[Dict[str, List[str]], Dict[str, int]]:
    """
    Parse atlas DataFrame into cascading dropdown structures.

    Args:
        atlas_df: pd.DataFrame with columns [index, region_full_name, hemisphere, functional_system]

    Returns:
        system_to_rois: {SystemName: [SortedDisplayLabels]}
        label_to_index: {DisplayLabel: DataIndex}
    """
    assert all(col in atlas_df.columns for col in ['index', 'region_full_name', 'hemisphere', 'functional_system']), \
        "Atlas missing required columns"

    atlas_df = atlas_df.copy()
    atlas_df['display_label'] = atlas_df['region_full_name'] + " (" + atlas_df['hemisphere'].str[0] + ")"

    system_to_rois: Dict[str, List[str]] = {}
    for system in sorted(atlas_df['functional_system'].unique()):
        labels = atlas_df[atlas_df['functional_system'] == system]['display_label'].tolist()
        system_to_rois[system] = sorted(labels)

    label_to_index: Dict[str, int] = dict(
        zip(atlas_df['display_label'], atlas_df['index'].astype(int))
    )

    logger.info(f"Built cascading map: {len(system_to_rois)} systems, {len(label_to_index)} ROIs")
    return system_to_rois, label_to_index


def get_default_roi_state(
    system_to_rois: Dict[str, List[str]],
    label_to_index: Dict[str, int]
) -> Tuple[str, str, int]:
    """Return (default_system, default_roi_label, default_roi_index)."""
    systems = sorted(system_to_rois.keys())
    assert len(systems) > 0, "No functional systems found in atlas"
    default_system = systems[0]
    rois = system_to_rois[default_system]
    assert len(rois) > 0, f"No ROIs found in system '{default_system}'"
    default_roi = rois[0]
    default_index = label_to_index[default_roi]
    return default_system, default_roi, default_index

# =============================================================================
# 3. CORE COMPUTATION
# =============================================================================

def compute_aligned_psd(
    sig_d: np.ndarray,
    sig_n: np.ndarray,
    sig_p: np.ndarray,
    sfreq: float,
    ref_band: Tuple[float, float] = REFERENCE_BAND_HZ,
    nperseg_seconds: float = PSD_WINDOW_SECONDS,
    overlap_fraction: float = PSD_OVERLAP_FRACTION,
    eps: float = PSD_EPSILON
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Compute Welch PSD with delta-band alignment to neutral baseline.

    Returns:
        freqs, psd_d_db, psd_n_db, psd_p_db
    """
    nperseg = min(int(sfreq * nperseg_seconds), len(sig_d))
    noverlap = int(nperseg * overlap_fraction)

    freqs, psd_d_lin = signal.welch(sig_d, fs=sfreq, nperseg=nperseg, noverlap=noverlap, window='hann')
    _, psd_n_lin = signal.welch(sig_n, fs=sfreq, nperseg=nperseg, noverlap=noverlap, window='hann')
    _, psd_p_lin = signal.welch(sig_p, fs=sfreq, nperseg=nperseg, noverlap=noverlap, window='hann')

    mask_ref = (freqs >= ref_band[0]) & (freqs <= ref_band[1])
    ref_lin = psd_n_lin[mask_ref].mean()
    mean_d_ref = psd_d_lin[mask_ref].mean()
    mean_p_ref = psd_p_lin[mask_ref].mean()

    offset_d = 10 * np.log10((ref_lin + eps) / (mean_d_ref + eps))
    offset_p = 10 * np.log10((ref_lin + eps) / (mean_p_ref + eps))

    psd_d_db = 10 * np.log10(psd_d_lin + eps) + offset_d
    psd_n_db = 10 * np.log10(psd_n_lin + eps)
    psd_p_db = 10 * np.log10(psd_p_lin + eps) + offset_p

    return freqs, psd_d_db, psd_n_db, psd_p_db

# =============================================================================
# 4. VISUALIZATION
# =============================================================================

def build_psd_figure(
    freqs: np.ndarray,
    psd_d_db: np.ndarray,
    psd_n_db: np.ndarray,
    psd_p_db: np.ndarray,
    roi_label: str,
    freq_max: float = FREQ_MAX_PLOT_HZ
) -> go.Figure:
    """Construct the main PSD Plotly figure with band shading."""
    fig = go.Figure()

    for f_lo, f_hi, name, color in PLOT_BANDS:
        if f_hi <= freq_max:
            fig.add_vrect(x0=f_lo, x1=f_hi, fillcolor=color, opacity=0.08, layer="below", line_width=0)
            fig.add_annotation(
                x=(f_lo + f_hi) / 2, y=0.97, xref="x", yref="paper",
                text=f"<b>{name}</b>", showarrow=False,
                font=dict(size=10, color='#1E3A5F'), opacity=0.8
            )

    mask = freqs <= freq_max
    traces = [
        ('Drug / Craving Cue', psd_d_db, CONDITION_COLORS['Drug']),
        ('Neutral Cue (baseline)', psd_n_db, CONDITION_COLORS['Neutral']),
        ('Positive Cue', psd_p_db, CONDITION_COLORS['Positive']),
    ]

    for label, psd_db, color in traces:
        fig.add_trace(go.Scatter(
            x=freqs[mask], y=psd_db[mask], mode='lines',
            name=label, line=dict(color=color, width=2.5),
        ))

    fig.update_layout(
        legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99, font=dict(size=12)),
        template='plotly_white', margin=dict(t=80, b=60, l=70, r=30), height=500,
    )
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='rgba(0,0,0,0.08)')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='rgba(0,0,0,0.08)')
    return fig


def build_ratio_figure(
    roi_label: str,
    band_label: str,
    d_mean: float,
    n_mean: float,
    p_mean: float,
    eps: float = PSD_EPSILON
) -> go.Figure:
    """Construct horizontal bar chart for condition modulation indices."""
    comparisons = ['Positive', 'Neutral']

    ratio_dp = ((d_mean - p_mean) / (d_mean + p_mean + eps)) * 100
    ratio_dn = ((d_mean - n_mean) / (d_mean + n_mean + eps)) * 100

    values = [ratio_dp, ratio_dn]
    colors = [
        CONDITION_COLORS['Drug'] if ratio_dp >= 0 else CONDITION_COLORS['Positive'],
        CONDITION_COLORS['Drug'] if ratio_dn >= 0 else CONDITION_COLORS['Neutral']
    ]

    fig = go.Figure()
    fig.add_trace(go.Bar(
        y=comparisons, x=values, orientation='h', marker_color=colors,
        text=[f'{v:+.1f}%' for v in values], textposition='inside',
        textfont=dict(size=12, family='monospace', color='white'), insidetextanchor='middle',
    ))

    fig.add_annotation(x=1.02, y='Positive', xref='paper', yref='y',
                       text='<b>Drug</b>', showarrow=False, font=dict(size=11, color='#333'), xanchor='left')
    fig.add_annotation(x=1.02, y='Neutral', xref='paper', yref='y',
                       text='<b>Drug</b>', showarrow=False, font=dict(size=11, color='#333'), xanchor='left')

    fig.update_layout(
        xaxis=dict(tickfont=dict(size=10), zeroline=True, zerolinewidth=1,
                   zerolinecolor='#999', showgrid=True, gridwidth=1, gridcolor='rgba(0,0,0,0.06)'),
        yaxis=dict(tickfont=dict(size=11, weight='bold'), showgrid=False, zeroline=False, side='left'),
        template='plotly_white', height=200, margin=dict(t=50, b=30, l=80, r=60),
    )
    return fig

# =============================================================================
# 5. GRADIO CALLBACKS
# =============================================================================

def update_psd(
    roi_label: str,
    label_to_index: Dict[str, int],
    drug: dict,
    neutral: dict,
    positive: dict,
    sfreq: float
) -> go.Figure:
    """Callback for PSD plot update."""
    if roi_label not in label_to_index:
        raise ValueError(f"ROI label '{roi_label}' not found in index map")
    idx = label_to_index[roi_label]

    freqs, d_db, n_db, p_db = compute_aligned_psd(
        drug['data'][0, idx, :],
        neutral['data'][0, idx, :],
        positive['data'][0, idx, :],
        sfreq
    )
    return build_psd_figure(freqs, d_db, n_db, p_db, roi_label)


def update_ratio(
    roi_label: str,
    band_label: str,
    label_to_index: Dict[str, int],
    drug: dict,
    neutral: dict,
    positive: dict,
    sfreq: float
) -> go.Figure:
    """Callback for ratio plot update using delta-aligned power."""
    if roi_label not in label_to_index:
        raise ValueError(f"ROI label '{roi_label}' not found in index map")
    if band_label not in BAND_OPTIONS:
        raise ValueError(f"Band label '{band_label}' not found in BAND_OPTIONS")

    idx = label_to_index[roi_label]
    f_lo, f_hi = BAND_OPTIONS[band_label]

    freqs, d_db, n_db, p_db = compute_aligned_psd(
        drug['data'][0, idx, :],
        neutral['data'][0, idx, :],
        positive['data'][0, idx, :],
        sfreq
    )

    band_mask = (freqs >= f_lo) & (freqs <= f_hi)
    d_mean = float(np.mean(10 ** (d_db[band_mask] / 10)))
    n_mean = float(np.mean(10 ** (n_db[band_mask] / 10)))
    p_mean = float(np.mean(10 ** (p_db[band_mask] / 10)))

    return build_ratio_figure(roi_label, band_label, d_mean, n_mean, p_mean)


def on_system_change(system: str, system_to_rois: Dict[str, List[str]]) -> gr.Dropdown:
    """Update ROI dropdown choices when functional system changes."""
    rois = system_to_rois.get(system, [])
    new_default = rois[0] if rois else None
    return gr.update(choices=rois, value=new_default)

# =============================================================================
# 6. APP INITIALIZATION
# =============================================================================

drug_data, neutral_data, positive_data, sampling_freq = load_study_data()
atlas_df = ls.get_cimt_labels()
SYSTEM_TO_ROIS, LABEL_TO_INDEX = build_cascading_roi_map(atlas_df)
DEFAULT_SYS, DEFAULT_ROI, DEFAULT_IDX = get_default_roi_state(SYSTEM_TO_ROIS, LABEL_TO_INDEX)

initial_fig = update_psd(DEFAULT_ROI, LABEL_TO_INDEX, drug_data, neutral_data, positive_data, sampling_freq)
initial_ratio = update_ratio(DEFAULT_ROI, 'Theta (4-8 Hz)', LABEL_TO_INDEX, drug_data, neutral_data, positive_data, sampling_freq)

with gr.Blocks(title="OCD Full Atlas Explorer") as app:
    gr.Markdown("# OCD Cue-Reactivity: Full Atlas Explorer\n"
                "Interactive delta-aligned PSD analysis across Drug / Neutral / Positive conditions")

    with gr.Row():
        with gr.Column(scale=1):
            sys_dropdown = gr.Dropdown(
                choices=sorted(SYSTEM_TO_ROIS.keys()),
                value=DEFAULT_SYS,
                label="Functional System",
                info="Select brain network to filter ROIs"
            )
            roi_dropdown = gr.Dropdown(
                choices=SYSTEM_TO_ROIS[DEFAULT_SYS],
                value=DEFAULT_ROI,
                label="Region of Interest",
                info="Select specific anatomical region"
            )
            band_dropdown = gr.Dropdown(
                choices=list(BAND_OPTIONS.keys()),
                value='Theta (4-8 Hz)',
                label="Frequency Band",
                info="Band-averaged power comparison"
            )
            ratio_output = gr.Plot(label="Condition Comparison", value=initial_ratio)

        with gr.Column(scale=2):
            psd_plot = gr.Plot(label="Delta-Aligned PSD", value=initial_fig)

    sys_dropdown.change(
        fn=lambda s: on_system_change(s, SYSTEM_TO_ROIS),
        inputs=sys_dropdown,
        outputs=roi_dropdown
    )

    roi_dropdown.change(
        fn=lambda r: update_psd(r, LABEL_TO_INDEX, drug_data, neutral_data, positive_data, sampling_freq),
        inputs=roi_dropdown,
        outputs=psd_plot
    )

    roi_dropdown.change(
        fn=lambda r, b: update_ratio(r, b, LABEL_TO_INDEX, drug_data, neutral_data, positive_data, sampling_freq),
        inputs=[roi_dropdown, band_dropdown],
        outputs=ratio_output
    )

    band_dropdown.change(
        fn=lambda r, b: update_ratio(r, b, LABEL_TO_INDEX, drug_data, neutral_data, positive_data, sampling_freq),
        inputs=[roi_dropdown, band_dropdown],
        outputs=ratio_output
    )

if __name__ == "__main__":
    app.launch(
        theme=gr.themes.Soft(),
        css=".gradio-container { max-width: 1200px !important; }"
    )