#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (c) 2025 Any2MIF Project
# All rights reserved.

"""
Any2MIF - 图像处理器
负责图像的灰度化、二值化等操作
"""

import os
import numpy as np
from PIL import Image, ImageFilter

class ImageProcessor:
    """图像处理器"""
    
    def __init__(self):
        """初始化图像处理器"""
        # 支持的图像格式
        self.supported_formats = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff']
    
    def is_image_file(self, file_path):
        """
        检查文件是否为图像文件
        
        参数:
            file_path (str): 文件路径
        
        返回:
            bool: 是否为图像文件
        """
        _, ext = os.path.splitext(file_path.lower())
        return ext in self.supported_formats
    
    def grayscale(self, image, method='weighted'):
        """
        将图像转换为灰度图
        
        参数:
            image (PIL.Image): 输入图像
            method (str): 灰度化方法
                'average': 平均值法
                'weighted': 加权平均法
                'max': 最大值法
        
        返回:
            PIL.Image: 灰度图像
        """
        # 确保图像是RGB模式
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # 转换为numpy数组
        img_array = np.array(image)
        
        # 根据方法进行灰度化
        if method == 'average':
            # 平均值法: (R + G + B) / 3
            gray_array = np.mean(img_array, axis=2).astype(np.uint8)
        elif method == 'weighted':
            # 加权平均法: 0.299 * R + 0.587 * G + 0.114 * B
            gray_array = (0.299 * img_array[:, :, 0] + 
                          0.587 * img_array[:, :, 1] + 
                          0.114 * img_array[:, :, 2]).astype(np.uint8)
        elif method == 'max':
            # 最大值法: max(R, G, B)
            gray_array = np.max(img_array, axis=2).astype(np.uint8)
        else:
            # 默认使用加权平均法
            gray_array = (0.299 * img_array[:, :, 0] + 
                          0.587 * img_array[:, :, 1] + 
                          0.114 * img_array[:, :, 2]).astype(np.uint8)
        
        # 转换回PIL图像
        return Image.fromarray(gray_array, mode='L')
    
    def threshold(self, image, threshold=128, method='fixed'):
        """
        对图像进行二值化处理
        
        参数:
            image (PIL.Image): 输入图像
            threshold (int): 阈值 (0-255)
            method (str): 二值化方法
                'fixed': 固定阈值
                'adaptive': 自适应阈值
                'otsu': Otsu阈值
        
        返回:
            PIL.Image: 二值化图像
        """
        # 确保图像是灰度模式
        if image.mode != 'L':
            image = image.convert('L')
        
        # 转换为numpy数组
        img_array = np.array(image)
        
        # 根据方法进行二值化
        if method == 'fixed':
            # 固定阈值
            binary_array = (img_array > threshold).astype(np.uint8) * 255
        elif method == 'adaptive':
            # 自适应阈值 (简化版本)
            # 使用局部区域的平均值作为阈值
            kernel_size = 15
            # 创建一个平均滤波器
            kernel = np.ones((kernel_size, kernel_size)) / (kernel_size * kernel_size)
            # 计算局部平均值
            local_mean = np.zeros_like(img_array, dtype=np.float32)
            
            # 简单的局部平均实现
            padded = np.pad(img_array, kernel_size // 2, mode='reflect')
            for i in range(img_array.shape[0]):
                for j in range(img_array.shape[1]):
                    local_mean[i, j] = np.mean(padded[i:i+kernel_size, j:j+kernel_size])
            
            # 二值化
            binary_array = (img_array > local_mean - threshold // 2).astype(np.uint8) * 255
        elif method == 'otsu':
            # Otsu阈值
            # 计算直方图
            hist, bins = np.histogram(img_array.flatten(), 256, [0, 256])
            
            # 计算累积和
            cum_sum = hist.cumsum()
            cum_sum_sq = (hist * np.arange(256)).cumsum()
            
            # 初始化
            max_var = 0
            otsu_threshold = 0
            
            # 遍历所有可能的阈值
            for t in range(1, 256):
                # 前景和背景的权重
                w_bg = cum_sum[t]
                w_fg = cum_sum[-1] - w_bg
                
                # 如果前景或背景为空，跳过
                if w_bg == 0 or w_fg == 0:
                    continue
                
                # 前景和背景的均值
                mean_bg = cum_sum_sq[t] / w_bg
                mean_fg = (cum_sum_sq[-1] - cum_sum_sq[t]) / w_fg
                
                # 计算类间方差
                var_between = w_bg * w_fg * (mean_bg - mean_fg) ** 2
                
                # 更新最大方差和阈值
                if var_between > max_var:
                    max_var = var_between
                    otsu_threshold = t
            
            # 使用Otsu阈值进行二值化
            binary_array = (img_array > otsu_threshold).astype(np.uint8) * 255
        else:
            # 默认使用固定阈值
            binary_array = (img_array > threshold).astype(np.uint8) * 255
        
        # 转换回PIL图像
        return Image.fromarray(binary_array, mode='L')
    
    def denoise(self, image, method='median', strength=3):
        """
        对图像进行降噪处理
        
        参数:
            image (PIL.Image): 输入图像
            method (str): 降噪方法
                'median': 中值滤波
                'gaussian': 高斯滤波
                'bilateral': 双边滤波 (简化版本)
            strength (int): 降噪强度 (1-10)
        
        返回:
            PIL.Image: 降噪后的图像
        """
        # 根据方法进行降噪
        if method == 'median':
            # 中值滤波
            kernel_size = strength * 2 - 1  # 确保是奇数
            return image.filter(ImageFilter.MedianFilter(size=kernel_size))
        elif method == 'gaussian':
            # 高斯滤波
            radius = strength / 2
            return image.filter(ImageFilter.GaussianBlur(radius=radius))
        elif method == 'bilateral':
            # 双边滤波 (简化版本)
            # 由于PIL没有内置的双边滤波，我们使用高斯滤波代替
            # 在实际应用中，可以使用OpenCV的双边滤波
            radius = strength / 2
            return image.filter(ImageFilter.GaussianBlur(radius=radius))
        else:
            # 默认使用中值滤波
            kernel_size = strength * 2 - 1  # 确保是奇数
            return image.filter(ImageFilter.MedianFilter(size=kernel_size))
    
    def resize(self, image, width, height, keep_aspect=True):
        """
        调整图像大小
        
        参数:
            image (PIL.Image): 输入图像
            width (int): 目标宽度
            height (int): 目标高度
            keep_aspect (bool): 是否保持纵横比
        
        返回:
            PIL.Image: 调整大小后的图像
        """
        if keep_aspect:
            # 计算缩放比例
            width_ratio = width / image.width
            height_ratio = height / image.height
            ratio = min(width_ratio, height_ratio)
            
            # 计算新尺寸
            new_width = int(image.width * ratio)
            new_height = int(image.height * ratio)
            
            # 调整图像大小
            resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # 创建新图像
            new_image = Image.new(image.mode, (width, height), (0, 0, 0))
            
            # 计算粘贴位置
            paste_x = (width - new_width) // 2
            paste_y = (height - new_height) // 2
            
            # 粘贴调整大小后的图像
            new_image.paste(resized_image, (paste_x, paste_y))
            
            return new_image
        else:
            # 直接调整到指定大小
            return image.resize((width, height), Image.Resampling.LANCZOS)