data_analyzer.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import sys
import os
import glob

class SensorDataAnalyzer:
    """Analyze and visualize 4-sensor ultrasonic detection data"""
    
    def __init__(self, csv_file):
        self.csv_file = csv_file
        self.df = None
        # Directory for saving results
        self.image_dir = 'images'
        # Match your CSV column names
        self.sensor_cols = ['sensor1', 'sensor2', 'sensor3', 'sensor4']
        
        self._setup_folders()
        self.load_data()
    
    def _setup_folders(self):
        """Ensure the images directory exists"""
        if not os.path.exists(self.image_dir):
            os.makedirs(self.image_dir)
            print(f"Created folder: {self.image_dir}")

    def _get_save_path(self, plot_type):
        """Generate a unique filename for the plot"""
        base_name = os.path.splitext(os.path.basename(self.csv_file))[0]
        timestamp = datetime.now().strftime('%H%M%S')
        filename = f"{plot_type}_{base_name}_{timestamp}.png"
        return os.path.join(self.image_dir, filename)

    def load_data(self):
        """Load and prepare CSV data"""
        try:
            self.df = pd.read_csv(self.csv_file)
            
            # Convert sensor columns to numeric
            for col in self.sensor_cols:
                if col in self.df.columns:
                    self.df[col] = pd.to_numeric(self.df[col], errors='coerce')
            
            print(f"\nLoaded {len(self.df)} records from {os.path.basename(self.csv_file)}")
            print()
        except Exception as e:
            print(f"Error loading file: {e}")
            sys.exit(1)
    
    # ─── SUMMARY STATISTICS ──────────────────────────────────────────────
    def show_summary(self):
        """Display summary statistics"""
        print("="*80)
        print(f"DATA SUMMARY - {os.path.basename(self.csv_file)}")
        print("="*80)
        
        print(f"\nTotal Samples: {len(self.df)}")
        sensor_names = ['(S1)', '(S2)', '(S3)', '(S4)']
        
        print("\nSensor Distance Statistics (cm):")
        for col, name in zip(self.sensor_cols, sensor_names):
            if col in self.df.columns:
                data = self.df[col].dropna()
                if len(data) > 0:
                    print(f"  {name:12s}: mean={data.mean():6.2f}, std={data.std():6.2f}, "
                          f"min={data.min():6.2f}, max={data.max():6.2f}")
        
        if 'object_type' in self.df.columns:
            print("\nObject Type Distribution:")
            print(self.df['object_type'].value_counts())
            
        print("\n" + "="*80 + "\n")

    # ─── GRAPH 1: TIME SERIES ────────────────────────────────────────────
    def plot_sensor_timeseries(self):
        fig, axes = plt.subplots(4, 1, figsize=(14, 10))
        fig.suptitle(f'Distances Over Time: {os.path.basename(self.csv_file)}', fontsize=16, fontweight='bold')
        colors = ['#2196F3', '#4CAF50', '#F44336', '#FF9800']
        
        for ax, col, color in zip(axes, self.sensor_cols, colors):
            if col in self.df.columns:
                data = pd.to_numeric(self.df[col], errors='coerce')
                ax.plot(data.index, data.values, color=color, linewidth=1, alpha=0.7)
                ax.fill_between(data.index, data.values, alpha=0.3, color=color)
                ax.set_ylabel('Distance (cm)')
                ax.set_title(col.capitalize(), fontweight='bold')
                ax.grid(True, alpha=0.3)
                if not data.dropna().empty:
                    ax.set_ylim(0, max(data.max() + 10, 200))

        plt.tight_layout()
        save_path = self._get_save_path("timeseries")
        plt.savefig(save_path)
        print(f"Graph saved to: {save_path}")
        plt.show()

    # ─── GRAPH 2: CORRELATION HEATMAP ────────────────────────────────────
    def plot_sensor_correlation(self):
        valid_cols = [c for c in self.sensor_cols if c in self.df.columns]
        if not valid_cols: return
            
        plt.figure(figsize=(10, 8))
        corr = self.df[valid_cols].apply(pd.to_numeric, errors='coerce').corr()
        sns.heatmap(corr, annot=True, fmt='.3f', cmap='coolwarm', center=0, square=True)
        plt.title('Sensor Correlation Matrix')
        
        plt.tight_layout()
        save_path = self._get_save_path("correlation")
        plt.savefig(save_path)
        print(f"Graph saved to: {save_path}")
        plt.show()

    # ─── GRAPH 3: FULL DASHBOARD ─────────────────────────────────────────
    def plot_full_dashboard(self):
        fig = plt.figure(figsize=(16, 12))
        fig.suptitle(f'Dashboard: {os.path.basename(self.csv_file)}', fontsize=18, fontweight='bold')
        gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
        
        # Main Overview
        ax_ov = fig.add_subplot(gs[0, :])
        for col in self.sensor_cols:
            if col in self.df.columns:
                ax_ov.plot(self.df[col], alpha=0.7, label=col)
        ax_ov.legend()
        ax_ov.set_title('All Sensors Overview')
        ax_ov.grid(True, alpha=0.3)
        
        # Object Type
        if 'object_type' in self.df.columns:
            ax = fig.add_subplot(gs[1, 0])
            obj_counts = self.df['object_type'].value_counts()
            ax.pie(obj_counts.values, labels=obj_counts.index, autopct='%1.1f%%')
            ax.set_title('Objects Detected')

        # Movement Type
        if 'movement_type' in self.df.columns:
            ax = fig.add_subplot(gs[1, 1])
            mov_counts = self.df['movement_type'].value_counts()
            ax.bar(mov_counts.index, mov_counts.values, color='steelblue')
            ax.set_title('Movement Distribution')
            plt.setp(ax.get_xticklabels(), rotation=45)

        # Confidence
        ax = fig.add_subplot(gs[1, 2])
        if 'object_conf' in self.df.columns:
            ax.plot(self.df['object_conf'], label='Object Conf', alpha=0.6)
        if 'movement_conf' in self.df.columns:
            ax.plot(self.df['movement_conf'], label='Movement Conf', alpha=0.6)
        ax.set_title('Detection Confidence')
        ax.set_ylim(0, 1)
        ax.legend()

        plt.tight_layout()
        save_path = self._get_save_path("dashboard")
        plt.savefig(save_path)
        print(f"Dashboard saved to: {save_path}")
        plt.show()

