#!/usr/bin/env python3
"""
Script 45: Full 5-Way Document-Level FAR (HC / HC+MC / HC+MC+HSC)
==================================================================
Codex round-31 noted: Script 43 reports HC-only document-level FAR
(17.97% any-pair). The actual deployed five-way classifier treats
the MC band (cos>0.95 AND 5<dh<=15) as "non-hand-signed" too, with
worst-case document-level priority HC > MC > HSC > UN > LH. The
paper must report doc-level FAR for each alarm definition.

This script reuses Script 43's per-signature simulation but tracks
the full five-way category each source signature would receive
under the random-inter-CPA pool, then aggregates to document level
under three alarm definitions:
  D1: HC only
  D2: HC + MC
  D3: HC + MC + HSC ("any non-hand-signed verdict")

For each definition we report:
  - Per-signature FAR (fraction of source sigs that fall into the
    alarm category against random pool)
  - Document-level FAR (any sig in doc triggers alarm)

The five-way rule used (inherited from v3.20.0 §III-K):
  HC : cos > 0.95 AND dh <= 5
  MC : cos > 0.95 AND 5 < dh <= 15
  HSC: cos > 0.95 AND dh > 15
  UN : 0.837 < cos <= 0.95
  LH : cos <= 0.837

We compute these on the realised (max_cos, min_dh) pair (any-pair
semantic, which matches the deployed v3/v4 rule per codex).

Outputs:
  reports/v4_big4/doc_level_far_full/
    doc_far_full_results.json
    doc_far_full_report.md
"""

import json
import sqlite3
import numpy as np
from pathlib import Path
from datetime import datetime
from collections import defaultdict

DB = '/Volumes/NV2/PDF-Processing/signature-analysis/signature_analysis.db'
OUT = Path('/Volumes/NV2/PDF-Processing/signature-analysis/reports/'
           'v4_big4/doc_level_far_full')
OUT.mkdir(parents=True, exist_ok=True)

BIG4 = ('勤業眾信聯合', '安侯建業聯合', '資誠聯合', '安永聯合')
ALIAS = {'勤業眾信聯合': 'Firm A',
         '安侯建業聯合': 'Firm B',
         '資誠聯合': 'Firm C',
         '安永聯合': 'Firm D'}
SEED = 42

COS_HIGH = 0.95
COS_LOW = 0.837
DH_HC = 5
DH_MC_UPPER = 15


def hamming_vec(query_bytes, cand_bytes_list):
    q = int.from_bytes(query_bytes, 'big')
    out = np.empty(len(cand_bytes_list), dtype=np.int32)
    for i, c in enumerate(cand_bytes_list):
        out[i] = (q ^ int.from_bytes(c, 'big')).bit_count()
    return out


def classify_five_way(max_cos, min_dh):
    if max_cos > COS_HIGH and min_dh <= DH_HC:
        return 'HC'
    if max_cos > COS_HIGH and DH_HC < min_dh <= DH_MC_UPPER:
        return 'MC'
    if max_cos > COS_HIGH and min_dh > DH_MC_UPPER:
        return 'HSC'
    if COS_LOW < max_cos <= COS_HIGH:
        return 'UN'
    return 'LH'


def wilson_ci(k, n, z=1.96):
    if n == 0:
        return (None, None)
    phat = k / n
    denom = 1 + z * z / n
    centre = (phat + z * z / (2 * n)) / denom
    half = z * np.sqrt(phat * (1 - phat) / n + z * z / (4 * n * n)) / denom
    return (max(0.0, centre - half), min(1.0, centre + half))


def load_big4():
    conn = sqlite3.connect(f'file:{DB}?mode=ro', uri=True)
    cur = conn.cursor()
    cur.execute('''
        SELECT s.signature_id, s.assigned_accountant, a.firm,
               s.source_pdf,
               s.feature_vector, s.dhash_vector
        FROM signatures s
        JOIN accountants a ON s.assigned_accountant = a.name
        WHERE s.assigned_accountant IS NOT NULL
          AND s.feature_vector IS NOT NULL
          AND s.dhash_vector IS NOT NULL
          AND a.firm IN (?, ?, ?, ?)
    ''', BIG4)
    rows = cur.fetchall()
    conn.close()
    return rows


