import os
import torch
import torch.nn as nn
import segmentation_models_pytorch as smp
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from app.config import config

# Set cache directory to a writable location for Hugging Face Spaces
os.environ['TORCH_HOME'] = '/tmp/.cache/torch'
os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache/huggingface/transformers'

# Create cache directories if they don't exist
os.makedirs('/tmp/.cache/torch', exist_ok=True)
os.makedirs('/tmp/.cache/huggingface', exist_ok=True)

def build_model():
    """Build and return the U-Net model"""
    return smp.Unet(
        encoder_name='efficientnet-b4',
        encoder_weights='imagenet',
        in_channels=3,
        classes=1,
        activation=None,
    )

def get_transforms():
    """Get image preprocessing transforms"""
    return A.Compose([
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])

class ModelManager:
    def __init__(self):
        self.model = None
        self.transforms = None
        self.is_loaded = False

    def load_model(self):
        """Load the model and transforms"""
        try:
            self.model = build_model().to(config.DEVICE)
            state_dict = torch.load(config.MODEL_PATH, map_location=config.DEVICE)
            self.model.load_state_dict(state_dict)
            self.model.eval()
            self.transforms = get_transforms()
            self.is_loaded = True
            print(f"Model loaded successfully from {config.MODEL_PATH}")
            return True
        except Exception as e:
            print(f"Model loading failed: {e}")
            self.is_loaded = False
            return False

    def predict(self, image_np: np.ndarray):
        """Make prediction on preprocessed image"""
        if not self.is_loaded:
            raise RuntimeError("Model not loaded")

        # Prepare input tensor
        dummy_mask = np.zeros((config.IMG_HEIGHT, config.IMG_WIDTH), dtype=np.float32)
        transformed = self.transforms(image=image_np, mask=dummy_mask)
        input_tensor = transformed['image'].unsqueeze(0).to(config.DEVICE)

        # Model inference
        with torch.no_grad():
            prediction_logits = self.model(input_tensor)[0, 0].cpu().numpy()

        return prediction_logits

def postprocess_prediction(prediction: np.ndarray, threshold: float = 0.5) -> tuple:
    """Convert model output to visualization images"""
    # Apply sigmoid to get probabilities
    prob_map = 1 / (1 + np.exp(-prediction))

    # Create binary mask
    binary_mask = (prob_map > threshold).astype(np.uint8) * 255

    # Calculate confidence and area
    confidence_score = float(np.max(prob_map))
    predicted_area_percentage = float(np.sum(prob_map > threshold) / (config.IMG_HEIGHT * config.IMG_WIDTH) * 100)

    return prob_map, binary_mask, confidence_score, predicted_area_percentage

# Global model manager instance
model_manager = ModelManager()