#!/usr/bin/env python3
"""Test attribute conditioning: generate the SAME composer/key with CONTRASTING attribute sets
and check the output shifts in the requested direction (feature adherence). MPS, load once."""
import os, sys, io, json
import torch, pretty_midi
ROOT = '/Users/alhill/projects/music/tmp/melody'
sys.path.insert(0, ROOT); sys.path.insert(0, f'{ROOT}/aria-repo')
from model_qwen import GPTQwen
from ariautils.tokenizer import AbsTokenizer
from attributes import extract, ATTR_ORDER

CK = torch.load(f'{ROOT}/ckpt_attr.pt', map_location='cpu'); A = CK['arch']
META = json.load(open(f'{ROOT}/attr_meta.json'))
DEV = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
TOK = AbsTokenizer()
MODEL = GPTQwen(vocab=META['vocab'], d=A['d'], nh=A['nh'], nkv=A['nkv'], nl=A['nl'], ctx=A['ctx']).to(DEV)
MODEL.load_state_dict(CK['model']); MODEL.eval()
print(f"loaded ckpt_attr step {CK['step']} ppl {2.718281828**CK['val']:.2f} on {DEV}", flush=True)

INSTR, BOS, EOS = META['instr_prefix_id'], META['bos'], META['eos']
COMP, KEYS, ATOK, AANY = META['composers'], META['keys'], META['attr_tokens'], META['attr_any']
PREFIX = set(COMP.values()) | set(KEYS.values()) | {INSTR, BOS, EOS, META['pad_id']}
for at in ATTR_ORDER:
    PREFIX |= set(ATOK[at].values()); PREFIX.add(AANY[at])
MASK = torch.zeros(META['vocab'])
for i in PREFIX:
    MASK[i] = float('-inf')                               # never GENERATE a prefix/attr token


def attr_ids(attrs):
    return [ATOK[at][attrs[at]] if at in attrs else AANY[at] for at in ATTR_ORDER]


def sample(logits, temp, topp):
    lg = (logits.to('cpu').float() + MASK) / temp
    p = torch.softmax(lg, dim=-1); sp, si = torch.sort(p, descending=True)
    sp = sp * ((torch.cumsum(sp, 0) - sp) < topp)
    return si[torch.multinomial(sp / sp.sum(), 1)].item()


@torch.no_grad()
def gen(composer, keymeta, attrs, n_tokens=400, temp=1.0, topp=0.96):
    seed = [INSTR, COMP.get(composer, COMP['unknown']), KEYS[keymeta], BOS] + attr_ids(attrs)
    toks = list(seed); notes = []
    logits, cache = MODEL.infer(torch.tensor([toks], device=DEV), cache=None, pos=0)
    for _ in range(n_tokens):
        nxt = sample(logits[0], temp, topp)
        if nxt == EOS or len(toks) >= MODEL.ctx - 1:
            break
        toks.append(nxt)
        if nxt not in PREFIX:
            notes.append(nxt)
        logits, cache = MODEL.infer(torch.tensor([[nxt]], device=DEV), cache=cache, pos=len(toks) - 1)
    return notes


def to_notes(note_ids):
    md = TOK.detokenize(TOK.decode([INSTR, BOS] + note_ids)); buf = io.BytesIO(); md.to_midi().save(file=buf)
    pm = pretty_midi.PrettyMIDI(io.BytesIO(buf.getvalue()))
    return sorted((n.start, n.pitch, n.end - n.start, n.velocity) for inst in pm.instruments for n in inst.notes)


TESTS = {
    'HIGH / DENSE / LOUD / ASCENDING': {'register': 'brilliant', 'density': 'torrential', 'dynamics': 'thunderous',
                                        'contour': 'ascending', 'motion': 'agitated', 'color': 'bright'},
    'LOW / SPARSE / SOFT / DESCENDING': {'register': 'bass', 'density': 'sparse', 'dynamics': 'hushed',
                                         'contour': 'descending', 'motion': 'still', 'color': 'dark'},
}
for name, attrs in TESTS.items():
    ni = gen('chopin', '0_maj', attrs)
    notes = to_notes(ni)
    got = extract(notes, spb=0.55) if len(notes) >= 3 else None
    print(f'\n### {name}  ({len(notes)} notes)')
    if got:
        for k in attrs:
            hit = '✓' if got[k] == attrs[k] else ('~' if k in ('register', 'density', 'dynamics') else '✗')
            print(f'    {k:11} asked {attrs[k]:11} -> got {got[k]:11} {hit}')