from sklearn.preprocessing import StandardScaler
import pandas as pd
import numpy as np

def MakeFeatures(df, i):
    # Remove the ticker column
    df.columns = df.columns.get_level_values(0)
    
    # Make sure Date is a number object
    df = df.reset_index()
    df['Date'] = pd.to_numeric(pd.to_datetime(df['Date']))

    # Add the Symbol column for tracking | as an int 1 hot encoded
    df['Symbol'] = i

    # Add feature Spread
    df['Spread'] = abs( df['High'] - df['Low'] )

    # Add feature for Returns
    df['Return'] = df['Close'].pct_change()

    # Add feature for volitility last 5
    df['Volatility_5'] = df['Return'].transform(lambda x: x.rolling(5).std())

    # Add feature for volitility last 20
    df['Volatility_20'] = df['Return'].transform(lambda x: x.rolling(20).std())

    # RSI (Relative Strength Index)
    delta = df['Close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    df['RSI'] = 100 - (100 / (1 + rs))

    # Moving Average Crossover (Golden/Death Cross logic)
    df['Moving_Average_5'] = df['Close'].rolling(window=5).mean()
    df['Moving_Average_20'] = df['Close'].rolling(window=20).mean()
    # if short term > long term (bullish), else 0
    df['Trend_Signal'] = (df['Moving_Average_5'] > df['Moving_Average_20']).astype(int)

    # Distance from MA (How overextended are we?)
    df['Dist_From_MA20'] = (df['Close'] / df['Moving_Average_20']) - 1

    # Bollinger Band Position (Where are we relative to volatility?)
    std_20 = df['Close'].rolling(20).std()
    upper_band = df['Moving_Average_20'] + (std_20 * 2)
    lower_band = df['Moving_Average_20'] - (std_20 * 2)
    df['BB_Pos'] = (df['Close'] - lower_band) / (upper_band - lower_band)

    # Log Returns (Better for AI than pct_change for statistical normality)
    df['Log_Return'] = np.log(df['Close'] / df['Close'].shift(1))

    # This is our training metric of 5 days ahead
    df['Target_Close'] = df['Close'].shift(-5).pct_change()

    # Volume Change
    df['Volume_Chg'] = df['Volume'].pct_change()

    # Return new df with new features
    return df

def Prepare(df):

    # Remove indicators and set the target
    X = df.drop('Target_Close', axis=1)
    Y = df['Target_Close']

    # Scale the features to the same size
    feature_scaler = StandardScaler()
    X_scaled = feature_scaler.fit_transform(X)

    # Safe for the Y
    target_scaler = StandardScaler()
    y_scaled = target_scaler.fit_transform(Y.values.reshape(-1, 1))

    return X_scaled, y_scaled, feature_scaler, target_scaler

def CleanDF(df):
    # Make date the index so it doesnt influence the training
    df.set_index('Date', inplace=True)

    # Drop rows with null values
    df.dropna(inplace=True)

    # Replace Infinity with 0 -> This fixes the AI mental breakdown
    df['Volume_Chg'] = df['Volume_Chg'].replace([np.inf, -np.inf], 0)

    # Return new df thats cleaned
    return df