import torch
import os
import json
import re
import unicodedata
from typing import Optional
from src.trainer.CtrlBase import ControlBase
from PIL import Image
import numpy as np
import torchvision.transforms as T
from tqdm import tqdm
from src.MuSA.GaMuSA_app import text_editing
from src.MuSA.GaMuSA import GaMuSA
from argparse import ArgumentParser
from pytorch_lightning import seed_everything
from utils import create_model, load_state_dict
from PIL import ImageDraw

try:
    # Lazy optional import; only required for --ocr_mode
    from paddleocr import PaddleOCR  # type: ignore
except Exception:  # pragma: no cover
    PaddleOCR = None

# Optional: docTR for DBNet detection backend
try:
    from doctr.io import DocumentFile  # type: ignore
    from doctr.models import ocr_predictor as doctr_ocr_predictor  # type: ignore
except Exception:
    DocumentFile = None
    doctr_ocr_predictor = None


def load_image(image_path, image_height=256, image_width=256):
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    img = Image.open(image_path)
    image = T.ToTensor()(T.Resize((image_height, image_width))(img.convert("RGB")))
    image = image.to(device)
    return image.unsqueeze(0)


def create_parser():
    parser = ArgumentParser()
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--target_height", default=256)
    parser.add_argument("--teaget_width", default=256)
    parser.add_argument("--style_height", default=256)
    parser.add_argument("--style_width", default=256)
    parser.add_argument("--ckpt_path", type=str, default="weights/model.pth")
    parser.add_argument("--dataset_dir", type=str, default="example/")
    parser.add_argument("--output_dir", type=str, default="example_result/")
    parser.add_argument("--starting_layer", default=10, type=int)
    parser.add_argument("--num_inference_steps", type=int, default=50)
    parser.add_argument("--num_sample_per_image", default=1, type=int)
    parser.add_argument("--guidance_scale", default=2, type=float)
    parser.add_argument("--benchmark", default=True)
    # OCR + rule-based editing
    parser.add_argument("--ocr_mode", action="store_true", help="Enable OCR-driven rule-based text replacement")
    parser.add_argument("--input_image", type=str, default=None, help="Path to a single input image when using --ocr_mode")
    parser.add_argument("--rules_json", type=str, default=None, help="JSON file specifying replacements, e.g. {\"NOM\": \"NGUYEN\", \"GIVEN_NAMES\": \"AndrewNguyen\"}")
    # OCR controls
    parser.add_argument("--ocr_lang", type=str, default="fr", help="OCR language (e.g., fr, en, vi)")
    parser.add_argument("--ocr_det_model", type=str, default="", help="Override detection model name (e.g., PP-OCRv5_server_det for DBNet)")
    parser.add_argument("--ocr_rec_model", type=str, default="", help="Override recognition model name (e.g., PP-OCRv5_server_rec)")
    parser.add_argument("--ocr_version", type=str, default="", help="Optional OCR version hint (e.g., PP-OCRv5)")
    parser.add_argument("--detector_backend", type=str, default="paddle", choices=["paddle", "doctr"], help="Detector backend: paddle (default) or doctr DBNet")
    parser.add_argument("--doctr_det_arch", type=str, default="db_resnet50", help="docTR detection architecture (e.g., db_resnet50)")
    return parser


def _strip_accents(text: str) -> str:
    """Remove accents for robust matching."""
    return "".join(c for c in unicodedata.normalize("NFD", text) if unicodedata.category(c) != "Mn")


def _sanitize(text: str) -> str:
    """Uppercase and remove non-letters for coarse rule matching."""
    text = _strip_accents(text)
    text = re.sub(r"[^A-Za-z,\-\s]", "", text)
    return text.strip().upper()


def _array_to_tensor(arr: np.ndarray, device: torch.device) -> torch.Tensor:
    """Convert a HxWxC float[0..1] array to a 1xCxHxW tensor on device."""
    if arr.max() > 1.0:
        arr = arr / 255.0
    tensor = torch.from_numpy(arr).float().permute(2, 0, 1).unsqueeze(0).to(device)
    return tensor


