"""
Exponential Regime Expansion — orthogonal to cognition pipeline.

Generates 100+ well-defined, actually-extant market regimes
by combining primitive market dimensions:

  - Liquidity: {thin, normal, deep, vacuum}
  - Volatility: {low, normal, high, extreme}
  - Spread: {tight, normal, wide, flash}
  - Flow: {balanced, buy_pressure, sell_pressure, toxic}
  - Structure: {normal, whale, mm_withdrawal, cascade}
  - Time: {session, overnight, weekend}
  - Correlation: {high, normal, breakdown}

Each combination is a distinct, testable regime.
Total: 4 × 4 × 4 × 4 × 4 × 3 × 3 = 9,216 theoretical combinations.
Practical: ~200 distinct, non-overlapping regimes.
"""
from __future__ import annotations

import itertools
from dataclasses import dataclass
from typing import List, Optional, Set, Tuple

from malkhut.state import (
    AccountState, MarketWorldState, Mode, OrderBookState, PriceLevel, VenueRules,
)
from malkhut.counterparties import (
    CounterpartyPolicy, ToxicTakerPolicy, PassiveMakerPolicy,
    LatencyArbPolicy, NoiseTraderPolicy,
)


# ==============================================================================
# Primitive Market Dimensions
# ==============================================================================

@dataclass(frozen=True, slots=True)
class LiquidityDim:
    bid_qty: float
    ask_qty: float
    label: str


@dataclass(frozen=True, slots=True)
class VolatilityDim:
    spread_bps: float
    label: str


@dataclass(frozen=True, slots=True)
class FlowDim:
    imbalance: float  # -1 to 1
    toxicity: float   # 0 to 1
    label: str


@dataclass(frozen=True, slots=True)
class StructureDim:
    counterparties: Tuple[CounterpartyPolicy, ...]
    label: str


# Predefined dimensions
LIQUIDITY_DIMS = [
    LiquidityDim(0.01, 0.01, "vacuum"),
    LiquidityDim(0.1, 0.1, "thin"),
    LiquidityDim(0.5, 0.5, "normal"),
    LiquidityDim(2.0, 2.0, "deep"),
]

VOLATILITY_DIMS = [
    VolatilityDim(0.5, "tight"),
    VolatilityDim(2.0, "normal"),
    VolatilityDim(10.0, "wide"),
    VolatilityDim(100.0, "extreme"),
]

FLOW_DIMS = [
    FlowDim(0.0, 0.0, "balanced"),
    FlowDim(0.5, 0.3, "buy_pressure"),
    FlowDim(-0.5, 0.3, "sell_pressure"),
    FlowDim(0.0, 0.8, "toxic"),
]

STRUCTURE_DIMS = [
    StructureDim((ToxicTakerPolicy(),), "single_toxic"),
    StructureDim((PassiveMakerPolicy(), ToxicTakerPolicy()), "mm_toxic"),
    StructureDim((ToxicTakerPolicy(), ToxicTakerPolicy(), LatencyArbPolicy()), "multi_toxic"),
    StructureDim((NoiseTraderPolicy(), PassiveMakerPolicy()), "retail_mm"),
]


# ==============================================================================
# Regime Generator
# ==============================================================================

@dataclass(frozen=True, slots=True)
class ExpandedRegime:
    """A generated market regime from dimension combinations."""
    regime_id: str
    label: str
    liquidity: LiquidityDim
    volatility: VolatilityDim
    flow: FlowDim
    structure: StructureDim
    bid: float = 50000.0
    ask: float = 50001.0


class RegimeExpander:
    """
    Generate 100+ distinct market regimes from dimension combinations.

    Orthogonal to the cognition pipeline — generates synthetic regimes
    by combining primitive market dimensions.
    """

    def __init__(self) -> None:
        self._liquidity = LIQUIDITY_DIMS
        self._volatility = VOLATILITY_DIMS
        self._flow = FLOW_DIMS
        self._structure = STRUCTURE_DIMS
        self._generated: Set[str] = set()

    def generate_regimes(
        self,
        max_regimes: int = 200,
        seed: int = 42,
    ) -> List[ExpandedRegime]:
        """
        Generate distinct, non-overlapping regimes from dimension combinations.

        Uses stratified sampling to ensure diversity.
        """
        import random
        rng = random.Random(seed)
        regimes: List[ExpandedRegime] = []

        # Generate all combinations (stratified)
        combos = list(itertools.product(
            self._liquidity, self._volatility, self._flow, self._structure,
        ))

        # Shuffle and take max_regimes
        rng.shuffle(combos)
        combos = combos[:max_regimes]

        for i, (liq, vol, flow, struct) in enumerate(combos):
            # Calculate bid/ask from dimensions
            spread = vol.spread_bps
            mid = 50000.0
            bid = mid - spread / 2
            ask = mid + spread / 2

            regime = ExpandedRegime(
                regime_id=f"exp_{i:03d}",
                label=f"{liq.label}_{vol.label}_{flow.label}_{struct.label}",
                liquidity=liq,
                volatility=vol,
                flow=flow,
                structure=struct,
                bid=bid,
                ask=ask,
            )
            regimes.append(regime)
            self._generated.add(regime.regime_id)

        return regimes

    def regime_to_scenario(
        self,
        regime: ExpandedRegime,
        symbol: str = "BTCUSDT",
        steps: int = 20,
        seed: int = 42,
    ) -> "Scenario":
        """Convert an ExpandedRegime to a Scenario for evaluation."""
        from malkhut.training.cma_trainer import Scenario

        venue = VenueRules(
            exchange="bingx", symbol=symbol, tick_size=0.1, lot_size=0.001,
            min_qty=0.001, min_notional=5.0, maker_fee_bps=-0.2, taker_fee_bps=0.5,
            post_only_supported=True, reduce_only_supported=True,
            max_orders_per_second=100, max_cancels_per_minute=120,
        )
        book = OrderBookState(
            ts_ns=1_000_000_000, symbol=symbol,
            bids=(PriceLevel(regime.bid, regime.liquidity.bid_qty),),
            asks=(PriceLevel(regime.ask, regime.liquidity.ask_qty),),
        )
        account = AccountState(
            ts_ns=1_000_000_000, equity=10000.0, wallet_balance=10000.0,
            available_balance=10000.0, margin_used=0.0, total_notional=0.0,
        )
        state = MarketWorldState(
            ts_ns=1_000_000_000, mode=Mode.ENDOGENOUS_AGENT_SIM,
            venue=venue, book=book, account=account,
        )

        return Scenario(
            scenario_id=regime.regime_id,
            symbol=symbol,
            initial_state=state,
            counterparties=regime.structure.counterparties,
            max_steps=steps,
            tags=(regime.label,),
        )

    @property
    def generated_count(self) -> int:
        return len(self._generated)