# -*- coding: utf-8 -*-
"""
마스크 모듈 - OCR 결과를 기반으로 마스크 생성 및 처리
확장, 블러, 모폴로지 연산 등 다양한 마스크 처리 기능을 제공합니다.
"""

import cv2
import numpy as np
from typing import List, Dict, Any, Tuple
import os
from shapely.geometry import Polygon

class MaskModule:
    def __init__(self):
        """마스크 모듈 초기화"""
        print("마스크 모듈 초기화 완료")
    
    def create_masks(self, image_path: str, ocr_results: List[Dict], 
                    expansion_size: int = 10, blur_size: int = 15, mask_option: str = "basic") -> np.ndarray:
        """
        OCR 결과를 기반으로 마스크 생성 (블러 확장 강화)
        
        Args:
            image_path (str): 원본 이미지 경로
            ocr_results (List[Dict]): OCR 결과
            expansion_size (int): 마스크 확장 크기 (기본값 증가)
            blur_size (int): 블러 크기 (기본값 증가)
            mask_option (str): 마스크 옵션 (기본값 "basic")
        Returns:
            np.ndarray: 생성된 마스크 (0-255)
        """
        # 원본 이미지 로드
        image = cv2.imread(image_path)
        if image is None:
            print(f"이미지를 읽을 수 없습니다: {image_path}")
            return None
            
        height, width = image.shape[:2]
        
        # 빈 마스크 생성
        mask = np.zeros((height, width), dtype=np.uint8)
        
        print(f"\n🎭 마스크 생성 시작 - 이미지 크기: {width}x{height}")
        print(f"📏 확장 크기: {expansion_size}px, 블러 크기: {blur_size}px")
        
        # OCR 결과로부터 마스크 생성
        for i, result in enumerate(ocr_results, 1):
            polygon = result['polygon']
            expanded_poly = self.expand_polygon(polygon, offset=5)
            text = result['text']
            
            print(f"   🔷 영역 #{i}: '{text}' 마스크 생성 중...")
            
            # 폴리곤을 마스크에 그리기
            cv2.fillPoly(mask, [expanded_poly], 255)
        
        # 마스크 후처리 (블러 확장 강화)
        if mask_option == "basic": 
            processed_mask = self.process_mask(mask, expansion_size, blur_size)
        # elif mask_option == "enhanced":
        #     processed_mask = self.process_mask_with_enhanced_blur(processed_mask)
        elif mask_option == "morphology":
            processed_mask = self.create_advanced_mask(image_path, ocr_results, method="morphology")
        elif mask_option == "adaptive":
            processed_mask = self.create_advanced_mask(image_path, ocr_results, method="adaptive")
        elif mask_option == "contour":
            processed_mask = self.create_advanced_mask(image_path, ocr_results, method="contour")
        elif mask_option == "ultra_blur":
            processed_mask = self.create_ultra_blur_mask(image_path, ocr_results)   
        # elif mask_option == "feathered":
        #     processed_mask = self.create_feathered_mask(processed_mask)

        print(f"✅ 마스크 생성 완료: {len(ocr_results)}개 영역")
        return processed_mask

    def expand_polygon(self, polygon, offset=15):
        poly = Polygon(polygon)
        expanded = poly.buffer(offset)
        if expanded.is_empty:
            return polygon
        return np.array(expanded.exterior.coords, dtype=np.int32)

    def process_mask(self, mask: np.ndarray, expansion_size: int = 5, 
                    blur_size: int = 3) -> np.ndarray:
        """
        기본 마스크 후처리 (확장, 블러 등)
        
        Args:
            mask (np.ndarray): 원본 마스크
            expansion_size (int): 확장 크기
            blur_size (int): 블러 크기
            
        Returns:
            np.ndarray: 처리된 마스크
        """
        processed_mask = mask.copy()
        
        # 1. 마스크 확장 (Dilation)
        if expansion_size > 0:
            kernel = np.ones((expansion_size, expansion_size), np.uint8)
            processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
        
        # 2. 가우시안 블러 적용
        if blur_size > 0:
            blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1  # 홀수로 만들기
            processed_mask = cv2.GaussianBlur(processed_mask, (blur_size, blur_size), 0)
        
        return processed_mask
    
    def process_mask_with_enhanced_blur(self, mask: np.ndarray, expansion_size: int = 10, 
                                      blur_size: int = 15) -> np.ndarray:
        """
        마스크 후처리 (강화된 블러 확장)
        
        Args:
            mask (np.ndarray): 원본 마스크
            expansion_size (int): 확장 크기
            blur_size (int): 블러 크기
            
        Returns:
            np.ndarray: 처리된 마스크
        """
        print(f"🔄 강화된 마스크 후처리 시작...")
        processed_mask = mask.copy()
        
        # 1. 다단계 마스크 확장 (Dilation)
        if expansion_size > 0:
            print(f"   📈 1단계: 기본 확장 ({expansion_size}px)")
            kernel = np.ones((expansion_size, expansion_size), np.uint8)
            processed_mask = cv2.dilate(processed_mask, kernel, iterations=1)
            
            # 추가 확장 (더 부드러운 경계를 위해)
            print(f"   📈 2단계: 추가 확장 ({expansion_size//2}px)")
            kernel2 = np.ones((expansion_size//2, expansion_size//2), np.uint8)
            processed_mask = cv2.dilate(processed_mask, kernel2, iterations=2)
        
        # 2. 강화된 가우시안 블러 적용
        if blur_size > 0:
            blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1  # 홀수로 만들기
            print(f"   🌫️  1단계 블러: 가우시안 블러 ({blur_size}px)")
            processed_mask = cv2.GaussianBlur(processed_mask, (blur_size, blur_size), 0)
            
            # 추가 블러 (더 부드러운 효과)
            additional_blur = max(5, blur_size // 2)
            additional_blur = additional_blur if additional_blur % 2 == 1 else additional_blur + 1
            print(f"   🌫️  2단계 블러: 추가 블러 ({additional_blur}px)")
            processed_mask = cv2.GaussianBlur(processed_mask, (additional_blur, additional_blur), 0)
        
        # 3. 모폴로지 연산으로 마무리
        print(f"   🔧 3단계: 모폴로지 연산")
        kernel_morph = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        processed_mask = cv2.morphologyEx(processed_mask, cv2.MORPH_CLOSE, kernel_morph)
        
        print(f"✅ 강화된 마스크 후처리 완료")
        return processed_mask
    
    def create_ultra_blur_mask(self, image_path: str, ocr_results: List[Dict],
                              expansion_size: int = 20, blur_size: int = 25) -> np.ndarray:
        """
        초강력 블러 확장 마스크 생성
        
        Args:
            image_path (str): 원본 이미지 경로
            ocr_results (List[Dict]): OCR 결과
            expansion_size (int): 대형 확장 크기
            blur_size (int): 대형 블러 크기
            
        Returns:
            np.ndarray: 초블러 마스크
        """
        image = cv2.imread(image_path)
        if image is None:
            return None
            
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        
        print(f"\n🌟 초강력 블러 마스크 생성 시작")
        print(f"📏 초대형 확장: {expansion_size}px, 초대형 블러: {blur_size}px")
        
        # 기본 마스크 생성
        for result in ocr_results:
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            cv2.fillPoly(mask, [polygon_np], 255)
        
        # 다단계 확장
        for i in range(3):
            kernel_size = expansion_size - (i * 3)
            if kernel_size > 0:
                kernel = np.ones((kernel_size, kernel_size), np.uint8)
                mask = cv2.dilate(mask, kernel, iterations=1)
                print(f"   📈 확장 단계 {i+1}: {kernel_size}px")
        
        # 다단계 블러
        blur_steps = [blur_size, blur_size//2, blur_size//4]
        for i, blur in enumerate(blur_steps):
            if blur > 0:
                blur = blur if blur % 2 == 1 else blur + 1
                mask = cv2.GaussianBlur(mask, (blur, blur), 0)
                print(f"   🌫️  블러 단계 {i+1}: {blur}px")
        
        print(f"✅ 초강력 블러 마스크 완료")
        return mask
    
    def create_advanced_mask(self, image_path: str, ocr_results: List[Dict],
                           method: str = "morphology") -> np.ndarray:
        """
        고급 마스크 생성 방법들
        
        Args:
            image_path (str): 원본 이미지 경로
            ocr_results (List[Dict]): OCR 결과
            method (str): 마스크 생성 방법
                - "morphology": 모폴로지 연산
                - "adaptive": 적응적 확장
                - "contour": 컨투어 기반
                
        Returns:
            np.ndarray: 생성된 마스크
        """
        image = cv2.imread(image_path)
        if image is None:
            return None
            
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        
        if method == "morphology":
            return self._create_morphology_mask(image, ocr_results)
        elif method == "adaptive":
            return self._create_adaptive_mask(image, ocr_results)
        elif method == "contour":
            return self._create_contour_mask(image, ocr_results)
        else:
            return self.create_masks(image_path, ocr_results)
    
    def _create_morphology_mask(self, image: np.ndarray, 
                               ocr_results: List[Dict]) -> np.ndarray:
        """모폴로지 연산을 사용한 마스크 생성"""
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        
        for result in ocr_results:
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            cv2.fillPoly(mask, [polygon_np], 255)
        
        # 모폴로지 연산
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.dilate(mask, kernel, iterations=1)
        
        return mask
    
    def _create_adaptive_mask(self, image: np.ndarray, 
                             ocr_results: List[Dict]) -> np.ndarray:
        """적응적 확장을 사용한 마스크 생성"""
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        
        for result in ocr_results:
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            
            # 텍스트 크기에 따라 확장 크기 조정
            bbox = result['bbox']
            text_area = bbox[2] * bbox[3]
            expansion_size = max(3, int(np.sqrt(text_area) / 20))
            
            # 임시 마스크 생성
            temp_mask = np.zeros((height, width), dtype=np.uint8)
            cv2.fillPoly(temp_mask, [polygon_np], 255)
            
            # 적응적 확장
            kernel = np.ones((expansion_size, expansion_size), np.uint8)
            temp_mask = cv2.dilate(temp_mask, kernel, iterations=1)
            
            # 메인 마스크에 추가
            mask = cv2.bitwise_or(mask, temp_mask)
        
        return mask
    
    def _create_contour_mask(self, image: np.ndarray, 
                            ocr_results: List[Dict]) -> np.ndarray:
        """컨투어 기반 마스크 생성"""
        height, width = image.shape[:2]
        mask = np.zeros((height, width), dtype=np.uint8)
        
        for result in ocr_results:
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            
            # 컨투어로 변환
            contour = polygon_np.reshape(-1, 1, 2)
            
            # 컨투어 확장
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 확장된 컨투어로 마스크 생성
            cv2.fillPoly(mask, [approx], 255)
        
        return mask
    
    def create_feathered_mask(self, mask: np.ndarray, feather_size: int = 10) -> np.ndarray:
        """
        페더링된 마스크 생성 (부드러운 경계)
        
        Args:
            mask (np.ndarray): 원본 마스크
            feather_size (int): 페더링 크기
            
        Returns:
            np.ndarray: 페더링된 마스크
        """
        # 거리 변환
        dist_transform = cv2.distanceTransform(mask, cv2.DIST_L2, 5)
        
        # 페더링 적용
        feathered = np.clip(dist_transform / feather_size * 255, 0, 255).astype(np.uint8)
        
        return feathered
    
    def visualize_mask(self, image_path: str, mask: np.ndarray, 
                      output_path: str = None, alpha: float = 0.5) -> np.ndarray:
        """
        마스크를 원본 이미지에 오버레이하여 시각화
        
        Args:
            image_path (str): 원본 이미지 경로
            mask (np.ndarray): 마스크
            output_path (str): 저장할 경로
            alpha (float): 투명도
            
        Returns:
            np.ndarray: 시각화된 이미지
        """
        image = cv2.imread(image_path)
        if image is None:
            return None
        
        # 마스크를 3채널로 변환
        if len(mask.shape) == 2:
            mask_colored = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        else:
            mask_colored = mask
        
        # 마스크에 색상 적용 (빨간색)
        mask_colored[:, :, 0] = 0  # Blue
        mask_colored[:, :, 1] = 0  # Green
        # Red 채널은 그대로 유지
        
        # 오버레이
        overlay = cv2.addWeighted(image, 1-alpha, mask_colored, alpha, 0)
        
        if output_path:
            cv2.imwrite(output_path, overlay)
            print(f"마스크 시각화 결과 저장: {output_path}")
        
        return overlay
    
    def get_mask_statistics(self, mask: np.ndarray) -> Dict[str, Any]:
        """
        마스크 통계 정보 반환
        
        Args:
            mask (np.ndarray): 마스크
            
        Returns:
            Dict: 통계 정보
        """
        total_pixels = mask.shape[0] * mask.shape[1]
        mask_pixels = np.sum(mask > 0)
        coverage = mask_pixels / total_pixels * 100
        
        # 컨투어 찾기
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        num_regions = len(contours)
        
        # 가장 큰 영역 크기
        if contours:
            largest_area = max(cv2.contourArea(contour) for contour in contours)
        else:
            largest_area = 0
        
        stats = {
            'total_pixels': total_pixels,
            'mask_pixels': mask_pixels,
            'coverage_percent': coverage,
            'num_regions': num_regions,
            'largest_region_area': largest_area
        }
        
        return stats
    
    def test_module(self):
        """마스크 모듈 테스트 (강화된 블러 확장 포함)"""
        print("🎭 마스크 모듈 테스트 시작...")
        
        # 테스트용 가상 OCR 결과 생성
        fake_ocr_results = [
            {
                'text': '测试文本1',
                'confidence': 0.95,
                'polygon': [[100, 50], [200, 50], [200, 80], [100, 80]],
                'bbox': (100, 50, 100, 30)
            },
            {
                'text': '测试文本2',
                'confidence': 0.88,
                'polygon': [[150, 120], [300, 120], [300, 160], [150, 160]],
                'bbox': (150, 120, 150, 40)
            }
        ]
        
        # 테스트 이미지 생성 (없을 경우)
        test_image_path = "test_output/test_mask_image.jpg"
        os.makedirs("test_output", exist_ok=True)
        
        # 간단한 테스트 이미지 생성
        test_image = np.ones((200, 400, 3), dtype=np.uint8) * 255
        cv2.putText(test_image, "Test Image for Mask", (50, 100), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
        cv2.imwrite(test_image_path, test_image)
        
        # 1. 기본 마스크 생성 테스트
        print("\n📋 1. 기본 마스크 테스트")
        mask = self.create_masks(test_image_path, fake_ocr_results)
        if mask is not None:
            print("✅ 기본 마스크 생성 성공")
            
            # 마스크 통계
            stats = self.get_mask_statistics(mask)
            print(f"📊 마스크 통계:")
            print(f"   - 총 픽셀: {stats['total_pixels']:,}")
            print(f"   - 마스크 픽셀: {stats['mask_pixels']:,}")
            print(f"   - 커버리지: {stats['coverage_percent']:.2f}%")
            print(f"   - 영역 수: {stats['num_regions']}")
            print(f"   - 최대 영역 크기: {stats['largest_region_area']:,}")
            
            # 시각화
            self.visualize_mask(test_image_path, mask, "test_output/mask_basic.jpg")
        
        # 2. 강화된 블러 확장 마스크 테스트
        print("\n🌟 2. 강화된 블러 확장 마스크 테스트")
        image = cv2.imread(test_image_path)
        height, width = image.shape[:2]
        base_mask = np.zeros((height, width), dtype=np.uint8)
        
        # 기본 마스크 생성
        for result in fake_ocr_results:
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            cv2.fillPoly(base_mask, [polygon_np], 255)
        
        # 강화된 블러 적용
        enhanced_mask = self.process_mask_with_enhanced_blur(base_mask, 12, 20)
        self.visualize_mask(test_image_path, enhanced_mask, "test_output/mask_enhanced_blur.jpg")
        print("✅ 강화된 블러 마스크 완료")
        
        # 3. 초강력 블러 마스크 테스트
        print("\n🚀 3. 초강력 블러 마스크 테스트")
        ultra_mask = self.create_ultra_blur_mask(test_image_path, fake_ocr_results, 25, 35)
        if ultra_mask is not None:
            self.visualize_mask(test_image_path, ultra_mask, "test_output/mask_ultra_blur.jpg")
            print("✅ 초강력 블러 마스크 완료")
        
        # 4. 고급 마스크 테스트
        print("\n🔧 4. 고급 마스크 방법 테스트")
        methods = ["morphology", "adaptive", "contour"]
        for method in methods:
            print(f"   🔹 {method} 방법 테스트 중...")
            advanced_mask = self.create_advanced_mask(
                test_image_path, fake_ocr_results, method
            )
            if advanced_mask is not None:
                output_path = f"test_output/mask_{method}.jpg"
                self.visualize_mask(test_image_path, advanced_mask, output_path)
                print(f"   ✅ {method} 마스크 완료")
        
        # 5. 페더링 마스크 테스트
        print("\n✨ 5. 페더링 마스크 테스트")
        if mask is not None:
            feathered_mask = self.create_feathered_mask(mask, 15)
            self.visualize_mask(test_image_path, feathered_mask, 
                              "test_output/mask_feathered.jpg")
            print("✅ 페더링 마스크 완료")
        
        print("\n🎉 마스크 모듈 테스트 완료!")
        print("📁 결과 파일들이 test_output/ 폴더에 저장되었습니다.")
        print("   - mask_basic.jpg: 기본 마스크")
        print("   - mask_enhanced_blur.jpg: 강화된 블러 마스크")
        print("   - mask_ultra_blur.jpg: 초강력 블러 마스크")
        print("   - mask_morphology.jpg: 모폴로지 마스크")
        print("   - mask_adaptive.jpg: 적응적 마스크")
        print("   - mask_contour.jpg: 컨투어 마스크")
        print("   - mask_feathered.jpg: 페더링 마스크")

    def create_masks_with_save(self, image_path: str, ocr_results: List[Dict], 
                              output_dir: str = None, expansion_size: int = 10, 
                              blur_size: int = 15, mask_option: str = "basic") -> tuple:
        """
        OCR 결과를 기반으로 마스크 생성 및 저장
        
        Args:
            image_path (str): 원본 이미지 경로
            ocr_results (List[Dict]): OCR 결과
            output_dir (str): 출력 디렉토리 (마스크 이미지 저장용)
            expansion_size (int): 마스크 확장 크기
            blur_size (int): 블러 크기
            mask_option (str): 마스크 옵션 (기본값 "basic")
        Returns:
            tuple: (mask, mask_image_path)
        """
        # 기본 마스크 생성
        mask = self.create_masks(image_path, ocr_results, expansion_size, blur_size, mask_option)
        
        mask_image_path = None
        if output_dir and mask is not None:
            os.makedirs(output_dir, exist_ok=True)
            
            # 파일명 생성
            base_name = os.path.splitext(os.path.basename(image_path))[0]
            mask_filename = f"{base_name}_mask_combined.jpg"
            mask_image_path = os.path.join(output_dir, mask_filename)
            
            # 마스크 이미지 저장
            cv2.imwrite(mask_image_path, mask)
            print(f"💾 마스크 이미지 저장: {mask_image_path}")
            
            # 개별 마스크들도 저장
            self._save_individual_masks(image_path, ocr_results, output_dir)
            
            # 마스크 시각화 이미지도 저장
            visualization_path = os.path.join(output_dir, f"{base_name}_mask_visualization.jpg")
            self.visualize_mask(image_path, mask, visualization_path)
            print(f"💾 마스크 시각화 이미지 저장: {visualization_path}")
        
        return mask, mask_image_path
    
    def _save_individual_masks(self, image_path: str, ocr_results: List[Dict], output_dir: str):
        """개별 텍스트 영역의 마스크를 각각 저장"""
        image = cv2.imread(image_path)
        if image is None:
            return
            
        height, width = image.shape[:2]
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        
        for i, result in enumerate(ocr_results, 1):
            # 개별 마스크 생성
            individual_mask = np.zeros((height, width), dtype=np.uint8)
            polygon = result['polygon']
            polygon_np = np.array(polygon, dtype=np.int32)
            cv2.fillPoly(individual_mask, [polygon_np], 255)
            
            # 개별 마스크 저장
            individual_filename = f"{base_name}_mask_area_{i:02d}.jpg"
            individual_path = os.path.join(output_dir, individual_filename)
            cv2.imwrite(individual_path, individual_mask)
            
            print(f"   💾 개별 마스크 #{i} 저장: {individual_filename}")

if __name__ == "__main__":
    mask_module = MaskModule()
    mask_module.test_module()