utils.py

"""
================================================================================
Utility Functions for Traffic Accident Analysis
================================================================================
Module: utils.py
Purpose: Reusable functions for data loading, cleaning, analysis
Tech Stack: PySpark, Pandas
Author: Data Analytics Team
================================================================================
"""

import os
import sys
import warnings
warnings.filterwarnings('ignore')

# Import PySpark libraries
try:
    from pyspark.sql import SparkSession
    from pyspark.sql.functions import (
        col, to_date, hour, dayofweek, month, year, 
        count, max as spark_max, min as spark_min, 
        round as spark_round, when, coalesce, upper, trim,
        to_timestamp
    )
except ImportError as e:
    print(f"Error: PySpark not installed. Install with: pip install pyspark")
    sys.exit(1)

# Import data manipulation and visualization libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns


# ================================================================================
# CONFIGURATION
# ================================================================================

class Config:
    """Configuration class for utilities"""
    
    # Spark Configuration
    SPARK_MEMORY = "4g"
    SPARK_CORES = "4"
    
    # Visualization Configuration
    FIG_SIZE_LARGE = (14, 8)
    FIG_SIZE_MEDIUM = (12, 6)
    FIG_SIZE_SMALL = (10, 6)
    DPI = 300
    
    # Data Validation
    REQUIRED_COLUMNS = [
        'id', 'date', 'time', 'city', 'weather', 'visibility', 
        'road_type', 'vehicles_involved', 'casualties', 'severity', 
        'latitude', 'longitude'
    ]


# ================================================================================
# SPARK SESSION INITIALIZATION
# ================================================================================

def initialize_spark():
    """
    Initialize and configure Spark Session with optimizations for Big Data
    
    Returns:
        SparkSession: Configured Spark session
    """
    try:
        spark = SparkSession.builder \
            .appName("TrafficAccidentAnalysis") \
            .config("spark.driver.memory", Config.SPARK_MEMORY) \
            .config("spark.executor.memory", Config.SPARK_MEMORY) \
            .config("spark.sql.shuffle.partitions", "200") \
            .config("spark.default.parallelism", "200") \
            .getOrCreate()
        
        # Set log level to reduce verbose output
        spark.sparkContext.setLogLevel("WARN")
        
        return spark
    
    except Exception as e:
        print(f"Error initializing Spark: {str(e)}")
        sys.exit(1)


# ================================================================================
# DATA LOADING
# ================================================================================

def load_data(spark, csv_path):
    """
    Load large CSV dataset into Spark DataFrame with error handling
    
    Args:
        spark (SparkSession): Spark session object
        csv_path (str): Path to CSV file
    
    Returns:
        DataFrame: Spark DataFrame with raw data
    """
    try:
        if not os.path.exists(csv_path):
            raise FileNotFoundError(f"CSV file not found: {csv_path}")
        
        # Load CSV with automatic schema inference
        df = spark.read \
            .option("header", "true") \
            .option("inferSchema", "true") \
            .option("mode", "PERMISSIVE") \
            .csv(csv_path)
        
        # Cache the dataframe for repeated operations
        df.cache()
        
        return df
    
    except FileNotFoundError as e:
        print(f"File Error: {str(e)}")
        sys.exit(1)
    except Exception as e:
        print(f"Error loading data: {str(e)}")
        sys.exit(1)


# ================================================================================
# DATA CLEANING
# ================================================================================

def clean_data(df):
    """
    Clean and preprocess data: handle missing values, format conversions,
    remove duplicates, normalize city names
    
    Args:
        df (DataFrame): Raw Spark DataFrame
    
    Returns:
        DataFrame: Cleaned Spark DataFrame
    """
    try:
        # Handle missing values
        df = df.fillna({
            'casualties': 0,
            'vehicles_involved': 1,
            'visibility': 'Unknown',
            'weather': 'Clear',
            'road_type': 'Unknown',
            'latitude': 0.0,
            'longitude': 0.0
        })
        
        # Normalize city names
        df = df.withColumn('city', upper(trim(col('city'))))
        
        # Convert date and time to proper formats
        df = df.withColumn('date', to_date(col('date'), 'yyyy-MM-dd')) \
                .withColumn('hour', hour(to_timestamp(col('time'), 'HH:mm:ss')))
        
        # Extract year and month for trend analysis
        df = df.withColumn('year', year(col('date'))) \
                .withColumn('month', month(col('date')))
        
        # Remove duplicates
        df = df.dropDuplicates(subset=['id'])
        
        # Cache cleaned data
        df.cache()
        
        return df
    
    except Exception as e:
        print(f"Error during data cleaning: {str(e)}")
        sys.exit(1)


# ================================================================================
# DATA ANALYSIS
# ================================================================================

