import torch
import torchvision
import torchvision.models as models

from torch import nn


class ConvModel(nn.Module):
    def __init__(self, input_shape, output_shape, hidden_units):
        super().__init__()
        
        self.block_1 = nn.Sequential(
                                nn.Conv2d(in_channels = input_shape,
                                         out_channels = hidden_units,
                                         kernel_size = 3,
                                         stride = 1,
                                         padding = 1),
                                nn.ReLU(),
                                nn.Conv2d(in_channels = hidden_units,
                                        out_channels = hidden_units,
                                        kernel_size = 3,
                                        stride = 1,
                                        padding = 1),
                                nn.ReLU(),
                                nn.MaxPool2d(kernel_size = 2,
                                            stride = 2))
        
        self.block_2 = nn.Sequential(
                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
                                nn.ReLU(),
                                nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
                                nn.ReLU(),
                                nn.MaxPool2d(2))
        
        self.classifier = nn.Sequential(
                                nn.Flatten(),
                                nn.Linear(in_features = 2560, out_features = output_shape))
        
        
    def forward(self, x):
        x = self.block_1(x)
        x = self.block_2(x)
        x = self.classifier(x)
        return x
    
# class ConvModel(nn.Module):
#     def __init__(self, input_shape, output_shape, hidden_units):
#         super().__init__()
        
#         self.block_1 = nn.Sequential(
#                                 nn.Conv2d(in_channels = input_shape,
#                                          out_channels = hidden_units,
#                                          kernel_size = 3,
#                                          stride = 1,
#                                          padding = 1),
#                                 nn.ReLU(),
#                                 nn.Conv2d(in_channels = hidden_units,
#                                         out_channels = hidden_units,
#                                         kernel_size = 3,
#                                         stride = 1,
#                                         padding = 1),
#                                 nn.ReLU(),
#                                 nn.MaxPool2d(kernel_size = 2,
#                                             stride = 2),
# #                                 nn.BatchNorm2d()
#         )
        
#         self.block_2 = nn.Sequential(
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.MaxPool2d(2),
# #                                 nn.BatchNorm2d()
#         )
        
#         self.block_3 = nn.Sequential(
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.MaxPool2d(2),
# #                                 nn.BatchNorm2d()
#         )
        
#         self.block_4 = nn.Sequential(
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.Conv2d(hidden_units, hidden_units, 3, padding = 1),
#                                 nn.ReLU(),
#                                 nn.MaxPool2d(2),
# #                                 nn.BatchNorm2d()
#         )
        
#         self.classifier = nn.Sequential(
#                                 nn.Flatten(),
#                                 nn.Linear(in_features = 2560, out_features = output_shape))
        
        
#     def forward(self, x):
#         x = self.block_1(x)
#         x = self.block_2(x)
#         x = self.block_3(x)
#         x = self.block_4(x)
# #         print(x.shape)
#         x = self.classifier(x)
#         return x
    
def create_custom_model(seed:int=42):
    """Creates an EfficientNetB2 feature extractor model and transforms.

    Args:
        num_classes (int, optional): number of classes in the classifier head. 
            Defaults to 3.
        seed (int, optional): random seed value. Defaults to 42.

    Returns:
        model (torch.nn.Module): EffNetB2 feature extractor model. 
        transforms (torchvision.transforms): EffNetB2 image transforms.
    """

    model = ConvModel(input_shape=3, hidden_units=10, output_shape = 1)

    # model_path = '/kaggle/input/cnn-models/model_cnn_proj_version_2.pt'
    # model.load_state_dict(torch.load(f=model_path))

    return model

def create_resnet_model(seed:int=42):
    """Creates an EfficientNetB2 feature extractor model and transforms.

    Args:
        num_classes (int, optional): number of classes in the classifier head. 
            Defaults to 3.
        seed (int, optional): random seed value. Defaults to 42.

    Returns:
        model (torch.nn.Module): EffNetB2 feature extractor model. 
        transforms (torchvision.transforms): EffNetB2 image transforms.
    """
    # Create EffNetB2 pretrained weights, transforms and model

    model = models.resnet50(pretrained=True)
    # Freeze all layers in base model
    for param in model.parameters():
        param.requires_grad = False

    # Change classifier head with random seed for reproducibility
    torch.manual_seed(seed)

    model.fc = nn.Sequential(
        nn.Linear(2048, 512),  # Change the input size as per your model
        nn.ReLU(),
        nn.Dropout(0.5),  # Dropout layer to reduce overfitting
        nn.Linear(512, 1)  # Replace 'num_classes' with the number of your output classes
    )

    return model