import os
import time
import shutil

import numpy as np
import gradio as gr
import plotly.graph_objs as go

glob_k = 0.0025
glob_a = -2.0
glob_b = 4.0
glob_c = 7.5

n_x, n_t = 10, 10


def clear_npz():
    current_directory = os.getcwd()  # Get the current working directory
    for filename in os.listdir(current_directory):
        if filename.endswith(".npz"):  # Check if the file ends with .npz
            file_path = os.path.join(current_directory, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)  # Remove the file or symbolic link
                else:
                    print(f"Skipping {file_path}, not a file or symbolic link.")
            except Exception as e:
                print(f"Failed to delete {file_path}. Reason: {e}")


def complex_heat_eq_solution(x, t, k, a, b, c):
    global glob_k, glob_a, glob_b, glob_c
    glob_k, glob_a, glob_b, glob_c = k, a, b, c
    return (
        np.exp(-glob_k * (glob_a * np.pi) ** 2 * t) * np.cos(glob_a * np.pi * x)
        + np.exp(-glob_k * (glob_b * np.pi) ** 2 * t) * np.sin(glob_b * np.pi * x)
        + np.exp(-glob_k * (glob_c * np.pi) ** 2 * t) * np.sin(glob_c * np.pi * x)
    )


def plot_heat_equation(m, approx_type, quality, rand_or_det):
    global glob_k, glob_a, glob_b, glob_c, n_x, n_t

    # Plot with more points than it was calculated
    new_nx = 1 * n_x
    new_nt = 1 * n_t

    try:
        loaded_values = np.load(
            f"{approx_type}_m{m}_{str.lower(quality)}_{str.lower(rand_or_det)}.npz"
        )
    except:
        raise gr.Error(f"First train the coefficients for {approx_type} and m = {m}")
    alpha = loaded_values["alpha"]
    Phi = loaded_values["Phi"]

    # Create grids for x and t
    x = np.linspace(0, 1, new_nx)  # Spatial grid
    t = np.linspace(0, 5, new_nt)  # Temporal grid
    X, T = np.meshgrid(x, t)

    # Compute the real solution over the grid
    U_real = complex_heat_eq_solution(X, T, glob_k, glob_a, glob_b, glob_c)

    # Compute the selected approximation
    # Compute the approximations as a single matrix multiplication
    Phi_reshaped = Phi.reshape(n_t, n_x, -1)
    # The result will be of shape (n_t, n_x), as U_approx should match U_real's shape
    U_approx = np.einsum("ijk,k->ij", Phi_reshaped, alpha)

    # Create the 3D plot with Plotly
    traces = []

    # Real solution surface with a distinct color (e.g., 'Viridis')
    traces.append(
        go.Surface(
            z=U_real,
            x=X,
            y=T,
            colorscale="Blues",
            showscale=False,
            name="Real Solution",
            showlegend=True,
        )
    )

    # Approximation surface with a distinct color (e.g., 'Plasma')
    traces.append(
        go.Surface(
            z=U_approx,
            x=X,
            y=T,
            colorscale="Reds",
            reversescale=True,
            showscale=False,
            name=f"{approx_type} Approximation",
            showlegend=True,
        )
    )

    # Layout for the Plotly plot without controls
    layout = go.Layout(
        scene=dict(
            camera=dict(
                eye=dict(x=0, y=-2, z=0),  # Front view
            ),
            xaxis_title="x",
            yaxis_title="t",
            zaxis_title="u",
        ),
        margin=dict(l=0, r=0, t=0, b=0),  # Reduce margins
    )

    # Create the figure
    fig = go.Figure(data=traces, layout=layout)

    fig.update_layout(
        modebar_remove=[
            "pan",
            "resetCameraLastSave",
            "hoverClosest3d",
            "hoverCompareCartesian",
            "zoomIn",
            "zoomOut",
            "select2d",
            "lasso2d",
            "zoomIn2d",
            "zoomOut2d",
            "sendDataToCloud",
            "zoom3d",
            "orbitRotation",
            "tableRotation",
            "toImage",
            "resetCameraDefault3d",
        ],
        legend=dict(yanchor="bottom", y=0.01, xanchor="left", x=0.01),
    )

    return fig


