"""
Ship or hold?  Steps 3 to 5.
============================
Data: tau-bench historical trajectories (Yao et al., 2024; github.com/sierra-research/tau-bench, MIT),
the same runs as Lesson 8. The question: should a support agent running on GPT-4o
(each task run 4 times) switch to Claude 3.5 Sonnet (each task run 8 times)?

Setup (once):  nothing to install beyond Python 3
Run:           python release_lesson.py                     (old vs new, paired by task, with a range)
               python release_lesson.py regressions retail  (tasks the new model does worse on)
               python release_lesson.py failures            (how often each kind of failure happens, per run)
               python release_lesson.py show retail sonnet 3 0
                                                            (read one run: domain, model, task, trial)
The first run downloads about 50 MB into ./data.
"""
import json
import os
import random
import sys
import urllib.request
from collections import Counter, defaultdict
from math import comb

BASE = 'https://raw.githubusercontent.com/sierra-research/tau-bench/main/historical_trajectories/'
RUNS = {('airline', 'gpt-4o'): 'gpt-4o-airline', ('retail', 'gpt-4o'): 'gpt-4o-retail',
        ('airline', 'sonnet'): 'sonnet-35-new-airline', ('retail', 'sonnet'): 'sonnet-35-new-retail'}
WRITES = ('book', 'cancel', 'modify', 'update', 'exchange', 'return', 'send')   # tools that change the database


def load(domain, model):
    os.makedirs('data', exist_ok=True)
    name = RUNS[(domain, model)]
    path = f'data/tau_{name}.json'
    if not os.path.exists(path):
        print('downloading', name, '...')
        urllib.request.urlretrieve(BASE + name + '.json', path)
    return json.load(open(path, encoding='utf-8'))


def by_task(runs):
    tasks = defaultdict(list)
    for r in runs:
        tasks[r['task_id']].append(r['reward'] >= 0.999)
    return tasks


def all_k(tasks, k):
    """Chance that k runs of a task, picked at random, ALL succeed (often written pass^k)."""
    return sum(comb(sum(v), k) / comb(len(v), k) for v in tasks.values()) / len(tasks)


def any_k(tasks, k):
    """Chance that at least one of k runs succeeds (often written pass@k)."""
    return sum(1 - comb(len(v) - sum(v), k) / comb(len(v), k) for v in tasks.values()) / len(tasks)




def rate(v):
    return sum(v) / len(v)


def compare():
    random.seed(0)
    print(f"  {'domain':<8} {'tasks':>5} {'old, per run':>13} {'new, per run':>13} {'new minus old':>14}"
          f" {'95% range':>16} {'resamples new ahead':>20} {'old, all 4':>11} {'new, all 4':>11}")
    for domain in ('retail', 'airline'):
        old, new = by_task(load(domain, 'gpt-4o')), by_task(load(domain, 'sonnet'))
        diffs = [rate(new[t]) - rate(old[t]) for t in old]           # paired: same task, two models
        mean = sum(diffs) / len(diffs)
        boots = sorted(sum(random.choice(diffs) for _ in diffs) / len(diffs) for _ in range(5000))
        lo, hi = boots[125], boots[4874]
        ahead = sum(b > 0 for b in boots) / len(boots)
        print(f"  {domain:<8} {len(diffs):>5} {all_k(old, 1):>13.0%} {all_k(new, 1):>13.0%} {mean:>+14.1%}"
              f"   {lo:>+6.1%} to {hi:>+6.1%} {ahead:>20.0%} {all_k(old, 4):>11.0%} {all_k(new, 4):>11.0%}")


