import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objs as go
import streamlit as st
import umap
from sklearn.manifold import TSNE

from nhssynth.modules.dashboard.utils import (
    hide_streamlit_content,
    id_selector,
    subset_selector,
)


def triangle(matrix):
    return matrix.where(np.tril(np.ones(matrix.shape), k=-1).astype(bool))


def distribution_plots(real_dataset, synthetic_datasets) -> None:
    column = st.sidebar.selectbox("Select column to display", real_dataset.columns)
    st.write(f"## Distribution Plot for `{column}`")
    synthetic_datasets = subset_selector(synthetic_datasets)
    fig = px.histogram(
        pd.concat(
            [
                pd.DataFrame({"Data": real_dataset[column], "Type": len(real_dataset) * ["Real"]}),
                pd.concat(
                    [
                        pd.DataFrame(
                            {
                                "Data": sd.iloc[0][column],
                                "Type": len(sd.iloc[0]) * [f"Synthetic: {idx[0]} (Repeat {idx[1]}, Config {idx[2]})"],
                            }
                        )
                        for idx, sd in synthetic_datasets.iterrows()
                    ]
                ),
            ]
        ),
        x="Data",
        color="Type",
        barmode="overlay",
    )
    log_scale = st.sidebar.checkbox("Log scale", value=False)
    if log_scale:
        fig.update_layout(yaxis_type="log")
    if real_dataset[column].dtype == "object":
        sort_xaxis = st.sidebar.checkbox("Sort x-axis", value=True)
        if sort_xaxis:
            fig.update_layout(xaxis={"categoryorder": "total descending"})
    st.plotly_chart(fig, use_container_width=True)


def correlation_plots(real_dataset, synthetic_datasets) -> None:
    show_continuous = st.sidebar.checkbox("Only show continuous columns", value=True)
    selection = real_dataset.select_dtypes(exclude=["object"]).columns
    if not show_continuous:
        non_continuous_columns = [c for c in real_dataset.columns if c not in selection]
        selection = st.sidebar.selectbox("Select categorical column to display", non_continuous_columns)
        if not non_continuous_columns:
            st.error("No categorical columns found!")
            return go.Figure()

    synthetic_dataset = id_selector(synthetic_datasets).iloc[0][selection]
    real_dataset = real_dataset[selection]
    if not show_continuous:
        synthetic_dataset = pd.get_dummies(synthetic_dataset)
        real_dataset = pd.get_dummies(real_dataset)

    correlation_type = st.sidebar.selectbox("Select correlation type", ["Pearson", "Spearman", "Kendall"]).lower()
    correlation_to_show = st.sidebar.selectbox("Select correlation to display", ["Real", "Synthetic", "Difference"])

    real_corr_matrix = real_dataset.corr(method=correlation_type)
    synthetic_corr_matrix = synthetic_dataset.corr(method=correlation_type)
    zmin = min(real_corr_matrix.min().min(), synthetic_corr_matrix.min().min())
    zmax = max(real_corr_matrix.max().max(), synthetic_corr_matrix.max().max())
    if correlation_to_show == "Difference":
        corr_matrix = triangle(abs(synthetic_corr_matrix - real_corr_matrix))
        zmin = 0
        zmax = zmax - zmin
    elif correlation_to_show == "Real":
        corr_matrix = triangle(real_corr_matrix)
    else:
        corr_matrix = triangle(synthetic_corr_matrix)
    fig = go.Figure(
        data=go.Heatmap(
            z=corr_matrix,
            x=corr_matrix.columns,
            y=corr_matrix.columns,
            colorscale=px.colors.diverging.RdYlGn,
            reversescale=True,
            zmin=zmin,
            zmax=zmax,
        ),
        layout=go.Layout(xaxis_showgrid=False, yaxis_showgrid=False, yaxis_autorange="reversed"),
    )
    st.plotly_chart(fig, use_container_width=True)


