"""
Full pipeline for biomass pyrolysis ML + optimization

- Loads dataset(pyrolysis.xlsx)
- Cleans data and computes:
    * oil_org = bio_liquid - water
    * Thermal Severity Index (TSI)
- Cross-validation for:
    * bio_liquid
    * oil_org
- Trains final LightGBM models
- Saves feature importances
- SHAP:
    * Summary plots (bio_liquid & oil_org)
    * 2×2 dependence panel for oil_org
    * Rank stability analysis (Oil vs Bio-liquid)
- Partial Dependence Plots (PDP) for key process variables
- Accumulated Local Effects (ALE) for key process variables
- One-at-a-time Sensitivity analysis (±10%) for oil_org
- Pareto exploration:
    * Maximise oil_org, minimise water
- Bayesian optimisation (optional, needs scikit-optimize):
    * Maximise oil_org
    * Minimise water
    * Scalarised objective: oil_org - λ * water

Outputs (in out_mechml/):
    cleaned_data.csv
    cv_results.csv
    feature_importances_bio_liquid.csv
    feature_importances_oil_org.csv
    shap_summary_bio_liquid.png
    shap_summary_oil_org.png
    shap_2x2_oil_org.png
    shap_rank_stability_oil_vs_liquid.png
    pdp_oil_org.png
    pdp_bio_liquid.png
    ale_oil_org.png              (if alibi installed)
    ale_bio_liquid.png           (if alibi installed)
    sensitivity_oil_org.png
    pareto_candidates.csv
    pareto_front.csv
    pareto_oil_vs_water.png
    bo_single_results.csv        (if skopt installed)
    bo_scalar_results.csv        (if skopt installed)
"""

import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")  # non-interactive backend
import matplotlib.pyplot as plt

# Publication-quality settings for Renewable Energy journal (Elsevier)
plt.rcParams['font.size'] = 11
plt.rcParams['axes.labelsize'] = 12
plt.rcParams['axes.titlesize'] = 13
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['legend.fontsize'] = 10
plt.rcParams['figure.titlesize'] = 14
plt.rcParams['lines.linewidth'] = 2.0
plt.rcParams['axes.linewidth'] = 1.2
plt.rcParams['grid.linewidth'] = 0.5
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman', 'DejaVu Serif']
plt.rcParams['mathtext.fontset'] = 'dejavuserif'
plt.rcParams['axes.grid'] = False
plt.rcParams['savefig.dpi'] = 600
plt.rcParams['savefig.bbox'] = 'tight'
plt.rcParams['savefig.pad_inches'] = 0.05


from sklearn.model_selection import KFold
from sklearn.metrics import r2_score, mean_absolute_error
from sklearn.inspection import partial_dependence

from lightgbm import LGBMRegressor
import shap

# Helper function to get feature labels with units
def get_label_with_unit(feat):
    """Return feature name with appropriate unit."""
    units = {
        "temp": "temp (°C)",
        "HR": "HR (°C/min)",
        "PS": "PS (mm)",
        "Flow-N2": "Flow-N2 (ml/min)",
        "cellulose": "cellulose (wt%)",
        "hemicellulose": "hemicellulose (wt%)",
        "lignin": "lignin (wt%)",
        "Ash": "Ash (wt%)",
        "vol": "vol (wt%)",
        "fc": "fc (wt%)",
        "C": "C (wt%)",
        "H": "H (wt%)",
        "O": "O (wt%)",
        "N": "N (wt%)",
        "hceff": "hceff (wt%)",
        "TSI": "TSI",
    }
    return units.get(feat, feat)

# Optional: scikit-optimize for Bayesian optimization
try:
    from skopt import gp_minimize
    from skopt.space import Real
except ImportError:
    gp_minimize = None
    Real = None
    print("[WARN] scikit-optimize not installed. "
          "Bayesian optimisation will be skipped. "
          "Install via: pip install scikit-optimize")



# -------------------------------------------------
# 0. CONFIG
# -------------------------------------------------
EXCEL_PATH = r"C:\Users\admin\Desktop\Bio-ML\V14122025\pyrolysis.xlsx"  # Change your dataset path

# Set output directory to same location as Excel file
import os.path as osp
OUTDIR = osp.join(osp.dirname(EXCEL_PATH), "output6")
os.makedirs(OUTDIR, exist_ok=True)

