from datetime import datetime
import numpy as np
import os

from collections import Iterable
from termcolor import colored
from tqdm import tqdm

from experiment.collect import debug_show
from pca.sparse_pca import SparsePCA
from symbols.file_utils import make_path, make_dir
from symbols.logger.precondition_logger import PreconditionLogger
from symbols.logger.transition_logger import TransitionLogger
import matplotlib.pyplot as plt
import pickle




if __name__ == '__main__':

    TASK_ID = 0
    ACTION = 1
    dir_name = datetime.today().strftime('%Y%m%d') + "_raw"
    dir_name = "20190912_raw"
    directory = make_path(dir_name)
    pca_path = os.path.abspath('new_pca.dat')
    task_dir = make_path(directory, str(TASK_ID))

    dir = make_path(task_dir, 'transition_data')

    obs = list()
    next_obs = list()

    for i in range(15):

        with open(make_path(dir, str(i)), 'rb') as file:
            transitions = pickle.load(file)
            for x in transitions:

                if x[2] == ACTION:
                    obs.append(x[1])
                    next_obs.append(x[-1])

    from pca import PCA
    pca = PCA(40)
    pca.load(pca_path)

    for x, y in zip(obs, next_obs):
        debug_show(x, pca)
        debug_show(y, pca)
        q = 0
    # for x in next_obs:
    #     debug_show(x, pca)

