from __future__ import annotations

from collections import Counter
from datetime import datetime, timedelta
from typing import cast

import polars as pl

from uforia.storage import DatasetName, DatasetStore

from .common import encode_feature_payload, horizon_to_seconds, safe_div

SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0
FEATURE_SOURCE = "tcn_boost_features_v2"
LABEL_SOURCE = "tcn_boost_cost_aware_v2"
OPTIONAL_CONTEXT_COLUMNS = {
    "basis_bps",
    "funding_rate",
    "oi_value",
    "liquidation_volume",
    "iv_atm",
    "iv_skew_25d",
    "iv_rv_gap",
}


def build_microstructure_feature_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    horizons: tuple[str, ...] = ("30s", "1m", "2m", "5m"),
    source: str = FEATURE_SOURCE,
) -> pl.DataFrame:
    bars = store.read(
        DatasetName.MICROSTRUCTURE_BAR,
        start=start,
        end=end,
        filters={"asset": asset},
    ).sort("ts")
    if bars.is_empty():
        return pl.DataFrame()
    min_bar_size = int(cast(int, bars.select(pl.col("bar_size_secs").min()).item(0, 0)))
    bars = bars.filter(pl.col("bar_size_secs") == min_bar_size)
    consolidated = (
        bars.group_by("ts", maintain_order=True)
        .agg(
            [
                pl.first("date").alias("date"),
                pl.first("bar_size_secs").alias("bar_size_secs"),
                pl.mean("close").alias("mid_price"),
                pl.mean("microprice").alias("microprice"),
                pl.mean("spread").alias("spread"),
                pl.sum("volume").alias("volume"),
                pl.sum("buy_volume").alias("buy_volume"),
                pl.sum("sell_volume").alias("sell_volume"),
                pl.mean("depth_bid_5").alias("depth_bid_5"),
                pl.mean("depth_ask_5").alias("depth_ask_5"),
                pl.when(pl.col("venue").str.contains("spot"))
                .then(pl.col("close"))
                .otherwise(None)
                .mean()
                .alias("spot_mid"),
                pl.when(
                    pl.col("venue").str.contains("futures") | pl.col("venue").str.contains("perps")
                )
                .then(pl.col("close"))
                .otherwise(None)
                .mean()
                .alias("futures_mid"),
                (pl.max("close") - pl.min("close")).alias("cross_venue_spread"),
            ]
        )
        .sort("ts")
        .with_columns(
            [
                pl.col("microprice").log().diff().alias("return_1"),
                pl.col("microprice").log().diff(5).alias("return_5"),
                (
                    (pl.col("depth_bid_5") - pl.col("depth_ask_5"))
                    / (pl.col("depth_bid_5") + pl.col("depth_ask_5") + 1e-9)
                ).alias("book_imbalance"),
                (
                    (pl.col("microprice") - pl.col("mid_price")) / (pl.col("mid_price") + 1e-9)
                ).alias("microprice_tilt"),
                (
                    (pl.col("buy_volume") - pl.col("sell_volume")) / (pl.col("volume") + 1e-9)
                ).alias("order_flow_imbalance"),
                (pl.col("buy_volume") - pl.col("sell_volume")).alias("signed_volume"),
                (pl.col("spread") / (pl.col("mid_price") + 1e-9)).alias("spread_rel"),
                (pl.col("depth_bid_5") + pl.col("depth_ask_5")).alias("depth_total"),
                (
                    (pl.col("mid_price") - pl.col("microprice"))
                    / (pl.col("microprice") + 1e-9)
                ).alias("microprice_gap"),
            ]
        )
        .with_columns(
            [
                pl.col("return_1").rolling_std(window_size=30).alias("realized_vol"),
                pl.col("signed_volume").rolling_sum(window_size=10).alias("cum_delta_10"),
                pl.col("signed_volume").rolling_sum(window_size=30).alias("cum_delta_30"),
                pl.col("signed_volume").rolling_sum(window_size=50).alias("cum_delta_50"),
                pl.col("order_flow_imbalance").rolling_mean(window_size=12).alias("ofi_mean_12"),
                pl.col("order_flow_imbalance").rolling_std(window_size=12).alias("ofi_std_12"),
                pl.col("book_imbalance").rolling_mean(window_size=12).alias("imbalance_mean_12"),
                pl.col("spread_rel").rolling_mean(window_size=12).alias("spread_rel_mean_12"),
                pl.col("depth_total").pct_change().alias("depth_change_1"),
                pl.col("microprice_tilt").diff(5).alias("microprice_tilt_change_5"),
                pl.when(
                    pl.col("spot_mid").is_not_null() & pl.col("futures_mid").is_not_null()
                )
                .then(
                    ((pl.col("futures_mid") - pl.col("spot_mid")) / (pl.col("spot_mid") + 1e-9))
                    * 10_000.0
                )
                .otherwise(None)
                .alias("basis_bps"),
            ]
        )
    )
    raw_events = store.read(
        DatasetName.MICROSTRUCTURE_RAW_EVENT,
        start=start - timedelta(days=1),
        end=end,
        filters={"asset": asset},
    ).sort("ts")
    oi_context = (
        raw_events.filter(pl.col("open_interest").is_not_null())
        .select(["ts", pl.col("open_interest").alias("oi_value")])
        .sort("ts")
    )
    liquidation_context = (
        raw_events.filter(pl.col("event_type") == "liquidation")
        .group_by_dynamic("ts", every=f"{min_bar_size}s", closed="left")
        .agg(pl.col("size").sum().alias("liquidation_volume"))
        .rename({"ts": "bucket_ts"})
        .sort("bucket_ts")
    )
    funding = store.read(
        DatasetName.FUNDING_RATE_SNAPSHOT,
        start=start - timedelta(days=1),
        end=end,
        filters={"underlying": asset},
    ).sort("ts")
    surface = store.read(
        DatasetName.SURFACE_FEATURES,
        start=start - timedelta(days=7),
        end=end,
        filters={"underlying": asset, "source": "ssvi"},
    ).sort("ts")
    enriched = consolidated.join_asof(
        funding.select(["ts", pl.col("interest_1h").alias("funding_rate")]).sort("ts"),
        on="ts",
        strategy="backward",
    ).join_asof(
        oi_context,
        on="ts",
        strategy="backward",
    ).join_asof(
        liquidation_context,
        left_on="ts",
        right_on="bucket_ts",
        strategy="backward",
    ).join_asof(
        surface.select(
            [
                "ts",
                pl.col("atm_iv_7d").alias("iv_atm"),
                pl.col("rr_25_30d").alias("iv_skew_25d"),
            ]
        ).sort("ts"),
        on="ts",
        strategy="backward",
    )
    enriched = enriched.with_columns(
        [
            pl.col("funding_rate").fill_null(strategy="forward"),
            pl.col("iv_atm").fill_null(strategy="forward"),
            pl.col("iv_skew_25d").fill_null(strategy="forward"),
            (
                pl.col("realized_vol")
                * pl.when(pl.col("bar_size_secs") > 0)
                .then((SECONDS_PER_YEAR / pl.col("bar_size_secs")).sqrt())
                .otherwise(None)
            ).alias("annualized_realized_vol"),
        ]
    ).with_columns(
        [
            (pl.col("iv_atm") - pl.col("annualized_realized_vol")).alias("iv_rv_gap"),
            pl.col("cum_delta_30").alias("cumulative_delta"),
        ]
    )
    required_columns = [
        "return_1",
        "return_5",
        "realized_vol",
        "book_imbalance",
        "microprice_tilt",
        "order_flow_imbalance",
        "cum_delta_10",
        "cum_delta_30",
        "cum_delta_50",
        "cross_venue_spread",
        "ofi_mean_12",
        "ofi_std_12",
        "imbalance_mean_12",
        "spread_rel_mean_12",
        "depth_change_1",
        "microprice_tilt_change_5",
    ]
    context_columns = sorted(OPTIONAL_CONTEXT_COLUMNS)
    rows: list[pl.DataFrame] = []
    for horizon in horizons:
        horizon_seconds = horizon_to_seconds(horizon)
        with_payload = []
        quality_values: list[str] = []
        for row in enriched.iter_rows(named=True):
            payload = {
                "spread_rel": safe_div(row.get("spread"), row.get("mid_price")),
                "depth_total": (row.get("depth_bid_5") or 0.0) + (row.get("depth_ask_5") or 0.0),
                "return_15": row.get("return_1"),
                "volume_delta": (row.get("buy_volume") or 0.0) - (row.get("sell_volume") or 0.0),
                "ofi_mean_12": row.get("ofi_mean_12"),
                "ofi_std_12": row.get("ofi_std_12"),
                "imbalance_mean_12": row.get("imbalance_mean_12"),
                "spread_rel_mean_12": row.get("spread_rel_mean_12"),
                "depth_change_1": row.get("depth_change_1"),
                "microprice_tilt_change_5": row.get("microprice_tilt_change_5"),
                "microprice_gap": row.get("microprice_gap"),
                "cum_delta_10": row.get("cum_delta_10"),
                "cum_delta_30": row.get("cum_delta_30"),
                "cum_delta_50": row.get("cum_delta_50"),
                "basis_bps": row.get("basis_bps"),
                "funding_rate": row.get("funding_rate"),
                "oi_value": row.get("oi_value"),
                "liquidation_volume": row.get("liquidation_volume"),
                "iv_atm": row.get("iv_atm"),
                "iv_skew_25d": row.get("iv_skew_25d"),
                "iv_rv_gap": row.get("iv_rv_gap"),
                "bar_size_secs": row.get("bar_size_secs"),
                "horizon_secs": float(horizon_seconds),
            }
            with_payload.append(encode_feature_payload(payload))
            if any(row.get(column) is None for column in required_columns):
                quality_values.append("insufficient_history")
            elif any(row.get(column) is None for column in context_columns):
                quality_values.append("partial_context")
            else:
                quality_values.append("ok")
        rows.append(
            enriched.select(
                [
                    "ts",
                    "date",
                    pl.lit(asset).alias("asset"),
                    pl.lit(horizon).alias("horizon"),
                    pl.lit(source).alias("source"),
                    "bar_size_secs",
                    "mid_price",
                    "return_1",
                    "return_5",
                    "realized_vol",
                    "book_imbalance",
                    "microprice_tilt",
                    "order_flow_imbalance",
                    "cumulative_delta",
                    "cum_delta_10",
                    "cum_delta_30",
                    "cum_delta_50",
                    "basis_bps",
                    pl.col("funding_rate"),
                    "oi_value",
                    "liquidation_volume",
                    pl.col("iv_atm"),
                    pl.col("iv_skew_25d"),
                    pl.col("cross_venue_spread"),
                    pl.col("iv_rv_gap"),
                ]
            ).with_columns(
                [
                    pl.Series("feature_count", [20] * enriched.height, dtype=pl.Int64),
                    pl.Series("feature_payload", with_payload, dtype=pl.String),
                    pl.Series("quality_flag", quality_values, dtype=pl.String),
                ]
            )
        )
    return pl.concat(rows, how="vertical_relaxed") if rows else pl.DataFrame()