RANDOM_STATE = 42
N_SPLITS = 5
N_SHAP_SAMPLES = 300
N_PARETO_SAMPLES = 20000
N_BOOT_RANK = 30      # for SHAP rank stability

# -------------------------------------------------
# 1. LOAD & CLEAN DATA
# -------------------------------------------------
df = pd.read_excel(EXCEL_PATH)

# Ensure outputs are present
df = df.dropna(subset=["bio-liquid yield(wt%)", "Water content"])

# Convert all numeric columns to float to avoid integer data warnings
for col in df.select_dtypes(include=['int64', 'int32']).columns:
    df[col] = df[col].astype(np.float64)

# Standardise column names - UPDATED WITH CAPITAL LETTERS
col_map = {
    "Ash(wt%)": "Ash",
    "FixedCarbon(wt%)": "fc",
    "Volatiles(wt%)": "vol",
    "C(wt%)": "C",
    "H/Ceff.(wt%)": "hceff",
    "H(wt%)": "H",
    "O(wt%)": "O",
    "N(wt%)": "N",
    "Cellulose(wt%)": "cellulose",
    "Hemicellulose(wt%)": "hemicellulose",
    "Lignin(wt%)": "lignin",
    "T(°C)": "temp",
    "HeatingRate(°C/min)": "HR",
    "ParticleSize(mm)": "PS",
    "FlowRate-Nitrogen(ml/min)": "Flow-N2",
    "bio-liquid yield(wt%)": "bio_liquid",
    "Water content": "water",
}
df = df.rename(columns=col_map)

required_cols = [
    "Ash", "fc", "vol", "C", "H", "O", "N", "hceff",
    "cellulose", "hemicellulose", "lignin",
    "temp", "HR", "PS", "Flow-N2",
    "bio_liquid", "water",
]
df = df.dropna(subset=required_cols).reset_index(drop=True)

# compute organic oil
df["oil_org"] = df["bio_liquid"] - df["water"]

# Thermal Severity Index (simple proposal)
eps = 1e-6
df["TSI"] = df["temp"] * np.log(df["HR"] + 1.0) / (df["PS"] + eps)

# save cleaned data
df.to_csv(os.path.join(OUTDIR, "cleaned_data.csv"), index=False)

print(f"Data loaded: {df.shape[0]} samples, {df.shape[1]} columns")

# -------------------------------------------------
# 2. FEATURES & TARGETS - UPDATED WITH CAPITAL LETTERS
# -------------------------------------------------
BASE_FEATURES = [
    "Ash", "fc", "vol", "C", "H", "O", "N", "hceff",
    "cellulose", "hemicellulose", "lignin",
    "temp", "HR", "PS", "Flow-N2",
]
FEATURES_WITH_TSI = BASE_FEATURES + ["TSI"]

TARGETS = {
    "bio_liquid": "Total bio-liquid yield (wt%)",
    "oil_org": "Organic oil yield (wt% = bio-liquid - water)",
}

# -------------------------------------------------
# 3. CROSS-VALIDATION HELPER
# -------------------------------------------------
def run_cv(X, y, model, n_splits=N_SPLITS, name="", feature_names=None):
    kf = KFold(n_splits=n_splits, shuffle=True, random_state=RANDOM_STATE)
    r2s, maes = [], []

    fold = 0
    for train_idx, test_idx in kf.split(X):
        fold += 1
        X_train, X_test = X[train_idx], X[test_idx]
        y_train, y_test = y[train_idx], y[test_idx]

        if feature_names is not None:
            model.fit(X_train, y_train, feature_name=feature_names)
        else:
            model.fit(X_train, y_train)
        y_pred = model.predict(X_test)

        r2 = r2_score(y_test, y_pred)
        mae = mean_absolute_error(y_test, y_pred)
        r2s.append(r2)
        maes.append(mae)

        print(f"[{name}] Fold {fold}: R2={r2:.3f}, MAE={mae:.3f}")

    r2_mean, r2_std = np.mean(r2s), np.std(r2s)
    mae_mean, mae_std = np.mean(maes), np.std(maes)
    print(
        f"== {name} CV summary: R2={r2_mean:.3f}±{r2_std:.3f}, "
        f"MAE={mae_mean:.3f}±{mae_std:.3f}\n"
    )
    return r2_mean, r2_std, mae_mean, mae_std


