"""Agent policies for Model C. A policy maps an observation to a *proposed* action; the simulation validates and clips
it against the ledger (cash, holdings, margin). Policies never touch the ledger.

Observation (dict): price, dividend, rf, signal (forecast of next dividend innovation, already Bayes-weighted into
`expected_dividend`), own cash/holdings/loan/equity, margin, history (last prices), target (for Bayesian agents),
own_valuation, period.

Action (dict): side ('buy'|'sell'|'hold'), quantity (>=0), order_type ('limit'|'market'), limit_price, borrow (>=0),
rationale (string; ignored by accounting).
"""
from __future__ import annotations

import json

import numpy as np

ACTION_SCHEMA = {
    "type": "object",
    "properties": {
        "side": {"type": "string", "enum": ["buy", "sell", "hold"]},
        "quantity": {"type": "number", "minimum": 0},
        "order_type": {"type": "string", "enum": ["limit", "market"]},
        "limit_price": {"type": "number", "minimum": 0},
        "borrow": {"type": "number", "minimum": 0},
        "rationale": {"type": "string"},
    },
    "required": ["side", "quantity", "order_type"],
    "additionalProperties": True,
}


# --------------------------------------------------------------------------- Bayesian mean-variance (control)
class BayesianMV:
    """h* = (E[p' + d'] - (1 + rf) p) / (gamma_a sigma^2); limit order halfway between price and valuation."""

    def __init__(self, gamma_a: float, sigma2: float, p_anchor: float, aggressiveness: float = 0.5,
                 max_leverage: float = 0.0):
        self.gamma_a, self.sigma2, self.p_anchor = gamma_a, sigma2, p_anchor
        self.aggr = aggressiveness
        self.max_leverage = max_leverage

    def valuation(self, obs):
        # next payoff = expected dividend + resale value; the resale value prices the persistence of the expected
        # dividend deviation: p_anchor + rho_d (E[d'] - dbar)/(1 + rf - rho_d)
        rho, dbar = obs.get('rho_d', 0.0), obs.get('dbar', obs['expected_dividend'])
        resale = self.p_anchor + rho * (obs['expected_dividend'] - dbar) / (1 + obs['rf'] - rho)
        return (obs['expected_dividend'] + resale) / (1 + obs['rf'])

    def act(self, obs):
        p = obs['price']
        v = self.valuation(obs)
        h_star = (v * (1 + obs['rf']) - (1 + obs['rf']) * p) / (self.gamma_a * self.sigma2)
        h_star = max(h_star, 0.0)
        q = h_star - obs['holdings']
        side = 'buy' if q > 1e-9 else ('sell' if q < -1e-9 else 'hold')
        limit = p + self.aggr * (v - p) if side == 'buy' else p - self.aggr * (p - v)
        limit = max(limit, 1e-6)
        borrow = 0.0
        if side == 'buy' and self.max_leverage > 0:
            need = abs(q) * limit - obs['cash']
            borrow = max(min(need, self.max_leverage * obs['equity']), 0.0)
        return dict(side=side, quantity=abs(q), order_type='limit', limit_price=limit, borrow=borrow,
                    rationale='bayesian mean-variance')


# --------------------------------------------------------------------------- rule-based
class Momentum:
    def __init__(self, lookback=3, size=0.1):
        self.lb, self.size = lookback, size

    def act(self, obs):
        hist = obs['history']
        if len(hist) <= self.lb:
            return dict(side='hold', quantity=0.0, order_type='limit', limit_price=obs['price'], borrow=0.0, rationale='warm-up')
        ret = hist[-1] / hist[-1 - self.lb] - 1
        side = 'buy' if ret > 0 else 'sell'
        q = self.size * (obs['holdings'] if side == 'sell' else max(obs['cash'] / obs['price'], 0))
        return dict(side=side, quantity=q, order_type='market', limit_price=obs['price'], borrow=0.0, rationale='momentum')


class Noise:
    def __init__(self, rng, size=0.05):
        self.rng, self.size = rng, size

    def act(self, obs):
        side = 'buy' if self.rng.random() < 0.5 else 'sell'
        q = self.size * (obs['holdings'] if side == 'sell' else max(obs['cash'] / obs['price'], 0))
        lim = obs['price'] * (1 + 0.02 * self.rng.standard_normal())
        return dict(side=side, quantity=q, order_type='limit', limit_price=max(lim, 1e-6), borrow=0.0, rationale='noise')


