import os
import cv2
import torch
import torchvision
from torch import nn
from torchvision import transforms
from PIL import Image

class_names = ["Glioma Tumor", "Meningitis Tumor", "No Tumor", "Pituitary Tumor"]

class CFG:
    DEVICE = 'cpu'
    NUM_DEVICES = torch.cuda.device_count()
    NUM_WORKERS = os.cpu_count()
    NUM_CLASSES = 4
    EPOCHS = 16
    BATCH_SIZE = 32
    LR = 0.001
    APPLY_SHUFFLE = True
    SEED = 768
    HEIGHT = 224
    WIDTH = 224
    CHANNELS = 3
    IMAGE_SIZE = (224, 224, 3)


class EfficientNetV2Model(nn.Module):
    def __init__(self, backbone_model, name='efficientnet-v2-large',
                 num_classes=CFG.NUM_CLASSES, device=CFG.DEVICE):
        super(EfficientNetV2Model, self).__init__()

        self.backbone_model = backbone_model
        self.device = device
        self.num_classes = num_classes
        self.name = name

        classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=0.2, inplace=True),
            nn.Linear(in_features=1280, out_features=256, bias=True),
            nn.GELU(),
            nn.Dropout(p=0.2, inplace=True),
            nn.Linear(in_features=256, out_features=num_classes, bias=False)
        ).to(device)

        self._set_classifier(classifier)

    def _set_classifier(self, classifier:nn.Module) -> None:
        self.backbone_model.classifier = classifier

    def forward(self, image):
        return self.backbone_model(image)


def get_effiecientnetv2_model(
    device: torch.device=CFG.NUM_CLASSES) -> nn.Module:
    # Set the manual seeds
    torch.manual_seed(CFG.SEED)
    torch.cuda.manual_seed(CFG.SEED)

    # Get model weights
    model_weights = (
        torchvision
        .models
        .EfficientNet_V2_L_Weights
        .DEFAULT
    )

    # Get model and push to device
    model = (
        torchvision.models.efficientnet_v2_l(
            weights=model_weights
        )
    ).to(device)

    # Freeze Model Parameters
    for param in model.features.parameters():
        param.requires_grad = False

    return model


# Get EfficientNet v2 model
backbone_model = get_effiecientnetv2_model(CFG.DEVICE)

efficientnetv2_params = {
    'backbone_model'    : backbone_model,
    'name'              : 'efficientnet-v2-large',
    'device'            : CFG.DEVICE
}

# Generate Model
efficientnet_model = EfficientNetV2Model(**efficientnetv2_params)

efficientnet_model.load_state_dict(
    torch.load('efficientnetV2.pth', map_location=torch.device('cpu'))
)


# def predict_eff(image_path):
#     # Define the image transformation
#     transform = transforms.Compose([
#         transforms.Resize((224, 224)),
#         transforms.ToTensor()
#     ])

#     # Load and preprocess the image
#     # image_path = "glioma.jpg"  # Replace with the path to your image
#     image = Image.open(image_path)
#     input_tensor = transform(image)
#     input_batch = input_tensor.unsqueeze(0).to("cuda")  # Add batch dimension

#     # Perform inference
#     with torch.no_grad():
#         output = efficientnet_model(input_batch).to("cuda")

#     # You can now use the 'output' tensor as needed (e.g., get predictions)
#     # print(output)
#     return output


def predict_eff(image_path):
    # Define the image transformation
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])

    # Load and preprocess the image
    # image_path = "glioma.jpg"  # Replace with the path to your image
    image = Image.open(image_path)
    input_tensor = transform(image)
    input_batch = input_tensor.unsqueeze(0).to(CFG.DEVICE)  # Add batch dimension

    # Perform inference
    with torch.no_grad():
        output = efficientnet_model(input_batch).to(CFG.DEVICE)

    # You can now use the 'output' tensor as needed (e.g., get predictions)
    # print(output)
    res = torch.softmax(output, dim=1)

    probs = {class_names[i]: float(res[0][i]) for i in range(len(class_names))}

    return probs