# -------------------------------------------------
# 4. RUN CV EXPERIMENTS
# -------------------------------------------------
cv_records = []

for tgt_key, tgt_desc in TARGETS.items():
    y = df[tgt_key].values

    # Base features
    X_base = df[BASE_FEATURES].values.astype(np.float64)
    model_base = LGBMRegressor(
        random_state=RANDOM_STATE,
        n_estimators=500,
        learning_rate=0.05,
        num_leaves=31,
        subsample=0.9,
        colsample_bytree=0.9,
    )
    res_base = run_cv(X_base, y, model_base, name=f"{tgt_key}_base", feature_names=BASE_FEATURES)
    cv_records.append(
        {
            "target": tgt_key,
            "features": "base",
            "R2_mean": res_base[0],
            "R2_std": res_base[1],
            "MAE_mean": res_base[2],
            "MAE_std": res_base[3],
        }
    )

    # Features + TSI
    X_tsi = df[FEATURES_WITH_TSI].values.astype(np.float64)
    model_tsi = LGBMRegressor(
        random_state=RANDOM_STATE,
        n_estimators=500,
        learning_rate=0.05,
        num_leaves=31,
        subsample=0.9,
        colsample_bytree=0.9,
    )
    res_tsi = run_cv(X_tsi, y, model_tsi, name=f"{tgt_key}_with_tsi", feature_names=FEATURES_WITH_TSI)
    cv_records.append(
        {
            "target": tgt_key,
            "features": "with_tsi",
            "R2_mean": res_tsi[0],
            "R2_std": res_tsi[1],
            "MAE_mean": res_tsi[2],
            "MAE_std": res_tsi[3],
        }
    )

cv_df = pd.DataFrame(cv_records)
cv_df.to_csv(os.path.join(OUTDIR, "cv_results.csv"), index=False)
print("CV results saved to cv_results.csv")

# -------------------------------------------------
# 5. TRAIN FINAL MODELS (WITH TSI)
# -------------------------------------------------
final_models = {}
X_all = df[FEATURES_WITH_TSI].values.astype(np.float64)

for tgt_key in TARGETS.keys():
    y = df[tgt_key].values
    model = LGBMRegressor(
        random_state=RANDOM_STATE,
        n_estimators=800,
        learning_rate=0.03,
        num_leaves=31,
        subsample=0.9,
        colsample_bytree=0.9,
    )
    model.fit(X_all, y, feature_name=FEATURES_WITH_TSI)
    final_models[tgt_key] = model

    # save feature importances
    imp = pd.Series(model.feature_importances_, index=FEATURES_WITH_TSI)
    imp.sort_values(ascending=False).to_csv(
        os.path.join(OUTDIR, f"feature_importances_{tgt_key}.csv")
    )

print("Final models trained and feature importances saved.")

# Separate model for water (for Pareto + BO)
y_water = df["water"].values
water_model = LGBMRegressor(
    random_state=RANDOM_STATE,
    n_estimators=800,
    learning_rate=0.03,
    num_leaves=31,
    subsample=0.9,
    colsample_bytree=0.9,
)
water_model.fit(X_all, y_water)
print("Water model trained.")

# -------------------------------------------------
# 6. SHAP SUMMARY + 2×2 DEPENDENCE
# -------------------------------------------------
# shap.initjs()  # not required when saving PNGs only

