import os
import torch
import numpy as np

def extract_tensor_from_jit(obj):
    # Try to extract a tensor from a TorchScript module
    if isinstance(obj, torch.Tensor):
        return obj
    # Try .tensor attribute (used in some containers)
    if hasattr(obj, 'tensor') and isinstance(obj.tensor, torch.Tensor):
        return obj.tensor
    # Try named_parameters
    if hasattr(obj, 'named_parameters'):
        params = list(obj.named_parameters())
        if params:
            return params[0][1]
    # Try named_buffers
    if hasattr(obj, 'named_buffers'):
        bufs = list(obj.named_buffers())
        if bufs:
            return bufs[0][1]
    return None

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")
            cpp_tensor = None
            py_tensor = None
            if os.path.exists(cpp_file):
                try:
                    obj = torch.load(cpp_file, map_location="cpu", weights_only=False)
                    cpp_tensor = extract_tensor_from_jit(obj)
                except Exception as e:
                    print(f"[WARN] Could not load {cpp_file}: {e}")
            else:
                print(f"[WARN] Missing file: {cpp_file}")
            if os.path.exists(py_file):
                try:
                    obj = torch.load(py_file, map_location="cpu", weights_only=False)
                    py_tensor = extract_tensor_from_jit(obj)
                except Exception as e:
                    print(f"[WARN] Could not load {py_file}: {e}")
            else:
                print(f"[WARN] Missing file: {py_file}")
            if cpp_tensor is None or py_tensor is None:
                print(f"Warning: Cannot compare 'BBReg Debug_{layer.capitalize()}' for sample {sample_idx}, one or both tensors are None.")
                continue
            # Flatten and compare
            cpp_flat = cpp_tensor.detach().cpu().numpy().flatten()
            py_flat = py_tensor.detach().cpu().numpy().flatten()
            cos_sim = np.dot(cpp_flat, py_flat) / (np.linalg.norm(cpp_flat) * np.linalg.norm(py_flat) + 1e-12)
            results[(layer, stage)] = cos_sim
            if verbose:
                print(f"Cosine similarity for {layer} {stage}: {cos_sim:.6f}")
    return results