"""Optimization entry points for fitted bond curves.

The fitter optimizes curve parameters and additive regression coefficients
against dirty-price observations. Residuals are computed in dirty price, while
the reported ``bp_residual`` uses ``(observed_yield - fitted_yield) * 10000``.
"""

from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal
from typing import Iterable

import numpy as np
from numpy.typing import NDArray

from fuggers_py.bonds.pricing import BondPricer
from fuggers_py.core.types import Price
from fuggers_py.math.optimization import OptimizationConfig, OptimizationResult, levenberg_marquardt

from .fair_value import dirty_price_from_curve
from .model import (
    CubicSplineDiscountFactorCurveModel,
    ExponentialSplineCurveModel,
    FittedBondDiagnostics,
    FittedBondFitResult,
    FittedBondObservation,
    FittedBondObjective,
    FittedBondPointResult,
    RegressionCoefficient,
)
from .regression import build_design_matrix


def _to_decimal(value: object) -> Decimal:
    if isinstance(value, Decimal):
        return value
    return Decimal(str(value))


CurveModelSpec = ExponentialSplineCurveModel | CubicSplineDiscountFactorCurveModel


@dataclass(frozen=True, slots=True)
class FittedBondFitConfig:
    """Configuration for fitted-bond optimization."""

    objective: FittedBondObjective = FittedBondObjective.L2
    optimization: OptimizationConfig = OptimizationConfig()
    variable_names: tuple[str, ...] | None = None
    metadata_fields: tuple[str, ...] = ()
    initial_curve_parameters: tuple[Decimal, ...] | None = None
    initial_coefficients: tuple[Decimal, ...] | None = None

    def __post_init__(self) -> None:
        object.__setattr__(
            self,
            "metadata_fields",
            tuple(field.strip().lower().replace("-", "_").replace(" ", "_") for field in self.metadata_fields),
        )