for tgt_key, tgt_desc in TARGETS.items():
    model = final_models[tgt_key]

    # sample for SHAP
    df_shap = df.sample(
        min(N_SHAP_SAMPLES, len(df)), random_state=RANDOM_STATE
    )
    X_shap = df_shap[FEATURES_WITH_TSI]

    explainer = shap.TreeExplainer(model)
    shap_values = explainer.shap_values(X_shap)

    # --- summary plot ---
    plt.figure()
    shap.summary_plot(
        shap_values,
        X_shap,
        feature_names=FEATURES_WITH_TSI,
        show=False,
    )
    plt.title(f"SHAP summary – {tgt_key}")
    plt.tight_layout()
    plt.savefig(
        os.path.join(OUTDIR, f"shap_summary_{tgt_key}.png"), dpi=600
    )
    plt.close()

    # --- combined 2×2 dependence figure (only for oil_org) ---
    if tgt_key == "oil_org":
        # 4 panels for oil_org:
        # (a) temp vs SHAP(oil_org), coloured by vol
        # (b) TSI vs SHAP(oil_org), coloured by lignin
        # (c) lignin vs SHAP(oil_org), coloured by cellulose
        # (d) Ash vs SHAP(oil_org), coloured by Flow_N2
        panels = [
            ("temp",   "vol"),
            ("TSI",    "lignin"),
            ("lignin", "cellulose"),
            ("Ash",    "Flow-N2"),
        ]

        fig, axes = plt.subplots(2, 2, figsize=(12, 10))

        for ax, (feat_x, feat_color) in zip(axes.ravel(), panels):
            shap.dependence_plot(
                feat_x,
                shap_values,
                X_shap,
                feature_names=FEATURES_WITH_TSI,
                interaction_index=feat_color,
                ax=ax,
                show=False,
            )
            # Update axis labels with units
            ax.set_xlabel(get_label_with_unit(feat_x), fontsize=12)
            ax.set_ylabel(f"SHAP value for {tgt_key}", fontsize=12)
            
            # Update title with units for both features
            ax.set_title(f"{get_label_with_unit(feat_x)} (color: {get_label_with_unit(feat_color)})", fontsize=11)
            ax.axhline(0, ls="--", c="gray", linewidth=0.8)

        plt.tight_layout()
        plt.savefig(
            os.path.join(OUTDIR, "shap_2x2_oil_org.png"), dpi=600
        )
        plt.close()

print("SHAP summary and 2×2 dependence plots saved.")

# -------------------------------------------------
# 6B. SHAP RANK STABILITY ANALYSIS (Oil vs Bio-liquid)
# -------------------------------------------------
print("Running SHAP rank stability analysis...")

targets_for_rank = ["oil_org", "bio_liquid"]
rank_results = {t: [] for t in targets_for_rank}

for tgt in targets_for_rank:
    model = final_models[tgt]

    for _ in range(N_BOOT_RANK):
        # bootstrap sample
        df_boot = df.sample(frac=0.6, replace=True, random_state=None)
        Xb = df_boot[FEATURES_WITH_TSI]

        expl = shap.TreeExplainer(model)
        sv = expl.shap_values(Xb)          # (n_samples, n_features)

        # rank features by mean(|SHAP|)
        mean_abs = np.abs(sv).mean(axis=0)
        ranks = mean_abs.argsort().argsort() + 1  # 1 = most important
        rank_results[tgt].append(ranks)

# Convert to arrays
rank_oil = np.array(rank_results["oil_org"])        # (N_BOOT, n_features)
rank_liq = np.array(rank_results["bio_liquid"])

# Compute statistics
mean_oil = rank_oil.mean(axis=0)
std_oil = rank_oil.std(axis=0)

mean_liq = rank_liq.mean(axis=0)
std_liq = rank_liq.std(axis=0)

features = FEATURES_WITH_TSI

# Plot rank stability
x = np.arange(len(features))
width = 0.35

plt.figure(figsize=(10, 6))
plt.bar(x - width/2, mean_oil, width,
        label="Oil_org", yerr=std_oil, capsize=6, color="#3B82F6",
        error_kw={'linewidth': 2.5, 'capthick': 2.5})
plt.bar(x + width/2, mean_liq, width,
        label="Bio-liquid", yerr=std_liq, capsize=6, color="#F97316",
        error_kw={'linewidth': 2.5, 'capthick': 2.5})

plt.xticks(x, features, rotation=60, ha='right', fontsize=16)
plt.yticks(fontsize=16)
plt.ylabel("Mean rank (lower = more important)", fontsize=16)
plt.title("SHAP Rank Stability — Oil vs Bio-liquid", fontsize=16)
plt.legend(fontsize=16, frameon=True, shadow=True, loc='best')
plt.gca().invert_yaxis()  # rank 1 at the top (optional)
plt.grid(axis='y', alpha=0.3, linewidth=1.0)
plt.tight_layout()

plt.savefig(os.path.join(OUTDIR, "shap_rank_stability_oil_vs_liquid.png"),
            dpi=600)
