ml_new.py

"""
DIRECTION-AWARE TRAFFIC CONGESTION PREDICTION SYSTEM
Predicts congestion for North, South, East, West directions at four-way junctions
Using Enhanced CNN + Bidirectional LSTM with Attention
Version 1.0 - COMPLETE
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
import warnings
warnings.filterwarnings('ignore')

# Deep Learning imports
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import (
    Dense, LSTM, Conv1D, MaxPooling1D, Flatten, 
    Dropout, BatchNormalization, Input, Concatenate,
    Reshape, TimeDistributed, Bidirectional, GlobalAveragePooling1D,
    MultiHeadAttention
)
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, ModelCheckpoint
from tensorflow.keras.utils import to_categorical

# Traditional ML for comparison
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.utils.class_weight import compute_class_weight

class DirectionalTrafficPredictor:
    """Direction-aware Deep Learning system for traffic congestion prediction"""
    
    def __init__(self, csv_file):
        """Initialize with CSV file path"""
        self.csv_file = csv_file
        self.df = None
        self.models = {}  # Dictionary to store models for each direction
        self.scalers = {}
        self.encoders = {}
        self.features = None
        self.sequence_length = 20
        
        print("="*80)
        print("DIRECTIONAL TRAFFIC PREDICTION SYSTEM v1.0")
        print("Four-Way Junction: North, South, East, West")
        print("="*80)
        print(f"TensorFlow Version: {tf.__version__}")
        print(f"GPU Available: {len(tf.config.list_physical_devices('GPU')) > 0}")
        print(f"✓ Sequence Length: {self.sequence_length}")
        print("✓ Directions: North, South, East, West")
        print("✓ Predictions: LOW, MEDIUM, HIGH, CRITICAL")
    
    def load_data(self):
        """Step 1: Load CSV data with direction information"""
        print("\n[STEP 1] Loading CSV data...")
        
        self.df = pd.read_csv(self.csv_file)
        
        print(f"✓ Loaded dataset")
        print(f"  Rows: {len(self.df)}")
        print(f"  Columns: {len(self.df.columns)}")
        
        # Check if direction column exists, if not create from data
        if 'direction' not in self.df.columns:
            print("\n⚠ 'direction' column not found. Creating directional data...")
            # Create directional splits based on patterns in data
            self.df = self.create_directional_splits(self.df)
        
        # Remove duplicates
        duplicates = self.df.duplicated().sum()
        if duplicates > 0:
            print(f"\n⚠ WARNING: Found {duplicates} duplicate rows!")
            self.df = self.df.drop_duplicates()
            print(f"  ✓ Removed duplicates. New count: {len(self.df)}")
        
        # Show direction distribution
        print("\nDirection Distribution:")
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            count = len(self.df[self.df['direction'] == direction])
            print(f"  {direction:6s}: {count:5d} samples")
        
        return self.df
    
    def create_directional_splits(self, df):
        """Create directional data if not present in CSV"""
        print("  Creating directional assignments...")
        
        # Strategy: Split data into 4 equal parts or use modulo of index
        # This ensures each direction gets representative samples
        n = len(df)
        df = df.copy()
        
        # Assign directions cyclically or based on some feature
        directions = ['NORTH', 'SOUTH', 'EAST', 'WEST']
        df['direction'] = [directions[i % 4] for i in range(n)]
        
        print(f"  ✓ Assigned directions to {n} samples")
        return df
    
    def preprocess_data(self):
        """Step 2: Enhanced preprocessing with rolling features for each direction"""
        print("\n[STEP 2] Enhanced preprocessing with rolling features...")
        
        original_rows = len(self.df)
        print(f"Starting with {original_rows} rows")
        
        # Handle missing values
        missing = self.df.isnull().sum().sum()
        if missing > 0:
            print(f"⚠ Found {missing} missing values - filling with 0")
            self.df = self.df.fillna(0)
        
        print("\n🔥 Creating ENHANCED features...")
        
        # ============= ORIGINAL FEATURES =============
        self.df['traffic_density'] = self.df['total_vehicles'] / 1000
        
        self.df['heavy_vehicle_ratio'] = (
            (self.df['trucks'] + self.df['buses']) / 
            (self.df['total_vehicles'] + 1)
        )
        
        self.df['light_vehicle_ratio'] = (
            (self.df['cars'] + self.df['motorcycles']) / 
            (self.df['total_vehicles'] + 1)
        )
        
        self.df['has_emergency'] = (self.df['emergency_present'] > 0).astype(int)
        self.df['priority_vehicle_total'] = (
            self.df['emergency_vehicles'] + 
            self.df['official_vehicles']
        )
        
        self.df['flow_efficiency'] = (
            self.df['avg_speed_kmh'] / (self.df['avg_waiting_time'] + 1)
        )
        
        self.df['throughput'] = (
            self.df['total_vehicles'] / (self.df['halting_count'] + 1)
        )
        
        self.df['speed_variance'] = np.abs(
            self.df['avg_speed_kmh'] - self.df['avg_speed_kmh'].rolling(5, min_periods=1).mean()
        ).fillna(0)
        
        self.df['traffic_momentum'] = self.df['total_vehicles'].diff().fillna(0)
        
        # ============= ROLLING FEATURES =============
        print("  ✓ Adding rolling statistics (window=5)...")
        
        self.df['vehicles_rolling_mean_5'] = self.df['total_vehicles'].rolling(5, min_periods=1).mean()
        self.df['vehicles_rolling_std_5'] = self.df['total_vehicles'].rolling(5, min_periods=1).std().fillna(0)
        self.df['speed_rolling_mean_5'] = self.df['avg_speed_kmh'].rolling(5, min_periods=1).mean()
        self.df['speed_rolling_std_5'] = self.df['avg_speed_kmh'].rolling(5, min_periods=1).std().fillna(0)
        self.df['vehicles_ema'] = self.df['total_vehicles'].ewm(span=5).mean()
        self.df['speed_ema'] = self.df['avg_speed_kmh'].ewm(span=5).mean()
        self.df['traffic_acceleration'] = self.df['traffic_momentum'].diff().fillna(0)
        self.df['congestion_rate_change'] = self.df['congestion_level'].diff().fillna(0)
        self.df['congestion_rolling_mean_5'] = self.df['congestion_level'].rolling(5, min_periods=1).mean()
        self.df['waiting_time_variance'] = np.abs(
            self.df['avg_waiting_time'] - self.df['avg_waiting_time'].rolling(5, min_periods=1).mean()
        ).fillna(0)
        
        # Create congestion categories
        self.df['congestion_category'] = pd.cut(
            self.df['congestion_level'],
            bins=[-1, 30, 60, 80, 100],
            labels=['LOW', 'MEDIUM', 'HIGH', 'CRITICAL']
        )
        
        print(f"\n✓ Created {len(self.df.columns)} total features")
        
        # Show distribution by direction
        print("\nCongestion Distribution by Direction:")
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            dir_data = self.df[self.df['direction'] == direction]
            print(f"\n  {direction}:")
            dist = dir_data['congestion_category'].value_counts()
            for level in ['LOW', 'MEDIUM', 'HIGH', 'CRITICAL']:
                count = dist.get(level, 0)
                print(f"    {level:8s}: {count:4d}")
        
        return self.df
    
    def visualize_directional_data(self):
        """Step 3: Create visualizations for each direction"""
        print("\n[STEP 3] Creating directional visualizations...")
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        directions = ['NORTH', 'SOUTH', 'EAST', 'WEST']
        colors = ['green', 'yellow', 'orange', 'red']
        
        for idx, (ax, direction) in enumerate(zip(axes.flat, directions)):
            dir_data = self.df[self.df['direction'] == direction]
            congestion_counts = dir_data['congestion_category'].value_counts()
            
            ordered_labels = ['LOW', 'MEDIUM', 'HIGH', 'CRITICAL']
            ordered_counts = [congestion_counts.get(label, 0) for label in ordered_labels]
            
            bars = ax.bar(range(len(ordered_labels)), ordered_counts, 
                         color=colors, tick_label=ordered_labels)
            ax.set_title(f'{direction} Direction\n(Total: {sum(ordered_counts)} samples)', 
                        fontsize=12, fontweight='bold')
            ax.set_ylabel('Count')
            
            for bar, val in zip(bars, ordered_counts):
                if val > 0:
                    ax.text(bar.get_x() + bar.get_width()/2, val + 5, 
                           str(val), ha='center', fontweight='bold', fontsize=9)
        
        plt.suptitle('Congestion Distribution by Direction - Four-Way Junction', 
                    fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig('directional_congestion_analysis.png', dpi=300, bbox_inches='tight')
        print("✓ Saved: directional_congestion_analysis.png")
        plt.show()
        
        # Create a comprehensive comparison plot
        self.plot_directional_comparison()
    
    def plot_directional_comparison(self):
        """Create comparison visualization across all directions"""
        print("\n[PLOTTING] Directional comparison...")
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 10))
        
        # 1. Average congestion by direction
        ax1 = axes[0, 0]
        avg_congestion = self.df.groupby('direction')['congestion_level'].mean()
        bars = ax1.bar(avg_congestion.index, avg_congestion.values, 
                      color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728'])
        ax1.set_title('Average Congestion Level by Direction', fontsize=12, fontweight='bold')
        ax1.set_ylabel('Avg Congestion Level')
        for bar in bars:
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width()/2., height,
                    f'{height:.1f}', ha='center', va='bottom', fontweight='bold')
        
        # 2. Traffic volume by direction
        ax2 = axes[0, 1]
        avg_vehicles = self.df.groupby('direction')['total_vehicles'].mean()
        ax2.bar(avg_vehicles.index, avg_vehicles.values,
               color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728'])
        ax2.set_title('Average Vehicle Count by Direction', fontsize=12, fontweight='bold')
        ax2.set_ylabel('Avg Vehicles')
        
        # 3. Speed comparison
        ax3 = axes[1, 0]
        avg_speed = self.df.groupby('direction')['avg_speed_kmh'].mean()
        ax3.bar(avg_speed.index, avg_speed.values,
               color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728'])
        ax3.set_title('Average Speed by Direction', fontsize=12, fontweight='bold')
        ax3.set_ylabel('Avg Speed (km/h)')
        
        # 4. Waiting time comparison
        ax4 = axes[1, 1]
        avg_waiting = self.df.groupby('direction')['avg_waiting_time'].mean()
        ax4.bar(avg_waiting.index, avg_waiting.values,
               color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728'])
        ax4.set_title('Average Waiting Time by Direction', fontsize=12, fontweight='bold')
        ax4.set_ylabel('Avg Waiting Time (s)')
        
        plt.tight_layout()
        plt.savefig('directional_metrics_comparison.png', dpi=300, bbox_inches='tight')
        print("✓ Saved: directional_metrics_comparison.png")
        plt.show()
    
    def augment_sequences(self, X, y, noise_level=0.01):
        """Data augmentation"""
        X_augmented = []
        y_augmented = []
        
        for i in range(len(X)):
            X_augmented.append(X[i])
            y_augmented.append(y[i])
            
            noise = np.random.normal(0, noise_level, X[i].shape)
            X_augmented.append(X[i] + noise)
            y_augmented.append(y[i])
        
        return np.array(X_augmented), np.array(y_augmented)
    
    def create_sequences(self, X, y):
        """Create sequences for LSTM"""
        X_seq = []
        y_seq = []
        
        for i in range(len(X) - self.sequence_length):
            X_seq.append(X[i:i+self.sequence_length])
            y_seq.append(y[i+self.sequence_length])
        
        return np.array(X_seq), np.array(y_seq)
    
    def build_direction_model(self, input_shape, num_classes):
        """Build CNN-BiLSTM-Attention model for one direction"""
        inputs = Input(shape=input_shape)
        
        # CNN layers
        x = Conv1D(filters=64, kernel_size=3, activation='relu', padding='same')(inputs)
        x = BatchNormalization()(x)
        x = MaxPooling1D(pool_size=2)(x)
        x = Dropout(0.35)(x)
        
        x = Conv1D(filters=128, kernel_size=3, activation='relu', padding='same')(x)
        x = BatchNormalization()(x)
        x = MaxPooling1D(pool_size=2)(x)
        x = Dropout(0.45)(x)
        
        # Bidirectional LSTM
        lstm_out = Bidirectional(LSTM(128, return_sequences=True, recurrent_dropout=0.2))(x)
        lstm_out = Dropout(0.4)(lstm_out)
        
        # Multi-Head Attention
        attention_out = MultiHeadAttention(
            num_heads=4, 
            key_dim=128,
            dropout=0.3
        )(lstm_out, lstm_out)
        attention_out = Dropout(0.3)(attention_out)
        
        pooled = GlobalAveragePooling1D()(attention_out)
        
        # Dense layers
        x = Dense(128, activation='relu', kernel_regularizer=keras.regularizers.l2(0.001))(pooled)
        x = BatchNormalization()(x)
        x = Dropout(0.5)(x)
        
        x = Dense(64, activation='relu', kernel_regularizer=keras.regularizers.l2(0.001))(x)
        x = Dropout(0.4)(x)
        
        outputs = Dense(num_classes, activation='softmax')(x)
        
        model = Model(inputs=inputs, outputs=outputs)
        
        model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        return model
    
    def train_directional_models(self):
        """Step 4: Train separate model for each direction"""
        print("\n[STEP 4] Training directional models...")
        
        # Select features
        self.features = [
            'total_vehicles', 'cars', 'trucks', 'buses', 'motorcycles',
            'emergency_vehicles', 'official_vehicles', 'avg_speed_kmh',
            'avg_waiting_time', 'max_waiting_time', 'halting_count',
            'traffic_density', 'heavy_vehicle_ratio', 'light_vehicle_ratio',
            'flow_efficiency', 'throughput', 'speed_variance', 'traffic_momentum',
            'has_emergency', 'vehicles_rolling_mean_5', 'vehicles_rolling_std_5',
            'speed_rolling_mean_5', 'speed_rolling_std_5', 'vehicles_ema',
            'speed_ema', 'traffic_acceleration', 'congestion_rate_change',
            'congestion_rolling_mean_5', 'waiting_time_variance'
        ]
        
        directions = ['NORTH', 'SOUTH', 'EAST', 'WEST']
        
        for direction in directions:
            print(f"\n{'='*80}")
            print(f"Training model for {direction} direction")
            print(f"{'='*80}")
            
            # Filter data for this direction
            dir_data = self.df[self.df['direction'] == direction].copy()
            
            X = dir_data[self.features].values
            y = dir_data['congestion_category'].values
            
            # Remove NaN
            valid_idx = pd.notna(y)
            X = X[valid_idx]
            y = y[valid_idx]
            
            print(f"✓ Valid samples for {direction}: {len(X)}")
            
            # Encode labels
            encoder = LabelEncoder()
            y_encoded = encoder.fit_transform(y)
            y_categorical = to_categorical(y_encoded)
            
            # Compute class weights
            class_weights = compute_class_weight(
                'balanced',
                classes=np.unique(y_encoded),
                y=y_encoded
            )
            class_weight_dict = dict(enumerate(class_weights))
            
            # Scale features
            scaler = StandardScaler()
            X_scaled = scaler.fit_transform(X)
            
            # Create sequences
            X_seq, y_seq = self.create_sequences(X_scaled, y_categorical)
            
            # Split data
            X_train, X_test, y_train, y_test = train_test_split(
                X_seq, y_seq, test_size=0.2, random_state=42, 
                stratify=np.argmax(y_seq, axis=1)
            )
            
            # Augment training data
            X_train_aug, y_train_aug = self.augment_sequences(X_train, y_train)
            
            print(f"  Training samples: {len(X_train_aug)}")
            print(f"  Testing samples: {len(X_test)}")
            
            # Build model
            input_shape = (X_train_aug.shape[1], X_train_aug.shape[2])
            num_classes = y_train_aug.shape[1]
            
            model = self.build_direction_model(input_shape, num_classes)
            
            # Callbacks
            early_stop = EarlyStopping(
                monitor='val_loss', patience=15, restore_best_weights=True, verbose=0
            )
            reduce_lr = ReduceLROnPlateau(
                monitor='val_loss', factor=0.5, patience=5, min_lr=0.000001, verbose=0
            )
            
            # Train
            history = model.fit(
                X_train_aug, y_train_aug,
                validation_split=0.2,
                epochs=100,
                batch_size=32,
                class_weight=class_weight_dict,
                callbacks=[early_stop, reduce_lr],
                verbose=0
            )
            
            # Evaluate
            test_loss, test_accuracy = model.evaluate(X_test, y_test, verbose=0)
            print(f"\n✓ {direction} Model - Test Accuracy: {test_accuracy:.4f} ({test_accuracy*100:.2f}%)")
            
            # Predictions and classification report
            y_pred_probs = model.predict(X_test, verbose=0)
            y_pred = np.argmax(y_pred_probs, axis=1)
            y_test_labels = np.argmax(y_test, axis=1)
            
            print(f"\n{direction} Classification Report:")
            print(classification_report(y_test_labels, y_pred, 
                                       target_names=encoder.classes_, digits=3))
            
            # Store model and preprocessing objects
            self.models[direction] = model
            self.scalers[direction] = scaler
            self.encoders[direction] = encoder
        
        print(f"\n{'='*80}")
        print("✓ All directional models trained successfully!")
        print(f"{'='*80}")
    
    def save_models(self):
        """Step 5: Save all directional models"""
        print("\n[STEP 5] Saving directional models...")
        
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            # Save Keras model
            self.models[direction].save(f'model_{direction.lower()}.h5')
            self.models[direction].save(f'model_{direction.lower()}.keras')
            
            # Save preprocessing objects
            joblib.dump(self.scalers[direction], f'scaler_{direction.lower()}.pkl')
            joblib.dump(self.encoders[direction], f'encoder_{direction.lower()}.pkl')
        
        # Save common objects
        joblib.dump(self.features, 'features_directional.pkl')
        joblib.dump(self.sequence_length, 'sequence_length_directional.pkl')
        
        print("✓ Saved all directional models and preprocessors:")
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            print(f"  • model_{direction.lower()}.keras")
            print(f"  • scaler_{direction.lower()}.pkl")
            print(f"  • encoder_{direction.lower()}.pkl")
    
    def display_prediction_results(self, predictions):
        """Display prediction results in a clear format"""
        print(f"\n{'='*80}")
        print("CONGESTION PREDICTION RESULTS - FOUR-WAY JUNCTION")
        print(f"{'='*80}\n")
        
        # Create visualization
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
        
        # Bar chart of congestion levels
        directions = []
        congestion_values = []
        colors_map = {'LOW': 'green', 'MEDIUM': 'yellow', 'HIGH': 'orange', 'CRITICAL': 'red'}
        bar_colors = []
        
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            if direction in predictions:
                pred = predictions[direction]
                print(f"🧭 {direction} Direction:")
                print(f"   Predicted Congestion: {pred['most_common']}")
                print(f"   Confidence: {pred['avg_confidence']*100:.1f}%")
                print(f"   Total Predictions: {len(pred['predictions'])}\n")
                
                directions.append(direction)
                # Map congestion to numeric value for visualization
                cong_map = {'LOW': 1, 'MEDIUM': 2, 'HIGH': 3, 'CRITICAL': 4}
                congestion_values.append(cong_map.get(pred['most_common'], 0))
                bar_colors.append(colors_map.get(pred['most_common'], 'gray'))
        
        # Plot 1: Bar chart
        if directions:
            bars = ax1.bar(directions, congestion_values, color=bar_colors, edgecolor='black', linewidth=2)
            ax1.set_ylim(0, 5)
            ax1.set_yticks([1, 2, 3, 4])
            ax1.set_yticklabels(['LOW', 'MEDIUM', 'HIGH', 'CRITICAL'])
            ax1.set_title('Congestion Level by Direction', fontsize=14, fontweight='bold')
            ax1.set_xlabel('Direction', fontsize=12)
            ax1.set_ylabel('Congestion Level', fontsize=12)
            ax1.grid(axis='y', alpha=0.3)
            
            # Add value labels
            for bar, direction in zip(bars, directions):
                height = bar.get_height()
                pred = predictions[direction]
                ax1.text(bar.get_x() + bar.get_width()/2., height,
                        f"{pred['most_common']}\n{pred['avg_confidence']*100:.0f}%",
                        ha='center', va='bottom', fontweight='bold', fontsize=10)
        
        # Plot 2: Traffic light style visualization
        ax2.set_xlim(0, 4)
        ax2.set_ylim(0, 5)
        ax2.axis('off')
        ax2.set_title('Junction Traffic Status', fontsize=14, fontweight='bold')
        
        # Draw junction
        positions = {'NORTH': (2, 3.5), 'SOUTH': (2, 0.5), 'EAST': (3.5, 2), 'WEST': (0.5, 2)}
        
        for direction in ['NORTH', 'SOUTH', 'EAST', 'WEST']:
            if direction in predictions:
                x, y = positions[direction]
                pred = predictions[direction]
                color = colors_map.get(pred['most_common'], 'gray')
                
                # Draw circle
                circle = plt.Circle((x, y), 0.4, color=color, alpha=0.7, ec='black', linewidth=2)
                ax2.add_patch(circle)
                
                # Add text
                ax2.text(x, y, f"{direction[0]}\n{pred['most_common'][:3]}", 
                        ha='center', va='center', fontweight='bold', fontsize=10)
        
        # Draw center junction
        center_circle = plt.Circle((2, 2), 0.3, color='lightgray', alpha=0.5, ec='black', linewidth=2)
        ax2.add_patch(center_circle)
        ax2.text(2, 2, 'Junction', ha='center', va='center', fontsize=8, fontweight='bold')
        
        plt.tight_layout()
        plt.savefig('junction_prediction_results.png', dpi=300, bbox_inches='tight')
        print("\n✓ Saved visualization: junction_prediction_results.png")
        plt.show()
    
    def run_complete_pipeline(self):
        """Run entire pipeline"""
        self.load_data()
        self.preprocess_data()
        self.visualize_directional_data()
        self.train_directional_models()
        self.save_models()
        
        print("\n" + "="*80)
        print("✓ COMPLETE DIRECTIONAL PIPELINE FINISHED!")
        print("="*80)
        print(f"\n✓ Trained 4 separate models (N, S, E, W)")
        print(f"✓ Final dataset size: {len(self.df)} samples")
        print("\nGenerated Files:")
        print("  📊 directional_congestion_analysis.png")
        print("  📊 directional_metrics_comparison.png")
        print("  🤖 4 model files (north, south, east, west)")
        print("  🤖 4 scaler files")
        print("  🤖 4 encoder files")


# ============================================================================
# MAIN EXECUTION
# ============================================================================

if __name__ == "__main__":
    
    # Use your generated CSV file
    CSV_FILENAME = 'traffic_data_with_direction.csv'
    
    # Create predictor
    predictor = DirectionalTrafficPredictor(CSV_FILENAME)
    
    # Run complete pipeline
    predictor.run_complete_pipeline()
    
    print("\n" + "="*80)
    print("🚦 SYSTEM READY FOR PREDICTIONS!")
    print("="*80)
    print("\nTo make predictions on new data:")
    print("  predictor.predict_congestion('new_traffic_data.csv')")