def plot_errors(m, approx_type, quality, rand_or_det):
    global n_x, n_t

    try:
        loaded_values = np.load(
            f"{approx_type}_m{m}_{str.lower(quality)}_{str.lower(rand_or_det)}.npz"
        )
    except:
        raise gr.Error(f"First train the coefficients for {approx_type} and m = {m}")
    alpha = loaded_values["alpha"]
    Phi = loaded_values["Phi"]

    # Create grids for x and t
    x = np.linspace(0, 1, n_x)  # Spatial grid
    t = np.linspace(0, 5, n_t)  # Temporal grid
    X, T = np.meshgrid(x, t)

    # Compute the real solution over the grid
    U_real = complex_heat_eq_solution(X, T, glob_k, glob_a, glob_b, glob_c)

    # Compute the selected approximation
    U_approx = np.zeros_like(U_real)
    for i, t_val in enumerate(t):
        Phi_at_t = Phi[i * n_x : (i + 1) * n_x]
        U_approx[i, :] = np.dot(Phi_at_t, alpha)

    U_err = abs(U_approx - U_real)

    # Create the 3D plot with Plotly
    traces = []

    # Real solution surface with a distinct color (e.g., 'Viridis')
    traces.append(
        go.Surface(
            z=U_err,
            x=X,
            y=T,
            colorscale="Viridis",
            showscale=False,
            name=f"Absolute Error",
            showlegend=True,
        )
    )

    # Layout for the Plotly plot without controls
    layout = go.Layout(
        scene=dict(
            camera=dict(
                eye=dict(x=0, y=-2, z=0),  # Front view
            ),
            xaxis_title="x",
            yaxis_title="t",
            zaxis_title="u",
        ),
        margin=dict(l=0, r=0, t=0, b=0),  # Reduce margins
    )

    # Create the figure
    fig = go.Figure(data=traces, layout=layout)

    fig.update_layout(
        modebar_remove=[
            "pan",
            "resetCameraLastSave",
            "hoverClosest3d",
            "hoverCompareCartesian",
            "zoomIn",
            "zoomOut",
            "select2d",
            "lasso2d",
            "zoomIn2d",
            "zoomOut2d",
            "sendDataToCloud",
            "zoom3d",
            "orbitRotation",
            "tableRotation",
            "toImage",
            "resetCameraDefault3d",
        ],
        legend=dict(yanchor="bottom", y=0.01, xanchor="left", x=0.01),
    )

    return fig


def generate_data():
    global glob_k, glob_a, glob_b, glob_c, n_x, n_t
    """Generate training data."""
    x = np.linspace(0, 1, n_x)  # spatial points
    t = np.linspace(0, 5, n_t)  # temporal points
    X, T = np.meshgrid(x, t)
    a_train = np.c_[X.ravel(), T.ravel()]  # shape (n_x * n_t, 2)
    u_train = complex_heat_eq_solution(
        a_train[:, 0], a_train[:, 1], glob_k, glob_a, glob_b, glob_c
    )  # shape (n_x * n_t,)
    return a_train, u_train, x, t


def features(a, theta_j, m, method="SINE", k=1, eps=1e-8):
    """Compute random features with adjustable method width."""
    if method == "SINE":
        return np.sin(k * np.linalg.norm(a - theta_j, axis=-1) + eps)
    elif method == "GFF":
        return np.log(np.linalg.norm(a - theta_j, axis=-1) + eps) / (2 * np.pi)
    else:
        raise ValueError("Unsupported method type!")


def design_matrix(a, theta, method):
    """Construct design matrix."""
    return np.array(
        [features(a, theta_j, theta.shape[0], method) for theta_j in theta]
    ).T


def learn_coefficients(Phi, u):
    """Learn coefficients alpha via least squares."""
    return np.linalg.lstsq(Phi, u, rcond=None)[0]


def approximate_solution(a, alpha, theta, method):
    """Compute the approximation."""
    Phi = design_matrix(a, theta, method)
    return Phi @ alpha


