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
[((1905, 1920), 'pdb.set_trace', 'pdb.set_trace', ([], {}), '()\n', (1918, 1920), False, 'import pdb\n'), ((108, 124), 'numpy.array', 'np.array', (['matrix'], {}), '(matrix)\n', (116, 124), True, 'import numpy as np\n'), ((555, 584), 'copy.deepcopy', 'copy.deepcopy', (['self._cur_traj'], {}), '(self._cur_traj)\n', (568, 584), False, 'import copy\n'), ((1212, 1241), 'numpy.array', 'np.array', (['[tt[3] for tt in t]'], {}), '([tt[3] for tt in t])\n', (1220, 1241), True, 'import numpy as np\n'), ((1590, 1616), 'pickle.dump', 'pickle.dump', (['traj', 'outfile'], {}), '(traj, outfile)\n', (1601, 1616), False, 'import pickle\n'), ((1997, 2021), 'json.dump', 'json.dump', (['traj', 'outfile'], {}), '(traj, outfile)\n', (2006, 2021), False, 'import json\n'), ((318, 333), 'pdb.set_trace', 'pdb.set_trace', ([], {}), '()\n', (331, 333), False, 'import pdb\n'), ((1010, 1039), 'numpy.array', 'np.array', (['[tt[0] for tt in t]'], {}), '([tt[0] for tt in t])\n', (1018, 1039), True, 'import numpy as np\n'), ((1077, 1106), 'numpy.array', 'np.array', (['[tt[1] for tt in t]'], {}), '([tt[1] for tt in t])\n', (1085, 1106), True, 'import numpy as np\n'), ((1144, 1173), 'numpy.array', 'np.array', (['[tt[2] for tt in t]'], {}), '([tt[2] for tt in t])\n', (1152, 1173), True, 'import numpy as np\n')]