#!/usr/bin/env python3
"""FP-Fusion engine: 双池并行调度 + 归一化 + JSD/split-half 统计.

并行结构:
  pool D   : D 层 420 条单 token 采样 (semaphore=d_concurrency, temperature=1.0)
  pool TXT : 基线 T₀ (20 条极短请求) → I/K/C/S 文本层 36 条 (semaphore=text_concurrency,
             max_tokens=TEXT_MAX_TOKENS 截断)
两池互不依赖, asyncio.gather 同时跑。
"""

import asyncio
import json
import math
import random
import re
import time

import httpx

from battery import (ALL_CELL_DEFS, BASELINE_MAX_TOKENS, BASELINE_PROMPT,
                     BASELINE_SAMPLES, D_SAMPLES_PER_CELL, D_TEMPERATURE,
                     REFUSAL_STARTERS, TEXT_MAX_TOKENS, TEXT_TEMPERATURE)

# ---------------------------------------------------------------- 归一化 ----

_CN_DIGITS = {'零': 0, '一': 1, '二': 2, '两': 2, '三': 3, '四': 4, '五': 5,
              '六': 6, '七': 7, '八': 8, '九': 9, '十': 10}
_COLOR_EN = {'red': 'red', 'blue': 'blue', 'green': 'green', 'yellow': 'yellow',
             'black': 'black', 'white': 'white', 'purple': 'purple', 'violet': 'purple',
             'orange': 'orange', 'pink': 'pink', 'brown': 'brown', 'gray': 'gray',
             'grey': 'gray'}
_COIN_MAP = {'heads': 'heads', 'tails': 'tails', '正面': 'heads', '反面': 'tails',
             'head': 'heads', 'tail': 'tails'}
_WEEKDAY_MAP = {'monday': 'monday', 'tuesday': 'tuesday', 'wednesday': 'wednesday',
                'thursday': 'thursday', 'friday': 'friday', 'saturday': 'saturday',
                'sunday': 'sunday', '周一': 'monday', '星期一': 'monday', '礼拜一': 'monday',
                '周二': 'tuesday', '星期二': 'tuesday', '礼拜二': 'tuesday',
                '周三': 'wednesday', '星期三': 'wednesday', '礼拜三': 'wednesday',
                '周四': 'thursday', '星期四': 'thursday', '礼拜四': 'thursday',
                '周五': 'friday', '星期五': 'friday', '礼拜五': 'friday',
                '周六': 'saturday', '星期六': 'saturday', '礼拜六': 'saturday',
                '周日': 'sunday', '星期日': 'sunday', '星期天': 'sunday',
                '礼拜日': 'sunday', '礼拜天': 'sunday', '周末': 'sunday'}

_PUNCT_RE = re.compile(r'[\W_]+', re.UNICODE)


def _first_token(text):
    return text.split()[0] if text.split() else ''


def normalize_answer(raw, domain):
    """移植 detector normalizer 的主干规则, 返回 (canonical, category)."""
    if raw is None:
        return None, 'error'
    text = raw.strip()
    if not text:
        return None, 'empty'
    low = text.lower()
    if any(s in low for s in REFUSAL_STARTERS):
        return None, 'refusal'
    # NFC + 去标点/emoji + 全角数字转半角
    cleaned = re.sub(r'[\uFF01-\uFF5E]',
                     lambda m: chr(ord(m.group(0)) - 0xFEE0), text)
    cleaned = _PUNCT_RE.sub(' ', cleaned).strip().lower()
    if not cleaned:
        return None, 'empty'
    tok = _first_token(cleaned)

    kind = domain[0]
    if kind == 'int':
        digits = ''.join(ch if ch.isdigit() else str(_CN_DIGITS.get(ch, '')) for ch in tok)
        digits = re.sub(r'\s+', '', digits)
        if digits.isdigit():
            v = int(digits)
            if domain[1] <= v <= domain[2]:
                return str(v), 'valid'
        return tok, 'invalid'
    if kind == 'letter':
        if len(tok) == 1 and tok.isalpha():
            return tok, 'valid'
        m = re.fullmatch(r'[a-z]', tok) or re.match(r'^([a-z])', cleaned)
        if m:
            return m.group(1), 'valid'
        return tok, 'invalid'
    if kind == 'color':
        # 对齐 detector 参考约定: en 归一为英文标准色; zh 保留中文、去掉尾部'色'
        # (参考库实证: random-color:zh keys = {"蓝","蓝紫"}, en = {"blue",...})
        tok = _first_token(cleaned)
        if not tok:
            return None, 'empty'
        if all(ord(ch) < 128 for ch in tok):
            return _COLOR_EN.get(tok, tok), 'valid'
        if len(tok) > 1 and tok.endswith('色'):
            tok = tok[:-1]
        return tok, 'valid'
    if kind == 'coin':
        for k, v in _COIN_MAP.items():
            if k in cleaned:
                return v, 'valid'
        return tok, 'invalid'
    if kind == 'enum':
        for v in domain[1]:
            if v in cleaned:
                return v, 'valid'
        for k, v in _WEEKDAY_MAP.items():
            if k in cleaned:
                return v, 'valid'
        return tok, 'invalid'
    return tok, 'valid'  # word 域: 任意词有效


