# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Traits provider surface examples
#
# This notebook shows:
#
# - a composite market-data provider
# - a composite reference-data provider
# - output publisher stubs for engine orchestration

# %%
from __future__ import annotations

import sys
from dataclasses import dataclass, field
from datetime import datetime
from decimal import Decimal
from pathlib import Path

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"))

# %% [markdown]
# ## Imports

# %%
from fuggers_py.core import Currency, Date, Frequency  # noqa: E402
from fuggers_py.data import (
    BondReferenceData,
    BondQuoteOutput,
    BondType,
    CallScheduleEntry,
    CurrencyPair,
    CurveId,
    CurveInput,
    CurveInputs,
    CurveInstrumentType,
    CurvePoint,
    EtfAnalyticsOutput,
    EtfHolding,
    EtfId,
    EtfQuote,
    FloatingRateTerms,
    FxRate,
    InMemoryCurveSource,
    InMemoryEtfQuoteSource,
    InMemoryFixingSource,
    InMemoryFxRateSource,
    InMemoryInflationFixingSource,
    InMemoryQuoteSource,
    InMemoryVolatilitySource,
    IndexFixing,
    InflationFixing,
    InstrumentId,
    IssuerReferenceData,
    IssuerType,
    MarketDataProvider,
    OutputPublisher,
    PortfolioAnalyticsOutput,
    QuoteSide,
    RawQuote,
    RatingRecord,
    ReferenceDataProvider,
    VolPoint,
    VolSurfaceId,
    VolSurfaceType,
    VolatilitySurface,
    YearMonth,
)

# %% [markdown]
# ## Build market-data records

# %%
as_of = Date.from_ymd(2026, 3, 13)

quote = RawQuote(
    instrument_id=InstrumentId("EXAMPLE-BOND"),
    value=Decimal("101.25"),
    side=QuoteSide.MID,
    as_of=as_of,
    timestamp=datetime(2026, 3, 13, 9, 0),
    currency=Currency.USD,
    bid=Decimal("101.10"),
    ask=Decimal("101.40"),
    source="dealer-run",
)

curve_inputs = CurveInputs.from_points(
    CurveId("usd.discount"),
    as_of,
    [
        CurvePoint(Decimal("1.0"), Decimal("0.0425")),
        CurvePoint(Decimal("5.0"), Decimal("0.0390")),
    ],
    instruments=[
        CurveInput(
            instrument_type=CurveInstrumentType.BOND,
            instrument_id=InstrumentId("EXAMPLE-BOND"),
            tenor=Decimal("5.0"),
            rate=Decimal("0.0390"),
            label="5Y on-the-run",
        )
    ],
)

fixing = IndexFixing("SOFR", as_of, Decimal("0.0410"))
surface = VolatilitySurface(
    surface_id=VolSurfaceId("usd.swaption"),
    surface_type=VolSurfaceType.SWAPTION,
    as_of=as_of,
    points=(
        VolPoint(
            expiry=YearMonth(2026, 6),
            tenor=YearMonth(2028, 6),
            strike=Decimal("0.03"),
            volatility=Decimal("0.22"),
        ),
    ),
)
fx_rate = FxRate(
    currency_pair=CurrencyPair.parse("usd/eur"),
    rate=Decimal("0.9200"),
    bid=Decimal("0.9190"),
    ask=Decimal("0.9210"),
    as_of=as_of,
)
inflation = InflationFixing("CPI", "2026-02", Decimal("307.12"))
etf_quote = EtfQuote(
    etf_id=EtfId("agg-usd"),
    market_price=Decimal("101.20"),
    nav=Decimal("101.10"),
    as_of=as_of,
)

# %% [markdown]
# ## Construct a composite market-data provider

# %%
market_data_provider = MarketDataProvider(
    quote_source=InMemoryQuoteSource([quote]),
    curve_input_source=InMemoryCurveSource([curve_inputs]),
    index_fixing_source=InMemoryFixingSource([fixing]),
    volatility_source=InMemoryVolatilitySource([surface]),
    fx_rate_source=InMemoryFxRateSource([fx_rate]),
    inflation_fixing_source=InMemoryInflationFixingSource([inflation]),
    etf_quote_source=InMemoryEtfQuoteSource([etf_quote]),
)

# %%
market_data_provider.get_quote("EXAMPLE-BOND", QuoteSide.BID)

# %%
market_data_provider.get_curve_inputs("usd.discount")

# %%
market_data_provider.get_volatility_surface("usd.swaption"), market_data_provider.get_fx_rate("USD/EUR"), market_data_provider.get_etf_quote("agg-usd")

# %% [markdown]
# ## Construct lightweight reference-data provider stubs

# %%
bond_reference = BondReferenceData(
    instrument_id=InstrumentId("EXAMPLE-BOND"),
    bond_type=BondType.CALLABLE,
    issuer_type=IssuerType.CORPORATE,
    issue_date=Date.from_ymd(2024, 1, 15),
    maturity_date=Date.from_ymd(2031, 1, 15),
    currency=Currency.USD,
    coupon_rate=Decimal("0.0450"),
    frequency=Frequency.SEMI_ANNUAL,
    issuer_name="Example Corp",
    rating="A",
    call_schedule=(CallScheduleEntry(Date.from_ymd(2029, 1, 15), Decimal("100.00")),),
)

