data_loader.py

# app/data_loader.py

import os
import sys
import glob
import polars as pl
import pandas as pd
import numpy as np
import streamlit as st

# ---------------------------------------------------------
# DIRECTORY SCANNING & METADATA LOADER
# ---------------------------------------------------------
def get_preloaded_datasets() -> dict:
    """Scans CLI arguments and local project folders for supported dataset files."""
    datasets = {}
    
    # 1. Parse command line arguments after '--'
    if len(sys.argv) > 1:
        for arg in sys.argv[1:]:
            clean_arg = arg.strip('\'"')
            if os.path.isfile(clean_arg) and clean_arg.lower().endswith(('.csv', '.xlsx', '.xls', '.tsv', '.parquet', '.ipc', '.json')):
                name = os.path.basename(clean_arg)
                datasets[f"CLI: {name}"] = clean_arg
                
    # 2. Scan specific search directories
    search_dirs = ['.', 'datasets']
    exts = ['*.csv', '*.xlsx', '*.xls', '*.tsv', '*.parquet', '*.ipc', '*.json']
    
    for d in search_dirs:
        if os.path.exists(d):
            for ext in exts:
                found_files = glob.glob(os.path.join(d, ext))
                for f in found_files:
                    name = os.path.basename(f)
                    key = f if d == 'datasets' else name
                    if key not in datasets:  # Prevent duplicate keys
                        datasets[key] = f
                        
    return datasets

def scan_file_columns(filepath: str) -> list[str]:
    """Scans dataset file header lazily using Polars to retrieve column names instantly without loading full dataset into RAM."""
    ext = filepath.split('.')[-1].lower()
    if ext == 'csv':
        return pl.scan_csv(filepath).collect_schema().names()
    elif ext in ['ipc', 'arrow', 'feather']:
        return pl.scan_ipc(filepath).collect_schema().names()
    elif ext == 'parquet':
        return pl.scan_parquet(filepath).collect_schema().names()
    elif ext in ['xlsx', 'xls']:
        df = pl.read_excel(filepath, read_options={"n_rows": 1})
        return list(df.columns)
    elif ext == 'tsv':
        return pl.scan_csv(filepath, separator='\t').collect_schema().names()
    else:
        df = pl.read_csv(filepath, n_rows=1)
        return list(df.columns)

def load_dataset_from_path(file_obj, columns: list[str] = None) -> pl.DataFrame:
    """Loads a dataset from file path or file-like object into a Polars DataFrame.
    If `columns` is specified, uses Polars lazy column projection to load ONLY those columns from disk.
    """
    if isinstance(file_obj, str):
        filepath = file_obj
        ext = filepath.split('.')[-1].lower()
        
        if ext == 'csv':
            lf = pl.scan_csv(filepath)
            if columns:
                avail = lf.collect_schema().names()
                valid = [c for c in columns if c in avail]
                if valid:
                    lf = lf.select(valid)
            return lf.collect()
        elif ext == 'parquet':
            lf = pl.scan_parquet(filepath)
            if columns:
                avail = lf.collect_schema().names()
                valid = [c for c in columns if c in avail]
                if valid:
                    lf = lf.select(valid)
            return lf.collect()
        elif ext in ['ipc', 'arrow', 'feather']:
            lf = pl.scan_ipc(filepath)
            if columns:
                avail = lf.collect_schema().names()
                valid = [c for c in columns if c in avail]
                if valid:
                    lf = lf.select(valid)
            return lf.collect()
        elif ext in ['xlsx', 'xls']:
            df = pl.read_excel(filepath)
            if columns:
                valid = [c for c in columns if c in df.columns]
                if valid:
                    df = df.select(valid)
            return df
        elif ext == 'tsv':
            lf = pl.scan_csv(filepath, separator='\t')
            if columns:
                avail = lf.collect_schema().names()
                valid = [c for c in columns if c in avail]
                if valid:
                    lf = lf.select(valid)
            return lf.collect()
        else:
            df = pl.read_csv(filepath)
            if columns:
                valid = [c for c in columns if c in df.columns]
                if valid:
                    df = df.select(valid)
            return df
    else:
        # File-like object (UploadedFile)
        content = file_obj.read()
        df = pl.read_csv(content)
        if columns:
            valid = [c for c in columns if c in df.columns]
            if valid:
                df = df.select(valid)
        return df

# ---------------------------------------------------------
# METADATA & LABEL DETECTION UTILITIES
# ---------------------------------------------------------
def get_label_type(df: pl.DataFrame | pd.DataFrame, label_col: str) -> tuple[str, str]:
    """Detects if the target label is continuous (numerical) or categorical using Polars or Pandas."""
    if label_col not in df.columns:
        return "categorical", "Unknown"

    if isinstance(df, pl.DataFrame):
        is_label_numeric = df[label_col].dtype.is_numeric()
        unique_labels = df[label_col].n_unique()
    else:
        is_label_numeric = pd.api.types.is_numeric_dtype(df[label_col])
        unique_labels = df[label_col].nunique()
        
    if is_label_numeric and unique_labels > 10:
        return "numeric", "Continuous Numerical Data (Regression)"
    else:
        return "categorical", f"Categorical / Discrete ({unique_labels} groups)"