def _run_paddle_ocr(image_path: str, ocr_kwargs: Optional[dict] = None):
    if PaddleOCR is None:
        raise RuntimeError("PaddleOCR is not installed. Please install paddleocr to use --ocr_mode.")
    ocr_kwargs = ocr_kwargs or {}
    # Always enable textline orientation if available
    ocr_kwargs.setdefault("use_textline_orientation", True)
    ocr = PaddleOCR(**ocr_kwargs)
    try:
        return ocr.predict(image_path)
    except Exception:
        # Fallback for older builds
        return ocr.ocr(image_path)


def _run_doctr_dbnet_detect(image_path: str, det_arch: str = "db_resnet50") -> list[list[list[int]]]:
    """Run docTR DBNet detection and return list of polygons [[x,y],...].
    Coordinates returned are absolute pixel values.
    """
    if doctr_ocr_predictor is None or DocumentFile is None:
        raise RuntimeError("python-doctr is not installed. pip install python-doctr")
    try:
        predictor = doctr_ocr_predictor(det_arch=det_arch, reco_arch="crnn_vgg16_bn", pretrained=True)
        doc = DocumentFile.from_images([image_path])
        result = predictor(doc)
        # docTR uses relative geometry [ (x0,y0), ... ] in [0,1]
        im = Image.open(image_path)
        W, H = im.size
        polys: list[list[list[int]]] = []
        # Try multiple result shapes robustly
        try:
            pages = getattr(result, 'pages', None)
            if pages and len(pages) > 0:
                page = pages[0]
                # Walk blocks → lines → words
                for blk in getattr(page, 'blocks', []) or []:
                    for ln in getattr(blk, 'lines', []) or []:
                        for wd in getattr(ln, 'words', []) or []:
                            g = getattr(wd, 'geometry', None)
                            if g is None:
                                continue
                            gl = list(g)
                            if len(gl) == 2:
                                # geometry provided as (xmin,ymin),(xmax,ymax)
                                (xmin, ymin), (xmax, ymax) = gl
                                pts = [
                                    [int(round(xmin * W)), int(round(ymin * H))],
                                    [int(round(xmax * W)), int(round(ymin * H))],
                                    [int(round(xmax * W)), int(round(ymax * H))],
                                    [int(round(xmin * W)), int(round(ymax * H))],
                                ]
                            else:
                                pts = [[int(round(x * W)), int(round(y * H))] for (x, y) in gl]
                            if len(pts) >= 3:
                                polys.append(pts)
        except Exception:
            pass
        try:
            # 2) Export dict may have 'boxes'
            if not polys:
                exp = result.export() if hasattr(result, 'export') else None
                if isinstance(exp, dict):
                    pages = exp.get('pages') or []
                    if pages:
                        p0 = pages[0]
                        boxes = p0.get('boxes') or (p0.get('predictions', {}) or {}).get('boxes')
                        if boxes:
                            for (xmin, ymin, xmax, ymax) in list(boxes):
                                poly = [
                                    [int(round(xmin * W)), int(round(ymin * H))],
                                    [int(round(xmax * W)), int(round(ymin * H))],
                                    [int(round(xmax * W)), int(round(ymax * H))],
                                    [int(round(xmin * W)), int(round(ymax * H))],
                                ]
                                polys.append(poly)
        except Exception:
            pass
        try:
            # 3) Older exports list of polygons [[(x,y)...]]
            if not polys:
                pages = getattr(result, 'pages', None)
                if pages and len(pages) > 0:
                    page = pages[0]
                    if hasattr(page, 'export'):
                        pexp = page.export()
                        # Look for nested geometries
                        for blk in pexp.get('blocks', []) or []:
                            for ln in blk.get('lines', []) or []:
                                for wd in ln.get('words', []) or []:
                                    g = wd.get('geometry')
                                    if g is None:
                                        continue
                                    gl = list(g)
                                    if len(gl) == 2:
                                        (xmin, ymin), (xmax, ymax) = gl
                                        pts = [
                                            [int(round(xmin * W)), int(round(ymin * H))],
                                            [int(round(xmax * W)), int(round(ymin * H))],
                                            [int(round(xmax * W)), int(round(ymax * H))],
                                            [int(round(xmin * W)), int(round(ymax * H))],
                                        ]
                                    else:
                                        pts = [[int(round(x * W)), int(round(y * H))] for (x, y) in gl]
                                    if len(pts) >= 3:
                                        polys.append(pts)
        except Exception:
            pass
        return polys
    except Exception:
        return []