# ---------------------------------------------------------------- 统计 ----

def jsd_bits(counts_p, counts_q):
    """Jensen-Shannon divergence, base 2, 范围 [0,1]."""
    tp, tq = sum(counts_p.values()), sum(counts_q.values())
    if tp <= 0 or tq <= 0:
        return 0.0
    support = set(counts_p) | set(counts_q)
    hm = hp = hq = 0.0
    for k in support:
        p = counts_p.get(k, 0) / tp
        q = counts_q.get(k, 0) / tq
        m = (p + q) / 2
        if m > 0:
            hm -= m * math.log2(m)
        if p > 0:
            hp -= p * math.log2(p)
        if q > 0:
            hq -= q * math.log2(q)
    return min(1.0, max(0.0, hm - (hp + hq) / 2))


def distributions_by_cell(samples):
    """samples: [{'cell':, 'norm':, 'cat':, 'arrival':}] → {cell: Counter}"""
    out = {}
    for s in samples:
        if s['cat'] == 'valid' and s['norm'] is not None:
            out.setdefault(s['cell'], {})
            out[s['cell']][s['norm']] = out[s['cell']].get(s['norm'], 0) + 1
    return out


def compare_cells(dist_a, dist_b, min_valid=10):
    """逐 cell JSD(双方 ≥min_valid 才可比), 返回按 JSD 降序的 entries + meanJsd."""
    entries = []
    for cell in sorted(set(dist_a) & set(dist_b)):
        if sum(dist_a[cell].values()) < min_valid or sum(dist_b[cell].values()) < min_valid:
            continue
        entries.append({'cell': cell, 'jsd': jsd_bits(dist_a[cell], dist_b[cell]),
                        'valid_a': sum(dist_a[cell].values()),
                        'valid_b': sum(dist_b[cell].values())})
    entries.sort(key=lambda e: e['jsd'], reverse=True)
    mean = (sum(e['jsd'] for e in entries) / len(entries)) if entries else None
    return entries, mean


def split_half_jsd(samples, min_per_half=5):
    """按到达顺序奇偶对半分, 逐 cell JSD 后取平均."""
    halves = {}
    for s in samples:
        if s['cat'] != 'valid' or s['norm'] is None:
            continue
        key = (s['cell'], s['arrival'] % 2)
        halves.setdefault(key, {})
        halves[key][s['norm']] = halves[key].get(s['norm'], 0) + 1
    jsds = []
    for cell in {k[0] for k in halves}:
        even, odd = halves.get((cell, 0)), halves.get((cell, 1))
        if even and odd and sum(even.values()) >= min_per_half and sum(odd.values()) >= min_per_half:
            jsds.append(jsd_bits(even, odd))
    return (sum(jsds) / len(jsds)) if jsds else None


# ---------------------------------------------------------------- 引擎 ----

