import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import datapuller
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
from sklearn.model_selection import train_test_split
from keras import Sequential, layers, optimizers, losses

def TrainAI(include_pull):

    if (include_pull):
        # Pull New Data
        datapuller.pull()

    # 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"))

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

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

    # 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)

    # Create a normalizer to nomralize the data
    normalizer = layers.Normalization(axis=-1)
    normalizer.adapt(np.array(train_features))

    # Create the DNN
    dnn_model = Sequential([
        layers.Input(shape=(train_features.shape[1],)), # Load the feature count dynamically
        normalizer,
        layers.Dense(64, activation='elu'),
        layers.Dense(64, activation='elu'),
        layers.Dense(1)
    ])

    # Configure the model
    dnn_model.compile(
        optimizer=optimizers.Adam(learning_rate=0.00001, clipvalue=1.0),
        loss=losses.Huber()
    )

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

    # Train the model
    Training_Data = dnn_model.fit( 
        train_features,
        train_labels,
        batch_size=64,
        epochs=39,                  # Tuned to the point before overfitting
        verbose=1,                  # Show progress
        validation_split = 0.2      # Calculate validation results on 20% of the training data.
    )

    # Predict
    test_predictions = dnn_model.predict(test_features).flatten()
    a = plt.axes(aspect='equal')
    plt.scatter(test_labels, test_predictions)
    plt.xlabel('True Values')
    plt.ylabel('Predictions')
    lims = [0, 50]
    plt.xlim(lims)
    plt.ylim(lims)
    _ = plt.plot(lims, lims)


    # Current Test Results: 1.221876300405711e-05
    test_results = dnn_model.evaluate(
        test_features, test_labels, verbose=0
    )
    print(f"Test Results: {test_results}")

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

if __name__ == "__main__":
    TrainAI(False)