import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models
import time

# Base model
from architectures.Model import Model

class Inception(Model):
    def get(self):
        model = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
        #for param in model.parameters():
        #    param.requires_grad = False
        model.fc = nn.Linear(model.fc.in_features, self.size)
        model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, self.size)
        return model.to(self.device)

    def train(self, epochs, loader, rate):
        # Override because Inception returns a tuple (main, aux)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)

        print(f"Starting training on {self.device}...")
        start_time = time.time()

        self.model.train()

        for epoch in range(epochs):
            total_loss = 0.0
            for inputs, labels in loader:
                
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                optimizer.zero_grad()
                
                outputs, aux_outputs = self.model(inputs)
                loss = criterion(outputs, labels) + 0.3 * criterion(aux_outputs, labels)
                
                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}")
        
        if self.device.type == 'cuda': torch.cuda.synchronize()
        print(f"Training completed in: {time.time() - start_time:.2f}s")