import argparse
import yaml
import os
import numpy as np
import cv2
from src.data_loader import DeepGlobeDataset
from src.pipeline import EcoPulsePipeline
from src.metrics import calculate_iou, calculate_dice
from tqdm import tqdm

def evaluate_deepglobe(config_path):
    print("Initializing Pipeline...")
    pipeline = EcoPulsePipeline(config_path)
    config = pipeline.config
    
    deepglobe_dir = config['paths']['deepglobe_dir']
    if not os.path.exists(os.path.join(deepglobe_dir, 'images')):
        print("DeepGlobe dataset not found. Please run download_datasets.py first.")
        return

    dataset = DeepGlobeDataset(deepglobe_dir)
    print(f"Found {len(dataset)} images in DeepGlobe validation set.")
    
    ious = []
    dices = []
    
    # We evaluate on a small subset for demonstration purposes
    subset_size = min(10, len(dataset))
    
    print("Evaluating Segmentation Accuracy on DeepGlobe subset...")
    for i in tqdm(range(subset_size)):
        img_name = dataset.image_files[i]
        img_path = os.path.join(dataset.images_dir, img_name)
        
        _, results = pipeline.process_image(img_path)
        
        # Get image dimensions from the processed image directly
        raw_img = cv2.imread(img_path)
        if raw_img is None:
            raise FileNotFoundError(f"Could not read image at {img_path} — file may be missing or corrupt")
        h, w = raw_img.shape[:2]
        
        # Build predicted green mask
        pred_mask = np.zeros((h, w), dtype=bool)
        for item in results['mask_classifications']:
            if item['is_green']:
                pred_mask = np.logical_or(pred_mask, item['segmentation'])
                
        # Build true green mask from DeepGlobe color-coded annotations
        _, true_mask_rgb = dataset[i]
        true_mask_rgb = np.array(true_mask_rgb)
        
        # Flatten condition for greenery in DeepGlobe
        true_green = ((true_mask_rgb[:,:,1] == 255) & (true_mask_rgb[:,:,0] == 0) & (true_mask_rgb[:,:,2] == 0)) | \
                     ((true_mask_rgb[:,:,0] == 255) & (true_mask_rgb[:,:,1] == 255) & (true_mask_rgb[:,:,2] == 0)) | \
                     ((true_mask_rgb[:,:,0] == 255) & (true_mask_rgb[:,:,1] == 0) & (true_mask_rgb[:,:,2] == 255))
        
        iou = calculate_iou(pred_mask, true_green)
        dice = calculate_dice(pred_mask, true_green)
        
        ious.append(iou)
        dices.append(dice)
        
    print(f"Mean IoU (Greenery): {np.mean(ious):.4f}")
    print(f"Mean Dice Score: {np.mean(dices):.4f}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config/config.yaml")
    args = parser.parse_args()
    evaluate_deepglobe(args.config)