import torch
import numpy as np
from PIL import Image

from src.transforms import EUROSAT_TRANSFORM

class GreeneryEstimator:
    """
    Core engine for quantifying vegetation in satellite imagery.
    Integrates Instance Segmentation (SAM) and Land-Cover Classification (CNN)
     to identify and measure the pixel footprint of greenery.
    """
    def __init__(self, sam_generator, cnn_model, greenery_classes, all_classes, device="cuda"):
        self.sam_generator = sam_generator
        self.cnn_model = cnn_model
        self.greenery_classes = greenery_classes
        self.all_classes = all_classes
        self.device = device if torch.cuda.is_available() else "cpu"
        
        self.cnn_model.to(self.device)
        self.cnn_model.eval()

        self.transform = EUROSAT_TRANSFORM

    def estimate(self, image_np):
        """
        Segments an image, classifies individual components, and calculates vegetation density.

        Args:
            image_np (np.ndarray): The input image in RGB format.

        Returns:
            dict: A result dictionary containing:
                - 'greenery_percentage' (float): Total vegetated area percentage.
                - 'green_pixels' (int): Total count of vegetated pixels.
                - 'total_pixels' (int): Total image resolution.
                - 'mask_classifications' (list): Detailed metadata for every segmented object.
        """
        # 1. Generate SAM masks
        masks = self.sam_generator.generate(image_np)
        
        total_pixels = image_np.shape[0] * image_np.shape[1]
        green_pixels = 0
        mask_classifications = []
        
        # 2. Iterate and classify each mask
        for idx, mask_data in enumerate(masks):
            bbox = mask_data['bbox'] # [x, y, w, h]
            x, y, w, h = [int(v) for v in bbox]
            h_img, w_img = image_np.shape[:2]

            # Clamp bounding box coordinates to image boundaries
            # to prevent silently wrong crop sizes from out-of-bounds masks
            x = max(0, min(x, w_img - 1))
            y = max(0, min(y, h_img - 1))
            w = min(w, w_img - x)
            h = min(h, h_img - y)

            # Avoid completely empty or tiny boxes
            if w <= 0 or h <= 0:
                continue

            # Crop the original image using the bounding box
            crop_np = image_np[y:y+h, x:x+w]
            if crop_np.size == 0:
                continue
                
            crop_img = Image.fromarray(crop_np)
            input_tensor = self.transform(crop_img).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                output = self.cnn_model(input_tensor)
                _, predicted = torch.max(output, 1)
                predicted_class = self.all_classes[predicted.item()]
            
            # Check if classified as greenery
            is_green = predicted_class in self.greenery_classes
            
            # Count pixels strictly within the boolean mask, not just the bounding box
            segmentation = mask_data['segmentation']
            mask_pixels = np.sum(segmentation)
            
            if is_green:
                green_pixels += mask_pixels
                
            mask_classifications.append({
                'mask_id': idx,
                'class': predicted_class,
                'is_green': is_green,
                'pixels': mask_pixels,
                'segmentation': segmentation,
                'bbox': bbox
            })

        greenery_percentage = (green_pixels / total_pixels) * 100
        
        return {
            'greenery_percentage': greenery_percentage,
            'green_pixels': green_pixels,
            'total_pixels': total_pixels,
            'mask_classifications': mask_classifications
        }