plt.close()

print("Saved SHAP rank stability plot.")

# -------------------------------------------------
# 7. PARTIAL DEPENDENCE PLOTS (PDP)
# -------------------------------------------------
print("Computing partial dependence plots (PDP)...")

PDP_FEATURES = ["temp", "HR", "PS", "Flow-N2"]

for tgt in ["oil_org", "bio_liquid"]:
    model = final_models[tgt]
    X_all_pd = df[FEATURES_WITH_TSI]

    fig, axs = plt.subplots(2, 2, figsize=(12, 10))

    for ax, feat in zip(axs.ravel(), PDP_FEATURES):
        pdp_res = partial_dependence(
            model,
            X_all_pd,
            [FEATURES_WITH_TSI.index(feat)],
            grid_resolution=50
        )
        xs = pdp_res["grid_values"][0]
        ys = pdp_res["average"][0]
        ax.plot(xs, ys, linewidth=2)
        ax.set_title(f"{feat} → {tgt}")
        ax.set_xlabel(get_label_with_unit(feat), fontsize=14, fontweight='bold')
        ax.set_ylabel(f"PDP({tgt}) (wt%)", fontsize=14, fontweight='bold')
        ax.tick_params(labelsize=12)
        ax.set_title(f"{feat}", fontsize=15, fontweight='bold')

    plt.tight_layout()
    plt.savefig(os.path.join(OUTDIR, f"pdp_{tgt}.png"), dpi=600, format="png")
    plt.savefig(os.path.join(OUTDIR, f"pdp_{tgt}.eps"), format="eps")
    plt.close()

print("PDP figures saved.")

# -------------------------------------------------
# 8. ACCUMULATED LOCAL EFFECTS (ALE) – COMBINED FIGURES
# -------------------------------------------------
try:
    from alibi.explainers import ALE
except ImportError:
    ALE = None
    print("[WARN] alibi not installed. ALE plots will be skipped. "
          "Install via: pip install alibi")

if ALE is None:
    print("[INFO] Skipping ALE plots (alibi not installed).")
