"""Model C period loop (temporary trading equilibrium per period; bounded-rationality simulated outcome).

Timing per period t (docs/model_specification.md §3):
 1. dividend d_t realised:  d_t = dbar + rho (d_{t-1} - dbar) + sigma_d u_t ;  dividends and interest paid
 2. mark-to-market at the last price: margin calls, forced sales, defaults, bail-in
 3. signals about u_{t+1}: private (theta + e_i) or model-m (theta + eta_m); Bayesian expected dividend
 4. agents propose actions; deterministic clipping (cash, holdings, margin, bank capital)
 5. trading under the chosen design; settlement through the ledger (infeasible matches are dropped and counted)
 6. metrics: price, return, order-flow correlation, spread/depth (CDA), price dispersion (bilateral)
No look-ahead: an agent's observation at t contains only t-dated information and its own signal about t+1.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from age.agents.policies import BayesianMV, LLMAgent, Momentum, Noise, clip_action
from age.banking.ledger import Ledger
from age.markets.designs import bilateral_bargaining, continuous_double_auction, walrasian_call


@dataclass
class MarketConfig:
    n_agents: int = 20
    T: int = 60
    burn_in: int = 10
    design: str = 'walrasian'          # walrasian | cda | bilateral
    dbar: float = 0.05
    rho_d: float = 0.5
    sigma_d: float = 0.01
    rf: float = 0.01
    sigma2: float = 0.01               # agents' perceived variance of next-period payoff p'+d' (goods^2)
    S: float | None = None             # asset supply (default: one share per agent)
    cash0: float = 1.0
    bank_equity0: float = 5.0
    margin: float = 0.5                # Reg-T style haircut
    fee_rate: float = 0.0005
    bank_capital_ratio: float = 0.0
    liquidation_impact: float = 0.5
    # information
    delegation_share: float = 0.0      # fraction of agents acting on a model signal
    M: int = 1                         # number of distinct foundation models among delegators
    sig_H: float = 0.02                # private noise s.d. (dividend units)
    sig_A: float = 0.005               # model noise s.d.
    # agents
    gamma: float = 2.0
    gamma_dispersion: float = 0.5      # agents' gamma_i = gamma * (1 + dispersion * z_i), z_i in [-1,1] (gains from trade)
    policy: str = 'bayesian'           # bayesian | momentum_mix | llm | llm_mix | surrogate
    max_leverage: float = 0.0          # Bayesian agents' borrowing as a multiple of equity (0 = unlevered)
    llm_clients: dict = field(default_factory=dict)   # name -> LLMClient (shared across agents)
    llm_assignment: list = field(default_factory=list) # per-agent client name (for llm policies)
    llm_variants: list = field(default_factory=list)   # per-agent prompt variant id
    llm_temperature: float = 0.0
    llm_memory: int = 0
    # market maker (liquidity provider): ladder of limit orders on both sides of its reference price
    mm_on: bool = True
    mm_depth: float = 0.5              # shares per level
    mm_tick: float = 0.0005            # relative price step between levels
    mm_levels: int = 10
    mm_cash: float = 50.0
    mm_inventory_max: float = 10.0     # beyond this the ladder on the exhausted side is withdrawn (depth under stress)
    mm_inv_sensitivity: float = 0.01   # reference price moves by this fraction per unit of inventory deviation
    shock_period: int | None = None    # period of an out-of-distribution common forecast shock
    shock_size: float = 3.0            # in units of sig_A (common model error)
    shock_persist: int = 1             # number of consecutive periods the common model error persists
    preload_leverage: float = 0.0      # fraction of margin capacity borrowed at t=0 (stress tests: agents start levered)
    auto_repay: bool = True            # repay loans with idle cash before acting (False in stress tests with preloaded leverage)
    seed: int = 1


def run_market(cfg: MarketConfig) -> dict:
    rng = np.random.default_rng(cfg.seed)
    n = cfg.n_agents
    if cfg.S is None:
        cfg.S = float(n)
    holdings0 = np.full(n, cfg.S / n)
    # steady-state price with a risk premium consistent with supply: p = (dbar + p - gamma_a sigma2 hbar)/(1+rf),
    # gamma_a = gamma/(cash0 + p hbar)  ->  rf hbar p^2 + (rf cash0 - dbar hbar) p - (dbar cash0 - gamma sigma2 hbar) = 0
    hbar = cfg.S / n
    a_, b_, c_ = cfg.rf * hbar, cfg.rf * cfg.cash0 - cfg.dbar * hbar, -(cfg.dbar * cfg.cash0 - cfg.gamma * cfg.sigma2 * hbar)
    disc = b_ ** 2 - 4 * a_ * c_
    assert disc > 0 and cfg.dbar * cfg.cash0 > cfg.gamma * cfg.sigma2 * hbar, 'risk premium too large: lower sigma2/gamma'
    p0 = (-b_ + np.sqrt(disc)) / (2 * a_)
    n_led = n + 1 if cfg.mm_on else n
    cash_init = np.full(n_led, cfg.cash0)
    hold_init = np.full(n_led, cfg.S / n)
    if cfg.mm_on:
        cash_init[n] = cfg.mm_cash
        hold_init[n] = cfg.mm_inventory_max / 2
    led = Ledger(n_led, cash_init, hold_init, cfg.bank_equity0, cfg.margin, cfg.fee_rate, p0, cfg.rf,
                 cfg.bank_capital_ratio)
    mm_target = cfg.mm_inventory_max / 2
    mm_ref = p0
    if cfg.preload_leverage > 0:
        for i in range(n):
            led.borrow(i, cfg.preload_leverage * (1 - cfg.margin) * p0 * hold_init[i])
    # information structure: delegators first
    nD = int(round(cfg.delegation_share * n))
    model_of = np.full(n, -1)
    model_of[:nD] = np.arange(nD) % cfg.M
    tau_t = 1.0 / cfg.sigma_d ** 2
    tau_H, tau_A = cfg.sig_H ** -2, cfg.sig_A ** -2
    bH, bA = tau_H / (tau_t + tau_H), tau_A / (tau_t + tau_A)
    # agents
    sigma2 = cfg.sigma2
    W0 = cfg.cash0 + holdings0[0] * p0
    policies = []
    gammas = cfg.gamma * (1 + cfg.gamma_dispersion * np.linspace(-1, 1, n))
    for i in range(n):
        if cfg.policy == 'bayesian':
            policies.append(BayesianMV(gammas[i] / W0, sigma2, p0, aggressiveness=0.8, max_leverage=cfg.max_leverage))
        elif cfg.policy == 'momentum_mix':
            policies.append(BayesianMV(gammas[i] / W0, sigma2, p0, aggressiveness=0.8, max_leverage=cfg.max_leverage) if i % 2 == 0
                            else (Momentum() if i % 4 == 1 else Noise(rng)))
        elif cfg.policy in ('llm', 'llm_mix', 'surrogate'):
            name = cfg.llm_assignment[i] if cfg.llm_assignment else list(cfg.llm_clients)[0]
            var = cfg.llm_variants[i] if cfg.llm_variants else 0
            policies.append(LLMAgent(cfg.llm_clients[name], 'investor', var, cfg.llm_temperature, cfg.seed * 1000 + i,
                                     cfg.llm_memory))
        else:
            raise ValueError(cfg.policy)
    # shocks (common random numbers: drawn up-front from the seed)
    u = rng.standard_normal(cfg.T + 1)
    e = rng.standard_normal((cfg.T + 1, n)) * cfg.sig_H
    eta = rng.standard_normal((cfg.T + 1, max(cfg.M, 1))) * cfg.sig_A
    if cfg.shock_period is not None:
        for k in range(cfg.shock_persist):                                 # signal received from period shock_period on
            if cfg.shock_period + 1 + k <= cfg.T:
                eta[cfg.shock_period + 1 + k, :] = cfg.shock_size * cfg.sig_A
    d = cfg.dbar
    prices, dividends, flows, returns = [p0], [], [], []
    rec = dict(spread=[], depth=[], dispersion=[], defaults=[], forced=[], losses=[], leverage=[], n_trades=[],
               invalid=[], clipped=[], excess=[], expected_div=[], signals=[])
    order_flow = np.zeros((cfg.T, n))
    for t in range(cfg.T):
        d = cfg.dbar + cfg.rho_d * (d - cfg.dbar) + cfg.sigma_d * u[t]
        led.pay_dividend(d)
        led.accrue_interest()
        mm = led.mark_to_market(led.price, cfg.liquidation_impact)
        theta = cfg.sigma_d * u[t + 1]                        # next innovation (unknown to agents)
        # signals and Bayesian expected dividend for next period
        exp_div = np.empty(n)
        for i in range(n):
            if model_of[i] >= 0:
                s = theta + eta[t + 1, model_of[i]]
                exp_div[i] = cfg.dbar + cfg.rho_d * (d - cfg.dbar) + bA * s
            else:
                s = theta + e[t + 1, i]
                exp_div[i] = cfg.dbar + cfg.rho_d * (d - cfg.dbar) + bH * s
        rec['expected_div'].append(exp_div.copy())
        orders = []
        n_invalid = n_clipped = 0
        for i in rng.permutation(n):
            if not led.alive[i]:
                continue
            if cfg.auto_repay and led.loan[i] > 0 and led.cash[i] > 0:   # repay loans with idle cash before acting
                led.repay(i, min(led.loan[i], led.cash[i]))
            eq = led.equity()[i]
            obs = dict(period=t, price=led.price, dividend=d, rf=cfg.rf, expected_dividend=float(exp_div[i]),
                       cash=float(led.cash[i]), holdings=float(led.holdings[i]), loan=float(led.loan[i]),
                       equity=float(eq), margin=cfg.margin, loan_capacity=float(led.loan_capacity(i)),
                       history=prices[-6:], rho_d=cfg.rho_d, dbar=cfg.dbar,
                       own_valuation=float((exp_div[i] + p0 + cfg.rho_d * (exp_div[i] - cfg.dbar) / (1 + cfg.rf - cfg.rho_d)) / (1 + cfg.rf)))
            pol = policies[i]
            if isinstance(pol, LLMAgent):
                # the surrogate/LLM also receives a target so that the surrogate can be a faithful control
                obs['target_holdings'] = float(max((obs['own_valuation'] - led.price) * (1 + cfg.rf) / (cfg.gamma / W0 * sigma2), 0))
                act = pol.act(obs, tags=dict(period=t, agent=int(i), run=cfg.seed))
                n_invalid += int(act.get('invalid', False))
            else:
                act = pol.act(obs)
            act = clip_action(act, obs, led, int(i))
            n_clipped += int(act['clipped'])
            if act['side'] in ('buy', 'sell') and act['quantity'] > 1e-9:
                orders.append(dict(agent=int(i), side=act['side'], qty=act['quantity'],
                                   limit=act['limit_price'] if act['order_type'] == 'limit' else None, type=act['order_type']))
                order_flow[t, i] = act['quantity'] * (1 if act['side'] == 'buy' else -1)
        if cfg.mm_on:
            inv = led.holdings[n]
            mm_ref = mm_ref * (1 - cfg.mm_inv_sensitivity * (inv - mm_target) / max(mm_target, 1e-9)) \
                if False else mm_ref                      # reference anchored; inventory skew applied to quotes below
            skew = cfg.mm_inv_sensitivity * (inv - mm_target) / max(mm_target, 1e-9)
            base = led.price * (1 - skew)
            for k in range(1, cfg.mm_levels + 1):
                ask = base * (1 + k * cfg.mm_tick)
                bid = base * (1 - k * cfg.mm_tick)
                if inv - (k - 1) * cfg.mm_depth > 0:                   # can sell only what it holds
                    orders.append(dict(agent=n, side='sell', qty=min(cfg.mm_depth, max(inv - (k - 1) * cfg.mm_depth, 0)), limit=ask, type='limit'))
                if inv + k * cfg.mm_depth <= cfg.mm_inventory_max and led.cash[n] > bid * cfg.mm_depth * k:
                    orders.append(dict(agent=n, side='buy', qty=cfg.mm_depth, limit=bid, type='limit'))
        if cfg.design == 'walrasian':
            trades, p_new, info = walrasian_call(orders, led.price)
        elif cfg.design == 'cda':
            trades, p_new, info = continuous_double_auction(orders, led.price, rng)
        else:
            trades, p_new, info = bilateral_bargaining(orders, led.price, rng)
        n_ok = 0
        for tr in trades:
            n_ok += int(led.settle_trade(tr['buyer'], tr['seller'], tr['qty'], tr['price']))
        if trades:
            led.price = float(np.average([tr['price'] for tr in trades], weights=[tr['qty'] for tr in trades]))
        prices.append(led.price)
        dividends.append(d)
        returns.append(np.log(prices[-1] / prices[-2]))
        rec['spread'].append(info.get('spread', np.nan)); rec['depth'].append(info.get('depth', np.nan))
        rec['dispersion'].append(info.get('dispersion', np.nan)); rec['excess'].append(info.get('excess', np.nan))
        rec['defaults'].append(mm['defaults']); rec['forced'].append(mm['forced']); rec['losses'].append(mm['losses'])
        rec['leverage'].append(led.snapshot()['leverage']); rec['n_trades'].append(n_ok)
        rec['invalid'].append(n_invalid); rec['clipped'].append(n_clipped)
    # ------------------------------------------------------------------ metrics (post burn-in)
    b = cfg.burn_in
    r = np.array(returns[b:])
    of = order_flow[b:]
    # average pairwise correlation of order flow across agents, *conditional on public information*:
    # each agent's flow is residualised on [1, d_t, p_t] (dividend and price are public), then correlated.
    pub = np.column_stack([np.ones(cfg.T - b), np.array(dividends[b:]), np.array(prices[b:-1])])
    beta_pub, *_ = np.linalg.lstsq(pub, of, rcond=None)
    res_of = of - pub @ beta_pub
    with np.errstate(invalid='ignore', divide='ignore'):
        C = np.corrcoef(res_of.T)
    C = np.nan_to_num(C)
    iu = np.triu_indices(n, 1)
    of_corr = float(C[iu].mean()) if n > 1 else np.nan
    # same correlation within delegators only and within humans only
    def _grp(mask):
        idx = np.where(mask)[0]
        if len(idx) < 2:
            return np.nan
        sub = C[np.ix_(idx, idx)]
        return float(sub[np.triu_indices(len(idx), 1)].mean())
    of_corr_deleg = _grp(model_of >= 0)
    of_corr_human = _grp(model_of < 0)
    # Kyle lambda: regress returns on net order flow
    net = of.sum(axis=1)
    lam = float(np.polyfit(net, r, 1)[0]) if np.std(net) > 1e-12 and len(r) > 3 else np.nan
    # price pressure: regression of returns on the average model error of delegators (common component)
    if nD > 0:
        common = np.array([np.mean([eta[t + 1, model_of[i]] for i in range(nD)]) for t in range(cfg.T)])[b:]
        pressure = float(np.polyfit(common, r, 1)[0]) if np.std(common) > 1e-12 else np.nan
        with np.errstate(invalid='ignore'):
            r2_common = float(np.corrcoef(common, r)[0, 1] ** 2) if np.std(common) > 1e-12 else np.nan
    else:
        pressure, r2_common = np.nan, 0.0
    out = dict(pressure=pressure, r2_common=r2_common, of_corr_deleg=of_corr_deleg, of_corr_human=of_corr_human, mm_inventory=float(led.holdings[n]) if cfg.mm_on else np.nan,prices=np.array(prices), dividends=np.array(dividends), returns=np.array(returns), order_flow=order_flow,
               vol=float(np.std(r)), of_corr=of_corr, kyle_lambda=lam,
               spread=float(np.nanmean(rec['spread'][b:])) if cfg.design == 'cda' else np.nan,
               depth=float(np.nanmean(rec['depth'][b:])) if cfg.design == 'cda' else np.nan,
               dispersion=float(np.nanmean(rec['dispersion'][b:])) if cfg.design == 'bilateral' else np.nan,
               defaults=int(np.sum(rec['defaults'])), forced=float(np.sum(rec['forced'])), losses=float(np.sum(rec['losses'])),
               leverage=float(np.mean(rec['leverage'][b:])), n_trades=float(np.mean(rec['n_trades'][b:])),
               invalid=int(np.sum(rec['invalid'])), clipped=int(np.sum(rec['clipped'])),
               final=led.snapshot(), eta=eta, u=u, model_of=model_of, cfg=cfg)
    if cfg.shock_period is not None and cfg.shock_period < cfg.T:
        sp = cfg.shock_period
        w = cfg.shock_persist + 2
        out['shock_return'] = float(np.sum(returns[sp:sp + w]))
        out['shock_forced'] = float(np.sum(rec['forced'][sp:sp + w]))
        out['shock_defaults'] = int(np.sum(rec['defaults'][sp:sp + w]))
        out['shock_losses'] = float(np.sum(rec['losses'][sp:sp + w]))
        out['shock_min_return'] = float(np.min(returns[sp:sp + w]))
        out['shock_trough'] = float(np.min(prices[sp:sp + w + 1]) / prices[sp] - 1)
    return out