import os
import warnings
import numpy as np
import pandas as pd
from flask import Flask, request, jsonify
from flask_cors import CORS
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score

try:
    from xgboost import XGBClassifier
    HAS_XGB = True
except Exception:  # pragma: no cover
    HAS_XGB = False

warnings.filterwarnings("ignore")

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

NUMERIC = [
    "age", "years_at_company", "years_in_current_role", "salary_thb", "salary_hike_pct",
    "distance_from_home_km", "overtime_hours_monthly", "monthly_absent_days",
    "satisfaction_score", "performance_rating", "num_projects", "training_hours_yearly",
    "promotion_last_3years", "work_life_balance", "job_involvement", "manager_rating",
]
CATEGORICAL = ["department", "gender", "education"]

# direction = +1 means "higher value increases resign risk"
DIRECTION = {
    "age": -1, "years_at_company": -1, "years_in_current_role": +1, "salary_thb": -1,
    "salary_hike_pct": -1, "distance_from_home_km": +1, "overtime_hours_monthly": +1,
    "monthly_absent_days": +1, "satisfaction_score": -1, "performance_rating": -1,
    "num_projects": +1, "training_hours_yearly": -1, "promotion_last_3years": -1,
    "work_life_balance": -1, "job_involvement": -1, "manager_rating": -1,
}


def build_models(name):
    name = (name or "xgboost").lower()
    defs = {
        "random_forest": lambda: RandomForestClassifier(
            n_estimators=300, min_samples_leaf=2, random_state=42, class_weight="balanced"),
        "logistic": lambda: make_pipeline(
            StandardScaler(), LogisticRegression(max_iter=2000, class_weight="balanced")),
    }
    if HAS_XGB:
        defs["xgboost"] = lambda: XGBClassifier(
            n_estimators=350, max_depth=4, learning_rate=0.08, subsample=0.9,
            colsample_bytree=0.9, reg_lambda=1.0, eval_metric="logloss", random_state=42)
    else:
        defs["xgboost"] = lambda: GradientBoostingClassifier(random_state=42)

    if name == "all":
        return {k: v() for k, v in defs.items()}
    if name in defs:
        return {name: defs[name]()}
    return {"xgboost": defs["xgboost"]()}


def prepare(df):
    """Return (X, feature_names, df) with numeric + one-hot encoded features."""
    for c in NUMERIC:
        if c not in df.columns:
            df[c] = 0
        df[c] = pd.to_numeric(df[c], errors="coerce")
        df[c] = df[c].fillna(df[c].median() if df[c].notna().any() else 0.0)
    for c in CATEGORICAL:
        if c not in df.columns:
            df[c] = "unknown"
        df[c] = df[c].astype(str).fillna("unknown")

    X = pd.get_dummies(df[NUMERIC + CATEGORICAL], columns=CATEGORICAL, dummy_na=False)
    X = X.astype(float)
    return X, list(X.columns), df


def rule_score(df):
    """Weak label / fallback risk score in 0..1 when `resigned` is not provided."""
    z = pd.DataFrame(index=df.index)
    for c in NUMERIC:
        s = df[c].astype(float)
        sd = s.std(ddof=0)
        z[c] = 0.0 if sd == 0 or np.isnan(sd) else (s - s.mean()) / sd
        z[c] = z[c] * DIRECTION.get(c, 1)
    weights = {
        "satisfaction_score": 2.0, "overtime_hours_monthly": 1.6, "salary_hike_pct": 1.4,
        "work_life_balance": 1.3, "manager_rating": 1.2, "monthly_absent_days": 1.1,
        "promotion_last_3years": 1.1, "distance_from_home_km": 0.9, "salary_thb": 0.9,
        "job_involvement": 0.9, "years_at_company": 0.7, "training_hours_yearly": 0.6,
        "performance_rating": 0.5, "num_projects": 0.5, "age": 0.4,
        "years_in_current_role": 0.4,
    }
    raw = sum(z[c] * w for c, w in weights.items())
    raw = raw / max(1e-9, sum(weights.values())) * 3.0
    return 1 / (1 + np.exp(-raw))


