"""UNI Task PSD Explorer: EC / EO / SM condition comparison."""
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from scipy import signal
import gradio as gr
import lcmv_xtra as lx
import logging
from pathlib import Path
from typing import Dict, List, Tuple

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

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

TENSOR_DIR = Path("./data")

CONDITION_LABELS = {
    "ec": "Eyes Closed",
    "eo": "Eyes Open",
    "sm": "Motor Task",
}

CONDITION_COLORS = {
    "ec": "#1F77B4",   # Blue
    "eo": "#2CA02C",   # Green
    "sm": "#D62728",   # Red
}

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 = 40.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),
}

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

def load_psd_cache() -> dict:
    """Load single precomputed PSD cache file."""
    cache_path = TENSOR_DIR / "psd_cache.npz"
    cache = np.load(cache_path, allow_pickle=True)
    logger.info(
        f"Loaded PSD cache: {len(cache['entries'])} entries × "
        f"{cache['n_rois']} ROIs × {len(cache['freqs'])} freq bins"
    )
    return cache


def build_cascading_roi_map(atlas_df: pd.DataFrame) -> Tuple[Dict[str, List[str]], Dict[str, int]]:
    """Parse CIMT atlas DataFrame into cascading dropdown structures."""
    required_cols = ['index', 'region_full_name', 'hemisphere', 'functional_system']
    assert all(col in atlas_df.columns for col in required_cols), \
        f"Atlas missing required columns: {set(required_cols) - set(atlas_df.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 (REMOVED — now served from cache)
# =============================================================================

# compute_aligned_psd is no longer needed at runtime.
# All PSDs are precomputed in psd_cache.npz.

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

def build_psd_figure(
    freqs: np.ndarray,
    psd_ec_db: np.ndarray,
    psd_eo_db: np.ndarray,
    psd_sm_db: np.ndarray,
    roi_label: str,
    freq_max: float = FREQ_MAX_PLOT_HZ
) -> go.Figure:
    """Construct PSD Plotly figure with band shading (original visual style)."""
    fig = go.Figure()

    # Band shading with annotations (identical to original)
    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 = [
        (CONDITION_LABELS["ec"], psd_ec_db, CONDITION_COLORS["ec"]),
        (CONDITION_LABELS["eo"], psd_eo_db, CONDITION_COLORS["eo"]),
        (CONDITION_LABELS["sm"], psd_sm_db, CONDITION_COLORS["sm"]),
    ]

    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),
        ))

    # Identical layout to original
    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,
    ec_mean: float,
    eo_mean: float,
    sm_mean: float,
    eps: float = PSD_EPSILON
) -> go.Figure:
    """Horizontal bar chart: modulation index relative to EC baseline (original visual style)."""
    comparisons = ['Motor Task', 'Eyes Open']

    ratio_sm = ((sm_mean - ec_mean) / (sm_mean + ec_mean + eps)) * 100
    ratio_eo = ((eo_mean - ec_mean) / (eo_mean + ec_mean + eps)) * 100

    values = [ratio_sm, ratio_eo]
    colors = [
        CONDITION_COLORS["sm"] if ratio_sm >= 0 else CONDITION_COLORS["ec"],
        CONDITION_COLORS["eo"] if ratio_eo >= 0 else CONDITION_COLORS["ec"],
    ]

    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',
    ))

    # Baseline reference annotations (mirrors original Drug annotation style)
    fig.add_annotation(x=1.02, y='MT', xref='paper', yref='y',
                       text='<b>EC</b>', showarrow=False, font=dict(size=11, color='#333'), xanchor='left')
    fig.add_annotation(x=1.02, y='EO', xref='paper', yref='y',
                       text='<b>EC</b>', showarrow=False, font=dict(size=11, color='#333'), xanchor='left')

    # Identical axis/layout styling to original
    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 (ZERO COMPUTATION — pure cache lookup)
# =============================================================================

def update_psd(roi_label, entry_name, label_to_index, cache):
    """Callback for PSD plot update from precomputed cache."""
    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]
    entry_idx = np.where(cache['entries'] == entry_name)[0][0]

    freqs = cache['freqs']
    ec_db = np.nan_to_num(cache['ec_db'][entry_idx, idx, :], nan=0.0)
    eo_db = np.nan_to_num(cache['eo_db'][entry_idx, idx, :], nan=0.0)
    sm_db = np.nan_to_num(cache['sm_db'][entry_idx, idx, :], nan=0.0)

    return build_psd_figure(freqs, ec_db, eo_db, sm_db, roi_label)


