import sys
from pathlib import Path
import numpy as np
import requests
from io import BytesIO

from PIL import Image

from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status

from notifications.serializers import ImageClassifierSerializer, ClassificationResultSerializer

import torch
import torch.nn as nn
from torchvision import transforms

import tensorflow as tf


# PATHS 

BASE_DIR = Path(__file__).resolve().parent.parent
PYTORCH_PATH = BASE_DIR / "models" / "pytorch_model.pth"
TENSORFLOW_PATH = BASE_DIR / "models" / "model_best.keras"

CLASSES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]

CONFIDENCE_THRESHOLD = 0.6  


#  PYTORCH MODEL 

class CNN(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()

        self.block1 = self._block(3, 32)
        self.block2 = self._block(32, 64)
        self.block3 = self._block(64, 128)
        self.block4 = self._block(128, 256)

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Linear(256, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.5)

    def _block(self, in_c, out_c):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(),
            nn.Conv2d(out_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)

        x = self.gap(x)
        x = x.view(x.size(0), -1)

        x = self.dropout(torch.relu(self.fc1(x)))
        x = self.fc2(x)
        return x


#  TRANSFORM 

pytorch_transform = transforms.Compose([
    transforms.Resize((150, 150)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])




_pytorch_model = None
_tensorflow_model = None


def get_pytorch_model():
    global _pytorch_model
    if _pytorch_model is None:
        model = CNN(num_classes=len(CLASSES))
        model.load_state_dict(torch.load(str(PYTORCH_PATH), map_location="cpu"))
        model.eval()
        _pytorch_model = model
    return _pytorch_model


def get_tensorflow_model():
    global _tensorflow_model
    if _tensorflow_model is None:
        _tensorflow_model = tf.keras.models.load_model(str(TENSORFLOW_PATH), compile=False)
    return _tensorflow_model



class ImageClassifierAPIView(APIView):

    def post(self, request):
        serializer = ImageClassifierSerializer(data=request.data)

        if not serializer.is_valid():
            return Response(serializer.errors, status=400)

        try:
           
            if "image" in serializer.validated_data:
                image = Image.open(serializer.validated_data["image"]).convert("RGB")

            else:
                image_url = serializer.validated_data["image_url"]

                if not image_url.startswith("http"):
                    return Response({"error": "Invalid URL"}, status=400)

                resp = requests.get(image_url, timeout=10)
                resp.raise_for_status()

                image = Image.open(BytesIO(resp.content)).convert("RGB")

            model_name = serializer.validated_data.get("model", "pytorch")

            if model_name == "pytorch":
                result = self._predict_pytorch(image)
            else:
                result = self._predict_tensorflow(image)

            result["model_used"] = model_name

            return Response(result, status=200)

        except Exception as e:
            print("ERROR:", e)
            return Response({"error": str(e)}, status=400)



    def _predict_pytorch(self, image):
        model = get_pytorch_model()
        tensor = pytorch_transform(image).unsqueeze(0)

        with torch.no_grad():
            outputs = model(tensor)
            probs = torch.nn.functional.softmax(outputs, dim=1)

        values, indices = torch.topk(probs, k=len(CLASSES), dim=1)

        values = values.squeeze().cpu().numpy()
        indices = indices.squeeze().cpu().numpy()

        confidence = float(values[0])

        all_probs = [
            {"class": CLASSES[indices[i]], "probability": float(values[i])}
            for i in range(len(CLASSES))
        ]

  
        if confidence < CONFIDENCE_THRESHOLD:
            return {
                "predicted_class": "unknown",
                "confidence": confidence,
                "all_probabilities": all_probs
            }

        return {
            "predicted_class": CLASSES[indices[0]],
            "confidence": confidence,
            "all_probabilities": all_probs
        }

    #  TENSORFLOW 

    def _predict_tensorflow(self, image):

        model = get_tensorflow_model()

        img = image.resize((130, 130)) 
        
        arr = np.array(img, dtype=np.float32)  
        arr = arr[:, :, ::-1]         
        arr = np.expand_dims(arr, 0) 
        

        preds = model.predict(arr, verbose=0)[0]
        sorted_idx = np.argsort(preds)[::-1]

        confidence = float(preds[sorted_idx[0]])

        all_probs = [
        {"class": CLASSES[i], "probability": float(preds[i])}
        for i in sorted_idx
        ]

        if confidence < CONFIDENCE_THRESHOLD:

            return {
            "predicted_class": "unknown",
            "confidence": confidence,
            "all_probabilities": all_probs
            }

        return {
        "predicted_class": CLASSES[sorted_idx[0]],
        "confidence": confidence,
        "all_probabilities": all_probs
    }