import torch
import torch.nn as nn
import torch.nn.functional as F

class MKClassCNNModel(nn.Module):
    def __init__(self, num_spec_classes, num_lum_classes):
        super(MKClassCNNModel, self).__init__()
        
        # 1D Convolutional layers
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=9, padding=4)
        self.conv2 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=15, padding=7)
        self.pool = nn.MaxPool1d(kernel_size=4)
        #self.conv1 = nn.Conv1d(in_channels=1, out_channels=4, kernel_size=5)  #Debug values
        #self.conv2 = nn.Conv1d(in_channels=4, out_channels=8, kernel_size=3)  #Debug values
        #self.pool = nn.MaxPool1d(kernel_size=2)   #Debug values
        
        # Shared Hidden Layer (The "Neck")
        # You'll calculate this input size once, e.g., 32 * 398
        #DANGEROULSY FINICKY:
        # 390 = ((int(1569 - 4)/2) - 2)/2 - convolution with kernal5, maxpool2, then convolve kernal3, maxpool2
        #self.fc_shared = nn.Linear(32 * 390, 256)
        # 94 = ((int(1569 - 8)/4) - 14)/4 - convolution with kernal5, maxpool2, then convolve kernal3, maxpool2
        self.fc_shared = nn.Linear(32 * 98, 256)
        #self.fc_shared = nn.Linear(8 * 390, 256)    #Debug values w. MaxPool4
        #self.fc_shared = nn.Linear(8 * 97, 256)    #Debug values w. MaxPool4
        
        # The Two Heads
        self.head_spec = nn.Linear(256, num_spec_classes)
        self.head_lum = nn.Linear(256, num_lum_classes)

    def forward(self, x):
        # x starts as (Batch, 1, 1600)
        x = self.pool(F.gelu(self.conv1(x)))
        ##Unpack for debugging:
        #temp1 = self.conv1(x)
        #temp2 = F.gelu(temp1)
        #x = self.pool(temp2)
        x = self.pool(F.gelu(self.conv2(x)))
        
        # Flatten for the FC layers
        x = x.view(x.size(0), -1) 
        x = F.gelu(self.fc_shared(x))
        #Unpack for debugging:
        #temp1 = self.fc_shared(x)
        #x = F.gelu(temp1)
        
        # Return both outputs
        spec_out = self.head_spec(x)
        lum_out = self.head_lum(x)
        return spec_out, lum_out  #These are logits

"""
# During the training loop on Graham:
optimizer.zero_grad()

# Forward pass returns two vectors
pred_spec, pred_lum = model(X_batch)

# Calculate two separate losses
loss_spec = criterion(pred_spec, Y_spec_batch)
loss_lum = criterion(pred_lum, Y_lum_batch)

# Combine them (you can even weight them if one is harder to learn)
total_loss = loss_spec + loss_lum

# One call handles the entire backprop chain for BOTH heads
total_loss.backward()
optimizer.step()
"""
