import pandas as pd
import numpy as np
import yfinance as yf
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

class DataProcessor:
    def __init__(self):
        pass
    
    def download_stock_data(self, symbol, period="1y", interval="1d"):
        """Download stock data using yfinance"""
        try:
            ticker = yf.Ticker(symbol)
            data = ticker.history(period=period, interval=interval)
            
            if data.empty:
                return None
            
            # Clean column names
            data.columns = [col.title() for col in data.columns]
            
            # Remove timezone info for consistency
            if data.index.tz is not None:
                data.index = data.index.tz_localize(None)
            
            return data
            
        except Exception as e:
            print(f"Error downloading data for {symbol}: {str(e)}")
            return None

    def prepare_features_and_labels(self, data, future_days=3, threshold=0.02):
        """Prepare features and labels for machine learning"""
        try:
            # Calculate future returns for labeling
            data = data.copy()
            data['Future_Return'] = data['Close'].pct_change(periods=future_days).shift(-future_days)
            
            # Create labels based on future returns
            # 1 = BUY (expected gain > threshold)
            # -1 = SELL (expected loss > threshold) 
            # 0 = HOLD (neutral)
            conditions = [
                data['Future_Return'] > threshold,
                data['Future_Return'] < -threshold
            ]
            choices = [1, -1]
            data['Label'] = np.select(conditions, choices, default=0)
            
            # Select feature columns (all technical indicators)
            feature_columns = [col for col in data.columns if col not in 
                             ['Open', 'High', 'Low', 'Close', 'Volume', 'Future_Return', 'Label']]
            
            # Remove rows with NaN values
            clean_data = data.dropna()
            
            if len(clean_data) == 0:
                return pd.DataFrame(), pd.Series()
            
            # Prepare features and labels
            features = clean_data[feature_columns]
            labels = clean_data['Label']
            
            return features, labels
            
        except Exception as e:
            print(f"Error preparing features and labels: {str(e)}")
            return pd.DataFrame(), pd.Series()
    
    def get_latest_features(self, data, feature_columns):
        """Extract latest features for prediction"""
        try:
            if data.empty or not feature_columns:
                return None
            
            # Get the latest row with all required features
            latest_data = data[feature_columns].dropna()
            
            if latest_data.empty:
                return None
            
            # Return the latest row as a list
            return latest_data.iloc[-1].values.tolist()
            
        except Exception as e:
            print(f"Error extracting latest features: {str(e)}")
            return None
    
    def calculate_returns(self, data, periods=[1, 3, 5]):
        """Calculate returns for different periods"""
        try:
            returns_data = data.copy()
            
            for period in periods:
                returns_data[f'Return_{period}d'] = data['Close'].pct_change(periods=period)
            
            return returns_data
            
        except Exception as e:
            print(f"Error calculating returns: {str(e)}")
            return data
    
    def add_price_features(self, data):
        """Add additional price-based features"""
        try:
            enhanced_data = data.copy()
            
            # Price position within daily range
            enhanced_data['Price_Position'] = (data['Close'] - data['Low']) / (data['High'] - data['Low'])
            enhanced_data['Price_Position'].fillna(0.5, inplace=True)
            
            # Gap from previous close
            enhanced_data['Gap'] = (data['Open'] - data['Close'].shift(1)) / data['Close'].shift(1)
            
            # Intraday return
            enhanced_data['Intraday_Return'] = (data['Close'] - data['Open']) / data['Open']
            
            # High-Low spread
            enhanced_data['HL_Spread'] = (data['High'] - data['Low']) / data['Close']
            
            return enhanced_data
            
        except Exception as e:
            print(f"Error adding price features: {str(e)}")
            return data