import pandas as pd
from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score
import features
import joblib
import os

# Suppress TensorFlow INFO and WARNING logs
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 

# Suppress specialized XLA and autotuning logs
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices=false'
os.environ['TF_CPP_MAX_VLOG_LEVEL'] = '0'

# CPU only
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from sklearn.model_selection import train_test_split
from keras import Sequential, layers, optimizers
from keras.callbacks import ReduceLROnPlateau

def TrainAI():

    # Get the CWD for pathing due to being called from C#
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
    DATA_DIR = os.path.join(SCRIPT_DIR, "data")

    # Load the dataset
    dataset = pd.read_parquet(os.path.join(DATA_DIR, "stocks.parquet"))

    # Use external featuers to make sure loaded is the same
    dataset = features.MakeFeatures(dataset)

    # Make sure empty columns dont exist
    dataset = dataset.dropna()

    # Create the X, Y vareables
    X, Y, X_Scaler, Y_Scaler = features.Prepare(dataset)

    # Save the scalers for future use
    joblib.dump(X_Scaler, os.path.join(DATA_DIR, "feature_scaler.pkl"))
    joblib.dump(Y_Scaler, os.path.join(DATA_DIR, "target_scaler.pkl"))

    # Show the datatypes
    print(dataset.dtypes)

    # Split out the test and train
    train_features, test_features, train_labels, test_labels = train_test_split(X, Y, test_size=0.2, shuffle=False) # Keep the training and test data in order as its sequential

    # Create the DNN
    dnn_model = Sequential([
        layers.Input(shape=(train_features.shape[1],)), # Load the feature count dynamically
        layers.Dense(256, activation='elu'),            # DNN layer
        layers.BatchNormalization(),                    # Nomralize
        layers.Dropout(0.1),                            # Small dropout to prevent it from memorizing noise
        layers.Dense(128, activation='elu'),            # DNN layer
        layers.BatchNormalization(),                    # Nomralize
        layers.Dropout(0.1),                            # Small dropout to prevent it from memorizing noise
        layers.Dense(24, activation='elu'),             # DNN layer
        layers.Dense(1, activation='linear')            # DNN layer
    ])

    # Configure the model
    dnn_model.compile(
        optimizer=optimizers.Adam(learning_rate=0.0001, clipvalue=1.0),
        loss="mse",
        metrics=['mae']             # See it train while it runs
    )

    # Show the summary before training the model
    dnn_model.summary()

    # Learning rate reducer
    reduce_lr = ReduceLROnPlateau(monitor='loss', factor=0.5, patience=5, min_lr=0.0000001)

    # Train the model
    Training_Data = dnn_model.fit( 
        train_features,
        train_labels,
        batch_size=64,
        epochs=50,                  # Tuned to the point before overfitting
        verbose=1,                  # Show progress
        shuffle=False,              # Time series data
        callbacks=[reduce_lr]       # Reduce the learning_rate every run
    )

    # Current Test Results: 0.3382711112499237
    test_results = dnn_model.evaluate(
        test_features, test_labels, verbose=0
    )

    # Perform test on test data split earlier
    predictions = dnn_model.predict(test_features)

    # Convert to Binary Direction (The "Signal")
    y_true_binary = (test_labels > 0).astype(int)
    y_pred_binary = (predictions > 0).astype(int)

    # Calculate the metrics
    acc = accuracy_score(y_true_binary, y_pred_binary)
    prec = precision_score(y_true_binary, y_pred_binary)
    rec = recall_score(y_true_binary, y_pred_binary)
    f1 = f1_score(y_true_binary, y_pred_binary)

    # Output the meterics; this gets spit out when training from the python file directly. The C# interop does not output these
    print(f"Test Loss: {test_results[0]:.4f}")
    print(f"Test MAE: {test_results[1]:.2%}")
    print(f"Test Accuracy: {acc:.2%}")
    print(f"Test Precision: {prec:.2%}")
    print(f"Test Recall: {rec:.2%}")
    print(f"F1 Score: {f1:.2%}")

    # Save the model
    dnn_model.save(os.path.join(DATA_DIR, "model.keras"))

if __name__ == "__main__":
    TrainAI()