from helpers.maths import Vector
from helpers.ui import interactive_figure, line_scatter, Button, Slider, label
from helpers.gpr.gpr import GPRCreator
from helpers.gpr.types import Kernel, RBFKernelCreator
from helpers.gpr.updater import (
    CovarianceCalculator,
    RandomDrawingGenerator,
    TrainingDataUpdater,
    HyperparameterUpdater,
)
from helpers.gpr.training import TrainingData
from helpers.gpr.plot import create_traces_for

import numpy as np


def interactive_rbf_kernel_figure(
    kernel_creator: RBFKernelCreator, x_values: Vector
) -> None:
    """Creates an interactive figure for exploring the Radial Basis Function kernel.

    Args:
        kernel_creator: An object that can create the RBF kernel.
        x_values: The x values to calculate the output of the kernel for.

    Example:
        ```python
        class MyRBFKernel: # Some implementation of an RBF kernel
            ...

        x_values = np.linspace(-10, 10, 100)
        interactive_rbf_kernel_figure(MyRBFKernel, x_values)

        ```

    Note:
        Check out the [`RBFKernelCreator`](exercise_2.md#helpers.gpr.RBFKernelCreator) interface to see the exact
        requirements for the `kernel_creator`, but in general, your class most likely has a constructor that takes
        the length scale and signal variance as arguments. In such cases, it is sufficient to pass the class itself
        as the `kernel_creator` argument.
    """

    calculate_covariance = CovarianceCalculator(x_values, kernel_creator)

    interactive_figure(
        data=[
            line_scatter(
                x=x_values,
                y=np.zeros_like(x_values),
                name="Covariance",
            ),
            label(x=0, y=0, text="$x_2$"),
        ],
        title="Radial Basis Function Kernel",
        x_title="$x_1$",
        y_title="$RBF(x_1, x_2)$",
        sliders=[
            Slider(
                "l",
                min=0.1,
                max=3,
                step=0.1,
                default=1.0,
                description=r"Length scale (l)",
            ),
            Slider(
                "sigma",
                min=0.1,
                max=2,
                step=0.01,
                default=1.0,
                description=r"Signal variance (sigma)",
            ),
            Slider(
                "x_2",
                min=-10,
                max=10,
                step=0.1,
                default=0.0,
                description=r"Second input (x_2)",
            ),
        ],
        on_update=calculate_covariance,
    )


def interactive_random_functions_figure(
    gpr: GPRCreator, training: TrainingData, kernel: Kernel, x: Vector
) -> None:
    """Creates an interactive figure for exploring random functions generated by a Gaussian process.

    Args:
        gpr: The Gaussian process regression model creator.
        training: The training data used to create the model.
        kernel: The kernel used to create the model.
        x: The x values to calculate the output of the model for.

    Example:
        ```python
        class MyGPR: # Some implementation of a GPR model
            ...

        training = ... # Some training data
        kernel = ... # Some kernel function
        x = np.linspace(-10, 10, 100)

        interactive_random_functions_figure(MyGPR, training, kernel, x)
        # This will show an interactive plot of random functions generated by the GPR model.
        ```

        Note:
        The `GPRCreator` interface specifies the requirements for the `gpr` argument. In general, if you
        add a static method to your class called `create`, that takes the training data and kernel as arguments,
        and returns a GPR model, you can pass the class itself as the `gpr` argument.
    """
    model = gpr.create(training, kernel)
    draw_random_function = RandomDrawingGenerator(model.predict(x), x)

    interactive_figure(
        data=create_traces_for(model, x, show_mean=False),
        title=f"Random functions of the Gaussian process with {kernel}",
        x_title="x",
        y_title="f(x)",
        buttons=[Button("Add random function", on_click=draw_random_function)],
    )


def interactive_training_data_figure(
    gpr: GPRCreator, training_data: TrainingData, kernel: Kernel, x: Vector
) -> None:
    """Creates an interactive figure for exploring the effect of adding and removing training data points.

    Args:
        gpr: The Gaussian process regression model creator.
        training_data: The initial training data to use.
        kernel: The kernel used to create the model.
        x: The x values to calculate the output of the model for.

    Example:
        ```python
        class MyGPR: # Some implementation of a GPR model
            ...

        training_data = ... # Some training data
        kernel = ... # Some kernel function
        x = np.linspace(-10, 10, 100)

        interactive_training_data_figure(MyGPR, training_data, kernel, x)
        # This will show an interactive plot of the GPR model with varying training data.
        ```

        See notes on [`interactive_random_functions_figure`](exercise_2.md#helpers.gpr.interactive_random_functions_figure)
        for more information on the `gpr` argument.
    """
    update_data_points = TrainingDataUpdater(gpr, training_data, kernel, x)

    interactive_figure(
        data=update_data_points.traces(),
        title="GPR with varying number of training points",
        x_title="x",
        y_title="f(x)",
        buttons=[
            Button("Add data point", on_click=update_data_points.add),
            Button("Remove data point", on_click=update_data_points.remove),
        ],
    )


def interactive_hyperparameter_figure(
    gpr: GPRCreator,
    kernel_creator: RBFKernelCreator,
    training_data: TrainingData,
    x: Vector,
) -> None:
    """Creates an interactive figure for exploring the effect of hyperparameters on the GPR model.

    Args:
        gpr: The Gaussian process regression model creator.
        kernel_creator: The kernel creator used to create the model.
        training_data: The training data used to create the model.
        x: The x values to calculate the output of the model for.

    Example:
        ```python
        class MyGPR: # Some implementation of a GPR model
            ...

        class MyRBFKernel: # Some implementation of an RBF kernel
            ...

        training_data = ... # Some training data
        x = np.linspace(-10, 10, 100)

        interactive_hyperparameter_figure(MyGPR, MyRBFKernel, training_data, x)
        # This will show an interactive plot of the GPR model with varying hyperparameters.
        ```

        See notes on [`interactive_random_functions_figure`](exercise_2.md#helpers.gpr.interactive_random_functions_figure)
        for more information on the `gpr` argument.
    """
    update_hyperparameters = HyperparameterUpdater(
        gpr, kernel_creator, training_data, x
    )

    interactive_figure(
        data=update_hyperparameters.plot(1, 0.5, 0.1),
        title="GPR with varying hyperparameters and noise",
        x_title="x",
        y_title="f(x)",
        sliders=[
            Slider(
                "sigma",
                min=0.01,
                max=3,
                step=0.01,
                default=1,
                description=r"Signal variance (sigma)",
            ),
            Slider(
                "l",
                min=0.01,
                max=3,
                step=0.01,
                default=0.5,
                description=r"Length scale (l)",
            ),
            Slider(
                "noise",
                min=0,
                max=0.5,
                step=0.025,
                default=0.1,
                description=r"Noise",
            ),
        ],
        on_update=update_hyperparameters,
    )