def train_coefficients(m, method, quality, rand_or_det):
    global glob_k, glob_a, glob_b, glob_c, n_x, n_t
    # Start time for training
    start_time = time.time()

    # Generate data
    a_train, u_train, x, t = generate_data()

    # Define random features
    if rand_or_det == "Random":
        theta = np.column_stack(
            (
                np.random.uniform(-1, 1, size=m),  # First dimension: [-1, 1]
                np.random.uniform(-5, 5, size=m),  # Second dimension: [-5, 5]
            )
        )
    else:
        theta = np.column_stack(
            (
                np.linspace(-5, 5, m),  # Linear spacing for x
                np.linspace(-25, 25, m),  # Linear spacing for y
            )
        )

    # Construct design matrix and learn coefficients
    Phi = design_matrix(a_train, theta, method)
    alpha = learn_coefficients(Phi, u_train)

    end_time = f"{time.time() - start_time:.2f}"

    # Save values to the npz folder
    np.savez(
        f"{method}_m{m}_{str.lower(quality)}_{str.lower(rand_or_det)}.npz",
        alpha=alpha,
        method=method,
        Phi=Phi,
        theta=theta,
    )

    # Compute the average error
    #
    x = np.linspace(0, 1, n_x)  # Spatial grid
    t = np.linspace(0, 5, n_t)  # Temporal grid
    X, T = np.meshgrid(x, t)

    # Compute the real solution over the grid
    U_real = complex_heat_eq_solution(X, T, glob_k, glob_a, glob_b, glob_c)

    # Compute the selected approximation
    U_approx = np.zeros_like(U_real)
    for i, t_val in enumerate(t):
        Phi_at_t = Phi[i * n_x : (i + 1) * n_x]
        U_approx[i, :] = np.dot(Phi_at_t, alpha)

    # Compute average error
    avg_err = np.mean(np.abs(U_real - U_approx))

    return (
        f"Training completed in {end_time} seconds. The average error is {avg_err}."
    )


def plot_function(k, a, b, c):
    global glob_k, glob_a, glob_b, glob_c

    glob_k, glob_a, glob_b, glob_c = k, a, b, c

    x = np.linspace(0, 1, 100)
    t = np.linspace(0, 5, 500)
    X, T = np.meshgrid(x, t)  # Create the mesh grid
    Z = complex_heat_eq_solution(X, T, glob_k, glob_a, glob_b, glob_c)

    traces = []
    traces.append(
        go.Surface(
            z=Z,
            x=X,
            y=T,
            colorscale="Viridis",
            showscale=False,
            showlegend=False,
        )
    )

    # Layout for the Plotly plot without controls
    layout = go.Layout(
        scene=dict(
            camera=dict(
                eye=dict(x=1.25, y=-1.75, z=0.3),  # Front view
            ),
            xaxis_title="x",
            yaxis_title="t",
            zaxis_title="u",
        ),
        margin=dict(l=0, r=0, t=0, b=0),  # Reduce margins
    )

    # Create the figure
    fig = go.Figure(data=traces, layout=layout)

    fig.update_layout(
        modebar_remove=[
            "pan",
            "resetCameraLastSave",
            "hoverClosest3d",
            "hoverCompareCartesian",
            "zoomIn",
            "zoomOut",
            "select2d",
            "lasso2d",
            "zoomIn2d",
            "zoomOut2d",
            "sendDataToCloud",
            "zoom3d",
            "orbitRotation",
            "tableRotation",
            "toImage",
            "resetCameraDefault3d",
        ],
        legend=dict(yanchor="bottom", y=0.01, xanchor="left", x=0.01),
    )

    return fig


def plot_all(m, method, quality, rand_or_det):
    # Generate the plot content (replace this with your actual plot logic)
    approx_fig = plot_heat_equation(
        m, method, quality, rand_or_det
    )  # Replace with your function for approx_plot
    error_fig = plot_errors(
        m, method, quality, rand_or_det
    )  # Replace with your function for error_plot

    # Return the figures and make the plots visible
    return (
        gr.update(visible=True, value=approx_fig),
        gr.update(visible=True, value=error_fig),
    )


def change_quality(quality):
    global n_x, n_t

    if quality == "Low":
        n_x, n_t = 10, 10
    elif quality == "Mid":
        n_x, n_t = 20, 20
    elif quality == "High":
        n_x, n_t = 40, 40


