import os
import torch
import numpy as np

def extract_tensor_from_jit(obj):
    """Extract a tensor from a TorchScript RecursiveScriptModule or return the tensor if already a tensor."""
    import torch
    if isinstance(obj, torch.Tensor):
        return obj
    # Try common attribute names
    for attr in ["tensor", "weight", "value"]:
        if hasattr(obj, attr):
            t = getattr(obj, attr)
            if isinstance(t, torch.Tensor):
                return t
    # Try named_parameters and named_buffers
    if hasattr(obj, "named_parameters"):
        for name, param in obj.named_parameters():
            if isinstance(param, torch.Tensor):
                return param
    if hasattr(obj, "named_buffers"):
        for name, buf in obj.named_buffers():
            if isinstance(buf, torch.Tensor):
                return buf
    # Try children recursively
    if hasattr(obj, "children"):
        for child in obj.children():
            t = extract_tensor_from_jit(child)
            if t is not None:
                return t
    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")
                    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")
                    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

def compare_resnet_debug_tensors(cpp_dir, py_dir, sample_idx=0, verbose=True):
    """
    Compare C++ and Python debug tensors for ResNet (after conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4).
    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 stage.
    """
    # Map stage to (cpp_filename, py_filename)
    stage_map = {
        "after_conv1":      (f"sample_{sample_idx}_conv1_output.pt",      f"sample_{sample_idx}_after_conv1.pt"),
        "after_bn1":        (f"sample_{sample_idx}_bn1_output.pt",        f"sample_{sample_idx}_after_bn1.pt"),
        "after_relu1":      (f"sample_{sample_idx}_relu1_output.pt",      f"sample_{sample_idx}_after_relu1.pt"),
        "after_maxpool":    (f"sample_{sample_idx}_maxpool_output.pt",    f"sample_{sample_idx}_after_maxpool.pt"),
        "after_layer1":     (f"sample_{sample_idx}_layer1.pt",            f"sample_{sample_idx}_after_layer1.pt"),
        "after_layer2":     (f"sample_{sample_idx}_layer2.pt",            f"sample_{sample_idx}_after_layer2.pt"),
        "after_layer3":     (f"sample_{sample_idx}_layer3.pt",            f"sample_{sample_idx}_after_layer3.pt"),
        "after_layer4":     (f"sample_{sample_idx}_layer4.pt",            f"sample_{sample_idx}_after_layer4.pt"),
    }
    results = {}
    import torch, os
    for stage, (cpp_file, py_file) in stage_map.items():
        cpp_path = os.path.join(cpp_dir, cpp_file)
        py_path = os.path.join(py_dir, py_file)
        cpp_tensor = None
        py_tensor = None
        if os.path.exists(cpp_path):
            cpp_obj = torch.load(cpp_path, map_location='cpu', weights_only=False)
            cpp_tensor = extract_tensor_from_jit(cpp_obj)
        else:
            if verbose:
                print(f"[WARN] Missing file: {cpp_path}")
        if os.path.exists(py_path):
            py_obj = torch.load(py_path, map_location='cpu')
            py_tensor = extract_tensor_from_jit(py_obj)
        else:
            if verbose:
                print(f"[WARN] Missing file: {py_path}")
        if cpp_tensor is None or py_tensor is None:
            results[stage] = {
                'cosine_similarity': None,
                'allclose': None,
                'max_abs_diff': None,
                'cpp_shape': None if cpp_tensor is None else tuple(cpp_tensor.shape),
                'py_shape': None if py_tensor is None else tuple(py_tensor.shape),
            }
            if verbose:
                print(f"Warning: Cannot compare 'ResNet Debug {stage}' for sample {sample_idx}, one or both tensors are None.")
            continue
        # Flatten for cosine similarity
        cpp_flat = cpp_tensor.flatten().float()
        py_flat = py_tensor.flatten().float()
        if cpp_flat.shape != py_flat.shape:
            results[stage] = {
                'cosine_similarity': None,
                'allclose': False,
                'max_abs_diff': None,
                'cpp_shape': tuple(cpp_tensor.shape),
                'py_shape': tuple(py_tensor.shape),
            }
            if verbose:
                print(f"[WARN] Shape mismatch for {stage}: C++ {cpp_tensor.shape}, Py {py_tensor.shape}")
            continue
        cos_sim = torch.nn.functional.cosine_similarity(cpp_flat, py_flat, dim=0, eps=1e-8).item()
        allclose = torch.allclose(cpp_tensor, py_tensor, atol=1e-5, rtol=1e-3)
        max_abs_diff = (cpp_tensor - py_tensor).abs().max().item()
        results[stage] = {
            'cosine_similarity': cos_sim,
            'allclose': allclose,
            'max_abs_diff': max_abs_diff,
            'cpp_shape': tuple(cpp_tensor.shape),
            'py_shape': tuple(py_tensor.shape),
        }
        if verbose:
            print(f"ResNet Debug {stage}: cos_sim={cos_sim:.6f}, allclose={allclose}, max_abs_diff={max_abs_diff:.3e}, shape={cpp_tensor.shape}")
    return results