def main():
    print('=' * 72)
    print('Script 45: Full 5-Way Doc-Level FAR (HC / HC+MC / HC+MC+HSC)')
    print('=' * 72)
    rows = load_big4()
    n_sigs = len(rows)
    print(f'\nLoaded {n_sigs:,} Big-4 signatures')

    cpas = np.array([r[1] for r in rows])
    firms = np.array([ALIAS[r[2]] for r in rows])
    source_pdfs = np.array([r[3] for r in rows])
    feats = np.stack([np.frombuffer(r[4], dtype=np.float32)
                       for r in rows]).astype(np.float32)
    norms = np.linalg.norm(feats, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    feats = feats / norms
    dhashes = [r[5] for r in rows]

    cpa_to_idx = defaultdict(list)
    for i, c in enumerate(cpas):
        cpa_to_idx[c].append(i)
    cpa_to_idx = {c: np.array(v, dtype=np.int64)
                  for c, v in cpa_to_idx.items()}
    pool_sizes = {c: len(v) - 1 for c, v in cpa_to_idx.items()}
    all_idx = np.arange(n_sigs, dtype=np.int64)

    rng = np.random.default_rng(SEED)
    print('\nSimulating per-signature category under random inter-CPA pool...')
    categories = np.empty(n_sigs, dtype=object)
    max_cos_arr = np.zeros(n_sigs, dtype=np.float32)
    min_dh_arr = np.zeros(n_sigs, dtype=np.int32)
    for si in range(n_sigs):
        if si % 5000 == 0:
            print(f'  {si:,}/{n_sigs:,} ({si/n_sigs*100:.1f}%)')
        s_cpa = cpas[si]
        n_pool = pool_sizes[s_cpa]
        if n_pool <= 0:
            categories[si] = 'LH'
            continue
        same_cpa = cpa_to_idx[s_cpa]
        need = n_pool
        cand_indices = []
        attempts = 0
        while need > 0 and attempts < 10:
            draw = rng.choice(n_sigs, size=need * 2, replace=True)
            same_mask = np.isin(draw, same_cpa)
            ok = draw[~same_mask]
            cand_indices.extend(ok[:need].tolist())
            need -= len(ok[:need])
            attempts += 1
        if need > 0:
            pool_mask = np.ones(n_sigs, dtype=bool)
            pool_mask[same_cpa] = False
            pool_idx = all_idx[pool_mask]
            fb = rng.choice(pool_idx, size=need, replace=False)
            cand_indices.extend(fb.tolist())
        cand_indices = np.array(cand_indices[:n_pool], dtype=np.int64)
        cos_vec = feats[cand_indices] @ feats[si]
        dh_vec = hamming_vec(dhashes[si],
                             [dhashes[c] for c in cand_indices])
        max_cos = float(cos_vec.max())
        min_dh = int(dh_vec.min())
        max_cos_arr[si] = max_cos
        min_dh_arr[si] = min_dh
        categories[si] = classify_five_way(max_cos, min_dh)

    print('  Done.')

    # Per-signature FAR by category
    print('\n[Per-signature FAR by 5-way category]')
    cat_counts = defaultdict(int)
    for c in categories:
        cat_counts[c] += 1
    for cat in ['HC', 'MC', 'HSC', 'UN', 'LH']:
        n_c = cat_counts[cat]
        far = n_c / n_sigs
        lo, hi = wilson_ci(n_c, n_sigs)
        print(f'  {cat}: n={n_c:,}, FAR={far:.4f}, '
              f'CI=[{lo:.4f}, {hi:.4f}]')

    # Per-signature FAR under three alarm definitions
    print('\n[Per-signature FAR under alarm definitions]')
    alarm_d1 = (categories == 'HC')
    alarm_d2 = np.isin(categories, ['HC', 'MC'])
    alarm_d3 = np.isin(categories, ['HC', 'MC', 'HSC'])
    persig_fars = {
        'D1_HC_only': {
            'far': float(alarm_d1.mean()),
            'hits': int(alarm_d1.sum()),
            'n': int(n_sigs),
            'ci95': wilson_ci(int(alarm_d1.sum()), n_sigs),
        },
        'D2_HC_plus_MC': {
            'far': float(alarm_d2.mean()),
            'hits': int(alarm_d2.sum()),
            'n': int(n_sigs),
            'ci95': wilson_ci(int(alarm_d2.sum()), n_sigs),
        },
        'D3_HC_plus_MC_plus_HSC': {
            'far': float(alarm_d3.mean()),
            'hits': int(alarm_d3.sum()),
            'n': int(n_sigs),
            'ci95': wilson_ci(int(alarm_d3.sum()), n_sigs),
        },
    }
    for k, v in persig_fars.items():
        print(f'  {k}: FAR={v["far"]:.4f}, '
              f'CI=[{v["ci95"][0]:.4f}, {v["ci95"][1]:.4f}], '
              f'{v["hits"]:,}/{v["n"]:,}')

    # Document-level FAR under three alarm definitions
    print('\n[Document-level FAR under alarm definitions]')
    doc_idx = defaultdict(list)
    for i, pdf in enumerate(source_pdfs):
        doc_idx[pdf].append(i)
    n_docs = len(doc_idx)
    doc_d1 = 0
    doc_d2 = 0
    doc_d3 = 0
    for pdf, idxs in doc_idx.items():
        idxs_a = np.array(idxs, dtype=np.int64)
        if alarm_d1[idxs_a].any():
            doc_d1 += 1
        if alarm_d2[idxs_a].any():
            doc_d2 += 1
        if alarm_d3[idxs_a].any():
            doc_d3 += 1
    print(f'  n_documents: {n_docs:,}')
    print(f'  D1 (HC only):           FAR={doc_d1/n_docs:.4f} '
          f'({doc_d1:,}/{n_docs:,})')
    print(f'  D2 (HC+MC):             FAR={doc_d2/n_docs:.4f} '
          f'({doc_d2:,}/{n_docs:,})')
    print(f'  D3 (HC+MC+HSC):         FAR={doc_d3/n_docs:.4f} '
          f'({doc_d3:,}/{n_docs:,})')

    # Per-firm doc-level FAR (D2 = HC+MC, the operational alarm)
    print('\n[Per-firm doc-level FAR D2 (HC+MC)]')
    # Map each doc to its dominant firm (mode of its signatures' firms)
    doc_firm = {}
    for pdf, idxs in doc_idx.items():
        fs = firms[idxs]
        vals, counts = np.unique(fs, return_counts=True)
        doc_firm[pdf] = str(vals[np.argmax(counts)])
    per_firm_doc = {}
    for f in ['Firm A', 'Firm B', 'Firm C', 'Firm D']:
        pdfs_f = [pdf for pdf, fr in doc_firm.items() if fr == f]
        n_f = len(pdfs_f)
        if n_f == 0:
            continue
        d1_h = sum(1 for pdf in pdfs_f
                   if alarm_d1[np.array(doc_idx[pdf])].any())
        d2_h = sum(1 for pdf in pdfs_f
                   if alarm_d2[np.array(doc_idx[pdf])].any())
        d3_h = sum(1 for pdf in pdfs_f
                   if alarm_d3[np.array(doc_idx[pdf])].any())
        per_firm_doc[f] = {
            'n_docs': n_f,
            'D1_HC': d1_h / n_f,
            'D2_HC_MC': d2_h / n_f,
            'D3_HC_MC_HSC': d3_h / n_f,
        }
        print(f'  {f} (n={n_f:,}): D1={d1_h/n_f:.4f}, '
              f'D2={d2_h/n_f:.4f}, D3={d3_h/n_f:.4f}')

    results = {
        'meta': {
            'script': '45',
            'timestamp': datetime.now().isoformat(timespec='seconds'),
            'n_signatures': n_sigs,
            'n_documents': n_docs,
            'seed': SEED,
            'note': ('Full 5-way doc-level FAR under three alarm '
                     'definitions, with per-firm stratification.'),
        },
        'persig_category_counts': dict(cat_counts),
        'persig_far_by_alarm': persig_fars,
        'doc_far_by_alarm': {
            'D1_HC_only': doc_d1 / n_docs,
            'D2_HC_plus_MC': doc_d2 / n_docs,
            'D3_HC_plus_MC_plus_HSC': doc_d3 / n_docs,
            'n_docs': n_docs,
            'hits': {'D1': doc_d1, 'D2': doc_d2, 'D3': doc_d3},
        },
        'per_firm_doc_far': per_firm_doc,
    }
    json_path = OUT / 'doc_far_full_results.json'
    json_path.write_text(json.dumps(results, indent=2, ensure_ascii=False),
                         encoding='utf-8')
    print(f'\n[json] {json_path}')

    md = [
        '# Full 5-Way Doc-Level FAR (Script 45)',
        '', f'Generated: {results["meta"]["timestamp"]}',
        f'Big-4 signatures: {n_sigs:,}; documents: {n_docs:,}',
        '',
        ('Per signature, simulate a random inter-CPA candidate pool of '
         'size n_pool, compute deployed (max-cos, min-dh), assign 5-way '
         'category, then aggregate to document level under three alarm '
         'definitions.'),
        '',
        '## 5-Way category distribution under random inter-CPA pool',
        '',
        '| Category | n | %  |',
        '|---|---|---|',
    ]
    for cat in ['HC', 'MC', 'HSC', 'UN', 'LH']:
        n_c = cat_counts[cat]
        md.append(f'| {cat} | {n_c:,} | {n_c/n_sigs:.4f} |')
    md += ['',
           '## Per-signature FAR by alarm definition',
           '',
           '| Definition | rule | FAR | 95% CI | hits / n |',
           '|---|---|---|---|---|',
           f'| D1 | HC only | {persig_fars["D1_HC_only"]["far"]:.4f} | '
           f'[{persig_fars["D1_HC_only"]["ci95"][0]:.4f}, '
           f'{persig_fars["D1_HC_only"]["ci95"][1]:.4f}] | '
           f'{persig_fars["D1_HC_only"]["hits"]:,} / {n_sigs:,} |',
           f'| D2 | HC + MC | {persig_fars["D2_HC_plus_MC"]["far"]:.4f} | '
           f'[{persig_fars["D2_HC_plus_MC"]["ci95"][0]:.4f}, '
           f'{persig_fars["D2_HC_plus_MC"]["ci95"][1]:.4f}] | '
           f'{persig_fars["D2_HC_plus_MC"]["hits"]:,} / {n_sigs:,} |',
           f'| D3 | HC + MC + HSC | '
           f'{persig_fars["D3_HC_plus_MC_plus_HSC"]["far"]:.4f} | '
           f'[{persig_fars["D3_HC_plus_MC_plus_HSC"]["ci95"][0]:.4f}, '
           f'{persig_fars["D3_HC_plus_MC_plus_HSC"]["ci95"][1]:.4f}] | '
           f'{persig_fars["D3_HC_plus_MC_plus_HSC"]["hits"]:,} / {n_sigs:,} |',
           '',
           '## Document-level FAR by alarm definition',
           '',
           '| Definition | rule | FAR | hits / n_docs |',
           '|---|---|---|---|',
           f'| D1 | any sig HC | {doc_d1/n_docs:.4f} | {doc_d1:,} / {n_docs:,} |',
           f'| D2 | any sig HC or MC | {doc_d2/n_docs:.4f} | '
           f'{doc_d2:,} / {n_docs:,} |',
           f'| D3 | any sig HC, MC, or HSC | {doc_d3/n_docs:.4f} | '
           f'{doc_d3:,} / {n_docs:,} |',
           '',
           '## Per-firm doc-level FAR',
           '',
           '| Firm | n_docs | D1 (HC) | D2 (HC+MC) | D3 (HC+MC+HSC) |',
           '|---|---|---|---|---|']
    for f, s in per_firm_doc.items():
        md.append(f'| {f} | {s["n_docs"]:,} | {s["D1_HC"]:.4f} | '
                  f'{s["D2_HC_MC"]:.4f} | {s["D3_HC_MC_HSC"]:.4f} |')
    md.append('')
    md_path = OUT / 'doc_far_full_report.md'
    md_path.write_text('\n'.join(md), encoding='utf-8')
    print(f'[md  ] {md_path}')


if __name__ == '__main__':
    main()