def update_ratio(roi_label, band_label, entry_name, label_to_index, cache):
    """Callback for ratio plot from precomputed cache."""
    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]
    entry_idx = np.where(cache['entries'] == entry_name)[0][0]

    freqs = cache['freqs']
    ec_db = np.nan_to_num(cache['ec_db'][entry_idx, idx, :], nan=0.0)
    eo_db = np.nan_to_num(cache['eo_db'][entry_idx, idx, :], nan=0.0)
    sm_db = np.nan_to_num(cache['sm_db'][entry_idx, idx, :], nan=0.0)

    band_mask = (freqs >= f_lo) & (freqs <= f_hi)
    ec_mean = float(np.mean(10 ** (ec_db[band_mask] / 10)))
    eo_mean = float(np.mean(10 ** (eo_db[band_mask] / 10)))
    sm_mean = float(np.mean(10 ** (sm_db[band_mask] / 10)))

    return build_ratio_figure(roi_label, band_label, ec_mean, eo_mean, sm_mean)


def on_system_change(system, system_to_rois):
    """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
# =============================================================================

def create_app():
    """Build and return the Gradio Blocks app. Importable entry point."""
    cache = load_psd_cache()

    # Load CIMT labels from bundled atlas
    import lcmv_xtra
    labels_path = Path(lcmv_xtra.__file__).parent / 'data' / 'cimt_atlas' / 'cimt_atlas_labels.csv'
    atlas_df = pd.read_csv(labels_path)

    SYSTEM_TO_ROIS, LABEL_TO_INDEX = build_cascading_roi_map(atlas_df)
    DEFAULT_SYS, DEFAULT_ROI, _ = get_default_roi_state(SYSTEM_TO_ROIS, LABEL_TO_INDEX)

    entries = list(cache['entries'])  # ["Group Average", "sub-01", "sub-02", ...]

    initial_fig = update_psd(DEFAULT_ROI, "Group Average", LABEL_TO_INDEX, cache)
    initial_ratio = update_ratio(DEFAULT_ROI, 'Alpha (8-13 Hz)', "Group Average", LABEL_TO_INDEX, cache)

    with gr.Blocks(title="UNI Task Atlas Explorer") as app:
        gr.Markdown(
            "# UNI Task: Full Atlas PSD Explorer\n"
            "Interactive delta-aligned PSD analysis across Eyes Closed / Eyes Open / Motor Task conditions"
        )

        with gr.Row():
            with gr.Column(scale=1):
                subject_dropdown = gr.Dropdown(
                    choices=entries,
                    value="Group Average",
                    label="Subject",
                    info="Select individual subject or group average"
                )
                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='Alpha (8-13 Hz)',
                    label="Frequency Band",
                    info="Band-averaged power comparison"
                )
                ratio_output = gr.Plot(label="Condition Modulation", 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, subj: update_psd(r, subj, LABEL_TO_INDEX, cache),
            inputs=[roi_dropdown, subject_dropdown], outputs=psd_plot
        )
        subject_dropdown.change(
            fn=lambda r, subj: update_psd(r, subj, LABEL_TO_INDEX, cache),
            inputs=[roi_dropdown, subject_dropdown], outputs=psd_plot
        )
        roi_dropdown.change(
            fn=lambda r, b, subj: update_ratio(r, b, subj, LABEL_TO_INDEX, cache),
            inputs=[roi_dropdown, band_dropdown, subject_dropdown], outputs=ratio_output
        )
        band_dropdown.change(
            fn=lambda r, b, subj: update_ratio(r, b, subj, LABEL_TO_INDEX, cache),
            inputs=[roi_dropdown, band_dropdown, subject_dropdown], outputs=ratio_output
        )
        subject_dropdown.change(
            fn=lambda r, b, subj: update_ratio(r, b, subj, LABEL_TO_INDEX, cache),
            inputs=[roi_dropdown, band_dropdown, subject_dropdown], outputs=ratio_output
        )

    return app


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