import torch
from PIL import Image, ImageDraw
from typing import List, Dict, Any
from doclayout_yolo import YOLOv10
from huggingface_hub import hf_hub_download
from pix2tex import cli as pix2tex
from pathlib import Path
import json
import fitz



class DetectionOCR:
    def __init__(self, device=None):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        print(f"[DetectionOCR] Using {self.device}")

        # ----------------------------
        # Load DocLayout-YOLO weight
        # ----------------------------
        print("[DetectionOCR] Loading DocLayout-YOLO weight...")

        weight_path = hf_hub_download(
            repo_id="juliozhao/DocLayout-YOLO-DocStructBench",
            filename="doclayout_yolo_docstructbench_imgsz1024.pt",
        )

        self.yolo = YOLOv10(weight_path).to(self.device)
        self.names = self.yolo.names

        # ----------------------------
        # Load Pix2TeX
        # ----------------------------
        print("[DetectionOCR] Loading Pix2TeX...")
        self.pix2tex = pix2tex.LatexOCR()

    # ------------------------------------------------------------
    # 1. DETECTION
    # ------------------------------------------------------------

    def detect_formula_regions(self, image_path: str) -> List[Dict[str, Any]]:
        """
        Detect formula regions using YOLO
        """

        det = self.yolo.predict(
            image_path,
            imgsz=1024,
            conf=0.25,
            device=self.device,
            verbose=False
        )[0]

        xyxy = det.boxes.xyxy.cpu().numpy()
        cls_ids = det.boxes.cls.cpu().numpy().astype(int)
        scores = det.boxes.conf.cpu().numpy()

        results = []
        for i, cid in enumerate(cls_ids):
            label = self.names[cid]
            if "formula" in label.lower():
                x1, y1, x2, y2 = xyxy[i]
                results.append({
                    "bbox": [float(x1), float(y1), float(x2), float(y2)],
                    "label": label,
                    "score": float(scores[i])
                })
        return results

    # ------------------------------------------------------------
    # 2. OCR
    # ------------------------------------------------------------

    def extract_latex_from_region(self, page_img: Image.Image, region) -> str:
        x1, y1, x2, y2 = map(int, region["bbox"])
        crop = page_img.crop((x1, y1, x2, y2))

        try:
            result = self.pix2tex(crop)
            if isinstance(result, dict):
                return result.get("latex", "")
            return str(result)
        except Exception as e:
            print("[OCR] Error:", e)
            return ""
        
    # ------------------------------------------------------------
    # 3. DRAW BBOX
    # ------------------------------------------------------------

    def draw_bounding_boxes(self, image_path, regions, output_path):
        img = Image.open(image_path).convert("RGB")
        draw = ImageDraw.Draw(img)

        for i, region in enumerate(regions):
            x1, y1, x2, y2 = map(int, region["bbox"])
            draw.rectangle([x1, y1, x2, y2], outline="red", width=4)
            draw.text((x1, max(0, y1-20)), f"Formula {i+1}", fill="red")

        img.save(output_path)
        print(f"[BBox] Saved:", output_path)

    # ------------------------------------------------------------
    # 4. SAVE TEX
    # ------------------------------------------------------------

    def save_latex_to_tex(self, regions, output_path):

        lines = [
            r"\documentclass{article}",
            r"\usepackage{amsmath, amssymb}",
            r"\begin{document}"
        ]

        for i, r in enumerate(regions):
            lines.append(r"\begin{equation*}")
            lines.append(r.get("raw_latex", ""))
            lines.append(r"\end{equation*}")
            lines.append("")

        lines.append(r"\end{document}")
        Path(output_path).write_text("\n".join(lines), encoding="utf-8")
        print("[TeX] Saved:", output_path)

    # ------------------------------------------------------------
    # 5. SAVE JSON
    # ------------------------------------------------------------

    def save_results_to_json(self, regions, output_path):
        Path(output_path).write_text(
            json.dumps({"formulas": regions, "total": len(regions)}, indent=2, ensure_ascii=False),
            encoding="utf-8"
        )
        print("[JSON] Saved:", output_path)

    # ------------------------------------------------------------
    # 6. PROCESS IMAGE
    # ------------------------------------------------------------

    def process_image(self, image_path, output_dir):

        Path(output_dir).mkdir(exist_ok=True)
        img_name = Path(image_path).stem

        page = Image.open(image_path).convert("RGB")
        regions = self.detect_formula_regions(image_path)

        results = []
        for region in regions:
            latex = self.extract_latex_from_region(page, region)
            region["raw_latex"] = latex
            results.append(region)

        # Save bbox image
        bbox_path = f"{output_dir}/{img_name}_bbox.png"
        self.draw_bounding_boxes(image_path, results, bbox_path)

        # Save tex
        tex_path = f"{output_dir}/{img_name}_formulas.tex"
        self.save_latex_to_tex(results, tex_path)

        # Save json
        json_path = f"{output_dir}/{img_name}_results.json"
        self.save_results_to_json(results, json_path)

        return results

    # ------------------------------------------------------------
    # 7. PROCESS PDF
    # ------------------------------------------------------------

    def process_pdf(self, pdf_path, output_dir="output"):
        Path(output_dir).mkdir(exist_ok=True)
        tmp_dir = Path(output_dir) / "pages"
        tmp_dir.mkdir(exist_ok=True)

        doc = fitz.open(pdf_path)
        all_results = {}

        for i in range(len(doc)):
            page = doc[i]
            pix = page.get_pixmap(dpi=300)

            img_file = tmp_dir / f"page_{i+1}.png"
            pix.save(str(img_file))

            print(f"\n[PDF] Page {i+1}")
            page_results = self.process_image(str(img_file), output_dir)

            all_results[f"page_{i+1}"] = {
                "total": len(page_results),
                "formulas": page_results
            }

        overall_json = Path(output_dir) / "overall_ocr_results.json"
        overall_json.write_text(json.dumps(all_results, indent=2, ensure_ascii=False), "utf-8")

        print("\n== Done ==")
        print("Output:", output_dir)
        return all_results