@dataclass(frozen=True, slots=True)
class FittedBondCurveFitter:
    """Fit a fitted-bond curve model to clean-price observations."""

    curve_model: CurveModelSpec = ExponentialSplineCurveModel()
    reference_date: object | None = None
    config: FittedBondFitConfig = FittedBondFitConfig()

    def fit(self, observations: Iterable[FittedBondObservation]) -> FittedBondFitResult:
        """Fit the curve model and regression coefficients to observations.

        Observations must share a common settlement date unless a reference
        date is supplied explicitly. Clean prices are taken in percent of par
        and converted to dirty prices for the objective.
        """
        observation_list = tuple(observations)
        if not observation_list:
            raise ValueError("FittedBondCurveFitter requires at least one observation.")
        reference_date = observation_list[0].settlement_date if self.reference_date is None else self.reference_date
        for observation in observation_list:
            if observation.settlement_date != reference_date:
                raise ValueError("FittedBondCurveFitter requires a common settlement/reference date.")

        design_matrix = build_design_matrix(
            observation_list,
            variable_names=self.config.variable_names,
            metadata_fields=self.config.metadata_fields,
        )
        x_matrix = design_matrix.as_numpy()
        weights = np.asarray(
            [float(observation.weight if observation.weight is not None else Decimal(1)) for observation in observation_list],
            dtype=float,
        )
        if np.any(weights <= 0.0):
            raise ValueError("FittedBondCurveFitter weights must be positive.")

        observed_dirty = np.asarray([float(observation.observed_dirty_price()) for observation in observation_list], dtype=float)
        observed_yields = np.asarray([float(observation.observed_yield()) for observation in observation_list], dtype=float)
        max_t = max(float(observation.maturity_years()) for observation in observation_list)

        if self.config.initial_curve_parameters is None:
            initial_curve = self.curve_model.initial_parameters(observed_yields=observed_yields, max_t=max_t)
        else:
            initial_curve = np.asarray([float(value) for value in self.config.initial_curve_parameters], dtype=float)
        coefficient_count = x_matrix.shape[1]
        if self.config.initial_coefficients is None:
            initial_coefficients = np.zeros(coefficient_count, dtype=float)
        else:
            initial_coefficients = np.asarray([float(value) for value in self.config.initial_coefficients], dtype=float)
            if initial_coefficients.size != coefficient_count:
                raise ValueError("initial_coefficients length must match the regression design matrix width.")
        initial_parameters = np.concatenate((initial_curve, initial_coefficients))

        def _modeled(params: NDArray[np.float64]) -> tuple[object, NDArray[np.float64], NDArray[np.float64]]:
            curve_parameter_count = len(self.curve_model.parameter_names())
            curve_parameters = params[:curve_parameter_count]
            coefficient_vector = params[curve_parameter_count:]
            curve = self.curve_model.build_curve(reference_date, curve_parameters, max_t=max_t)
            curve_dirty = np.asarray(
                [float(dirty_price_from_curve(observation.bond, curve, reference_date)) for observation in observation_list],
                dtype=float,
            )
            fair_dirty = curve_dirty + (x_matrix @ coefficient_vector if coefficient_vector.size else 0.0)
            return curve, curve_dirty, fair_dirty

        def residual_vector(params: NDArray[np.float64]) -> NDArray[np.float64]:
            _, _, fair_dirty = _modeled(params)
            return fair_dirty - observed_dirty

        if self.config.objective is FittedBondObjective.L2:
            def residuals_fn(params: NDArray[np.float64]) -> NDArray[np.float64]:
                return np.sqrt(weights) * residual_vector(params)

            optimization_result = levenberg_marquardt(
                residuals_fn,
                initial_parameters,
                config=self.config.optimization,
            )
        else:
            epsilon = max(self.config.optimization.step_size, 1e-4)
            current = np.asarray(initial_parameters, dtype=float)
            converged = False
            total_iterations = 0
            objective_value = float(np.sum(weights * np.abs(residual_vector(current))))
            outer_iterations = max(3, min(12, self.config.optimization.max_iterations))
            lm_max_iterations = max(5, self.config.optimization.max_iterations // outer_iterations)
            lm_config = OptimizationConfig(
                tolerance=self.config.optimization.tolerance,
                max_iterations=lm_max_iterations,
                step_size=self.config.optimization.step_size,
                armijo_c=self.config.optimization.armijo_c,
                backtracking_beta=self.config.optimization.backtracking_beta,
                min_step=self.config.optimization.min_step,
                lm_initial_damping=self.config.optimization.lm_initial_damping,
                lm_damping_increase=self.config.optimization.lm_damping_increase,
                lm_damping_decrease=self.config.optimization.lm_damping_decrease,
            )
            for _ in range(outer_iterations):
                base_residuals = residual_vector(current)
                residual_scale = max(float(np.median(np.abs(base_residuals))), epsilon)
                adaptive_weights = weights / np.sqrt(np.square(base_residuals) + residual_scale * residual_scale)
                adaptive_weights = adaptive_weights / float(np.mean(adaptive_weights))

                def residuals_fn(params: NDArray[np.float64]) -> NDArray[np.float64]:
                    return np.sqrt(adaptive_weights) * residual_vector(params)

                lm_result = levenberg_marquardt(
                    residuals_fn,
                    current,
                    config=lm_config,
                )
                total_iterations += lm_result.iterations
                new_objective_value = float(np.sum(weights * np.abs(residual_vector(lm_result.parameters))))
                if abs(new_objective_value - objective_value) <= self.config.optimization.tolerance:
                    current = lm_result.parameters
                    objective_value = new_objective_value
                    converged = True
                    break
                if float(np.linalg.norm(lm_result.parameters - current)) <= self.config.optimization.tolerance:
                    current = lm_result.parameters
                    objective_value = new_objective_value
                    converged = True
                    break
                current = lm_result.parameters
                objective_value = new_objective_value

            optimization_result = OptimizationResult(
                parameters=current,
                objective_value=objective_value,
                iterations=total_iterations,
                converged=converged,
            )

        curve, curve_dirty, fair_dirty = _modeled(optimization_result.parameters)
        curve_parameter_count = len(self.curve_model.parameter_names())
        curve_parameters = optimization_result.parameters[:curve_parameter_count]
        coefficient_vector = optimization_result.parameters[curve_parameter_count:]

        pricer = BondPricer()
        point_results: list[FittedBondPointResult] = []
        price_residuals: list[float] = []
        bp_residuals: list[float] = []

        for index, observation in enumerate(observation_list):
            observed_dirty_price = observation.observed_dirty_price()
            observed_clean_price = observation.clean_price
            curve_dirty_price = Decimal(str(curve_dirty[index]))
            accrued_interest = observation.bond.accrued_interest(reference_date)
            curve_clean_price = curve_dirty_price - accrued_interest
            fair_dirty_price = Decimal(str(fair_dirty[index]))
            fair_clean_price = fair_dirty_price - accrued_interest
            observed_yield = observation.observed_yield(pricer=pricer)
            fitted_yield = pricer.yield_from_price(
                observation.bond,
                Price.new(fair_clean_price, observation.bond.currency()),
                reference_date,
            ).ytm.value()
            regression_adjustment = fair_dirty_price - curve_dirty_price
            price_residual = observed_dirty_price - fair_dirty_price
            bp_residual = (observed_yield - fitted_yield) * Decimal("10000")
            price_residuals.append(float(price_residual))
            bp_residuals.append(float(bp_residual))
            point_results.append(
                FittedBondPointResult(
                    instrument_id=observation.instrument_id,
                    bond=observation.bond,
                    settlement_date=reference_date,
                    maturity_date=observation.maturity_date(),
                    maturity_years=observation.maturity_years(),
                    coupon_rate=observation.coupon_rate(),
                    weight=observation.weight or Decimal(1),
                    observed_clean_price=observed_clean_price,
                    observed_dirty_price=observed_dirty_price,
                    observed_yield=observed_yield,
                    curve_clean_price=curve_clean_price,
                    curve_dirty_price=curve_dirty_price,
                    fitted_clean_price=fair_clean_price,
                    fitted_dirty_price=fair_dirty_price,
                    fitted_yield=fitted_yield,
                    fair_value_clean_price=fair_clean_price,
                    fair_value_dirty_price=fair_dirty_price,
                    regression_adjustment=regression_adjustment,
                    price_residual=price_residual,
                    bp_residual=bp_residual,
                    reference_data=observation.reference_data,
                )
            )

        weight_sum = float(np.sum(weights))
        weighted_rmse = Decimal(str(float(np.sqrt(np.sum(weights * np.square(price_residuals)) / weight_sum))))
        weighted_mae = Decimal(str(float(np.sum(weights * np.abs(price_residuals)) / weight_sum)))
        weighted_bp_mae = Decimal(str(float(np.sum(weights * np.abs(bp_residuals)) / weight_sum)))
        diagnostics = FittedBondDiagnostics(
            objective_value=_to_decimal(optimization_result.objective_value),
            iterations=optimization_result.iterations,
            converged=optimization_result.converged,
            observation_count=len(observation_list),
            curve_parameter_count=curve_parameter_count,
            regression_parameter_count=coefficient_vector.size,
            weighted_rmse_price=weighted_rmse,
            weighted_mae_price=weighted_mae,
            max_abs_price_residual=Decimal(str(float(np.max(np.abs(price_residuals))))),
            weighted_mean_abs_bp_residual=weighted_bp_mae,
            max_abs_bp_residual=Decimal(str(float(np.max(np.abs(bp_residuals))))),
        )
        return FittedBondFitResult(
            reference_date=reference_date,
            curve_family=self.curve_model.family,
            objective=self.config.objective,
            metadata_fields=self.config.metadata_fields,
            curve=curve,
            curve_parameter_names=self.curve_model.parameter_names(),
            curve_parameters=tuple(Decimal(str(float(value))) for value in curve_parameters),
            coefficients=tuple(
                RegressionCoefficient(name=name, value=Decimal(str(float(value))))
                for name, value in zip(design_matrix.variable_names, coefficient_vector, strict=True)
            ),
            bonds=tuple(point_results),
            diagnostics=diagnostics,
        )


__all__ = ["FittedBondCurveFitter", "FittedBondFitConfig"]
