#!/usr/bin/env python3
"""Cell-by-cell: Tables/main_exp.tex vs published.ts; Average vs mean of the four; ID/OOD vs main; other cross-checks.
argv: <v3 dir> <published.ts>"""
import hashlib, json, os, re, sys
from decimal import Decimal as D

v3, pub = sys.argv[1], sys.argv[2]
tex_path = os.path.join(v3, 'Tables/main_exp.tex')
tex = open(tex_path, encoding='utf-8').read()
print('main_exp.tex sha256', hashlib.sha256(tex.encode()).hexdigest())
src = open(pub, encoding='utf-8').read()
print('published.ts SOURCE_SHA256', re.search(r'SOURCE_SHA256 = "(\w+)"', src).group(1))
rows_pub = json.loads(re.search(r'PAPER_ROWS: PaperRow\[\] = (\[.*\]);', src, re.S).group(1))
KEYS = ['sc', 'an', 'cf', 'sq', 'avg', 'bleu', 'rouge']
CELL = re.compile(r'(\d+\.\d+)\\scriptsize\{\$\\pm(\d+\.\d+)\$\}')


def parse_main():
    rows = []
    section = None
    for n, line in enumerate(tex.split('\n'), 1):
        if 'Close-source' in line:
            section = 'closed'
        if 'Open-source' in line:
            section = 'open'
        cells = CELL.findall(line)
        if len(cells) == 7:
            rows.append({'line': n, 'name': line.split('&')[0].strip(), 'section': section, 'cells': cells})
    return rows


rows = parse_main()
print('rows main_exp', len(rows), 'rows published', len(rows_pub))
bad = 0
for r, p in zip(rows, rows_pub):
    if r['name'] != p['name'] or r['section'] != p['section']:
        print('NAME/SECTION', r['name'], p['name'])
        bad += 1
    for k, (m, s) in zip(KEYS, r['cells']):
        if D(m) != D(str(p['scores'][k])) or D(s) != D(str(p['spread'][k])):
            print('CELL', r['name'], k, m, s, p['scores'][k], p['spread'][k])
            bad += 1
print('cells compared', sum(len(r['cells']) * 2 for r in rows), 'mismatches', bad)

print('\nAverage vs mean of four dims (printed at 2 decimals; a difference above 0.005+0.005 cannot be rounding)')
for r, p in zip(rows, rows_pub):
    four = [D(c[0]) for c in r['cells'][:4]]
    mean = sum(four) / 4
    avg = D(r['cells'][4][0])
    d = avg - mean
    print(f"main_exp.tex:{r['line']:>3} {r['name']:<20} printed {avg}  mean4 {mean:.4f}  diff {d:+.4f} {'**' if abs(d) > D('0.01') else ''}  avgMismatch field={p['avgMismatch']}")

# In-domain / out-of-domain
idood = open(os.path.join(v3, 'Tables/exp_idood.tex'), encoding='utf-8').read()
CELL2 = re.compile(r'(\d+\.\d+)\\tiny\{\$\\pm(\d+\.\d+)\$\}')
print('\nmain vs (ID+OOD)/2 per cell (2-decimal print, tolerance 0.015)')
main_by = {r['name']: r for r in rows}
worst = []
for n, line in enumerate(idood.split('\n'), 1):
    c = CELL2.findall(line)
    if len(c) == 10:
        name = line.split('&')[0].strip()
        for i in range(5):
            x = (D(c[i][0]) + D(c[i + 5][0])) / 2
            y = D(main_by[name]['cells'][i][0])
            worst.append((abs(x - y), name, ['sc', 'an', 'cf', 'sq', 'avg'][i], x, y, n))
worst.sort(reverse=True)
print(' cells', len(worst), ' above tolerance:', sum(1 for w in worst if w[0] > D('0.015')))
for w in worst[:12]:
    print(' ', f'{w[0]:.3f}', w[1], w[2], 'mean(ID,OOD)=', w[3], 'main=', w[4], f'exp_idood.tex:{w[5]}')

