#!/usr/bin/env python3
"""End-to-end Hangul -> mixed-script conversion (CP6).

Pipeline per sentence
  1. Kiwi tokenises; noun-like tokens (NNG/NNP/NNB/NR/XR/XPN/XSN) define convertible regions.
  2. Inside each region every substring of 2..8 syllables that is a lexicon key becomes a
     candidate span (lattice); a region with no lexicon hit at all goes to the Composer fallback
     (character-level candidates ranked by a Hanja bigram LM).
  3. The cross-encoder scores every (span, candidate) pair; softmax over the span's candidates
     gives p(best).
  4. Dynamic programming picks a non-overlapping set of spans maximising
     sum(len(span) * (p_best - tau)), so short low-confidence spans are left in Hangul.

Usage:
  convert.py --ckpt results/run5/ckpt --lexicon processed/lexicon_v3_sto.tsv "문장" ["문장" ...]
  convert.py --ckpt ... --lexicon ... --stdin        (one sentence per line, JSON out)
Options: --tau 0.5 --max-cands 12 --compose (enable Composer fallback, needs --chars/--counts)
"""
import argparse
import collections
import json
import math
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.candidates import CandidateGenerator, Composer, _std_initial  # noqa: E402
from src.train_scorer import Scorer, mark  # noqa: E402
from transformers import AutoTokenizer  # noqa: E402

NOUN_TAGS = {'NNG', 'NNP', 'NNB', 'NR', 'XR', 'XPN', 'XSN'}
AFFIX_TAGS = {'XSN', 'XPN', 'NNB'}          # single-syllable Sino affixes: 的 側 黨 長 性 化 ...
NUM = dict(zip('영일이삼사오육칠팔구십백천만억', '零一二三四五六七八九十百千萬億'))
NUM_UNIT = dict(zip('년월일시분초', '年月日時分秒'))
import re as _re
NUM_RE = _re.compile(r'(?<![가-힣])((?:[영일이삼사오육칠팔구십백천만억]){2,}(?:[년월일시분초])?|[영일이삼사오육칠팔구십백천만억][년월일시분초])(?![가-힣])')


