import os

import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.mixture import GaussianMixture
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KernelDensity
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder


def train_profile_brain(
    processed_dir: str = "data/processed",
    artifacts_dir: str = "outputs/artifacts",
    out_csv_name: str = "bd_profile_with_risk.csv",
    seed: int = 42,
    gmm_components: int = 3,
    kde_bandwidth: float = 2.0,
    w_gmm: float = 0.6,
    w_kde: float = 0.4,
):
    """
    Branch 2 — Profile Brain (Density Estimation)

    - Loads structured profile dataset (no labels required)
    - Preprocesses numeric + categorical features (impute + one-hot)
    - Fits GMM + KDE on training split (unsupervised)
    - Calibrates log-likelihood to risk score using calibration split percentiles
    - Produces final profile risk = w_gmm*risk_gmm + w_kde*risk_kde
    - Writes an output CSV with risk columns and saves artifacts
    """
    os.makedirs(artifacts_dir, exist_ok=True)

    bd_path = os.path.join(processed_dir, "bd_suicide_clean_structured.csv")
    bd = pd.read_csv(bd_path)

    # Keep a raw copy (unsupervised)
    X_raw = bd.copy()

    # Split numeric vs categorical
    num_cols = [c for c in X_raw.columns if pd.api.types.is_numeric_dtype(X_raw[c])]
    cat_cols = [c for c in X_raw.columns if c not in num_cols]

    print("Numeric:", num_cols)
    print("Categorical:", cat_cols)

    # ---------------- Preprocessor ----------------
    # IMPORTANT: make OneHotEncoder output dense to satisfy GMM/KDE requirements
    # sklearn >= 1.2 uses sparse_output; older versions use sparse
    try:
        oh = OneHotEncoder(handle_unknown="ignore", sparse_output=False)
    except TypeError:
        oh = OneHotEncoder(handle_unknown="ignore", sparse=False)

    pre = ColumnTransformer(
        transformers=[
            ("num", Pipeline([("imp", SimpleImputer(strategy="median"))]), num_cols),
            (
                "cat",
                Pipeline(
                    [
                        ("imp", SimpleImputer(strategy="most_frequent")),
                        ("oh", oh),
                    ]
                ),
                cat_cols,
            ),
        ],
        remainder="drop",
    )

    # Train / calibration split (unsupervised)
    X_train_raw, X_cal_raw = train_test_split(X_raw, test_size=0.3, random_state=seed)

    X_train = pre.fit_transform(X_train_raw)
    X_cal = pre.transform(X_cal_raw)

    # Ensure numeric dtype (helps stability / speed)
    X_train = np.asarray(X_train, dtype=np.float32)
    X_cal = np.asarray(X_cal, dtype=np.float32)

    print("X_train:", X_train.shape, "X_cal:", X_cal.shape)

    # ---------------- GMM ----------------
    gmm = GaussianMixture(
        n_components=gmm_components,
        covariance_type="diag",
        reg_covar=1e-3,
        random_state=seed,
    )
    gmm.fit(X_train)

    # Calibration: map log-likelihood to risk via percentile scaling
    cal_ll = gmm.score_samples(X_cal)
    low, high = np.percentile(cal_ll, [5, 95])

    def ll_to_risk(ll: np.ndarray) -> np.ndarray:
        # Higher LL => more "normal" => lower risk
        norm = np.clip((ll - low) / (high - low + 1e-9), 0.0, 1.0)
        return 1.0 - norm

    X_all = pre.transform(X_raw)
    X_all = np.asarray(X_all, dtype=np.float32)

    all_ll = gmm.score_samples(X_all)
    bd["risk_gmm"] = ll_to_risk(all_ll)

    # ---------------- KDE ----------------
    kde = KernelDensity(kernel="gaussian", bandwidth=kde_bandwidth)
    kde.fit(X_train)

    cal_ll_kde = kde.score_samples(X_cal)
    low_kde, high_kde = np.percentile(cal_ll_kde, [5, 95])

    def ll_to_risk_kde(ll: np.ndarray) -> np.ndarray:
        norm = np.clip((ll - low_kde) / (high_kde - low_kde + 1e-9), 0.0, 1.0)
        return 1.0 - norm

    all_ll_kde = kde.score_samples(X_all)
    bd["risk_kde"] = ll_to_risk_kde(all_ll_kde)

    # ---------------- Final Profile Risk ----------------
    bd["risk_profile"] = w_gmm * bd["risk_gmm"] + w_kde * bd["risk_kde"]

    # Risk levels by percentiles (dataset-relative)
    p70, p90 = np.percentile(bd["risk_profile"], [70, 90])

    def risk_to_level(r: float) -> str:
        if r >= p90:
            return "High"
        if r >= p70:
            return "Medium"
        return "Low"

    bd["risk_level_profile"] = bd["risk_profile"].apply(risk_to_level)

    # ---------------- Save CSV ----------------
    out_csv_path = os.path.join(processed_dir, out_csv_name)
    bd.to_csv(out_csv_path, index=False)
    print("✅ Saved:", out_csv_path)

    # ---------------- Save Artifacts ----------------
    joblib.dump(pre, os.path.join(artifacts_dir, "profile_preprocessor.joblib"))
    joblib.dump(gmm, os.path.join(artifacts_dir, "gmm.joblib"))
    joblib.dump(kde, os.path.join(artifacts_dir, "kde.joblib"))

    meta = {
        "gmm_components": int(gmm_components),
        "kde_bandwidth": float(kde_bandwidth),
        "w_gmm": float(w_gmm),
        "w_kde": float(w_kde),
        "p70": float(p70),
        "p90": float(p90),
        "calibration_ll_percentiles_gmm": [float(low), float(high)],
        "calibration_ll_percentiles_kde": [float(low_kde), float(high_kde)],
        "num_features_after_preprocess": int(X_train.shape[1]),
        "train_rows": int(X_train.shape[0]),
        "calibration_rows": int(X_cal.shape[0]),
    }
    joblib.dump(meta, os.path.join(artifacts_dir, "profile_risk_meta.joblib"))

    return {"out_csv": out_csv_path, "meta": meta}