def _extract_text_conf_from_result(result: object) -> list[tuple[str, float]]:
    """Normalize PaddleOCR result to a flat list of (text, conf)."""
    items: list[tuple[str, float]] = []
    try:
        if isinstance(result, dict):
            texts = result.get("rec_texts") or result.get("rec_text") or result.get("texts") or []
            scores = result.get("rec_scores") or result.get("scores") or []
            if isinstance(texts, str):
                texts = [texts]
            if not hasattr(scores, "__len__"):
                scores = [None] * len(texts)
            for i, t in enumerate(list(texts)):
                sc = 0.0
                try:
                    sc = float(scores[i]) if i < len(scores) and scores[i] is not None else 0.0
                except Exception:
                    sc = 0.0
                items.append((str(t), sc))
            return items

        if isinstance(result, list):
            if len(result) == 1 and isinstance(result[0], dict):
                return _extract_text_conf_from_result(result[0])
            # Legacy nested [[[poly], (text, score)], ...]
            if len(result) > 0 and isinstance(result[0], list) and len(result[0]) > 0 and isinstance(result[0][0], list):
                for line in result[0]:
                    if not isinstance(line, (list, tuple)) or len(line) < 2:
                        continue
                    payload = line[1]
                    text = ""
                    score = 0.0
                    if isinstance(payload, (list, tuple)) and len(payload) >= 1:
                        text = str(payload[0])
                        if len(payload) >= 2:
                            try:
                                score = float(payload[1])
                            except Exception:
                                score = 0.0
                    elif isinstance(payload, str):
                        text = payload
                        if len(line) >= 3:
                            try:
                                score = float(line[2])
                            except Exception:
                                score = 0.0
                    items.append((text, score))
                return items
            # Dict list
            if len(result) > 0 and isinstance(result[0], dict):
                for obj in result:
                    text = obj.get("text") or obj.get("rec_text") or ""
                    score = obj.get("score") or obj.get("rec_score") or 0.0
                    try:
                        score = float(score)
                    except Exception:
                        score = 0.0
                    items.append((str(text), float(score)))
                return items
            # Flat list [poly, text, score]
            for line in result:
                if not isinstance(line, (list, tuple)) or len(line) < 2:
                    continue
                text = str(line[1])
                score = 0.0
                if len(line) >= 3:
                    try:
                        score = float(line[2])
                    except Exception:
                        score = 0.0
                items.append((text, score))
            return items
    except Exception:
        return items
    return items


def _extract_items_with_boxes(result: object) -> list[dict]:
    """Return a list of {'text': str, 'score': float, 'box': [[x,y], ...]}.
    Tries to support multiple PaddleOCR output shapes.
    """
    def to_pts(poly) -> list[list[int]]:
        try:
            if poly is None:
                return []
            if hasattr(poly, "tolist"):
                poly = poly.tolist()
            return [[int(p[0]), int(p[1])] for p in list(poly)]
        except Exception:
            return []

    items: list[dict] = []
    try:
        if isinstance(result, dict):
            texts = result.get("rec_texts") or result.get("texts") or []
            scores = result.get("rec_scores") or result.get("scores") or []
            boxes = result.get("dt_polys") or result.get("det_polygons") or result.get("boxes") or []
            if isinstance(texts, str):
                texts = [texts]
            if not hasattr(scores, "__len__"):
                scores = [None] * len(texts)
            n = len(texts)
            for i in range(n):
                txt = str(texts[i])
                try:
                    sc = float(scores[i]) if i < len(scores) and scores[i] is not None else 0.0
                except Exception:
                    sc = 0.0
                bx = to_pts(boxes[i]) if i < len(boxes) else []
                items.append({"text": txt, "score": sc, "box": bx})
            return items
        if isinstance(result, list):
            if len(result) == 1 and isinstance(result[0], dict):
                return _extract_items_with_boxes(result[0])
            if len(result) > 0 and isinstance(result[0], list) and len(result[0]) > 0 and isinstance(result[0][0], list):
                for line in result[0]:
                    if not isinstance(line, (list, tuple)) or len(line) < 2:
                        continue
                    box = to_pts(line[0])
                    payload = line[1]
                    text = ""
                    score = 0.0
                    if isinstance(payload, (list, tuple)) and len(payload) >= 1:
                        text = str(payload[0])
                        if len(payload) >= 2:
                            try:
                                score = float(payload[1])
                            except Exception:
                                score = 0.0
                    elif isinstance(payload, str):
                        text = payload
                        if len(line) >= 3:
                            try:
                                score = float(line[2])
                            except Exception:
                                score = 0.0
                    items.append({"text": text, "score": score, "box": box})
                return items
            if len(result) > 0 and isinstance(result[0], dict):
                for obj in result:
                    text = obj.get("text") or obj.get("rec_text") or ""
                    score = obj.get("score") or obj.get("rec_score") or 0.0
                    box = obj.get("box") or obj.get("poly") or obj.get("det_polygons") or []
                    try:
                        score = float(score)
                    except Exception:
                        score = 0.0
                    items.append({"text": str(text), "score": float(score), "box": to_pts(box)})
                return items
    except Exception:
        return items
    return items


