"""Raw 5-minute OHLCV loading and cleaning."""
import numpy as np
import pandas as pd


class GoldDataProcessor:
    """
    Data processor for 5-minute gold price data.
    Handles data loading, cleaning, and preprocessing.
    """
    def __init__(self, fill_gaps=True, remove_outliers=True, handle_after_hours=True):
        self.fill_gaps = fill_gaps
        self.remove_outliers = remove_outliers
        self.handle_after_hours = handle_after_hours

    def load_data(self, file_path):
        """
        Load data from CSV file and set appropriate index
        """
        # Load data from CSV
        data = pd.read_csv(file_path)

        # Ensure timestamp column exists
        if 'timestamp' not in data.columns:
            raise ValueError("CSV must contain a 'timestamp' column")

        # Convert timestamp to datetime and set as index
        data['timestamp'] = pd.to_datetime(data['timestamp'])
        data = data.drop_duplicates(subset=['timestamp'])
        data.set_index('timestamp', inplace=True)

        # Ensure all necessary columns exist
        required_columns = ['open', 'high', 'low', 'close', 'volume']
        missing_columns = [col for col in required_columns if col not in data.columns]
        if missing_columns:
            raise ValueError(f"Missing required columns: {missing_columns}")

        # Make column names lowercase if they aren't already
        data.columns = [col.lower() for col in data.columns]

        return self._preprocess_data(data)

    def _preprocess_data(self, data):
        """
        Apply preprocessing steps to clean and prepare the data
        """
        # Sort data by timestamp to ensure chronological order
        data = data.sort_index()

        # Fill gaps in 5-minute data
        if self.fill_gaps:
            data = self._fill_time_gaps(data)

        # Remove outliers
        if self.remove_outliers:
            data = self._remove_price_outliers(data)

        # Handle after-hours and weekend data
        if self.handle_after_hours:
            data = self._handle_after_hours_data(data)

        return data

    def _fill_time_gaps(self, data):
        """
        Fill gaps in 5-minute data by reindexing with complete 5-minute intervals
        during trading hours
        """
        # Create complete 5-minute intervals
        start_date = data.index.min()
        end_date = data.index.max()
        full_range = pd.date_range(start=start_date, end=end_date, freq='5min')

        # Reindex the data
        reindexed_data = data.reindex(full_range)

        # Forward fill price data (OHLC)
        for col in ['open', 'high', 'low', 'close']:
            reindexed_data[col] = reindexed_data[col].ffill()

        # Fill volume with zeros
        reindexed_data['volume'] = reindexed_data['volume'].fillna(0)

        return reindexed_data

    def _remove_price_outliers(self, data, z_threshold=10):
        """
        Remove extreme price outliers based on z-score of returns
        """
        # Calculate returns
        returns = data['close'].pct_change()

        # Calculate rolling z-score (20 periods ~ 100 minutes)
        rolling_mean = returns.rolling(window=20).mean()
        rolling_std = returns.rolling(window=20).std()
        z_scores = (returns - rolling_mean) / rolling_std

        # Identify outliers
        outliers = abs(z_scores) > z_threshold

        if outliers.sum() > 0:
            print(f"Identified {outliers.sum()} outliers out of {len(data)} records")

            # Replace outlier prices with interpolated values
            clean_data = data.copy()
            outlier_indices = outliers[outliers].index

            for col in ['open', 'high', 'low', 'close']:
                clean_data.loc[outlier_indices, col] = np.nan
                clean_data[col] = clean_data[col].interpolate(method='linear')

            return clean_data

        return data

    def _handle_after_hours_data(self, data):
        """
        Handle after-hours and weekend data in gold trading
        """
        # Create day of week and hour features
        data['day_of_week'] = data.index.dayofweek
        data['hour'] = data.index.hour

        # Filter out weekends (Saturday and Sunday)
        # For gold, market typically closes Friday ~5PM ET and opens Sunday ~6PM ET
        weekend_mask = (data['day_of_week'] == 5) & (data['hour'] >= 17)  # After Friday 5PM
        weekend_mask |= (data['day_of_week'] == 6)  # All of Saturday
        weekend_mask |= (data['day_of_week'] == 0) & (data['hour'] < 18)  # Before Sunday 6PM

        # Mark weekend data
        data['is_weekend'] = weekend_mask

        # Forward fill prices for weekend gaps
        # We don't remove weekend data to maintain continuity
        # but we mark it so we can filter it later if needed

        # Remove temporary columns if needed
        data = data.drop(['day_of_week', 'hour'], axis=1)

        return data