#!/usr/bin/env python3
"""
Self-contained inference example for invoice-layoutlmv3-multidomain.

Downloads the model from HuggingFace, runs EasyOCR + LayoutLMv3, and prints
extracted fields as JSON.

Requirements:
    pip install transformers torch easyocr huggingface_hub Pillow

Usage:
    python inference_example.py path/to/invoice.png               # auto-detect domain
    python inference_example.py path/to/invoice.png --domain general

Domains: general, receipt, medical, insurance, logistics
"""

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from huggingface_hub import snapshot_download
from transformers import AutoModelForTokenClassification, LayoutLMv3Processor

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

REPO_ID = "rhlprj/invoice-layoutlmv3-multidomain"
MAX_SEQ_LENGTH = 512

# Domain routing keywords (used when --domain is not specified)
DOMAIN_KEYWORDS = {
    "insurance": [
        "eob", "claim", "policy", "insurer", "deductible", "copay",
        "eligible", "approved", "not covered", "patient liability",
        "insurance pays", "explanation of benefits",
    ],
    "medical": [
        "patient", "hospital", "doctor", "diagnosis", "uhid", "blood group",
        "department", "procedure", "attending", "discharge", "admission",
    ],
    "logistics": [
        "waybill", "awb", "bol", "freight", "carrier", "consignee",
        "shipper", "container", "incoterms", "port of loading",
        "origin port", "destination port", "gross weight",
    ],
    "receipt": [
        "receipt", "cash", "change", "subtotal", "cashier", "store",
    ],
    "general": [
        "invoice", "due date", "payment terms", "bank", "account",
    ],
}

# Anchor keywords per domain+field for stripping label prefixes from OCR text.
# When EasyOCR merges "Invoice #: FRT-2025-5862" into one box, the model tags
# the whole box correctly but the extracted text includes the label prefix.
# These anchors let us strip "Invoice #:" to get just "FRT-2025-5862".
ANCHORS = {
    "general": {
        "invoice_number": ["invoice no", "invoice number", "invoice #", "invoice"],
        "invoice_date": ["invoice date", "date"],
        "due_date": ["due date", "due"],
        "subtotal": ["subtotal", "sub total"],
        "total": ["total", "grand total", "amount due"],
        "tax": ["tax", "gst", "vat"],
        "payment_terms": ["payment terms", "terms"],
        "bank_name": ["bank", "bank name"],
        "account_number": ["account", "account number", "a/c"],
        "vendor_email": ["email", "e-mail"],
        "vendor_phone": ["phone", "tel", "mobile", "contact"],
    },
    "receipt": {
        "subtotal": ["subtotal", "sub total"],
        "tax": ["tax", "gst", "vat", "service"],
        "total": ["total", "grand total"],
        "cash_paid": ["cash", "paid", "cash paid"],
        "change": ["change"],
        "invoice_date": ["date"],
    },
    "medical": {
        "bill_number": ["bill no", "bill number", "bill #", "invoice"],
        "bill_date": ["bill date", "date"],
        "patient_name": ["patient", "patient name", "name"],
        "patient_age": ["age"],
        "patient_gender": ["gender", "sex"],
        "hospital_name": ["hospital"],
        "attending_doctor": ["doctor", "physician", "consultant", "dr"],
        "department": ["department", "dept"],
        "diagnosis": ["diagnosis"],
        "patient_type": ["patient type", "type"],
        "subtotal": ["subtotal", "sub total"],
        "total_due": ["total due", "total", "amount due"],
        "uhid": ["uhid"],
        "gst": ["gst", "tax"],
        "insurance_covered": ["insurance", "insurance covered", "covered"],
        "patient_blood_group": ["blood group", "blood"],
    },
    "insurance": {
        "eob_number": ["eob no", "eob number", "eob"],
        "claim_number": ["claim no", "claim number", "claim"],
        "policy_number": ["policy no", "policy number", "policy"],
        "claim_date": ["claim date"],
        "process_date": ["process date", "processed"],
        "claim_type": ["claim type", "type"],
        "claim_status": ["status", "claim status"],
        "insurer_name": ["insurer", "insurance company"],
        "policy_holder": ["policy holder", "holder", "insured"],
        "patient_name": ["patient", "patient name"],
        "patient_age": ["age"],
        "hospital_name": ["hospital", "provider"],
        "diagnosis": ["diagnosis"],
        "total_billed": ["total billed", "billed", "claimed"],
        "eligible_amount": ["eligible"],
        "approved_amount": ["approved"],
        "insurance_pays": ["insurance pays", "payable", "insurer pays"],
        "patient_liability": ["patient liability", "liability", "patient pays"],
        "not_covered": ["not covered", "disallowed"],
        "deductible": ["deductible"],
        "admission_date": ["admission", "admission date"],
        "discharge_date": ["discharge", "discharge date"],
    },
    "logistics": {
        "invoice_number": ["invoice no", "invoice number", "invoice"],
        "invoice_date": ["invoice date", "date"],
        "due_date": ["due date", "due"],
        "waybill_number": ["waybill", "awb", "bol", "tracking"],
        "freight_type": ["freight type", "service", "mode"],
        "carrier_name": ["carrier", "shipper"],
        "cargo_description": ["cargo", "commodity", "goods"],
        "gross_weight": ["weight", "gross weight"],
        "num_packages": ["packages", "pieces", "pkgs"],
        "origin_port": ["origin", "from", "port of loading"],
        "destination_port": ["destination", "to", "port of discharge"],
        "shipper_name": ["shipper", "shipper name", "from"],
        "consignee_name": ["consignee", "consignee name", "to"],
        "container_number": ["container", "container no", "container number", "cntr"],
        "notify_party": ["notify", "notify party", "also notify"],
        "volume": ["volume", "cbm", "measurement"],
        "payment_mode": ["payment", "payment mode", "terms"],
        "incoterms": ["incoterms", "incoterm"],
        "subtotal": ["subtotal", "sub total"],
        "total": ["total", "grand total"],
        "po_number": ["po no", "po number", "po", "purchase order"],
        "gst": ["gst", "tax"],
    },
}