else:
    print("Computing ALE plots (combined multi-panel figures)...")

    X_train = df[FEATURES_WITH_TSI].values.astype(float)

    # Split features into operating conditions vs composition - UPDATED WITH CAPITAL LETTERS
    operating_feats = ["temp", "HR", "PS", "Flow-N2"]
    composition_feats = [
        "cellulose", "hemicellulose", "lignin",
        "Ash", "vol", "fc", "C", "H", "O", "N",
        "hceff", "TSI"   # you can drop TSI here if you prefer
    ]

    # Helper: get index of feature name in FEATURES_WITH_TSI
    feat_idx = {name: FEATURES_WITH_TSI.index(name)
                for name in FEATURES_WITH_TSI}

    for tgt_key in TARGETS.keys():
        model = final_models[tgt_key]

        # One explainer per target model
        ale_explainer = ALE(
            predictor=model.predict,
            feature_names=FEATURES_WITH_TSI
        )

        # ---------- Figure A: operating parameters (2x2) ----------
        fig, axes = plt.subplots(
            nrows=2, ncols=2,
            figsize=(12, 10)
        )
        axes = axes.ravel()

        for ax, feat in zip(axes, operating_feats):
            if feat not in feat_idx:
                ax.set_visible(False)
                continue

            f_idx = feat_idx[feat]
            try:
                explanation = ale_explainer.explain(
                    X_train,
                    features=[f_idx]
                )
                xs = explanation.feature_values[0]
                ys = explanation.ale_values[0]
            except Exception as e:
                print(f"[ALE ERROR] {tgt_key} – {feat}: {e}")
                ax.set_visible(False)
                continue

            ax.plot(xs, ys, marker="o", linewidth=1.2, markersize=3)
            ax.set_xlabel(get_label_with_unit(feat), fontsize=14)
            ax.set_ylabel("ALE (wt%)", fontsize=14)
            ax.grid(True, linewidth=0.8, alpha=0.3)
            ax.tick_params(labelsize=12, width=1.5, length=6)
            ax.set_title(f"{feat}", fontsize=15)

        fig.suptitle(f"ALE – Operating parameters – {tgt_key}",
                     fontsize=17, fontweight='bold', y=0.995)
        fig.tight_layout(rect=[0, 0, 1, 0.97], h_pad=2.5, w_pad=2.0)
        fig.savefig(
            os.path.join(OUTDIR,
                         f"ale_operating_{tgt_key}.png"),
            dpi=600
        )
        plt.close(fig)

        # ---------- Figure B: composition features (3x4) ----------
        n_rows, n_cols = 3, 4
        fig, axes = plt.subplots(
            nrows=n_rows, ncols=n_cols,
            figsize=(14, 10)
        )
        axes = axes.ravel()

        for ax, feat in zip(axes, composition_feats):
            if feat not in feat_idx:
                ax.set_visible(False)
                continue

            f_idx = feat_idx[feat]
            try:
                explanation = ale_explainer.explain(
                    X_train,
                    features=[f_idx]
                )
                xs = explanation.feature_values[0]
                ys = explanation.ale_values[0]
            except Exception as e:
                print(f"[ALE ERROR] {tgt_key} – {feat}: {e}")
                ax.set_visible(False)
                continue

            ax.plot(xs, ys, marker="o", linewidth=1.2, markersize=3)
            ax.set_xlabel(get_label_with_unit(feat), fontsize=14)
            ax.set_ylabel("ALE (wt%)", fontsize=14)
            ax.grid(True, linewidth=0.6, alpha=0.3)
            ax.tick_params(labelsize=12, width=1.2, length=6)
            ax.set_title(f"{feat}", fontsize=15)

        # Hide any leftover empty axes if composition_feats < n_rows*n_cols
        for k in range(len(composition_feats), len(axes)):
            axes[k].set_visible(False)

        fig.suptitle(f"ALE – Biomass composition – {tgt_key}",
                     fontsize=17, fontweight='bold', y=0.995)
        fig.tight_layout(rect=[0, 0, 1, 0.97], h_pad=2.5, w_pad=2.0)
        fig.savefig(
            os.path.join(OUTDIR,
                         f"ale_composition_{tgt_key}.png"),
            dpi=600
        )
        plt.close(fig)

    print("ALE multi-panel figures saved.")



# -------------------------------------------------
# 9. SENSITIVITY ANALYSIS (±10%, tornado plot)
# -------------------------------------------------
print("Running sensitivity analysis (±10%) for oil_org...")

SENS_FEATURES = ["temp", "HR", "PS", "Flow-N2"]

for tgt in ["oil_org"]:
    model = final_models[tgt]
    base_series = df[FEATURES_WITH_TSI].mean()

    sensitivities = []

    for feat in SENS_FEATURES:
        # baseline prediction
        X0 = np.array([base_series.values])
        y0 = model.predict(X0)[0]

        # +10%
        base_pos = base_series.copy()
        base_pos[feat] = base_series[feat] * 1.10
        X_pos = np.array([base_pos.values])
        y_pos = model.predict(X_pos)[0]

        # -10%
        base_neg = base_series.copy()
        base_neg[feat] = base_series[feat] * 0.90
        X_neg = np.array([base_neg.values])
        y_neg = model.predict(X_neg)[0]

        delta = max(abs(y_pos - y0), abs(y_neg - y0))
        sensitivities.append((feat, delta))

    # sort by sensitivity magnitude
    sensitivities = sorted(sensitivities, key=lambda x: x[1], reverse=True)

    feats = [x[0] for x in sensitivities]
    vals = [x[1] for x in sensitivities]

    plt.figure(figsize=(8, 6))
    plt.barh(feats, vals, color="steelblue")
    plt.xlabel("Sensitivity (|Δ predicted oil_org|, wt%)")
    plt.title("One-At-A-Time Sensitivity Analysis (±10%)")
    plt.gca().invert_yaxis()
    plt.tight_layout()
    plt.savefig(os.path.join(OUTDIR, "sensitivity_oil_org.png"), dpi=600)
    plt.close()

print("Sensitivity analysis saved.")

# -------------------------------------------------
# 10. PARETO EXPLORATION: ORGANIC OIL vs WATER
# -------------------------------------------------
print("Running Pareto exploration (oil_org vs water)...")
rng = np.random.default_rng(RANDOM_STATE)

