"""
models/cnn.py
CNN pour images RGB 3 canaux — Intel Image Classification (228×228, 6 classes).

RÈGLE DE NORMALISATION :
    La normalisation est faite UNIQUEMENT dans utils/prep.py (pipeline de données).
    Les modèles reçoivent des images déjà normalisées — il n'y a PAS de couche
    Rescaling à l'intérieur des modèles. Cela garantit un comportement identique
    entre training, evaluation et production (Flask).
"""

# ── PyTorch ───────────────────────────────────────────────────────────────────
import torch.nn as nn
import torch.nn.functional as F


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)))


# ── TensorFlow / Keras ────────────────────────────────────────────────────────
def build_cnn_tf(num_classes: int = 6, input_shape: tuple = (228, 228, 3)):
    """
    CNN TF reproduisant l'architecture du notebook de référence hassanraof.
    Source : https://www.kaggle.com/code/hassanraof/intel-image-classification

    Entrée  : (B, 228, 228, 3) — valeurs [0, 1] normalisées par prep.py
    Sortie  : (B, num_classes) — softmax

    Architecture (5 blocs conv) :
        Block 1 : Conv(32, 5×5) → ReLU → MaxPool(2,2)
        Block 2 : Conv(32, 5×5) → ReLU → MaxPool(2,2)
        Block 3 : Conv(32, 3×3) → ReLU → MaxPool(2,2)
        Block 4 : Conv(64, 3×3) → ReLU → MaxPool(2,2)
        Block 5 : Conv(64, 3×3) → ReLU → MaxPool(2,2)
        Head    : Flatten → Dense(1024) → Dropout(0.20)
                           → Dense(124)  → Dropout(0.20)
                           → Dense(num_classes, softmax)

    ⚠️  PAS de couche Rescaling ici — la normalisation est faite dans prep.py.
        Ajouter Rescaling ici causerait une double normalisation.
    """
    from tensorflow.keras import layers, models

    return models.Sequential([
        layers.Input(shape=input_shape),
        # ← PAS de Rescaling ici

        # Block 1
        layers.Conv2D(32, kernel_size=(5, 5), activation="relu"),
        layers.MaxPooling2D(2, 2),

        # Block 2
        layers.Conv2D(32, kernel_size=(5, 5), activation="relu"),
        layers.MaxPooling2D(2, 2),

        # Block 3
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        # Block 4
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        # Block 5
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),

        # Head
        layers.Flatten(),
        layers.Dense(1024, activation="relu"),
        layers.Dropout(0.20),
        layers.Dense(124, activation="relu"),
        layers.Dropout(0.20),
        layers.Dense(num_classes, activation="softmax"),

    ], name="CNN_TF_hassanraof")