train.py

import os
import argparse
import logging
from typing import Dict, List, Tuple, Any
from pathlib import Path

import numpy as np
import pandas as pd
import joblib
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC

from utils import (
    compute_metrics, get_classification_report,
    plot_confusion_matrix, plot_accuracy_comparison,
    plot_feature_importance, save_model, save_scaler,
    save_results_to_csv, save_best_model, get_best_model,
    ensure_directory
)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class Config:
    TEST_SIZE = 0.15
    RANDOM_STATE = 89
    RF_N_ESTIMATORS = 100
    RF_MAX_DEPTH = 10
    LR_MAX_ITER = 1000
    LR_C = 1.0
    SVM_KERNEL = 'rbf'
    SVM_C = 1.0
    SVM_GAMMA = 'scale'
    MODELS_DIR = './models'
    PLOTS_DIR = './plots'


class ModelTrainer:
    def __init__(self, model: Any, name: str):
        self.model = model
        self.name = name
        self.metrics = {}
        self.y_pred = None
        self.y_test = None

    def train(self, X_train: np.ndarray, y_train: np.ndarray) -> None:
        logger.info(f"  Training {self.name}...")
        self.model.fit(X_train, y_train)
        logger.info(f"  ✓ {self.name} training complete")

    def predict(self, X_test: np.ndarray, y_test: np.ndarray) -> Dict[str, float]:
        self.y_test = y_test
        self.y_pred = self.model.predict(X_test)
        self.metrics = compute_metrics(y_test, self.y_pred)
        return self.metrics

    def get_classification_report(self, target_names: List[str] = None) -> str:
        if target_names is None:
            target_names = ['Healthy', 'Parkinson']
        return get_classification_report(self.y_test, self.y_pred, target_names)

    def save(self, filepath: str) -> None:
        save_model(self.model, filepath)

    def get_feature_importances(self) -> np.ndarray:
        if hasattr(self.model, 'feature_importances_'):
            return self.model.feature_importances_
        return None