def build_microstructure_label_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    horizons: tuple[str, ...] = ("30s", "1m", "2m", "5m"),
    tau_bps: float = 2.5,
    source: str = LABEL_SOURCE,
) -> pl.DataFrame:
    bars = store.read(
        DatasetName.MICROSTRUCTURE_BAR,
        start=start,
        end=end,
        filters={"asset": asset},
    ).sort("ts")
    if bars.is_empty():
        return pl.DataFrame()
    min_bar_size = int(cast(int, bars.select(pl.col("bar_size_secs").min()).item(0, 0)))
    bars = bars.filter(pl.col("bar_size_secs") == min_bar_size)
    consolidated = (
        bars.group_by("ts", maintain_order=True)
        .agg(
            [
                pl.first("date").alias("date"),
                pl.mean("close").alias("price"),
                pl.mean("microprice").alias("microprice"),
            ]
        )
        .sort("ts")
    )
    rows: list[pl.DataFrame] = []
    for horizon in horizons:
        horizon_seconds = horizon_to_seconds(horizon)
        if horizon_seconds % min_bar_size != 0:
            labelled = consolidated.with_columns(
                [
                    pl.lit(None, dtype=pl.Float64).alias("forward_return_bps"),
                    pl.lit(None, dtype=pl.Float64).alias("forward_microprice_return_bps"),
                    pl.lit(None, dtype=pl.Float64).alias("forward_move_abs_bps"),
                    pl.lit(None, dtype=pl.Int64).alias("label"),
                    pl.lit(min_bar_size).cast(pl.Int64).alias("bar_size_secs"),
                    pl.lit(None, dtype=pl.Int64).alias("label_horizon_steps"),
                    pl.lit(horizon_seconds).cast(pl.Int64).alias("horizon_seconds"),
                    pl.lit("incompatible_bar_size").alias("quality_flag"),
                ]
            )
        else:
            step = max(1, horizon_seconds // min_bar_size)
            labelled = consolidated.with_columns(
                [
                    ((pl.col("price").shift(-step) / pl.col("price")).log() * 10_000).alias(
                        "forward_return_bps"
                    ),
                    (
                        (pl.col("microprice").shift(-step) / pl.col("microprice")).log() * 10_000
                    ).alias("forward_microprice_return_bps"),
                ]
            ).with_columns(
                [
                    pl.col("forward_microprice_return_bps").abs().alias("forward_move_abs_bps"),
                    pl.when(pl.col("forward_microprice_return_bps") < -tau_bps)
                    .then(pl.lit(0))
                    .when(pl.col("forward_microprice_return_bps") > tau_bps)
                    .then(pl.lit(2))
                    .otherwise(pl.lit(1))
                    .alias("label"),
                    pl.lit(min_bar_size).cast(pl.Int64).alias("bar_size_secs"),
                    pl.lit(step).cast(pl.Int64).alias("label_horizon_steps"),
                    pl.lit(horizon_seconds).cast(pl.Int64).alias("horizon_seconds"),
                    pl.when(pl.col("forward_microprice_return_bps").is_null())
                    .then(pl.lit("insufficient_forward_window"))
                    .otherwise(pl.lit("ok"))
                    .alias("quality_flag"),
                ]
            )
        rows.append(
            labelled.select(
                [
                    "ts",
                    "date",
                    pl.lit(asset).alias("asset"),
                    pl.lit(horizon).alias("horizon"),
                    pl.lit(source).alias("source"),
                    "forward_microprice_return_bps",
                    "forward_return_bps",
                    "forward_move_abs_bps",
                    "label",
                    pl.lit(tau_bps).alias("tau_bps"),
                    "bar_size_secs",
                    "label_horizon_steps",
                    "horizon_seconds",
                    "quality_flag",
                ]
            )
        )
    return pl.concat(rows, how="vertical_relaxed") if rows else pl.DataFrame()


def build_microstructure_regime_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    source: str = "tcn_boost_regime_v1",
) -> pl.DataFrame:
    features = store.read(
        DatasetName.MICROSTRUCTURE_FEATURE_SERIES,
        start=start,
        end=end,
        filters={"asset": asset, "horizon": "1m"},
    ).sort("ts")
    if features.is_empty():
        return pl.DataFrame()
    frame = features.with_columns(
        [
            pl.col("realized_vol")
            .rank(method="average")
            .truediv(pl.len())
            .alias("vol_percentile"),
            pl.col("order_flow_imbalance")
            .abs()
            .rank(method="average")
            .truediv(pl.len())
            .alias("flow_percentile"),
            pl.col("basis_bps")
            .abs()
            .rank(method="average")
            .truediv(pl.len())
            .alias("basis_percentile"),
        ]
    )
    regime_labels: list[str] = []
    persistence: list[float] = []
    current = "calm_trend"
    current_length = 0
    for row in frame.iter_rows(named=True):
        vol = float(row.get("vol_percentile") or 0.0)
        flow = float(row.get("flow_percentile") or 0.0)
        if vol >= 0.8:
            next_state = "vol_expansion"
        elif flow >= 0.8:
            next_state = "thin_liquidity_stressed"
        elif vol <= 0.25 and flow <= 0.4:
            next_state = "calm_mean_reversion"
        else:
            next_state = "calm_trend"
        if next_state == current:
            current_length += 1
        else:
            current = next_state
            current_length = 1
        regime_labels.append(next_state)
        persistence.append(float(current_length))
    return frame.select(
        [
            "ts",
            "date",
            pl.lit(asset).alias("asset"),
            pl.lit(source).alias("source"),
            "vol_percentile",
            "flow_percentile",
            "basis_percentile",
        ]
    ).with_columns(
        [
            pl.Series("regime_state", regime_labels, dtype=pl.String),
            pl.Series("persistence_hours", persistence, dtype=pl.Float64),
            pl.Series("hysteresis_state", regime_labels, dtype=pl.String),
            pl.lit("ok").alias("quality_flag"),
        ]
    )


