import torch
import torch.nn as nn
import torchvision.models as models
from huggingface_hub import hf_hub_download
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Global variables to cache models
_wbc_model = None
_skin_processor = None
_skin_model = None

# Custom layers for WBC Model (from fastai)
class AdaptiveConcatPool2d(nn.Module):
    def __init__(self, sz=None):
        super().__init__()
        self.ap = nn.AdaptiveAvgPool2d(sz or 1)
        self.mp = nn.AdaptiveMaxPool2d(sz or 1)
    def forward(self, x):
        return torch.cat([self.ap(x), self.mp(x)], 1)

def get_wbc_model_architecture():
    # Instantiate standard resnet18 architecture
    backbone = models.resnet18()
    backbone.avgpool = nn.Identity()
    backbone.fc = nn.Identity()
    
    # Sequential layout to match the fastai sequential wrappers
    backbone_seq = nn.Sequential(
        backbone.conv1,
        backbone.bn1,
        backbone.relu,
        backbone.maxpool,
        backbone.layer1,
        backbone.layer2,
        backbone.layer3,
        backbone.layer4
    )
    
    # Custom head structure
    head_seq = nn.Sequential(
        AdaptiveConcatPool2d(),
        nn.Flatten(),
        nn.BatchNorm1d(1024),
        nn.Dropout(0.25),
        nn.Linear(1024, 512, bias=False),
        nn.ReLU(inplace=True),
        nn.BatchNorm1d(512),
        nn.Dropout(0.5),
        nn.Linear(512, 8, bias=True)  # Modified below if state_dict does not contain bias
    )
    
    model = nn.Sequential(backbone_seq, head_seq)
    return model

def load_wbc_model():
    global _wbc_model
    if _wbc_model is not None:
        return _wbc_model
        
    # Check for custom local keras model
    import os
    local_keras_path = os.path.join(os.path.dirname(__file__), "model", "my_model.keras")
    if os.path.exists(local_keras_path):
        print(f"Loading custom Keras WBC model from {local_keras_path}...")
        import tensorflow as tf
        model = tf.keras.models.load_model(local_keras_path)
        _wbc_model = {"framework": "tensorflow", "model": model}
        print("Custom Keras WBC Model loaded successfully.")
        return _wbc_model

    print("Loading fallback PyTorch WBC Blood Cell Model (ResNet-18)...")
    model = get_wbc_model_architecture()
    
    # Download weights
    weights_path = hf_hub_download(repo_id="esab/pbc-cell-classifier", filename="cell_classifier_weights.pth")
    state_dict = torch.load(weights_path, map_location="cpu")
    
    # Adjust output layer bias if not present in weights
    if "1.8.bias" not in state_dict:
        model[1][8] = nn.Linear(512, 8, bias=False)
        
    model.load_state_dict(state_dict)
    model.eval()
    _wbc_model = {"framework": "pytorch", "model": model}
    print("WBC Model loaded successfully.")
    return _wbc_model

def load_skin_model():
    global _skin_processor, _skin_model
    if _skin_model is not None and _skin_processor is not None:
        return _skin_processor, _skin_model
        
    print("Loading Skin Cancer Model (Vision Transformer)...")
    model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
    
    try:
        _skin_processor = AutoImageProcessor.from_pretrained(model_name)
        _skin_model = AutoModelForImageClassification.from_pretrained(model_name, attn_implementation="eager")
    except Exception as e:
        print(f"Network error preloading ViT ({str(e)}). Attempting to load from local cache...")
        try:
            _skin_processor = AutoImageProcessor.from_pretrained(model_name, local_files_only=True)
            _skin_model = AutoModelForImageClassification.from_pretrained(model_name, attn_implementation="eager", local_files_only=True)
        except Exception as e_local:
            print(f"Failed to load from local Hugging Face cache: {str(e_local)}")
            raise e_local
            
    _skin_model.eval()
    print("Skin Cancer Model loaded successfully.")
    return _skin_processor, _skin_model

def preload_all_models():
    load_wbc_model()
    load_skin_model()