"""Exploratory SONAR language-pivot pilot; no fitted transforms or training.

Run with the retained night8 environment on 4l. Outputs include every input,
generation, embedding, environment version and script hash. Hand-authored
translations are provisional references, not independent bilingual annotation.
"""
import argparse
import hashlib
import importlib.metadata
import json
from pathlib import Path
import time

import numpy as np
import torch
from sonar.inference_pipelines.text import (
    TextToEmbeddingModelPipeline, EmbeddingToTextModelPipeline,
)

LANG = {'fr': 'fra_Latn', 'en': 'eng_Latn', 'zh': 'zho_Hans'}
# Adjacent rows are controlled contrasts; the final four are ambiguous tu/vous
# contrasts (vous may mean formal singular OR plural), not unique gold targets.
DATA = [
 ('role', 'Le chat poursuit le chien.', 'The cat chases the dog.', '猫追狗。'),
 ('role', 'Le chien poursuit le chat.', 'The dog chases the cat.', '狗追猫。'),
 ('negation', 'Marie a ouvert la fenêtre.', 'Marie opened the window.', '玛丽打开了窗户。'),
 ('negation', "Marie n’a pas ouvert la fenêtre.", 'Marie did not open the window.', '玛丽没有打开窗户。'),
 ('number', 'Il y a trois livres sur la table.', 'There are three books on the table.', '桌子上有三本书。'),
 ('number', 'Il y a cinq livres sur la table.', 'There are five books on the table.', '桌子上有五本书。'),
 ('time', 'Paul arrivera demain.', 'Paul will arrive tomorrow.', '保罗明天会到。'),
 ('time', 'Paul est arrivé hier.', 'Paul arrived yesterday.', '保罗昨天到了。'),
 ('spatial', 'La lampe est au-dessus de la table.', 'The lamp is above the table.', '灯在桌子上方。'),
 ('spatial', 'La lampe est au-dessous de la table.', 'The lamp is below the table.', '灯在桌子下方。'),
 ('ownership', 'Alice a donné son livre à Paul.', 'Alice gave her book to Paul.', '爱丽丝把她的书给了保罗。'),
 ('ownership', 'Paul a donné son livre à Alice.', 'Paul gave his book to Alice.', '保罗把他的书给了爱丽丝。'),
 ('comparison', 'Le vélo rouge est plus cher que le vélo bleu.', 'The red bicycle is more expensive than the blue bicycle.', '红色自行车比蓝色自行车贵。'),
 ('comparison', 'Le vélo bleu est plus cher que le vélo rouge.', 'The blue bicycle is more expensive than the red bicycle.', '蓝色自行车比红色自行车贵。'),
 ('quantifier', 'Tous les enfants ont mangé une pomme.', 'All the children ate an apple.', '所有孩子都吃了苹果。'),
 ('quantifier', 'Aucun enfant n’a mangé de pomme.', 'No child ate an apple.', '没有孩子吃苹果。'),
 ('relation', 'La sœur de Paul connaît Marie.', 'Paul’s sister knows Marie.', '保罗的姐妹认识玛丽。'),
 ('relation', 'La sœur de Marie connaît Paul.', 'Marie’s sister knows Paul.', '玛丽的姐妹认识保罗。'),
 ('modality', 'Tu dois partir maintenant.', 'You must leave now.', '你现在必须离开。'),
 ('modality', 'Tu peux partir maintenant.', 'You may leave now.', '你现在可以离开。'),
 ('tu_vous_door', 'Peux-tu fermer la porte ?', 'Can you close the door?', '你能关门吗？'),
 ('tu_vous_door', 'Pouvez-vous fermer la porte ?', 'Can you close the door?', '您能关门吗？'),
 ('tu_vous_bag', 'Tu as oublié ton sac.', 'You forgot your bag.', '你忘了你的包。'),
 ('tu_vous_bag', 'Vous avez oublié votre sac.', 'You forgot your bag.', '您忘了您的包。'),
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', type=Path, required=True)
    ap.add_argument('--beam-size', type=int, default=1)
    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    torch.manual_seed(0)
    torch.set_num_threads(4)
    dev = torch.device('cuda:0')
    items = [dict(id=i, family=r[0], fr=r[1], en=r[2], zh=r[3]) for i, r in enumerate(DATA)]
    result = dict(items=items, outputs={}, metrics={}, meta=dict(
        status='running', kind='exploratory hand-authored pilot', seed=0,
        encoder='text_sonar_basic_encoder', decoder='text_sonar_basic_decoder',
        dtype='float32', beam_size=args.beam_size, max_gen_len=[0, 64], batch_size=8,
        gpu=torch.cuda.get_device_name(0),
        script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        versions={p: importlib.metadata.version(p) for p in ['torch', 'sonar-space', 'fairseq2']},
        reference_caveat='Assistant-authored. tu/vous Chinese references select formal singular; plural is also valid for vous. English loses that distinction.'))
    def save():
        result['meta']['elapsed_s'] = round(time.time() - t0, 2)
        (args.out / 'results.json').write_text(json.dumps(result, ensure_ascii=False, indent=2))
    save()
    enc = TextToEmbeddingModelPipeline(encoder='text_sonar_basic_encoder', tokenizer='text_sonar_basic_encoder', device=dev, dtype=torch.float32)
    dec = EmbeddingToTextModelPipeline(decoder='text_sonar_basic_decoder', tokenizer='text_sonar_basic_encoder', device=dev, dtype=torch.float32)
    enc.model.eval(); dec.model.eval()
    print('Models loaded', time.time()-t0, flush=True)
    Z = {}
    def encode(texts, lang):
        return enc.predict(texts, source_lang=LANG[lang], batch_size=8).float().cpu()
    def decode(z, lang):
        return dec.predict(z, target_lang=LANG[lang], batch_size=8, beam_size=args.beam_size, max_gen_len=(0, 64))
    with torch.inference_mode():
        for lang in LANG:
            Z[lang] = encode([it[lang] for it in items], lang)
        for source in LANG:
            for target in LANG:
                key = source + '_to_' + target
                result['outputs'][key] = decode(Z[source], target)
                save(); print(key, round(time.time()-t0, 1), flush=True)
        # Control the extra decode/re-encode step with French and Chinese pivots.
        for pivot in LANG:
            key = 'fr_via_' + pivot + '_to_zh'
            Z['fr_via_' + pivot] = encode(result['outputs']['fr_to_' + pivot], pivot)
            result['outputs'][key] = decode(Z['fr_via_' + pivot], 'zh')
            save(); print(key, round(time.time()-t0, 1), flush=True)
        repeat = decode(Z['fr'][:1], 'zh')[0]
        result['meta']['repeat_first_fr_to_zh_exact'] = repeat == result['outputs']['fr_to_zh'][0]
    # Geometry is descriptive, not an independent correctness oracle.
    norm = {k: torch.nn.functional.normalize(v, dim=1).numpy() for k, v in Z.items()}
    for a, b in [('fr', 'en'), ('fr', 'zh'), ('en', 'zh')]:
        sims = norm[a] @ norm[b].T
        same = sims.diagonal()
        contrast = sims[np.arange(len(items)), np.arange(len(items)) ^ 1]
        result['metrics'][a + '_' + b] = dict(
            paired_cos_mean=float(same.mean()), paired_cos= same.tolist(),
            minimal_pair_cos=contrast.tolist(),
            correct_above_contrast=int((same[:20] > contrast[:20]).sum()),
            unambiguous_n=20,
            unrelated_cos_mean=float(np.mean([sims[i, j] for i in range(24) for j in range(24) if i//2 != j//2])))
    for pivot in LANG:
        result['metrics']['fr_via_' + pivot] = dict(
            cosine_to_original_mean=float((norm['fr'] * norm['fr_via_' + pivot]).sum(1).mean()),
            chinese_exact_agreement_with_direct=sum(a==b for a,b in zip(result['outputs']['fr_to_zh'], result['outputs']['fr_via_'+pivot+'_to_zh'])),
            n=24)
    np.savez_compressed(args.out / 'embeddings.npz', **{k: v.numpy() for k,v in Z.items()})
    result['meta']['status'] = 'complete'
    save()
    print(json.dumps(result['metrics'], ensure_ascii=False), flush=True)


if __name__ == '__main__':
    main()