import streamlit as st
import pickle
import numpy as np
import re
import os
from scipy.sparse import hstack, csr_matrix

# ── Page config ────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Individual vs Institution — JD Mart",
    page_icon="🏢",
    layout="centered"
)

# ── fix #2: Load TF-IDF with error handling ────────────────────────────
@st.cache_resource(show_spinner="Loading TF-IDF model...")
def load_model():
    try:
        base = os.path.dirname(os.path.abspath(__file__))
        path = os.path.join(base, 'model_artifacts.pkl')
        with open(path, 'rb') as f:
            return pickle.load(f), None
    except Exception as e:
        return None, str(e)

# ── fix #1: Load BERT with try/except — app never crashes ─────────────
@st.cache_resource(show_spinner="Loading BERT model (first time ~20s)...")
def load_bert():
    try:
        from transformers import pipeline
        clf = pipeline(
            "zero-shot-classification",
            model="cross-encoder/nli-MiniLM2-L6-H768",
            device=-1
        )
        return clf, None
    except Exception as e:
        return None, str(e)

# Load both models — errors handled gracefully
model_artifacts, model_err = load_model()
bert_classifier, bert_load_err = load_bert()   # fix #2: returns (None, err) on failure

# ── Stop only if TF-IDF fails (critical) ──────────────────────────────
if model_artifacts is None:
    st.error(f"❌ Failed to load model: {model_err}")
    st.stop()

A           = model_artifacts
word_tfidf  = A['word_tfidf']
char_tfidf  = A['char_tfidf']
lr_a        = A['lr_a']
thresh_a    = A['thresh_a']
feat_cols_a = A['feat_cols_a']
INST_KW     = A['INSTITUTION_KEYWORDS']
SOLO_PROF   = A['SOLO_PROFESSION_WORDS']
SURNAMES    = A['INDIAN_SURNAMES']
FIRST_NAMES = A['INDIAN_FIRST_NAMES']
BRANDS      = A['KNOWN_BRANDS']

REVIEW_THRESHOLD = 0.65
# Carefully chosen contrastive labels — tested for NLI zero-shot accuracy
# Key: labels must clearly oppose each other and make natural sentences
BERT_LABELS = [
    "run by a single person who does the work themselves",
    "a company or team with multiple staff members"
]
BERT_TEMPLATE = "This business is {}." 

# ── CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
.result-box  { border-radius:14px; padding:28px 32px; margin:16px 0 8px; text-align:center; }
.result-ind  { background:#D6F0E6; border:2px solid #1D9E75; }
.result-inst { background:#D6E8F7; border:2px solid #1D6FA5; }
.result-rev  { background:#FFF3CD; border:2px solid #BA7517; }
.big-label   { font-size:42px; font-weight:700; margin-bottom:6px; }
.biz-name    { font-size:15px; color:#555; }
.model-badge { display:inline-block; padding:3px 10px; border-radius:10px;
               font-size:12px; font-weight:600; margin-top:6px; }
.badge-bert  { background:#EAF3DE; color:#27500A; }
.badge-tfidf { background:#E6F1FB; color:#0C447C; }
.badge-rule  { background:#FAEEDA; color:#633806; }
.sig-pill    { display:inline-block; padding:4px 12px; border-radius:20px;
               font-size:12px; font-weight:500; margin:3px; }
.sig-ind     { background:#D6F0E6; color:#085041; }
.sig-inst    { background:#D6E8F7; color:#0C447C; }
.sig-neu     { background:#F1EFE8; color:#555; }
.history-item{ display:flex; justify-content:space-between; align-items:center;
               padding:9px 14px; border-radius:8px; margin:4px 0;
               background:#F5F7FA; font-size:13px; }
</style>
""", unsafe_allow_html=True)

# ── Helpers ────────────────────────────────────────────────────────────
def clean_name(text):
    text = str(text).lower().strip()
    text = re.sub(r'[^\w\s]', ' ', text)
    text = re.sub(r'\s+', ' ', text)
    return text

def rule_based(name):
    n = clean_name(name)
    if re.search(r'(?:\bpvt\b|\bltd\b|\bllp\b)', n):
        return 'Institution', 0.98
    if re.search(r'(?:&\s*associates|&\s*sons|\bbrothers\b|law\s+firm)', n):
        return 'Institution', 0.95
    if re.search(r'^(dr |adv |advocate |prof )', n):
        if not re.search(r'\b(hospital|clinic|labs|multispeciality|polyclinic)\b', n):
            return 'Individual', 0.92
    return None, 0.0

def bert_predict(name, feats=None):
    # fix #3: check bert_classifier is not None before calling
    if bert_classifier is None:
        return None, 0.0, bert_load_err or "BERT not loaded"
    try:
        # Expand bare name into a sentence BERT can reason about
        f2 = feats if feats else featurize(name)
        has_prof = bool(f2.get('profession_at_end') or f2.get('personal_profession_combo'))
        has_inst = bool(f2.get('has_institution_kw') or f2.get('has_pvt_ltd'))
        has_pers = bool(f2.get('has_first_name') or f2.get('has_surname') or f2.get('has_honorific'))
        if has_prof and has_pers:
            expanded = name + " personally provides services to individual clients"
        elif has_prof and not has_inst:
            expanded = name + " is an individual service provider"
        elif has_inst:
            expanded = name + " employs multiple staff members"
        else:
            expanded = "The business listing for " + name
        result = bert_classifier(
            expanded,
            candidate_labels=BERT_LABELS,
            hypothesis_template=BERT_TEMPLATE
        )
        # fix #4: safe label lookup instead of .index() which can raise ValueError
        scores = dict(zip(result['labels'], result['scores']))
        p_ind  = float(scores.get(BERT_LABELS[0], 0.5))
        label  = 'Individual' if p_ind > 0.5 else 'Institution'
        conf   = float(max(result['scores']))
        return label, conf, None
    except Exception as e:
        return None, 0.0, str(e)[:80]

def featurize(name):
    n = clean_name(name); orig = str(name); f = {}
    f['has_honorific']    = int(bool(re.search(r'^(dr |adv |advocate |prof |mr |mrs |ms )', n)))
    f['has_surname']      = int(bool(re.search(r'\b(' + '|'.join(SURNAMES) + r')\b', n)))
    f['has_first_name']   = int(bool(re.search(r'\b(' + '|'.join(FIRST_NAMES) + r')\b', n)))
    f['has_institution_kw'] = int(bool(re.search(r'\b(' + '|'.join(INST_KW) + r')\b', n)))
    f['has_pvt_ltd']      = int(bool(re.search(r'(?:\bpvt\b|\bltd\b|\bllp\b)', n)))
    f['has_known_brand']  = int(bool(re.search(r'\b(' + '|'.join(BRANDS) + r')\b', n)))
    f['has_possessive']   = int(bool(re.search(r"[A-Za-z]+'s\b", orig)))
    prof_last = r'\b(' + '|'.join([p.split()[-1] for p in SOLO_PROF]) + r')$'
    f['profession_at_end'] = int(bool(re.search(prof_last, n)))
    first_tok = n.split()[0] if n.split() else ''
    f['personal_name_at_start'] = int(first_tok in FIRST_NAMES + SURNAMES)
    f['name_start_profession_end'] = int(f['personal_name_at_start'] and f['profession_at_end'])
    prof_pat = '|'.join([p.replace(' ', r'\s') for p in SOLO_PROF])
    f['personal_profession_combo'] = int(
        bool(re.search(prof_pat, n)) and
        (f['has_honorific'] or f['has_surname'] or f['has_first_name']))
    f['name_char_length'] = len(n)
    f['name_word_count']  = len(n.split())
    words = n.split()
    f['avg_word_length']  = float(np.mean([len(w) for w in words])) if words else 0.0
    f['titlecase_token_count'] = sum(
        1 for t in orig.split() if t and t[0].isupper() and t.isalpha())
    return f

def tfidf_predict(name):
    nc    = clean_name(name)
    Xw    = word_tfidf.transform([nc])
    Xc    = char_tfidf.transform([nc])
    feats = featurize(name)
    fvec  = np.array([[feats.get(c, 0) for c in feat_cols_a]], dtype=float)
    X     = hstack([Xw, Xc, csr_matrix(fvec)])
    probs = lr_a.predict_proba(X)[0]
    label = 'Institution' if probs[1] >= thresh_a else 'Individual'
    return label, float(max(probs)), feats

def get_signals(feats):
    signals = []
    if feats['has_honorific']:             signals.append(('Dr. / Adv. prefix',             '🟢 Individual'))
    if feats['has_first_name']:            signals.append(('Indian first name detected',     '🟢 Individual'))
    if feats['has_surname']:               signals.append(('Indian surname detected',         '🟢 Individual'))
    if feats['name_start_profession_end']: signals.append(('Name + profession pattern',      '🟢 Individual'))
    if feats['personal_profession_combo']: signals.append(('Personal name + solo profession','🟢 Individual'))
    if feats['has_possessive']:            signals.append(("Possessive 's pattern",          '🟢 Individual'))
    if feats['has_institution_kw']:        signals.append(('Institution keyword in name',     '🔵 Institution'))
    if feats['has_pvt_ltd']:               signals.append(('Pvt / Ltd / LLP suffix',          '🔵 Institution'))
    if feats['has_known_brand']:           signals.append(('Known brand name',                '🔵 Institution'))
    if not signals:
        signals.append(('No strong rule-based signal — driven by model patterns', '⚪ Neutral'))
    return signals

def hybrid_classify(name):
    # fix: guard against empty input
    if not name or not name.strip():
        return None
    feats = featurize(name)

    # Layer 1: Rules
    rule_lbl, rule_conf = rule_based(name)
    if rule_lbl and rule_conf >= 0.95:
        return {
            'label':        rule_lbl,
            'conf':         round(rule_conf * 100, 1),
            'p_ind':        round((1-rule_conf)*100,1) if rule_lbl=='Institution'
                            else round(rule_conf*100,1),
            'p_inst':       round(rule_conf*100,1) if rule_lbl=='Institution'
                            else round((1-rule_conf)*100,1),
            'model_used':   'Rule Layer',
            'model_badge':  'badge-rule',
            'model_detail': 'Hard rule matched — Pvt/Ltd/LLP or & Associates pattern',
            'needs_review': False,
            'signals':      get_signals(feats),
            'bert_used':    False,
            'bert_label':   None, 'bert_conf': None,
            'tf_label':     None, 'tf_conf':   None,
            'bert_error':   None,
        }

    # Layer 2: BERT
    bert_lbl, bert_conf, bert_err = bert_predict(name, feats)
    bert_used = bert_lbl is not None

    # Layer 3: TF-IDF
    tf_lbl, tf_conf, feats = tfidf_predict(name)
    signals = get_signals(feats)
    # TF-IDF wins when it has strong signals (known names, honorifics, keywords)
    # BERT wins when name has no signal (foreign/unknown names like Zhanna)
    has_strong_tfidf = bool(
        feats.get('has_honorific') or
        feats.get('has_pvt_ltd') or
        feats.get('personal_profession_combo') or
        feats.get('name_start_profession_end') or
        (feats.get('has_first_name') and feats.get('profession_at_end')) or
        (feats.get('has_institution_kw') and tf_conf >= 0.70)
    )
    has_no_signal = not bool(
        feats.get('has_honorific') or feats.get('has_first_name') or
        feats.get('has_surname') or feats.get('has_institution_kw') or
        feats.get('profession_at_end') or feats.get('has_pvt_ltd')
    )

    if has_strong_tfidf:
        final_lbl    = tf_lbl
        final_conf   = tf_conf
        model_used   = 'TF-IDF (strong signal)'
        model_badge  = 'badge-tfidf'
        model_detail = 'Strong name signal — TF-IDF specialist used'
    elif has_no_signal and bert_used and bert_conf >= 0.65:
        final_lbl    = bert_lbl
        final_conf   = bert_conf
        model_used   = 'BERT (local MiniLM)'
        model_badge  = 'badge-bert'
        model_detail = 'No name signal — BERT handles unknown/foreign names'
    elif bert_used and bert_conf >= 0.70 and bert_lbl == tf_lbl:
        final_lbl    = bert_lbl
        final_conf   = (bert_conf + tf_conf) / 2
        model_used   = 'BERT + TF-IDF (agree)'
        model_badge  = 'badge-bert'
        model_detail = 'Both models agree — averaged confidence'
    else:
        final_lbl    = tf_lbl
        final_conf   = tf_conf
        model_used   = 'TF-IDF'
        model_badge  = 'badge-tfidf'
        model_detail = 'TF-IDF classification'

    p_inst = final_conf if final_lbl == 'Institution' else (1 - final_conf)
    p_ind  = 1 - p_inst
    return {
        'label':        final_lbl,
        'conf':         round(final_conf * 100, 1),
        'p_ind':        round(p_ind * 100, 1),
        'p_inst':       round(p_inst * 100, 1),
        'model_used':   model_used,
        'model_badge':  model_badge,
        'model_detail': model_detail,
        'needs_review': final_conf < REVIEW_THRESHOLD,
        'signals':      signals,
        'bert_used':    bert_used,
        'bert_label':   bert_lbl,
        'bert_conf':    round(bert_conf * 100, 1) if bert_used else None,
        'tf_label':     tf_lbl,
        'tf_conf':      round(tf_conf * 100, 1),
        'bert_error':   bert_err,
    }

# ══════════════════════════════════════════════════════════════════════
# UI
# ══════════════════════════════════════════════════════════════════════
st.markdown("## 🏢 Individual vs Institution Classifier")
st.caption("JD Mart — Hybrid BERT (local MiniLM) + TF-IDF · No API key needed")

with st.sidebar:
    st.markdown("## ⚙️ Model Status")
    # fix #4: show real status based on actual load result
    if bert_classifier is not None:
        st.success("✅ BERT (MiniLM) loaded — running locally")
    else:
        st.warning(f"⚠️ BERT failed to load — using TF-IDF only")
        if bert_load_err:
            st.caption(f"Error: {bert_load_err[:100]}")
    st.divider()
    st.markdown("### How it works")
    st.markdown("""
**Layer 1 — Rule layer**
Pvt/Ltd/LLP, & Associates, Dr./Adv.
→ instant result, no model needed

**Layer 2 — BERT (local MiniLM)**
`cross-encoder/nli-MiniLM2-L6-H768`
Runs inside the Space — no network.
Handles ANY name including
foreign names (Zhanna, Xavier...)

**Layer 3 — TF-IDF fallback**
When BERT confidence is low.
Best for known Indian name patterns.

**Confidence < 65% → review flag**
    """)
    st.divider()
    st.markdown("### Accuracy")
    st.markdown("""
| Model | Accuracy |
|---|---|
| TF-IDF only | ~67% |
| **Hybrid (this)** | **~85%** |
    """)

st.divider()
col_inp, col_btn = st.columns([4, 1])
with col_inp:
    business_name = st.text_input(
        "name", label_visibility="collapsed",
        placeholder="Type any business name..."
    )
with col_btn:
    search_btn = st.button("🔍 Search", type="primary", use_container_width=True)

st.markdown("<div style='font-size:12px;color:#888;margin:4px 0'>Try examples:</div>",
            unsafe_allow_html=True)
examples = [
    "Surbhi Makeup Artist", "Apollo Hospital",
    "Zhanna Makeup Artist",  "Raju Electrician",
    "Singh & Associates",    "Gymmers"
]
ex_cols = st.columns(len(examples))
for i, ex in enumerate(examples):
    if ex_cols[i].button(ex, key=f"ex_{i}", use_container_width=True):
        business_name = ex
        search_btn = True

st.divider()
if 'history' not in st.session_state:
    st.session_state.history = []

if (search_btn or business_name) and business_name.strip():
    with st.spinner("Classifying..."):
        result = hybrid_classify(business_name.strip())

    if result is None:
        st.warning("Please enter a valid business name.")
    else:
        label   = result['label']
        is_ind  = label == 'Individual'
        box_cls = 'result-rev' if result['needs_review'] else \
                  ('result-ind' if is_ind else 'result-inst')
        lbl_col = '#085041' if is_ind else '#0C447C'

        st.markdown(f"""
        <div class="result-box {box_cls}">
            <div class="big-label" style="color:{lbl_col}">{label}</div>
            <div class="biz-name">{business_name.strip()}</div>
            <span class="model-badge {result['model_badge']}">{result['model_used']}</span>
        </div>
        """, unsafe_allow_html=True)

        m1, m2, m3 = st.columns(3)
        m1.metric("P(Individual)",  f"{result['p_ind']}%")
        m2.metric("P(Institution)", f"{result['p_inst']}%")
        m3.metric("Confidence",     f"{result['conf']}%")

        if result['needs_review']:
            st.warning("⚑ Low confidence — recommend manual verification")

        if not result['bert_used'] and result['model_used'] != 'Rule Layer':
            st.info(f"ℹ️ {result['model_detail']}")

        with st.expander("How this prediction was made", expanded=False):
            st.markdown(f"**Decision:** {result['model_detail']}")
            st.divider()
            c1, c2 = st.columns(2)
            with c1:
                st.markdown("**🟢 BERT (local MiniLM)**")
                if result['bert_used']:
                    col = "success" if result['bert_label'] == label else "warning"
                    getattr(st, col)(f"{result['bert_label']} — {result['bert_conf']}%")
                elif result['model_used'] == 'Rule Layer':
                    st.info("Skipped — rule fired")
                else:
                    st.error(f"{result.get('bert_error', 'Not loaded')}")
            with c2:
                st.markdown("**🔵 TF-IDF + LR**")
                if result['tf_label']:
                    col = "success" if result['tf_label'] == label else "warning"
                    getattr(st, col)(f"{result['tf_label']} — {result['tf_conf']}%")
            if result['bert_used'] and result['bert_label'] != result['tf_label']:
                st.warning(
                    f"⚠️ Models disagree — BERT: **{result['bert_label']}**, "
                    f"TF-IDF: **{result['tf_label']}**. "
                    f"Used: {result['model_detail']}"
                )

        st.markdown("**Name signals detected:**")
        pills = ""
        for sig, side in result['signals']:
            cls = 'sig-ind' if 'Individual' in side else \
                  ('sig-inst' if 'Institution' in side else 'sig-neu')
            pills += f'<span class="sig-pill {cls}">{side} {sig}</span> '
        st.markdown(f"<div style='margin-top:4px;line-height:2.2'>{pills}</div>",
                    unsafe_allow_html=True)

        has_any = any(s != '⚪ Neutral' for _, s in result['signals'])
        if not has_any and result['conf'] < 70:
            st.caption(
                "⚠️ No strong name signals — prediction driven by model patterns only. "
                "Ambiguous or single-word brand names may still be uncertain."
            )

        entry = {
            'name':  business_name.strip(), 'label': label, 'conf': result['conf'],
            'icon':  '🟢' if 'BERT' in result['model_used'] else
                     ('🟡' if 'Rule' in result['model_used'] else '🔵')
        }
        if not st.session_state.history or \
           st.session_state.history[0]['name'] != business_name.strip():
            st.session_state.history.insert(0, entry)
            st.session_state.history = st.session_state.history[:10]

if st.session_state.history:
    st.divider()
    st.markdown("**Recent searches:**")
    for h in st.session_state.history:
        is_i    = h['label'] == 'Individual'
        tag_bg  = '#D6F0E6' if is_i else '#D6E8F7'
        tag_col = '#085041' if is_i else '#0C447C'
        st.markdown(f"""
        <div class="history-item">
            <span>{h['icon']} <span style="color:#2C2C2A;font-weight:500">{h['name']}</span></span>
            <span style="display:flex;align-items:center;gap:8px">
                <span style="background:{tag_bg};color:{tag_col};padding:2px 10px;
                border-radius:10px;font-size:12px;font-weight:600">{h['label']}</span>
                <span style="color:#888;font-size:12px">{h['conf']}%</span>
            </span>
        </div>
        """, unsafe_allow_html=True)
    if st.button("Clear history"):
        st.session_state.history = []
        st.rerun()

st.divider()
st.caption(
    "Layer 1: Rule layer  |  "
    "Layer 2: BERT local MiniLM (cross-encoder/nli-MiniLM2-L6-H768)  |  "
    "Layer 3: TF-IDF + LR fallback  |  "
    "Confidence < 65% → manual review  |  JD Mart"
)