#!/usr/bin/env python3
"""GGML vs Python cosine similarity comparison for ACE-Step DiT.

Run from tests/ directory. All paths relative to CWD.

Usage:
    cd tests/
    ./debug-dit-cossim.py                 # turbo BF16
    ./debug-dit-cossim.py --quant Q6_K    # turbo Q6_K
    ./debug-dit-cossim.py --mode sft      # SFT BF16
    ./debug-dit-cossim.py --mode xl-turbo # XL turbo BF16
    ./debug-dit-cossim.py --mode all      # all 4 models
"""
import os, sys, subprocess, struct, shutil, argparse, json
import numpy as np

SEED = 42

MODE_CONFIG = {
    "turbo": {
        "gguf_base": "acestep-v15-turbo",
        "config_path": "acestep-v15-turbo",
        "steps": 8, "shift": 3.0, "guidance": 0.0, "n_layers": 24,
    },
    "sft": {
        "gguf_base": "acestep-v15-sft",
        "config_path": "acestep-v15-sft",
        "steps": 50, "shift": 1.0, "guidance": 1.0, "n_layers": 24,
    },
    "xl-turbo": {
        "gguf_base": "acestep-v15-xl-turbo",
        "config_path": "acestep-v15-xl-turbo",
        "steps": 8, "shift": 3.0, "guidance": 0.0, "n_layers": 32,
    },
    "xl-sft": {
        "gguf_base": "acestep-v15-xl-sft",
        "config_path": "acestep-v15-xl-sft",
        "steps": 50, "shift": 1.0, "guidance": 1.0, "n_layers": 32,
    },
}

def load_request():
    if not os.path.isfile("request0.json"):
        print("[Error] request0.json not found in CWD")
        sys.exit(1)
    with open("request0.json") as f:
        req = json.load(f)
    print(f"[Request] Loaded request0.json")
    return req

def save_dump(path, data):
    import torch
    if isinstance(data, torch.Tensor):
        data = data.detach().float().cpu().numpy()
    data = np.ascontiguousarray(data.astype(np.float32))
    shape = data.shape
    header = struct.pack("i", len(shape))
    for s in shape:
        header += struct.pack("i", s)
    with open(path, "wb") as f:
        f.write(header)
        f.write(data.tobytes())

def load_dump(path):
    raw = np.fromfile(path, dtype=np.float32)
    ndim = int(struct.unpack("i", struct.pack("f", raw[0]))[0])
    shape = [int(struct.unpack("i", struct.pack("f", raw[1+i]))[0]) for i in range(ndim)]
    data = raw[1 + ndim:]
    return data, shape

def _cos_flat(a, b):
    n = min(len(a), len(b))
    if n == 0:
        return 0.0
    a, b = a[:n], b[:n]
    d = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / d) if d > 1e-10 else 0.0

def cos(a, b, shape_a=None, shape_b=None):
    if shape_a and shape_b and len(shape_a) == 2 and len(shape_b) == 2:
        if shape_a[0] == shape_b[1] and shape_a[1] == shape_b[0]:
            ra = a.reshape(shape_a)
            rb = b.reshape(shape_b)
            c_normal = _cos_flat(ra.flatten(), rb.flatten())
            c_transposed = _cos_flat(ra.T.flatten(), rb.flatten())
            if c_transposed > c_normal:
                return c_transposed
            return c_normal
    return _cos_flat(a, b)

