import os
import json
import joblib
import numpy as np
import pandas as pd
import gradio as gr

# Plotting backend for Spaces (prevents runtime plotting issues)
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import shap
import dice_ml
from dice_ml import Dice

# ----------------------------
# Load artifacts
# ----------------------------
PIPELINE_PATH = "artifacts/pipeline.joblib"
SCHEMA_PATH = "artifacts/schema.json"

clf = joblib.load(PIPELINE_PATH)
schema = json.load(open(SCHEMA_PATH))

SKILLS_RAW = schema["skills"]  # like Skill_Python, Skill_SQL...
SKILLS_PRETTY = [s.replace("Skill_", "").replace("_", " ") for s in SKILLS_RAW]

# Your dataset choices
EDU_CHOICES = ["B.Sc", "B.Tech", "M.Tech", "MBA", "PhD"]
CERT_CHOICES = ["None", "AWS Certified", "Google ML", "Deep Learning Specialization"]
ROLE_CHOICES = ["AI Researcher", "Cybersecurity Analyst", "Data Scientist", "Software Engineer"]

# ----------------------------
# Helpers
# ----------------------------
def build_row(experience, salary, projects, education, certification, job_role, selected_skills_pretty):
    # Map pretty skills back to raw Skill_ columns
    selected_skills_raw = set("Skill_" + s.replace(" ", "_") for s in selected_skills_pretty)

    row = {k: 0 for k in SKILLS_RAW}
    for sk in SKILLS_RAW:
        row[sk] = 1 if sk in selected_skills_raw else 0

    row.update({
        "Experience (Years)": float(experience),
        "Salary Expectation ($)": float(salary),
        "Projects Count": int(projects),
        "Education": education,
        "Certifications": certification,
        "Job Role": job_role,
    })
    return pd.DataFrame([row])

def format_decision(proba: float):
    decision = "Hire ✅" if proba >= 0.75 else "Reject ❌"
    confidence = "High" if (proba >= 0.5 or proba <= 0.2) else "Medium"
    return decision, confidence

# ----------------------------
# Core functions
# ----------------------------
def predict_fn(experience, salary, projects, education, certification, job_role, selected_skills):
    X = build_row(experience, salary, projects, education, certification, job_role, selected_skills)
    proba = float(clf.predict_proba(X)[:, 1][0])
    decision, confidence = format_decision(proba)
    return decision, proba, confidence

def explain_shap_fn(experience, salary, projects, education, certification, job_role, selected_skills):
    X = build_row(experience, salary, projects, education, certification, job_role, selected_skills)

    pre = clf.named_steps["preprocess"]
    model = clf.named_steps["model"]
    Xt = pre.transform(X)

    explainer = shap.TreeExplainer(model)

    # New SHAP API is more stable across versions
    sv = explainer(Xt)

    plt.figure(figsize=(9, 5))
    shap.plots.waterfall(sv[0], show=False)
    plt.tight_layout()
    return plt.gcf()

def dice_recourse_fn(experience, salary, projects, education, certification, job_role, selected_skills):
    """
    Lightweight demo-style DiCE:
    - Education & Job Role are immutable (ethical recourse)
    - Actionable: skills, projects, certifications, salary
    """
    try:
        X = build_row(experience, salary, projects, education, certification, job_role, selected_skills)

        # Create a tiny pool for DiCE to operate on
        df_pool = pd.concat([X] * 50, ignore_index=True)
        df_pool["target"] = (clf.predict_proba(df_pool)[:, 1] >= 0.5).astype(int)

        immutable = ["Education", "Job Role"]
        continuous = [
            c for c in df_pool.columns
            if df_pool[c].dtype != "object" and c not in immutable and c != "target"
        ]

        if len(continuous) == 0:
            return pd.DataFrame({"error": ["No continuous features detected for counterfactual search."]})

        d = dice_ml.Data(dataframe=df_pool, continuous_features=continuous, outcome_name="target")
        m = dice_ml.Model(model=clf, backend="sklearn")
        dice = Dice(d, m, method="random")

        actionable = [c for c in X.columns if c not in immutable]

        cfs = dice.generate_counterfactuals(
            query_instances=X,
            total_CFs=3,
            desired_class=1,
            features_to_vary=actionable
        )
        cf_df = cfs.cf_examples_list[0].final_cfs_df

        # Make it readable: remove Skill_ prefix
        rename_map = {c: c.replace("Skill_", "") for c in cf_df.columns if c.startswith("Skill_")}
        cf_df = cf_df.rename(columns=rename_map)

        return cf_df

    except Exception as e:
        return pd.DataFrame({"error": [str(e)]})

