from abc import ABC, abstractmethod
import torch
import torch.nn as nn
import torch.optim as optim
import time
import numpy as np
from sklearn.metrics import classification_report

class Model(ABC):
    def __init__(self, device, size):
        self.device = device
        self.size = size
        self.model = self.get()

    @abstractmethod
    def get(self):
        # return the model
        return self.model

    def train(self, epochs, loader, rate):
        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 = self.model(inputs)
                loss = criterion(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")

    def evaluate(self, loader):
        self.model.eval()
        all_preds, all_labels = [], []
        print("\nEvaluating...")
        
        with torch.no_grad():
            for inputs, labels in loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                outputs = self.model(inputs)
                _, predicted = torch.max(outputs, 1)
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())

        accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels)
        print(f"Test Accuracy: {accuracy:.2f}%")
        print(classification_report(all_labels, all_preds, zero_division=0))


    # Using the factory patern here
    @staticmethod
    def create(arch, device, size):
        print(f'>> MODEL ARCHITECTURE  >> {arch.name}.')

        match arch:

            # ResNet18
            case Architecture.RESNET18: 
                from architectures.ResNet18 import ResNet18
                return ResNet18(device, size)
            
            # ResNet50
            case Architecture.RESNET50:
                from architectures.ResNet18 import ResNet18
                return ResNet18(device, size)
            
            # INCEPTION
            case Architecture.INCEPTION:
                from architectures.Inception import Inception
                return Inception(device, size)
            
            # DENSENET121
            case Architecture.DENSENET121:
                from architectures.DenseNet121 import DenseNet121
                return DenseNet121(device, size)
            case _:
                raise ValueError(f"Unknown model: {arch}")

 
# model architectures
from enum import Enum, auto

class Architecture(Enum):
    RESNET18 = auto()
    RESNET50 = auto()
    INCEPTION = auto()
    DENSENET121 = auto()