import torch
import torch.nn as nn
import segmentation_models_pytorch as smp
from segmentation_models_pytorch import utils

from SemanticModel.encoder_management import initialize_encoder
from SemanticModel.custom_losses import FocalLossFunction, TverskyLossFunction, EnhancedCrossEntropy
from SemanticModel.image_preprocessing import get_preprocessing_pipeline

class SegmentationModel:
    def __init__(self, classes=['background', 'foreground'], architecture='unet',
                 encoder='timm-regnety_120', weights='imagenet', loss=None):
        self._initialize_classes(classes)
        self.architecture = architecture
        self.encoder = encoder
        self.weights = weights
        self._setup_loss_function(loss)
        self._initialize_model()

    def _initialize_classes(self, classes):
        """Sets up class configuration."""
        if len(classes) <= 2:
            self.classes = [c for c in classes if c.lower() != 'background']
            self.class_values = [i for i, c in enumerate(classes) if c.lower() != 'background']
            self.background_flag = 'background' in classes
        else:
            self.classes = classes
            self.class_values = list(range(len(classes)))
            self.background_flag = False
        self.n_classes = len(self.classes)

    def _setup_loss_function(self, loss):
        """Configures model's loss function."""
        if not loss:
            loss = 'bce_with_logits' if self.n_classes > 1 else 'dice'
        
        if loss.lower() not in ['dice', 'bce_with_logits', 'focal', 'tversky']:
            print(f'Invalid loss: {loss}, defaulting to dice')
            loss = 'dice'

        loss_configs = {
            'bce_with_logits': {
                'activation': None,
                'loss': EnhancedCrossEntropy() if self.n_classes > 1 else utils.losses.BCEWithLogitsLoss()
            },
            'dice': {
                'activation': 'softmax' if self.n_classes > 1 else 'sigmoid',
                'loss': utils.losses.DiceLoss()
            },
            'focal': {
                'activation': None,
                'loss': FocalLossFunction()
            },
            'tversky': {
                'activation': None,
                'loss': TverskyLossFunction()
            }
        }

        config = loss_configs[loss.lower()]
        self.activation = config['activation']
        self.loss = config['loss']
        self.loss_name = loss

    def _initialize_model(self):
        """Initializes the segmentation model architecture."""
        if self.weights.endswith('pth'):
            self._load_pretrained_model()
        else:
            self._create_new_model()

    def _load_pretrained_model(self):
        """Loads model from pretrained weights."""
        print('Loading pretrained model...')
        self.model = torch.load(self.weights)
        if isinstance(self.model, torch.nn.DataParallel):
            self.model = self.model.module

        try:
            preprocessing_fn = smp.encoders.get_preprocessing_fn(self.encoder, 'imagenet')
            self.preprocessing = get_preprocessing_pipeline(preprocessing_fn)
        except:
            print('Failed to configure preprocessing. Setting to None.')
            self.preprocessing = None

    def _create_new_model(self):
        """Creates new model with specified architecture."""
        preprocessing_fn = smp.encoders.get_preprocessing_fn(self.encoder, 'imagenet')
        self.preprocessing = get_preprocessing_pipeline(preprocessing_fn)
        initialize_encoder(name=self.encoder, weights=self.weights)

        architectures = {
            'unet': smp.Unet,
            'unet++': smp.UnetPlusPlus,
            'deeplabv3': smp.DeepLabV3,
            'deeplabv3+': smp.DeepLabV3Plus,
            'fpn': smp.FPN,
            'linknet': smp.Linknet,
            'manet': smp.MAnet,
            'pan': smp.PAN,
            'pspnet': smp.PSPNet
        }

        if self.architecture not in architectures:
            raise ValueError(f'Unsupported architecture: {self.architecture}')

        self.model = architectures[self.architecture](
            encoder_name=self.encoder,
            encoder_weights=self.weights,
            classes=self.n_classes,
            activation=self.activation
        )

    @property
    def config_data(self):
        """Returns model configuration data."""
        return {
            'architecture': self.architecture,
            'encoder': self.encoder,
            'weights': self.weights,
            'activation': self.activation,
            'loss': self.loss_name,
            'classes': ['background'] + self.classes if self.background_flag else self.classes
        }

def list_architectures():
    """Returns available architecture options."""
    return ['unet', 'unet++', 'deeplabv3', 'deeplabv3+', 'fpn', 
            'linknet', 'manet', 'pan', 'pspnet']