"""
NetraLink Entity Resolution — deployment app
Loads:
  1. The Siamese-BERT sentence-transformers encoder (this repo)
  2. xgboost_matcher.joblib + feature_config.json (this repo)
Reproduces the exact feature pipeline from the training notebook
(pair_features / FEATURE_NAMES) so predictions match training.
"""

import json
import re

import joblib
import numpy as np
from huggingface_hub import hf_hub_download
from rapidfuzz import fuzz
from sentence_transformers import SentenceTransformer, util as st_util

import gradio as gr

REPO_ID = "ankoor123/Netralink-entity-resolution"

# ---------------------------------------------------------------------------
# Load artifacts once at startup
# ---------------------------------------------------------------------------

print("Loading Siamese-BERT encoder...")
siamese_model = SentenceTransformer(REPO_ID)

print("Loading XGBoost matcher + feature config...")
xgb_path = hf_hub_download(REPO_ID, "xgboost_matcher.joblib")
xgb_model = joblib.load(xgb_path)

config_path = hf_hub_download(REPO_ID, "feature_config.json")
with open(config_path, "r", encoding="utf-8") as f:
    feature_config = json.load(f)
FEATURE_NAMES = feature_config["feature_names"]

# ---------------------------------------------------------------------------
# Exact feature pipeline from the training notebook (cell 8)
# ---------------------------------------------------------------------------


def normalize_text(text):
    text = str(text).lower()
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def get_norm(m):
    if m.get("norm"):
        return str(m["norm"])
    return normalize_text(m["text"])


def mention_input_text(m):
    context = m.get("context", "") or ""
    if len(context) > 200:
        context = context[:200]
    return f"[{m.get('type', '')}] {m['text']} :: {context}"


def pair_features(m1, m2):
    emb1 = siamese_model.encode(mention_input_text(m1), convert_to_tensor=True)
    emb2 = siamese_model.encode(mention_input_text(m2), convert_to_tensor=True)
    cos_sim = float(st_util.cos_sim(emb1, emb2)[0][0])

    norm1, norm2 = get_norm(m1), get_norm(m2)

    ratio = fuzz.ratio(norm1, norm2) / 100.0
    token_sort_ratio = fuzz.token_sort_ratio(norm1, norm2) / 100.0

    soundex1 = m1.get("soundex", "")
    soundex2 = m2.get("soundex", "")
    soundex_match = 1.0 if soundex1 and soundex1 == soundex2 else 0.0

    len_diff = abs(len(norm1) - len(norm2))

    values = {
        "cosine_sim": cos_sim,
        "char_ratio": ratio,
        "token_sort_ratio": token_sort_ratio,
        "soundex_match": soundex_match,
        "len_diff": len_diff,
    }
    # Order strictly by FEATURE_NAMES from feature_config.json
    return [values[name] for name in FEATURE_NAMES], values


# ---------------------------------------------------------------------------
# Prediction function
# ---------------------------------------------------------------------------


def predict(text_a, type_a, context_a, text_b, type_b, context_b):
    if not text_a.strip() or not text_b.strip():
        return "Enter text for both mentions.", {}

    m1 = {"text": text_a.strip(), "type": type_a.strip(), "context": context_a.strip() or text_a.strip()}
    m2 = {"text": text_b.strip(), "type": type_b.strip(), "context": context_b.strip() or text_b.strip()}

    feat_vector, feat_dict = pair_features(m1, m2)
    prob = float(xgb_model.predict_proba(np.array([feat_vector], dtype=np.float32))[0, 1])
    label = "MATCH" if prob >= 0.5 else "NO MATCH"

    result = f"{label}  (match probability: {prob:.3f})"
    return result, feat_dict


# ---------------------------------------------------------------------------
# Gradio UI (also auto-exposes a REST API at /predict via the "api" tab)
# ---------------------------------------------------------------------------

with gr.Blocks(title="NetraLink Entity Resolution") as demo:
    gr.Markdown(
        "# NetraLink Entity Resolution\n"
        "Siamese-BERT (IndicBERTv2) similarity + XGBoost matcher.\n"
        "Enter two entity mentions to check whether they refer to the same entity."
    )
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Mention A**")
            text_a = gr.Textbox(label="Text")
            type_a = gr.Textbox(label="Entity type (e.g. P, ACC, PH)", value="")
            context_a = gr.Textbox(label="Context (optional, defaults to Text)", lines=2)
        with gr.Column():
            gr.Markdown("**Mention B**")
            text_b = gr.Textbox(label="Text")
            type_b = gr.Textbox(label="Entity type (e.g. P, ACC, PH)", value="")
            context_b = gr.Textbox(label="Context (optional, defaults to Text)", lines=2)

    btn = gr.Button("Resolve", variant="primary")
    output_label = gr.Textbox(label="Result")
    output_features = gr.JSON(label="Feature breakdown")

    btn.click(
        predict,
        inputs=[text_a, type_a, context_a, text_b, type_b, context_b],
        outputs=[output_label, output_features],
    )

if __name__ == "__main__":
    demo.launch()