import torch
import os
from ltr.models.bbreg.atom_iou_net import AtomIoUNet

def compare_tensors(a, b, label):
    a = a.float().cpu().contiguous().view(-1)
    b = b.float().cpu().contiguous().view(-1)
    if a.shape != b.shape:
        print(f"{label}: Shape mismatch: {a.shape} vs {b.shape}")
        return
    cos_sim = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
    mae = torch.mean(torch.abs(a - b)).item()
    max_abs = torch.max(torch.abs(a - b)).item()
    print(f"{label}: cos_sim={cos_sim:.8f}, MAE={mae:.8e}, max_abs={max_abs:.8e}")

def main():
    # Layer names and file patterns
    layers = [
        ("conv3_1t", 0),
        ("conv3_2t", 0),
        ("conv4_1t", 0),
        ("conv4_2t", 0),
    ]
    model_dir = "exported_weights/bb_regressor"
    # Load Python model
    model = AtomIoUNet(input_dim=(512, 1024), pred_input_dim=(256, 256), pred_inter_dim=(256, 256))
    # Load weights from exported files
    state_dict = {}
    for lname, idx in layers:
        wfile = os.path.join(model_dir, f"{lname}_{idx}_weight.pt")
        bfile = os.path.join(model_dir, f"{lname}_{idx}_bias.pt")
        state_dict[f"{lname}.{idx}.weight"] = torch.load(wfile, map_location='cpu', weights_only=False)
        state_dict[f"{lname}.{idx}.bias"] = torch.load(bfile, map_location='cpu', weights_only=False)
    model.load_state_dict(state_dict, strict=False)
    # Compare each weight and bias
    for lname, idx in layers:
        py_w = getattr(getattr(model, lname)[idx], 'weight').data
        py_b = getattr(getattr(model, lname)[idx], 'bias').data
        cpp_w = state_dict[f"{lname}.{idx}.weight"]
        cpp_b = state_dict[f"{lname}.{idx}.bias"]
        compare_tensors(py_w, cpp_w, f"{lname}_weight")
        compare_tensors(py_b, cpp_b, f"{lname}_bias")

if __name__ == "__main__":
    main()