"""
ALWAS ML Models — Train and evaluate all 3 models:
1. Hours Estimation (XGBoost Regressor)
2. Complexity Classification (XGBoost + LightGBM Ensemble)
3. Bottleneck Risk Prediction (Gradient Boosting Classifier)
"""
import numpy as np
import pandas as pd
import json
import joblib
import os
from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold
from sklearn.preprocessing import LabelEncoder, OrdinalEncoder
from sklearn.metrics import (
    mean_absolute_error, mean_squared_error, r2_score,
    classification_report, confusion_matrix, accuracy_score,
    f1_score, roc_auc_score
)
from sklearn.calibration import CalibratedClassifierCV
import xgboost as xgb
import lightgbm as lgb

# === Load Data ===
print("=" * 60)
print("ALWAS ML MODEL TRAINING")
print("=" * 60)

df = pd.read_csv('/app/alwas_blocks_dataset.csv')
print(f"\nLoaded {len(df)} blocks ({df['is_completed'].sum()} completed, {(~df['is_completed'].astype(bool)).sum()} in-progress)")

# === Feature Engineering ===
print("\n--- Feature Engineering ---")

# Encode categoricals
tech_node_encoder = LabelEncoder()
block_type_encoder = LabelEncoder()
priority_encoder = OrdinalEncoder(categories=[['P4-Low', 'P3-Medium', 'P2-High', 'P1-Critical']])
engineer_encoder = LabelEncoder()

df['tech_node_encoded'] = tech_node_encoder.fit_transform(df['tech_node'])
df['block_type_encoded'] = block_type_encoder.fit_transform(df['block_type'])
df['priority_encoded'] = priority_encoder.fit_transform(df[['priority']]).astype(int).flatten()
df['engineer_encoded'] = engineer_encoder.fit_transform(df['engineer_id'])

# Create interaction features
df['type_node_interaction'] = df['tech_node_encoded'] * 10 + df['block_type_encoded']
df['complexity_score'] = df['constraint_complexity'] * df['transistor_count_log']
df['size_priority_interaction'] = df['transistor_count_log'] * df['priority_numeric']

# Encode targets early so slices have them
complexity_encoder = LabelEncoder()
df['complexity_encoded'] = complexity_encoder.fit_transform(df['complexity'])
bottleneck_encoder = LabelEncoder()
df['bottleneck_encoded'] = bottleneck_encoder.fit_transform(df['bottleneck_risk'])

# === MODEL 1: Hours Estimation ===
print("\n" + "=" * 60)
print("MODEL 1: Hours Estimation (XGBoost Regressor)")
print("=" * 60)

# Use completed blocks for training hours estimation
completed = df[df['is_completed'] == 1].copy()

HOURS_FEATURES = [
    'tech_node_encoded', 'block_type_encoded', 'priority_encoded',
    'transistor_count', 'transistor_count_log', 'has_dependencies',
    'num_dependencies', 'constraint_complexity', 'drc_iterations',
    'engineer_skill_factor', 'type_node_interaction', 'complexity_score',
    'size_priority_interaction'
]

X_hours = completed[HOURS_FEATURES]
y_hours = completed['actual_hours']

X_train_h, X_test_h, y_train_h, y_test_h = train_test_split(
    X_hours, y_hours, test_size=0.2, random_state=42
)

# XGBoost model
hours_model = xgb.XGBRegressor(
    n_estimators=500,
    learning_rate=0.05,
    max_depth=7,
    subsample=0.8,
    colsample_bytree=0.8,
    min_child_weight=3,
    reg_alpha=0.1,
    reg_lambda=1.0,
    objective='reg:squarederror',
    tree_method='hist',
    random_state=42,
    early_stopping_rounds=50,
)
hours_model.fit(
    X_train_h, y_train_h,
    eval_set=[(X_test_h, y_test_h)],
    verbose=False
)

# Evaluate
y_pred_h = hours_model.predict(X_test_h)
mae = mean_absolute_error(y_test_h, y_pred_h)
rmse = np.sqrt(mean_squared_error(y_test_h, y_pred_h))
r2 = r2_score(y_test_h, y_pred_h)
mape = np.mean(np.abs((y_test_h - y_pred_h) / y_test_h)) * 100

print(f"\nHours Estimation Results:")
print(f"  MAE:  {mae:.2f} hours")
print(f"  RMSE: {rmse:.2f} hours")
print(f"  R²:   {r2:.4f}")
print(f"  MAPE: {mape:.1f}%")

