import numpy as np
import gradio as gr

from sklearn.datasets import make_checkerboard
from sklearn.cluster import SpectralBiclustering
from sklearn.metrics import consensus_score

import plotly.express as px


score = [0.0]


def dataset(n_cluster_rows, n_cluster_cols, noise=5, n_rows=300, n_cols=300):
    data, rows, columns = make_checkerboard(
        shape=(n_rows, n_cols),
        n_clusters=(n_cluster_rows, n_cluster_cols),
        noise=noise,
        shuffle=False,
        random_state=0,
    )

    fig = px.imshow(data, title="Original Data")
    return fig


def shuffle_dataset(n_cluster_rows, n_cluster_cols, noise=5, n_rows=300, n_cols=300):
    data, rows, columns = make_checkerboard(
        shape=(n_rows, n_cols),
        n_clusters=(n_cluster_rows, n_cluster_cols),
        noise=noise,
        shuffle=False,
        random_state=0,
    )
    rng = np.random.RandomState(0)
    row_idx = rng.permutation(data.shape[0])
    col_idx = rng.permutation(data.shape[1])
    data = data[row_idx][:, col_idx]
    fig = px.imshow(data, title="Shuffled Data")
    return fig


def model_fit(
    n_cluster_rows,
    n_cluster_cols,
    noise,
    n_rows,
    n_cols,
    n_cluster_rows_,
    n_cluster_cols_,
    svd_method,
):
    data, rows, columns = make_checkerboard(
        shape=(n_rows, n_cols),
        n_clusters=(n_cluster_rows, n_cluster_cols),
        noise=noise,
        shuffle=False,
        random_state=0,
    )
    fig_original = px.imshow(data, title="Original Data")
    rng = np.random.RandomState(0)
    row_idx = rng.permutation(data.shape[0])
    col_idx = rng.permutation(data.shape[1])
    data = data[row_idx][:, col_idx]
    fig_shuffled = px.imshow(data, title="Shuffled Data")
    model = SpectralBiclustering(
        n_clusters=(n_cluster_rows_, n_cluster_cols_),
        method="log",
        random_state=0,
        svd_method=svd_method,
    )
    model.fit(data)
    score.append(
        consensus_score(model.biclusters_, (rows[:, row_idx], columns[:, col_idx]))
    )
    fit_data = data[np.argsort(model.row_labels_)]
    fit_data = fit_data[:, np.argsort(model.column_labels_)].T
    fig = px.imshow(fit_data, title="After Bi-Clustering")

    fig_1 = px.imshow(
        np.outer(np.sort(model.row_labels_) + 1, np.sort(model.column_labels_) + 1),
        title="Checkerboard structure of rearranged data",
    )

    return fig_original, fig_shuffled, fig, fig_1


def get_score():
    return score[-1].__format__(".3f")


with gr.Blocks() as demo:
    gr.Markdown("## Spectral Bi-Clustering")
    gr.Markdown(
        "Demo is based on the [Spectral Bi-Clustering](https://scikit-learn.org/stable/auto_examples/bicluster/plot_spectral_coclustering.html) example from scikit-learn. The goal of co-clustering is to find subgroups of rows and columns that are highly correlated. The data is first shuffled, then the rows and columns are reordered to match the biclusters. The consensus score is a measure of how well the biclusters found by the model match the true biclusters. The score is between 0 and 1, with 1 being a perfect match."
    )

    with gr.Tab("Data"):
        gr.Markdown("## Play with the parameters to see how the data changes")
        gr.Markdown("### Parameters")
        with gr.Row():
            n_rows = gr.Slider(1, 500, label="Number of Rows", value=300, step=1)
            n_cols = gr.Slider(1, 500, label="Number of Columns", value=300, step=1)
            n_cluster_rows = gr.Slider(
                1, 50, label="Number of Clusters Rows", value=5, step=1
            )
            n_cluster_cols = gr.Slider(
                1, 50, label="Number of Clusters Columns", value=5, step=1
            )
            noise = gr.Slider(0, 10, label="Noise", value=5, step=1)
        with gr.Row():
            gen_btn = gr.Button("Generate Data")
            shu_btn = gr.Button("Shuffle Data")
        with gr.Row():
            gen_btn.click(
                fn=dataset,
                inputs=[n_cluster_rows, n_cluster_cols, noise, n_rows, n_cols],
                outputs=gr.Plot(),
            )
            shu_btn.click(
                fn=shuffle_dataset,
                inputs=[n_cluster_rows, n_cluster_cols, noise, n_rows, n_cols],
                outputs=gr.Plot(),
            )

    with gr.Tab("Model"):
        gr.Markdown("## Model")
        gr.Markdown("### Data Parameters")
        with gr.Row():
            n_rows = gr.Slider(1, 500, label="Number of Rows", value=300, step=1)
            n_cols = gr.Slider(1, 500, label="Number of Columns", value=300, step=1)
            n_cluster_rows = gr.Slider(
                1, 50, label="Number of Clusters Rows", value=5, step=1
            )
            n_cluster_cols = gr.Slider(
                1, 50, label="Number of Clusters Columns", value=5, step=1
            )
            noise = gr.Slider(0, 10, label="Noise", value=5, step=1)
        gr.Markdown("### Model Parameters")
        with gr.Row():
            n_cluster_rows_ = gr.Slider(
                1, 50, label="Number of Clusters Rows", value=5, step=1
            )
            n_cluster_cols_ = gr.Slider(
                1, 50, label="Number of Clusters Columns", value=5, step=1
            )
            svd_method = gr.Dropdown(
                ["randomized", "arpack"], label="SVD Method", value="randomized"
            )
        model_btn = gr.Button("Fit Model")
        with gr.Row():
            model_btn.click(
                fn=model_fit,
                inputs=[
                    n_cluster_rows,
                    n_cluster_cols,
                    noise,
                    n_rows,
                    n_cols,
                    n_cluster_rows_,
                    n_cluster_cols_,
                    svd_method,
                ],
                outputs=[gr.Plot(), gr.Plot(), gr.Plot(), gr.Plot()],
            )
        gr.Markdown("### Consensus Score")
        score_btn = gr.Button("Get Score")
        with gr.Row():
            score_btn.click(fn=get_score, outputs=gr.Text())


demo.launch()