def metrics(y_true, y_prob, thr):
    yp = (np.asarray(y_prob) >= thr).astype(int)
    yt = np.asarray(y_true).astype(int)
    p, r, f1, _ = precision_recall_fscore_support(yt, yp, average="binary", zero_division=0)
    out = {
        "accuracy": round(float(accuracy_score(yt, yp)), 4),
        "precision": round(float(p), 4),
        "recall": round(float(r), 4),
        "f1": round(float(f1), 4),
    }
    try:
        if len(np.unique(yt)) > 1:
            out["auc"] = round(float(roc_auc_score(yt, y_prob)), 4)
    except Exception:
        pass
    return out


def importances_of(model, feature_names, X, y):
    est = model[-1] if hasattr(model, "steps") else model
    if hasattr(est, "feature_importances_"):
        vals = np.asarray(est.feature_importances_, dtype=float)
    elif hasattr(est, "coef_"):
        vals = np.abs(np.asarray(est.coef_, dtype=float)).ravel()
    else:
        vals = np.zeros(len(feature_names))
    total = vals.sum() or 1.0
    pairs = sorted(zip(feature_names, vals / total), key=lambda kv: kv[1], reverse=True)
    return [{"feature": f, "importance": round(float(v), 4)} for f, v in pairs[:20]]


def key_factors(row, X_row, importance_map):
    out = []
    for feat, imp in sorted(importance_map.items(), key=lambda kv: kv[1], reverse=True):
        base = feat.split("_")[0]
        if feat in NUMERIC:
            direction = DIRECTION.get(feat, 1)
            val = float(row[feat])
            out.append({"feature": feat, "value": val,
                        "impact": round(float(imp) * direction, 4)})
        elif feat in X_row.index and X_row[feat] > 0:
            out.append({"feature": feat, "value": 1, "impact": round(float(imp), 4)})
        if len(out) >= 6:
            break
    return out


def threshold_analysis(y_true, y_prob):
    if y_true is None:
        return []
    yt = np.asarray(y_true).astype(int)
    yp = np.asarray(y_prob, dtype=float)
    rows = []
    for thr in np.round(np.arange(0.1, 0.95, 0.05), 2):
        flag = (yp >= thr).astype(int)
        tp = int(((flag == 1) & (yt == 1)).sum())
        fp = int(((flag == 1) & (yt == 0)).sum())
        fn = int(((flag == 0) & (yt == 1)).sum())
        tn = int(((flag == 0) & (yt == 0)).sum())
        prec = tp / (tp + fp) if tp + fp else 0.0
        rec = tp / (tp + fn) if tp + fn else 0.0
        f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0
        rows.append({
            "threshold": float(thr),
            "precision_pct": round(prec * 100, 1),
            "recall_pct": round(rec * 100, 1),
            "f1_pct": round(f1 * 100, 1),
            "flagged": int(flag.sum()), "tp": tp, "fp": fp, "fn": fn, "tn": tn,
        })
    return rows


def risk_level(p, thr):
    if p >= thr:
        return "High"
    if p >= max(0.15, thr * 0.6):
        return "Medium"
    return "Low"


