#!/usr/bin/env python3
"""KOCR v4 — 图像预处理 + 自动图像分析调参"""

import cv2
import numpy as np
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class ImageQuality:
    """图像质量分析结果"""
    brightness: float = 0.0          # 平均亮度 [0-255]
    contrast: float = 0.0            # 标准偏差
    blur_score: float = 0.0          # 拉普拉斯方差（越小越模糊）
    shadow_ratio: float = 0.0        # 阴影面积占比
    photo_type: str = 'normal'       # flash / dark / blurry / normal
    needs_shadow_removal: bool = True
    needs_sharpen: bool = False
    needs_deskew: bool = True


def analyze_image(img: np.ndarray) -> ImageQuality:
    """
    自动分析图像质量，返回各项指标和推荐的照片类型
    
    Args:
        img: 输入图像（BGR格式）
    
    Returns:
        ImageQuality: 包含所有分析结果
    """
    h, w = img.shape[:2]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 1. 亮度分析（取直方图均值）
    brightness = float(np.mean(gray))
    
    # 2. 对比度（灰度标准偏差）
    contrast = float(np.std(gray))
    
    # 3. 模糊检测（拉普拉斯方差）
    laplacian = cv2.Laplacian(gray, cv2.CV_64F)
    blur_score = float(laplacian.var())
    
    # 4. 阴影检测
    # 用大核膨胀分离背景和前景
    dilated = cv2.dilate(gray, np.ones((15, 15), np.uint8))
    bg = cv2.medianBlur(dilated, 31)
    diff = cv2.absdiff(gray.astype(np.float32), bg.astype(np.float32))
    shadow_mask = diff < 20  # 与背景差异小的区域视为阴影
    shadow_ratio = float(np.sum(shadow_mask) / (h * w))
    
    # 5. 判定照片类型
    if brightness >= 200:
        photo_type = 'flash'
    elif brightness <= 60:
        photo_type = 'dark'
    elif blur_score < 80:
        photo_type = 'blurry'
    else:
        photo_type = 'normal'
    
    # 6. 预处理策略建议
    needs_shadow_removal = shadow_ratio > 0.3 and photo_type != 'flash'
    needs_sharpen = blur_score < 150
    needs_deskew = photo_type not in ('flash',)  # 闪光灯照片也可能歪
    
    return ImageQuality(
        brightness=round(brightness, 1),
        contrast=round(contrast, 1),
        blur_score=round(blur_score, 1),
        shadow_ratio=round(shadow_ratio, 3),
        photo_type=photo_type,
        needs_shadow_removal=needs_shadow_removal,
        needs_sharpen=needs_sharpen,
        needs_deskew=needs_deskew,
    )


class VoucherPreprocessor:
    """记账凭证图像预处理器，支持自动调参"""
    
    def __init__(self, max_side=1600, denoise=True, deskew=True,
                 remove_shadow=True, sharpen=True):
        self.max_side = max_side
        self.denoise = denoise
        self.deskew = deskew
        self.remove_shadow = remove_shadow
        self.sharpen = sharpen
    
    @classmethod
    def from_config(cls, config: dict, photo_type: str = 'normal'):
        """根据配置和照片类型创建预处理器实例
        
        Args:
            config: config.yaml 中 preprocessing 章节
            photo_type: 照片类型（用于覆盖默认参数）
        """
        params = {
            'max_side': config.get('max_side', 1600),
            'denoise': config.get('denoise', True),
            'deskew': config.get('deskew', True),
            'remove_shadow': config.get('remove_shadow', True),
            'sharpen': config.get('sharpen', True),
        }
        # 应用照片类型覆盖
        overrides = config.get('photo_type_overrides', {}).get(photo_type, {})
        params.update(overrides)
        return cls(**params)
    
    def _resize(self, img):
        """统一缩放，长边不超过max_side，保持比例"""
        h, w = img.shape[:2]
        if max(h, w) > self.max_side:
            scale = self.max_side / max(h, w)
            new_w, new_h = int(w * scale), int(h * scale)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        return img
    
    def _remove_shadow(self, img):
        """去除拍照阴影、纸张背景不均"""
        rgb_planes = cv2.split(img)
        result_planes = []
        for plane in rgb_planes:
            dilated_img = cv2.dilate(plane, np.ones((7, 7), np.uint8))
            bg_img = cv2.medianBlur(dilated_img, 21)
            diff_img = 255 - cv2.absdiff(plane, bg_img)
            norm_img = cv2.normalize(diff_img, None, alpha=0, beta=255,
                                     norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
            result_planes.append(norm_img)
        return cv2.merge(result_planes)
    
    def _deskew(self, img):
        """基于霍夫直线的倾斜矫正"""
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
        lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)
        if lines is None:
            return img
        angles = []
        for line in lines:
            _, theta = line[0]
            angle = (theta * 180 / np.pi) - 90
            if abs(angle) < 15:
                angles.append(angle)
        if not angles:
            return img
        median_angle = np.median(angles)
        h, w = img.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
        rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC,
                                 borderMode=cv2.BORDER_REPLICATE)
        return rotated
    
    def _denoise(self, img):
        """高斯去噪"""
        return cv2.GaussianBlur(img, (3, 3), 0)
    
    def _sharpen_op(self, img):
        """USM 锐化"""
        blur = cv2.GaussianBlur(img, (0, 0), 3)
        sharp = cv2.addWeighted(img, 1.5, blur, -0.5, 0)
        return sharp
    
    def process(self, img_or_path):
        """完整预处理入口
        
        Args:
            img_or_path: 图像路径或 numpy 数组
        
        Returns:
            预处理后的图像
        """
        if isinstance(img_or_path, str):
            img = cv2.imread(img_or_path)
            if img is None:
                raise ValueError(f"无法读取图片：{img_or_path}")
        else:
            img = img_or_path.copy()
        
        img = self._resize(img)
        
        if self.remove_shadow:
            img = self._remove_shadow(img)
        if self.deskew:
            img = self._deskew(img)
        if self.denoise:
            img = self._denoise(img)
        if self.sharpen:
            img = self._sharpen_op(img)
        
        return img