def build_microstructure_signal_series(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    horizons: tuple[str, ...] = ("30s", "1m", "2m", "5m"),
    source: str = "tcn_boost_paper_signal",
) -> pl.DataFrame:
    outputs = store.read(
        DatasetName.MICROSTRUCTURE_MODEL_OUTPUT,
        start=start,
        end=end,
        filters={"asset": asset},
    ).sort("ts")
    if outputs.is_empty():
        return pl.DataFrame()
    rows: list[pl.DataFrame] = []
    for horizon in horizons:
        frame = outputs.filter(pl.col("horizon") == horizon)
        if frame.is_empty():
            continue
        rows.append(
            frame.select(
                [
                    "ts",
                    "date",
                    pl.lit(asset).alias("asset"),
                    pl.lit(horizon).alias("horizon"),
                    pl.lit(source).alias("source"),
                    pl.when(
                        (pl.col("blocked_reason").is_null())
                        & (pl.col("prob_up") > pl.col("prob_down"))
                        & (pl.col("prob_up") > pl.col("prob_flat"))
                    )
                    .then(pl.lit("long_bias"))
                    .when(
                        (pl.col("blocked_reason").is_null())
                        & (pl.col("prob_down") > pl.col("prob_up"))
                        & (pl.col("prob_down") > pl.col("prob_flat"))
                    )
                    .then(pl.lit("short_bias"))
                    .otherwise(pl.lit("flat"))
                    .alias("signal_state"),
                    pl.when(
                        (pl.col("blocked_reason").is_null())
                        & (pl.col("prob_up") > pl.col("prob_down"))
                        & (pl.col("prob_up") > pl.col("prob_flat"))
                    )
                    .then(pl.lit("long"))
                    .when(
                        (pl.col("blocked_reason").is_null())
                        & (pl.col("prob_down") > pl.col("prob_up"))
                        & (pl.col("prob_down") > pl.col("prob_flat"))
                    )
                    .then(pl.lit("short"))
                    .otherwise(pl.lit("flat"))
                    .alias("paper_position"),
                    pl.col("expected_edge_bps"),
                    (pl.col("blocked_reason").is_null()).alias("cost_filter_passed"),
                    (pl.col("quality_flag") == "ok").alias("quality_filter_passed"),
                    pl.col("blocked_reason"),
                    pl.col("confidence"),
                    pl.col("tradability_score"),
                    pl.col("quality_flag"),
                ]
            )
        )
    return pl.concat(rows, how="vertical_relaxed") if rows else pl.DataFrame()


def summarize_transition_probabilities(states: list[str]) -> tuple[dict[str, float], float]:
    if len(states) < 2:
        return {}, 0.0
    transitions: Counter[tuple[str, str]] = Counter(zip(states[:-1], states[1:], strict=False))
    origin_counts: Counter[str] = Counter(states[:-1])
    latest = states[-1]
    probs = {
        target: transitions[(latest, target)] / origin_counts[latest]
        for target in {"long_bias", "flat", "short_bias"}
        if origin_counts[latest] > 0
    }
    durations: list[int] = []
    run = 1
    for left, right in zip(states[:-1], states[1:], strict=False):
        if left == right:
            run += 1
        else:
            durations.append(run)
            run = 1
    durations.append(run)
    median = float(sorted(durations)[len(durations) // 2])
    return probs, median