# ---------------------------------------------------------------------------
# OCR (EasyOCR)
# ---------------------------------------------------------------------------

_easyocr_reader = None


def _get_reader():
    global _easyocr_reader
    if _easyocr_reader is None:
        import easyocr
        _easyocr_reader = easyocr.Reader(["en"], gpu=torch.cuda.is_available())
    return _easyocr_reader


def run_ocr(image: Image.Image) -> List[Dict[str, Any]]:
    """Run EasyOCR on a PIL image, return list of {text, bbox, confidence}.

    bbox is [x0, y0, x1, y1] in pixel coordinates.
    """
    reader = _get_reader()
    np_img = np.array(image)
    raw = reader.readtext(np_img, canvas_size=1280, min_size=5)

    words = []
    for coords, text, conf in raw:
        text = text.strip()
        if not text:
            continue
        xs = [p[0] for p in coords]
        ys = [p[1] for p in coords]
        words.append({
            "text": text,
            "bbox": [min(xs), min(ys), max(xs), max(ys)],
            "confidence": float(conf),
        })
    return words


# ---------------------------------------------------------------------------
# Domain routing
# ---------------------------------------------------------------------------

def route_domain(ocr_words: List[Dict[str, Any]]) -> str:
    """Score OCR text against domain keywords, return best domain."""
    full_text = " ".join(w["text"] for w in ocr_words).lower()
    scores = {}
    for domain, keywords in DOMAIN_KEYWORDS.items():
        scores[domain] = sum(1 for kw in keywords if kw in full_text)
    best = max(scores, key=scores.get)
    return best if scores[best] > 0 else "general"


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

_model_cache: Dict[str, Dict[str, Any]] = {}


