"""
Intel Scene Classifier — Flask App
"""

import io
import os
import urllib.request
import numpy as np
from flask import Flask, jsonify, render_template, request
from PIL import Image
import tensorflow as tf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

app = Flask(__name__)

CLASSES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
PYTORCH_IMG_SIZE = 150
TF_IMG_SIZE = 228

_pytorch_model = None
_tf_model = None


class CNN_Torch(nn.Module):
    """
    CNN PyTorch 4 blocs pour images RGB (3 canaux, 150×150).
    Entrée  : (B, 3, 150, 150) — normalisée ImageNet (mean/std)
    Sortie  : (B, num_classes) — logits bruts (CrossEntropyLoss)

    Architecture :
        Block 1 : Conv(3→32)×2   + BN + ReLU + MaxPool(2)            150→75
        Block 2 : Conv(32→64)×2  + BN + ReLU + MaxPool(2) + Drop2d   75→37
        Block 3 : Conv(64→128)×2 + BN + ReLU + MaxPool(2) + Drop2d   37→18
        Block 4 : Conv(128→256)×2+ BN + ReLU + MaxPool(2) + Drop2d   18→9
        GAP     : AdaptiveAvgPool2d(1)                                  →(B,256)
        Head    : Linear(256→256) + ReLU + Dropout + Linear(256→C)
    """
    def __init__(self, num_classes: int = 6):
        super().__init__()

        self.features = nn.Sequential(
            # Block 1 — 150×150 → 75×75
            nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(32), nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            # Block 2 — 75×75 → 37×37
            nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
            nn.MaxPool2d(2), nn.Dropout2d(0.10),

            # Block 3 — 37×37 → 18×18
            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.MaxPool2d(2), nn.Dropout2d(0.15),

            # Block 4 — 18×18 → 9×9
            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
            nn.MaxPool2d(2), nn.Dropout2d(0.20),
        )

        # (B,256,9,9) → (B,256,1,1) → (B,256)
        self.gap = nn.AdaptiveAvgPool2d(1)

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.30),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.gap(self.features(x)))


def load_pytorch():
    global _pytorch_model
    if _pytorch_model is not None:
        return _pytorch_model

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CNN_Torch(num_classes=6).to(device)

    state_dict = torch.load("parfait_model.pth", map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

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

    _pytorch_model = (model, device, tf_transform)
    return _pytorch_model


def load_tensorflow():
    global _tf_model
    if _tf_model is None:
        import tensorflow as tf
        _tf_model = tf.keras.models.load_model("parfait_model.keras", compile=False)
    return _tf_model
    

def read_input_image():
    if "image" in request.files and request.files["image"].filename:
        return Image.open(io.BytesIO(request.files["image"].read())).convert("RGB")

    image_url = request.form.get("image_url", "").strip()
    if image_url:
        with urllib.request.urlopen(image_url) as response:
            return Image.open(io.BytesIO(response.read())).convert("RGB")

    raise ValueError("No image provided")


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/predict", methods=["POST"])
def predict():
    framework = request.form.get("model", "pytorch")

    try:
        pil_img = read_input_image()
    except Exception:
        return jsonify({"error": "Fichier image invalide"}), 400

    try:
        if framework == "pytorch":
            model, device, tf_transform = load_pytorch()
            tensor = tf_transform(pil_img).unsqueeze(0).to(device)

            with torch.no_grad():
                out = model(tensor)
                probs = torch.softmax(out, dim=1).cpu().numpy()[0]
                #probs = torch.exp(out).cpu().numpy()[0]

        elif framework == "tensorflow":
            model = load_tensorflow()
            arr = np.array(
                pil_img.resize((TF_IMG_SIZE, TF_IMG_SIZE)),
                dtype=np.float32
            ) / 255.0
            arr = np.expand_dims(arr, axis=0)
            probs = model.predict(arr, verbose=0)[0]

        else:
            return jsonify({"error": "Framework non supporté"}), 400

        pred_idx = int(np.argmax(probs))
        return jsonify({
            "class": CLASSES[pred_idx],
            #"confidence": float(probs[pred_idx]),
            "confidence": float(probs[pred_idx]),
            "probabilities": {
                c: float(p) for c, p in zip(CLASSES, probs)
            },
        })

    except FileNotFoundError as e:
        return jsonify({
            "error": f"Modèle introuvable : {e}. Placez les fichiers .pth et .keras à la racine."
        }), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port, debug=False)