import json
import copy
import pdb
import numpy as np
import pickle


def listify_mat(matrix):
    matrix = np.array(matrix).astype(str)
    if len(matrix.shape) > 1:
        matrix_list = []
        for row in matrix:
            try:
                matrix_list.append(list(row))
            except:
                pdb.set_trace()
        return matrix_list
    else:

        return list(matrix)


class Recorder():
    def __init__(self):
        self._traj, self._cur_traj = [], []
        return

    def pack_traj(self):
        self._traj.append(copy.deepcopy(self._cur_traj))
        self._cur_traj = []

        return

    def add(self, o, a, r, d):
        # self._cur_traj.append((o, a, r, d))
        self._cur_traj.append(
            (listify_mat(o), listify_mat(a), listify_mat(r), d))
        return

    def export_pickle(self, filename='traj'):
        if filename == '':
            raise ValueError('incorrect file name')
        traj = []
        for t in self._traj:
            obs = np.array([tt[0] for tt in t]).astype(np.float32)
            act = np.array([tt[1] for tt in t]).astype(np.float32)
            rwd = np.array([tt[2] for tt in t]).astype(np.float32)
            done = np.array([tt[3] for tt in t])

            # pdb.set_trace()
            traj.append({
                'observations': obs[:-1],
                'next_observations': obs[1:],
                'actions': act[:-1],
                'rewards': rwd[:-1],
                'terminals': done[:-1]
            })
        with open('{}.pkl'.format(filename), 'wb') as outfile:
            pickle.dump(traj, outfile)
        return

    def export(self, filename='traj'):
        if filename == '':
            raise ValueError('incorrect file name')
        traj = {'traj': []}
        for t in self._traj:
            traj['traj'].append(t)
        # json.dumps(traj, sort_keys=True, indent=4)

        pdb.set_trace()

        with open('{}.json'.format(filename), 'w') as outfile:
            json.dump(traj, outfile)
        return
[(70, 'pdb.set_trace', 'pdb.set_trace', False, 'import pdb\n'), (9, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (29, 'copy.deepcopy', 'copy.deepcopy', False, 'import copy\n'), (48, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (59, 'pickle.dump', 'pickle.dump', False, 'import pickle\n'), (73, 'json.dump', 'json.dump', False, 'import json\n'), (16, 'pdb.set_trace', 'pdb.set_trace', False, 'import pdb\n'), (45, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (46, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (47, 'numpy.array', 'np.array', True, 'import numpy as np\n')]