def sample_uniform(col, n):
    return rng.uniform(df[col].min(), df[col].max(), size=n)

candidates = pd.DataFrame()

# resample compositional / structural features - UPDATED WITH CAPITAL LETTERS
for col in [
    "Ash",
    "fc",
    "vol",
    "C",
    "H",
    "O",
    "N",
    "hceff",
    "cellulose",
    "hemicellulose",
    "lignin",
]:
    candidates[col] = rng.choice(df[col].values, size=N_PARETO_SAMPLES, replace=True)

# process variables
for col in ["temp", "HR", "PS", "Flow-N2"]:
    candidates[col] = sample_uniform(col, N_PARETO_SAMPLES)

# recompute TSI
candidates["TSI"] = (
    candidates["temp"] * np.log(candidates["HR"] + 1.0) / (candidates["PS"] + eps)
)

# predict organic oil & water
X_cand = candidates[FEATURES_WITH_TSI].values

oil_model = final_models["oil_org"]

candidates["oil_org_pred"] = oil_model.predict(X_cand)
candidates["water_pred"] = water_model.predict(X_cand)

# Pareto front: maximise oil_org_pred, minimise water_pred
def pareto_front(df_points, oil_col="oil_org_pred", water_col="water_pred"):
    vals = df_points[[oil_col, water_col]].values
    n = vals.shape[0]
    dominated = np.zeros(n, dtype=bool)

    for i in range(n):
        if dominated[i]:
            continue
        oil_i, w_i = vals[i]
        better_oil = vals[:, 0] >= oil_i
        lower_water = vals[:, 1] <= w_i
        mask = better_oil & lower_water
        mask[i] = False
        if np.any(mask & ((vals[:, 0] > oil_i) | (vals[:, 1] < w_i))):
            dominated[i] = True

    return np.where(~dominated)[0]

pareto_idx = pareto_front(candidates)
pareto_df = candidates.iloc[pareto_idx].copy()

candidates.to_csv(os.path.join(OUTDIR, "pareto_candidates.csv"), index=False)
pareto_df.to_csv(os.path.join(OUTDIR, "pareto_front.csv"), index=False)

print(f"Pareto front size: {len(pareto_df)}")

# scatter plot
plt.figure()
plt.scatter(
    candidates["water_pred"],
    candidates["oil_org_pred"],
    s=5,
    alpha=0.2,
    label="Candidates",
)
plt.scatter(
    pareto_df["water_pred"],
    pareto_df["oil_org_pred"],
    s=15,
    alpha=0.9,
    label="Pareto front",
)
plt.xlabel("Predicted water content (wt%)")
plt.ylabel("Predicted organic oil yield (wt%)")
plt.legend()
plt.tight_layout()
plt.savefig(os.path.join(OUTDIR, "pareto_oil_vs_water.png"), dpi=600)
plt.close()

print("Pareto plot and CSVs saved in", OUTDIR)

# -------------------------------------------------
# 11. BAYESIAN OPTIMISATION: OIL_org vs WATER
# -------------------------------------------------
if gp_minimize is None:
    print("[INFO] Skipping Bayesian optimisation (scikit-optimize not installed).")
