import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import cv2
import numpy as np
import os

class CustomCnnModel(nn.Module):
    def __init__(self, input_dim=128, num_classes=3):
        super(CustomCnnModel, self).__init__()
        self.input_dim = input_dim
        self.num_classes = num_classes

        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        # Compute dynamically the input size for fully connected layers
        self._to_linear = None
        self._get_conv_output(self.input_dim)

        self.fc_layers = nn.Sequential(
            nn.Linear(self._to_linear, 512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, self.num_classes)
        )

    def _get_conv_output(self, input_dim=128):
        with torch.no_grad():
            dummy_input = torch.zeros(1, 3, input_dim, input_dim)
            output = self.conv_layers(dummy_input)
            self._to_linear = output.view(1, -1).size(1)

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)  # Flatten
        x = self.fc_layers(x)
        return x


class ImageClassifier:
    def __init__(self, model_path, class_name=None):
        # Default to CPU on startup to safely support ZeroGPU environments
        self.device = torch.device("cpu")
        self.model = CustomCnnModel(input_dim=128, num_classes=3)
        self.model.load_state_dict(torch.load(model_path, map_location="cpu"))
        
        # If CUDA is immediately available, move model
        if torch.cuda.is_available():
            try:
                self.device = torch.device("cuda")
                self.model.to(self.device)
            except Exception:
                self.device = torch.device("cpu")

        self.model.eval()

        if class_name is None:
            self.class_name = {0: 'Cat', 1: 'Dog', 2: 'Person'}
        else:
            self.class_name = class_name

        self.transform = transforms.Compose([
            transforms.Resize((128, 128)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        ])

    def predict(self, image_input):
        """
        Accepts either a file path string or a PIL Image/numpy array.
        Returns:
            label (str): Name of predicted class
            labeled_image: Annotated PIL Image (or file path if a string path was supplied)
        """
        is_path_input = isinstance(image_input, str)
        if is_path_input:
            image = Image.open(image_input).convert("RGB")
        elif isinstance(image_input, Image.Image):
            image = image_input.convert("RGB")
        elif isinstance(image_input, np.ndarray):
            image = Image.fromarray(image_input).convert("RGB")
        else:
            raise ValueError(f"Unsupported image input type: {type(image_input)}")

        # Dynamically attach to GPU if active (supports ZeroGPU during @spaces.GPU calls)
        current_cuda_available = torch.cuda.is_available()
        target_device = torch.device("cuda" if current_cuda_available else "cpu")
        if self.device != target_device:
            try:
                self.model.to(target_device)
                self.device = target_device
            except Exception:
                self.device = torch.device("cpu")
                self.model.to(self.device)

        image_tensor = self.transform(image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            output = self.model(image_tensor)
            _, predicted = torch.max(output, 1)

        label = self.class_name.get(predicted.item(), f"Class {predicted.item()}")

        # Annotate image with the predicted label
        img_np = np.array(image)
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

        h, w = img_bgr.shape[:2]
        font_scale = max(0.6, min(w, h) / 300.0)
        thickness = max(1, int(min(w, h) / 150.0))
        font = cv2.FONT_HERSHEY_SIMPLEX

        (text_w, text_h), baseline = cv2.getTextSize(label, font, font_scale, thickness)
        
        box_padding = 8
        cv2.rectangle(
            img_bgr,
            (10, 10),
            (10 + text_w + 2 * box_padding, 10 + text_h + 2 * box_padding),
            (0, 0, 0),
            -1
        )
        cv2.putText(
            img_bgr,
            label,
            (10 + box_padding, 10 + text_h + box_padding),
            font,
            font_scale,
            (0, 255, 0),
            thickness
        )

        annotated_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        labeled_image = Image.fromarray(annotated_rgb)

        if is_path_input:
            output_path = os.path.abspath("labeled_image.jpg")
            labeled_image.save(output_path)
            return label, output_path

        return label, labeled_image