def analyze_data(df):
    """
    Perform comprehensive Big Data analysis on cleaned dataset
    
    Args:
        df (DataFrame): Cleaned Spark DataFrame
    
    Returns:
        dict: Dictionary containing analysis results
    """
    try:
        analysis_results = {}
        
        # Analysis 1: Accidents per city
        accidents_per_city = df.groupBy('city') \
                              .agg(count('*').alias('accident_count')) \
                              .orderBy(col('accident_count').desc()) \
                              .limit(15)
        analysis_results['accidents_per_city'] = accidents_per_city
        
        # Analysis 2: Accidents by hour of day
        accidents_by_hour = df.groupBy('hour') \
                             .agg(count('*').alias('accident_count')) \
                             .orderBy('hour')
        analysis_results['accidents_by_hour'] = accidents_by_hour
        
        # Analysis 3: Accidents by weather condition
        accidents_by_weather = df.groupBy('weather') \
                                .agg(count('*').alias('accident_count')) \
                                .orderBy(col('accident_count').desc())
        analysis_results['accidents_by_weather'] = accidents_by_weather
        
        # Analysis 4: Severity distribution
        severity_dist = df.groupBy('severity') \
                         .agg(count('*').alias('count')) \
                         .orderBy(col('count').desc())
        analysis_results['severity_distribution'] = severity_dist
        
        # Analysis 5: Monthly and yearly trends
        monthly_trends = df.groupBy('year', 'month') \
                          .agg(count('*').alias('accident_count')) \
                          .orderBy('year', 'month')
        analysis_results['monthly_trends'] = monthly_trends
        
        # Analysis 6: Top 10 accident-prone locations (lat/long clusters)
        top_locations = df.groupBy('latitude', 'longitude') \
                         .agg(count('*').alias('accident_count')) \
                         .orderBy(col('accident_count').desc()) \
                         .limit(10)
        analysis_results['top_locations'] = top_locations
        
        # Additional Analysis: Vehicles and casualties statistics
        statistics = df.agg({
            'vehicles_involved': 'avg',
            'casualties': 'avg'
        }).collect()[0].asDict()
        analysis_results['statistics'] = statistics
        
        return analysis_results
    
    except Exception as e:
        print(f"Error during analysis: {str(e)}")
        sys.exit(1)


# ================================================================================
# DATA CONVERSION FOR VISUALIZATION
# ================================================================================

def convert_to_pandas(spark_df):
    """
    Convert Spark DataFrame to Pandas for visualization
    
    Args:
        spark_df (DataFrame): Spark DataFrame
    
    Returns:
        DataFrame: Pandas DataFrame
    """
    try:
        pandas_df = spark_df.toPandas()
        return pandas_df
    except Exception as e:
        print(f"Error converting to Pandas: {str(e)}")
        return None


# ================================================================================
# GRAPH SAVING
# ================================================================================

def save_graph(filename, output_dir, figsize=Config.FIG_SIZE_MEDIUM):
    """
    Save matplotlib figure as PNG
    
    Args:
        filename (str): Name of output file (without extension)
        output_dir (str): Output directory path
        figsize (tuple): Figure size
    """
    try:
        os.makedirs(output_dir, exist_ok=True)
        filepath = os.path.join(output_dir, f"{filename}.png")
        plt.tight_layout()
        plt.savefig(filepath, dpi=Config.DPI, bbox_inches='tight')
        return filepath
    except Exception as e:
        print(f"Error saving graph: {str(e)}")
        return None


# ================================================================================
# DATA STATISTICS
# ================================================================================

def get_data_stats(df):
    """
    Get basic statistics about the dataset
    
    Args:
        df (DataFrame): Spark DataFrame
    
    Returns:
        dict: Statistics dictionary
    """
    try:
        stats = {
            'total_rows': df.count(),
            'total_columns': len(df.columns),
            'column_names': df.columns,
            'schema': df.schema
        }
        return stats
    except Exception as e:
        print(f"Error getting statistics: {str(e)}")
        return None


# ================================================================================
# DATA VALIDATION
# ================================================================================

def validate_data(df):
    """
    Validate dataset structure and required columns
    
    Args:
        df (DataFrame): Spark DataFrame
    
    Returns:
        tuple: (is_valid, error_messages)
    """
    try:
        errors = []
        
        # Check required columns
        missing_columns = set(Config.REQUIRED_COLUMNS) - set(df.columns)
        if missing_columns:
            errors.append(f"Missing columns: {', '.join(missing_columns)}")
        
        # Check row count
        if df.count() == 0:
            errors.append("DataFrame is empty")
        
        is_valid = len(errors) == 0
        return is_valid, errors
    
    except Exception as e:
        return False, [f"Validation error: {str(e)}"]