class Converter:
    def __init__(self, ckpt, lexicon, device=None, max_cands=12, tau=0.5, composer=None, max_len=128,
                 hangul_counts=None, prior_floor=0.0, prior_gate=0.0, prior_mode='gate',
                 region='noun', affix=True, numerals=True, unknown_prior=0.0):
        self.device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu'))
        self.tok = AutoTokenizer.from_pretrained(ckpt)
        self.model = Scorer(ckpt)
        head = os.path.join(ckpt, 'head.pt')
        if os.path.exists(head):
            self.model.head.load_state_dict(torch.load(head, map_location='cpu'))
        self.model.to(self.device).eval()
        self.drop_tt = getattr(self.model.enc.config, 'type_vocab_size', 2) < 2
        self.gen = CandidateGenerator(lexicon, composer=composer)
        # candidate order = lexicon count order (CandidateGenerator keeps file order = count desc)
        self.max_cands, self.tau, self.max_len = max_cands, tau, max_len
        # conversion prior p_sino(form) = hanja_count / (hanja_count + hangul_count): how often the
        # form was written in Hanja vs Hangul in the mixed-script corpus (native words -> ~0)
        self.hanja_count = collections.Counter()
        with open(lexicon, encoding='utf-8') as fh:
            next(fh)
            for line in fh:
                cols = line.rstrip('\n').split('\t')
                if len(cols) > 3 and cols[3].isdigit():
                    self.hanja_count[_std_initial(cols[0])] += int(cols[3])
        self.hangul_count = {}
        if hangul_counts and os.path.exists(hangul_counts):
            with open(hangul_counts, encoding='utf-8') as fh:
                next(fh)
                for line in fh:
                    w, n = line.rstrip('\n').split('\t')[:2]
                    self.hangul_count[w] = int(n)
        self.prior_floor = prior_floor
        self.prior_gate, self.prior_mode = prior_gate, prior_mode
        self.region, self.affix, self.numerals = region, affix, numerals
        self.unknown_prior = unknown_prior

    def p_sino(self, form):
        if not self.hangul_count:
            return 1.0
        h = self.hanja_count.get(_std_initial(form), 0)
        g = self.hangul_count.get(form, 0)
        if h == 0 and g == 0:
            return self.unknown_prior   # form unknown to the historical corpus (modern terms, OCR fragments)
        return max(self.prior_floor, h / (h + g + 1.0))

    # ---- candidate spans -------------------------------------------------------------
    def spans(self, sentence):
        out = []
        toks = self.gen.kiwi.tokenize(sentence)
        if self.region in ('all', 'token'):
            regions = [(m.start(), m.end()) for m in _re.finditer(r'[가-힣]+', sentence)]
        else:
            regions = [(t.start, t.start + t.len) for t in toks if t.tag in NOUN_TAGS]
        # token mode: a span must start at a Kiwi token start and end at a token end (any tag),
        # which blocks matches that straddle word boundaries in unspaced old text (파견책 -> 견책)
        starts = {t.start for t in toks}; ends = {t.start + t.len for t in toks}
        # every token a span covers must be noun-like: blocks 건설과 -> 建設課 (과 is a particle) and
        # 위해 -> 危害 (verb stem + ending) while still allowing multi-token compounds (자연+과학)
        bad_pos = set()
        for t in toks:
            if t.tag not in NOUN_TAGS and t.tag not in ('SN', 'SL', 'SH'):
                bad_pos.update(range(t.start, t.start + t.len))
        if self.region == 'token':
            aligned = lambda a, b: a in starts and b in ends and not any(i in bad_pos for i in range(a, b))
        else:
            aligned = lambda a, b: True
        for s, e in regions:
            region = sentence[s:e]
            hit = False
            for a in range(len(region)):
                for L in range(2, min(8, len(region) - a) + 1):
                    if not aligned(s + a, s + a + L):
                        continue
                    w = region[a:a + L]
                    cands = self.gen.table.get(_std_initial(w))
                    if cands:
                        hit = True
                        out.append({'start': s + a, 'end': s + a + L, 'surface': w, 'cands': cands[: self.max_cands], 'src': 'lex'})
            if not hit and self.gen.composer and 2 <= len(region) <= 6:
                comp = self.gen.composer.compose(region, k=self.max_cands)
                if comp:
                    out.append({'start': s, 'end': e, 'surface': region, 'cands': comp, 'src': 'compose'})
        if self.affix:   # single-syllable bound Sino morphemes tagged by Kiwi, looked up as 1-char keys
            for t in toks:
                if t.tag in AFFIX_TAGS and t.len == 1:
                    cands = self.gen.table.get(t.form)
                    if cands:
                        out.append({'start': t.start, 'end': t.start + 1, 'surface': t.form, 'cands': cands[: self.max_cands], 'src': 'affix'})
        if self.numerals:   # rule: Sino-Korean numeral strings (+ time unit) -> Hanja, fixed p=1
            for m in NUM_RE.finditer(sentence):
                w = m.group(1)
                hj = ''.join(NUM.get(ch) or NUM_UNIT.get(ch) or ch for ch in w)
                out.append({'start': m.start(1), 'end': m.end(1), 'surface': w, 'cands': [hj], 'src': 'numeral'})
        return out

    # ---- scoring -----------------------------------------------------------------------
    @torch.no_grad()
    def score(self, sentence, spans, bs=96):
        texts, cands, owner = [], [], []
        for i, sp in enumerate(spans):
            ctx = mark(sentence, sp['start'], sp['end'])
            for c in sp['cands']:
                texts.append(ctx); cands.append(c); owner.append(i)
        scores = []
        for b in range(0, len(texts), bs):
            enc = self.tok(texts[b:b + bs], cands[b:b + bs], truncation='only_first', max_length=self.max_len,
                           padding=True, return_tensors='pt')
            if self.drop_tt:
                enc.pop('token_type_ids', None)
            enc = {k: v.to(self.device) for k, v in enc.items()}
            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.device.type == 'cuda'):
                scores.extend(self.model(**enc).float().cpu().tolist())
        per = collections.defaultdict(list)
        for i, s in zip(owner, scores):
            per[i].append(s)
        for i, sp in enumerate(spans):
            ss = per[i]
            m = max(ss); ex = [math.exp(x - m) for x in ss]; z = sum(ex)
            probs = [x / z for x in ex]
            j = max(range(len(ss)), key=lambda k: probs[k])
            sp['hanja'], sp['p'] = sp['cands'][j], probs[j]
            sp['ranked'] = sorted(zip(sp['cands'], probs), key=lambda kv: -kv[1])[:5]
            sp['p_sino'] = round(self.p_sino(sp['surface']), 3)
        return spans

    # ---- overlap resolution --------------------------------------------------------------
    def select(self, sentence, spans):
        n = len(sentence)
        best = [0.0] * (n + 1); choice = [None] * (n + 1)
        by_end = collections.defaultdict(list)
        for sp in spans:
            by_end[sp['end']].append(sp)
        for e in range(1, n + 1):
            best[e], choice[e] = best[e - 1], None
            for sp in by_end[e]:
                ps = 1.0 if sp['src'] == 'numeral' else self.p_sino(sp['surface'])
                if ps < self.prior_gate:
                    continue                      # native/rarely-Hanja form: leave in Hangul
                conf = sp['p'] * ps if self.prior_mode == 'product' else sp['p']
                gain = (sp['end'] - sp['start']) * (conf - self.tau)
                if gain > 0 and best[sp['start']] + gain > best[e]:
                    best[e], choice[e] = best[sp['start']] + gain, sp
        chosen, e = [], n
        while e > 0:
            if choice[e] is None:
                e -= 1
            else:
                chosen.append(choice[e]); e = choice[e]['start']
        return sorted(chosen, key=lambda s: s['start'])

    def convert(self, sentence):
        spans = self.spans(sentence)
        if not spans:
            return {'input': sentence, 'output': sentence, 'spans': []}
        spans = self.score(sentence, spans)
        chosen = self.select(sentence, spans)
        out, pos = [], 0
        for sp in chosen:
            out.append(sentence[pos:sp['start']]); out.append(sp['hanja']); pos = sp['end']
        out.append(sentence[pos:])
        return {'input': sentence, 'output': ''.join(out),
                'spans': [{k: sp[k] for k in ('start', 'end', 'surface', 'hanja', 'p', 'p_sino', 'src', 'ranked')} for sp in chosen]}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True); ap.add_argument('--lexicon', required=True)
    ap.add_argument('--tau', type=float, default=0.7); ap.add_argument('--max-cands', type=int, default=12)
    ap.add_argument('--compose', action='store_true'); ap.add_argument('--chars', nargs='*', default=[])
    ap.add_argument('--counts'); ap.add_argument('--stdin', action='store_true'); ap.add_argument('--device')
    ap.add_argument('--hangul-counts', help='okhc_hangul_counts.tsv for the conversion prior')
    ap.add_argument('--prior-gate', type=float, default=0.3); ap.add_argument('--prior-mode', default='gate', choices=['gate', 'product'])
    ap.add_argument('--region', default='token', choices=['noun', 'all', 'token']); ap.add_argument('--no-affix', action='store_true'); ap.add_argument('--no-numerals', action='store_true')
    ap.add_argument('sentences', nargs='*')
    ap.add_argument('--preset', choices=['historical', 'modern'], help='historical: gate 0.3/tau 0.7 (1920-60s newspapers); modern: gate 0.05/tau 0.5')
    a = ap.parse_args()
    if a.preset == 'modern':
        a.prior_gate, a.tau = 0.05, 0.5
    elif a.preset == 'historical':
        a.prior_gate, a.tau = 0.3, 0.7
    composer = Composer(a.chars, a.counts) if a.compose and a.chars else None
    cv = Converter(a.ckpt, a.lexicon, device=a.device, max_cands=a.max_cands, tau=a.tau, composer=composer,
                   hangul_counts=a.hangul_counts, prior_gate=a.prior_gate, prior_mode=a.prior_mode,
                   region=a.region, affix=not a.no_affix, numerals=not a.no_numerals)
    sents = [l.rstrip('\n') for l in sys.stdin] if a.stdin else a.sentences
    for s in sents:
        r = cv.convert(s)
        if a.stdin:
            print(json.dumps(r, ensure_ascii=False))
        else:
            print(r['output'])
            for sp in r['spans']:
                print(f"   {sp['surface']}→{sp['hanja']} p={sp['p']:.2f} p_sino={sp['p_sino']} {sp['src']} {[(h, round(p, 2)) for h, p in sp['ranked'][:3]]}")


if __name__ == '__main__':
    main()