import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import SGD, lr_scheduler

torch.backends.cudnn.benchmark = False  # You can set it to True if you experience performance gains
torch.backends.cudnn.deterministic = False
from src.loss_functions.losses import AsymmetricLoss, ASLSingleLabel

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import torch.nn.functional as F

class MyCNN(nn.Module):
    def __init__(self, num_classes=12, dropout_prob=0.2, in_channels=3):
        super(MyCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=128, kernel_size=3, padding=1)
        self.global_avg_pooling = nn.AdaptiveAvgPool2d(1)
        self.conv2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64 *3* 3, 1024)
        self.fc2 = nn.Linear(1024, 256)
        self.fc3 = nn.Linear(256, num_classes)

        # Dropout layers
        self.dropout1 = nn.Dropout(p=dropout_prob)
        self.dropout2 = nn.Dropout(p=dropout_prob)
        
    
    def forward(self, x_input):
        # Apply convolutional and pooling layers
        x = F.leaky_relu(self.conv1(x_input))
        x = F.max_pool2d(x, 2)
        x = F.leaky_relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = F.leaky_relu(self.conv3(x))
        x = F.max_pool2d(x, 2)

        # Flatten the output for the fully connected layers
        x = x.view(x.size(0), -1)
        x = self.dropout1(x)
        x = F.leaky_relu(self.fc1(x))
        x = self.dropout2(x)
        x = F.leaky_relu(self.fc2(x))

        # Apply fully connected layers
    
        x = self.dropout2(x)
        x = self.fc3(x)
        return x

# Rest of the code remains unchanged

# Initialize the model
cell_attribute_model = MyCNN(num_classes=12, dropout_prob=0.5, in_channels=256).to(device)
cell_attribute_model.train()  # Set the model in training mode

# Initialize optimizer, criterion, and scheduler
optimizer_cell_model = torch.optim.SGD(cell_attribute_model.parameters(), lr=0.01, weight_decay=0.01)
step_size = 5
gamma = 0.1
scheduler_cell_model = lr_scheduler.StepLR(optimizer_cell_model, step_size=step_size, gamma=gamma)
#criterion = nn.CrossEntropyLoss()
criterion = AsymmetricLoss(gamma_neg=4, gamma_pos=1, clip=0.08, disable_torch_grad_focal_loss=True)
# criterion = ASLSingleLabel()


# /num_classes = 2
#criterion = nn.BCEWithLogitsLoss()  # Binary Cross-Entropy Loss

def cell_training(cell_attribute_model_main,cell_datas, labels):
    obj_batch_size = len(cell_datas)
     # Set the model in training mode
    #optimizer_cell_model.zero_grad()

        # Filter out instances with label=2 and their corresponding cell_datas
    # Filter out rows where any element in the row (excluding the first column) is equal to 2
    valid_indices = [i for i, row in enumerate(labels[:,1:]) if not torch.any(row[1:] == 2).item()]
    
    if not valid_indices:
        # print("No valid instances, skipping training.")
        object_batch_loss = torch.tensor(0.0, requires_grad=True, device=device)  # Initialize as a torch.Tensor

        return  object_batch_loss

    filtered_cell_datas = [cell_datas[i] for i in valid_indices]
    filtered_labels = labels[:,1:][valid_indices]

    # Assuming each element in filtered_cell_datas is a tensor of shape (in_channels, height, width)
    cell_images = torch.stack(filtered_cell_datas).to(device)
    cell_datas_batch = cell_images.squeeze(1)
    filtered_labels = filtered_labels.to(device)

    # Initialize the model with the dynamically determined in_channels
    # in_channels = filtered_cell_datas[0].size(1)  # Assuming the first element in filtered_cell_datas defines in_channels
    # cell_attribute_model_main.conv1.in_channels = in_channels

    # Forward pass
    outputs_my = cell_attribute_model_main(cell_datas_batch.float())
    outputs_my = outputs_my.view(len(valid_indices), -1)

    # Process labels to create target_tensor
    # label_att = filtered_labels[:, 5].float()  # Assuming label[5] contains 0 or 1
    # target_tensor = label_att.view(-1, 1)

    # Compute the loss
    num_classes = 2
    one_hot_encoded_tensors = []

    # Perform one-hot encoding for each column separately
    for i in range(filtered_labels.size(1)):
        # Extract the current column
        column_values = filtered_labels[:, i].long()

        # Generate one-hot encoded tensor for the current column
        one_hot_encoded_col = torch.eye(num_classes, device=filtered_labels.device)[column_values]

        # Reshape to match the original shape
        one_hot_encoded_col = one_hot_encoded_col.unsqueeze(1)

        one_hot_encoded_tensors.append(one_hot_encoded_col)

    # Concatenate the one-hot encoded tensors along the second dimension (axis=1)
    one_hot_encoded_result = torch.cat(one_hot_encoded_tensors, dim=1)
    outputs_my = outputs_my.view(outputs_my.size(0), 6,2)

    object_batch_loss = criterion(outputs_my, one_hot_encoded_result)

    # Check if the loss contains NaN
    if torch.isnan(object_batch_loss):
        # If NaN, trigger a breakpoint to inspect variables
        breakpoint()

    torch.use_deterministic_algorithms(False, warn_only=True)

    # Backward pass and optimization
    object_batch_loss = object_batch_loss/len(filtered_labels)
   # object_batch_loss.backward(retain_graph=True)
   # optimizer_cell_model.step()
    #scheduler_cell_model.step()

    # Explicitly release tensors
    #del cell_images, target_tensor
    #torch.cuda.empty_cache()

    return object_batch_loss