"""LLM agent laboratory (Model C with live local models).

Cells (pre-registered): policy ∈ {bayesian control, surrogate (rule-based mock with the LLM interface), single shared
LLM (qwen7b), diverse LLMs (qwen7b / mistral7b / llama3b assigned round-robin), shared LLM with varied prompts},
signal structure ∈ {private (D=0), shared model (D=1, M=1), diverse models (D=1, M=3)}, temperature ∈ {0, 0.7},
seeds 1..S. Design: CDA. All calls cached and audited (results/llm_audit/*.jsonl); invalid actions counted.
Run: python3 scripts/35_llm_lab.py [--quick] [--providers qwen7b,mistral7b,llama3b]
Outputs: results/llm_lab/cells.csv, results/llm_lab/summary.csv, results/llm_lab/audit_stats.json
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'src'))
from age.llm.client import LLMClient, LOCAL_PROVIDERS, MockLLM  # noqa: E402
from age.markets.simulation import MarketConfig, run_market  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'results' / 'llm_lab'
OUT.mkdir(parents=True, exist_ok=True)
QUICK = '--quick' in sys.argv
N, T, BURN = (8, 24, 4) if QUICK else (10, 40, 8)
SEEDS = range(1, 4) if QUICK else range(1, 7)
provs = ['qwen7b', 'mistral7b', 'llama3b']
POLICIES = ['bayesian', 'surrogate', 'llm_shared', 'llm_diverse', 'llm_shared_varied_prompts']
TAG = ''
for a in sys.argv:
    if a.startswith('--providers'):
        provs = a.split('=')[1].split(',')
    if a.startswith('--policies='):
        POLICIES = a.split('=')[1].split(',')
    if a.startswith('--tag='):
        TAG = '_' + a.split('=')[1]

clients = {k: LLMClient(LOCAL_PROVIDERS[k], run_id=f'lab_{k}') for k in provs}
clients['mock'] = MockLLM(run_id='lab_mock')
KEYS = ['vol', 'of_corr', 'of_corr_deleg', 'pressure', 'r2_common', 'n_trades', 'invalid', 'clipped', 'defaults', 'leverage']


def run_cell(policy, D, M, temp, seed, memory=0):
    kw = dict(n_agents=N, T=T, burn_in=BURN, design='cda', delegation_share=D, M=M, seed=seed, llm_temperature=temp,
              llm_memory=memory)
    if policy == 'bayesian':
        cfg = MarketConfig(policy='bayesian', **kw)
    elif policy == 'surrogate':
        cfg = MarketConfig(policy='surrogate', llm_clients={'mock': clients['mock']}, llm_assignment=['mock'] * N, **kw)
    elif policy == 'llm_shared':
        cfg = MarketConfig(policy='llm', llm_clients={provs[0]: clients[provs[0]]}, llm_assignment=[provs[0]] * N, **kw)
    elif policy == 'llm_diverse':
        assign = [provs[i % len(provs)] for i in range(N)]
        cfg = MarketConfig(policy='llm_mix', llm_clients={k: clients[k] for k in provs}, llm_assignment=assign, **kw)
    elif policy == 'llm_shared_varied_prompts':
        cfg = MarketConfig(policy='llm', llm_clients={provs[0]: clients[provs[0]]}, llm_assignment=[provs[0]] * N,
                           llm_variants=[i % 5 for i in range(N)], **kw)
    elif policy == 'llm_shared_memory':
        cfg = MarketConfig(policy='llm', llm_clients={provs[0]: clients[provs[0]]}, llm_assignment=[provs[0]] * N, **kw)
        cfg.llm_memory = 3
    else:
        raise ValueError(policy)
    t0 = time.time()
    out = run_market(cfg)
    row = dict(policy=policy, D=D, M=M, temperature=temp, seed=seed, runtime_s=time.time() - t0)
    row.update({k: out.get(k, np.nan) for k in KEYS})
    # action homogeneity: average pairwise correlation of *raw* order flow among agents (not residualised)
    of = out['order_flow'][BURN:]
    with np.errstate(invalid='ignore', divide='ignore'):
        C = np.nan_to_num(np.corrcoef(of.T))
    row['raw_of_corr'] = float(C[np.triu_indices(N, 1)].mean())
    row['price_end'] = float(out['prices'][-1])
    return row


rows = []
cells = []
for policy in POLICIES:
    for (D, M) in [(0.0, 1), (1.0, 1), (1.0, 3)]:
        temps = [0.0] if policy in ('bayesian', 'surrogate') else [0.0, 0.7]
        for temp in temps:
            for s in SEEDS:
                cells.append((policy, D, M, temp, s))
if not QUICK:
    for s in SEEDS:
        cells.append(('llm_shared_memory', 1.0, 1, 0.0, s))
print(f'{len(cells)} cells, N={N}, T={T}')
for i, c in enumerate(cells):
    r = run_cell(*c)
    rows.append(r)
    print(f"[{i+1}/{len(cells)}] {c} vol={r['vol']:.4f} ofc={r['of_corr']:.2f} raw={r['raw_of_corr']:.2f} trades={r['n_trades']:.1f} invalid={r['invalid']} ({r['runtime_s']:.0f}s)", flush=True)
    pd.DataFrame(rows).to_csv(OUT / f'cells{TAG}.csv', index=False)
df = pd.DataFrame(rows)
summ = df.groupby(['policy', 'D', 'M', 'temperature'])[KEYS + ['raw_of_corr', 'runtime_s']].agg(['mean', 'std']).reset_index()
summ.columns = ['_'.join(c).strip('_') for c in summ.columns]
summ.to_csv(OUT / f'summary{TAG}.csv', index=False)
stats = {k: v.stats() for k, v in clients.items()}
json.dump(stats, open(OUT / f'audit_stats{TAG}.json', 'w'), indent=1)
print(summ.to_string())
print(stats)