from __future__ import annotations

import polars as pl
import numpy as np
from typing import Iterable, Optional


# Core aggregation
def aggregate_chain_greeks(
    df: pl.DataFrame,
) -> dict:
    """
    Aggregates per-option Greeks into chain-level exposure features.
    """

    required = {
        "delta",
        "gamma",
        "vega",
        "theta",
        "charm",
        "vanna",
        "vomma",
        "contract_size",
    }

    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"Missing required columns: {sorted(missing)}")

    # scale by contract size
    w = df["contract_size"]

    delta_exp = (df["delta"] * w).sum()
    gamma_exp = (df["gamma"] * w).sum()
    vega_exp = (df["vega"] * w).sum()
    theta_exp = (df["theta"] * w).sum()
    charm_exp = (df["charm"] * w).sum()
    vanna_exp = (df["vanna"] * w).sum()
    vomma_exp = (df["vomma"] * w).sum()

    # distributional summaries (delta-weighted)
    def _moments(x, w):
        wsum = w.sum()
        if wsum == 0:
            return 0.0, 0.0, 0.0

        mean = (x * w).sum() / wsum
        var = ((x - mean) ** 2 * w).sum() / wsum
        std = np.sqrt(var)

        if std == 0:
            skew = 0.0
        else:
            skew = ((x - mean) ** 3 * w).sum() / (wsum * std**3)

        return mean, std, skew

    gamma_mean_d, gamma_std_d, gamma_skew_d = _moments(df["delta"], df["gamma"].abs())
    vega_mean_d, vega_std_d, vega_skew_d = _moments(df["delta"], df["vega"].abs())
    theta_mean_d, theta_std_d, theta_skew_d = _moments(df["delta"], df["theta"].abs())

    return {
        # totals
        "delta_exposure_total": float(delta_exp),
        "gamma_exposure_total": float(gamma_exp),
        "vega_exposure_total": float(vega_exp),
        "theta_exposure_total": float(theta_exp),
        "charm_exposure_total": float(charm_exp),
        "vanna_exposure_total": float(vanna_exp),
        "vomma_exposure_total": float(vomma_exp),

        # distributional structure
        "gamma_mean_delta": float(gamma_mean_d),
        "gamma_std_delta": float(gamma_std_d),
        "gamma_skew_delta": float(gamma_skew_d),

        "vega_mean_delta": float(vega_mean_d),
        "vega_std_delta": float(vega_std_d),
        "vega_skew_delta": float(vega_skew_d),

        "theta_mean_delta": float(theta_mean_d),
        "theta_std_delta": float(theta_std_d),
        "theta_skew_delta": float(theta_skew_d),
    }


# Optional delta-bucket view
def aggregate_chain_greeks_by_delta(
    df: pl.DataFrame,
    *,
    delta_buckets: Optional[Iterable[float]] = None,
) -> pl.DataFrame:
    """
    Aggregates Greeks by delta buckets.
    """

    if delta_buckets is None:
        delta_buckets = [-1.0, -0.75, -0.50, -0.25, 0.0, 0.25, 0.50, 0.75, 1.0]

    labels = [
        f"{delta_buckets[i]:.2f}_{delta_buckets[i+1]:.2f}"
        for i in range(len(delta_buckets) - 1)
    ]

    df = df.with_columns(
        pl.cut(
            pl.col("delta"),
            bins=list(delta_buckets),
            labels=labels,
        ).alias("delta_bucket")
    )

    grouped = (
        df
        .group_by("delta_bucket")
        .agg([
            (pl.col("delta") * pl.col("contract_size")).sum().alias("DEX"),
            (pl.col("gamma") * pl.col("contract_size")).sum().alias("GEX"),
            (pl.col("vega") * pl.col("contract_size")).sum().alias("VEX"),
            (pl.col("theta") * pl.col("contract_size")).sum().alias("TEX"),
            (pl.col("charm") * pl.col("contract_size")).sum().alias("CEX"),
            (pl.col("vanna") * pl.col("contract_size")).sum().alias("VANEX"),
            (pl.col("vomma") * pl.col("contract_size")).sum().alias("VOMEX"),
            pl.col("contract_size").sum().alias("total_open_interest"),
        ])
        .sort("delta_bucket")
    )

    totals = (
        grouped
        .select([
            pl.lit("ALL").alias("delta_bucket"),
            pl.sum("DEX"),
            pl.sum("GEX"),
            pl.sum("VEX"),
            pl.sum("TEX"),
            pl.sum("CEX"),
            pl.sum("VANEX"),
            pl.sum("VOMEX"),
            pl.sum("total_open_interest"),
        ])
    )

    return pl.concat([grouped, totals], how="vertical")