def _bbox_center(box: list[list[int]]) -> tuple[float, float]:
    xs = [p[0] for p in box]
    ys = [p[1] for p in box]
    return (float(sum(xs)) / max(len(xs), 1), float(sum(ys)) / max(len(ys), 1))


def _distance(p1: tuple[float, float], p2: tuple[float, float]) -> float:
    return ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5


def _expand_box(box: list[list[int]], img_w: int, img_h: int, pad_ratio: float = 0.15) -> tuple[int, int, int, int]:
    xs = [p[0] for p in box]
    ys = [p[1] for p in box]
    x0, y0, x1, y1 = min(xs), min(ys), max(xs), max(ys)
    w = x1 - x0
    h = y1 - y0
    pad_x = int(w * pad_ratio)
    pad_y = int(h * pad_ratio)
    x0 = max(0, x0 - pad_x)
    y0 = max(0, y0 - pad_y)
    x1 = min(img_w, x1 + pad_x)
    y1 = min(img_h, y1 + pad_y)
    return x0, y0, x1, y1


def _poly_to_rect(box: list[list[int]]) -> list[int]:
    if not box:
        return [0, 0, 0, 0]
    xs = [p[0] for p in box]
    ys = [p[1] for p in box]
    return [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]


def _pil_to_tensor(img: Image.Image, device: torch.device, image_height: int = 256, image_width: int = 256) -> torch.Tensor:
    arr = T.ToTensor()(T.Resize((image_height, image_width))(img.convert("RGB")))
    return arr.unsqueeze(0).to(device)


def _run_local_edit(pipeline: GaMuSA, crop_img: Image.Image, old_text: str, new_text: str,
                    starting_layer: int, num_inference_steps: int, guidance_scale: float,
                    device: torch.device) -> Image.Image:
    source_tensor = _pil_to_tensor(crop_img, device)
    style_tensor = _pil_to_tensor(crop_img, device)
    result = text_editing(
        pipeline,
        source_tensor,
        style_tensor,
        old_text,
        new_text,
        starting_layer=starting_layer,
        ddim_steps=num_inference_steps,
        scale=guidance_scale,
    )
    _, out_np = result[:]
    out_img = Image.fromarray((out_np * 255).astype(np.uint8))
    return out_img


