"""genderize v2 «fuso» — un solo modello, due reti interne (v1 ultra + v2 distillata D2).

Uso:
    g = GenderizeFuso('models/genderize_fuso_v2')
    g.predict(['Maria Rossi', 'Nguyen Van An'])   # -> [{'gender','p_male','country','p_country','top5'}]

Regole (DECISIONI 25/09): p(M) = media delle p(M) calibrate delle due reti; «M» se p(M) >= soglia
(0,60: scelta sui banchi, verificata su 6 insiemi freschi); paese = 0,7·v1 + 0,3·v2 sulle probabilità.
Ogni rete riceve il nome con la propria codifica (v1 byte maxlen 48, v2 codepoint maxlen 32).

T1 — instradamento per scrittura con traslitterazione (DOM-102, 27/09): la scrittura del nome
viene riconosciuta dai blocchi Unicode e si applica la regola misurata in
reports/trainer/g3_traslit_T1.json (scelta solo sulla metà di validazione, nessun parametro
appreso; totale test 85,82 contro 78,11 fuso e 74,31 v1): «nativo» classifica il nome originale,
«anyascii» la sua traslitterazione, «mix» la media delle probabilità delle due forme; le scritture
latine e quelle non mappate restano nativo. Il paese è sempre calcolato sul nome originale.
anyascii è una dipendenza runtime (licenza ISC) richiesta solo quando serve traslitterare.
Disambiguazioni: kana → giapponese (senza kana il solo Han va in cinese); lettere persiane
specifiche (پ چ ژ گ ک ی) distinguono il farsi dall'arabo; lettere kazake specifiche
(ә ғ қ ң ө ұ ү һ) e ucraine (ї є ґ і) distinguono i cirillici, altrimenti russo.

v2.1 (DOM-112, 28/09) — testa paese e selettore per paese:
• logit adjustment β=0,75 della testa paese: log(p_c) − β·log(π), dove p_c è la probabilità
  fusa (0,7·v1 + 0,3·v2) sul nome originale e π è la frequenza empirica del corpus v2 letta da
  models/genderize_fuso_v2/country_prior.json (accanto ai pesi, SHA nel manifest); β scelto solo
  sulla validation in DOM-105. La distribuzione restituita è rinormalizzata; l'argmax decide
  il paese.
• selettore per paese: se il paese predetto (dopo l'aggiustamento) è in S {AM GE MM TG TD WS},
  la decisione di genere usa il ramo scelto solo sui banchi VECCHI (residuo 226 v1 + nativo;
  mai il banco 226 v2): v1@0,5 oppure fuso a soglia diversa. Per quelle righe il p(M) restituito
  è il punteggio del ramo rimappato in modo che il taglio 0,60 riproduca esattamente la
  decisione del ramo (p_ramo + 0,60 − soglia_ramo, limitato a [0,1]); per tutte le altre righe
  resta fuso@0,60. «gender» e «p_male» restano sempre coerenti con la stessa regola di taglio.

v2.2 (DOM-113, 28/09) — rimedio «tratti + massa» per GE MM TG TD (WS escluso):
• se la decisione v2.1 e v1@0,50 discordano E il nome ha tratti tipici di quei paesi, la
  decisione di genere (e il punteggio, rimappato come il selettore: p_v1 + 0,60 − 0,50) viene
  dal ramo v1. Ambito: SOLO nomi di scritture latine (la validazione che ha scelto il rimedio
  è latina; sulle scritture native il ramo v1 leggerebbe byte fuori distribuzione e il trigger
  non è mai stato misurato). Tratti scelti SOLO sul banco di validazione DOM-113 (Wikidata CC0,
  congelato; mai il 226 v2) e vincolati a non danneggiare il residuo 226 v1 (−0,023 pp
  misurati):
  (i) token esatto del lessico del paese (supporto ≥ 2 in validazione, ovunque nel nome);
  (ii) onorifico birmano iniziale {Maung, Mg, Saw, Sao, U} con almeno 2 token;
  (iii) somma delle p_paese aggiustate su {AM GE MM TG TD WS} ≥ 0,15 con il paese a p massima
  tra {GE MM TG TD WS} attivo. Configurazione in models/genderize_fuso_v2/dom113_rimedio.json
  (SHA nel manifest del verdetto); se il file manca, il comportamento resta v2.1.
"""
from __future__ import annotations
import json
from pathlib import Path
import torch, torch.nn.functional as F
from inference import NameModel, encode, normalise

