# src/infer.py
from __future__ import annotations
import numpy as np
import pandas as pd

def _ensure_cols(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
    """Add any missing columns in one concat (avoids fragmentation)."""
    missing = [c for c in cols if c not in df.columns]
    if not missing:
        return df

    add = pd.DataFrame(np.nan, index=df.index, columns=missing)
    out = pd.concat([df, add], axis=1)
    return out.copy()  # defragment

def _base_model_preds(stacking_model, X: np.ndarray) -> np.ndarray:
    """
    Extract base estimator predictions from sklearn StackingRegressor (robustly).
    Handles both:
      - estimators_ : list of fitted estimators (most common)
      - estimators  : list of (name, estimator) pairs (pre-fit)
    """
    preds = []

    # Prefer fitted estimators_
    ests = getattr(stacking_model, "estimators_", None)
    if ests is None:
        # fallback to pre-fit definition (name, estimator)
        ests = [e for e in getattr(stacking_model, "estimators", [])]

    for item in ests or []:
        # item may be an estimator OR (name, estimator)
        est = item[1] if isinstance(item, (tuple, list)) and len(item) >= 2 else item

        if est is None or est == "drop":
            continue

        try:
            p = est.predict(X)
            preds.append(np.asarray(p).reshape(-1))
        except Exception:
            pass

    if not preds:
        return np.zeros((X.shape[0], 1))

    return np.column_stack(preds)  # (n_samples, n_base)




def _safe_minmax_norm(x: np.ndarray) -> np.ndarray:
    """Min-max normalize, but if constant vector -> zeros."""
    x = np.asarray(x, dtype=float)
    xmin = np.nanmin(x)
    xmax = np.nanmax(x)
    if not np.isfinite(xmin) or not np.isfinite(xmax) or (xmax - xmin) < 1e-12:
        return np.zeros_like(x, dtype=float)
    return (x - xmin) / (xmax - xmin)


def _norm_with_training_scale(x: np.ndarray, scale: float | None) -> np.ndarray:
    """
    Normalize with a training-derived scale (e.g., p95).
    If scale missing/invalid -> fallback to minmax (but stable for 1 row).
    """
    x = np.asarray(x, dtype=float)
    if scale is not None and np.isfinite(scale) and scale > 1e-12:
        return np.clip(x / scale, 0.0, 1.0)
    return _safe_minmax_norm(x)


def predict_with_confidence(bundle: dict, new_df: pd.DataFrame, interval: str = "q90") -> pd.DataFrame:
    """
    Returns a DF with prediction, interval, confidence_score, confidence_label, flags.
    interval: 'q90' or 'q95'
    """
    model = bundle["model"]
    imputer = bundle["imputer"]
    impute_cols = bundle["impute_cols"]
    feature_cols = bundle["feature_cols"]

    calib = bundle.get("calibration", {})          # contains q90/q95, etc.
    missing_rate = bundle.get("missing_rate", None)

    # Replace sentinel missing
    df = new_df.replace(-1, np.nan)

    # One-hot encode sex to match training (astype Int64 -> get_dummies).
    # Build all expected dummy columns explicitly so absent values get 0, not NaN.
    if "sex" in df.columns:
        sex_vals = df["sex"].astype("Int64")
        for v in [0, 1]:
            df[f"sex_{v}"] = (sex_vals == v).astype(int)
        df["sex_<NA>"] = sex_vals.isna().astype(int)
        df = df.drop(columns=["sex"])

    df = _ensure_cols(df, impute_cols)

    # Impute
    imputed = pd.DataFrame(
        imputer.transform(df[impute_cols]),
        columns=impute_cols,
        index=df.index,
    )

    X = imputed[feature_cols].values
    pred = model.predict(X)

    # Conformal interval half-width
    q = float(calib.get(interval, np.nan))
    lower = pred - q if np.isfinite(q) else np.full_like(pred, np.nan)
    upper = pred + q if np.isfinite(q) else np.full_like(pred, np.nan)

    # ---- Confidence components ----

    # 1) Disagreement across base estimators
    base_preds = _base_model_preds(model, X)
    disagreement = np.std(base_preds, axis=1) if base_preds.shape[1] > 1 else np.zeros(len(pred))

    dis_stats = bundle.get("disagreement_stats", {})
    dis_p95 = dis_stats.get("p95", None)
    d_norm = _norm_with_training_scale(disagreement, dis_p95)

    # 2) Missing input fraction — only morphological inputs.
    # Targets are always absent at inference (that's what we're predicting),
    # so including them would artificially inflate the missing penalty.
    morph_cols = bundle.get("morph_cols", impute_cols)
    missing_input_frac = df[morph_cols].isna().mean(axis=1).values
    m_norm = np.clip(missing_input_frac, 0.0, 1.0)

    # 3) Interval width penalty (calibrated + conformal)

    q50 = float(calib.get("q50", np.nan))
    q75 = float(calib.get("q75", np.nan))
    q90 = float(calib.get("q90", np.nan))
    q95 = float(calib.get("q95", np.nan))

    # interval width
    interval_width = upper - lower

    if np.isfinite(q50) and np.isfinite(q90) and q90 > q50:
        q_norm = np.clip((interval_width - q50) / (q90 - q50), 0.0, 1.0)
    else:
        target_std = bundle.get("target_std", None)
        if np.isfinite(q) and target_std is not None and target_std > 1e-12:
            q_norm = np.clip(interval_width / (2.0 * target_std), 0.0, 1.0)
        else:
            q_norm = np.zeros(len(pred))

    score = 1.0 - np.mean(
    np.vstack([m_norm, q_norm, d_norm]),
    axis=0
)

    # Apply training missingness penalty (your existing logic)
    sparse_flag = (
        missing_rate is not None
        and np.isfinite(missing_rate)
        and missing_rate >= 0.90
    )
    if sparse_flag:
        score = score * 0.8

    score = np.clip(score, 0.0, 1.0)

    # Labels
    # label = np.where(score >= 0.75, "High", np.where(score >= 0.50, "Medium", "Low"))
    label = np.where(score >= 0.60, "High",
        np.where(score >= 0.35, "Medium", "Low"))
    # if sparse_flag:
    #     label[:] = "Low"

    out = pd.DataFrame(
        {
            "prediction": pred,
            "lower": lower,
            "upper": upper,
            "confidence_score": score,
            "confidence_label": label,
        },
        index=df.index,
    )

    out["note"] = "⚠️ Target had ~90%+ missing data." if sparse_flag else ""
    return out