import torch
import os
from ltr.models.backbone import resnet50

def save_tensor(tensor, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save(tensor.cpu(), path)

# Load the preprocessed input tensor exported from C++
preprocessed_input_path = 'test/input_samples/common/sample_0_image.pt'
input_obj = torch.load(preprocessed_input_path, map_location='cpu', weights_only=False)
if isinstance(input_obj, torch.Tensor):
    input_tensor = input_obj
else:
    # Try to extract tensor from JIT module
    if hasattr(input_obj, 'named_parameters'):
        params = list(input_obj.named_parameters())
        if params:
            input_tensor = params[0][1]
        else:
            raise RuntimeError('No tensor found in JIT module')
    else:
        raise RuntimeError('Unknown input object type')

# Add batch dimension if missing
if input_tensor.dim() == 3:
    input_tensor = input_tensor.unsqueeze(0)

# Load Python ResNet model
model = resnet50(output_layers=['layer1', 'layer2', 'layer3', 'layer4'], pretrained=False)
# Optionally load weights if needed (uncomment and set path if required)
# model.load_state_dict(torch.load('backbone_pure_tensors/state_dict.pt'))
model.eval()

# Forward pass and save intermediate outputs
x = input_tensor
out_dir = 'test/output_py/resnet_debug/'

# After conv1
x1 = model.conv1(x)
save_tensor(x1[0], os.path.join(out_dir, 'sample_0_after_conv1.pt'))
# After bn1
x2 = model.bn1(x1)
save_tensor(x2[0], os.path.join(out_dir, 'sample_0_after_bn1.pt'))
# After relu1
x3 = model.relu(x2)
save_tensor(x3[0], os.path.join(out_dir, 'sample_0_after_relu1.pt'))
# After maxpool
x4 = model.maxpool(x3)
save_tensor(x4[0], os.path.join(out_dir, 'sample_0_after_maxpool.pt'))
# After layer1
x5 = model.layer1(x4)
save_tensor(x5[0], os.path.join(out_dir, 'sample_0_after_layer1.pt'))
# After layer2
x6 = model.layer2(x5)
save_tensor(x6[0], os.path.join(out_dir, 'sample_0_after_layer2.pt'))
# After layer3
x7 = model.layer3(x6)
save_tensor(x7[0], os.path.join(out_dir, 'sample_0_after_layer3.pt'))
# After layer4
x8 = model.layer4(x7)
save_tensor(x8[0], os.path.join(out_dir, 'sample_0_after_layer4.pt'))

print('Saved all intermediate outputs for sample 0 using C++-preprocessed input.')