example_usage.py

"""
Example Usage Script for ML Agent

This script demonstrates how to use the ML Agent and includes
sample dataset generators for testing.
"""

import argparse
import pandas as pd
import numpy as np
from pathlib import Path
from ml_agent import MLAgent


def generate_classification_dataset(n_samples=1000, n_features=10):
    """
    Generate a sample classification dataset

    Args:
        n_samples: Number of samples
        n_features: Number of features

    Returns:
        DataFrame with features and binary target
    """
    print(f"Generating classification dataset: {n_samples} samples, {n_features} features")

    np.random.seed(42)

    # Generate features
    data = {}

    # Numerical features
    for i in range(n_features // 2):
        data[f'num_feature_{i + 1}'] = np.random.randn(n_samples)

    # Categorical features
    categories = ['A', 'B', 'C', 'D']
    for i in range(n_features // 2):
        data[f'cat_feature_{i + 1}'] = np.random.choice(categories, n_samples)

    # Create target with some correlation to features
    # Target is based on first numerical feature + some noise
    target_prob = 1 / (1 + np.exp(-data['num_feature_1']))
    data['target'] = (np.random.random(n_samples) < target_prob).astype(int)

    df = pd.DataFrame(data)

    # Add some missing values (10% randomly)
    for col in df.columns[:-1]:  # Don't add missing to target
        mask = np.random.random(len(df)) < 0.1
        df.loc[mask, col] = np.nan

    return df


def generate_regression_dataset(n_samples=1000, n_features=10):
    """
    Generate a sample regression dataset

    Args:
        n_samples: Number of samples
        n_features: Number of features

    Returns:
        DataFrame with features and continuous target
    """
    print(f"Generating regression dataset: {n_samples} samples, {n_features} features")

    np.random.seed(42)

    # Generate features
    data = {}

    # Numerical features
    for i in range(n_features):
        data[f'feature_{i + 1}'] = np.random.randn(n_samples)

    # Create target as linear combination + noise
    target = sum(data[f'feature_{i + 1}'] * (i + 1) for i in range(3))
    target += np.random.randn(n_samples) * 2  # Add noise

    data['target'] = target

    df = pd.DataFrame(data)

    # Add some missing values (5% randomly)
    for col in df.columns[:-1]:
        mask = np.random.random(len(df)) < 0.05
        df.loc[mask, col] = np.nan

    return df


def generate_iris_dataset():
    """
    Generate the classic Iris dataset for classification

    Returns:
        DataFrame with iris features and species target
    """
    print("Generating Iris dataset (classic multiclass classification)")

    try:
        from sklearn.datasets import load_iris

        iris = load_iris()
        df = pd.DataFrame(iris.data, columns=iris.feature_names)
        df['species'] = iris.target

        return df
    except ImportError:
        print("⚠️  sklearn.datasets not available, using dummy data")
        return generate_classification_dataset(n_samples=150, n_features=4)


def generate_boston_dataset():
    """
    Generate a housing price regression dataset

    Returns:
        DataFrame with housing features and price target
    """
    print("Generating housing price dataset (regression)")

    np.random.seed(42)
    n_samples = 500

    data = {
        'rooms': np.random.randint(3, 10, n_samples),
        'age': np.random.randint(5, 100, n_samples),
        'distance_to_city': np.random.uniform(1, 20, n_samples),
        'crime_rate': np.random.uniform(0, 10, n_samples),
        'tax_rate': np.random.uniform(200, 700, n_samples),
        'student_teacher_ratio': np.random.uniform(12, 22, n_samples),
    }

    # Price based on features with some noise
    price = (
            data['rooms'] * 30000 +
            -data['age'] * 500 +
            -data['distance_to_city'] * 5000 +
            -data['crime_rate'] * 10000 +
            -data['tax_rate'] * 100 +
            -data['student_teacher_ratio'] * 2000 +
            200000  # Base price
    )

    # Add noise
    price += np.random.randn(n_samples) * 20000

    data['price'] = price

    return pd.DataFrame(data)


def run_example_classification():
    """Run example with classification dataset"""
    print("\n" + "=" * 60)
    print("EXAMPLE 1: Classification")
    print("=" * 60)

    # Generate data
    df = generate_classification_dataset(n_samples=1000, n_features=8)

    # Save to file
    data_dir = Path("sample_data")
    data_dir.mkdir(exist_ok=True)

    data_path = data_dir / "classification_example.csv"
    df.to_csv(data_path, index=False)
    print(f"✅ Data saved to: {data_path}")

    # Run agent
    agent = MLAgent(
        data_path=str(data_path),
        output_dir="outputs/classification_example",
        target_col="target",
        problem_type="classification"
    )

    agent.run()


def run_example_regression():
    """Run example with regression dataset"""
    print("\n" + "=" * 60)
    print("EXAMPLE 2: Regression")
    print("=" * 60)

    # Generate data
    df = generate_regression_dataset(n_samples=800, n_features=10)

    # Save to file
    data_dir = Path("sample_data")
    data_dir.mkdir(exist_ok=True)

    data_path = data_dir / "regression_example.csv"
    df.to_csv(data_path, index=False)
    print(f"✅ Data saved to: {data_path}")

    # Run agent
    agent = MLAgent(
        data_path=str(data_path),
        output_dir="outputs/regression_example",
        target_col="target",
        problem_type="regression"
    )

    agent.run()


def run_example_iris():
    """Run example with Iris dataset"""
    print("\n" + "=" * 60)
    print("EXAMPLE 3: Iris Dataset (Multiclass Classification)")
    print("=" * 60)

    # Generate data
    df = generate_iris_dataset()

    # Save to file
    data_dir = Path("sample_data")
    data_dir.mkdir(exist_ok=True)

    data_path = data_dir / "iris.csv"
    df.to_csv(data_path, index=False)
    print(f"✅ Data saved to: {data_path}")

    # Run agent (auto-detect everything)
    agent = MLAgent(
        data_path=str(data_path),
        output_dir="outputs/iris_example"
    )

    agent.run()


def run_example_housing():
    """Run example with housing price dataset"""
    print("\n" + "=" * 60)
    print("EXAMPLE 4: Housing Prices (Regression)")
    print("=" * 60)

    # Generate data
    df = generate_boston_dataset()

    # Save to file
    data_dir = Path("sample_data")
    data_dir.mkdir(exist_ok=True)

    data_path = data_dir / "housing.csv"
    df.to_csv(data_path, index=False)
    print(f"✅ Data saved to: {data_path}")

    # Run agent
    agent = MLAgent(
        data_path=str(data_path),
        output_dir="outputs/housing_example"
    )

    agent.run()


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description="ML Agent Examples - Generate and run sample datasets"
    )

    parser.add_argument(
        '--example',
        type=str,
        choices=['classification', 'regression', 'iris', 'housing', 'all'],
        default='iris',
        help='Which example to run'
    )

    parser.add_argument(
        '--generate-data',
        action='store_true',
        help='Only generate data, do not run agent'
    )

    args = parser.parse_args()

    # Create sample data directory
    Path("sample_data").mkdir(exist_ok=True)

    if args.generate_data:
        print("Generating sample datasets...")

        # Generate all datasets
        df_class = generate_classification_dataset()
        df_class.to_csv("sample_data/classification_example.csv", index=False)
        print("✅ classification_example.csv")

        df_reg = generate_regression_dataset()
        df_reg.to_csv("sample_data/regression_example.csv", index=False)
        print("✅ regression_example.csv")

        df_iris = generate_iris_dataset()
        df_iris.to_csv("sample_data/iris.csv", index=False)
        print("✅ iris.csv")

        df_housing = generate_boston_dataset()
        df_housing.to_csv("sample_data/housing.csv", index=False)
        print("✅ housing.csv")

        print("\n✅ All sample datasets generated in 'sample_data/' directory")
        print("\nRun examples with:")
        print("  python example_usage.py --example iris")
        print("  python example_usage.py --example classification")
        print("  python example_usage.py --example regression")
        print("  python example_usage.py --example housing")

        return

    # Run examples
    if args.example == 'classification':
        run_example_classification()
    elif args.example == 'regression':
        run_example_regression()
    elif args.example == 'iris':
        run_example_iris()
    elif args.example == 'housing':
        run_example_housing()
    elif args.example == 'all':
        run_example_classification()
        run_example_regression()
        run_example_iris()
        run_example_housing()

    print("\n" + "=" * 60)
    print("✅ ALL EXAMPLES COMPLETE")
    print("=" * 60)
    print("\nCheck the outputs/ directory for results!")


if __name__ == "__main__":
    main()