def load_model(domain: str, models_dir: str = "models") -> Dict[str, Any]:
    """Load model + processor + label maps for a domain. Caches in memory."""
    if domain in _model_cache:
        return _model_cache[domain]

    model_path = Path(models_dir) / domain
    if not model_path.exists():
        print(f"Downloading models from {REPO_ID}...")
        snapshot_download(REPO_ID, local_dir=models_dir)

    # Label maps
    with open(model_path / "label_maps.json", "r", encoding="utf-8") as f:
        lm = json.load(f)
    id2label = {int(k): v for k, v in lm["id2label"].items()}

    # Processor (apply_ocr=False — we supply our own OCR words + boxes)
    processor = LayoutLMv3Processor.from_pretrained(str(model_path), apply_ocr=False)

    # Model
    model = AutoModelForTokenClassification.from_pretrained(str(model_path))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device).eval()

    entry = {"model": model, "processor": processor, "id2label": id2label, "device": device}
    _model_cache[domain] = entry
    return entry


# ---------------------------------------------------------------------------
# Subword → word alignment (CRITICAL — this is what the reviewer got wrong)
# ---------------------------------------------------------------------------

def map_token_preds_to_words(
    encoding,
    token_logits: torch.Tensor,  # [seq_len, num_labels]
    n_words: int,
) -> Tuple[List[int], List[float]]:
    """Map subword predictions back to word-level predictions.

    LayoutLMv3's tokenizer splits words into subword tokens. For example:
        "INV-2025-00782" -> ["IN", "V", "-", "2025", "-", "007", "82"]  (7 tokens)

    The model predicts one label per subword token. We need ONE label per
    original word. Strategy: take the FIRST subword's prediction for each word.

    encoding.word_ids(0) returns a list mapping each token position to its
    source word index (None for special tokens like [CLS], [SEP], [PAD]).

    WRONG approach (what the reviewer did):
        preds[1:len(words)+1]  # Assumes 1 token per word — BROKEN

    CORRECT approach (what we do):
        Use word_ids() to find the first subword for each word, take that
        subword's prediction.

    Returns:
        pred_label_ids: List[int] of length n_words
        pred_confs: List[float] of length n_words (softmax confidence)
    """
    word_ids = encoding.word_ids(0)
    probs = F.softmax(token_logits, dim=-1)

    # Find first subword token index for each word
    first_subword: Dict[int, int] = {}
    for tok_idx, w_id in enumerate(word_ids):
        if w_id is not None and w_id not in first_subword:
            first_subword[w_id] = tok_idx

    pred_label_ids = []
    pred_confs = []

    for w_idx in range(n_words):
        if w_idx in first_subword:
            tok_idx = first_subword[w_idx]
            label_id = int(token_logits[tok_idx].argmax().item())
            conf = float(probs[tok_idx, label_id].item())
        else:
            # Word was truncated (beyond 512 tokens) — default to "O"
            label_id = 0
            conf = 1.0
        pred_label_ids.append(label_id)
        pred_confs.append(conf)

    return pred_label_ids, pred_confs


# ---------------------------------------------------------------------------
# BIO span merging
# ---------------------------------------------------------------------------

def merge_spans(
    words: List[str],
    bboxes: List[List[int]],
    pred_label_ids: List[int],
    id2label: Dict[int, str],
    confidences: List[float],
) -> List[Dict[str, Any]]:
    """Walk word-level BIO predictions and merge consecutive tokens into spans.

    BIO scheme:
        B-field_name = beginning of a new span for field_name
        I-field_name = inside/continuation of the current span
        O            = outside any span

    Returns list of span dicts with keys:
        label_field, text, bbox, confidence, word_idxs
    """
    spans = []
    current = None

    for idx, (word, bbox, label_id) in enumerate(zip(words, bboxes, pred_label_ids)):
        label = id2label.get(label_id, "O")

        if label == "O":
            if current:
                spans.append(_finalise_span(current))
                current = None
            continue

        if label.startswith("B-"):
            if current:
                spans.append(_finalise_span(current))
            current = {
                "field": label[2:],
                "words": [word], "bboxes": [bbox],
                "confs": [confidences[idx]], "idxs": [idx],
            }

        elif label.startswith("I-"):
            field = label[2:]
            if current and current["field"] == field:
                current["words"].append(word)
                current["bboxes"].append(bbox)
                current["confs"].append(confidences[idx])
                current["idxs"].append(idx)
            else:
                # Stray I- without matching B- — start new span
                if current:
                    spans.append(_finalise_span(current))
                current = {
                    "field": field,
                    "words": [word], "bboxes": [bbox],
                    "confs": [confidences[idx]], "idxs": [idx],
                }

    if current:
        spans.append(_finalise_span(current))

    return spans