class FusionEngine:
    def __init__(self, api_url, model, timeout=120, d_samples=D_SAMPLES_PER_CELL,
                 baseline_samples=BASELINE_SAMPLES, text_limit=0,
                 d_concurrency=4, text_concurrency=3,
                 text_max_tokens=TEXT_MAX_TOKENS, thinking=False):
        self.base = api_url.rstrip('/')
        self.model = model
        self.timeout = timeout
        self.d_samples = d_samples
        self.baseline_samples = baseline_samples
        self.text_limit = text_limit        # >0 时只跑前 N 条文本探针(冒烟用)
        self.d_sem = asyncio.Semaphore(d_concurrency)
        self.text_sem = asyncio.Semaphore(text_concurrency)
        self.text_max_tokens = text_max_tokens
        self.extra = {'chat_template_kwargs': {'thinking': thinking}}
        self.records = []                   # 所有请求的原始记录
        self.baseline_p50 = None
        self.tokens_in = self.tokens_out = 0

    # -- 单次请求 ---------------------------------------------------------
    async def _chat(self, sem, user_prompt, max_tokens, temperature, system_prompt,
                    tag, rec):
        body = {'model': self.model, 'temperature': temperature,
                'max_tokens': max_tokens, 'stream': False,
                'messages': [{'role': 'system', 'content': system_prompt},
                             {'role': 'user', 'content': user_prompt}],
                **self.extra}
        t0 = time.perf_counter()
        err = None
        content, usage = None, None
        # 慢端点排队不可控: 所有请求统一放宽下限 300s(含 D 层单 token 请求)
        eff_timeout = max(self.timeout, 300)
        try:
            async with sem:
                async with httpx.AsyncClient(timeout=eff_timeout) as client:
                    r = await client.post(f'{self.base}/chat/completions', json=body)
                    r.raise_for_status()
                    data = r.json()
                    content = data['choices'][0]['message'].get('content') or ''
                    usage = data.get('usage') or {}
        except Exception as e:
            # httpx 超时异常的 str() 可能为空串 → 永远保证 error 非空(fail-closed 依赖它)
            err = str(e) or type(e).__name__
        latency = (time.perf_counter() - t0) * 1000
        if usage:
            self.tokens_in += usage.get('prompt_tokens') or 0
            self.tokens_out += usage.get('completion_tokens') or 0
        rec.update({'latency_ms': round(latency, 1), 'response': content,
                    'error': err, 'tag': tag,
                    'completion_tokens': (usage or {}).get('completion_tokens')})
        self.records.append(rec)
        return rec

    SYS_ONE_WORD = {'en': 'Answer with exactly one word. No punctuation, no explanation.',
                    'zh': '只回答一个词，不要标点，不要解释。'}
    SYS_DEFAULT = 'You are a helpful assistant.'

    # -- 基线 --------------------------------------------------------------
    async def _run_baseline(self):
        sem = asyncio.Semaphore(3)
        tasks = [self._chat(sem, BASELINE_PROMPT, BASELINE_MAX_TOKENS, 0.0,
                            self.SYS_DEFAULT, 'baseline',
                            {'id': f'baseline_{i}', 'layer': 'BASE'})
                 for i in range(self.baseline_samples)]
        done = [r for r in await asyncio.gather(*tasks) if not r['error']]
        lats = sorted(r['latency_ms'] for r in done)
        self.baseline_p50 = lats[len(lats) // 2] if lats else None

    # -- D 层 ---------------------------------------------------------------
    def _d_jobs(self):
        jobs = []
        for c in ALL_CELL_DEFS:
            for lang in ('en', 'zh'):
                cell_id = f"{c['id']}:{lang}"
                pool = c['par'][lang]
                for i in range(self.d_samples):
                    jobs.append((cell_id, c['domain'], random.choice(pool), lang))
        random.shuffle(jobs)   # 防单 cell 突发(触发缓存/限流偏差)
        return jobs

    async def _run_d_layer(self):
        sem = self.d_sem
        tasks = []
        for idx, (cell_id, domain, prompt, lang) in enumerate(self._d_jobs()):
            rec = {'id': f'd_{cell_id}_{idx}', 'layer': 'D', 'cell': cell_id,
                   'lang': lang, 'prompt': prompt, 'arrival': idx}
            tasks.append(self._chat(sem, prompt, 16, D_TEMPERATURE,
                                    self.SYS_ONE_WORD[lang], 'dist', rec))
        await asyncio.gather(*tasks)

    # -- 文本层 -------------------------------------------------------------
    async def _run_text_layer(self, probes):
        sem = self.text_sem
        tasks = []
        for p in probes:
            rec = {'id': p['id'], 'layer': p['layer'], 'prompt': p['text'],
                   'meta': {k: v for k, v in p.items()
                            if k in ('pair', 'lang', 'metacog', 'refusal_grad', 'len_ctrl')}}
            tasks.append(self._chat(sem, p['text'], self.text_max_tokens,
                                    TEXT_TEMPERATURE, self.SYS_DEFAULT, 'text', rec))
        await asyncio.gather(*tasks)

    # -- 主入口 --------------------------------------------------------------
    async def run(self, text_probes):
        if self.text_limit > 0:
            text_probes = text_probes[:self.text_limit]
        # 基线必须独占测量: 若与 D 池并发, CPU 端点的基线会被排队延迟污染
        await self._run_baseline()
        await asyncio.gather(self._run_d_layer(), self._run_text_layer(text_probes))
        return self.records

    # -- 结果整理 -------------------------------------------------------------
    def d_samples_normalized(self):
        """返回 [{cell, norm, cat, arrival}]"""
        out = []
        for r in self.records:
            if r.get('tag', r.get('layer')) != 'dist':
                continue
            # _chat 不知 domain; 由调用方(cell)反查 —— 在 run_fp_fusion 里完成
            out.append(r)
        return out


def build_d_normalized(records):
    """records(D层) → 归一化样本列表"""
    from battery import ALL_CELL_DEFS
    dom = {}
    for c in ALL_CELL_DEFS:
        dom[f'{c["id"]}:en'] = c['domain']
        dom[f'{c["id"]}:zh'] = c['domain']
    out = []
    for r in records:
        if r['layer'] != 'D':
            continue
        norm, cat = (None, 'error') if r['error'] else normalize_answer(
            r.get('response'), dom[r['cell']])
        out.append({'cell': r['cell'], 'norm': norm, 'cat': cat,
                    'arrival': r['arrival'], 'error': r['error']})
    return out


def load_reference(path):
    """加载 detector schema 的单模型参考指纹 → {cellId: counts}"""
    with open(path, encoding='utf-8') as f:
        ref = json.load(f)
    if ref.get('formatVersion') != 1 or not isinstance(ref.get('cells'), dict):
        raise ValueError(f'unsupported reference format: {path}')
    cells = {}
    for cid, c in ref['cells'].items():
        counts = c.get('counts', {})
        cells[cid] = {str(k): v for k, v in counts.items()}
    return {'model': ref.get('model'), 'cells': cells}