def _match_rules_to_regions(ocr_items: list[dict], rules: dict[str, str]) -> list[dict]:
    """Return a list of {old_text, new_text, box} for localized edits.
    Uses label anchors like NOM/SURNAME and PRENOMS/GIVEN NAMES to choose nearby values.
    """
    # Build anchors
    def matches_any(text: str, keywords: list[str]) -> bool:
        nt = _sanitize(text)
        return any(k in nt for k in keywords)

    anchors: dict[str, list[dict]] = {
        "NOM": [],
        "GIVEN_NAMES": [],
    }
    for it in ocr_items:
        t = it.get("text", "")
        if matches_any(t, ["NOM", "SURNAME"]):
            anchors["NOM"].append(it)
        if matches_any(t, ["PRENOM", "GIVEN", "GIVENNAMES", "PRENOMS", "GIVEN NAMES"]):
            anchors["GIVEN_NAMES"].append(it)

    # Candidates for values
    def is_surname_candidate(t: str) -> bool:
        nt = _sanitize(t)
        if any(k in nt for k in ["NOM", "SURNAME", "PRENOM", "GIVEN", "GIVENNAMES", "PRENOMS"]):
            return False
        return len(nt) >= 3 and nt == nt.upper() and re.fullmatch(r"[A-Z\-\s]+", nt or "") is not None

    def is_given_candidate(t: str) -> bool:
        nt = _sanitize(t)
        if any(k in nt for k in ["NOM", "SURNAME", "PRENOM", "GIVEN", "GIVENNAMES", "PRENOMS"]):
            return False
        # mixed or lowercase with separators (avoid all-caps)
        return any(c.islower() for c in t) or ("-" in t or "," in t)

    tasks: list[dict] = []
    for key, new_val in rules.items():
        key_u = key.strip().upper()
        if key_u not in ("NOM", "GIVEN_NAMES"):
            continue
        anchor_list = anchors.get(key_u, [])
        best_item = None
        best_dist = 1e9
        for it in ocr_items:
            t = it.get("text", "")
            if key_u == "NOM" and not is_surname_candidate(t):
                continue
            if key_u == "GIVEN_NAMES" and not is_given_candidate(t):
                continue
            # compute distance to nearest anchor if exists; otherwise prefer high score
            dist = 0.0
            if anchor_list:
                c = _bbox_center(it.get("box") or [])
                dists = []
                for an in anchor_list:
                    dists.append(_distance(c, _bbox_center(an.get("box") or [])))
                dist = min(dists) if dists else 0.0
            else:
                dist = 0.0
            if dist < best_dist:
                best_dist = dist
                best_item = it
        if best_item is not None and best_item.get("box"):
            tasks.append({"old_text": best_item["text"], "new_text": new_val, "box": best_item["box"]})
    return tasks


def _build_form_entries(ocr_items: list[dict]) -> list[dict]:
    """Build a DocVQA-like form list with question/answer/other and linking.
    We focus on key pairs: NOM ↔ value; GIVEN NAMES ↔ value.
    """
    def matches_any(text: str, keywords: list[str]) -> bool:
        nt = _sanitize(text)
        return any(k in nt for k in keywords)

    items = []
    for it in ocr_items:
        items.append({
            "text": it.get("text", ""),
            "rect": _poly_to_rect(it.get("box") or []),
            "score": float(it.get("score") or 0.0),
            "box": it.get("box") or [],
        })

    questions_idx: dict[str, list[int]] = {"NOM": [], "GIVEN_NAMES": []}
    for idx, obj in enumerate(items):
        t = obj["text"]
        if matches_any(t, ["NOM", "SURNAME"]):
            questions_idx["NOM"].append(idx)
        if matches_any(t, ["PRENOM", "GIVEN", "GIVENNAMES"]):
            questions_idx["GIVEN_NAMES"].append(idx)

    # find nearest answers
    def pick_answer(candidates: list[int], predicate) -> Optional[int]:
        best = None
        bestd = 1e9
        # candidate list of value boxes
        value_ids = [i for i in range(len(items)) if predicate(items[i]["text"])]
        for qi in candidates:
            qc = _bbox_center(items[qi]["box"]) if items[qi]["box"] else (0.0, 0.0)
            for vi in value_ids:
                vc = _bbox_center(items[vi]["box"]) if items[vi]["box"] else (0.0, 0.0)
                d = _distance(qc, vc)
                if d < bestd:
                    bestd = d
                    best = (qi, vi)
        if best is None:
            return None
        return best[1]

    def is_surname_candidate(t: str) -> bool:
        nt = _sanitize(t)
        return len(nt) >= 3 and nt == nt.upper() and re.fullmatch(r"[A-Z\-\s]+", nt or "") is not None

    def is_given_candidate(t: str) -> bool:
        return any(c.islower() for c in t) or ("-" in t or "," in t)

    links: list[tuple[int, int]] = []
    if questions_idx["NOM"]:
        ans = pick_answer(questions_idx["NOM"], is_surname_candidate)
        if ans is not None:
            links.append((questions_idx["NOM"][0], ans))
    if questions_idx["GIVEN_NAMES"]:
        ans = pick_answer(questions_idx["GIVEN_NAMES"], is_given_candidate)
        if ans is not None:
            links.append((questions_idx["GIVEN_NAMES"][0], ans))

    # Build form list with ids and labels
    form: list[dict] = []
    id_map: dict[int, int] = {}
    next_id = 0
    question_set = set(q for q, _ in links)
    answer_set = set(a for _, a in links)
    for i, obj in enumerate(items):
        label = "other"
        if i in question_set:
            label = "question"
        elif i in answer_set:
            label = "answer"
        entry = {
            "box": obj["rect"],
            "text": obj["text"],
            "label": label,
            "words": [{"box": obj["rect"], "text": obj["text"]}],
            "linking": [],
            "id": next_id,
        }
        id_map[i] = next_id
        next_id += 1
        form.append(entry)

    # apply linking
    for q_i, a_i in links:
        if q_i in id_map and a_i in id_map:
            form[id_map[q_i]]["linking"].append([id_map[q_i], id_map[a_i]])
            form[id_map[a_i]]["linking"].append([id_map[q_i], id_map[a_i]])

    return form


