#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Optimize Seal Text Recognition (v15 - Refined Search)
Systematically searches for unwarping parameters to match target text.
"""
import sys
import cv2
import numpy as np
import os
import difflib
import json
import time

# Target text: "威凯检测技术有限公司" using unicode escapes for safety on Windows
TARGET_TEXT = "\u5a01\u51ef\u68c0\u6d4b\u6280\u672f\u6709\u9650\u516c\u53f8"
MIN_SIMILARITY = 0.70

def similarity(s1, s2):
    if not s1 or not s2: return 0.0
    # Strip any extra whitespace
    s1 = s1.strip().replace(" ", "")
    s2 = s2.strip().replace(" ", "")
    return difflib.SequenceMatcher(None, s1, s2).ratio()

def optimize_params(image_path, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    print(f"Loading {image_path}...", file=sys.stderr)
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Could not load {image_path}", file=sys.stderr)
        return
    h, w = img.shape[:2]
    
    # Pre-processing (Hough circle only)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    saturation = hsv[:,:,1]
    mask = cv2.bitwise_or(cv2.inRange(hsv, (0, 40, 40), (10, 255, 255)),
                          cv2.inRange(hsv, (160, 40, 40), (180, 255, 255)))
    red_intensity = cv2.bitwise_and(saturation, saturation, mask=mask)
    
    # Initialize OCR
    try:
        os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
        from paddleocr import PaddleOCR
        import logging
        logging.getLogger('ppocr').setLevel(logging.ERROR)
        ocr_model = PaddleOCR(use_angle_cls=False, lang='ch', show_log=False)
    except Exception as e:
        # Some versions don't support show_log in constructor
        from paddleocr import PaddleOCR
        ocr_model = PaddleOCR(use_angle_cls=False, lang='ch')

    # Centers
    centers = [((w // 2, h // 2), "img_center")]
    blur = cv2.medianBlur(red_intensity, 5)
    circles = cv2.HoughCircles(blur, cv2.HOUGH_GRADIENT, 1, h / 8,
                               param1=100, param2=30,
                               minRadius=int(h / 4), maxRadius=int(h / 1.5))
    radius_hough = min(w, h) // 2
    if circles is not None:
        best_c = circles[0][0]
        centers.append(((int(best_c[0]), int(best_c[1])), f"hough_r{int(best_c[2])}"))
        radius_hough = int(best_c[2])

    # Focussed search
    start_angles = np.arange(0, 360, 15)
    factors = [1.0, 1.5, 2.0, 2.5, 3.0]
    r_scales = [0.95, 1.0, 1.05]
    flips = [("none", None), ("vh", -1)]
    modes = [("bgr", img)]

    best_sim = 0.0
    combination_count = 0
    total_total = len(centers) * len(start_angles) * len(r_scales) * len(factors) * 2 * len(flips) * len(modes)
    print(f"Starting focused search across {total_total} combinations...", file=sys.stderr)
    
    log_path = os.path.join(output_dir, "v5_optimization_log.txt")
    debug_candidates_dir = os.path.join(output_dir, "debug_candidates")
    if not os.path.exists(debug_candidates_dir): os.makedirs(debug_candidates_dir)

    with open(log_path, "w", encoding="utf-8") as f_log:
        for m_name, work_img in modes:
            for (cx, cy), c_name in centers:
                radius_base = radius_hough if "hough" in c_name else min(cx, cy, w-cx, h-cy)
                for angle in start_angles:
                    M_rot = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
                    rotated_img = cv2.warpAffine(work_img, M_rot, (w, h), borderValue=(255, 255, 255))
                    for r_scale in r_scales:
                        rr = int(radius_base * r_scale)
                        for factor in factors:
                            out_w = int(rr * 2 * np.pi * factor)
                            out_h = rr
                            unwarped = cv2.warpPolar(rotated_img, (out_w, out_h), (cx, cy), rr, cv2.WARP_POLAR_LINEAR)
                            strips = [
                                ("outer", unwarped[int(out_h * 0.7):out_h, :]),
                                ("inner", unwarped[int(out_h * 0.4):int(out_h * 0.7), :])
                            ]
                            for s_name, strip in strips:
                                for flip_label, flip_code in flips:
                                    final = strip
                                    if flip_code is not None: final = cv2.flip(strip, flip_code)
                                    combination_count += 1
                                    param_str = f"{m_name}_{c_name}_a{angle}_r{r_scale}_f{factor}_{s_name}_{flip_label}"
                                    try:
                                        res = ocr_model.ocr(final)
                                        recognized = ""
                                        if res and len(res) > 0:
                                            if isinstance(res[0], dict) and 'rec_texts' in res[0]:
                                                for text in res[0]['rec_texts']: recognized += str(text)
                                            elif isinstance(res[0], list):
                                                for line in res[0]:
                                                    if line and len(line) > 1: recognized += str(line[1][0])
                                        sim = similarity(recognized, TARGET_TEXT)
                                        if recognized:
                                            f_log.write(f"Sim: {sim:.4f} | lenR:{len(recognized)} | {param_str} | {recognized}\n")
                                            f_log.flush()
                                            if sim > 0.1:
                                                print(f"  FOUND: {sim:.4f} | R:'{recognized}' | {param_str}", flush=True)
                                        if sim > 0.4 or combination_count % 100 == 0:
                                            cv2.imwrite(os.path.join(debug_candidates_dir, f"{param_str}_s{sim:.2f}.png"), final)
                                        if sim > best_sim:
                                            best_sim = sim
                                            cv2.imwrite(os.path.join(output_dir, f"best_sim_{sim:.2f}.png"), final)
                                            if sim >= MIN_SIMILARITY:
                                                print(f"WINNER: {param_str} -> {recognized}", file=sys.stderr)
                                                return
                                    except: pass
                                    if combination_count % 50 == 0:
                                        print(f"Progress: {combination_count}/{total_total}...", flush=True)
    print(f"Done. Best Sim: {best_sim:.2f}", file=sys.stderr)

if __name__ == "__main__":
    if len(sys.argv) < 3: sys.exit(1)
    optimize_params(sys.argv[1], sys.argv[2])