@app.route("/api/hr/classify", methods=["POST", "OPTIONS"])
def classify():
    if request.method == "OPTIONS":
        return ("", 204)
    try:
        body = request.get_json(force=True) or {}
        rows = body.get("data", [])
        model_name = body.get("model", "xgboost")
        thr = float(body.get("threshold", 0.5) or 0.5)
        if not rows:
            return jsonify({"error": "no data"}), 400

        df = pd.DataFrame(rows)
        df.columns = [str(c).strip().lower() for c in df.columns]
        X, feature_names, df = prepare(df)

        has_label = "resigned" in df.columns and pd.to_numeric(
            df["resigned"], errors="coerce").notna().any()
        y = None
        if has_label:
            y = pd.to_numeric(df["resigned"], errors="coerce").fillna(0)
            y = (y >= 1).astype(int)
            has_label = y.nunique() >= 2 and len(df) >= 10

        model_results = {}
        prob_by_model = {}

        if has_label:
            strat = y if y.value_counts().min() >= 2 else None
            Xtr, Xte, ytr, yte = train_test_split(
                X, y, test_size=0.25, random_state=42, stratify=strat)
            for name, mdl in build_models(model_name).items():
                try:
                    mdl.fit(Xtr, ytr)
                    pte = mdl.predict_proba(Xte)[:, 1]
                    m = metrics(yte, pte, thr)
                    mdl.fit(X, y)
                    prob_by_model[name] = mdl.predict_proba(X)[:, 1]
                    m["feature_importance"] = importances_of(mdl, feature_names, X, y)
                    model_results[name] = m
                except Exception as e:
                    model_results[name] = {"error": str(e)[:200]}

        if not prob_by_model:
            # unsupervised fallback: train on rule-based weak label
            weak = rule_score(df)
            y_weak = (weak >= np.quantile(weak, 0.75)).astype(int)
            for name, mdl in build_models(model_name).items():
                try:
                    mdl.fit(X, y_weak)
                    p = mdl.predict_proba(X)[:, 1]
                    prob_by_model[name] = 0.5 * p + 0.5 * weak.values
                    model_results[name] = {
                        **metrics(y_weak, p, thr),
                        "feature_importance": importances_of(mdl, feature_names, X, y_weak),
                        "note": "weak-label (no `resigned` column)",
                    }
                except Exception:
                    prob_by_model[name] = weak.values
                    model_results[name] = {"note": "rule-based fallback"}
            has_label = False

        def score_of(k):
            m = model_results.get(k, {})
            return (m.get("auc") or 0) * 2 + (m.get("f1") or 0)

        best_model = max(prob_by_model.keys(), key=score_of)
        probs = np.asarray(prob_by_model[best_model], dtype=float)

        imp_map = {d["feature"]: d["importance"]
                   for d in model_results.get(best_model, {}).get("feature_importance", [])}

        predictions = []
        for i, (_, row) in enumerate(df.iterrows()):
            p = float(np.clip(probs[i], 0.0, 1.0))
            item = {
                "employee_id": str(row.get("employee_id", f"E{i+1}")),
                "employee_name": str(row.get("employee_name", "")),
                "department": str(row.get("department", "")),
                "position": str(row.get("position", "")),
                "resign_probability": round(p, 4),
                "risk_level": risk_level(p, thr),
                "key_factors": key_factors(row, X.iloc[i], imp_map),
            }
            for c in NUMERIC + ["gender", "education"]:
                v = row.get(c)
                item[c] = float(v) if c in NUMERIC else str(v)
            if has_label:
                item["resigned"] = int(pd.to_numeric(row.get("resigned", 0), errors="coerce") or 0)
            predictions.append(item)

        high = [p for p in predictions if p["risk_level"] == "High"]
        med = [p for p in predictions if p["risk_level"] == "Medium"]
        low = [p for p in predictions if p["risk_level"] == "Low"]
        n = max(1, len(predictions))

        dept = []
        for name, g in df.assign(_p=probs).groupby(df.get("department", "unknown")):
            dept.append({
                "department": str(name),
                "avg_resign_prob": round(float(g["_p"].mean()) * 100, 1),
                "high_risk": int((g["_p"] >= thr).sum()),
                "total": int(len(g)),
            })
        dept.sort(key=lambda d: d["avg_resign_prob"], reverse=True)

        top = sorted(predictions, key=lambda p: p["resign_probability"], reverse=True)[:10]

        return jsonify({
            "total_employees": len(predictions),
            "risk_summary": {
                "high_risk": len(high), "high_risk_pct": round(len(high) / n * 100, 1),
                "medium_risk": len(med), "medium_risk_pct": round(len(med) / n * 100, 1),
                "low_risk": len(low), "low_risk_pct": round(len(low) / n * 100, 1),
            },
            "department_summary": dept,
            "top_at_risk": top,
            "predictions": predictions,
            "model_results": model_results,
            "best_model": best_model,
            "has_label": bool(has_label),
            "threshold_used": thr,
            "threshold_analysis": threshold_analysis(
                y.values if has_label else None, probs),
            "meta": {"requested_model": model_name, "xgboost_available": HAS_XGB,
                     "n_features": len(feature_names)},
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/")
def root():
    return jsonify({"status": "ok", "service": "entgroup-hr-analytics",
                    "endpoints": ["/health", "/api/hr/classify"]})


@app.route("/health")
def health():
    return jsonify({"status": "ok"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))