# ----------------------------
# Creative presets (presentation boost)
# ----------------------------
PRESETS = {
    "Strong Candidate (Hire)": dict(
        experience=5, salary=80000, projects=3, education="M.Tech",
        certification="AWS Certified", job_role="Data Scientist",
        skills=["Python", "SQL", "MachineLearning", "TensorFlow", "Pytorch"]
    ),
    "Borderline Candidate": dict(
        experience=2, salary=70000, projects=2, education="B.Sc",
        certification="None", job_role="Data Scientist",
        skills=["Python", "SQL"]
    ),
    "Weak Candidate (Reject)": dict(
        experience=0, salary=95000, projects=0, education="B.Sc",
        certification="None", job_role="Data Scientist",
        skills=[]
    ),
    "Cybersecurity Profile": dict(
        experience=4, salary=90000, projects=3, education="B.Tech",
        certification="None", job_role="Cybersecurity Analyst",
        skills=["Cybersecurity", "EthicalHacking", "Linux", "Networking"]
    ),
}

def load_preset(preset_name):
    p = PRESETS[preset_name]
    return (
        p["experience"], p["salary"], p["projects"],
        p["education"], p["certification"], p["job_role"],
        p["skills"]
    )

# ----------------------------
# UI Styling (simple, clean)
# ----------------------------
CUSTOM_CSS = """
#title { font-size: 34px; font-weight: 800; margin-bottom: 0.2rem; }
#subtitle { font-size: 15px; opacity: 0.85; margin-bottom: 1rem; }
.badge { display:inline-block; padding: 6px 10px; border-radius: 999px; font-size: 12px; margin-right: 6px; }
.badge1 { background: rgba(0,255,150,0.12); border: 1px solid rgba(0,255,150,0.25); }
.badge2 { background: rgba(90,180,255,0.12); border: 1px solid rgba(90,180,255,0.25); }
.badge3 { background: rgba(255,200,90,0.12); border: 1px solid rgba(255,200,90,0.25); }
"""

# ----------------------------
# Build Gradio App
# ----------------------------
with gr.Blocks(css=CUSTOM_CSS, title="XAI Recruitment Screening (XGBoost + SHAP + DiCE)") as demo:
    gr.Markdown('<div id="title">XAI Recruitment Screening System</div>')
    gr.Markdown(
        '<div id="subtitle">'
        '<span class="badge badge1">Predict</span>'
        '<span class="badge badge2">Explain (SHAP)</span>'
        '<span class="badge badge3">Recourse (DiCE)</span>'
        '<br/>A transparent screening demo using XGBoost + SHAP + DiCE counterfactual recourse.'
        '</div>'
    )

    with gr.Accordion("📌 How to use this demo (for assessment)", open=True):
        gr.Markdown(
            """
**What you’ll see**
- **Predict:** Hire/Reject + probability  
- **Explain (SHAP):** Which features pushed the decision up/down  
- **Recourse (DiCE):** Actionable counterfactual suggestions  

**Ethical design**
- `Education` and `Job Role` are treated as **immutable** in counterfactuals (we don’t recommend changing identity/context).
- Dataset does not include protected attributes (e.g., gender/race), so fairness is audited by available groups (Education / Job Role).
            """
        )

    with gr.Row():
        with gr.Column(scale=1):
            preset = gr.Dropdown(
                choices=list(PRESETS.keys()),
                value="Strong Candidate (Hire)",
                label="🎛️ Quick Test Presets"
            )
            load_btn = gr.Button("Load Preset")

            exp = gr.Number(label="Experience (Years)", value=5)
            sal = gr.Number(label="Salary Expectation ($)", value=80000)
            proj = gr.Number(label="Projects Count", value=3)

            edu = gr.Dropdown(EDU_CHOICES, label="Education", value="B.Sc")
            cert = gr.Dropdown(CERT_CHOICES, label="Certifications", value="None")
            role = gr.Dropdown(ROLE_CHOICES, label="Job Role", value="Data Scientist")

            skills = gr.CheckboxGroup(
                choices=SKILLS_PRETTY,
                label="Skills (select all that apply)"
            )

        with gr.Column(scale=1):
            with gr.Tab("Predict"):
                run_btn = gr.Button("✅ Run Prediction", variant="primary")
                out_decision = gr.Text(label="Decision")
                out_prob = gr.Number(label="Hire Probability")
                out_conf = gr.Text(label="Confidence")

            with gr.Tab("Explain (SHAP)"):
                gr.Markdown("**Local explanation:** red features decrease hire probability, blue features increase it.")
                shap_plot = gr.Plot()
                shap_btn = gr.Button("🔎 Generate SHAP Explanation")

            with gr.Tab("Recourse (DiCE)"):
                gr.Markdown("**Recourse suggestions:** actionable changes that may flip the decision to *Hire*.")
                cf_table = gr.Dataframe()
                dice_btn = gr.Button("🧭 Generate Counterfactuals (DiCE)")

    # Preset loader
    load_btn.click(
        load_preset,
        inputs=[preset],
        outputs=[exp, sal, proj, edu, cert, role, skills]
    )

    # Wire buttons
    run_btn.click(
        predict_fn,
        inputs=[exp, sal, proj, edu, cert, role, skills],
        outputs=[out_decision, out_prob, out_conf]
    )

    shap_btn.click(
        explain_shap_fn,
        inputs=[exp, sal, proj, edu, cert, role, skills],
        outputs=[shap_plot]
    )

    dice_btn.click(
        dice_recourse_fn,
        inputs=[exp, sal, proj, edu, cert, role, skills],
        outputs=[cf_table]
    )

demo.launch()