from dataclasses import dataclass
from typing import Union

import jax
import numpy as np
from jax import numpy as jnp

from ..datatypes import Periodogram
from .initialisation import init_basis_and_penalty, init_knots, init_weights


@dataclass
class LogPSplines:
    """Model for log power splines using a B-spline basis and a penalty matrix."""

    degree: int
    diffMatrixOrder: int
    n: int
    basis: jnp.ndarray
    penalty_matrix: jnp.ndarray
    knots: np.ndarray
    weights: jnp.ndarray = None
    parametric_model: Union[jnp.ndarray, None] = None

    def __post_init__(self):
        if self.degree < self.diffMatrixOrder:
            raise ValueError("Degree must be larger than diffMatrixOrder.")
        if self.degree not in [0, 1, 2, 3, 4, 5]:
            raise ValueError("Degree must be between 0 and 5.")
        if self.diffMatrixOrder not in [0, 1, 2]:
            raise ValueError("diffMatrixOrder must be 0, 1, or 2.")
        if self.n_knots < self.degree:
            raise ValueError(f"#knots: {self.n_knots}, degree: {self.degree}")
        assert (
            self.log_parametric_model is not None
        ), "parametric_model must be provided or initialized."

    def __repr__(self):
        return f"LogPSplines(knots={self.n_knots}, degree={self.degree}, n={self.n})"  # , sparsity={self.basis_sparsity:.2f}, penalty_sparsity={self.penalty_sparsity:.2f})"

    @classmethod
    def from_periodogram(
        cls,
        periodogram: Periodogram,
        n_knots: int,
        degree: int,
        diffMatrixOrder: int = 2,
        parametric_model: jnp.ndarray = None,
        knot_kwargs: dict = {},
    ):
        knots = init_knots(
            n_knots, periodogram, parametric_model, **knot_kwargs
        )
        # compute degree based on the number of knots
        basis, penalty_matrix = init_basis_and_penalty(
            knots, degree, periodogram.n, diffMatrixOrder
        )
        model = cls(
            knots=knots,
            degree=degree,
            diffMatrixOrder=diffMatrixOrder,
            n=periodogram.n,
            basis=basis,
            penalty_matrix=penalty_matrix,
            weights=jnp.zeros(basis.shape[1]),
            parametric_model=parametric_model,
        )
        weights = init_weights(jnp.log(periodogram.power), model)
        model.weights = weights
        return model

    @property
    def log_parametric_model(self) -> jnp.ndarray:
        if not hasattr(self, "_log_parametric_model"):
            if self.parametric_model is None:
                self.parametric_model = jnp.ones(self.n)
            self._log_parametric_model = jnp.log(self.parametric_model)
        return self._log_parametric_model

    @property
    def order(self) -> int:
        return self.degree + 1

    @property
    def n_knots(self) -> int:
        return len(self.knots)

    @property
    def n_basis(self) -> int:
        return self.n_knots + self.degree - 1

    def __call__(
        self, weights: jnp.ndarray = None, use_parametric_model=True
    ) -> jnp.ndarray:
        """Compute the weighted sum of the B-spline basis functions minus a constant."""
        if weights is None:
            weights = self.weights
        ln_para = self.log_parametric_model
        if not use_parametric_model:
            ln_para = jnp.zeros_like(ln_para)
        return build_spline(self.basis, weights, ln_para)


@jax.jit
def build_spline(
    ln_spline_basis: jnp.ndarray,
    weights: jnp.ndarray,
    log_parametric: jnp.ndarray,
) -> jnp.ndarray:
    return (ln_spline_basis @ weights) + log_parametric


#     @property
#     def basis_sparsity(self) -> jnp.ndarray:
#         """Compute the sparsity of the spline basis."""
#         return _compute_sparsity(self.basis)
#
#     @property
#     def penalty_sparsity(self) -> jnp.ndarray:
#         """Compute the sparsity of the penalty matrix."""
#         return _compute_sparsity(self.penalty_matrix)
#
# def _compute_sparsity(x) -> jnp.ndarray:
#     # copy x to avoid modifying the original array
#     return jnp.count_nonzero(x) / x.size
