import torch
import numpy as np

class PilotNet(torch.nn.Module):
    def __init__(self):
        super(PilotNet, self).__init__()
        # Capas convolucionales con BatchNorm
        self.conv1 = torch.nn.Conv2d(3, 24, 5, stride=2)
        self.bn1 = torch.nn.BatchNorm2d(24)

        self.conv2 = torch.nn.Conv2d(24, 36, 5, stride=2)
        self.bn2 = torch.nn.BatchNorm2d(36)

        self.conv3 = torch.nn.Conv2d(36, 48, 5, stride=2)
        self.bn3 = torch.nn.BatchNorm2d(48)

        self.conv4 = torch.nn.Conv2d(48, 64, 3)
        self.bn4 = torch.nn.BatchNorm2d(64)

        self.conv5 = torch.nn.Conv2d(64, 64, 3)
        self.bn5 = torch.nn.BatchNorm2d(64)

        # Capas fully connected con BatchNorm
        self.fc1 = torch.nn.Linear(1152, 100)
        self.bn_fc1 = torch.nn.BatchNorm1d(100)

        self.fc2 = torch.nn.Linear(100, 50)
        self.bn_fc2 = torch.nn.BatchNorm1d(50)

        self.fc3 = torch.nn.Linear(50, 10)
        self.fc4 = torch.nn.Linear(10, 1)

    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        x = torch.relu(self.bn2(self.conv2(x)))
        x = torch.relu(self.bn3(self.conv3(x)))
        x = torch.relu(self.bn4(self.conv4(x)))
        x = torch.relu(self.bn5(self.conv5(x)))

        x = x.view(x.size(0), -1)

        x = torch.relu(self.bn_fc1(self.fc1(x)))
        x = torch.relu(self.bn_fc2(self.fc2(x)))
        x = torch.relu(self.fc3(x))
        x = self.fc4(x)
        return x

print("Cargando checkpoint...")
checkpoint = torch.load('pilotnet_weights.pth')

print("Información del entrenamiento:")
print(f"  Epoch: {checkpoint['epoch']}")
print(f"  Normalización: {checkpoint.get('normalization_params')}")
print(f"  Torque min: {checkpoint.get('min_torque_ref')}")
print(f"  Torque max: {checkpoint.get('max_torque_ref')}")
print(f"  Dataset: {checkpoint.get('data_file')}")
print(f"  Error de test: {checkpoint.get('test_er')}")

# Carga el modelo
model = PilotNet().cuda()
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

print("\n✓ Modelo cargado correctamente en GPU")

# Test rápido
test_input = torch.randn(1, 3, 66, 200).cuda()
with torch.no_grad():
    output = model(test_input)
print(f"✓ Inferencia OK - Output de prueba: {output.cpu().item():.4f}")

# ← PEGA AQUÍ EL CÓDIGO DE DIAGNÓSTICO ↓

print("\n--- DIAGNÓSTICO DEL MODELO ---")

# 1. Verifica rango de salida
test_inputs = torch.randn(100, 3, 66, 200).cuda()
with torch.no_grad():
    outputs = model(test_inputs)

outputs_cpu = outputs.cpu().numpy()
print(f"Rango de salidas del modelo:")
print(f"  Min: {outputs_cpu.min():.4f}")
print(f"  Max: {outputs_cpu.max():.4f}")
print(f"  Mean: {outputs_cpu.mean():.4f}")
print(f"  Std: {outputs_cpu.std():.4f}")

print(f"\nRango esperado según checkpoint:")
print(f"  Min torque: {checkpoint.get('min_torque_ref')}")
print(f"  Max torque: {checkpoint.get('max_torque_ref')}")

# 2. Verifica si las salidas están en rango razonable
if checkpoint.get('min_torque_ref') and checkpoint.get('max_torque_ref'):
    min_ref = checkpoint['min_torque_ref']
    max_ref = checkpoint['max_torque_ref']

    in_range = ((outputs_cpu >= min_ref) & (outputs_cpu <= max_ref)).mean() * 100
    print(f"\n% de predicciones dentro del rango: {in_range:.1f}%")

    if in_range < 50:
        print("⚠️  ADVERTENCIA: Muchas predicciones fuera de rango esperado")
    else:
        print("✓ Predicciones en rango razonable")