"""
Contains the data generation from PKBD functionality of the UI
"""

import importlib
import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import streamlit as st
from plotly.subplots import make_subplots

pkbd = importlib.import_module("QuadratiK.spherical_clustering").PKBD

st.title("Data generation from PKBD Models")
st.write('Generates samples from the supported PKBD Models - "rejvmf" and "rejacg"')

with st.expander("Click to view code"):
    code_python = """
    from QuadratiK.spherical_clustering import PKBD
    rho = specify your rho value here
    n_samples = specify the number of samples here
    mu = specify a list of location parameters
    data1 = PKBD().rpkb(n_samples,mu,rho,method = "rejvmf")
    data2 = PKBD().rpkb(n_samples,mu,rho,method = "rejacg")
    """
    st.code(code_python, language="python")

    code_R = """
    library(QuadratiK)
    rho = specify your rho value here
    n_samples = specify the number of samples here
    mu = specify a list of location parameters
    dat1 <- rpkb(n_samples, rho=rho, mu=mu, method="rejvmf")$x
    dat2 <- rpkb(n_samples, rho=rho, mu=mu, method="rejacg")$x
    """
    st.code(code_R, language="r")

n_samples = int(
    st.number_input("Enter the total number of samples to be generated", value=100)
)

rho = float(
    st.number_input(r"Enter the value of concentration parameter ($\rho$)", value=0.85)
)


def collect_numbers(x):
    return [float(i) for i in re.split("[^.0-9]", x) if i != ""]


mu_list = collect_numbers(
    st.text_input(
        r"Please enter the location parameter ($\mu$) separated by space",
        value="0 0 0.5",
    )
)


@st.cache_data
def generate_samples(n_samples, mu_list, rho, method):
    return pkbd().rpkb(n_samples, mu_list, rho, method)


data1 = generate_samples(n_samples, mu_list, rho, "rejvmf")
with st.expander('**Data generated using "rejvmf"**'):
    st.write(data1)
csv_res1 = pd.DataFrame(data1).to_csv().encode()
download_rejvmf = st.download_button(
    "Click to Download the generated samples using rejvmf",
    csv_res1,
    "rejvmf_generated_samples.csv",
    "text/csv",
    key="download-rejvmf",
)

data2 = generate_samples(n_samples, mu_list, rho, "rejacg")
with st.expander('**Data generated using "rejacg"**'):
    st.write(data2)
csv_res2 = pd.DataFrame(data2).to_csv().encode()
download_rejacg = st.download_button(
    "Click to Download the generated samples using rejacg",
    csv_res2,
    "rejvmf_generated_samples.csv",
    "text/csv",
    key="download-rejacg",
)

if st.button("Rerun"):
    st.cache_data.clear()

if len(mu_list) == 2:
    fig1 = plt.figure(figsize=(6, 6))
    plt.scatter(data1[:, 0], data1[:, 1], label="rejvmf")
    plt.scatter(data2[:, 0], data2[:, 1], label="rejacg")

    theta = np.linspace(0, 2 * np.pi, 100)
    unit_circle_x = np.cos(theta)
    unit_circle_y = np.sin(theta)

    plt.plot(unit_circle_x, unit_circle_y, linestyle="dashed", color="red")
    plt.xlabel("Dimension 1")
    plt.ylabel("Dimension 2")
    st.pyplot(fig1)

elif len(mu_list) == 3:
    r = 1
    pi = np.pi
    cos = np.cos
    sin = np.sin
    phi, theta = np.mgrid[0.0:pi:100j, 0.0 : 2.0 * pi : 100j]
    x1 = r * sin(phi) * cos(theta)
    y1 = r * sin(phi) * sin(theta)
    z1 = r * cos(phi)
    xx1, yy1, zz1 = data1[:, 0], data1[:, 1], data1[:, 2]
    xx2, yy2, zz2 = data2[:, 0], data2[:, 1], data2[:, 2]

    fig = make_subplots(rows=1, cols=1, specs=[[{"type": "scatter3d"}]])
    fig.append_trace(
        go.Surface(
            x=x1,
            y=y1,
            z=z1,
            colorscale=[[0, "#DCDCDC"], [1, "#DCDCDC"]],
            opacity=0.5,
            showscale=False,
        ),
        row=1,
        col=1,
    )
    fig.append_trace(
        go.Scatter3d(
            x=xx1,
            y=yy1,
            z=zz1,
            mode="markers",
            marker=dict(size=5, color="red", showscale=False),
            name="rejvmf",
        ),
        row=1,
        col=1,
    )
    fig.append_trace(
        go.Scatter3d(
            x=xx2,
            y=yy2,
            z=zz2,
            mode="markers",
            marker=dict(size=5, color="blue", showscale=False),
            name="rejacg",
        ),
        row=1,
        col=1,
    )

    fig.update_layout(
        title="",
        scene=dict(
            xaxis=dict(range=[-1, 1]),
            yaxis=dict(range=[-1, 1]),
            zaxis=dict(range=[-1, 1]),
            aspectmode="data",
        ),
    )
    fig.update_layout(showlegend=True)
    st.plotly_chart(fig, use_container_width=True)

else:
    st.error(
        "Visualization cannot be generated as number of dimensions are greater than 3"
    )

st.markdown(
    r"""
    <style>
        .reportview-container {
            margin-top: -2em;
        }
        #MainMenu {visibility: hidden;}
        .stDeployButton {display:none;}
        footer {visibility: hidden;}
        #stDecoration {display:none;}
    </style>
""",
    unsafe_allow_html=True,
)