def _pick_field_values(ocr_result):
    """Heuristic extraction of key fields from OCR output.

    Returns a dict with potential values for keys: NOM, GIVEN_NAMES.
    """
    # Normalize OCR results
    texts = _extract_text_conf_from_result(ocr_result)

    # Candidates
    blacklist = {
        "REPUBLIQUE FRANCAISE",
        "IDENTITY CARD",
        "IDENTITY",
        "FR",
        "FRA",
        "PARIS",
        "SIGNATURE",
        "DATE DE NAISS",
        "DATE D EXPIR",
        "NUM DU DOCUMENT",
        "NOM D USAGE",
        "NATIONALITE",
        "LIEU DE NAISSANCE",
        "SEXE",
        "SURNAME",
        "IDENTITYCARD",
    }

    cleaned = [(_sanitize(t), t, c) for t, c in texts]

    # Heuristic for surname (NOM): all-uppercase word, not in blacklist
    surname_candidates = []
    for norm, original, conf in cleaned:
        if len(norm) >= 3 and norm == norm.upper() and re.fullmatch(r"[A-Z\-\s]+", norm or ""):
            if norm not in blacklist and not re.search(r"\b(FR|FRA|PARIS)\b", norm):
                surname_candidates.append((conf, original))
    surname = max(surname_candidates, default=(0.0, ""))[1]

    # Heuristic for given names: contains at least one lowercase in original and long enough
    given_candidates = []
    for norm, original, conf in cleaned:
        if any(ch.islower() for ch in original) and len(original) >= 5:
            if "/" not in original and not original.isdigit():
                given_candidates.append((conf, original))
    # Prefer the longest, then highest confidence
    if given_candidates:
        given_candidates = sorted(given_candidates, key=lambda x: (len(x[1]), x[0]), reverse=True)
        given_names = given_candidates[0][1]
    else:
        given_names = ""

    return {"NOM": surname.strip(), "GIVEN_NAMES": given_names.strip()}


def _load_rules(path: Optional[str]):
    if path is None or not os.path.exists(path):
        # Demo defaults
        return {"NOM": "NGUYEN", "GIVEN_NAMES": "AndrewNguyen"}
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    # Normalize keys
    norm = {k.strip().upper(): v for k, v in data.items()}
    return norm


