"""
Model Loading Utilities
Loads all three trained models (Fundus, OCT, Multimodal Fusion)
"""

import torch
import torch.nn as nn
from torchvision import models
from pathlib import Path
import os

class FundusModel(nn.Module):
    """ResNet50-based CNN for Fundus image classification"""
    def __init__(self, num_classes=2):
        super(FundusModel, self).__init__()
        resnet = models.resnet50(weights=None)
        self.features = nn.Sequential(*list(resnet.children())[:-1])
        num_features = resnet.fc.in_features
        self.classifier = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(num_features, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes)
        )
    
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x
    
    def extract_features(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return x

class OCTCNN_Optimized(nn.Module):
    """ResNet50-based CNN for OCT image classification"""
    def __init__(self, num_classes=7):
        super(OCTCNN_Optimized, self).__init__()
        self.resnet = models.resnet50(weights=None)
        in_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Dropout(0.4),
            nn.Linear(in_features, 512),
            nn.ReLU(),
            nn.BatchNorm1d(512),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes)
        )
        
    def forward(self, x):
        return self.resnet(x)
    
    def extract_features(self, x):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)
        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)
        x = self.resnet.layer4(x)
        x = self.resnet.avgpool(x)
        x = torch.flatten(x, 1)
        return x

class MultimodalFusionModel(nn.Module):
    """Multimodal fusion model combining Fundus and OCT"""
    def __init__(self, fundus_model, oct_model, num_classes=2):
        super(MultimodalFusionModel, self).__init__()
        self.fundus_model = fundus_model
        self.oct_model = oct_model
        
        fusion_input_dim = 2048 + 2048 + 2  # Fundus + OCT + Clinical
        
        self.fusion_head = nn.Sequential(
            nn.Linear(fusion_input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )
        
    def forward(self, fundus_img, oct_img, clinical_features):
        f_feat = self.fundus_model.extract_features(fundus_img)
        o_feat = self.oct_model.extract_features(oct_img)
        combined = torch.cat((f_feat, o_feat, clinical_features), dim=1)
        return self.fusion_head(combined)

class ModelLoader:
    """Loads and manages all models"""
    def __init__(self):
        self.device = torch.device('cpu')  # Force CPU for deployment availability
        print(f"Using device: {self.device}")
        
        # Paths to model checkpoints (Modified for deployment structure)
        # Assumes models are in the SAME directory as this script
        self.models_dir = Path(__file__).parent
        
        self.fundus_path = self.models_dir / 'fundus_cnn_best.pth'
        self.oct_path = self.models_dir / 'oct_resnet50_best.pth'
        self.fusion_path = self.models_dir / 'multimodal_fusion_best.pth'
        
        # Load models
        self.fundus_model = self._load_fundus_model()
        self.oct_model = self._load_oct_model()
        self.fusion_model = self._load_fusion_model()
        
        # Class names
        self.fundus_classes = ['GON-', 'GON+']
        self.oct_classes = ['AMD', 'DME', 'ERM', 'NO', 'RAO', 'RVO', 'VID']
        self.fusion_classes = ['NORMAL', 'GLAUCOMA']
    
    def _load_fundus_model(self):
        """Load Fundus model"""
        print(f"Loading Fundus model from {self.fundus_path}")
        model = FundusModel(num_classes=2)
        
        if self.fundus_path.exists():
            checkpoint = torch.load(self.fundus_path, map_location=self.device, weights_only=False)
            if 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'], strict=False)
            else:
                model.load_state_dict(checkpoint, strict=False)
            print("✓ Fundus model loaded")
        else:
            print(f"⚠ Warning: Fundus model not found at {self.fundus_path}")
        
        model = model.to(self.device)
        model.eval()
        return model
    
    def _load_oct_model(self):
        """Load OCT model"""
        print(f"Loading OCT model from {self.oct_path}")
        model = OCTCNN_Optimized(num_classes=7)
        
        if self.oct_path.exists():
            checkpoint = torch.load(self.oct_path, map_location=self.device, weights_only=False)
            if 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'], strict=False)
            else:
                model.load_state_dict(checkpoint, strict=False)
            print("✓ OCT model loaded")
        else:
            print(f"⚠ Warning: OCT model not found at {self.oct_path}")
        
        model = model.to(self.device)
        model.eval()
        return model
    
    def _load_fusion_model(self):
        """Load Multimodal Fusion model"""
        print(f"Loading Fusion model from {self.fusion_path}")
        
        # Create base models for fusion
        fundus_base = FundusModel(num_classes=2)
        oct_base = OCTCNN_Optimized(num_classes=7)
        
        model = MultimodalFusionModel(fundus_base, oct_base, num_classes=2)
        
        if self.fusion_path.exists():
            checkpoint = torch.load(self.fusion_path, map_location=self.device, weights_only=False)
            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'], strict=False)
            else:
                model.load_state_dict(checkpoint, strict=False)
            print("✓ Fusion model loaded")
        else:
            print(f"⚠ Warning: Fusion model not found at {self.fusion_path}")
        
        model = model.to(self.device)
        model.eval()
        return model