# ---------------------------------------------------------
# SYNTHETIC DATA GENERATORS (CACHED FOR APP SPEED)
# ---------------------------------------------------------
@st.cache_data
def generate_synthetic_classification() -> pl.DataFrame:
    """Generates synthetic Polars dataset for Customer Segmentation (Classification task)."""
    np.random.seed(42)
    n_samples = 300
    groups = ['VIP Customer', 'Saver Customer', 'Frequent Shopper', 'Regular Customer']
    group_choices = np.random.choice(groups, n_samples, p=[0.15, 0.25, 0.30, 0.30])
    
    data = []
    for g in group_choices:
        if g == 'VIP Customer':
            age = np.random.normal(45, 8)
            income = np.random.normal(120, 15)
            spending = np.random.normal(85, 10)
            satisfaction = np.random.normal(4.5, 0.3)
        elif g == 'Saver Customer':
            age = np.random.normal(38, 10)
            income = np.random.normal(40, 8)
            spending = np.random.normal(20, 8)
            satisfaction = np.random.normal(3.8, 0.5)
        elif g == 'Frequent Shopper':
            age = np.random.normal(28, 6)
            income = np.random.normal(75, 12)
            spending = np.random.normal(78, 9)
            satisfaction = np.random.normal(4.1, 0.4)
        else:
            age = np.random.normal(35, 11)
            income = np.random.normal(55, 10)
            spending = np.random.normal(48, 11)
            satisfaction = np.random.normal(3.5, 0.6)
            
        data.append([round(age), round(income, 1), round(spending), round(satisfaction, 1), g])
        
    df_pd = pd.DataFrame(data, columns=['Age', 'Income (kUSD)', 'Spending Score', 'Satisfaction Rate', 'Customer Group'])
    df_pd['Age'] = df_pd['Age'].clip(18, 80)
    df_pd['Spending Score'] = df_pd['Spending Score'].clip(1, 100)
    df_pd['Satisfaction Rate'] = df_pd['Satisfaction Rate'].clip(1.0, 5.0)
    return pl.from_pandas(df_pd)

@st.cache_data
def generate_synthetic_regression() -> pl.DataFrame:
    """Generates synthetic Polars dataset for House Pricing (Regression task)."""
    np.random.seed(100)
    n_samples = 400
    area = np.random.uniform(30, 250, n_samples)
    bedrooms = np.random.choice([1, 2, 3, 4, 5], n_samples, p=[0.1, 0.4, 0.3, 0.15, 0.05])
    distance = np.random.uniform(0.5, 20.0, n_samples)
    year_built = np.random.choice(range(1995, 2026), n_samples)
    
    price = (area * 0.08) + (bedrooms * 0.5) - (distance * 0.25) + ((year_built - 1995) * 0.03) + np.random.normal(0, 1.5, n_samples)
    price = np.clip(price, 0.5, None)
    
    df_pd = pd.DataFrame({
        'Area (m2)': np.round(area, 1),
        'Bedrooms': bedrooms,
        'Distance to City Center (km)': np.round(distance, 2),
        'Year Built': year_built,
        'House Price (kUSD)': np.round(price, 2)
    })
    return pl.from_pandas(df_pd)

@st.cache_data
def generate_synthetic_diabetes() -> pl.DataFrame:
    """Generates synthetic medical Polars dataset for Diabetes Diagnosis."""
    np.random.seed(999)
    n_samples = 350
    bmi = np.random.normal(26, 5, n_samples)
    blood_pressure = np.random.normal(120, 15, n_samples)
    glucose = np.random.normal(100, 25, n_samples)
    age = np.random.randint(20, 80, n_samples)
    
    z = (bmi - 25)*0.1 + (blood_pressure - 120)*0.03 + (glucose - 100)*0.05 + (age - 40)*0.02 - 1.5
    prob = 1 / (1 + np.exp(-z))
    risk = np.random.binomial(1, prob, n_samples)
    risk_labels = np.where(risk == 1, 'High Risk', 'Normal Risk')
    
    df_pd = pd.DataFrame({
        'BMI Index': np.round(bmi, 1),
        'Blood Pressure': np.round(blood_pressure),
        'Glucose Level': np.round(glucose),
        'Age': age,
        'Diagnosis Result': risk_labels
    })
    return pl.from_pandas(df_pd)

