"""Flat-forward interpolation for zero rates (piecewise-constant forwards)."""

from __future__ import annotations

from dataclasses import dataclass

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

from ..errors import ExtrapolationNotAllowed, InsufficientData, InvalidInput
from ..utils import assert_finite, assert_finite_array, assert_same_length, assert_strictly_increasing, bisect_segment


@dataclass(frozen=True, slots=True)
class FlatForward:
    _tenors: NDArray[np.float64]
    _zero_rates: NDArray[np.float64]
    _forwards: NDArray[np.float64]
    _allow_extrapolation: bool = False

    def __init__(self, tenors: ArrayLike, zero_rates: ArrayLike, allow_extrapolation: bool = False) -> None:
        t = assert_strictly_increasing(tenors, name="tenors")
        r = assert_finite_array(zero_rates, name="zero_rates").astype(float, copy=False)
        assert_same_length(t, r, a_name="tenors", b_name="zero_rates")
        if t.size < 2:
            raise InsufficientData(required=2, actual=int(t.size))
        if np.any(t <= 0.0):
            raise InvalidInput("tenors must be strictly positive.")

        forwards = np.empty(int(t.size - 1), dtype=float)
        for i in range(forwards.size):
            t0 = float(t[i])
            t1 = float(t[i + 1])
            r0 = float(r[i])
            r1 = float(r[i + 1])
            forwards[i] = (r1 * t1 - r0 * t0) / (t1 - t0)

        object.__setattr__(self, "_tenors", t)
        object.__setattr__(self, "_zero_rates", r)
        object.__setattr__(self, "_forwards", forwards)
        object.__setattr__(self, "_allow_extrapolation", bool(allow_extrapolation))

    @classmethod
    def new(cls, tenors: ArrayLike, zero_rates: ArrayLike) -> "FlatForward":
        return cls(tenors, zero_rates)

    @property
    def allows_extrapolation(self) -> bool:
        return bool(self._allow_extrapolation)

    def with_extrapolation(self) -> "FlatForward":
        return FlatForward(self._tenors, self._zero_rates, allow_extrapolation=True)

    def tenors(self) -> NDArray[np.float64]:
        return self._tenors.copy()

    def zero_rates(self) -> NDArray[np.float64]:
        return self._zero_rates.copy()

    def forward_rates_vec(self) -> NDArray[np.float64]:
        return self._forwards.copy()

    def min_x(self) -> float:
        return float(self._tenors[0])

    def max_x(self) -> float:
        return float(self._tenors[-1])

    def in_range(self, x: float) -> bool:
        x = float(x)
        return self.min_x() <= x <= self.max_x()

    def _validate_t(self, t: float) -> float:
        t = assert_finite(t, name="t")
        if t <= 0.0:
            raise InvalidInput("t must be positive.")
        if not self.allows_extrapolation and not self.in_range(t):
            raise ExtrapolationNotAllowed(x=t, min=self.min_x(), max=self.max_x())
        return t

    def forward_rate(self, t: float) -> float:
        t = self._validate_t(t)
        if t < self._tenors[0]:
            return float(self._zero_rates[0])
        if t > self._tenors[-1]:
            return float(self._forwards[-1])
        i = bisect_segment(self._tenors, t)
        return float(self._forwards[i])

    def interpolate(self, t: float) -> float:
        t = self._validate_t(t)
        if t <= self._tenors[0]:
            return float(self._zero_rates[0])
        if t >= self._tenors[-1]:
            if t == self._tenors[-1]:
                return float(self._zero_rates[-1])
            f_last = float(self._forwards[-1])
            r_last = float(self._zero_rates[-1])
            t_last = float(self._tenors[-1])
            return (r_last * t_last + f_last * (t - t_last)) / t

        i = bisect_segment(self._tenors, t)
        t0 = float(self._tenors[i])
        r0 = float(self._zero_rates[i])
        f = float(self._forwards[i])
        return (r0 * t0 + f * (t - t0)) / t

    def derivative(self, t: float) -> float:
        t = self._validate_t(t)
        if t <= self._tenors[0]:
            return 0.0

        if t >= self._tenors[-1]:
            if t == self._tenors[-1]:
                i = int(self._tenors.size - 2)
            else:
                i = int(self._tenors.size - 2)
        else:
            i = bisect_segment(self._tenors, t)

        t0 = float(self._tenors[i])
        r0 = float(self._zero_rates[i])
        f = float(self._forwards[i])
        return -(t0 * (r0 - f)) / (t**2)
