# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # `fuggers_py.math` examples
#
# This notebook demonstrates the **math** layer:
#
# - Root finding (Brent + Hybrid)
# - Interpolation (Linear / Cubic spline / Flat-forward)
# - Parametric curves (Nelson–Siegel)
# - Extrapolation (Smith–Wilson on toy discount factors)

# %%
from __future__ import annotations

import math
import sys
from pathlib import Path

import numpy as np

root = Path.cwd()
if (root / "src").exists():
    sys.path.insert(0, str(root / "src"))
elif (root.parent / "src").exists():
    sys.path.insert(0, str(root.parent / "src"))

from fuggers_py.math import (  # noqa: E402
    FlatForward,
    LinearInterpolator,
    MonotoneConvex,
    NelsonSiegel,
    SmithWilson,
    SolverConfig,
    brent,
    hybrid_numerical,
)

np.set_printoptions(precision=8, suppress=True)

# %% [markdown]
# ## Root finding

# %%
def f(x: float) -> float:
    return x * x - 2.0


config = SolverConfig(tolerance=1e-12, max_iterations=200)

res_brent = brent(f, 1.0, 2.0, config=config)
res_hybrid = hybrid_numerical(f, 1.0, 2.0, 1.5, config=config)

sqrt2 = math.sqrt(2.0)
print("sqrt(2)=", sqrt2)
print("brent:", res_brent)
print("hybrid:", res_hybrid)

# %% [markdown]
# ## Interpolation: Linear / Cubic spline
#
# Build a small toy curve and evaluate it on a grid.

# %%
times = np.array([0.0, 1.0, 2.0, 5.0, 10.0])
zeros = np.array([0.0100, 0.0125, 0.0150, 0.0200, 0.0225])

lin = LinearInterpolator(times, zeros, allow_extrapolation=False)
spl = LinearInterpolator(times, zeros, allow_extrapolation=False)  # baseline comparison

grid = np.linspace(times.min(), times.max(), 11)
lin_vals = np.array([lin.interpolate(float(t)) for t in grid])
lin_d1 = np.array([lin.derivative(float(t)) for t in grid])

print("grid:", grid)
print("linear zero:", lin_vals)
print("linear dz/dt:", lin_d1)

# %%
from fuggers_py.math.interpolation import CubicSpline  # noqa: E402

cub = CubicSpline(times, zeros, allow_extrapolation=True)
cub_vals = np.array([cub.interpolate(float(t)) for t in grid])
cub_d1 = np.array([cub.derivative(float(t)) for t in grid])

print("cubic zero:", cub_vals)
print("cubic dz/dt:", cub_d1)

# %% [markdown]
# ## Interpolation: Flat-forward + Monotone convex

# %%
pillars = np.array([0.5, 1.0, 2.0, 5.0, 10.0])
pillar_zeros = np.array([0.0100, 0.0120, 0.0150, 0.0200, 0.0220])

ff = FlatForward(pillars, pillar_zeros, allow_extrapolation=True)
mc = MonotoneConvex(pillars, pillar_zeros, allow_extrapolation=True)

grid2 = np.array([0.25, 0.5, 0.75, 1.0, 3.0, 7.0, 12.0])
ff_z = np.array([ff.interpolate(float(t)) for t in grid2])
ff_f = np.array([ff.forward_rate(float(t)) for t in grid2])

mc_z = np.array([mc.interpolate(float(t)) for t in grid2])
mc_f = np.array([mc.forward_rate(float(t)) for t in grid2])

print("grid:", grid2)
print("flat-forward zero:", ff_z)
print("flat-forward fwd :", ff_f)
print("monotone-convex zero:", mc_z)
print("monotone-convex fwd :", mc_f)

# %% [markdown]
# ## Parametric: Nelson–Siegel

# %%
ns = NelsonSiegel.new(beta0=0.02, beta1=-0.01, beta2=0.03, tau=2.0)
ns_vals = np.array([ns.interpolate(float(t)) for t in grid2])
ns_d1 = np.array([ns.derivative(float(t)) for t in grid2])

print("Nelson–Siegel yields:", ns_vals)
print("Nelson–Siegel dy/dt:", ns_d1)

# %% [markdown]
# ## Extrapolation: Smith–Wilson (toy discount factors)
#
# We construct a Smith–Wilson extrapolator from a few observed discount factors
# and confirm the long end converges to the UFR.

# %%
ufr = 0.03
alpha = 0.10

obs_maturities = np.array([1.0, 2.0, 3.0, 5.0, 10.0])
obs_zeros = np.array([0.015, 0.0175, 0.020, 0.022, 0.024])
obs_dfs = np.exp(-obs_zeros * obs_maturities)

sw = SmithWilson(obs_maturities, obs_dfs, ufr=ufr, alpha=alpha)

check = np.array([sw.discount_factor(float(t)) for t in obs_maturities])
print("pillar dfs (obs):", obs_dfs)
print("pillar dfs (sw) :", check)

long_maturity = 60.0
df_long = sw.discount_factor(long_maturity)
z_long = -math.log(df_long) / long_maturity
print("60y df:", df_long)
print("60y zero rate:", z_long, "(ufr=", ufr, ")")