# Feature importance
importance = pd.DataFrame({
    'feature': HOURS_FEATURES,
    'importance': hours_model.feature_importances_
}).sort_values('importance', ascending=False)
print(f"\nTop features for hours estimation:")
print(importance.to_string(index=False))

# Cross-validation
cv_scores = cross_val_score(
    xgb.XGBRegressor(n_estimators=500, learning_rate=0.05, max_depth=7, 
                      subsample=0.8, colsample_bytree=0.8, tree_method='hist', random_state=42),
    X_hours, y_hours, cv=5, scoring='r2'
)
print(f"\n5-Fold CV R²: {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")

# === MODEL 2: Complexity Classification ===
print("\n" + "=" * 60)
print("MODEL 2: Complexity Classification (XGBoost + LightGBM Ensemble)")
print("=" * 60)

COMPLEXITY_FEATURES = [
    'tech_node_encoded', 'block_type_encoded', 'priority_encoded',
    'transistor_count', 'transistor_count_log', 'has_dependencies',
    'num_dependencies', 'constraint_complexity', 'drc_iterations',
    'type_node_interaction', 'complexity_score', 'size_priority_interaction'
]

X_comp = completed[COMPLEXITY_FEATURES]
y_comp = completed['complexity_encoded']

X_train_c, X_test_c, y_train_c, y_test_c = train_test_split(
    X_comp, y_comp, test_size=0.2, random_state=42, stratify=y_comp
)

# XGBoost classifier
xgb_clf = xgb.XGBClassifier(
    n_estimators=500,
    learning_rate=0.05,
    max_depth=6,
    subsample=0.8,
    colsample_bytree=0.8,
    objective='multi:softprob',
    num_class=3,
    tree_method='hist',
    random_state=42,
    early_stopping_rounds=50,
)
xgb_clf.fit(X_train_c, y_train_c, eval_set=[(X_test_c, y_test_c)], verbose=False)

# LightGBM classifier
lgb_clf = lgb.LGBMClassifier(
    n_estimators=500,
    learning_rate=0.05,
    num_leaves=63,
    subsample=0.8,
    colsample_bytree=0.8,
    random_state=42,
    verbose=-1,
)
lgb_clf.fit(X_train_c, y_train_c)

# Ensemble predictions (average probabilities)
xgb_proba = xgb_clf.predict_proba(X_test_c)
lgb_proba = lgb_clf.predict_proba(X_test_c)
ensemble_proba = (xgb_proba + lgb_proba) / 2
y_pred_c = np.argmax(ensemble_proba, axis=1)

accuracy = accuracy_score(y_test_c, y_pred_c)
f1 = f1_score(y_test_c, y_pred_c, average='weighted')

print(f"\nComplexity Classification Results (Ensemble):")
print(f"  Accuracy: {accuracy:.4f}")
print(f"  F1 (weighted): {f1:.4f}")
print(f"\nClassification Report:")
target_names = complexity_encoder.classes_
print(classification_report(y_test_c, y_pred_c, target_names=target_names))

# Per-model scores
xgb_acc = accuracy_score(y_test_c, xgb_clf.predict(X_test_c))
lgb_acc = accuracy_score(y_test_c, lgb_clf.predict(X_test_c))
print(f"  XGBoost alone: {xgb_acc:.4f}")
print(f"  LightGBM alone: {lgb_acc:.4f}")
print(f"  Ensemble:       {accuracy:.4f}")

# === MODEL 3: Bottleneck Risk Prediction ===
print("\n" + "=" * 60)
print("MODEL 3: Bottleneck Risk Prediction (Gradient Boosting)")
print("=" * 60)

# For bottleneck prediction, use ALL blocks (including in-progress)
BOTTLENECK_FEATURES = [
    'tech_node_encoded', 'block_type_encoded', 'priority_encoded',
    'transistor_count_log', 'has_dependencies', 'num_dependencies',
    'constraint_complexity', 'estimated_hours', 'hours_logged',
    'hours_over_estimate_ratio', 'drc_iterations', 'drc_violations_total',
    'lvs_mismatches_total', 'current_stage_idx', 'days_in_current_stage',
    'engineer_skill_factor', 'is_overdue', 'complexity_score'
]

X_bn = df[BOTTLENECK_FEATURES]
y_bn = df['bottleneck_encoded']