try:
    from anyascii import anyascii as _anyascii
except ImportError:
    _anyascii = None

REGOLA_T1 = {
    'zh': 'nativo', 'ja': 'nativo', 'ko': 'nativo', 'ar': 'nativo', 'he': 'nativo',
    'th': 'nativo', 'el': 'nativo', 'uk': 'nativo',
    'ka': 'anyascii', 'hy': 'anyascii', 'kk': 'anyascii', 'ru': 'anyascii',
    'fa': 'mix', 'bn': 'mix', 'hi': 'mix',
}

_BLOCCHI = (
    ('ka', 0x10A0, 0x10FF), ('hy', 0x0530, 0x058F), ('he', 0x0590, 0x05FF),
    ('ar', 0x0600, 0x06FF), ('bn', 0x0980, 0x09FF), ('hi', 0x0900, 0x097F),
    ('th', 0x0E00, 0x0E7F), ('el', 0x0370, 0x03FF), ('han', 0x3400, 0x9FFF),
    ('hangul', 0xAC00, 0xD7AF), ('hiragana', 0x3040, 0x309F), ('katakana', 0x30A0, 0x30FF),
    ('cyrillic', 0x0400, 0x052F), ('latino', 0x0041, 0x024F),
)
_PERSIANO = frozenset('پچژگکی')
_KAZAKO = frozenset('әғқңөұүһӘҒҚҢӨҰҮҺ')
_UCRAINO = frozenset('їєґіЇЄҐІ')


def script_t1(testo: str):
    conteggi: dict[str, int] = {}
    for c in testo:
        cp = ord(c)
        for nome, a, b in _BLOCCHI:
            if a <= cp <= b:
                conteggi[nome] = conteggi.get(nome, 0) + 1
                break
    if not conteggi:
        return None
    if 'hiragana' in conteggi or 'katakana' in conteggi:
        return 'ja'
    if 'hangul' in conteggi:
        return 'ko'
    if 'han' in conteggi:
        return 'zh'
    dominante = max(conteggi, key=lambda k: (conteggi[k], k != 'latino'))
    if dominante == 'latino':
        return None
    if dominante == 'ar':
        return 'fa' if _PERSIANO & set(testo) else 'ar'
    if dominante == 'cyrillic':
        if _KAZAKO & set(testo):
            return 'kk'
        if _UCRAINO & set(testo):
            return 'uk'
        return 'ru'
    return dominante


def _cap(s: str) -> str:
    return ' '.join(w[:1].upper() + w[1:] for w in s.split())


def _traslittera(nome: str) -> str:
    if _anyascii is None:
        raise RuntimeError('anyascii è richiesto per le scritture traslitterate: pip install anyascii (licenza ISC)')
    return _cap(_anyascii(nome)) or nome


def varianti_t1(nome: str) -> list[str]:
    regola = REGOLA_T1.get(script_t1(nome))
    if regola in (None, 'nativo'):
        return [nome]
    tr = _traslittera(nome)
    return [nome, tr] if regola == 'mix' else [tr]


