"""
Feature extraction for planner and reward functions.

Rule: every human-obvious feature is allowed, but the CMA-ES optimiser
must be allowed to discover non-obvious interactions (queue churn, time
since MFE, recovery velocity, cross-venue lead, etc.).
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Mapping, Protocol

from malkhut.state import MarketWorldState


@dataclass(frozen=True, slots=True)
class FeatureVector:
    values: Mapping[str, float]


class FeatureExtractor(Protocol):
    def extract(self, state: MarketWorldState) -> FeatureVector: ...


class DefaultFeatureExtractor:
    def extract(self, state: MarketWorldState) -> FeatureVector:
        b = state.book
        bid_qty = sum(x.qty for x in b.bids[:5])
        ask_qty = sum(x.qty for x in b.asks[:5])
        imbalance = (bid_qty - ask_qty) / max(bid_qty + ask_qty, 1e-12)

        path = state.trade_path

        values = {
            "mid": b.mid if b.bids and b.asks else 0.0,
            "spread_bps": b.spread_bps if b.bids and b.asks else 0.0,
            "top5_imbalance": imbalance,
            "funding_bps": state.funding_bps or 0.0,
            "volatility_state": state.volatility_state or 0.0,
            "pnl_bps": path.pnl_bps if path else 0.0,
            "mae_bps": path.mae_bps if path else 0.0,
            "mfe_bps": path.mfe_bps if path else 0.0,
            "distance_from_mfe_bps": path.distance_from_mfe_bps if path else 0.0,
            "seconds_held": path.seconds_held if path else 0.0,
            "time_in_loss_s": path.time_in_loss_s if path else 0.0,
            "time_since_deep_mae_s": path.time_since_deep_mae_s if path else 0.0,
            "recovery_velocity_bps_per_s": path.recovery_velocity_bps_per_s if path else 0.0,
            "adverse_velocity_bps_per_s": path.adverse_velocity_bps_per_s if path else 0.0,
            "orderflow_toxicity": path.orderflow_toxicity if path else 0.0,
            "queue_churn_score": path.queue_churn_score if path else 0.0,
            "cross_venue_lead_score": path.cross_venue_lead_score if path else 0.0,
        }
        return FeatureVector(values=values)