# Gradio interface
def create_gradio_ui():
    global glob_k, glob_a, glob_b, glob_c

    markdown_content = r"""
    ## Goal function:
    $$
    \begin{alignat*}{5}
    u(x, t)
    \coloneqq &\exp(-\textcolor{magenta}{k}(&\textcolor{cyan}{a}&\pi)^2t)\sin(&\textcolor{cyan}{a}&\pi x) \\
            + &\exp(-\textcolor{magenta}{k}(&\textcolor{lime}{b}&\pi)^2t)\sin(&\textcolor{lime}{b}&\pi x) \\
            + &\exp(-\textcolor{magenta}{k}(&\textcolor{orange}{c}&\pi)^2t)\sin(&\textcolor{orange}{c}&\pi x)
    \end{alignat*}
    $$

    Adjust the values for <span style='color: magenta;'>k</span>, <span style='color: cyan;'>a</span>, <span style='color: lime;'>b</span> and <span style='color: orange;'>c</span> with the sliders below.
    
    Pressing "Train Coefficients" aims to solve
    $$
    \argmin_{\alpha\in\mathbb{R}^m}\|{\alpha\Phi-u}\|_2^2,
    $$
    where $\Phi$ contains the features depending on the method.
    """

    # Get the initial available files
    with gr.Blocks() as demo:
        gr.Markdown("# Approximating a solution to the heat equation using RFM")

        # Function parameter inputs
        gr.Markdown(
            markdown_content,
            latex_delimiters=[
                {"left": "$$", "right": "$$", "display": True},
                {"left": "$", "right": "$", "display": False},
            ],
        )

        with gr.Row():
            with gr.Column(min_width=500):
                k_slider = gr.Slider(
                    minimum=0.0001, maximum=0.1, step=0.0001, value=glob_k, label="k"
                )
                a_slider = gr.Slider(
                    minimum=-10, maximum=10, step=0.1, value=glob_a, label="a"
                )
                b_slider = gr.Slider(
                    minimum=-10, maximum=10, step=0.1, value=glob_b, label="b"
                )
                c_slider = gr.Slider(
                    minimum=-10, maximum=10, step=0.1, value=glob_c, label="c"
                )
            with gr.Column(min_width=500):
                plot_output = gr.Plot()

        k_slider.change(
            fn=plot_function,
            inputs=[k_slider, a_slider, b_slider, c_slider],
            outputs=[plot_output],
        )
        a_slider.change(
            fn=plot_function,
            inputs=[k_slider, a_slider, b_slider, c_slider],
            outputs=[plot_output],
        )
        b_slider.change(
            fn=plot_function,
            inputs=[k_slider, a_slider, b_slider, c_slider],
            outputs=[plot_output],
        )
        c_slider.change(
            fn=plot_function,
            inputs=[k_slider, a_slider, b_slider, c_slider],
            outputs=[plot_output],
        )

        with gr.Column():
            with gr.Row():
                # method selection and slider for m
                quality_dropdown = gr.Dropdown(
                    label="Choose Quality", choices=["Low", "Mid", "High"], value="Low"
                )
                quality_dropdown.change(
                    fn=change_quality, inputs=quality_dropdown, outputs=None
                )
                method_dropdown = gr.Dropdown(
                    label="Choose Method", choices=["SINE", "GFF"], value="SINE"
                )
                m_slider = gr.Dropdown(
                    label="Number of Random Features (m)",
                    choices=[50, 250, 1000, 5000, 10000, 25000],
                    value=1000,
                )
                rand_det_dropdown = gr.Dropdown(
                    label="Choose Random / Deterministic",
                    choices=["Deterministic", "Random"],
                    value="Deterministic",
                )
            # Output to show status
            output = gr.Textbox(label="Status", interactive=False)

            with gr.Column():
                # Button to train coefficients
                train_button = gr.Button("Train Coefficients")
                # Function to trigger training and update dropdown
                train_button.click(
                    fn=train_coefficients,
                    inputs=[
                        m_slider,
                        method_dropdown,
                        quality_dropdown,
                        rand_det_dropdown,
                    ],
                    outputs=output,
                )
                approx_button = gr.Button("Plot Approximation")

        with gr.Row():
            with gr.Column(min_width=500):
                approx_plot = gr.Plot(visible=False)
            with gr.Column(min_width=500):
                error_plot = gr.Plot(visible=False)

        approx_button.click(
            fn=plot_all,
            inputs=[m_slider, method_dropdown, quality_dropdown, rand_det_dropdown],
            outputs=[approx_plot, error_plot],
        )

        demo.load(fn=clear_npz, inputs=None, outputs=None)
        demo.load(
            fn=plot_function,
            inputs=[k_slider, a_slider, b_slider, c_slider],
            outputs=[plot_output],
        )

    return demo


# Launch Gradio app
if __name__ == "__main__":
    interface = create_gradio_ui()
    interface.launch(share=False)