"""
EDEN Plant Disease Detection - Mob-Res Model
"""
import torch
import torch.nn as nn
from torchvision import models
import json
from pathlib import Path

class MobResModel(nn.Module):
    """
    Mob-Res: MobileNetV2 + ResNet Hybrid for Disease Detection
    """

    def __init__(self, num_classes=14, dropout=0.3):
        super(MobResModel, self).__init__()

        mobilenet = models.mobilenet_v2(weights=None)
        self.features = mobilenet.features

        in_features = 1280
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        self.fc1 = nn.Linear(in_features, 512)
        self.bn1 = nn.BatchNorm1d(512)
        self.relu1 = nn.ReLU(inplace=True)
        self.dropout1 = nn.Dropout(dropout)

        self.fc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.relu2 = nn.ReLU(inplace=True)
        self.dropout2 = nn.Dropout(dropout * 0.5)

        self.residual_adapter = nn.Linear(in_features, 256)
        self.fc3 = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)

        identity = x
        x = self.fc1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.dropout1(x)

        x = self.fc2(x)
        x = self.bn2(x)

        identity = self.residual_adapter(identity)
        x = x + identity

        x = self.relu2(x)
        x = self.dropout2(x)
        x = self.fc3(x)

        return x

def load_model(model_dir=".", device="cpu"):
    """Load the trained model"""
    with open(Path(model_dir) / "labels.json", 'r') as f:
        labels = json.load(f)
    num_classes = labels['num_classes']

    model = MobResModel(num_classes=num_classes)
    state_dict = torch.load(Path(model_dir) / "pytorch_model.bin", 
                           map_location=device, weights_only=True)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    return model, labels