"""Three market designs for Model C. Each takes a list of orders and returns executed trades.

Order: dict(agent=i, side='buy'|'sell', qty=float, limit=float|None, type='limit'|'market').
Trade: dict(buyer, seller, qty, price).

* `walrasian_call`: a single clearing price equates buy and sell quantities (limit orders define step demand/supply
  schedules; market orders are price-insensitive); rationing pro rata on the long side.
* `continuous_double_auction`: limit order book with price–time priority; orders arrive in a random sequence
  (seeded), marketable orders execute against resting ones at the resting price; records mid, spread, depth.
* `bilateral_bargaining`: random matching of a buyer and a seller for R rounds; trade if the buyer's limit >= the
  seller's limit at the Nash-bargaining price (weight `power` on the buyer), quantity = min of the two.
Feasibility (cash, holdings) is enforced by the ledger when trades are settled; designs only match.
"""
from __future__ import annotations

import numpy as np


def _split(orders):
    buys = [o for o in orders if o['side'] == 'buy' and o['qty'] > 0]
    sells = [o for o in orders if o['side'] == 'sell' and o['qty'] > 0]
    return buys, sells


def walrasian_call(orders, p_ref: float, p_lo: float = 1e-6, p_hi: float = None):
    buys, sells = _split(orders)
    if not buys or not sells:
        return [], p_ref, dict(excess=0.0)
    p_hi = p_hi or 100 * p_ref

    def demand(p):
        return sum(o['qty'] for o in buys if o['type'] == 'market' or o['limit'] >= p)

    def supply(p):
        return sum(o['qty'] for o in sells if o['type'] == 'market' or o['limit'] <= p)

    # price grid from limits: excess demand is a step function, find the price minimising |excess|
    cands = sorted(set([o['limit'] for o in buys + sells if o['limit'] is not None] + [p_ref]))
    best, best_ex = p_ref, float('inf')
    for p in cands:
        ex = demand(p) - supply(p)
        if abs(ex) < best_ex or (abs(ex) == best_ex and abs(p - p_ref) < abs(best - p_ref)):
            best, best_ex = p, abs(ex)
    p = best
    D, S = demand(p), supply(p)
    if D <= 0 or S <= 0:
        return [], p_ref, dict(excess=D - S)
    q = min(D, S)
    rb, rs = q / D, q / S
    bq = [(o['agent'], o['qty'] * rb) for o in buys if o['type'] == 'market' or o['limit'] >= p]
    sq = [(o['agent'], o['qty'] * rs) for o in sells if o['type'] == 'market' or o['limit'] <= p]
    trades = []
    i = j = 0
    bq = [list(x) for x in bq]
    sq = [list(x) for x in sq]
    while i < len(bq) and j < len(sq):
        x = min(bq[i][1], sq[j][1])
        if x > 1e-12:
            trades.append(dict(buyer=bq[i][0], seller=sq[j][0], qty=x, price=p))
        bq[i][1] -= x
        sq[j][1] -= x
        if bq[i][1] <= 1e-12:
            i += 1
        if sq[j][1] <= 1e-12:
            j += 1
    return trades, p, dict(excess=D - S, demand=D, supply=S)


def continuous_double_auction(orders, p_ref: float, rng: np.random.Generator):
    book_b, book_s = [], []       # resting orders: [price, time, agent, qty]
    trades = []
    seq = rng.permutation(len(orders))
    spreads, mids = [], []
    t = 0
    for k in seq:
        o = dict(orders[k])
        t += 1
        qty = o['qty']
        if qty <= 0:
            continue
        if o['side'] == 'buy':
            lim = np.inf if o['type'] == 'market' else o['limit']
            book_s.sort(key=lambda r: (r[0], r[1]))
            while qty > 1e-12 and book_s and book_s[0][0] <= lim:
                r = book_s[0]
                x = min(qty, r[3])
                trades.append(dict(buyer=o['agent'], seller=r[2], qty=x, price=r[0]))
                qty -= x
                r[3] -= x
                if r[3] <= 1e-12:
                    book_s.pop(0)
            if qty > 1e-12 and o['type'] == 'limit':
                book_b.append([o['limit'], t, o['agent'], qty])
        else:
            lim = -np.inf if o['type'] == 'market' else o['limit']
            book_b.sort(key=lambda r: (-r[0], r[1]))
            while qty > 1e-12 and book_b and book_b[0][0] >= lim:
                r = book_b[0]
                x = min(qty, r[3])
                trades.append(dict(buyer=r[2], seller=o['agent'], qty=x, price=r[0]))
                qty -= x
                r[3] -= x
                if r[3] <= 1e-12:
                    book_b.pop(0)
            if qty > 1e-12 and o['type'] == 'limit':
                book_s.append([o['limit'], t, o['agent'], qty])
        if book_b and book_s:
            bb = max(r[0] for r in book_b)
            ba = min(r[0] for r in book_s)
            if ba >= bb:
                spreads.append(ba - bb)
                mids.append(0.5 * (ba + bb))
    last = trades[-1]['price'] if trades else p_ref
    depth_b = sum(r[3] for r in book_b)
    depth_s = sum(r[3] for r in book_s)
    info = dict(spread=float(np.mean(spreads)) if spreads else np.nan, mid=float(np.mean(mids)) if mids else np.nan,
                depth=depth_b + depth_s, n_trades=len(trades))
    return trades, last, info


def bilateral_bargaining(orders, p_ref: float, rng: np.random.Generator, rounds: int = 3, power: float = 0.5):
    buys, sells = _split(orders)
    buys = [dict(o) for o in buys]
    sells = [dict(o) for o in sells]
    trades = []
    for _ in range(rounds):
        bl = [o for o in buys if o['qty'] > 1e-12]
        sl = [o for o in sells if o['qty'] > 1e-12]
        if not bl or not sl:
            break
        rng.shuffle(bl)
        rng.shuffle(sl)
        for b, s in zip(bl, sl):
            bl_lim = b['limit'] if b['type'] == 'limit' else p_ref * 1.5
            sl_lim = s['limit'] if s['type'] == 'limit' else p_ref * 0.5
            if bl_lim >= sl_lim and b['agent'] != s['agent']:
                price = power * sl_lim + (1 - power) * bl_lim   # buyer with power `power` pulls price toward seller's ask
                x = min(b['qty'], s['qty'])
                trades.append(dict(buyer=b['agent'], seller=s['agent'], qty=x, price=price))
                b['qty'] -= x
                s['qty'] -= x
    last = float(np.mean([t['price'] for t in trades])) if trades else p_ref
    return trades, last, dict(n_trades=len(trades), dispersion=float(np.std([t['price'] for t in trades])) if trades else np.nan)