"""
Layer 2: Model Training
- Train multi-target stacked LightGBM predictor.
- Base models (LOO CV) + Meta-models + Quantile regression.
- SHAP explainer for interpretability.
"""

import pandas as pd
import numpy as np
import pickle
import os
import sys
from lightgbm import LGBMRegressor
from sklearn.model_selection import LeaveOneOut, cross_val_predict
import shap
import warnings

warnings.filterwarnings('ignore')

# Add parent directory to path for config import
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import CFG

BOUNDS = {
    "Hardness": (0, 200),
    "Friability": (0, 5),
    "Dissolution_Rate": (50, 100),
    "Content_Uniformity": (80, 120),
    "Disintegration_Time": (0, 60)
}
def predict_batch(X_new: pd.DataFrame,
                  base_median: dict, base_q10: dict, base_q90: dict,
                  meta: dict) -> dict:
    """
    Exposed function for real-time quality prediction of one batch.
    """
    ranges = {
        "Hardness": [0, 200],
        "Friability": [0, 5],
        "Dissolution_Rate": [50, 100],
        "Content_Uniformity": [80, 120],
        "Disintegration_Time": [0, 60]
    }
    
    results = {}
    base_preds_df = pd.DataFrame(index=X_new.index)
    
    # Step 2: Meta Median Prediction
    for target in CFG.TARGET_COLS:
        # Base Prediction for this target
        b_pred = base_median[target].predict(X_new)
        
        # Meta features: X + base_pred
        X_meta = pd.concat([X_new.reset_index(drop=True), 
                            pd.Series(b_pred, name=f"base_{target}")], axis=1)
        X_meta.columns = [str(c) for c in X_meta.columns]
        
        median = meta[target].predict(X_meta.astype(float))[0]
        q10 = base_q10[target].predict(X_new.astype(float))[0]
        q90 = base_q90[target].predict(X_new.astype(float))[0]
        
        # Clip
        c_min, c_max = ranges.get(target, [0, 100])
        median = np.clip(median, c_min, c_max)
        q10 = np.clip(q10, c_min, c_max)
        q90 = np.clip(q90, c_min, c_max)
        
        results[target] = {
            "median": float(median),
            "lower": float(q10),
            "upper": float(q90)
        }
        
    return results

def main():
    print(">>> Starting Layer 2: Stacked Model Training")
    with open(os.path.join(CFG.PROC_DIR, "X_final.pkl"), "rb") as f:
        X = pickle.load(f)
    with open(os.path.join(CFG.PROC_DIR, "production_clean.pkl"), "rb") as f:
        dfy = pickle.load(f)
        
    y = dfy.set_index("Batch_ID").loc[X.index, CFG.TARGET_COLS]
    
    base_models = {}
    q10_models = {}
    q90_models = {}
    meta_models = {}
    loo_preds = pd.DataFrame(index=X.index, columns=CFG.TARGET_COLS)
    
    loo = LeaveOneOut()
    
    for target in CFG.TARGET_COLS:
        print(f"Training models for {target}...")
        
        # Step 1: Base Model (LOO CV for meta-features)
        base_model = LGBMRegressor(**CFG.LGB_PARAMS)
        oof_preds = cross_val_predict(base_model, X, y[target], cv=loo)
        base_model.fit(X, y[target])
        base_models[target] = base_model
        
        # Step 2: Quantile Models (Uncertainty)
        q10_params = CFG.LGB_PARAMS.copy()
        q10_params.update({"objective": "quantile", "alpha": 0.10})
        q10_model = LGBMRegressor(**q10_params)
        q10_model.fit(X, y[target])
        q10_models[target] = q10_model
        
        q90_params = CFG.LGB_PARAMS.copy()
        q90_params.update({"objective": "quantile", "alpha": 0.90})
        q90_model = LGBMRegressor(**q90_params)
        q90_model.fit(X, y[target])
        q90_models[target] = q90_model
        
        # Step 3: Meta-model
        # Use OOF preds as features
        X_meta = pd.concat([X, pd.Series(oof_preds, index=X.index, name=f"base_{target}")], axis=1)
        X_meta.columns = [str(c) for c in X_meta.columns]
        
        meta_model = LGBMRegressor(**CFG.LGB_PARAMS)
        # Final OOF for evaluation
        loo_preds[target] = cross_val_predict(meta_model, X_meta, y[target], cv=loo)
        meta_model.fit(X_meta, y[target])
        meta_models[target] = meta_model

    # SHAP Explainer (on Dissolution meta model)
    try:
        example_meta = pd.concat([X, loo_preds["Dissolution_Rate"]], axis=1)
        example_meta.columns = [str(c) for c in example_meta.columns]
        explainer = shap.TreeExplainer(meta_models["Dissolution_Rate"])
        with open(os.path.join(CFG.MODEL_DIR, "shap_explainer.pkl"), "wb") as f:
            pickle.dump(explainer, f)
    except Exception as e:
        print(f"SHAP Warning: {e}")

    # Clip LOO predictions for evaluation
    ranges = {
        "Hardness": [0, 200],
        "Friability": [0, 5],
        "Dissolution_Rate": [50, 100],
        "Content_Uniformity": [80, 120],
        "Disintegration_Time": [0, 60]
    }
    for target, r in ranges.items():
        loo_preds[target] = np.clip(loo_preds[target], r[0], r[1])

    # Save
    with open(os.path.join(CFG.MODEL_DIR, "base_models_median.pkl"), "wb") as f:
        pickle.dump(base_models, f)
    with open(os.path.join(CFG.MODEL_DIR, "base_models_q10.pkl"), "wb") as f:
        pickle.dump(q10_models, f)
    with open(os.path.join(CFG.MODEL_DIR, "base_models_q90.pkl"), "wb") as f:
        pickle.dump(q90_models, f)
    with open(os.path.join(CFG.MODEL_DIR, "meta_models.pkl"), "wb") as f:
        pickle.dump(meta_models, f)
    with open(os.path.join(CFG.MODEL_DIR, "loo_predictions.pkl"), "wb") as f:
        pickle.dump(loo_preds, f)
        
    print("="*60)
    print(f"✅ LAYER 2 COMPLETE")
    print(f"   Models trained: 20 (5 targets x 4 types)")
    print(f"   Output file:  {os.path.join(CFG.MODEL_DIR, 'meta_models.pkl')}")
    print("="*60)

if __name__ == "__main__":
    main()