def _finalise_span(acc: dict) -> dict:
    bboxes = acc["bboxes"]
    return {
        "label_field": acc["field"],
        "text": " ".join(acc["words"]),
        "bbox": [
            min(b[0] for b in bboxes), min(b[1] for b in bboxes),
            max(b[2] for b in bboxes), max(b[3] for b in bboxes),
        ],
        "confidence": sum(acc["confs"]) / len(acc["confs"]),
        "word_idxs": acc["idxs"],
    }


# ---------------------------------------------------------------------------
# Label-prefix stripping
# ---------------------------------------------------------------------------

def strip_label_prefix(text: str, domain: str, label_field: str) -> str:
    """Remove leading anchor/label keywords from extracted span text.

    EasyOCR often merges a printed label and its value into one box:
        "Invoice #: FRT-2025-5862"  ->  "FRT-2025-5862"
        "Date: 22/08/2025"          ->  "22/08/2025"
        "Total: Rs. 57,269.19"      ->  "Rs. 57,269.19"

    Matches the longest anchor phrase, then strips it plus trailing separators.
    """
    anchors = ANCHORS.get(domain, {}).get(label_field, [])
    if not anchors:
        return text

    lower = text.lower()
    best_end = 0
    for anchor in sorted(anchors, key=len, reverse=True):
        if lower.startswith(anchor.lower()):
            best_end = len(anchor)
            break

    if best_end == 0:
        return text

    rest = text[best_end:].lstrip(" \t:#-.|/")
    return rest if rest else text


# ---------------------------------------------------------------------------
# Amount parsing
# ---------------------------------------------------------------------------

_CURRENCY_RE = re.compile(r"\b(rs\.?|inr|gbp|eur|usd)\b|[£€₹$]", re.IGNORECASE)


def parse_amount(text: str) -> Dict[str, Any]:
    """Parse a currency/amount string into {value, currency, raw}."""
    currency_match = _CURRENCY_RE.search(text)
    currency = currency_match.group(0) if currency_match else None
    cleaned = _CURRENCY_RE.sub("", text)
    cleaned = re.sub(r"[^\d.\-]", "", cleaned)
    # Handle multiple dots (e.g. "1.000.00" -> "1000.00")
    parts = cleaned.split(".")
    if len(parts) > 2:
        cleaned = "".join(parts[:-1]) + "." + parts[-1]
    try:
        value = float(cleaned)
    except ValueError:
        value = None
    return {"value": value, "currency": currency, "raw": text}


# ---------------------------------------------------------------------------
# Main extract()
# ---------------------------------------------------------------------------

