"""
Train Hospital Readmission Prediction Model

Trains a Random Forest classifier on the Diabetes 130-US Hospitals dataset for 30-day readmission prediction.
Saves the trained model and feature names for use in examples.
"""

import pickle
import pandas as pd
import numpy as np
import os
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score


def load_diabetes_readmission_data():
    """Load UCI Diabetes 130-US Hospitals dataset from local file or UCI repository."""
    csv_path = "dataset_diabetes/diabetic_data.csv"

    # First, check if local dataset exists
    if os.path.exists(csv_path):
        try:
            print(f"Loading dataset from local file: {csv_path}")
            df = pd.read_csv(csv_path, na_values="?")
            print(f"✓ Loaded {len(df)} hospital encounters from local dataset")
            return df
        except Exception as e:
            print(f"Failed to load local dataset: {e}")

    # If local file doesn't exist, try downloading from UCI
    try:
        print("Local dataset not found. Downloading from UCI repository...")
        from ucimlrepo import fetch_ucirepo

        # Fetch dataset using official UCI client (dataset ID 296)
        dataset = fetch_ucirepo(id=296)

        # Combine features and targets into single dataframe
        df = pd.concat([dataset.data.features, dataset.data.targets], axis=1)

        print(f"✓ Loaded {len(df)} hospital encounters from UCI repository")
        return df

    except ImportError:
        print("ucimlrepo package not installed. Install with: pip install ucimlrepo")
        print("Using synthetic dataset instead...")
        return _generate_synthetic_data()

    except Exception as e:
        print(f"Could not fetch from UCI: {e}")
        print("Using synthetic dataset instead...")
        return _generate_synthetic_data()


def _generate_synthetic_data():
    """Generate synthetic diabetes readmission data for testing.

    Note: This is a fallback when the real dataset is unavailable.
    To use the real UCI Diabetes 130-US Hospitals dataset:
    1. Download from: https://archive.ics.uci.edu/dataset/296/diabetes+130-us+hospitals+for+years+1999-2008
    2. Extract: unzip diabetes*.zip -d dataset_diabetes/
    3. Re-run this script
    """
    print("\n⚠️  Using synthetic data - not representative of real clinical patterns")
    print("   For real dataset, download from UCI and extract to dataset_diabetes/\n")

    np.random.seed(42)
    n_samples = 10000
    df = pd.DataFrame(
        {
            "encounter_id": range(1, n_samples + 1),
            "patient_nbr": np.random.randint(1, 5000, n_samples),
            "age": np.random.choice(
                [
                    "[0-10)",
                    "[10-20)",
                    "[20-30)",
                    "[30-40)",
                    "[40-50)",
                    "[50-60)",
                    "[60-70)",
                    "[70-80)",
                    "[80-90)",
                    "[90-100)",
                ],
                n_samples,
            ),
            "time_in_hospital": np.random.randint(1, 14, n_samples),
            "num_lab_procedures": np.random.randint(1, 100, n_samples),
            "num_procedures": np.random.randint(0, 6, n_samples),
            "num_medications": np.random.randint(1, 80, n_samples),
            "number_diagnoses": np.random.randint(1, 16, n_samples),
            "readmitted": np.random.choice(["<30", ">30", "NO"], n_samples),
        }
    )
    return df