def regressions(domain):
    old, new = by_task(load(domain, 'gpt-4o')), by_task(load(domain, 'sonnet'))
    tasks = {r['task_id']: r['info']['task'] for r in load(domain, 'gpt-4o')}
    always = [t for t in old if all(old[t])]
    slipped = [t for t in always if not all(new[t])]
    print(f"{domain}: old model right on every run for {len(always)} tasks; new model fails at least once on {len(slipped)} of them.\n")
    print("Tasks where the new model's success rate is 50+ points lower:")
    for t in sorted(old):
        if rate(new[t]) - rate(old[t]) <= -0.5:
            changes = [a['name'] for a in tasks[t]['actions'] if a['name'].startswith(WRITES)]
            print(f"  task {t:>3}  old {sum(old[t])}/{len(old[t])}  new {sum(new[t])}/{len(new[t])}   required: {', '.join(changes) or 'none'}")
    better = sum(rate(new[t]) - rate(old[t]) >= 0.5 for t in old)
    print(f"\nTasks where it is 50+ points higher: {better}")


def failures():
    kinds = ['made a change it should not have', 'right kind of change, wrong details', 'skipped a required change',
             'made no change', 'changes right, reply to the customer wrong']
    print("Share of ALL runs that ended in each kind of failure (not share of failures)\n")
    for domain in ('retail', 'airline'):
        print(f"  {domain:<44} {'old':>6} {'new':>6}")
        runs = {m: load(domain, m) for m in ('gpt-4o', 'sonnet')}
        counts = {m: Counter(failure_kind(r) for r in rs if r['reward'] < 0.999) for m, rs in runs.items()}
        for k in kinds:
            print(f"    {k:<42} {counts['gpt-4o'][k] / len(runs['gpt-4o']):>6.1%} {counts['sonnet'][k] / len(runs['sonnet']):>6.1%}")
        print()


def tool_calls(run):
    out = []
    for m in run['traj']:
        for call in (m.get('tool_calls') or []):
            f = call.get('function', {})
            raw = f.get('arguments') or {}
            try:
                args = raw if isinstance(raw, dict) else json.loads(raw)   # the two models' logs store these differently
            except ValueError:
                args = {}
            out.append((f.get('name'), args))
    return out


def failure_kind(run):
    """A rough rule: compare the database-changing tool calls the agent made with the ones the task required."""
    need = [(a['name'], a['kwargs']) for a in run['info']['task']['actions'] if a['name'].startswith(WRITES)]
    made = [(n, a) for n, a in tool_calls(run) if n and n.startswith(WRITES)]
    need_n, made_n = Counter(n for n, _ in need), Counter(n for n, _ in made)
    if not made and need:
        return 'made no change'
    if need_n - made_n:
        return 'skipped a required change'
    if made_n - need_n:
        return 'made a change it should not have'
    key = lambda pairs: sorted(json.dumps(a, sort_keys=True) for _, a in pairs)
    if key(need) != key(made):
        return 'right kind of change, wrong details'
    return 'changes right, reply to the customer wrong'



def show(domain, model, task_id, trial):
    run = [r for r in load(domain, model) if r['task_id'] == task_id and r['trial'] == trial][0]
    print("TASK (what the simulated customer was told):", run['info']['task']['instruction'])
    print("REQUIRED CHANGES:", [a['name'] for a in run['info']['task']['actions'] if a['name'].startswith(WRITES)] or 'none')
    print("RESULT:", "success" if run['reward'] >= 0.999 else "failure", "|", failure_kind(run) if run['reward'] < 0.999 else "")
    print()
    for m in run['traj'][1:]:
        if m['role'] == 'tool':
            print('     [tool result]', str(m['content'])[:300].replace('\n', ' '))
        elif m.get('tool_calls'):
            for n, a in tool_calls({'traj': [m]}):
                print(f'     [agent calls {n}]', json.dumps(a)[:200])
        else:
            print(f"{m['role'].upper()}: {str(m['content'])[:500]}")


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'regressions':
        regressions(sys.argv[2] if len(sys.argv) > 2 else 'retail')
    elif cmd == 'failures':
        failures()
    elif cmd == 'show' and len(sys.argv) == 6:
        show(sys.argv[2], sys.argv[3], int(sys.argv[4]), int(sys.argv[5]))
    else:
        compare()