X_train_b, X_test_b, y_train_b, y_test_b = train_test_split(
    X_bn, y_bn, test_size=0.2, random_state=42, stratify=y_bn
)

# XGBoost classifier with calibration
base_bn_model = xgb.XGBClassifier(
    n_estimators=500,
    learning_rate=0.05,
    max_depth=6,
    subsample=0.8,
    colsample_bytree=0.8,
    scale_pos_weight=1,
    objective='multi:softprob',
    num_class=3,
    tree_method='hist',
    random_state=42,
)

# Calibrate probabilities
bn_model = CalibratedClassifierCV(base_bn_model, cv=3, method='isotonic')
bn_model.fit(X_train_b, y_train_b)

y_pred_b = bn_model.predict(X_test_b)
y_proba_b = bn_model.predict_proba(X_test_b)

bn_accuracy = accuracy_score(y_test_b, y_pred_b)
bn_f1 = f1_score(y_test_b, y_pred_b, average='weighted')

print(f"\nBottleneck Risk Prediction Results:")
print(f"  Accuracy: {bn_accuracy:.4f}")
print(f"  F1 (weighted): {bn_f1:.4f}")
print(f"\nClassification Report:")
bn_target_names = bottleneck_encoder.classes_
print(classification_report(y_test_b, y_pred_b, target_names=bn_target_names))

# === Save All Models & Artifacts ===
print("\n" + "=" * 60)
print("SAVING MODELS")
print("=" * 60)

os.makedirs('/app/models', exist_ok=True)

# Save models
joblib.dump(hours_model, '/app/models/hours_estimator.joblib')
joblib.dump(xgb_clf, '/app/models/complexity_xgb.joblib')
joblib.dump(lgb_clf, '/app/models/complexity_lgb.joblib')
joblib.dump(bn_model, '/app/models/bottleneck_predictor.joblib')

# Save encoders
joblib.dump(tech_node_encoder, '/app/models/tech_node_encoder.joblib')
joblib.dump(block_type_encoder, '/app/models/block_type_encoder.joblib')
joblib.dump(priority_encoder, '/app/models/priority_encoder.joblib')
joblib.dump(engineer_encoder, '/app/models/engineer_encoder.joblib')
joblib.dump(complexity_encoder, '/app/models/complexity_encoder.joblib')
joblib.dump(bottleneck_encoder, '/app/models/bottleneck_encoder.joblib')

# Save feature lists
feature_config = {
    'hours_features': HOURS_FEATURES,
    'complexity_features': COMPLEXITY_FEATURES,
    'bottleneck_features': BOTTLENECK_FEATURES,
    'tech_nodes': list(tech_node_encoder.classes_),
    'block_types': list(block_type_encoder.classes_),
    'priorities': ['P4-Low', 'P3-Medium', 'P2-High', 'P1-Critical'],
    'complexity_classes': list(complexity_encoder.classes_),
    'bottleneck_classes': list(bottleneck_encoder.classes_),
}
with open('/app/models/feature_config.json', 'w') as f:
    json.dump(feature_config, f, indent=2)

# Save model evaluation metrics
metrics = {
    'hours_estimation': {
        'mae': round(mae, 2),
        'rmse': round(rmse, 2),
        'r2': round(r2, 4),
        'mape_percent': round(mape, 1),
        'cv_r2_mean': round(cv_scores.mean(), 4),
        'cv_r2_std': round(cv_scores.std(), 4),
    },
    'complexity_classification': {
        'accuracy': round(accuracy, 4),
        'f1_weighted': round(f1, 4),
        'xgb_accuracy': round(xgb_acc, 4),
        'lgb_accuracy': round(lgb_acc, 4),
        'ensemble_accuracy': round(accuracy, 4),
    },
    'bottleneck_prediction': {
        'accuracy': round(bn_accuracy, 4),
        'f1_weighted': round(bn_f1, 4),
    },
    'training_data': {
        'total_samples': len(df),
        'completed_blocks': int(df['is_completed'].sum()),
        'in_progress_blocks': int((~df['is_completed'].astype(bool)).sum()),
    }
}
with open('/app/models/metrics.json', 'w') as f:
    json.dump(metrics, f, indent=2)

print(f"\nModels saved to /app/models/:")
for f in sorted(os.listdir('/app/models')):
    size = os.path.getsize(f'/app/models/{f}')
    print(f"  {f} ({size:,} bytes)")

print("\n" + "=" * 60)
print("ALL MODELS TRAINED SUCCESSFULLY")
print("=" * 60)