def train_model():
    """Train and save the hospital readmission prediction model."""
    print("=== Training Hospital Readmission Prediction Model ===\n")

    # Load and preprocess data
    df = load_diabetes_readmission_data()

    # Preprocessing: Convert target to binary (readmitted within 30 days)
    df["readmitted_30"] = (df["readmitted"] == "<30").astype(int)

    # Select key features for prediction
    feature_cols = [
        "time_in_hospital",
        "num_lab_procedures",
        "num_procedures",
        "num_medications",
        "number_diagnoses",
    ]

    # Handle age encoding (convert age ranges to numeric)
    if "age" in df.columns and df["age"].dtype == "object":
        age_mapping = {
            "[0-10)": 5,
            "[10-20)": 15,
            "[20-30)": 25,
            "[30-40)": 35,
            "[40-50)": 45,
            "[50-60)": 55,
            "[60-70)": 65,
            "[70-80)": 75,
            "[80-90)": 85,
            "[90-100)": 95,
        }
        df["age_numeric"] = df["age"].map(age_mapping)
        feature_cols.append("age_numeric")

    # Add encounter_id if not present
    if "encounter_id" not in df.columns:
        df["encounter_id"] = range(1, len(df) + 1)

    # Select subset of data for faster training
    df = df[feature_cols + ["readmitted_30", "encounter_id"]].dropna()
    # df = df.sample(n=min(10000, len(df)), random_state=42)

    print(f"Using {len(df)} encounters for training")
    print(
        f"Features: {', '.join([col for col in df.columns if col not in ['readmitted_30', 'encounter_id']])}"
    )
    print(f"Readmission rate (30-day): {df['readmitted_30'].mean():.1%}\n")

    # Split features and target
    X = df.drop(["readmitted_30", "encounter_id"], axis=1)
    y = df["readmitted_30"]
    encounter_ids = df["encounter_id"].values
    feature_names = X.columns.tolist()

    # Train/test split (stratified to maintain class balance)
    X_train, X_test, y_train, y_test, idx_train, idx_test = train_test_split(
        X, y, range(len(X)), test_size=0.2, random_state=42, stratify=y
    )

    # Scale features
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Train model
    print("Training Random Forest classifier...")
    model = RandomForestClassifier(
        n_estimators=100,
        max_depth=10,
        min_samples_split=5,
        min_samples_leaf=2,
        random_state=42,
        n_jobs=-1,
        class_weight="balanced",
    )
    model.fit(X_train_scaled, y_train)

    # Evaluate
    y_pred = model.predict(X_test_scaled)
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred, zero_division=0)
    recall = recall_score(y_test, y_pred, zero_division=0)
    f1 = f1_score(y_test, y_pred, zero_division=0)

    print("\n=== Model Performance ===")
    print(f"Accuracy:  {accuracy:.3f}")
    print(f"Precision: {precision:.3f}")
    print(f"Recall:    {recall:.3f}")
    print(f"F1 Score:  {f1:.3f}")

    # Feature importance
    feature_importance = pd.DataFrame(
        {"feature": feature_names, "importance": model.feature_importances_}
    ).sort_values("importance", ascending=False)

    print("\n=== Top 5 Features ===")
    for idx, row in feature_importance.head(5).iterrows():
        print(f"{row['feature']:22s}: {row['importance']:.3f}")

    # Save model artifacts
    print("\n=== Saving Model ===")
    artifacts = {
        "model": model,
        "scaler": scaler,
        "feature_names": feature_names,
        "metrics": {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1},
    }

    with open("readmission_model.pkl", "wb") as f:
        pickle.dump(artifacts, f)

    print("✓ Model saved to readmission_model.pkl")
    print(f"✓ Model size: {len(pickle.dumps(artifacts)) / 1024:.1f} KB")

    # Save training dataset
    print("\n=== Saving Training Dataset ===")

    # Create inference dataset with features and actual outcomes
    training_df = X_train.copy()
    training_df["actual_readmission"] = y_train.values
    training_df["encounter_id"] = encounter_ids[idx_train]

    # Save as CSV for easy upload to Vertex AI or other systems
    training_df.to_csv("training_dataset.csv", index=False)

    print("✓ Training dataset saved to training_dataset.csv")
    print(f"✓ Contains {len(training_df)} training records")
    print(f"✓ Columns: {', '.join(training_df.columns.tolist())}")

    # Save training dataset
    print("\n=== Saving Training Dataset ===")

    # Create inference dataset with features and actual outcomes
    inference_df = X_test.copy()
    inference_df["actual_readmission"] = y_test.values
    inference_df["encounter_id"] = encounter_ids[idx_test]

    # Save as CSV for easy upload to Vertex AI or other systems
    inference_df.to_csv("inference_dataset_with_actuals.csv", index=False)

    print("✓ Inference dataset saved to inference_dataset_with_actuals.csv")
    print(f"✓ Contains {len(inference_df)} test samples with actual outcomes")
    print(f"✓ Columns: {', '.join(inference_df.columns.tolist())}")

    # Save a summary of the datasets
    dataset_info = {
        "training_samples": len(X_train),
        "test_samples": len(X_test),
        "feature_names": feature_names,
        "readmission_rate_train": float(y_train.mean()),
        "readmission_rate_test": float(y_test.mean()),
        "model_metrics": {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1},
    }

    import json

    with open("dataset_info.json", "w") as f:
        json.dump(dataset_info, f, indent=2)

    print("✓ Dataset info saved to dataset_info.json")
    print("\nFiles ready for deployment:")
    print("  1. readmission_model.pkl - Trained model")
    print("  2. inference_dataset_with_actuals.csv - Test data with ground truth")
    print("  3. dataset_info.json - Dataset metadata")


if __name__ == "__main__":
    train_model()
