"""Embedding-compression study on cached eval embeddings.

For each cached task (eval_beir --save-emb output + the in-domain benchmark):
  fp16 @ d  ->  int8 @ d  ->  binary @ d (symmetric sign/hamming)
                          ->  binary @ d with fp32 queries (asymmetric)
A 256-dim binary vector is 32 bytes; 2560-dim binary is 320 bytes.
"""
import glob
import json
import os
import sys

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from eval_beir import load_beir, evaluate


def variants(q, d, dim):
    """Yield (name, q_emb, d_emb) at truncation `dim`, each row re-normalized."""
    qt = q[:, :dim] / np.linalg.norm(q[:, :dim], axis=1, keepdims=True)
    dt = d[:, :dim] / np.linalg.norm(d[:, :dim], axis=1, keepdims=True)
    yield f"fp16@{dim}", qt, dt
    qi = np.round(qt * 127).clip(-127, 127) / 127.0
    di = np.round(dt * 127).clip(-127, 127) / 127.0
    yield f"int8@{dim}", qi.astype(np.float32), di.astype(np.float32)
    qb = np.sign(qt).astype(np.float32) / np.sqrt(dim)
    db = np.sign(dt).astype(np.float32) / np.sqrt(dim)
    yield f"bin@{dim} ({dim//8}B)", qb, db
    yield f"bin-asym@{dim}", qt, db


def run(tag, q, d, score_fn, results):
    for dim in (256, 2560):
        for name, qv, dv in variants(q, d, dim):
            ndcg = score_fn(qv.astype(np.float32), dv.astype(np.float32))
            results.setdefault(tag, {})[name] = round(ndcg, 4)
            print(f"{tag} {name}: NDCG@10={ndcg:.4f}", flush=True)


def main():
    cache_dir = sys.argv[1] if len(sys.argv) > 1 else "/home/anon/pog/eval/emb_cache"
    label = sys.argv[2] if len(sys.argv) > 2 else "pog-v2"
    ckpt = sys.argv[3] if len(sys.argv) > 3 else "/home/anon/pog/checkpoints/pog-v2"
    out = "/home/anon/pog/eval/results_quant.json"
    results = {}

    for path in sorted(glob.glob(os.path.join(cache_dir, f"*_{label}.npz"))):
        task = os.path.basename(path).split("_")[0]
        z = np.load(path, allow_pickle=True)
        q, d = z["q"].astype(np.float32), z["d"].astype(np.float32)
        qids, dids = list(z["qids"]), list(z["dids"])
        _, _, rel = load_beir(task)
        run(task, q, d,
            lambda qv, dv: evaluate(qv, dv, qids, dids, rel)[0], results)

    # in-domain from cached features + v2 adapter
    sys.path.insert(0, "/home/anon/pog/adapter")
    from pog_adapter import POGAdapter
    from eval_indomain import load_setup, metrics
    N_LAYERS, N_SLOTS, D = 3, 6, 2560
    REC = N_LAYERS * N_SLOTS * D
    feat_path = "/home/anon/pog/features/train_v2.bin"
    n = os.path.getsize(feat_path) // (REC * 2)
    feats = np.memmap(feat_path, dtype=np.float16, mode="r", shape=(n, N_LAYERS, N_SLOTS, D))
    model = POGAdapter.load(ckpt, device="cuda").eval()
    q_rows, doc_rows, pos_col = load_setup()

    @torch.no_grad()
    def emb(rows):
        outs = []
        for i in range(0, len(rows), 8192):
            x = torch.from_numpy(np.ascontiguousarray(
                feats[np.asarray(rows[i:i + 8192])]).astype(np.float32)).cuda()
            outs.append(model.embed(x).cpu().numpy())
        return np.concatenate(outs).astype(np.float32)

    q, d = emb(list(q_rows)), emb(list(doc_rows))
    run("msmarco-dev", q, d, lambda qv, dv: metrics(qv, dv, pos_col)[0], results)

    json.dump(results, open(out, "w"), indent=2)
    print("saved", out)


if __name__ == "__main__":
    main()