class GenderizeFuso:
    def __init__(self, cartella: str, device: str = 'cpu', batch_size: int = 256):
        if batch_size < 1:
            raise ValueError('batch_size deve essere positivo')
        self.batch_size = batch_size
        base = Path(cartella); cfg = json.loads((base / 'fuso.json').read_text())
        self.soglia, self.peso_paese_v1, self.dev = float(cfg['soglia_M']), float(cfg['peso_paese_v1']), device
        self.reti = []
        for sub in ('v1', 'v2'):
            d = base / sub; c = json.loads((d / 'genderize_ultra.config.json').read_text())
            cal = json.loads((d / 'genderize_ultra.calibration.json').read_text()); mp = json.loads((d / 'maps.json').read_text())
            vc = mp.get('vocab_codepoint') if sub == 'v2' else None
            m = NameModel(len(mp['country']), ch=int(c['ch']), vocab_size=len(vc) + 2 if vc else 256)
            st = torch.load(d / 'genderize_ultra.pt', map_location='cpu', weights_only=True)
            m.load_state_dict(st.get('state_dict', st)); m.eval().to(device)
            paesi = mp['country'] if isinstance(mp['country'], list) else [k for k, _ in sorted(mp['country'].items(), key=lambda x: x[1])]
            self.reti.append((m, cal, vc, int(c.get('maxlen', mp.get('maxlen', 48))), paesi))
        self.paesi = self.reti[1][4]
        self.perm = torch.tensor([self.reti[0][4].index(p) for p in self.paesi])   # v1 -> ordine v2
        pri = json.loads((base / 'country_prior.json').read_text())
        self.beta = float(pri['beta'])
        pr = pri['prior']; assert len(pr) == len(self.paesi) and set(pr) == set(self.paesi)
        self.log_prior = torch.tensor([float(pr[p]) for p in self.paesi], dtype=torch.float64)
        sc = pri['selettore']['scelte']
        self._sel_v1 = torch.zeros(len(self.paesi), dtype=torch.bool)
        self._sel_off = torch.zeros(len(self.paesi))
        for c, b in sc.items():
            i = self.paesi.index(c)
            self._sel_v1[i] = b['ramo'] == 'v1'
            self._sel_off[i] = self.soglia - float(b['soglia'])
        # v2.2 (DOM-113): rimedio «tratti + massa»; senza config resta v2.1
        self._rim_lessico: frozenset[str] = frozenset()
        self._rim_onorifici: frozenset[str] = frozenset()
        self._rim_tau = 1.0
        rp = base / 'dom113_rimedio.json'
        if rp.exists():
            rim = json.loads(rp.read_text())
            lessico = set()
            for c in rim['paesi_attivi']:
                lessico.update(rim['lessico'][c])
            self._rim_lessico = frozenset(lessico)
            self._rim_onorifici = frozenset(rim['onorifici_mm'])
            self._rim_tau = float(rim['tau_massa'])
            s6 = ['AM', 'GE', 'MM', 'TG', 'TD', 'WS']   # somma di massa (come DOM-113)
            s5 = ['GE', 'MM', 'TG', 'TD', 'WS']          # argmax di massa (senza AM)
            self._rim_s6 = torch.tensor([self.paesi.index(c) for c in s6])
            self._rim_s5 = torch.tensor([self.paesi.index(c) for c in s5])
            self._rim_attivi = torch.zeros(len(self.paesi), dtype=torch.bool)
            for c in rim['paesi_attivi']:
                self._rim_attivi[self.paesi.index(c)] = True

    @torch.no_grad()
    def probabilita(self, nomi: list[str]):
        if not nomi:
            return torch.empty(0), torch.empty((0, len(self.paesi)))
        # One normalized work list shared by both networks. Only selected rows
        # add a transliteration; equal forms and repeated names reuse inference.
        piatti, righe, indice = [], [], {}
        def aggiungi(nome):
            testo = normalise(nome)
            if testo not in indice:
                indice[testo] = len(piatti)
                piatti.append(testo)
            return indice[testo]
        for n in nomi:
            regola = REGOLA_T1.get(script_t1(n))
            a = aggiungi(n)
            b = aggiungi(_traslittera(n)) if regola in ('anyascii', 'mix') else None
            righe.append((a, b, regola))
        pm_reti, pc_reti = [], []
        for m, cal, vc, ml, _ in self.reti:
            lg_all, lc_all = [], []
            for i in range(0, len(piatti), self.batch_size):
                lg, lc = m(encode(piatti[i:i + self.batch_size], maxlen=ml, vocab_codepoint=vc).to(self.dev))
                lg_all.append(lg / cal.get('temp_gender', 1.0)); lc_all.append(lc / cal.get('temp_country', 1.0))
            pm_reti.append(F.softmax(torch.cat(lg_all), -1)[:, 0].cpu())
            pc_reti.append(F.softmax(torch.cat(lc_all), -1).cpu())
        a_t = torch.tensor([a for a, _, _ in righe])
        b_t = torch.tensor([b if b is not None else 0 for _, b, _ in righe])
        aa_m = torch.tensor([regola == 'anyascii' for _, _, regola in righe], dtype=torch.bool)
        mix_m = torch.tensor([regola == 'mix' for _, _, regola in righe], dtype=torch.bool)
        pn = (pm_reti[0][a_t] + pm_reti[1][a_t]) / 2
        pt = (pm_reti[0][b_t] + pm_reti[1][b_t]) / 2
        p_m = torch.where(aa_m, pt, torch.where(mix_m, (pn + pt) / 2, pn))
        p_c = self.peso_paese_v1 * pc_reti[0][a_t][:, self.perm] + (1 - self.peso_paese_v1) * pc_reti[1][a_t]
        # v2.1: logit adjustment della testa paese — log(p_c) − β·log(π), float64 come DOM-105;
        # la distribuzione restituita è rinormalizzata e l'argmax decide il paese.
        lg_c = torch.log(p_c.to(torch.float64)) - self.beta * self.log_prior
        paese = lg_c.argmax(-1)
        p_c = torch.softmax(lg_c, -1).to(p_m.dtype)
        # v2.1: selettore per paese — per i paesi di S il punteggio del ramo scelto è rimappato
        # perché il taglio self.soglia (0,60) riproduca esattamente la decisione del ramo.
        base = torch.where(self._sel_v1[paese], pm_reti[0][a_t], p_m)
        p_m = torch.clamp(base + self._sel_off[paese], 0.0, 1.0)
        # v2.2: rimedio DOM-113 — se v2.1 e v1 discordano e il nome ha tratti tipici
        # (lessico / onorifico birmano / massa del paese), decide il ramo v1
        # (punteggio rimappato: p_v1 + 0,60 − 0,50), come il selettore DOM-112.
        if self._rim_lessico:
            s6 = p_c[:, self._rim_s6].sum(-1)
            cS = self._rim_s5[p_c[:, self._rim_s5].argmax(-1)]
            massa = (s6 >= self._rim_tau) & self._rim_attivi[cS]
            # il rimedio vale solo per nomi latini: la validazione DOM-113 che lo ha
            # scelto è latina; sulle scritture native il ramo v1 girerebbe sui byte
            # originali (fuori distribuzione) e il trigger non è stato mai misurato
            tratti = torch.tensor(
                [script_t1(n) is None and (bool(massa[j]) or
                 any(t in self._rim_lessico for t in n.split())
                 or (len(n.split()) >= 2 and n.split()[0] in self._rim_onorifici))
                 for j, n in enumerate(nomi)], dtype=torch.bool)
            pv1 = pm_reti[0][a_t]
            discorda = (p_m >= self.soglia) != (pv1 >= 0.5)
            p_m = torch.where(tratti & discorda,
                              torch.clamp(pv1 + self.soglia - 0.50, 0.0, 1.0), p_m)
        return p_m, p_c

    def predict(self, nomi: list[str]):
        out = []
        for i in range(0, len(nomi), self.batch_size):
            p_m, p_c = self.probabilita(nomi[i:i + self.batch_size]); top = p_c.topk(5, -1)
            for j in range(len(p_m)):
                out.append({'gender': 'M' if p_m[j] >= self.soglia else 'F', 'p_male': round(float(p_m[j]), 4),
                            'country': self.paesi[int(top.indices[j, 0])], 'p_country': round(float(top.values[j, 0]), 4),
                            'top5': [self.paesi[int(k)] for k in top.indices[j]]})
        return out