from torch.utils.data import DataLoader
from torchvision import transforms
import numpy as np
import pandas as pd
import os
import cv2
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch
import torch.nn as nn


class HybridCNNViT(nn.Module):
    def __init__(self, in_channels: int, num_classes: int):
        super(HybridCNNViT, self).__init__()

        self.conv1 = nn.Conv2d(
            in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(64, 128, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(128)

        self.conv3 = nn.Conv2d(128, 128, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(128)

        self.conv4 = nn.Conv2d(128, 256, kernel_size=3,
                               stride=2, padding=1, bias=False)
        self.bn4 = nn.BatchNorm2d(256)

        self.conv5 = nn.Conv2d(256, 256, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn5 = nn.BatchNorm2d(256)

        self.conv6 = nn.Conv2d(256, 512, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn6 = nn.BatchNorm2d(512)

        self.conv7 = nn.Conv2d(512, 512, kernel_size=3,
                               stride=2, padding=1, bias=False)
        self.bn7 = nn.BatchNorm2d(512)

        # Optional MaxPooling (can be removed if strictly no max pooling)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)

        self.classifier_conv = nn.Conv2d(
            512, num_classes, kernel_size=1, stride=1, padding=0, bias=False)

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Dropout(0.5)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))
        x = self.relu(self.bn4(self.conv4(x)))
        x = self.relu(self.bn5(self.conv5(x)))
        x = self.relu(self.bn6(self.conv6(x)))
        x = self.relu(self.bn7(self.conv7(x)))

        x = self.maxpool(x)  # Comment this line if no max pooling is needed

        x = self.classifier_conv(x)
        x = self.classifier(x)

        return x


def load_and_pad_single_image(image_path, img_size=(224, 224)):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"Could not read image: {image_path}")
    img = cv2.resize(img, img_size)
    return np.array(img)


def check_file(image_path):
    # image_path = "d/Control-Axial/C-A (2).png"

    # Load and preprocess the single image
    image = load_and_pad_single_image(image_path)
    image = np.expand_dims(image, axis=0)  # Convert to batch format

    # Duplicate the image 10 times
    data = np.repeat(image, 10, axis=0)

    # Normalize and transform the image
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
                             0.229, 0.224, 0.225])
    ])

    data = torch.tensor(data, dtype=torch.float32).permute(
        0, 3, 1, 2).to(device)
    # Placeholder labels for 10 images
    labels = torch.tensor([0] * 10, dtype=torch.long).to(device)

    data, labels = shuffle(data, labels, random_state=42)

    train_data, test_data, train_labels, test_labels = train_test_split(
        data, labels, test_size=0.2, random_state=42
    )

    train_labels = torch.tensor(train_labels, dtype=torch.long)
    test_labels = torch.tensor(test_labels, dtype=torch.long)

    batch_size = 1  # Since we are working with a single image
    train_dataset = list(zip(train_data, train_labels))
    test_dataset = list(zip(test_data, test_labels))
    test_loader = DataLoader(
        test_dataset, batch_size=batch_size, shuffle=False)

    # Simple test with a model
    output = ""

    def test_model(model, test_loader, device):
        global output
        model.to(device)
        model.eval()
        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                output = predicted
                # Convert logits to probabilities
                probabilities = F.softmax(outputs, dim=1)
                # Get confidence score and prediction
                confidence, d = torch.max(probabilities, 1)
                print(confidence)
                return predicted, confidence

    def remove_module_from_checkpoint(checkpoint):
        new_state_dict = {}
        for key, value in checkpoint["model_state_dict"].items():
            new_key = key.replace("module.", "")
            new_state_dict[new_key] = value
        checkpoint["model_state_dict"] = new_state_dict
        return checkpoint

    model = HybridCNNViT(3, 2)
    checkpoint = torch.load(
        "/home/user/app/checkpoint32.pth", weights_only=False, map_location=torch.device('cpu'))
    checkpoint = remove_module_from_checkpoint(checkpoint)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    model.to(device)
    model = nn.DataParallel(model)
    output, confidence = test_model(model, test_loader, device)
    return "No ms detected" if output.item() == 0 else "MS Detected", confidence.item()