from sklearn.preprocessing import MinMaxScaler
import numpy as np

def MakeFeatures(df):
    # Convert all F64 to F32 to save ram
    df = df.astype({col: 'float32' for col in df.select_dtypes('float64').columns})

    # Create Grouped columns by ticker
    grouped = df.groupby('Ticker')

    # Candle Wick's
    df['Spread'] = (df['High'] - df['Low']) / df['Close'] # in percentage of price
    candle_top = df[['Open', 'Close']].max(axis=1)
    df['Body_Size'] = (df['Close'] - df['Open']).abs() / (df['High'] - df['Low'])
    df['Upper_Shadow'] = (df['High'] - candle_top) / (df['High'] - df['Low'])

     # Is volume 2x higher than the 20-day average?
    df['Vol_Intensity'] = df['Volume'] / grouped['Volume'].transform(lambda x: x.rolling(20).mean())
    # Volume Change
    df['Volume_Chg'] = grouped['Volume'].pct_change()

    # Moving Average Crossover (Golden/Death Cross logic)
    df['Moving_Average_5'] = grouped['Close'].transform(lambda x: x.rolling(5).mean())
    df['Moving_Average_20'] = grouped['Close'].transform(lambda x: x.rolling(20).mean())
    # if short term > long term (bullish), else 0
    df['Trend_Signal'] = (df['Moving_Average_5'] > df['Moving_Average_20']).astype(np.float32)
    # 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 = grouped['Close'].transform(lambda x: x.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).replace(0, 1e-6)

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

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

    # RSI (Relative Strength Index)
    delta = grouped['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))

    # Return lagged
    for lag in range(1, 4):
        df[f'Return_Lag_{lag}'] = grouped['Return'].shift(lag)
        df[f'Vol_Lag_{lag}'] = grouped['Volume_Chg'].shift(lag)

    # This is our training metric of price difference 5 days ahead
    df['Target_Close'] = (np.log(df['Close'].shift(-5) / df['Close']) / df['Volatility_5']).clip(-10, 10)

    # Save the overall trend to predict based off of later
    with open("Target_Close_Average.txt", "w") as file:
        file.write(str(df["Target_Close"].mean()))

    # Make a feature for the S&P500 average for the day
    df['SP500_Market_Log_Return'] = df.groupby('Date')['Log_Return'].transform('mean')
    # Relative Strength agains the S&P500
    df['SP500_Relative_Performance'] = df['Log_Return'] - df['SP500_Market_Log_Return']
    # S&P500 market trend
    daily_trend = df.groupby('Date')['SP500_Market_Log_Return'].first().rolling(window=20).mean()
    daily_trend.name = 'SP500_Market_Trend_20'
    df = df.merge(daily_trend, on='Date', how='left')

    # Drop every column that is a raw price or an unscaled average
    cols_to_drop = ['Open', 'High', 'Low', 'Volume', 'Close', 'Moving_Average_5', 'Moving_Average_20', 'Ticker']
    for col in cols_to_drop:
        if col in df.columns:
            df.drop(col, axis=1, inplace=True)

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

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

    # 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 with new features
    return df

def Prepare(df):

    df = df.replace([np.inf, -np.inf], 0)

    # 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 = MinMaxScaler()
    X_scaled = feature_scaler.fit_transform(X)

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

    return X_scaled, y_scaled, feature_scaler, target_scaler