cleveragents-core/simulations/sim10/loader.py

"""CSV file loading with automatic type detection and encoding handling."""

import pandas as pd
import chardet
import numpy as np
from typing import Optional, Dict, Any
import os


class CSVLoader:
    """Handles CSV file loading with automatic encoding detection and type inference."""
    
    def __init__(self):
        self.encoding = None
        self.delimiter = None
        
    def detect_encoding(self, file_path: str) -> str:
        """Detect the encoding of a CSV file."""
        try:
            with open(file_path, 'rb') as file:
                raw_data = file.read(10000)  # Read first 10KB
                result = chardet.detect(raw_data)
                encoding = result.get('encoding', 'utf-8')
                
                if encoding is None or result.get('confidence', 0) < 0.7:
                    encoding = 'utf-8'
                    
                self.encoding = encoding
                return encoding
        except Exception as e:
            print(f"Warning: Could not detect encoding, using utf-8. Error: {e}")
            self.encoding = 'utf-8'
            return 'utf-8'
    
    def detect_delimiter(self, file_path: str, encoding: str) -> str:
        """Detect the delimiter used in the CSV file."""
        delimiters = [',', ';', '\t', '|']
        delimiter_counts = {}
        
        try:
            with open(file_path, 'r', encoding=encoding) as file:
                first_line = file.readline()
                for delimiter in delimiters:
                    delimiter_counts[delimiter] = first_line.count(delimiter)
                
                # Choose delimiter with highest count
                best_delimiter = max(delimiter_counts, key=delimiter_counts.get)
                self.delimiter = best_delimiter if delimiter_counts[best_delimiter] > 0 else ','
                return self.delimiter
        except Exception as e:
            print(f"Warning: Could not detect delimiter, using comma. Error: {e}")
            self.delimiter = ','
            return ','
    
    def infer_column_types(self, df: pd.DataFrame) -> pd.DataFrame:
        """Automatically infer and convert column types."""
        df_copy = df.copy()
        
        for column in df_copy.columns:
            # Skip if already numeric
            if pd.api.types.is_numeric_dtype(df_copy[column]):
                continue
                
            # Try to convert to numeric first
            numeric_series = pd.to_numeric(df_copy[column], errors='coerce')
            if not numeric_series.isna().all():
                # If more than 80% of values can be converted to numeric, treat as numeric
                non_null_count = df_copy[column].count()
                numeric_count = numeric_series.count()
                if non_null_count > 0 and (numeric_count / non_null_count) >= 0.8:
                    df_copy[column] = numeric_series
                    continue
            
            # Try to convert to datetime
            try:
                datetime_series = pd.to_datetime(df_copy[column], errors='coerce', infer_datetime_format=True)
                if not datetime_series.isna().all():
                    non_null_count = df_copy[column].count()
                    datetime_count = datetime_series.count()
                    if non_null_count > 0 and (datetime_count / non_null_count) >= 0.8:
                        df_copy[column] = datetime_series
                        continue
            except:
                pass
            
            # Keep as categorical/string
            df_copy[column] = df_copy[column].astype('string')
        
        return df_copy
    
    def load_csv(self, file_path: str, columns: Optional[list] = None) -> pd.DataFrame:
        """Load CSV file with automatic encoding detection and type inference."""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        # Detect encoding and delimiter
        encoding = self.detect_encoding(file_path)
        delimiter = self.detect_delimiter(file_path, encoding)
        
        try:
            # Load the CSV
            df = pd.read_csv(
                file_path,
                encoding=encoding,
                delimiter=delimiter,
                low_memory=False
            )
            
            # Filter columns if specified
            if columns:
                available_columns = [col for col in columns if col in df.columns]
                missing_columns = [col for col in columns if col not in df.columns]
                
                if missing_columns:
                    print(f"Warning: Columns not found in file: {missing_columns}")
                
                if not available_columns:
                    raise ValueError("None of the specified columns were found in the file")
                
                df = df[available_columns]
            
            # Infer and convert column types
            df = self.infer_column_types(df)
            
            print(f"Successfully loaded CSV with {len(df)} rows and {len(df.columns)} columns")
            print(f"Encoding: {encoding}, Delimiter: '{delimiter}'")
            
            return df
            
        except pd.errors.EmptyDataError:
            raise ValueError("The CSV file is empty")
        except pd.errors.ParserError as e:
            raise ValueError(f"Error parsing CSV file: {e}")
        except Exception as e:
            raise ValueError(f"Error loading CSV file: {e}")
    
    def apply_filters(self, df: pd.DataFrame, filter_conditions: list) -> pd.DataFrame:
        """Apply filtering conditions to the DataFrame."""
        if not filter_conditions:
            return df
        
        filtered_df = df.copy()
        
        for condition in filter_conditions:
            try:
                # Parse condition (format: column_name operator value)
                parts = condition.split(' ', 2)
                if len(parts) != 3:
                    print(f"Warning: Invalid filter format '{condition}'. Expected: 'column operator value'")
                    continue
                
                column, operator, value = parts
                
                if column not in filtered_df.columns:
                    print(f"Warning: Column '{column}' not found. Skipping filter.")
                    continue
                
                # Convert value to appropriate type
                if pd.api.types.is_numeric_dtype(filtered_df[column]):
                    try:
                        value = float(value)
                    except ValueError:
                        print(f"Warning: Cannot convert '{value}' to numeric for column '{column}'")
                        continue
                
                # Apply filter based on operator
                if operator == '==':
                    mask = filtered_df[column] == value
                elif operator == '!=':
                    mask = filtered_df[column] != value
                elif operator == '>':
                    mask = filtered_df[column] > value
                elif operator == '>=':
                    mask = filtered_df[column] >= value
                elif operator == '<':
                    mask = filtered_df[column] < value
                elif operator == '<=':
                    mask = filtered_df[column] <= value
                elif operator == 'contains':
                    mask = filtered_df[column].astype(str).str.contains(str(value), na=False)
                else:
                    print(f"Warning: Unsupported operator '{operator}'. Supported: ==, !=, >, >=, <, <=, contains")
                    continue
                
                filtered_df = filtered_df[mask]
                
            except Exception as e:
                print(f"Warning: Error applying filter '{condition}': {e}")
                continue
        
        print(f"Applied {len(filter_conditions)} filters. Rows remaining: {len(filtered_df)}")
        return filtered_df