# Means of ID/OOD Average vs mean of ID four dims
print('\nID / OOD Average vs mean of their four')
for n, line in enumerate(idood.split('\n'), 1):
    c = CELL2.findall(line)
    if len(c) == 10:
        name = line.split('&')[0].strip()
        for lo, lab in ((0, 'ID'), (5, 'OOD')):
            m = sum(D(c[lo + i][0]) for i in range(4)) / 4
            d = D(c[lo + 4][0]) - m
            if abs(d) > D('0.01'):
                print(f'  exp_idood.tex:{n} {name} {lab} printed {c[lo + 4][0]} mean4 {m:.4f} diff {d:+.4f}')

# Cross-checks of other tables
full = open(os.path.join(v3, 'Tables/exp_main_cot_full.tex'), encoding='utf-8').read().split('\n')
sect = None
print('\nexp_main_cot_full rows vs main_exp (only w/ I.T. rows) and Average vs mean4')
for n, l in enumerate(full, 1):
    if 'Test w/ I.T.' in l:
        sect = 'with'
    if 'Test w/o I.T.' in l:
        sect = 'without'
    nums = re.findall(r'(?<![\w.])(\d+\.\d+)(?![\w.])', l)
    if len(nums) == 7:
        name = l.split('&')[0].strip()
        if sect == 'with' and name in main_by:
            same = [D(a) == D(b[0]) for a, b in zip(nums, main_by[name]['cells'])]
            print(f'  exp_main_cot_full.tex:{n} {name}: all equal main = {all(same)}')
        m = sum(D(x) for x in nums[:4]) / 4
        d = D(nums[4]) - m
        if abs(d) > D('0.01'):
            print(f'  exp_main_cot_full.tex:{n} {sect} {name}: printed {nums[4]} mean4 {m:.4f} diff {d:+.4f}')

cf3 = open(os.path.join(v3, 'Tables/main_exp_cf3.tex'), encoding='utf-8').read().split('\n')
print('\nmain_exp_cf3 (k=3) Average vs mean4')
for n, l in enumerate(cf3, 1):
    nums = re.findall(r'(?<![\w.])(\d+\.\d+)(?![\w.])', l)
    if len(nums) == 7:
        m = sum(D(x) for x in nums[:4]) / 4
        d = D(nums[4]) - m
        if abs(d) > D('0.01'):
            print(f'  main_exp_cf3.tex:{n} {l.split("&")[0].strip()}: printed {nums[4]} mean4 {m:.4f} diff {d:+.4f}')

print('\nexp_diff_judges GPT-4o column vs main Average (same judge, same models)')
dj = open(os.path.join(v3, 'Tables/exp_diff_judges.tex'), encoding='utf-8').read().split('\n')
for n, l in enumerate(dj, 1):
    p = [x.strip() for x in l.replace('\\\\', '').split('&')]
    if len(p) == 4 and re.match(r'(\\textbf\{)?\d', p[1]):
        name = p[0].replace('\\method', 'CoSER')
        v = re.sub(r'[^\d.]', '', p[1])
        if name in main_by:
            print(f"  exp_diff_judges.tex:{n} {name}: table {v} main {main_by[name]['cells'][4][0]} gap {D(v) - D(main_by[name]['cells'][4][0]):+}")
print('\nBold/underline: max per column')
closed_open = rows
for i, k in enumerate(KEYS):
    allmax = max(rows, key=lambda r: D(r['cells'][i][0]))
    openmax = max([r for r in rows if r['section'] == 'open'], key=lambda r: D(r['cells'][i][0]))
    bold_line = tex.split('\n')[allmax['line'] - 1]
    und_line = tex.split('\n')[openmax['line'] - 1]
    val = allmax['cells'][i][0]
    print(f" {k}: max all {allmax['name']} {val} (bold in tex: {('textbf{' + val) in bold_line or ('textbf{' + val) in bold_line.replace(chr(92) + 'underline{' + chr(92), chr(92))}), max open {openmax['name']} {openmax['cells'][i][0]} (underlined: {('underline{' + openmax['cells'][i][0]) in und_line or ('underline{' + chr(92) + 'textbf{' + openmax['cells'][i][0]) in und_line})")
print('\nModel names of main table not listed in Sections/5_Experiments.tex Models paragraph: see report. Spread: Qwen-2-72B BLEU ±', main_by['Qwen-2-72B']['cells'][5][1])
