import time
import experiments.utils
import experiments.data.utils
import experiments.methods.utils
from experiments.plot_experiment import visualize_stocks_adjacency, visualize_root_causes
import pandas as pd
import os
import numpy as np
import experiments.evaluation.evaluation
import experiments.evaluation.utils
from experiments.draw_network import stocks_visualization
seed=42

### args
return_type = "Close"
filename_data = "X_{}".format(return_type)

def run_real():    
    # make directory to put results
    if not os.path.exists("results/{}/".format(path)):
        os.makedirs("results/{}/".format(path))

    for t in [50]:#, 10, 20, 25, 50, 100, 1000]:
        with open('results/{}.csv'.format(path), 'a') as f:

            # graph initialization
            start = time.time()
            X, _, _, _, _ = experiments.data.utils.get_data(args, 0, 0, dataset="S&P", filename_data=filename_data)
            print("\n\nData generation process done. Time: {:.3f}\n\n".format(time.time() - start))

            # normalizes or standardizes data if supposed to
            X = experiments.data.utils.data_transform(X, args) 

            # causal discovery algorithms
            for r in range(args.runs):
                for method in args.methods: # assumes that it is ran only with one method
                    samples, d = X.shape
                    # B_est, W_est, T = methods.utils.execute_method(X, method, f, args, pc, dataset="thames")
                    B_est, W_est, T = experiments.methods.utils.execute_method(X, method, args, 0, d, t, dataset="S&P")

                    df = pd.DataFrame(W_est)
                    df.to_csv('results/{}/W_est_{}_{}_{}_run_{}.csv'.format(path, method, t, return_type, r), header=None, index=False)
                    # save result for future reference
                    print(W_est.shape)
                    # Create two subplots 
                    np.random.seed(seed)
                    ind = np.random.choice(d, 100, replace=False)
                    ind = np.sort(ind)

                    visualize_stocks_adjacency(W_est, args.number_of_lags, t, args, method=method, ind=ind)

                    np.random.seed(seed)
                    indx = np.random.choice(samples - samples % t, 50, replace=False)
                    indx = np.sort(indx)
                    C_est, _, _, _ = experiments.evaluation.evaluation.rct_approximation(method, X, t, W_est)
                    visualize_root_causes(C_est, t, args, method=method, indx=indx, indy=ind)

if __name__ == '__main__':
    parser, args = experiments.utils.get_args()
    print(vars(args))

    path = "S&P500"
    run_real()
