"""t-SNE tab for Analysis plugin."""

from ._shared import (
    fac, html, dcc, go, px, pd, np, logger, os,
    TSNE_COMPONENT_OPTIONS, PLOTLY_HIGH_RES_CONFIG,
    get_physical_cores, create_invisible_figure, dash,
    Input, Output, PreventUpdate
)
from sklearn.manifold import TSNE


def create_layout():
    """Return the t-SNE tab layout component."""
    return html.Div(
        [
            # t-SNE-specific controls
            fac.AntdFlex(
                [
                    fac.AntdSpace(
                        [
                            html.Span("X axis:", style={'fontWeight': 500}),
                            fac.AntdSelect(
                                id='tsne-x-comp',
                                options=TSNE_COMPONENT_OPTIONS,
                                value='t-SNE-1',
                                allowClear=False,
                                style={'width': 100},
                            ),
                        ],
                        align='center',
                        size='small',
                    ),
                    fac.AntdSpace(
                        [
                            html.Span("Y axis:", style={'fontWeight': 500}),
                            fac.AntdSelect(
                                id='tsne-y-comp',
                                options=TSNE_COMPONENT_OPTIONS,
                                value='t-SNE-2',
                                allowClear=False,
                                style={'width': 100},
                            ),
                        ],
                        align='center',
                        size='small',
                    ),
                    fac.AntdDivider(direction='vertical', style={'height': '24px', 'margin': '0 12px'}),
                    fac.AntdSpace(
                        [
                            fac.AntdText("Perplexity:", style={'fontWeight': 500}),
                            fac.AntdSlider(
                                id='tsne-perplexity-slider',
                                min=1,
                                max=100,
                                step=1,
                                value=30,
                                style={'width': '200px'},
                                tooltipPrefix='Perplexity: '
                            ),
                        ],
                        align='center',
                        size='small',
                    ),
                    fac.AntdButton(
                        "Generate",
                        id="tsne-regenerate-btn",
                        icon=fac.AntdIcon(icon="antd-sync"),
                        style={'textTransform': 'uppercase'},
                    ),
                ],
                gap='middle',
                align='center',
                style={'marginBottom': '12px'},
            ),
            fac.AntdSpin(
                dcc.Graph(
                    id='tsne-graph',
                    config=PLOTLY_HIGH_RES_CONFIG,
                    style={'height': 'calc(100vh - 220px)', 'width': '80%', 'minHeight': '400px', 'margin': '0 auto'},
                    figure={
                        'data': [],
                        'layout': {
                            'xaxis': {'visible': False},
                            'yaxis': {'visible': False},
                            'paper_bgcolor': 'white',
                            'plot_bgcolor': 'white',
                            'margin': {'l': 0, 'r': 0, 't': 0, 'b': 0},
                            'autosize': True,
                        }
                    },
                ),
                id='tsne-spinner',
                spinning=False,
                text='Loading t-SNE...',
                style={'minHeight': '20vh', 'width': '100%'},
            ),
        ],
        style={'height': '100%'},
    )


def run_tsne_samples_in_cols(ndf, perplexity):
    """Run t-SNE on the sample matrix and return score coordinates."""
    if ndf.empty or ndf.shape[0] < 2:
        return None

    n_jobs = 1
    n_components = 3
    perplexity_value = perplexity if perplexity else 30
    if ndf.shape[0] > 0:
        perplexity_value = min(perplexity_value, max(1, ndf.shape[0] - 1))

    tsne_model = TSNE(
        n_components=n_components,
        perplexity=perplexity_value,
        n_jobs=n_jobs,
        random_state=42,
        init='pca',
    )
    embedded = tsne_model.fit_transform(ndf.to_numpy())
    tsne_cols = [f"t-SNE-{i+1}" for i in range(n_components)]
    return pd.DataFrame(embedded, index=ndf.index, columns=tsne_cols)


def generate_tsne_figure(ndf, color_labels, color_map, group_label, x_comp, y_comp, perplexity, tsne_scores=None):
    """Generate the t-SNE figure."""
    logger.info("Generating t-SNE...")
    
    if tsne_scores is None:
        scores_df = run_tsne_samples_in_cols(ndf, perplexity)
        if scores_df is None:
            return create_invisible_figure()
    else:
        scores_df = tsne_scores.copy()
    
    # Add metadata for plotting
    scores_df['color_group'] = color_labels
    scores_df['sample_label'] = scores_df.index
    
    x_axis = x_comp or 't-SNE-1'
    y_axis = y_comp or 't-SNE-2'
    
    # Determine actual columns to use (fallback to available if selected not present)
    if x_axis not in scores_df.columns and len(scores_df.columns) > 0:
        x_axis = scores_df.columns[0]
    if y_axis not in scores_df.columns and len(scores_df.columns) > 1:
        y_axis = scores_df.columns[1]
        
    fig = px.scatter(
        scores_df,
        x=x_axis,
        y=y_axis,
        color='color_group',
        symbol='color_group',
        color_discrete_map=color_map if color_map else None,
        hover_data={'sample_label': True},
    )
    
    fig.update_layout(
        autosize=True,
        margin=dict(l=140, r=80, t=60, b=50),
        legend_title_text=group_label,
        legend=dict(
            x=-0.1,
            y=1.05,
            xanchor='right',
            orientation='v',
            title=dict(text=f'{group_label}<br>', font=dict(size=12)),
            font=dict(size=11),
            itemsizing='constant',
            tracegroupgap=7.5,
        ),
        xaxis_title_font=dict(size=16),
        yaxis_title_font=dict(size=16),
        xaxis_tickfont=dict(size=12),
        yaxis_tickfont=dict(size=12),
        template='plotly_white',
        paper_bgcolor='white',
        plot_bgcolor='white',
    )
    return fig


def register_callbacks(app):
    """Register t-SNE callbacks."""
    @app.callback(
        Output('tsne-spinner', 'spinning'),
        Input('analysis-sidebar-menu', 'currentKey'),
        Input('tsne-regenerate-btn', 'nClicks'),
        Input('tsne-x-comp', 'value'),
        Input('tsne-y-comp', 'value'),
        Input('tsne-perplexity-slider', 'value'),
        Input('tsne-graph', 'figure'),
        prevent_initial_call=True,
    )
    def toggle_tsne_spinner(active_tab, _regen, _x_comp, _y_comp, _perplexity, _figure):
        if active_tab != 'tsne':
            return False
        ctx = dash.callback_context
        trigger = ctx.triggered[0]["prop_id"].split(".")[0] if ctx.triggered else ""
        if trigger == 'tsne-graph':
            return False
        if trigger in (
            'analysis-sidebar-menu',
            'tsne-regenerate-btn',
            'tsne-x-comp',
            'tsne-y-comp',
            'tsne-perplexity-slider',
        ):
            return True
        raise PreventUpdate
