"""
used for debug using tensor comparison
dump {name: tensor} into "log_hf.jsonl" and "log_srt.jsonl"
use the same name for two tensors that supposed to be close
recommend name like: "layer 2 after mlp"
"""

import json
import sys

import torch

if len(sys.argv) > 1:
    assert sys.argv[1] == "base"
    hf_log = "base_log_hf.jsonl"
    srt_log = "base_log_srt.jsonl"
else:
    hf_log = "log_hf.jsonl"
    srt_log = "log_srt.jsonl"


def load_data(filepath):
    tensors = {}
    with open(filepath, "r") as f:
        lines = f.readlines()
        for line in lines:
            data = json.loads(line)
            for k, v in data.items():
                tensors[k] = torch.tensor(v)
    return tensors


hf_tensors = load_data(hf_log)
srt_tensors = load_data(srt_log)


def get_diff(t1, t2):
    t1 = t1.reshape(t2.shape)
    max_diff = torch.max(abs(t1.reshape(t2.shape) - t2))
    l2_dis = torch.dist(t1, t2, p=2)
    return l2_dis, max_diff


for k, _ in srt_tensors.items():
    l2_dis, max_diff = get_diff(hf_tensors[k], srt_tensors[k])
    print(f"{k} {l2_dis=} {max_diff=}")
    if k == "layer 1 attn":
        print(hf_tensors[k])
        print(srt_tensors[k])
    if k == "layer 0 prefill k":
        print(srt_tensors[k].shape)
        print(hf_tensors[k].shape)