def stft_cos(a, b, win=2048, hop=512):
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]
    window = np.hanning(win)
    frames = (n - win) // hop + 1
    sa = np.zeros((frames, win // 2 + 1))
    sb = np.zeros((frames, win // 2 + 1))
    for i in range(frames):
        s = i * hop
        sa[i] = np.abs(np.fft.rfft(a[s:s+win] * window))
        sb[i] = np.abs(np.fft.rfft(b[s:s+win] * window))
    return _cos_flat(sa.flatten(), sb.flatten())

def codes_to_python_format(codes_csv):
    """Convert '43316,18426,...' to '<|audio_code_43316|><|audio_code_18426|>...'"""
    if not codes_csv:
        return ""
    return "".join(f"<|audio_code_{c.strip()}|>" for c in codes_csv.split(",") if c.strip())

# GGML runner

def run_ggml(dump_dir, req, cfg, gguf_path, adapter_dir=None):
    ggml_bin = "../build/ace-synth"
    if not os.path.isfile(ggml_bin):
        print(f"[GGML] binary not found: {ggml_bin}")
        return False
    os.makedirs(dump_dir, exist_ok=True)

    # Build request from input, override mode-specific params.
    # Model selection travels inside the JSON: synth_model is the GGUF filename,
    # adapter (if any) is the adapter file or PEFT directory name under its parent.
    merged = dict(req)
    merged["seed"] = SEED
    merged["inference_steps"] = cfg["steps"]
    merged["guidance_scale"] = cfg["guidance"]
    merged["shift"] = cfg["shift"]
    merged["thinking"] = False
    merged["synth_model"] = os.path.basename(gguf_path)
    merged["vae"] = "vae-BF16.gguf"

    adapters_root = None
    if adapter_dir:
        adapters_root = os.path.dirname(os.path.abspath(adapter_dir)) or "."
        merged["adapter"] = os.path.basename(os.path.normpath(adapter_dir))

    request_json = os.path.join(dump_dir, "request0.json")
    with open(request_json, "w") as f:
        json.dump(merged, f, indent=4)

    models_dir = os.path.dirname(os.path.abspath(gguf_path)) or "../models"
    cmd = [ggml_bin, "--models", models_dir]
    if adapters_root:
        cmd += ["--adapters", adapters_root]
    cmd += [
        "--request", request_json,
        "--dump", dump_dir,
        "--no-fa",
    ]
    print(f"[GGML] Running {os.path.basename(gguf_path)}...")
    r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=None, text=True)
    n = len([f for f in os.listdir(dump_dir) if f.endswith(".bin")])
    if r.returncode != 0:
        if n > 0:
            print(f"[GGML] WARNING: exit {r.returncode} but {n} dump files exist, continuing")
        else:
            print(f"[GGML] FAILED (exit {r.returncode})")
            if r.stdout:
                print(r.stdout[-500:])
            return False
    print(f"[GGML] Done, {n} dump files")
    return True

# Python runner

def run_python(dump_dir, req, cfg, adapter_dir=None):
    sys.path.insert(0, "../../ACE-Step-1.5")
    from acestep.handler import AceStepHandler

    os.makedirs(dump_dir, exist_ok=True)
    has_cfg = cfg["guidance"] > 1.0

    caption  = req["caption"]
    lyrics   = req.get("lyrics", "")
    bpm      = req.get("bpm", 0)
    duration = req["duration"]
    language = req.get("vocal_language", "en")

    print(f"[Python] Initializing {cfg['config_path']}...")
    handler = AceStepHandler()
    handler.initialize_service(
        project_root="..",
        config_path=cfg["config_path"],
        device="cuda",
    )

    if adapter_dir:
        # torch.nn forbids '.' in module names, PEFT derives the adapter name
        # from the directory basename. Sanitize so directory names like
        # 'ACE-Step-v1.5-chinese-new-year-LoRA' do not abort Python ref load.
        adapter_name = os.path.basename(os.path.normpath(adapter_dir)).replace(".", "_") or "default"
        lr = handler.add_lora(adapter_dir, adapter_name=adapter_name)
        print(f"[Python] LoRA: {lr}")

    model = handler.model
    _dumps = {}

    orig_text = handler.infer_text_embeddings
    def hooked_text(*a, **kw):
        r = orig_text(*a, **kw)
        _dumps["text_hidden"] = r[0].clone()
        return r
    handler.infer_text_embeddings = hooked_text

    orig_lyric = handler.infer_lyric_embeddings
    def hooked_lyric(*a, **kw):
        r = orig_lyric(*a, **kw)
        _dumps["lyric_embed"] = r[0].clone()
        return r
    handler.infer_lyric_embeddings = hooked_lyric

    orig_cond = model.prepare_condition
    def hooked_prepare(*a, **kw):
        r = orig_cond(*a, **kw)
        enc_hs, enc_mask, ctx = r
        _dumps["enc_hidden"] = enc_hs[0].clone()
        _dumps["context"] = ctx[0].clone()
        if has_cfg:
            null_expanded = model.null_condition_emb.expand_as(enc_hs)
            _dumps["null_enc_hidden"] = null_expanded[0].clone()
        return r
    model.prepare_condition = hooked_prepare

    orig_noise = model.prepare_noise
    def hooked_noise(*a, **kw):
        n = orig_noise(*a, **kw)
        _dumps["noise"] = n[0].clone()
        return n
    model.prepare_noise = hooked_noise

    decoder = model.decoder
    _step = [0]
    orig_fwd = decoder.forward
    def hooked_fwd(*args, **kwargs):
        xt_in = args[0] if args else kwargs.get('hidden_states')
        step = _step[0]
        if step > 0 and xt_in is not None:
            _dumps[f"dit_step{step - 1}_xt"] = xt_in[0].clone()
        out = orig_fwd(*args, **kwargs)
        vt = out[0]
        if has_cfg and vt.shape[0] == 2:
            _dumps[f"dit_step{step}_vt_cond"] = vt[0].clone()
            _dumps[f"dit_step{step}_vt_uncond"] = vt[1].clone()
        else:
            _dumps[f"dit_step{step}_vt_cond"] = vt[0].clone()
        if not has_cfg:
            _dumps[f"dit_step{step}_vt"] = vt[0].clone()
        _step[0] += 1
        return out
    decoder.forward = hooked_fwd

    if has_cfg:
        gen_globals = model.generate_audio.__func__.__globals__
        _apg_step = [0]
        orig_apg = gen_globals['apg_forward']
        def hooked_apg(*args, **kwargs):
            result = orig_apg(*args, **kwargs)
            _dumps[f"dit_step{_apg_step[0]}_vt"] = result[0].clone()
            _apg_step[0] += 1
            return result
        gen_globals['apg_forward'] = hooked_apg
        _dumps["null_condition_emb"] = model.null_condition_emb.squeeze().clone()

    _hooks = []
    def make_hook(name, step_filter=0):
        def hook(module, input, output):
            if _step[0] == step_filter:
                out = output[0] if isinstance(output, tuple) else output
                _dumps[name] = out[0].clone().float()
        return hook

    _hooks.append(decoder.proj_in.register_forward_hook(make_hook("hidden_after_proj_in")))
    _hooks.append(decoder.condition_embedder.register_forward_hook(make_hook("enc_after_cond_emb")))
    _hooks.append(decoder.layers[0].register_forward_hook(make_hook("hidden_after_layer0")))
    _hooks.append(decoder.layers[0].self_attn.register_forward_hook(make_hook("layer0_sa_output")))
    for li in [6, 12, 18, cfg["n_layers"] - 1]:
        _hooks.append(decoder.layers[li].register_forward_hook(make_hook(f"hidden_after_layer{li}")))

    _hooks.append(decoder.time_embed.register_forward_hook(make_hook("temb_t")))

    # Hook detokenizer (runs during prepare_condition, before diffusion)
    if hasattr(model, 'detokenizer'):
        def detok_hook(module, input, output):
            _dumps["detok_output"] = output[0].clone().float()
        _hooks.append(model.detokenizer.register_forward_hook(detok_hook))

    gen_kwargs = dict(
        captions=caption, lyrics=lyrics, bpm=bpm,
        audio_duration=float(duration), seed=str(SEED),
        use_random_seed=False, batch_size=1,
        inference_steps=cfg["steps"], shift=cfg["shift"],
        guidance_scale=cfg["guidance"],
        infer_method="ode", vocal_language=language,
        audio_code_string=codes_to_python_format(req.get("audio_codes", "")),
        key_scale=req.get("keyscale", ""),
        time_signature=req.get("timesignature", ""),
    )

    # When audio_codes are present, Python auto-sets is_covers=True via
    # conditioning_masks.py (instruction match + has_code_hint).
    # This makes it use decoded codes as context, matching C++ behavior.
    # Do NOT patch is_covers to False, that would use silence instead of codes.

    tag = f"{cfg['config_path']}, {cfg['steps']} steps"
    if has_cfg:
        tag += f", CFG {cfg['guidance']}"
    print(f"[Python] Generating ({tag})...")
    result = handler.generate_music(**gen_kwargs)

    if not result.get("success"):
        print(f"[Python] Generation failed: {result.get('error', 'unknown')}")
        return False

    for h in _hooks:
        h.remove()

    extra = result.get("extra_outputs", {})
    if extra.get("pred_latents") is not None:
        _dumps["dit_x0"] = extra["pred_latents"][0]

    audios = result.get("audios", [])
    if audios and "tensor" in audios[0]:
        _dumps["vae_audio"] = audios[0]["tensor"].squeeze(0)
        audio_np = audios[0]["tensor"].squeeze(0).cpu().numpy()
        wav_path = os.path.join(dump_dir, "output.wav")
        import wave
        n_samples = audio_np.shape[1]
        interleaved = np.empty(2 * n_samples, dtype=np.float32)
        interleaved[0::2] = audio_np[0]
        interleaved[1::2] = audio_np[1]
        pcm = (np.clip(interleaved, -1, 1) * 32767).astype(np.int16)
        with wave.open(wav_path, 'w') as wf:
            wf.setnchannels(2)
            wf.setsampwidth(2)
            wf.setframerate(48000)
            wf.writeframes(pcm.tobytes())
        print(f"[Python] Wrote {wav_path}: {n_samples} samples ({n_samples/48000:.2f}s @ 48kHz stereo)")

    for name, tensor in sorted(_dumps.items()):
        save_dump(os.path.join(dump_dir, f"{name}.bin"), tensor)

    print(f"[Python] Done, {len(_dumps)} dump files")
    return True

# comparison

def build_stages(cfg):
    has_cfg = cfg["guidance"] > 1.0
    steps = cfg["steps"]
    stages = [
        "text_hidden", "lyric_embed", "enc_hidden", "detok_output", "context", "noise",
        "temb_t", "hidden_after_proj_in", "enc_after_cond_emb",
        "layer0_sa_output", "hidden_after_layer0",
        "hidden_after_layer6", "hidden_after_layer12", "hidden_after_layer18",
        f"hidden_after_layer{cfg['n_layers'] - 1}",
    ]
    if has_cfg:
        stages += ["null_condition_emb", "null_enc_hidden"]
    if steps <= 8:
        step_indices = list(range(steps))
    else:
        step_indices = list(range(0, steps, 5))
        if (steps - 1) not in step_indices:
            step_indices.append(steps - 1)
    for si in step_indices:
        if has_cfg:
            stages.append(f"dit_step{si}_vt_cond")
            if si < 2:
                stages.append(f"dit_step{si}_vt_uncond")
        stages.append(f"dit_step{si}_vt")
        if si < steps - 1:
            stages.append(f"dit_step{si}_xt")
    stages += ["dit_x0", "vae_audio"]
    return stages

def compare(dirs, stages, tag):
    labels = sorted(dirs.keys())
    pairs = [(labels[i], labels[j]) for i in range(len(labels)) for j in range(i+1, len(labels))]

    print(f"[{tag}] Cosine similarities GGML vs Python")
    print(f"  {'stage':30s}", end="")
    for a, b in pairs:
        print(f" {a+' vs '+b:>14s}", end="")
    print()

    for stage in stages:
        data = {}
        for label, d in dirs.items():
            f = os.path.join(d, stage + ".bin")
            if os.path.isfile(f):
                data[label] = load_dump(f)
        if not data:
            continue
        print(f"  {stage:30s}", end="")
        for a, b in pairs:
            if a in data and b in data:
                da, sa = data[a]
                db, sb = data[b]
                c = cos(da, db, sa, sb)
                print(f" {c:>14.6f}", end="")
            else:
                print(f" {'N/A':>14s}", end="")
        print()

    vae_data = {}
    for label, d in dirs.items():
        f = os.path.join(d, "vae_audio.bin")
        if os.path.isfile(f):
            vae_data[label] = load_dump(f)
    if len(vae_data) >= 2:
        print(f"  {'vae_audio (STFT cosine)':30s}", end="")
        for a, b in pairs:
            if a in vae_data and b in vae_data:
                c = stft_cos(vae_data[a][0], vae_data[b][0])
                print(f" {c:>14.6f}", end="")
            else:
                print(f" {'N/A':>14s}", end="")
        print()

    if len(pairs) > 0:
        a_label, b_label = pairs[0]
        a_dir, b_dir = dirs[a_label], dirs[b_label]
        xt_stages = [s for s in stages if "_xt" in s]
        if xt_stages:
            print(f"[{tag}] Error growth GGML vs Python")
            print(f"  {'stage':22s} {'cos':>10s} {'max_err':>10s} {'mean_err':>10s}"
                  f" {'mean_A':>10s} {'std_A':>10s} {'mean_B':>10s} {'std_B':>10s}")
            for stage in xt_stages:
                fa = os.path.join(a_dir, stage + ".bin")
                fb = os.path.join(b_dir, stage + ".bin")
                if os.path.isfile(fa) and os.path.isfile(fb):
                    da_raw, sa = load_dump(fa)
                    db_raw, sb = load_dump(fb)
                    if len(sa) == 2 and len(sb) == 2 and sa[0] == sb[0] and sa[1] == sb[1]:
                        da = da_raw.reshape(sa)
                        db = db_raw.reshape(sb)
                        c_flat = _cos_flat(da.flatten(), db.flatten())
                        c_trans = _cos_flat(da.T.flatten(), db.flatten())
                        if c_trans > c_flat:
                            da = da.T
                        da, db = da.flatten(), db.flatten()
                    else:
                        da, db = da_raw, db_raw
                    n = min(len(da), len(db))
                    da, db = da[:n], db[:n]
                    c = _cos_flat(da, db)
                    diff = np.abs(da - db)
                    print(f"  {stage:22s} {c:10.6f} {diff.max():10.6f} {diff.mean():10.6f}"
                          f" {da.mean():10.6f} {da.std():10.6f} {db.mean():10.6f} {db.std():10.6f}")
                else:
                    missing = []
                    if not os.path.isfile(fa): missing.append(a_label)
                    if not os.path.isfile(fb): missing.append(b_label)
                    print(f"  {stage:22s} missing: {', '.join(missing)}")

# main

def run_mode(mode_name, cfg, req, gguf_path, adapter_dir=None):
    dump_ggml   = f"ggml-{mode_name}"
    dump_python = f"python-{mode_name}"

    tag = mode_name.upper() if mode_name == "sft" else mode_name.capitalize()
    cfg_str = f"steps={cfg['steps']}, shift={cfg['shift']}"
    if cfg['guidance'] > 1.0:
        cfg_str += f", CFG={cfg['guidance']}"
    print(f"[{tag}] {cfg_str} | {os.path.basename(gguf_path)}")

    if os.path.isdir(dump_ggml):
        shutil.rmtree(dump_ggml)
    if not run_ggml(dump_ggml, req, cfg, gguf_path, adapter_dir):
        print(f"[{tag}] GGML failed")
        return False

    if os.path.isdir(dump_python):
        shutil.rmtree(dump_python)
    if not run_python(dump_python, req, cfg, adapter_dir):
        print(f"[{tag}] Python failed")
        return False

    stages = build_stages(cfg)
    compare({"GGML": dump_ggml, "Python": dump_python}, stages, tag)
    return True

def main():
    ap = argparse.ArgumentParser(description="GGML vs Python cosine similarity comparison")
    ap.add_argument("--mode", default="turbo", choices=list(MODE_CONFIG.keys()) + ["all"],
                    help="which model to test (default: turbo)")
    ap.add_argument("--quant", default="BF16",
                    help="quantization suffix for GGUF (default: BF16, e.g. Q6_K, Q8_0)")
    ap.add_argument("--adapters", default=None,
                    help="path to adapter directory (optional)")
    args = ap.parse_args()

    req = load_request()

    modes = list(MODE_CONFIG.keys()) if args.mode == "all" else [args.mode]
    ok = True
    for m in modes:
        cfg = MODE_CONFIG[m]
        gguf_path = f"../models/{cfg['gguf_base']}-{args.quant}.gguf"
        if not os.path.isfile(gguf_path):
            print(f"[Error] GGUF not found: {gguf_path}")
            ok = False
            continue
        if not run_mode(m, cfg, req, gguf_path, args.adapters):
            ok = False

    return 0 if ok else 1

if __name__ == "__main__":
    sys.exit(main())