def get_file_selection():
    """Scan 'data' folder and let user choose a file"""
    data_dir = 'data'
    if not os.path.exists(data_dir):
        print(f"Folder '{data_dir}' not found. Looking in current directory...")
        search_path = "*.csv"
    else:
        search_path = os.path.join(data_dir, "*.csv")
        
    files = sorted(glob.glob(search_path))
    if not files:
        print(f"No CSV files found.")
        sys.exit(1)
        
    print("\n--- AVAILABLE CSV FILES ---")
    for i, f in enumerate(files):
        print(f"  {i + 1}. {os.path.basename(f)}")
    
    while True:
        try:
            choice = input("\nSelect file number (or 'q' to quit): ").strip().lower()
            if choice == 'q': sys.exit(0)
            idx = int(choice)
            if 1 <= idx <= len(files):
                return files[idx - 1]
            print("Invalid number.")
        except ValueError:
            print("Please enter a number.")

def show_menu():
    print("\n" + "="*70)
    print("  FILE ANALYSIS MENU (Auto-Save Active)")
    print("="*70)
    print("  1. Show Summary Statistics")
    print("  2. Plot Time Series (All Sensors)")
    print("  3. Plot Correlation Heatmap")
    print("  4. Plot Full Dashboard")
    print("  5. Select a Different File")
    print("  0. Exit")
    print("="*70)

def main():
    selected_file = get_file_selection()
    analyzer = SensorDataAnalyzer(selected_file)
    
    while True:
        show_menu()
        choice = input("\nSelect option: ").strip()
        
        if choice == '1': analyzer.show_summary()
        elif choice == '2': analyzer.plot_sensor_timeseries()
        elif choice == '3': analyzer.plot_sensor_correlation()
        elif choice == '4': analyzer.plot_full_dashboard()
        elif choice == '5':
            selected_file = get_file_selection()
            analyzer = SensorDataAnalyzer(selected_file)
        elif choice == '0': break
        else: print("Invalid option.")
        
        input("\nPress Enter to continue...")

if __name__ == "__main__":
    main()