import joblib, numpy as np
from pathlib import Path

def load_globalwellfm(model_dir="GlobalWellFM_v1"):
    p = Path(model_dir)
    model   = joblib.load(p / "xgb_model.joblib")
    imputer = joblib.load(p / "imputer.joblib")
    le      = joblib.load(p / "label_encoder.joblib")
    return model, imputer, le

def predict_lithofacies(df, model_dir="GlobalWellFM_v1"):
    """df: DataFrame with columns GR, RHOB, NPHI, RT (ohm.m)"""
    model, imputer, le = load_globalwellfm(model_dir)
    f = df.copy()
    f["GR_N"]   = f["GR"].clip(0,200)/200
    f["RHOB_N"] = (f["RHOB"].clip(1.5,3.0)-1.5)/1.5
    f["NPHI_N"] = f["NPHI"].clip(-0.1,0.6)
    f["RT_LOG"] = np.log10(f["RT"].clip(0.01,10000))
    f["VCL"]    = ((f["GR"].clip(0,200)-30)/(120-30)).clip(0,1)
    f["PHI_D"]  = ((2.65-f["RHOB"])/(2.65-1.0)).clip(0,0.45)
    f["PHI_N"]  = f["NPHI"].clip(0,0.45)
    f["PHI_AVG"]= (f["PHI_D"]+f["PHI_N"])/2
    f["PHI_EFF"]= (f["PHI_AVG"]*(1-f["VCL"])).clip(0,0.45)
    f["NPHI_RHOB_DIFF"] = f["NPHI_N"]-f["RHOB_N"]
    f["GR_RT_RATIO"]    = f["GR_N"]/(f["RT_LOG"]+2)
    import warnings; warnings.filterwarnings("ignore")
    Ro = 0.05/np.power(np.where(f["PHI_EFF"]>0,f["PHI_EFF"],np.nan),2)
    f["SW"] = np.clip(np.sqrt(Ro/f["RT"].clip(0.01)),0,1)
    FCOLS = ["GR_N","RHOB_N","NPHI_N","RT_LOG","VCL","PHI_D","PHI_N","PHI_AVG","PHI_EFF","NPHI_RHOB_DIFF","GR_RT_RATIO","SW"]
    X = imputer.transform(f[FCOLS].values)
    probs = model.predict_proba(X)
    preds = le.inverse_transform(model.predict(X))
    conf  = probs.max(axis=1)
    return preds, conf