#!/usr/bin/env python3
"""
Compare the inputs to the BB regressor (layer2, layer3, proposals, bbox) between C++ and Python for sample 0.
Print shape, min, max, mean, and check for equality.
"""
import torch
from pathlib import Path

def extract_tensor_from_jit(jit_module):
    try:
        params = list(jit_module.parameters())
        if params:
            return params[0]
        buffers = list(jit_module.buffers())
        if buffers:
            return buffers[0]
    except Exception:
        pass
    return None

def print_stats(label, t1, t2):
    print(f"{label}:")
    print(f"  C++: shape={t1.shape}, min={t1.min().item():.6f}, max={t1.max().item():.6f}, mean={t1.mean().item():.6f}")
    print(f"  Py:  shape={t2.shape}, min={t2.min().item():.6f}, max={t2.max().item():.6f}, mean={t2.mean().item():.6f}")
    if t1.shape == t2.shape:
        same = torch.allclose(t1, t2, atol=1e-6)
        print(f"  Allclose: {same}")
        print(f"  Max abs diff: {(t1-t2).abs().max().item():.6f}")
    else:
        print("  Shapes differ!")
    print()

def main():
    # Directories
    cpp_out = Path("test/output")
    py_in = Path("test/input_samples/common")

    # Sample index
    idx = 0

    # ResNet features (layer2, layer3)
    cpp_layer2 = extract_tensor_from_jit(torch.jit.load(str(cpp_out / "resnet" / f"sample_{idx}_layer2.pt")))
    cpp_layer3 = extract_tensor_from_jit(torch.jit.load(str(cpp_out / "resnet" / f"sample_{idx}_layer3.pt")))
    py_layer2 = extract_tensor_from_jit(torch.jit.load(str(py_in / f"sample_{idx}_layer2.pt"))) if (py_in / f"sample_{idx}_layer2.pt").exists() else None
    py_layer3 = extract_tensor_from_jit(torch.jit.load(str(py_in / f"sample_{idx}_layer3.pt"))) if (py_in / f"sample_{idx}_layer3.pt").exists() else None

    # Proposals
    cpp_proposals = extract_tensor_from_jit(torch.jit.load(str(cpp_out / "bb_regressor" / f"sample_{idx}_proposals.pt"))) if (cpp_out / "bb_regressor" / f"sample_{idx}_proposals.pt").exists() else None
    py_proposals = extract_tensor_from_jit(torch.jit.load(str(py_in / f"sample_{idx}_proposals.pt"))) if (py_in / f"sample_{idx}_proposals.pt").exists() else None

    # BBox
    cpp_bb = extract_tensor_from_jit(torch.jit.load(str(cpp_out / "bb_regressor" / f"sample_{idx}_bb.pt"))) if (cpp_out / "bb_regressor" / f"sample_{idx}_bb.pt").exists() else None
    py_bb = extract_tensor_from_jit(torch.jit.load(str(py_in / f"sample_{idx}_bb.pt"))) if (py_in / f"sample_{idx}_bb.pt").exists() else None

    print("=== BB Regressor Input Comparison (Sample 0) ===\n")
    if cpp_layer2 is not None and py_layer2 is not None:
        print_stats("Layer2", cpp_layer2, py_layer2)
    else:
        print("Layer2: Could not load both tensors\n")
    if cpp_layer3 is not None and py_layer3 is not None:
        print_stats("Layer3", cpp_layer3, py_layer3)
    else:
        print("Layer3: Could not load both tensors\n")
    if cpp_proposals is not None and py_proposals is not None:
        print_stats("Proposals", cpp_proposals, py_proposals)
    else:
        print("Proposals: Could not load both tensors\n")
    if cpp_bb is not None and py_bb is not None:
        print_stats("BBox", cpp_bb, py_bb)
    else:
        print("BBox: Could not load both tensors\n")

if __name__ == "__main__":
    main()