from __future__ import annotations

import numpy as np
import pandas as pd

from optpricing.atoms import Option, OptionType, Rate, Stock
from optpricing.models import BaseModel
from optpricing.techniques.base import BaseTechnique

from .vectorized_bsm_iv import BSMIVSolver

__doc__ = """
Defines the VolatilitySurface class, used to compute and store implied
volatility surfaces from both market and model prices.
"""


class VolatilitySurface:
    """
    A class to compute and hold market and model-implied volatility surfaces.

    This class takes a DataFrame of option market data and provides methods
    to calculate the Black-Scholes implied volatility (IV) for each option,
    either from its market price or from a price generated by a financial model.
    """

    def __init__(self, option_data: pd.DataFrame):
        """
        Initializes the VolatilitySurface.

        Parameters
        ----------
        option_data : pd.DataFrame
            A DataFrame containing market prices of options. Must include
            'strike', 'maturity', 'marketPrice', 'optionType', and 'expiry' columns.

        Raises
        ------
        ValueError
            If any of the required columns are missing from `option_data`.
        """
        required_cols = ["strike", "maturity", "marketPrice", "optionType", "expiry"]
        if not all(col in option_data.columns for col in required_cols):
            msg = (
                "Input option_data is missing one "
                f"of the required columns: {required_cols}"
            )
            raise ValueError(msg)

        self.data = option_data[required_cols].copy()
        self.surface: pd.DataFrame | None = None
        self.iv_solver = BSMIVSolver()

    def _calculate_ivs(
        self,
        stock: Stock,
        rate: Rate,
        prices_to_invert: pd.Series,
    ) -> np.ndarray:
        """Calculates IVs using the fast, vectorized BSM solver."""
        ivs = self.iv_solver.solve(prices_to_invert.values, self.data, stock, rate)
        return ivs

    def calculate_market_iv(
        self,
        stock: Stock,
        rate: Rate,
    ) -> VolatilitySurface:
        """
        Calculates the market implied volatility surface from market prices.

        This method inverts the Black-Scholes formula for each option's market
        price to find the corresponding implied volatility. The results are
        stored in the `self.surface` DataFrame.

        Parameters
        ----------
        stock : Stock
            The underlying asset's properties.
        rate : Rate
            The risk-free rate structure.

        Returns
        -------
        VolatilitySurface
            The same instance of the class, allowing for method chaining.
        """
        print("Calculating market implied volatility surface...")
        market_ivs = self._calculate_ivs(stock, rate, self.data["marketPrice"])
        self.surface = self.data.copy()
        self.surface["iv"] = market_ivs
        self.surface.dropna(inplace=True)
        self.surface = self.surface[
            (self.surface["iv"] > 1e-4) & (self.surface["iv"] < 2.0)
        ]
        return self

    def calculate_model_iv(
        self,
        stock: Stock,
        rate: Rate,
        model: BaseModel,
        technique: BaseTechnique,
    ) -> VolatilitySurface:
        """
        Calculates a model's implied volatility surface.

        This method first prices every option in the dataset using the provided
        model and technique. It then inverts the Black-Scholes formula for each
        of these model prices to generate the model-implied volatility surface.

        Parameters
        ----------
        stock : Stock
            The underlying asset's properties.
        rate : Rate
            The risk-free rate structure.
        model : BaseModel
            The financial model to generate prices from.
        technique : BaseTechnique
            The pricing technique to use with the model.

        Returns
        -------
        VolatilitySurface
            The same instance of the class, allowing for method chaining.
        """
        print(f"Calculating {model.name} implied volatility surface...")

        model_prices = np.array(
            [
                technique.price(
                    Option(
                        strike=row.strike,
                        maturity=row.maturity,
                        option_type=OptionType.CALL
                        if row.optionType == "call"
                        else OptionType.PUT,
                    ),
                    stock,
                    model,
                    rate,
                    **model.params,
                ).price
                for _, row in self.data.iterrows()
            ]
        )

        model_ivs = self._calculate_ivs(
            stock, rate, pd.Series(model_prices, index=self.data.index)
        )
        self.surface = self.data.copy()
        self.surface["iv"] = model_ivs
        self.surface.dropna(inplace=True)
        return self
