from __future__ import annotations

from decimal import Decimal

import numpy as np

from fuggers_py.bonds import FixedBondBuilder, YieldCalculationRules
from fuggers_py.core import Currency, Date, Frequency
from fuggers_py.curves import (
    CubicSplineDiscountFactorCurveModel,
    ExponentialSplineCurveModel,
    FittedBondCurveFitter,
    FittedBondFitConfig,
    FittedBondObservation,
    FittedBondObjective,
    RegressionExposure,
    dirty_price_from_curve,
)
from fuggers_py.data import BondReferenceData, BondType, InstrumentId, IssuerType


REFERENCE_DATE = Date.from_ymd(2026, 1, 15)
UNIVERSE = (
    ("UST2Y", 2, Decimal("0.0200"), Decimal("-1.2"), Decimal("400000000"), Decimal("0.80"), False),
    ("UST3Y", 3, Decimal("0.0225"), Decimal("-0.8"), Decimal("500000000"), Decimal("0.90"), False),
    ("UST4Y", 4, Decimal("0.0250"), Decimal("-0.4"), Decimal("700000000"), Decimal("1.00"), True),
    ("UST5Y", 5, Decimal("0.0275"), Decimal("0.0"), Decimal("900000000"), Decimal("1.10"), True),
    ("UST6Y", 6, Decimal("0.0300"), Decimal("0.4"), Decimal("1100000000"), Decimal("1.15"), True),
    ("UST8Y", 8, Decimal("0.0350"), Decimal("0.8"), Decimal("800000000"), Decimal("1.00"), True),
    ("UST10Y", 10, Decimal("0.0400"), Decimal("1.2"), Decimal("600000000"), Decimal("0.85"), False),
)


def exponential_model() -> ExponentialSplineCurveModel:
    return ExponentialSplineCurveModel((Decimal("0.40"), Decimal("1.20")))


def cubic_model() -> CubicSplineDiscountFactorCurveModel:
    return CubicSplineDiscountFactorCurveModel((Decimal("2.0"), Decimal("4.0"), Decimal("6.0"), Decimal("8.0"), Decimal("10.0")))


def _curve_parameters_for(model) -> np.ndarray:
    if isinstance(model, ExponentialSplineCurveModel):
        return np.asarray([0.032, -0.008, 0.004], dtype=float)
    if isinstance(model, CubicSplineDiscountFactorCurveModel):
        return np.asarray([-0.060, -0.118, -0.174, -0.238, -0.305], dtype=float)
    raise TypeError(f"Unsupported helper model: {type(model).__name__}.")


def make_observations(
    *,
    curve_model=None,
    regression_coefficient: Decimal = Decimal("0"),
    mispricings: dict[str, Decimal] | None = None,
    weights: dict[str, Decimal] | None = None,
) -> tuple[tuple[FittedBondObservation, ...], object]:
    model = curve_model or exponential_model()
    curve = model.build_curve(REFERENCE_DATE, _curve_parameters_for(model), max_t=10.0)
    observations: list[FittedBondObservation] = []
    for instrument_id, maturity_years, coupon_rate, exposure_value, amount_outstanding, liquidity_score, benchmark_flag in UNIVERSE:
        bond = (
            FixedBondBuilder.new()
            .with_issue_date(Date.from_ymd(2021, 1, 15))
            .with_maturity_date(REFERENCE_DATE.add_years(maturity_years))
            .with_coupon_rate(coupon_rate)
            .with_frequency(Frequency.SEMI_ANNUAL)
            .with_currency(Currency.USD)
            .with_rules(YieldCalculationRules.us_treasury())
            .build()
        )
        reference_data = BondReferenceData(
            instrument_id=instrument_id,
            bond_type=BondType.FIXED_RATE,
            issuer_type=IssuerType.SOVEREIGN,
            issue_date=Date.from_ymd(2021, 1, 15),
            maturity_date=REFERENCE_DATE.add_years(maturity_years),
            currency=Currency.USD,
            coupon_rate=coupon_rate,
            frequency=Frequency.SEMI_ANNUAL,
            amount_outstanding=amount_outstanding,
            benchmark_flag=benchmark_flag,
            futures_deliverable_flags=("USTFUT",) if maturity_years in {6, 8} else (),
            liquidity_score=liquidity_score,
        )
        curve_dirty = dirty_price_from_curve(bond, curve, REFERENCE_DATE)
        regression_adjustment = regression_coefficient * exposure_value
        mispricing = (mispricings or {}).get(instrument_id, Decimal(0))
        observed_dirty = curve_dirty + regression_adjustment + mispricing
        clean_price = observed_dirty - bond.accrued_interest(REFERENCE_DATE)
        observations.append(
            FittedBondObservation(
                instrument_id=InstrumentId(instrument_id),
                bond=bond,
                clean_price=clean_price,
                settlement_date=REFERENCE_DATE,
                weight=(weights or {}).get(instrument_id),
                regression_exposures=(RegressionExposure("liquidity", exposure_value),),
                reference_data=reference_data,
            )
        )
    return tuple(observations), curve


def fit_result(
    *,
    curve_model=None,
    objective: FittedBondObjective = FittedBondObjective.L2,
    regression_coefficient: Decimal = Decimal("0.25"),
    mispricings: dict[str, Decimal] | None = None,
    weights: dict[str, Decimal] | None = None,
) :
    observations, _ = make_observations(
        curve_model=curve_model,
        regression_coefficient=regression_coefficient,
        mispricings=mispricings,
        weights=weights,
    )
    fitter = FittedBondCurveFitter(
        curve_model=curve_model or exponential_model(),
        config=FittedBondFitConfig(objective=objective),
    )
    return fitter.fit(observations)