def extract(
    image_path: str,
    domain: Optional[str] = None,
    models_dir: str = "models",
) -> Dict[str, Any]:
    """Full pipeline: image -> OCR -> LayoutLMv3 -> canonical JSON.

    Args:
        image_path: path to invoice image (PNG/JPG).
        domain: one of general/receipt/medical/insurance/logistics.
                If None, auto-detected from OCR text.
        models_dir: directory containing downloaded model subdirectories.

    Returns:
        Dict with extracted fields. Each scalar field is either:
          {"value": ..., "confidence": float, "raw": str}  (text fields)
          {"value": float, "currency": str, "raw": str}    (amount fields)
        Plus "_meta": {"domain": str, "num_spans": int}
    """
    image = Image.open(image_path).convert("RGB")
    img_w, img_h = image.size

    # 1. OCR
    ocr_words = run_ocr(image)
    if not ocr_words:
        return {"_meta": {"domain": domain or "unknown", "error": "No text found"}}

    # 2. Route domain
    if domain is None:
        domain = route_domain(ocr_words)

    # 3. Load model
    entry = load_model(domain, models_dir)
    model = entry["model"]
    processor = entry["processor"]
    id2label = entry["id2label"]
    device = entry["device"]

    # 4. Normalise bboxes to 0-1000 (LayoutLMv3 convention)
    words_text = [w["text"] for w in ocr_words]
    boxes_1000 = [
        [
            int(min(max(w["bbox"][0] * 1000 / img_w, 0), 1000)),
            int(min(max(w["bbox"][1] * 1000 / img_h, 0), 1000)),
            int(min(max(w["bbox"][2] * 1000 / img_w, 0), 1000)),
            int(min(max(w["bbox"][3] * 1000 / img_h, 0), 1000)),
        ]
        for w in ocr_words
    ]
    ocr_confs = [w["confidence"] for w in ocr_words]

    # 5. Encode with processor
    encoding = processor(
        images=image,
        text=words_text,
        boxes=boxes_1000,
        truncation=True,
        padding="max_length",
        max_length=MAX_SEQ_LENGTH,
        return_tensors="pt",
    )

    # 6. Run model
    with torch.no_grad():
        inputs = {k: v.to(device) for k, v in encoding.items()}
        outputs = model(**inputs)
    token_logits = outputs.logits[0].cpu()

    # 7. Map subword predictions -> word predictions (CRITICAL STEP)
    pred_label_ids, pred_confs = map_token_preds_to_words(
        encoding, token_logits, len(words_text)
    )

    # 8. Merge BIO spans
    spans = merge_spans(words_text, boxes_1000, pred_label_ids, id2label, pred_confs)

    # 9. Build output — pick highest-confidence span per field
    best_spans: Dict[str, Dict[str, Any]] = {}
    for span in spans:
        lf = span["label_field"]
        if lf not in best_spans or span["confidence"] > best_spans[lf]["confidence"]:
            best_spans[lf] = span

    # 10. Format output with label-prefix stripping
    result: Dict[str, Any] = {}
    for lf, span in best_spans.items():
        raw_text = span["text"]
        clean_text = strip_label_prefix(raw_text, domain, lf)

        # Detect amounts by checking if the field name suggests a numeric value
        amount_fields = {
            "subtotal", "total", "tax", "gst", "total_due", "total_billed",
            "eligible_amount", "approved_amount", "insurance_pays",
            "patient_liability", "not_covered", "deductible", "insurance_covered",
            "cash_paid", "change", "li_amount", "charge_amount", "item_billed",
            "proc_amount",
        }
        if lf in amount_fields:
            result[lf] = parse_amount(clean_text)
        else:
            result[lf] = {
                "value": clean_text,
                "confidence": round(span["confidence"], 4),
                "raw": raw_text,
            }

    result["_meta"] = {"domain": domain, "num_spans": len(spans)}
    return result


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Extract fields from an invoice image using LayoutLMv3.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python inference_example.py invoice.png
    python inference_example.py invoice.png --domain general
    python inference_example.py invoice.png --models-dir ./my_models

The first run downloads ~2.5 GB of model weights from HuggingFace.
        """,
    )
    parser.add_argument("image", help="Path to invoice image (PNG/JPG)")
    parser.add_argument("--domain", default=None,
                        choices=["general", "receipt", "medical", "insurance", "logistics"],
                        help="Force domain (default: auto-detect)")
    parser.add_argument("--models-dir", default="models",
                        help="Directory for model weights (default: ./models)")
    args = parser.parse_args()

    if not Path(args.image).exists():
        print(f"Error: {args.image} not found", file=sys.stderr)
        sys.exit(1)

    result = extract(args.image, domain=args.domain, models_dir=args.models_dir)
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()