import cv2
import numpy as np

class StripeRemover:
    def __init__(self):
        pass
    
    def fourier_method(self, image):
        """傅里叶变换去除条纹"""
        # 转换为灰度图像
        if len(image.shape) > 2:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 傅里叶变换
        dft = cv2.dft(np.float32(gray), flags=cv2.DFT_COMPLEX_OUTPUT)
        dft_shift = np.fft.fftshift(dft)
        
        # 创建掩模
        rows, cols = gray.shape
        crow, ccol = rows // 2, cols // 2
        mask = np.ones((rows, cols, 2), np.uint8)
        mask[crow-5:crow+5, :] = 0  # 调整带宽以匹配条纹频率
        
        # 应用掩模和逆变换
        fshift = dft_shift * mask
        f_ishift = np.fft.ifftshift(fshift)
        img_back = cv2.idft(f_ishift)
        img_back = cv2.magnitude(img_back[:,:,0], img_back[:,:,1])
        
        # 归一化处理
        img_back = cv2.normalize(img_back, None, 0, 255, cv2.NORM_MINMAX)
        return np.uint8(img_back)
    
    def morphological_method(self, image):
        """形态学操作去除条纹"""
        # Convert to grayscale if needed
        if len(image.shape) > 2:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # Create horizontal kernel (adjust size if needed)
        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25,1))
        
        # Detect lines using morphological operations
        detected_lines = cv2.morphologyEx(gray, cv2.MORPH_OPEN, horizontal_kernel)
        
        # Remove detected lines and normalize
        result = cv2.subtract(gray, detected_lines)
        
        # Normalize and enhance contrast
        result = cv2.normalize(result, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        
        # Convert back to BGR for display
        return cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)
    
    def adaptive_threshold_method(self, image):
        """自适应阈值处理"""
        if len(image.shape) > 2:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 自适应阈值处理
        thresh = cv2.adaptiveThreshold(gray, 255, 
                                     cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                     cv2.THRESH_BINARY, 11, 2)
        
        # 中值滤波去噪
        result = cv2.medianBlur(thresh, 3)
        return result
    
    def enhance_image(self, image):
        """图像增强处理"""
        # CLAHE对比度增强
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        enhanced = clahe.apply(image)
        
        # 锐化处理
        kernel = np.array([[-1,-1,-1],
                         [-1, 9,-1],
                         [-1,-1,-1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel)
        
        return sharpened
    
    def adaptive_enhance(self, image):
        """自适应阈值 + 图像增强"""
        thresh_result = self.adaptive_threshold_method(image)
        return self.enhance_image(thresh_result)
        
    def fourier_enhance(self, image):
        """傅里叶变换 + 图像增强"""
        fourier_result = self.fourier_method(image)
        return self.enhance_image(fourier_result)
        
    def morphological_enhance(self, image):
        """形态学操作 + 图像增强"""
        morph_result = self.morphological_method(image)
        # Convert BGR to gray for enhance_image
        gray = cv2.cvtColor(morph_result, cv2.COLOR_BGR2GRAY)
        enhanced = self.enhance_image(gray)
        return cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
        
    def adaptive_fourier(self, image):
        """自适应阈值 + 傅里叶变换"""
        thresh_result = self.adaptive_threshold_method(image)
        return self.fourier_method(thresh_result)
        
    def morphological_adaptive(self, image):
        """形态学操作 + 自适应阈值"""
        morph_result = self.morphological_method(image)
        gray = cv2.cvtColor(morph_result, cv2.COLOR_BGR2GRAY)
        return self.adaptive_threshold_method(gray)
        
    def fourier_morphological(self, image):
        """傅里叶变换 + 形态学操作"""
        fourier_result = self.fourier_method(image)
        return self.morphological_method(fourier_result)