issuer_reference = IssuerReferenceData(
    issuer_name="Example Corp",
    issuer_type=IssuerType.CORPORATE,
    country="US",
    sector="Industrials",
    rating="A",
)

rating_record = RatingRecord(
    rating="A",
    agency="S&P",
    instrument_id=InstrumentId("EXAMPLE-BOND"),
    issuer_name="Example Corp",
)

etf_holdings = (
    EtfHolding(instrument_id=InstrumentId("EXAMPLE-BOND"), quantity=Decimal("10")),
    EtfHolding(instrument_id=InstrumentId("OTHER-BOND"), weight=Decimal("0.25")),
)


@dataclass(slots=True)
class ExampleBondSource:
    bonds: dict[InstrumentId, BondReferenceData]

    def get_bond_reference(self, instrument_id: InstrumentId | str) -> BondReferenceData | None:
        return self.bonds.get(InstrumentId.parse(instrument_id))


@dataclass(slots=True)
class ExampleIssuerSource:
    issuers: dict[str, IssuerReferenceData]

    def get_issuer_reference(self, issuer_name: str) -> IssuerReferenceData | None:
        return self.issuers.get(issuer_name)


@dataclass(slots=True)
class ExampleRatingSource:
    ratings: dict[InstrumentId, RatingRecord]

    def get_rating(self, *, instrument_id: InstrumentId | str | None = None, issuer_name: str | None = None) -> RatingRecord | None:
        if instrument_id is not None:
            return self.ratings.get(InstrumentId.parse(instrument_id))
        if issuer_name is None:
            return None
        for record in self.ratings.values():
            if record.issuer_name == issuer_name:
                return record
        return None


@dataclass(slots=True)
class ExampleEtfHoldingsSource:
    holdings: dict[EtfId, tuple[EtfHolding, ...]]

    def get_etf_holdings(self, etf_id: EtfId | str) -> tuple[EtfHolding, ...]:
        return self.holdings.get(EtfId.parse(etf_id), ())


reference_data_provider = ReferenceDataProvider(
    bond_source=ExampleBondSource({bond_reference.instrument_id: bond_reference}),
    issuer_source=ExampleIssuerSource({issuer_reference.issuer_name: issuer_reference}),
    rating_source=ExampleRatingSource({bond_reference.instrument_id: rating_record}),
    etf_holdings_source=ExampleEtfHoldingsSource({EtfId("agg-usd"): etf_holdings}),
)

# %%
reference_data_provider.get_bond_reference("EXAMPLE-BOND")

# %%
reference_data_provider.get_issuer_reference("Example Corp"), reference_data_provider.get_rating(instrument_id="EXAMPLE-BOND")

# %%
reference_data_provider.get_etf_holdings("agg-usd")

# %% [markdown]
# ## Publisher stubs for orchestration

# %%
@dataclass(slots=True)
class CollectingQuotePublisher:
    quotes: list[BondQuoteOutput] = field(default_factory=list)

    def publish_quote(self, quote: BondQuoteOutput) -> None:
        self.quotes.append(quote)


@dataclass(slots=True)
class CollectingEtfPublisher:
    etf_outputs: list[EtfAnalyticsOutput] = field(default_factory=list)

    def publish_etf(self, analytics: EtfAnalyticsOutput) -> None:
        self.etf_outputs.append(analytics)


@dataclass(slots=True)
class CollectingAnalyticsPublisher:
    analytics: list[object] = field(default_factory=list)

    def publish_analytics(self, analytics: object) -> None:
        self.analytics.append(analytics)


@dataclass(slots=True)
class CollectingAlertPublisher:
    alerts: list[tuple[str, str]] = field(default_factory=list)

    def publish_alert(self, message: str, *, severity: str = "info") -> None:
        self.alerts.append((message, severity))


quote_publisher = CollectingQuotePublisher()
etf_publisher = CollectingEtfPublisher()
analytics_publisher = CollectingAnalyticsPublisher()
alert_publisher = CollectingAlertPublisher()

publisher = OutputPublisher(
    quote_publisher=quote_publisher,
    etf_publisher=etf_publisher,
    analytics_publisher=analytics_publisher,
    alert_publisher=alert_publisher,
)

quote_output = BondQuoteOutput(
    instrument_id=InstrumentId("EXAMPLE-BOND"),
    clean_price=Decimal("101.25"),
    dv01=Decimal("0.042"),
    settlement_date=as_of,
    source="pricing-router",
)
etf_output = EtfAnalyticsOutput(nav=Decimal("101.10"), aggregate_dv01=Decimal("12.5"), settlement_date=as_of)
portfolio_output = PortfolioAnalyticsOutput(total_market_value=Decimal("1000000"), aggregate_dv01=Decimal("25.0"), settlement_date=as_of)

# %%
publisher.publish_quote(quote_output)
publisher.publish_etf(etf_output)
publisher.publish_analytics(portfolio_output)
publisher.publish_alert("Curve source stale", severity="warn")

# %%
quote_publisher.quotes, etf_publisher.etf_outputs

# %%
analytics_publisher.analytics, alert_publisher.alerts
