#!/usr/bin/env python3
"""
COVID-19 Prediction Pipeline
---------------------------
This script runs the complete pipeline from data preprocessing to model training
and launches the Gradio UI.
"""

import os
import argparse
import subprocess
import time

def clear_screen():
    """Clear the terminal screen"""
    os.system('cls' if os.name == 'nt' else 'clear')

def run_command(command, description):
    """Run a system command and print output"""
    print(f"\n{'=' * 80}")
    print(f"STEP: {description}")
    print(f"{'=' * 80}\n")
    print(f"Running: {command}\n")
    
    # Run the command
    start_time = time.time()
    result = subprocess.run(command, shell=True)
    end_time = time.time()
    
    # Check if command was successful
    if result.returncode == 0:
        print(f"\nSuccess! Completed in {end_time - start_time:.2f} seconds")
    else:
        print(f"\nError! Command failed with exit code {result.returncode}")
        exit(1)
    
    print(f"\n{'=' * 80}\n")
    time.sleep(1)

def parse_args():
    """Parse command-line arguments"""
    parser = argparse.ArgumentParser(description="COVID-19 Prediction Pipeline")
    parser.add_argument("--skip-preprocessing", action="store_true", help="Skip data preprocessing")
    parser.add_argument("--skip-training", action="store_true", help="Skip model training")
    parser.add_argument("--only-ui", action="store_true", help="Only launch the Gradio UI")
    return parser.parse_args()

def main():
    """Run the complete pipeline"""
    args = parse_args()
    
    # Display welcome banner
    clear_screen()
    print("\n" + "=" * 80)
    print("COVID-19 PREDICTION PIPELINE".center(80))
    print("=" * 80 + "\n")
    print("This script will run the complete pipeline:")
    print("1. Data preprocessing")
    print("2. Model training")
    print("3. Launch Gradio UI for predictions")
    print("\nPress Ctrl+C at any time to stop the pipeline.")
    print()
    
    try:
        # Step 1: Data Preprocessing
        if args.only_ui:
            print("Skipping preprocessing and training, launching UI only...")
        else:
            if not args.skip_preprocessing:
                run_command("python preprocess_data.py", "Data Preprocessing")
            else:
                print("Skipping preprocessing as requested.")
            
            # Step 2: Model Training
            if not args.skip_training:
                run_command("python train_models.py", "Model Training")
            else:
                print("Skipping model training as requested.")
        
        # Step 3: Launch Gradio UI
        print("\nLaunching Gradio UI for predictions...")
        run_command("python gradio_app.py", "Gradio UI Launch")
    
    except KeyboardInterrupt:
        print("\n\nPipeline interrupted by user. Exiting.")
        exit(0)
    
    except Exception as e:
        print(f"\n\nError in pipeline: {str(e)}")
        exit(1)

if __name__ == "__main__":
    main()