"""A primitive sochastic process generator."""
import numpy as np
import numpy.typing as npt
import numpy.linalg as npl
import typing as t
from numpy.random import default_rng


class Kernels:
    @classmethod
    def constant(_, c):
        def kernel(t, s):
            shp = np.max(t.shape)
            return np.ones((shp, shp)) * c

        return kernel

    @classmethod
    def squared_exp(_, l):
        def kernel(t, s):
            return np.exp(-np.abs(t - s) ** 2 / l)

        return kernel

    @classmethod
    def periodic(_, a, ω):
        def kernel(t, s):
            return np.exp(-np.abs(np.sin((t - s)) * ω) * a)

        return kernel

    @classmethod
    def squares(_):
        def kernel(t, s):
            return t ** 2 * s ** 2

        return kernel


def svd(k: t.Callable[(npt.ArrayLike, npt.ArrayLike), npt.NDArray], t):
    """Uses the svd decomposition of the covariance matrix to
    generate a naive sample.
    """
    t = np.asarray(t)
    tt, ss = np.meshgrid(t, t, sparse=True)
    Σ = k(tt, ss)

    (U, S, V) = npl.svd(Σ)

    rng = default_rng()
    X = rng.standard_normal(t.shape)

    return U @ np.diag(np.sqrt(S)) @ X