import time
import experiments.utils
import experiments.data.utils
import experiments.methods.utils
from experiments.plot_experiment import visualize
import os
import experiments.evaluation.evaluation
import experiments.evaluation.utils
seed=42

### args
filename_data = "dream3"

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

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

                # graph initialization
                start = time.time()
                X, _, _, B_true, W_true = experiments.data.utils.get_data(args, 0, 0, dataset="dream3", dataset_id=dataset_id, 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) 

                current = {}
                avgT = {}

                for key in args.methods:
                    current[key] = []
                    avgT[key] = []

                # 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
                        #### assuming number of lags is zero                        
                        B_est, W_est, T = experiments.methods.utils.execute_method(X, method, args, 0, d, t, dataset="dream3")

                        # Q = np.zeros((d, d))
                        # for i in range(args.algo_lags + 1):
                        #     Q += B_est[:, i * d : (i + 1) * d]
                        # Q = (Q!=0)
                        # # mask =  np.ones((d,d)) - np.eye(d)
                        # B_est = Q 
                        W_est = B_est
                        
                        # Create visualizations
                        for i in range(args.number_of_lags + 1):
                            visualize(B_true[:, i * d : (i + 1) * d], W_est[:, i * d : (i + 1) * d], method=method, filename="dream3_dataset_{}".format(dataset_id))

                        experiments.evaluation.utils.compute_metrics(method, current, filename_data, r, t, T, X, None, B_true, W_true, B_est, W_est, args)
                
                # save average results in csv
                experiments.evaluation.utils.save_results(current, filename_data, args, r)

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

    path = "dream3"
    run_real()