# ---------------------------------------------------------
# POLARS LAZY & SELECTIVE COLUMN LOADER
# ---------------------------------------------------------
@st.cache_data(show_spinner="Scanning dataset metadata & column schema...")
def get_dataset_schema_from_config(config_path="config.json") -> tuple:
    """Scans dataset column metadata lazily via Polars without reading data rows into RAM."""
    if not os.path.exists(config_path):
        return [], [], "", "Configuration file 'config.json' not found at project root."
        
    try:
        import json
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except Exception as e:
        return [], [], "", f"Error reading JSON configuration file: {e}"
        
    file_mode = config.get("file_mode", "1_file")
    
    if file_mode == "1_file":
        filepath = config.get("data_file")
        if not filepath or not os.path.exists(filepath):
            return [], [], "", f"Data file '{filepath}' specified in JSON does not exist."
        cols = scan_file_columns(filepath)
        return cols, cols, os.path.basename(filepath), None

    elif file_mode == "2_files":
        feat_path = config.get("features_file")
        label_path = config.get("labels_file")
        
        if not feat_path or not label_path:
            return [], [], "", "JSON configuration missing 'features_file' or 'labels_file' parameter."
        if not os.path.exists(feat_path) or not os.path.exists(label_path):
            return [], [], "", f"File '{feat_path}' or '{label_path}' not found."
            
        feat_cols_all = scan_file_columns(feat_path)
        label_cols_all = scan_file_columns(label_path)
        
        if not feat_cols_all or not label_cols_all:
            return [], [], "", "One of the dataset files has no columns."
            
        feat_key_col = feat_cols_all[0]
        label_key_col = label_cols_all[0]
        
        features_cols = [c for c in feat_cols_all if c != feat_key_col]
        labels_cols = [c for c in label_cols_all if c != label_key_col]
        
        dataset_name = f"Features ({os.path.basename(feat_path)}) + Labels ({os.path.basename(label_path)}) joined on '{feat_key_col}'"
        return features_cols, labels_cols, dataset_name, None

    else:
        return [], [], "", f"Invalid 'file_mode' value: '{file_mode}'."

@st.cache_data(show_spinner="Extracting selected columns with Polars...")
def load_selected_columns_from_config(config_path: str = "config.json", selected_cols: list[str] = None) -> tuple:
    """Uses Polars Lazy Evaluation to read ONLY the selected columns from disk on form submit."""
    if not os.path.exists(config_path):
        return None, "", "Configuration file not found."
        
    try:
        import json
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except Exception as e:
        return None, "", f"JSON Error: {e}"
        
    file_mode = config.get("file_mode", "1_file")
    
    if file_mode == "1_file":
        filepath = config.get("data_file")
        if not filepath or not os.path.exists(filepath):
            return None, "", f"Data file '{filepath}' does not exist."
            
        df = load_dataset_from_path(filepath, columns=selected_cols)
        return df, os.path.basename(filepath), None

    elif file_mode == "2_files":
        feat_path = config.get("features_file")
        label_path = config.get("labels_file")
        
        if not feat_path or not label_path or not os.path.exists(feat_path) or not os.path.exists(label_path):
            return None, "", "Merged dataset files do not exist."
            
        feat_cols_all = scan_file_columns(feat_path)
        label_cols_all = scan_file_columns(label_path)
        
        feat_key_col = feat_cols_all[0]
        label_key_col = label_cols_all[0]
        
        feat_lf = pl.scan_csv(feat_path)
        label_lf = pl.scan_csv(label_path)
        
        if selected_cols:
            feat_needed = [feat_key_col] + [c for c in selected_cols if c in feat_cols_all and c != feat_key_col]
            label_needed = [label_key_col] + [c for c in selected_cols if c in label_cols_all and c != label_key_col]
            
            # Deduplicate
            feat_needed = list(dict.fromkeys(feat_needed))
            label_needed = list(dict.fromkeys(label_needed))
            
            feat_lf = feat_lf.select(feat_needed)
            label_lf = label_lf.select(label_needed)
            
        merged_lf = feat_lf.join(label_lf, left_on=feat_key_col, right_on=label_key_col, how="inner")
        df = merged_lf.collect()
        
        dataset_name = f"Features ({os.path.basename(feat_path)}) + Labels ({os.path.basename(label_path)}) joined on '{feat_key_col}'"
        return df, dataset_name, None
    else:
        return None, "", "Invalid file mode."

def load_dataset_from_config(config_path="config.json", selected_cols: list[str] = None) -> tuple:
    """Wrapper function returning (df, features_cols, labels_cols, dataset_name, error_msg) for backward compatibility."""
    feat_cols, label_cols, dataset_name, err = get_dataset_schema_from_config(config_path)
    if err:
        return None, [], [], "", err
        
    df, name, err_load = load_selected_columns_from_config(config_path, selected_cols=selected_cols)
    if err_load:
        return None, feat_cols, label_cols, dataset_name, err_load
        
    return df, feat_cols, label_cols, dataset_name, None