utils/calculations.py

"""AI models and optimization calculations for the PFAS-SBEAD pipeline."""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from sklearn.ensemble import (
    GradientBoostingClassifier,
    GradientBoostingRegressor,
    IsolationForest,
    RandomForestRegressor,
)
from sklearn.model_selection import cross_val_score


FEATURE_COLUMNS = [
    "OLR_kg_m3_d",
    "HRT_days",
    "pH",
    "temperature_C",
    "COD_mg_L",
    "VFA_mg_L",
    "alkalinity_mg_CaCO3_L",
    "voltage_V",
    "current_A",
    "current_density_A_m2",
    "conductivity_mS_cm",
    "electrode_area_m2",
    "electrode_spacing_cm",
    "initial_PFAS_ug_L",
]

AI_SCORE_WEIGHTS = {
    "pfas_degradation": 0.40,
    "fluoride_release": 0.30,
    "short_chain_risk": -0.15,
    "energy_input": -0.10,
    "instability": -0.05,
}


def compute_ai_score(
    degradation_norm: float,
    fluoride_norm: float,
    short_chain_risk: float,
    energy_norm: float,
    instability: float,
) -> float:
    """Compute the composite AI optimization score."""
    score = (
        AI_SCORE_WEIGHTS["pfas_degradation"] * degradation_norm
        + AI_SCORE_WEIGHTS["fluoride_release"] * fluoride_norm
        + AI_SCORE_WEIGHTS["short_chain_risk"] * short_chain_risk
        + AI_SCORE_WEIGHTS["energy_input"] * energy_norm
        + AI_SCORE_WEIGHTS["instability"] * instability
    )
    return float(np.clip(score, 0, 1))


def train_degradation_model(df: pd.DataFrame) -> tuple[GradientBoostingRegressor, float]:
    """XGBoost-style regressor for PFAS degradation prediction."""
    X = df[FEATURE_COLUMNS].values
    y = df["PFAS_degradation_pct"].values
    model = GradientBoostingRegressor(
        n_estimators=100, max_depth=4, learning_rate=0.1, random_state=42
    )
    scores = cross_val_score(model, X, y, cv=5, scoring="r2")
    model.fit(X, y)
    return model, float(scores.mean())


def train_fluoride_model(df: pd.DataFrame) -> tuple[RandomForestRegressor, float]:
    """Random Forest for fluoride release prediction."""
    X = df[FEATURE_COLUMNS].values
    y = df["fluoride_release_mg_L"].values
    model = RandomForestRegressor(n_estimators=100, max_depth=5, random_state=42)
    scores = cross_val_score(model, X, y, cv=5, scoring="r2")
    model.fit(X, y)
    return model, float(scores.mean())


def train_short_chain_classifier(df: pd.DataFrame) -> tuple[GradientBoostingClassifier, float]:
    """XGBoost classifier for short-chain PFAS risk."""
    X = df[FEATURE_COLUMNS].values
    y = (df["short_chain_formation_ratio"] > 0.3).astype(int).values
    model = GradientBoostingClassifier(
        n_estimators=80, max_depth=3, learning_rate=0.1, random_state=42
    )
    scores = cross_val_score(model, X, y, cv=5, scoring="accuracy")
    model.fit(X, y)
    return model, float(scores.mean())


def train_instability_detector(df: pd.DataFrame) -> IsolationForest:
    """Isolation Forest for reactor instability detection."""
    stability_cols = ["pH_drop", "VFA_accumulation_mg_L", "ORP_drift_mV", "current_instability_index"]
    X = df[stability_cols].values
    model = IsolationForest(contamination=0.15, random_state=42)
    model.fit(X)
    return model


def compute_shap_importance(model: Any, df: pd.DataFrame) -> pd.DataFrame:
    """Compute feature importance as SHAP-like proxy using tree-based importances."""
    X = df[FEATURE_COLUMNS]
    importances = model.feature_importances_
    imp_df = pd.DataFrame({
        "feature": FEATURE_COLUMNS,
        "importance": importances,
    }).sort_values("importance", ascending=False).reset_index(drop=True)
    return imp_df


def bayesian_next_recommendation(df: pd.DataFrame) -> dict[str, Any]:
    """
    Simple Bayesian-inspired recommendation for the next experiment.
    Finds conditions that maximize expected AI score based on model trends.
    """
    top_pct = df.nlargest(10, "AI_score")
    rec = {}
    for col in FEATURE_COLUMNS:
        rec[col] = float(np.round(top_pct[col].mean(), 3))

    rec["predicted_degradation_pct"] = float(np.round(top_pct["PFAS_degradation_pct"].mean(), 1))
    rec["predicted_fluoride_release"] = float(np.round(top_pct["fluoride_release_mg_L"].mean(), 1))
    rec["expected_ai_score"] = float(np.round(top_pct["AI_score"].mean(), 3))
    rec["confidence"] = float(np.round(1.0 - top_pct["AI_score"].std(), 2))

    return rec


def sensitivity_analysis(df: pd.DataFrame) -> pd.DataFrame:
    """Compute correlation-based sensitivity of AI score to each input feature."""
    correlations = []
    for col in FEATURE_COLUMNS:
        corr = df[col].corr(df["AI_score"])
        correlations.append({"feature": col, "correlation_with_AI_score": round(corr, 4)})
    sens_df = pd.DataFrame(correlations)
    sens_df["abs_correlation"] = sens_df["correlation_with_AI_score"].abs()
    return sens_df.sort_values("abs_correlation", ascending=False).reset_index(drop=True)