"""
Simple test script - Test model on sample images without masks
Phiên bản đơn giản - test mô hình trên ảnh mẫu mà không cần mask
"""

import os
import argparse
from pathlib import Path
import numpy as np
from PIL import Image
import json
from tqdm import tqdm

import torch
import torch.nn.functional as F
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

class SimpleSegmentationTester:
    def __init__(self, model_path, device="auto"):
        self.device = torch.device("cuda" if device == "auto" and torch.cuda.is_available() else "cpu")
        
        print(f"🖥️  Device: {self.device}")
        print(f"📁 Loading model from: {model_path}")
        
        try:
            # Load model
            self.model = SegformerForSemanticSegmentation.from_pretrained(model_path)
            self.model.to(self.device)
            self.model.eval()
            
            # Create default processor (from nvidia/segformer-b0-finetuned-cityscapes-1024-1024)
            self.processor = SegformerImageProcessor(
                do_resize=True,
                size={"height": 512, "width": 512},
                do_normalize=True,
                image_mean=[0.485, 0.456, 0.406],
                image_std=[0.229, 0.224, 0.225],
                do_reduce_labels=False
            )
            
            print("✓ Model loaded successfully")
        except Exception as e:
            print(f"✗ Error loading model: {e}")
            raise
    
    def predict_single(self, image_path, return_probs=False):
        """Dự đoán trên một ảnh"""
        try:
            # Load image
            image = Image.open(image_path).convert("RGB")
            original_size = image.size[::-1]  # (H, W)
            
            # Process image
            inputs = self.processor(images=image, return_tensors="pt")
            
            # Inference
            with torch.no_grad():
                outputs = self.model(pixel_values=inputs["pixel_values"].to(self.device))
                logits = outputs.logits
            
            # Interpolate to original size
            upsampled_logits = F.interpolate(
                logits,
                size=original_size,
                mode="bilinear",
                align_corners=False
            )
            
            pred_mask = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
            
            if return_probs:
                probs = torch.softmax(upsampled_logits, dim=1)[0].cpu().numpy()
                return pred_mask, probs
            
            return pred_mask
        except Exception as e:
            print(f"✗ Error predicting on {image_path}: {e}")
            return None
    
    def process_images(self, image_dir, output_dir=None):
        """Xử lý tất cả ảnh trong thư mục"""
        image_dir = Path(image_dir)
        
        if not image_dir.exists():
            print(f"✗ Directory not found: {image_dir}")
            return False
        
        image_paths = sorted(list(image_dir.glob("*.png"))) + sorted(list(image_dir.glob("*.jpg")))
        
        if not image_paths:
            print(f"✗ No images found in {image_dir}")
            return False
        
        print(f"\n📊 Processing {len(image_paths)} images...")
        
        if output_dir:
            output_dir = Path(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)
        
        results = []
        
        for img_path in tqdm(image_paths):
            img_id = img_path.stem
            
            # Predict
            pred_mask = self.predict_single(img_path)
            
            if pred_mask is None:
                continue
            
            # Count detected organs
            total_pixels = pred_mask.size
            detected_organs = {
                'large_bowel': int((pred_mask == 1).sum()),
                'small_bowel': int((pred_mask == 2).sum()),
                'stomach': int((pred_mask == 3).sum()),
                'background': int((pred_mask == 0).sum()),
                'total_pixels': total_pixels
            }
            
            result = {
                'image_id': img_id,
                'detected_organs': detected_organs,
                'total_pixels': total_pixels
            }
            results.append(result)
            
            # Save prediction mask if output_dir provided
            if output_dir:
                # Colorize prediction
                pred_colored = np.zeros((*pred_mask.shape, 3), dtype=np.uint8)
                
                # Colors: 1=red, 2=green, 3=blue, 0=black
                pred_colored[pred_mask == 1] = [255, 0, 0]      # Large bowel - Red
                pred_colored[pred_mask == 2] = [0, 154, 23]     # Small bowel - Green
                pred_colored[pred_mask == 3] = [0, 127, 255]    # Stomach - Blue
                
                pred_img = Image.fromarray(pred_colored)
                pred_img.save(output_dir / f"{img_id}_pred.png")
        
        # Print summary
        print("\n" + "="*60)
        print("📈 Prediction Summary")
        print("="*60)
        
        if results:
            print(f"\nProcessed {len(results)} images successfully\n")
            
            # Statistics
            for idx, result in enumerate(results, 1):
                print(f"{idx}. {result['image_id']}")
                organs = result['detected_organs']
                total = organs['large_bowel'] + organs['small_bowel'] + organs['stomach']
                if total > 0:
                    print(f"   - Large bowel: {organs['large_bowel']:,} pixels")
                    print(f"   - Small bowel: {organs['small_bowel']:,} pixels")
                    print(f"   - Stomach: {organs['stomach']:,} pixels")
                    print(f"   - Total organs: {total:,} pixels ({100*total/organs['total_pixels']:.1f}%)")
                else:
                    print(f"   - No organs detected")
        
        # Save results
        if output_dir:
            with open(output_dir / "predictions.json", 'w') as f:
                json.dump(results, f, indent=2)
            print(f"\n✓ Predictions saved to {output_dir}")
            print(f"  - Colored masks: {output_dir}/*_pred.png")
            print(f"  - Results JSON: {output_dir}/predictions.json")
        
        return True

def main():
    parser = argparse.ArgumentParser(description="Simple test on sample images")
    parser.add_argument("--model", type=str, required=True,
                       help="Path to trained model")
    parser.add_argument("--images", type=str, required=True,
                       help="Path to images directory")
    parser.add_argument("--output-dir", type=str, default=None,
                       help="Output directory for results")
    
    args = parser.parse_args()
    
    # Initialize tester
    tester = SimpleSegmentationTester(args.model)
    
    # Process images
    success = tester.process_images(args.images, args.output_dir)
    
    return success

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)