class TrainingPipeline:
    def __init__(self, data_path: str, output_dir: str = Config.MODELS_DIR,
                 plots_dir: str = Config.PLOTS_DIR):
        self.data_path = data_path
        self.output_dir = Path(output_dir)
        self.plots_dir = Path(plots_dir)

        ensure_directory(str(self.output_dir / 'dummy'))
        ensure_directory(str(self.plots_dir / 'dummy'))

        self.X = None
        self.y = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None

        self.scaler = None

        self.models: List[ModelTrainer] = []
        self.results: List[Dict[str, Any]] = []

        self.feature_names = None

    def load_data(self) -> None:
        logger.info("=" * 60)
        logger.info("LOADING DATA")
        logger.info("=" * 60)

        df = pd.read_csv(self.data_path)
        logger.info(f"  Loaded {len(df)} samples from {self.data_path}")

        exclude_cols = ['filename', 'filepath']
        feature_cols = [c for c in df.columns if c not in exclude_cols + ['label']]

        self.feature_names = feature_cols
        self.X = df[feature_cols].values
        self.y = df['label'].values

        logger.info(f"  Features: {self.X.shape[1]}")
        logger.info(f"  Healthy samples: {np.sum(self.y == 0)}")
        logger.info(f"  Parkinson samples: {np.sum(self.y == 1)}")

    def split_data(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("SPLITTING DATA")
        logger.info("=" * 60)

        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            self.X, self.y,
            test_size=Config.TEST_SIZE,
            random_state=Config.RANDOM_STATE,
            stratify=self.y
        )

        logger.info(f"  Training samples: {len(self.X_train)}")
        logger.info(f"  Test samples: {len(self.X_test)}")
        logger.info(f"  Train/Test split: {100-Config.TEST_SIZE*100:.0f}%/{Config.TEST_SIZE*100:.0f}%")

    def preprocess(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("PREPROCESSING (StandardScaler)")
        logger.info("=" * 60)

        self.scaler = StandardScaler()
        self.X_train = self.scaler.fit_transform(self.X_train)
        self.X_test = self.scaler.transform(self.X_test)

        logger.info("  ✓ Scaler fitted on training data")
        logger.info("  ✓ Training data transformed")
        logger.info("  ✓ Test data transformed")

    def create_models(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("INITIALIZING MODELS")
        logger.info("=" * 60)

        rf = RandomForestClassifier(
            n_estimators=Config.RF_N_ESTIMATORS,
            max_depth=Config.RF_MAX_DEPTH,
            random_state=Config.RANDOM_STATE,
            n_jobs=-1
        )
        self.models.append(ModelTrainer(rf, 'RandomForest'))

        lr = LogisticRegression(
            max_iter=Config.LR_MAX_ITER,
            C=Config.LR_C,
            random_state=Config.RANDOM_STATE,
            solver='lbfgs'
        )
        self.models.append(ModelTrainer(lr, 'LogisticRegression'))

        svm = SVC(
            kernel=Config.SVM_KERNEL,
            C=Config.SVM_C,
            gamma=Config.SVM_GAMMA,
            random_state=Config.RANDOM_STATE,
            probability=True
        )
        self.models.append(ModelTrainer(svm, 'SVM'))

        for model in self.models:
            logger.info(f"  ✓ {model.name} initialized")

    def train_models(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("TRAINING MODELS")
        logger.info("=" * 60)

        for model in self.models:
            model.train(self.X_train, self.y_train)

    def evaluate_models(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("EVALUATING MODELS")
        logger.info("=" * 60)

        for model in self.models:
            logger.info(f"\n  {model.name}:")
            metrics = model.predict(self.X_test, self.y_test)

            self.results.append({
                'model_name': model.name,
                'pca': False,
                'accuracy': metrics['accuracy'],
                'precision': metrics['precision'],
                'recall': metrics['recall'],
                'f1': metrics['f1']
            })

            logger.info(f"    Accuracy:  {metrics['accuracy']:.4f}")
            logger.info(f"    Precision: {metrics['precision']:.4f}")
            logger.info(f"    Recall:    {metrics['recall']:.4f}")
            logger.info(f"    F1 Score:  {metrics['f1']:.4f}")

            logger.info(f"\n    Classification Report:")
            for line in model.get_classification_report().split('\n'):
                logger.info(f"      {line}")

    def save_models(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("SAVING MODELS")
        logger.info("=" * 60)

        scaler_path = self.output_dir / 'scaler_no_pca.joblib'
        save_scaler(self.scaler, str(scaler_path))

        for model in self.models:
            model_path = self.output_dir / f'{model.name}_no_pca.joblib'
            model.save(str(model_path))

    def save_results(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("SAVING RESULTS")
        logger.info("=" * 60)

        results_path = self.output_dir.parent / 'model_results.csv'
        save_results_to_csv(self.results, str(results_path))

        best_name, best_acc = get_best_model(self.results, 'accuracy')
        best_model_path = self.output_dir.parent / 'best_model_no_pca.txt'
        save_best_model(best_name, str(best_model_path))

    def generate_plots(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("GENERATING PLOTS")
        logger.info("=" * 60)

        for model in self.models:
            cm_path = self.plots_dir / f'confusion_matrix_{model.name}_no_pca.png'
            plot_confusion_matrix(
                model.y_test, model.y_pred,
                f'{model.name} (No PCA)',
                str(cm_path)
            )

        accuracy_path = self.plots_dir / 'accuracy_comparison_no_pca.png'
        plot_accuracy_comparison(
            {r['model_name']: r for r in self.results},
            str(accuracy_path),
            pca_suffix="(No PCA)"
        )

        rf_model = self.models[0]
        importances = rf_model.get_feature_importances()
        if importances is not None:
            fi_path = self.plots_dir / 'feature_importance_randomforest_no_pca.png'
            plot_feature_importance(
                self.feature_names, importances,
                'Random Forest (No PCA)',
                str(fi_path),
                top_n=15
            )

    def run(self) -> None:
        logger.info("\n" + "=" * 60)
        logger.info("PARKINSON'S DISEASE DETECTION - MODEL TRAINING")
        logger.info("WITHOUT PCA")
        logger.info("=" * 60)

        self.load_data()
        self.split_data()
        self.preprocess()
        self.create_models()
        self.train_models()
        self.evaluate_models()
        self.save_models()
        self.save_results()
        self.generate_plots()

        logger.info("\n" + "=" * 60)
        logger.info("TRAINING COMPLETE!")
        logger.info("=" * 60)
        logger.info(f"  Models saved to: {self.output_dir}")
        logger.info(f"  Plots saved to: {self.plots_dir}")
        logger.info(f"  Results saved to: model_results.csv")
        logger.info("=" * 60)


def train(data_path: str, output_dir: str, plots_dir: str) -> None:
    pipeline = TrainingPipeline(data_path, output_dir, plots_dir)
    pipeline.run()


def main():
    parser = argparse.ArgumentParser(
        description='Train Parkinson\'s detection models (without PCA)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python train.py --data features.csv --output_dir ./models
    python train.py -d data/features.csv -o models/no_pca
        """
    )

    parser.add_argument(
        '--data', '-d',
        type=str,
        default='features.csv',
        help='Path to features CSV file (default: features.csv)'
    )

    parser.add_argument(
        '--output_dir', '-o',
        type=str,
        default='./models',
        help='Directory to save models (default: ./models)'
    )

    parser.add_argument(
        '--plots_dir', '-p',
        type=str,
        default='./plots',
        help='Directory to save plots (default: ./plots)'
    )

    args = parser.parse_args()

    if not os.path.exists(args.data):
        logger.error(f"Data file not found: {args.data}")
        logger.error("Run feature_extraction.py first to create features.csv")
        return

    try:
        train(args.data, args.output_dir, args.plots_dir)
        logger.info("\n✅ Training complete!")
    except Exception as e:
        logger.error(f"Training failed: {str(e)}")
        raise


if __name__ == '__main__':
    main()