"""
Model Export Script
Regenerates and exports the trained model with current dependencies
Run this to fix numpy compatibility issues
"""

import os
import sys
import json
import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

print("\n" + "="*70)
print("📊 Model Export & Regeneration Script")
print("="*70)

# Step 1: Load data
print("\n[1/5] Loading data...")
try:
    df = pd.read_csv("../../cleaned_patient_data.csv")
    print(f"✅ Data loaded: {df.shape[0]} samples, {df.shape[1]} features")
except FileNotFoundError:
    print("❌ Error: cleaned_patient_data.csv not found!")
    print("Expected path: ../../cleaned_patient_data.csv")
    sys.exit(1)

# Step 2: Prepare features and target
print("\n[2/5] Preparing features and target...")
X = df.drop(columns=["workload_hours"])
y = df["workload_hours"]

print(f"✅ Original features: {X.shape[1]}")
print(f"   {X.columns.tolist()}")

# Compute interaction features from the 5 basic features
print("\n   Computing interaction features...")
X["alert_cognition_interaction"] = X["num_alerts_last_week"] * (30 - X["cognitive_score"])
X["activity_sleep_ratio"] = X["patient_activity_level"] / X["sleep_hours_avg"]
X["cognitive_activity_ratio"] = X["cognitive_score"] * X["patient_activity_level"] / 1000

print(f"✅ After interaction features: {X.shape[1]}")
print(f"   {X.columns.tolist()}")

# Step 3: Train/Test Split
print("\n[3/5] Splitting data...")
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
print(f"✅ Training: {X_train.shape[0]} samples")
print(f"✅ Testing: {X_test.shape[0]} samples")

# Step 4: Train Best Model (Linear Regression)
print("\n[4/5] Training Linear Regression model...")
model = LinearRegression()
model.fit(X_train, y_train)

# Evaluate
y_pred = model.predict(X_test)
r2 = r2_score(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))

print(f"✅ Model trained!")
print(f"   R² Score: {r2:.4f}")
print(f"   MAE: {mae:.4f}")
print(f"   RMSE: {rmse:.4f}")

if r2 >= 0.80:
    print(f"   ✅ DSO Target ACHIEVED (R² ≥ 0.80)")
else:
    print(f"   ⚠️  DSO Target PENDING (R² = {r2:.4f}, target 0.80)")

# Step 5: Export Model & Metadata
print("\n[5/5] Exporting model and metadata...")

# Export model using joblib (more robust than pickle)
model_path = os.path.join(os.path.dirname(__file__), "workload_predictor.pkl")
joblib.dump(model, model_path)
print(f"✅ Model exported: {model_path}")

# Export features
features_path = os.path.join(os.path.dirname(__file__), "features.json")
with open(features_path, 'w') as f:
    json.dump(X.columns.tolist(), f, indent=2)
print(f"✅ Features exported: {features_path}")

# Export metadata
metadata = {
    "model_type": "RandomForestRegressor",
    "r2_score": float(r2),
    "mae": float(mae),
    "rmse": float(rmse),
    "n_features": int(X.shape[1]),
    "features": X.columns.tolist(),
    "training_samples": int(X_train.shape[0]),
    "test_samples": int(X_test.shape[0]),
    "export_timestamp": pd.Timestamp.now().isoformat(),
    "numpy_version": np.__version__,
    "sklearn_version": __import__('sklearn').__version__
}

metadata_path = os.path.join(os.path.dirname(__file__), "model_metadata.json")
with open(metadata_path, 'w') as f:
    json.dump(metadata, f, indent=2)
print(f"✅ Metadata exported: {metadata_path}")

print("\n" + "="*70)
print("✅ Model Export Complete!")
print("="*70)
print("\n📋 Files created/updated:")
print(f"   • {model_path}")
print(f"   • {features_path}")
print(f"   • {metadata_path}")
print("\n🚀 Next: Restart Flask API")
print("   Run: python app.py")
print("="*70 + "\n")