data_loader.py

"""
Data Loading Module for Weyland-Yutani Mining Ops Dashboard
Handles loading data from Google Sheets and generating demo data
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import re


def load_data_from_sheets(
    spreadsheet_url: str,
    credentials: Optional[Dict[str, Any]] = None,
    sheet_name: str = "Data"
) -> pd.DataFrame:
    """
    Load data from Google Sheets
    
    Args:
        spreadsheet_url: Full URL or spreadsheet ID
        credentials: Service account credentials dict
        sheet_name: Name of the sheet to load
    
    Returns:
        DataFrame with mining data
    """
    try:
        import gspread
        from google.oauth2.service_account import Credentials
    except ImportError:
        raise ImportError(
            "Please install gspread and google-auth: "
            "pip install gspread google-auth"
        )
    
    # Extract spreadsheet ID from URL if needed
    spreadsheet_id = extract_spreadsheet_id(spreadsheet_url)
    
    # Setup credentials
    scopes = [
        'https://www.googleapis.com/auth/spreadsheets.readonly',
        'https://www.googleapis.com/auth/drive.readonly'
    ]
    
    if credentials:
        creds = Credentials.from_service_account_info(credentials, scopes=scopes)
    else:
        try:
            creds = Credentials.from_service_account_file(
                'credentials.json', 
                scopes=scopes
            )
        except FileNotFoundError:
            return load_public_sheet(spreadsheet_id, sheet_name)
    
    client = gspread.authorize(creds)
    
    try:
        spreadsheet = client.open_by_key(spreadsheet_id)
        worksheet = spreadsheet.worksheet(sheet_name)
    except gspread.exceptions.SpreadsheetNotFound:
        raise ValueError(f"Spreadsheet not found: {spreadsheet_id}")
    except gspread.exceptions.WorksheetNotFound:
        raise ValueError(f"Worksheet not found: {sheet_name}")
    
    data = worksheet.get_all_values()
    
    if not data:
        raise ValueError("No data found in spreadsheet")
    
    df = pd.DataFrame(data[1:], columns=data[0])
    
    df = process_mining_data(df)
    
    return df


def load_public_sheet(spreadsheet_id: str, sheet_name: str) -> pd.DataFrame:
    """
    Load data from a publicly accessible Google Sheet via CSV export
    
    Args:
        spreadsheet_id: The spreadsheet ID
        sheet_name: Name of the sheet (used to get gid)
    
    Returns:
        DataFrame with mining data
    """
    urls_to_try = [
        f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/gviz/tq?tqx=out:csv&sheet={sheet_name}",
        f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/export?format=csv&sheet={sheet_name}",
    ]
    
    for url in urls_to_try:
        try:
            df = pd.read_csv(url)
            df = process_mining_data(df)
            return df
        except Exception:
            continue
    
    raise ValueError(
        f"Could not load spreadsheet. Make sure it's publicly accessible "
        f"or provide service account credentials."
    )


def extract_spreadsheet_id(url_or_id: str) -> str:
    """
    Extract spreadsheet ID from URL or return as-is if already an ID
    
    Args:
        url_or_id: Full Google Sheets URL or just the ID
    
    Returns:
        Spreadsheet ID
    """
    if 'docs.google.com' in url_or_id or 'spreadsheets' in url_or_id:
        patterns = [
            r'/spreadsheets/d/([a-zA-Z0-9-_]+)',
            r'/d/([a-zA-Z0-9-_]+)',
            r'id=([a-zA-Z0-9-_]+)',
        ]
        
        for pattern in patterns:
            match = re.search(pattern, url_or_id)
            if match:
                return match.group(1)
        
        raise ValueError(f"Could not extract spreadsheet ID from URL: {url_or_id}")
    
    return url_or_id


def process_mining_data(df: pd.DataFrame) -> pd.DataFrame:
    """
    Process raw mining data from Google Sheets
    
    Args:
        df: Raw DataFrame from Google Sheets
    
    Returns:
        Processed DataFrame with proper types
    """
    df = df.copy()
    
    df = df.dropna(how='all')
    
    date_col = df.columns[0]
    
    try:
        df[date_col] = pd.to_datetime(df[date_col])
    except Exception:
        for fmt in ['%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%Y/%m/%d']:
            try:
                df[date_col] = pd.to_datetime(df[date_col], format=fmt)
                break
            except Exception:
                continue
    
    for col in df.columns[1:]:
        if col.lower() in ['dayofweek', 'day']:
            continue
        try:
            if df[col].dtype == 'object':
                df[col] = df[col].astype(str).str.replace(',', '.', regex=False)
            df[col] = pd.to_numeric(df[col], errors='coerce')
        except Exception:
            pass
    
    df = df.sort_values(date_col).reset_index(drop=True)
    
    return df


def load_demo_data(
    num_days: int = 40,
    num_mines: int = 3,
    start_date: Optional[datetime] = None,
    seed: Optional[int] = None
) -> pd.DataFrame:
    """
    Generate realistic demo mining data
    
    Args:
        num_days: Number of days to generate
        num_mines: Number of mines
        start_date: Starting date (default: 2099-11-02)
        seed: Random seed for reproducibility
    
    Returns:
        DataFrame with demo mining data
    """
    if seed is not None:
        np.random.seed(seed)
    
    if start_date is None:
        start_date = datetime(2099, 11, 2)
    
    mine_names = ['LV-426', 'Origae-6', 'Fiorina 151', 'Hadley\'s Hope', 'Fury 161', 'Prometheus'][:num_mines]
    base_outputs = [1.0, 1.0, 1.0, 0.8, 0.6, 1.2][:num_mines]
    day_factors = [1.0, 1.0, 1.0, 1.0, 1.0, 0.8, 0.6]
    mean = 50
    std_dev = 20
    correlation = 0.3
    daily_growth = 0.02
    
    events = [
        (8, 3, 0.4, 1.0),
        (18, 1, 1.4, 0.5),
        (25, 2, 0.6, 0.8),
        (32, 2, 1.3, 0.7),
    ]
    
    event_occurs = [np.random.random() < prob for _, _, _, prob in events]
    
    dates = []
    day_names = []
    mine_data = {name: [] for name in mine_names}
    
    prev_values = [mean] * num_mines
    
    for day in range(num_days):
        current_date = start_date + timedelta(days=day)
        dates.append(current_date)
        
        dow = current_date.weekday()
        day_names.append(['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun'][dow])
        
        day_factor = day_factors[dow]
        
        trend_factor = (1 + daily_growth) ** day
        
        event_factor = 1.0
        for i, (offset, duration, factor, _) in enumerate(events):
            if event_occurs[i] and offset <= day < offset + duration:
                position = (day - offset) / (duration - 1) if duration > 1 else 0.5
                bell = np.exp(-((position - 0.5) * 3) ** 2)
                event_factor *= 1 + (factor - 1) * bell
        
        for m, mine_name in enumerate(mine_names):
            random_value = np.random.normal(mean, std_dev)
            smoothed = correlation * prev_values[m] + (1 - correlation) * random_value
            value = smoothed * base_outputs[m] * day_factor * trend_factor * event_factor
           
            value = max(0.01, value)
            
            mine_data[mine_name].append(round(value, 2))
            prev_values[m] = smoothed
    
    df = pd.DataFrame({
        'Date': dates,
        'DayOfWeek': day_names,
        **mine_data
    })
    
    df['Total'] = df[mine_names].sum(axis=1).round(2)
    
    return df


def validate_data(df: pd.DataFrame) -> Dict[str, Any]:
    """
    Validate loaded data and return validation results
    
    Args:
        df: DataFrame to validate
    
    Returns:
        Dictionary with validation results
    """
    results = {
        'valid': True,
        'warnings': [],
        'errors': [],
        'info': {}
    }
    
    if df.empty:
        results['valid'] = False
        results['errors'].append("DataFrame is empty")
        return results
    
    date_col = df.columns[0]
    if not pd.api.types.is_datetime64_any_dtype(df[date_col]):
        results['warnings'].append(f"First column '{date_col}' is not datetime type")
    
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    if len(numeric_cols) == 0:
        results['valid'] = False
        results['errors'].append("No numeric columns found")
        return results
    
    missing = df.isnull().sum()
    if missing.any():
        cols_with_missing = missing[missing > 0].to_dict()
        results['warnings'].append(f"Missing values found: {cols_with_missing}")
    
    for col in numeric_cols:
        if (df[col] < 0).any():
            results['warnings'].append(f"Negative values found in column '{col}'")
    
    results['info'] = {
        'rows': len(df),
        'columns': len(df.columns),
        'date_range': f"{df[date_col].min()} to {df[date_col].max()}" if pd.api.types.is_datetime64_any_dtype(df[date_col]) else "N/A",
        'numeric_columns': numeric_cols
    }
    
    return results