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, 8, 70, 23), 'pdb.set_trace', 'pdb.set_trace', ({}, {}), '()', False, 'import pdb\n'), ((9, 13, 9, 29), 'numpy.array', 'np.array', ({(9, 22, 9, 28): 'matrix'}, {}), '(matrix)', True, 'import numpy as np\n'), ((29, 26, 29, 55), 'copy.deepcopy', 'copy.deepcopy', ({(29, 40, 29, 54): 'self._cur_traj'}, {}), '(self._cur_traj)', False, 'import copy\n'), ((48, 19, 48, 48), 'numpy.array', 'np.array', ({(48, 28, 48, 47): '[tt[3] for tt in t]'}, {}), '([tt[3] for tt in t])', True, 'import numpy as np\n'), ((59, 12, 59, 38), 'pickle.dump', 'pickle.dump', ({(59, 24, 59, 28): 'traj', (59, 30, 59, 37): 'outfile'}, {}), '(traj, outfile)', False, 'import pickle\n'), ((73, 12, 73, 36), 'json.dump', 'json.dump', ({(73, 22, 73, 26): 'traj', (73, 28, 73, 35): 'outfile'}, {}), '(traj, outfile)', False, 'import json\n'), ((16, 16, 16, 31), 'pdb.set_trace', 'pdb.set_trace', ({}, {}), '()', False, 'import pdb\n'), ((45, 18, 45, 47), 'numpy.array', 'np.array', ({(45, 27, 45, 46): '[tt[0] for tt in t]'}, {}), '([tt[0] for tt in t])', True, 'import numpy as np\n'), ((46, 18, 46, 47), 'numpy.array', 'np.array', ({(46, 27, 46, 46): '[tt[1] for tt in t]'}, {}), '([tt[1] for tt in t])', True, 'import numpy as np\n'), ((47, 18, 47, 47), 'numpy.array', 'np.array', ({(47, 27, 47, 46): '[tt[2] for tt in t]'}, {}), '([tt[2] for tt in t])', True, 'import numpy as np\n')]