else:
    print("Running Bayesian optimisation for operating condition selection...")

    # For BO, vary only process variables; fix composition at dataset mean - UPDATED WITH CAPITAL LETTERS
    comp_cols = [
        "Ash", "fc", "vol", "C", "H", "O", "N", "hceff",
        "cellulose", "hemicellulose", "lignin",
    ]
    comp_mean = df[comp_cols].mean()

    def make_feature_row(temp, hr, ps, Flow_N2):
        """Build a single-row ndarray with all features for prediction."""
        row = comp_mean.copy()
        row["temp"] = temp
        row["HR"] = hr
        row["PS"] = ps
        row["Flow-N2"] = Flow_N2
        row["TSI"] = temp * np.log(hr + 1.0) / (ps + eps)
        return row[FEATURES_WITH_TSI].values.reshape(1, -1)

    # Bounds from observed data
    def minmax(col):
        return float(df[col].min()), float(df[col].max())

    t_min, t_max = minmax("temp")
    hr_min, hr_max = minmax("HR")
    ps_min, ps_max = minmax("PS")
    f_min, f_max = minmax("Flow-N2")

    space = [
        Real(t_min,  t_max,  name="temp"),
        Real(hr_min, hr_max, name="HR"),
        Real(ps_min, ps_max, name="PS"),
        Real(f_min,  f_max,  name="Flow-N2"),
    ]

    # --- Objective 1: Maximise oil_org (minimise negative oil_org) ---
    def obj_max_oil(x):
        temp, hr, ps, Flow_N2 = x
        X = make_feature_row(temp, hr, ps, Flow_N2)
        oil_pred = oil_model.predict(X)[0]
        return -oil_pred  # gp_minimize -> minimisation

    res_oil = gp_minimize(
        obj_max_oil,
        space,
        n_calls=40,
        n_random_starts=10,
        random_state=RANDOM_STATE,
    )

    best_temp_o, best_hr_o, best_ps_o, best_flow_o = res_oil.x
    X_best_o = make_feature_row(best_temp_o, best_hr_o, best_ps_o, best_flow_o)
    best_oil = oil_model.predict(X_best_o)[0]
    best_water_at_o = water_model.predict(X_best_o)[0]

    # --- Objective 2: Minimise water ---
    def obj_min_water(x):
        temp, hr, ps, Flow_N2 = x
        X = make_feature_row(temp, hr, ps, Flow_N2)
        water_pred = water_model.predict(X)[0]
        return water_pred

    res_water = gp_minimize(
        obj_min_water,
        space,
        n_calls=40,
        n_random_starts=10,
        random_state=RANDOM_STATE + 1,
    )

    best_temp_w, best_hr_w, best_ps_w, best_flow_w = res_water.x
    X_best_w = make_feature_row(best_temp_w, best_hr_w, best_ps_w, best_flow_w)
    best_water = water_model.predict(X_best_w)[0]
    best_oil_at_w = oil_model.predict(X_best_w)[0]

    # Collect single-objective BO results
    bo_single = pd.DataFrame([
        {
            "Objective": "Maximise Oil_org",
            "temp": best_temp_o,
            "HR": best_hr_o,
            "PS": best_ps_o,
            "Flow-N2": best_flow_o,
            "Pred_Oil_org": best_oil,
            "Pred_Water": best_water_at_o,
        },
        {
            "Objective": "Minimise Water",
            "temp": best_temp_w,
            "HR": best_hr_w,
            "PS": best_ps_w,
            "Flow-N2": best_flow_w,
            "Pred_Oil_org": best_oil_at_w,
            "Pred_Water": best_water,
        },
    ])
    bo_single.to_csv(os.path.join(OUTDIR, "bo_single_results.csv"), index=False)
    print("Bayesian optimisation single-objective results saved to bo_single_results.csv")

    # --- Scalarised multi-objective: Oil_org - λ * Water ---
    lambdas = [0.15, 0.25, 0.35]
    bo_scalar_rows = []

    for lam in lambdas:
        def obj_scalar(x, lam=lam):
            temp, hr, ps, Flow_N2 = x
            X = make_feature_row(temp, hr, ps, Flow_N2)
            oil_pred = oil_model.predict(X)[0]
            water_pred = water_model.predict(X)[0]
            score = oil_pred - lam * water_pred
            return -score  # maximise Oil_org - λ * Water

        res_scalar = gp_minimize(
            obj_scalar,
            space,
            n_calls=40,
            n_random_starts=10,
            random_state=RANDOM_STATE + int(lam * 100),
        )

        bt, bhr, bps, bflow = res_scalar.x
        X_best = make_feature_row(bt, bhr, bps, bflow)
        oil_pred = oil_model.predict(X_best)[0]
        water_pred = water_model.predict(X_best)[0]
        score = oil_pred - lam * water_pred

        bo_scalar_rows.append({
            "lambda": lam,
            "temp": bt,
            "HR": bhr,
            "PS": bps,
            "Flow-N2": bflow,
            "Oil_org": oil_pred,
            "Water": water_pred,
            "Score": score,
        })

    bo_scalar = pd.DataFrame(bo_scalar_rows)
    bo_scalar.to_csv(os.path.join(OUTDIR, "bo_scalar_results.csv"), index=False)
    print("Bayesian optimisation scalarised results saved to bo_scalar_results.csv")

print("All done.")