def prepare_for_dimensionality(df: pd.DataFrame) -> pd.DataFrame:
    """Factorize all categorical columns in a dataframe and normalize values."""
    df = df.copy()
    for col in df.columns:
        if df[col].dtype == "object":
            df[col] = pd.factorize(df[col])[0]
        elif df[col].dtype == "datetime64[ns]":
            df[col] = pd.to_numeric(df[col])
        min_val = df[col].min()
        max_val = df[col].max()
        # Avoid division by zero when all values are the same
        if max_val != min_val:
            df[col] = (df[col] - min_val) / (max_val - min_val)
        else:
            df[col] = 0.0
    # Drop rows with NaN values (required for UMAP/t-SNE)
    df = df.dropna()
    return df


@st.cache_data(show_spinner=False)
def compute_umap(data: np.ndarray) -> np.ndarray:
    """Compute UMAP projection with caching."""
    reducer = umap.UMAP()
    return reducer.fit_transform(data)


@st.cache_data(show_spinner=False)
def compute_tsne(data: np.ndarray) -> np.ndarray:
    """Compute t-SNE projection with caching."""
    reducer = TSNE(n_components=2, init="pca")
    return reducer.fit_transform(data)


def plot_reducer(
    real_dataset: pd.DataFrame,
    synthetic_dataset: pd.DataFrame,
    method: str,
) -> go.Figure:
    with st.spinner(f"Running {method}..."):
        fig = go.Figure(layout=go.Layout(xaxis_title=f"{method} 1", yaxis_title=f"{method} 2"))

        # Use cached computation functions
        if method == "UMAP":
            proj_real = compute_umap(real_dataset.values)
            proj_synth = compute_umap(synthetic_dataset.values)
        else:  # t-SNE
            proj_real = compute_tsne(real_dataset.values)
            proj_synth = compute_tsne(synthetic_dataset.values)

        fig.add_scatter(
            x=proj_real[:, 0],
            y=proj_real[:, 1],
            mode="markers",
            marker=dict(size=5),
            opacity=0.75,
            name="Real data",
        )
        fig.add_scatter(
            x=proj_synth[:, 0],
            y=proj_synth[:, 1],
            mode="markers",
            marker=dict(size=5),
            opacity=0.75,
            name="Synthetic data",
        )
    return fig


def dimensionality_plots(real_dataset: pd.DataFrame, synthetic_datasets: pd.DataFrame) -> None:
    original_real_len = len(real_dataset)
    synthetic_raw = id_selector(synthetic_datasets).iloc[0].copy()
    original_synth_len = len(synthetic_raw)

    synthetic_dataset = prepare_for_dimensionality(synthetic_raw)
    real_dataset = prepare_for_dimensionality(real_dataset.copy())

    # Warn if significant data was dropped due to NaN
    real_dropped = original_real_len - len(real_dataset)
    synth_dropped = original_synth_len - len(synthetic_dataset)
    if real_dropped > 0 or synth_dropped > 0:
        st.warning(f"Rows with missing values were dropped: {real_dropped} real rows, {synth_dropped} synthetic rows.")

    dimensionality_method = st.sidebar.selectbox("Select dimensionality reduction method", ["UMAP", "t-SNE"])
    run = st.sidebar.button("Run dimensionality reduction")
    if run:
        if len(real_dataset) == 0 or len(synthetic_dataset) == 0:
            st.error("No data remaining after dropping NaN values. Cannot run dimensionality reduction.")
            return
        fig = plot_reducer(real_dataset, synthetic_dataset, dimensionality_method)
        st.plotly_chart(fig, use_container_width=True)


def page():
    # metric_groups = st.session_state["evaluations"].keys()
    plot_types = ["Distribution", "Correlation", "Dimensionality"]
    selected_plot_type = st.sidebar.selectbox("Select a plot type", plot_types)

    real_data = st.session_state["typed"]
    synthetic_datasets = st.session_state["synthetic_datasets"]

    if selected_plot_type == "Distribution":
        distribution_plots(real_data, synthetic_datasets)
    elif selected_plot_type == "Correlation":
        correlation_plots(real_data, synthetic_datasets)
    elif selected_plot_type == "Dimensionality":
        dimensionality_plots(real_data, synthetic_datasets)


if __name__ == "__main__":
    st.set_page_config(layout="wide")
    hide_streamlit_content()
    if "evaluations" not in st.session_state:
        st.error("Upload an evaluation bundle to get started!")
    else:
        page()
