"""Smith–Wilson extrapolation for discount factors."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from numpy.typing import ArrayLike, NDArray

from ..errors import InvalidInput
from ..linear_algebra import solve_linear_system
from ..numerical import finite_difference_derivative
from ..utils import assert_all_positive, assert_finite, assert_same_length, assert_strictly_increasing


def _wilson_kernel(t: float, u: float, *, alpha: float, ufr: float) -> float:
    """Wilson function W(t,u) used by Smith–Wilson."""

    m = min(t, u)
    M = max(t, u)
    exp_ufr = np.exp(-ufr * (t + u))
    term = alpha * m - np.exp(-alpha * M) * np.sinh(alpha * m)
    return float(exp_ufr * term)


@dataclass(frozen=True, slots=True)
class SmithWilson:
    ufr: float
    alpha: float
    maturities: NDArray[np.float64]
    discount_factors: NDArray[np.float64]
    weights: NDArray[np.float64]

    def __init__(
        self,
        maturities: ArrayLike,
        discount_factors: ArrayLike,
        *,
        ufr: float,
        alpha: float,
    ) -> None:
        t = assert_strictly_increasing(maturities, name="maturities")
        p = assert_all_positive(discount_factors, name="discount_factors")
        assert_same_length(t, p, a_name="maturities", b_name="discount_factors")

        ufr_f = assert_finite(ufr, name="ufr")
        alpha_f = assert_finite(alpha, name="alpha")
        if alpha_f <= 0.0:
            raise InvalidInput("alpha must be positive.")
        if np.any(t <= 0.0):
            raise InvalidInput("maturities must be strictly positive.")

        n = int(t.size)
        w_mat = np.empty((n, n), dtype=float)
        for i in range(n):
            for j in range(n):
                w_mat[i, j] = _wilson_kernel(float(t[i]), float(t[j]), alpha=alpha_f, ufr=ufr_f)

        u_vec = np.exp(-ufr_f * t)
        m_vec = p - u_vec
        weights = solve_linear_system(w_mat, m_vec)

        object.__setattr__(self, "ufr", ufr_f)
        object.__setattr__(self, "alpha", alpha_f)
        object.__setattr__(self, "maturities", t)
        object.__setattr__(self, "discount_factors", p)
        object.__setattr__(self, "weights", weights)

    def _u(self, t: float) -> float:
        return float(np.exp(-self.ufr * t))

    def discount_factor(self, t: float) -> float:
        t = assert_finite(t, name="t")
        if t < 0.0:
            raise InvalidInput("t must be non-negative.")
        base = self._u(t)
        k = np.array(
            [_wilson_kernel(t, float(u), alpha=self.alpha, ufr=self.ufr) for u in self.maturities],
            dtype=float,
        )
        return float(base + float(np.dot(self.weights, k)))

    def extrapolate(self, x: float) -> float:
        return self.discount_factor(float(x))

    def derivative(self, x: float) -> float:
        return float(finite_difference_derivative(self.discount_factor, float(x), h=1e-6, method="central"))