def main(opt):
    cfg_path = 'configs/inference.yaml'
    model = create_model(cfg_path).cuda()
    model.load_state_dict(load_state_dict(opt.ckpt_path), strict=False)
    model.eval()

    # Common pipeline pieces
    monitor_cfg = {
        "max_length": 25,
        "loss_weight": 1.,
        "attention": 'position',
        "backbone": 'transformer',
        "backbone_ln": 3,
        "checkpoint": "weights/vision_model.pth",
        "charset_path": "src/module/abinet/data/charset_36.txt"
    }
    pipeline = GaMuSA(model, monitor_cfg)
    output_dir = opt.output_dir
    os.makedirs(output_dir, exist_ok=True)
    seed = opt.seed
    starting_layer = opt.starting_layer
    guidance_scale = opt.guidance_scale
    num_sample_per_image = opt.num_sample_per_image
    num_inference_steps = opt.num_inference_steps
    seed_everything(seed)

    if opt.ocr_mode:
        if opt.input_image is None:
            raise ValueError("--input_image is required when --ocr_mode is enabled")
        image_path = opt.input_image
        image_name = os.path.basename(image_path)
        w, h = Image.open(image_path).size

        # OCR detection and rule mapping
        # Build OCR options based on CLI
        ocr_opts = {"lang": getattr(opt, "ocr_lang", "fr")}
        if getattr(opt, "ocr_version", ""):
            ocr_opts["ocr_version"] = opt.ocr_version
        if getattr(opt, "ocr_det_model", ""):
            ocr_opts["text_detection_model_name"] = opt.ocr_det_model
        if getattr(opt, "ocr_rec_model", ""):
            ocr_opts["text_recognition_model_name"] = opt.ocr_rec_model
        # Choose detector backend
        if getattr(opt, 'detector_backend', 'paddle') == 'doctr':
            # 1) detect with docTR DBNet
            dbnet_polys = _run_doctr_dbnet_detect(image_path, getattr(opt, 'doctr_det_arch', 'db_resnet50'))
            # 2) recognize with PaddleOCR on crops
            # Create Paddle recognizer instance (rec-only if supported)
            rec_kwargs = {"lang": ocr_opts.get("lang", "fr")}
            if ocr_opts.get("text_recognition_model_name"):
                rec_kwargs["text_recognition_model_name"] = ocr_opts["text_recognition_model_name"]
            rec_ocr = PaddleOCR(**rec_kwargs) if PaddleOCR is not None else None
            img = Image.open(image_path).convert("RGB")
            items = []
            for poly in dbnet_polys:
                rect = _poly_to_rect(poly)
                x0, y0, x1, y1 = rect
                crop = img.crop((x0, y0, x1, y1))
                text_val = ""
                score_val = 0.0
                if rec_ocr is not None:
                    try:
                        try:
                            rec_res = rec_ocr.predict(np.array(crop))
                        except Exception:
                            rec_res = rec_ocr.ocr(np.array(crop))
                        # extract best text
                        flat = _extract_text_conf_from_result(rec_res)
                        if flat:
                            text_val, score_val = flat[0]
                    except Exception:
                        pass
                items.append({"text": text_val, "score": score_val, "box": poly})
            ocr_items = items
            # For anchors, we still need global text; run a lightweight OCR on full image to find labels
            full_res = _run_paddle_ocr(image_path, ocr_opts)
            detected_values = _pick_field_values(full_res)
            # Save raw OCR like before
            ocr_result = {
                "det_polygons": dbnet_polys,
                "items": items,
            }
        else:
            ocr_result = _run_paddle_ocr(image_path, ocr_opts)
            detected_values = _pick_field_values(ocr_result)
            # Extract full items with boxes for logging/analysis
            ocr_items = _extract_items_with_boxes(ocr_result)
        # Save OCR JSON
        base = os.path.splitext(image_name)[0]
        ocr_json_path = os.path.join(output_dir, f"{base}_ocr.json")
        with open(ocr_json_path, "w", encoding="utf-8") as jf:
            json.dump({
                "image": os.path.abspath(image_path),
                "items": ocr_items,
                "detected_fields": detected_values,
            }, jf, ensure_ascii=False, indent=2)
        # Build DocVQA-like form and save
        form_entries = _build_form_entries(ocr_items)
        form_json_path = os.path.join(output_dir, f"{base}_form.json")
        with open(form_json_path, "w", encoding="utf-8") as fj:
            json.dump({"form": form_entries}, fj, ensure_ascii=False, indent=2)

        # Also save a quick overlay visualization
        try:
            im = Image.open(image_path).convert("RGB")
            draw = ImageDraw.Draw(im)
            for it in ocr_items:
                pts = it.get("box") or []
                if len(pts) >= 3:
                    poly = [tuple(p) for p in pts]
                    draw.polygon(poly, outline=(255, 0, 0), width=2)
            im.save(os.path.join(output_dir, f"{base}_ocr.png"))
        except Exception:
            pass
        rules = _load_rules(opt.rules_json)

        # Prepare initial tensors
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        current_tensor = load_image(image_path)

        # Determine localized edit regions for rules
        tasks = _match_rules_to_regions(ocr_items, rules)
        # If heuristics failed to link, force simple fallbacks using detected_values
        if not tasks:
            fallback = []
            # choose the closest all-caps box to NOM label
            for key_u, new_val in (("NOM", rules.get("NOM")), ("GIVEN_NAMES", rules.get("GIVEN_NAMES"))):
                if not new_val:
                    continue
                # rough selection by text similarity
                for it in ocr_items:
                    t = it.get("text", "")
                    if key_u == "NOM" and t and t.isupper():
                        fallback.append({"old_text": t, "new_text": new_val, "box": it.get("box")})
                        break
                    if key_u == "GIVEN_NAMES" and t and (any(c.islower() for c in t) or ("-" in t or "," in t)):
                        fallback.append({"old_text": t, "new_text": new_val, "box": it.get("box")})
                        break
            if fallback:
                tasks = fallback
        # Apply localized edits: crop → edit → paste back
        base_img = Image.open(image_path).convert("RGB")
        img_w, img_h = base_img.size
        composed = base_img.copy()
        for task in tasks:
            box = task.get("box") or []
            if len(box) < 3:
                continue
            x0, y0, x1, y1 = _expand_box(box, img_w, img_h, pad_ratio=0.15)
            crop = base_img.crop((x0, y0, x1, y1))
            edited_crop = _run_local_edit(
                pipeline,
                crop,
                task["old_text"],
                task["new_text"],
                starting_layer,
                num_inference_steps,
                guidance_scale,
                device,
            )
            edited_crop = edited_crop.resize((x1 - x0, y1 - y0))
            composed.paste(edited_crop, (x0, y0))

        # Save result
        composed.save(os.path.join(output_dir, image_name))
        return

    # Default dataset-driven flow (original behavior)
    dataset_dir = opt.dataset_dir
    style_dir = os.path.join(dataset_dir, 'i_s')
    style_images_path = {image_name: os.path.join(style_dir, image_name) for image_name in os.listdir(style_dir)}
    style_txt = os.path.join(dataset_dir, 'i_s.txt')
    style_dict = {}
    with open(style_txt, 'r') as f:
        for line in f.readlines():
            if line != '\n':
                image_name, text = line.strip().split(' ')[:]
                style_dict[image_name] = text
    target_txt = os.path.join(dataset_dir, 'i_t.txt')
    target_dict = {}
    with open(target_txt, 'r') as f:
        for line in f.readlines():
            if line != '\n':
                image_name, text = line.strip().split(' ')[:]
                target_dict[image_name] = text

    for i in tqdm(range(len(list(style_images_path.keys())))):
        image_name = list(style_images_path.keys())[i]
        image_path = style_images_path[image_name]
        style_text = style_dict[image_name]
        target_text = target_dict[image_name]
        w, h = Image.open(image_path).size
        source_image = load_image(image_path)
        style_image = load_image(image_path)
        result = text_editing(
            pipeline,
            source_image,
            style_image,
            style_text,
            target_text,
            starting_layer=starting_layer,
            ddim_steps=num_inference_steps,
            scale=guidance_scale,
        )
        if opt.benchmark:
            _, GaMuSA_image = result[:]
            GaMuSA_image = Image.fromarray((GaMuSA_image * 255).astype(np.uint8)).resize((w, h))
            GaMuSA_image.save(os.path.join(output_dir, image_name))

        else:
            save_dir = os.path.join(output_dir, image_name.split('.')[0])
            os.makedirs(save_dir, exist_ok=True)
            reconstruction_image, GaMuSA_image = result[:]
            reconstruction_image = Image.fromarray((reconstruction_image * 255).astype(np.uint8)).resize((w, h))
            GaMuSA_image = Image.fromarray((GaMuSA_image * 255).astype(np.uint8)).resize((w, h))
            reconstruction_image.save(os.path.join(save_dir, 'recons_' + style_text + '.png'))
            GaMuSA_image.save(os.path.join(save_dir, 'GaMUSA_' + target_text + '.png'))


if __name__ == "__main__":
    parser = create_parser()
    opt = parser.parse_args()
    main(opt)
