import os
import torch
import numpy as np

def compare_debug_tensors(cpp_dir, py_dir, sample_idx=0, verbose=True):
    """
    Compare C++ and Python debug tensors for BB regressor (BatchNorm and ReLU outputs).
    Args:
        cpp_dir (str): Directory with C++ debug tensors.
        py_dir (str): Directory with Python debug tensors.
        sample_idx (int): Sample index to compare.
        verbose (bool): Print detailed comparison results.
    Returns:
        dict: Comparison metrics for each layer/output.
    """
    layers = ["conv3_1t", "conv3_2t", "conv4_1t", "conv4_2t"]
    stages = ["bn", "relu"]
    results = {}
    for layer in layers:
        for stage in stages:
            cpp_file = os.path.join(cpp_dir, f"sample_{sample_idx}_debug_{layer}_{stage}.pt")
            py_file = os.path.join(py_dir, f"sample_{sample_idx}_debug_{layer}_{stage}_py.pt")
            if not os.path.exists(cpp_file) or not os.path.exists(py_file):
                results[f"{layer}_{stage}"] = None
                if verbose:
                    print(f"[WARN] Missing file: {cpp_file if not os.path.exists(cpp_file) else py_file}")
                continue
            cpp_tensor = torch.load(cpp_file, map_location="cpu")
            py_tensor = torch.load(py_file, map_location="cpu")
            cpp_np = cpp_tensor.cpu().detach().numpy().astype(np.float32)
            py_np = py_tensor.cpu().detach().numpy().astype(np.float32)
            if cpp_np.shape != py_np.shape:
                results[f"{layer}_{stage}"] = None
                if verbose:
                    print(f"[WARN] Shape mismatch for {layer}_{stage}: C++ {cpp_np.shape}, Py {py_np.shape}")
                continue
            flat_cpp = cpp_np.flatten()
            flat_py = py_np.flatten()
            mae = np.mean(np.abs(flat_cpp - flat_py))
            max_err = np.max(np.abs(flat_cpp - flat_py))
            l2_cpp = np.linalg.norm(flat_cpp)
            l2_py = np.linalg.norm(flat_py)
            cos_sim = np.dot(flat_cpp, flat_py) / (l2_cpp * l2_py) if l2_cpp > 0 and l2_py > 0 else float('nan')
            results[f"{layer}_{stage}"] = {
                "mae": mae,
                "max_err": max_err,
                "cos_sim": cos_sim,
                "l2_cpp": l2_cpp,
                "l2_py": l2_py,
                "shape": cpp_np.shape
            }
            if verbose:
                print(f"{layer}_{stage}: cos_sim={cos_sim:.6f}, mae={mae:.4e}, max_err={max_err:.4e}, shape={cpp_np.shape}")
    return results