# --------------------------------------------------------------------------- LLM agent
SYSTEM_PROMPTS = {
    'investor': ("You are an AI portfolio agent acting on behalf of a human principal in a simulated market with one risky "
                 "asset (pays a stochastic dividend each period) and a safe deposit paying rate rf. Your mandate: maximise the "
                 "principal's expected wealth next period while keeping risk moderate; you may borrow against holdings only "
                 "within the margin limit. You must reply with ONLY a JSON object with keys side (buy|sell|hold), quantity "
                 "(number of shares, >= 0), order_type (limit|market), limit_price (number), borrow (number >= 0, cash to "
                 "borrow), rationale (short string). Never exceed your cash or holdings; the exchange will reject infeasible orders."),
    'market_maker': ("You are a market-maker agent. Quote a limit order on the side where your inventory deviates most from "
                     "its target, with a price inside the recent range. Reply with ONLY the JSON object described."),
    'lender': ("You are a bank credit officer. Given the borrower's balance sheet, decide the fraction of the requested loan "
               "to grant (0..1). Reply ONLY with JSON {\"grant_fraction\": number, \"rationale\": string}."),
}

PROMPT_VARIANTS = [
    "", " Be concise.", " Think like a careful long-term investor.", " Trade actively when the signal is strong.",
    " Prefer limit orders near the current price.",
]


class LLMAgent:
    def __init__(self, client, role: str = 'investor', variant: int = 0, temperature: float = 0.0, seed: int = 0,
                 memory: int = 0):
        self.client, self.role, self.variant, self.T, self.seed, self.memory = client, role, variant, temperature, seed, memory
        self.history = []

    def observation_text(self, obs):
        o = {k: (round(v, 4) if isinstance(v, float) else v) for k, v in obs.items() if k != 'target'}
        o['history'] = [round(x, 4) for x in obs['history'][-5:]]
        return json.dumps(o)

    def act(self, obs, tags=None):
        sys_prompt = SYSTEM_PROMPTS[self.role] + PROMPT_VARIANTS[self.variant % len(PROMPT_VARIANTS)]
        msgs = [{'role': 'system', 'content': sys_prompt}]
        if self.memory and self.history:
            msgs.append({'role': 'user', 'content': 'Your previous actions: ' + json.dumps(self.history[-self.memory:])})
        msgs.append({'role': 'user', 'content': 'Observation: ' + self.observation_text(obs) + '\nReply with the JSON action.'})
        res = self.client.complete_json(msgs, ACTION_SCHEMA, temperature=self.T, seed=self.seed, tags=tags)
        if res.parsed is None:
            act = dict(side='hold', quantity=0.0, order_type='limit', limit_price=obs['price'], borrow=0.0,
                       rationale='INVALID-LLM-OUTPUT -> no action', invalid=True)
        else:
            act = dict(res.parsed)
            act.setdefault('limit_price', obs['price'])
            act.setdefault('borrow', 0.0)
            act['invalid'] = False
        if self.memory:
            self.history.append({k: act[k] for k in ('side', 'quantity', 'limit_price')})
        return act


def clip_action(act: dict, obs: dict, ledger, i: int) -> dict:
    """Deterministic enforcement: borrow within capacity, quantities within cash/holdings, prices positive."""
    a = dict(act)
    a['side'] = a.get('side', 'hold')
    a['quantity'] = float(max(a.get('quantity', 0.0) or 0.0, 0.0))
    a['order_type'] = a.get('order_type', 'limit') if a.get('order_type') in ('limit', 'market') else 'limit'
    lp = a.get('limit_price', obs['price'])
    a['limit_price'] = float(lp) if (lp is not None and lp > 0) else obs['price']
    a['borrow'] = float(max(a.get('borrow', 0.0) or 0.0, 0.0))
    if a['borrow'] > 0:
        a['borrow'] = ledger.borrow(i, a['borrow'])            # actually granted (margin + bank capital)
    if a['side'] == 'buy':
        px = a['limit_price'] if a['order_type'] == 'limit' else obs['price'] * 1.05
        a['quantity'] = min(a['quantity'], ledger.max_affordable(i, px))
    elif a['side'] == 'sell':
        a['quantity'] = min(a['quantity'], ledger.holdings[i])
    else:
        a['quantity'] = 0.0
    a['clipped'] = (a['quantity'] < float(act.get('quantity', 0.0) or 0.0) - 1e-9)
    return a