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.
glob_b = 4.
glob_c = 7.5


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=glob_k, a=glob_a, b=glob_b, c=glob_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):
    # Define grid dimensions
    n_x = 32  # Fixed spatial grid resolution
    n_t = 50

    try:
        loaded_values = np.load(f"{approx_type}_m{m}.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)

    # Compute the selected approximation
    U_approx = np.zeros_like(U_real)
    for i, t_val in enumerate(t):
        Phi_gff_at_t = Phi[i * n_x : (i + 1) * n_x]
        U_approx[i, :] = np.dot(Phi_gff_at_t, 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(
        title=f"Heat Equation Approximation | Kernel = {approx_type} | m = {m}",
        scene=dict(
            camera=dict(
                eye=dict(x=0, y=-2, z=0),  # Front view
            ),
            xaxis_title="x",
            yaxis_title="t",
            zaxis_title="u",
        ),
    )

    # Config to remove modebar buttons except the save image button
    config = {
        "modeBarButtonsToRemove": [
            "pan",
            "resetCameraLastSave",
            "hoverClosest3d",
            "hoverCompareCartesian",
            "zoomIn",
            "zoomOut",
            "select2d",
            "lasso2d",
            "zoomIn2d",
            "zoomOut2d",
            "sendDataToCloud",
            "zoom3d",
            "orbitRotation",
            "tableRotation",
        ],
        "displayModeBar": True,  # Keep the modebar visible
        "displaylogo": False,  # Hide the Plotly logo
    }

    # 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",
        ]
    )

    return fig


def plot_errors(m, approx_type):
    # Define grid dimensions
    n_x = 32  # Fixed spatial grid resolution
    n_t = 50

    try:
        loaded_values = np.load(f"{approx_type}_m{m}.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)

    # Compute the selected approximation
    U_approx = np.zeros_like(U_real)
    for i, t_val in enumerate(t):
        Phi_gff_at_t = Phi[i * n_x : (i + 1) * n_x]
        U_approx[i, :] = np.dot(Phi_gff_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(
        title=f"Heat Equation Approximation Error | Kernel = {approx_type} | m = {m}",
        scene=dict(
            camera=dict(
                eye=dict(x=0, y=-2, z=0),  # Front view
            ),
            xaxis_title="x",
            yaxis_title="t",
            zaxis_title="u",
        ),
    )

    # Config to remove modebar buttons except the save image button
    config = {
        "modeBarButtonsToRemove": [
            "pan",
            "resetCameraLastSave",
            "hoverClosest3d",
            "hoverCompareCartesian",
            "zoomIn",
            "zoomOut",
            "select2d",
            "lasso2d",
            "zoomIn2d",
            "zoomOut2d",
            "sendDataToCloud",
            "zoom3d",
            "orbitRotation",
            "tableRotation",
        ],
        "displayModeBar": True,  # Keep the modebar visible
        "displaylogo": False,  # Hide the Plotly logo
    }

    # 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",
        ]
    )

    return fig


def generate_data(n_x=32, n_t=50):
    """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]
    )  # shape (n_x * n_t,)
    return a_train, u_train, x, t


def random_features(a, theta_j, kernel="SINE", k=0.5, t=1.0):
    """Compute random features with adjustable kernel width."""
    if kernel == "SINE":
        return np.sin(t * np.linalg.norm(a - theta_j, axis=-1))
    elif kernel == "GFF":
        return np.log(np.linalg.norm(a - theta_j, axis=-1)) / (2 * np.pi)
    else:
        raise ValueError("Unsupported kernel type!")


def design_matrix(a, theta, kernel):
    """Construct design matrix."""
    return np.array([random_features(a, theta_j, kernel=kernel) 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, kernel):
    """Compute the approximation."""
    Phi = design_matrix(a, theta, kernel)
    return Phi @ alpha


def polyfit2d(x, y, z, kx=3, ky=3, order=None):
    # grid coords
    x, y = np.meshgrid(x, y)
    # coefficient array, up to x^kx, y^ky
    coeffs = np.ones((kx + 1, ky + 1))

    # solve array
    a = np.zeros((coeffs.size, x.size))

    # for each coefficient produce array x^i, y^j
    for index, (j, i) in enumerate(np.ndindex(coeffs.shape)):
        # do not include powers greater than order
        if order is not None and i + j > order:
            arr = np.zeros_like(x)
        else:
            arr = coeffs[i, j] * x**i * y**j
        a[index] = arr.ravel()

    # do leastsq fitting and return leastsq result
    return np.linalg.lstsq(a.T, np.ravel(z), rcond=None)


def train_coefficients(m, kernel):
    # Start time for training
    start_time = time.time()

    # Generate data
    n_x, n_t = 32, 50
    a_train, u_train, x, t = generate_data(n_x, n_t)

    # Define random features
    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]
        )
    )

    # Construct design matrix and learn coefficients
    Phi = design_matrix(a_train, theta, kernel)
    alpha = learn_coefficients(Phi, u_train)
    # Validate and animate results
    u_real = np.array([complex_heat_eq_solution(x, t_i) for t_i in t])
    a_test = np.c_[np.meshgrid(x, t)[0].ravel(), np.meshgrid(x, t)[1].ravel()]
    u_approx = approximate_solution(a_test, alpha, theta, kernel).reshape(n_t, n_x)

    # Save values to the npz folder
    np.savez(
        f"{kernel}_m{m}.npz",
        alpha=alpha,
        kernel=kernel,
        Phi=Phi,
        theta=theta,
    )

    # Compute average error
    avg_err = np.mean(np.abs(u_real - u_approx))

    return f"Training completed in {time.time() - start_time:.2f} 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",
        ]
    )

    return fig


def plot_all(m, kernel):
    # Generate the plot content (replace this with your actual plot logic)
    approx_fig = plot_heat_equation(
        m, kernel
    )  # Replace with your function for approx_plot
    error_fig = plot_errors(m, kernel)  # 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),
    )

# Gradio interface
def create_gradio_ui():
    global glob_k, glob_a, glob_b, glob_c
    
    # Get the initial available files
    with gr.Blocks() as demo:
        gr.Markdown("# Learn the Coefficients for the Heat Equation using the RFM")

        # Function parameter inputs
        gr.Markdown(
            """
        ## Function: $$u(x, t)\\coloneqq\\exp(-\\textcolor{magenta}{k}(\\textcolor{cyan}{a}\\pi)^2t)\\sin(\\textcolor{cyan}{a}\\pi x)+0.5\\exp(-\\textcolor{magenta}{k}(\\textcolor{lime}{b}\\pi)^2t)\\sin(\\textcolor{lime}{b}\\pi x)+0.25\\exp(-\\textcolor{magenta}{k}(\\textcolor{orange}{c}\\pi)^2t)\\sin(\\textcolor{orange}{c}\\pi x)$$
        
        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.
        """
        )

        with gr.Row():
            with gr.Column():
                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"
                )

            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():
                # Kernel selection and slider for m
                kernel_dropdown = gr.Dropdown(
                    label="Choose Kernel", choices=["SINE", "GFF"], value="SINE"
                )
                m_slider = gr.Dropdown(
                    label="Number of Random Features (m)",
                    choices=[50, 250, 1000, 5000, 10000, 25000],
                    value=1000,
                )

                # 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, kernel_dropdown],
                    outputs=output,
                )
                approx_button = gr.Button("Plot Approximation")

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

        approx_button.click(
            fn=plot_all,
            inputs=[m_slider, kernel_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)