#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Brute Force Seal Unwarping
Tests all permutations of rotation and flipping to find 70%+ similarity.
"""
import sys
import cv2
import numpy as np
import os
import difflib

# Target text
TARGET_TEXT = "威凯检测技术有限公司"
MIN_SIMILARITY = 0.70

def similarity(s1, s2):
    if not s1 or not s2: return 0.0
    return difflib.SequenceMatcher(None, s1, s2).ratio()

def brute_force_search(image_path, results_dir):
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)
        
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Could not load {image_path}")
        return

    # 1. Circle Detection
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.medianBlur(gray, 5)
    rows = gray.shape[0]
    circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, rows / 8,
                               param1=100, param2=30,
                               minRadius=int(rows / 4), maxRadius=int(rows / 1.5))
    
    if circles is not None:
        best_circle = np.uint16(np.around(circles))[0][0]
        cx, cy, radius = [int(x) for x in best_circle]
    else:
        h, w = img.shape[:2]
        cx, cy, radius = w // 2, h // 2, min(w, h) // 2

    print(f"Using Center ({cx}, {cy}), Radius {radius}", file=sys.stderr)

    # 2. Setup OCR
    try:
        os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
        from paddleocr import TextRecognition
        text_rec = TextRecognition(model_name="PP-OCRv4_server_rec")
    except Exception as e:
        print(f"OCR Error: {e}")
        return

    # 3. Permutations
    rotations = [
        (None, "no_rot"),
        (cv2.ROTATE_90_CLOCKWISE, "90_cw"),
        (cv2.ROTATE_180, "180"),
        (cv2.ROTATE_90_COUNTERCLOCKWISE, "90_ccw")
    ]
    
    flips = [
        (None, "no_flip"),
        (0, "flip_v"),
        (1, "flip_h"),
        (-1, "flip_both")
    ]
    
    # Text ring is usually outer 25% of radius
    factors = [1.5, 2.5] # Test two widths
    
    found_best = None
    max_sim = 0.0

    for rot_code, rot_name in rotations:
        # Rotate image or not
        if rot_code is not None:
            work_img = cv2.rotate(img, rot_code)
            # Re-detect center on rotated image to be exact
            gray_rot = cv2.cvtColor(work_img, cv2.COLOR_BGR2GRAY)
            circles_rot = cv2.HoughCircles(gray_rot, cv2.HOUGH_GRADIENT, 1, rows / 8,
                                           param1=100, param2=30,
                                           minRadius=int(radius * 0.8), maxRadius=int(radius * 1.2))
            if circles_rot is not None:
                bc_r = np.uint16(np.around(circles_rot))[0][0]
                rcx, rcy, rr = [int(x) for x in bc_r]
            else:
                rh, rw = work_img.shape[:2]
                rcx, rcy, rr = rw // 2, rh // 2, radius
        else:
            work_img = img.copy()
            rcx, rcy, rr = cx, cy, radius
            
        for factor in factors:
            # Warp
            out_w = int(rr * 2 * np.pi * factor)
            out_h = rr
            unwarped = cv2.warpPolar(work_img, (out_w, out_h), (rcx, rcy), rr, cv2.WARP_POLAR_LINEAR)
            
            # Crop outer ring
            crop_outer = unwarped[int(out_h * 0.7):out_h, :]
            
            for flip_code, flip_name in flips:
                final = crop_outer.copy()
                if flip_code is not None:
                    final = cv2.flip(final, flip_code)
                
                # Test OCR
                test_name = f"{rot_name}_{flip_name}_f{factor:.1f}"
                temp_path = os.path.join(results_dir, f"{test_name}.png")
                cv2.imwrite(temp_path, final)
                
                try:
                    res = text_rec.predict(temp_path, batch_size=1)
                    texts = [r.get("rec_text", "") for r in res]
                    recognized = "".join(texts)
                    
                    sim = similarity(recognized, TARGET_TEXT)
                    
                    print(f"Testing {test_name}: '{recognized}' (Sim: {sim:.2f})")
                    
                    if sim > max_sim:
                        max_sim = sim
                        found_best = (test_name, recognized, sim, temp_path)
                        
                    if sim >= MIN_SIMILARITY:
                        print(f"SUCCESS! Found match with {test_name}")
                        # Keep going to find best, but we found a winner
                    else:
                        # Clean up failed images to save space
                        # os.remove(temp_path) 
                        pass
                except:
                    pass

    print("\n--- Search Complete ---")
    if found_best:
        print(f"Best: {found_best[0]} | '{found_best[1]}' | Sim: {found_best[2]:.2f}")
        print(f"Best Image: {found_best[3]}")
    else:
        print("No matches found above 0.0 similarity.")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python brute_force_unwarp.py <image_path> <results_dir>")
        sys.exit(1)
    brute_force_search(sys.argv[1], sys.argv[2])