#!/usr/bin/env python3
"""Recompute exp/evaluation/test_set_gpt-4o.json (extracted with git show to a data directory) from its own flaws.
Tests the paper's formula and alternatives. argv: <eval json> <simulation json> <test_set json>.
Prints numbers only; never prints dialogue text."""
import json, math, statistics as st, sys
from collections import Counter

ev = json.load(open(sys.argv[1]))
sim = json.load(open(sys.argv[2]))
test = json.load(open(sys.argv[3]))
DIMS = ['Storyline Consistency', 'Anthropomorphism', 'Character Fidelity', 'Storyline Quality']
print('top keys', list(ev), 'scores', ev['scores'])
cases = ev['cases']
print('cases', len(cases), 'sim entries', len(sim), 'test_set entries', len(test))
k0 = next(iter(cases))
print('case keys', list(cases[k0]), 'critique keys', list(cases[k0]['critique']))
print('dim keys', list(cases[k0]['critique'][DIMS[0]]))

forms = {
    'paper main text: 100-5S+1.5R clip[0,100]': lambda S, R, A: max(0, min(100 - 5 * S + 1.5 * R, 100)),
    'official exact: 100-(S-0.3R)*5 clip[0,100]': lambda S, R, A: max(0, min(100 - (S - 0.3 * R) * 5, 100)),
    'same, no upper clip (paper-era)': lambda S, R, A: max(0, 100 - (S - 0.3 * R) * 5),
    'same, no clip at all': lambda S, R, A: 100 - (S - 0.3 * R) * 5,
    'appendix: 100-S+1.5R': lambda S, R, A: 100 - S + 1.5 * R,
    'no length correction: 100-5S': lambda S, R, A: 100 - 5 * S,
    'R counts env too: 100-5S+1.5A clip': lambda S, R, A: max(0, min(100 - 5 * S + 1.5 * A, 100)),
    'lambda=1.0': lambda S, R, A: max(0, min(100 - 5 * S + 1.0 * R, 100)),
    'lambda=2.0': lambda S, R, A: max(0, min(100 - 5 * S + 2.0 * R, 100)),
    'factor 4': lambda S, R, A: max(0, min(100 - 4 * S + 1.5 * R, 100)),
    'factor 6': lambda S, R, A: max(0, min(100 - 6 * S + 1.5 * R, 100)),
}
match = Counter()
total = 0
clipped_hi = Counter()
clipped_lo = Counter()
sev_vals = Counter()
sev_types = Counter()
rows = []
rs = []
for key, c in cases.items():
    msgs = c['simulation']
    R = sum(1 for m in msgs if m['role'] != 'Environment')
    A = len(msgs)
    rs.append(R)
    for d in DIMS:
        crit = c['critique'][d]
        fl = crit['flaws']
        ints = [f['severity'] for f in fl if isinstance(f.get('severity'), int)]
        for f in fl:
            sev_types[type(f.get('severity')).__name__] += 1
            if isinstance(f.get('severity'), int):
                sev_vals[f['severity']] += 1
        S = sum(ints)
        s = crit['score']
        total += 1
        for name, fn in forms.items():
            if abs(fn(S, R, A) - s) < 1e-9:
                match[name] += 1
        if s == 100 and 100 - 5 * S + 1.5 * R > 100:
            clipped_hi[d] += 1
        if s == 0 and 100 - 5 * S + 1.5 * R < 0:
            clipped_lo[d] += 1
        rows.append((key, d, S, R, A, s))
print('\ndimension scores recomputed', total)
for name in forms:
    print(f'  {name}: {match[name]}/{total}')
print('severity value histogram', sorted(sev_vals.items()))
print('severity types', sev_types)
print('upper clip hit', dict(clipped_hi), 'lower clip hit', dict(clipped_lo))
print('R per case: min/mean/max', min(rs), st.mean(rs), max(rs), 'histogram', sorted(Counter(rs).items()))
msgs_all = [len(c['simulation']) for c in cases.values()]
print('non-NSP messages per case: min/mean/max', min(msgs_all), st.mean(msgs_all), max(msgs_all))

# file's aggregates
for d in DIMS:
    print(' mean', d, st.mean(c['critique'][d]['score'] for c in cases.values()), 'file', ev['scores'][d])
print(' avg of dim means', st.mean(ev['scores'][d] for d in DIMS), 'file avg', ev['scores'].get('avg'))
print(' mean of case scores', st.mean(c['score'] for c in cases.values()))
print(' case score == mean of its four:', sum(abs(c['score'] - st.mean(c['critique'][d]['score'] for d in DIMS)) < 1e-9 for c in cases.values()), '/', len(cases))
print(' mean bleu', st.mean(c['critique']['bleu'] for c in cases.values()), 'file', ev['scores'].get('bleu'))
print(' mean rouge_l', st.mean(c['critique']['rouge_l'] for c in cases.values()), 'file', ev['scores'].get('rouge_l'))

# Vanilla regression on this sample
xs, ys = [], []
for key, c in cases.items():
    R = sum(1 for m in c['simulation'] if m['role'] != 'Environment')
    v = st.mean(100 - 5 * sum(f['severity'] for f in c['critique'][d]['flaws'] if isinstance(f.get('severity'), int)) for d in DIMS)
    xs.append(R)
    ys.append(v)
if len(set(xs)) > 1:
    mx, my = st.mean(xs), st.mean(ys)
    b = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sum((x - mx) ** 2 for x in xs)
    print('\nlinear fit vanilla (×5, no correction) on R, this sample only: slope', round(b, 4), 'intercept', round(my - b * mx, 4), 'n', len(xs))
    v1 = [st.mean(100 - sum(f['severity'] for f in c['critique'][d]['flaws'] if isinstance(f.get('severity'), int)) for d in DIMS) for c in cases.values()]
    print('mean of 100-S (no ×5) on this sample', round(st.mean(v1), 2), ' mean of 100-5S', round(my, 2), '; paper regression intercept 59.06 slope -1.59')

# Simulation file: shape only
s0 = sim[0]
print('\nsimulation keys', list(s0))
rolecount = Counter()
prefixed = 0
n = 0
nsp_end = 0
for e in sim:
    for m in e['simulation']:
        rolecount['NSP' if m['role'] == 'NSP' else ('Env' if m['role'] == 'Environment' else 'char')] += 1
        if m['role'] != 'NSP':
            n += 1
            if any(line.strip().startswith(m['role'] + ':') for line in m['content'].split('\n')):
                prefixed += 1
print('role counts', rolecount, 'speaker messages with own name prefix', prefixed, '/', n)
keys_sim = [f"{e['book_title']}-{e['i_p']}-{e['i_c']}" for e in sim]
print('case keys equal between sim and eval:', set(keys_sim) == set(cases))
kt = [f"{c['book']}-{c['plot']['i_p']}-{c['i_c']}" for c in test]
print('how many eval cases are in test_set.json', sum(k in set(kt) for k in cases), 'of', len(cases))
print('first indices of those cases in test_set', sorted(kt.index(k) for k in cases if k in kt)[:20])
c0 = sim[0]['circumstance']
print('key_characters fields in the simulation file circumstance', sorted({f for e in sim for ch in e['circumstance']['key_characters'] for f in ch}))