"""
Utility functions for the tensors notebook.
"""

import matplotlib.pyplot as plt
import numpy as np
import plotnine as pn
import plotnine.options as pn_options
import polars as pl
import torch


def plot_points1d(x: torch.Tensor, w: torch.Tensor | np.ndarray | int | float | None = None) -> None:
    """
    Plot 1D input points and a parameter point on a horizontal line.

    Parameters
    ----------
    x : torch.Tensor
        A 1D tensor representing the inut points to be plotted along the x-axis.
    w : torch.Tensor
        A single value tensor representing the parameter point to be plotted.

    """
    df_x = pl.DataFrame({"x": x.numpy(), "y": 0})
    if isinstance(w, torch.Tensor):
        w = w.detach().numpy()
    pn_options.figure_size = (10, 0.8)
    g = (
        pn.ggplot(df_x)
        + pn.geom_hline(yintercept=0, linetype="dotted")
        + pn.geom_point(pn.aes(x="x"), y=0, color="white", fill="yellowgreen", size=4)
        + pn.theme_minimal()
        + pn.theme(
            axis_title_x=pn.element_blank(),
            axis_ticks_y=pn.element_blank(),
            axis_text_y=pn.element_blank(),
            panel_grid_major_y=pn.element_blank(),
            panel_grid_minor_y=pn.element_blank(),
        )
    )
    if w is not None:
        g = g + pn.geom_point(x=w, y=0, color="white", fill="orchid", size=4)
    g.show()


def figure_gradient_1d():
    x = np.linspace(-10, 10, 100)
    y = np.square(x + 4)
    return (
        pn.ggplot(pl.DataFrame({"x": x, "y": y}))
        + pn.geom_line(pn.aes(x="x", y="y"), color="royalblue")
        + pn.geom_abline(slope=12, intercept=12, color="orchid")
        + pn.geom_hline(yintercept=0, linetype="solid")
        + pn.geom_vline(xintercept=0, linetype="solid")
        + pn.geom_segment(x=2, xend=2, y=0, yend=36, linetype="dotted")
        + pn.annotate(geom="label", x=7, y=190, label="(w + 4)²", color="royalblue")
        + pn.annotate(geom="label", x=9, y=85, label="12w + 12", color="orchid")
        + pn.labs(x="w", y="step2")
        + pn.scale_x_continuous(breaks=(-10, -5, 0, 2, 5, 10))
        + pn.theme_minimal()
    )


def figure_gradient_2d():
    x1 = np.linspace(1, 8, 20)
    x2 = np.linspace(1, 8, 20)
    x1, x2 = np.meshgrid(x1, x2)
    y = x1 / x2
    y_grad = 0.2 * x1 - 0.12 * x2 + 3 / 5

    def plot(ax):
        ax.plot_surface(x1, x2, y, alpha=0.5)  # ty:ignore[unresolved-attribute]
        ax.plot_wireframe(x1, x2, y_grad, linewidths=0.5, colors="orchid")  # ty:ignore[unresolved-attribute]
        ax.scatter(3, 5, 3 / 5, marker="o", color="red", s=20)  # ty:ignore[parameter-already-assigned]
        ax.set_xlabel("x1")
        ax.set_ylabel("x2")
        ax.set_zlabel("res")  # ty:ignore[unresolved-attribute]

    _, ax = plt.subplots(subplot_kw={"projection": "3d"}, nrows=1, ncols=3, figsize=(16, 6))
    plot(ax[0])
    ax[1].view_init(elev=0, azim=-90, roll=0)
    plot(ax[1])
    ax[2].view_init(elev=0, azim=0, roll=0)
    plot(ax[2])

    plt.show()
