import yaml
import cv2
import os
import torch
from .cnn_model import load_model
from .sam_utils import get_sam_generator
from .greenery_estimator import GreeneryEstimator

class EcoPulsePipeline:
    """
    Orchestration layer for the EcoPulse satellite analysis system.
    Handles the initialization of SAM and ResNet-50 models and manages
     the end-to-end processing of geographic imagery.
    """
    def __init__(self, config_path="config/config.yaml"):
        with open(config_path, "r") as f:
            self.config = yaml.safe_load(f)
            
        self.classes = self.config['classes']
        self.greenery_classes = self.config['greenery_classes']
        
        # Load models
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Loading CNN model on {device}...")
        self.cnn_model = load_model(
            weights_path=os.path.join(self.config['paths']['output_models'], 'resnet50_eurosat.pth'),
            num_classes=self.config['model']['num_classes'],
            device=device
        )
        
        print("Loading SAM model...")
        self.sam_generator = get_sam_generator(
            model_type=self.config['model']['sam_model_type'],
            checkpoint_path=self.config['paths']['sam_checkpoint']
        )
        
        self.estimator = GreeneryEstimator(
            sam_generator=self.sam_generator,
            cnn_model=self.cnn_model,
            greenery_classes=self.greenery_classes,
            all_classes=self.classes
        )
        print("Pipeline initialized successfully.")

    def process_image(self, image_path):
        """
        Runs the end-to-end pipeline on a single image.
        
        Args:
            image_path (str): File path to the input satellite image.
            
        Returns:
            tuple: (image_np, results) where image_np is the RGB image array 
                   and results is a dictionary containing greenery metrics and masks.
        """
        image_np = cv2.imread(image_path)
        if image_np is None:
            raise FileNotFoundError(f"Could not read image at {image_path} — file may be missing or corrupt")
            
        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
        
        results = self.estimator.estimate(image_np)
        return image_np, results