import zipfile
from collections import defaultdict
from zipfile import ZipFile
import os
import numpy as np

from symbols.file_utils import make_path, make_dir, save
import pandas as pd





class PreconditionLoggerPD:
    """
    A class for logging transition samples.
    """

    def __init__(self, base_data_dir):

        base_data_dir = make_path(base_data_dir, 'init_set_data')
        make_dir(base_data_dir)
        self.dir = base_data_dir
        self.initiation_data = pd.DataFrame(columns=['episode', 'state', 'object_state', 'option', 'object', 'can_execute'])

    def close(self):
        self.initiation_data.to_pickle('{}/init.pkl'.format(self.dir), compression='gzip')

    def save(self):
        self.initiation_data.to_pickle('{}/init.pkl'.format(self.dir), compression='gzip')

    def log_sample(self,
                   episode: int,
                   state: np.ndarray,
                   observation: np.ndarray,
                   option,
                   object_id,
                   can_exec: bool):
        """
        Log a single state-option sample to file.
        :param observation:
        :param episode:
        :param state: the state under consideration, represented as a feature vector
        :param option: the number of the option under consideration
        :param can_exec: whether the option could be executed in the given state
        """

        self.initiation_data.loc[len(self.initiation_data)] = [episode, state, observation, option,
                                                               object_id, can_exec]

class PreconditionLogger:
    """
    A class for logging precondition/initiation set samples.
    """

    def __init__(self, base_data_dir):

        base_data_dir = make_path(base_data_dir, 'init_set_data')
        make_dir(base_data_dir)

        self._output_files = dict()
        self.dir = base_data_dir

    def close(self):
        for k, f in self._output_files.items():
            f.close()


    def log_sample(self,
                   episode: int,
                   state: np.ndarray,
                   observation: np.ndarray,
                   option,
                   object_id,
                   can_exec: bool):
        """
        Log a single state-option sample to file.
        :param observation:
        :param episode:
        :param state: the state under consideration, represented as a feature vector
        :param option: the number of the option under consideration
        :param can_exec: whether the option could be executed in the given state
        """

        state = ', '.join(map(str, state)).replace("\n", "")
        observation = ', '.join(map(str, observation)).replace("\n", "")
        line = '{};{};{};{};{};{}\n'.format(episode, option, object_id, state, observation, can_exec)

        if option not in self._output_files:
            f_name = make_path(self.dir, 'option_' + str(option) + '.dat')
            fl = open(f_name, 'w')
            self._output_files[option] = fl

        self._output_files[option].write(line)  # option is a tuple ; option[0] is the act


class PreconditionLogger2:
    """
    A class for logging precondition/initiation set samples.
    """

    def __init__(self, base_data_dir):
        base_data_dir = make_path(base_data_dir, 'init_set_data')
        make_dir(base_data_dir)
        self.dir = base_data_dir
        self._output = defaultdict(list)

    def close(self):
        for option, data in self._output.items():
            save(data, filename=make_path(self.dir, 'option_' + str(option) + '.dat'))

    def log_sample(self,
                   episode: int,
                   state: np.ndarray,
                   observation: np.ndarray,
                   option,
                   object_id,
                   can_exec: bool):
        """
        Log a single state-option sample to file.
        :param observation:
        :param episode:
        :param state: the state under consideration, represented as a feature vector
        :param option: the number of the option under consideration
        :param can_exec: whether the option could be executed in the given state
        """
        self._output[option].append((episode, option, object_id, state, observation, can_exec))

class PreconditionLoggerZ:
    """
    A class for logging precondition/initiation set samples.
    """

    def __init__(self, base_data_dir):

        base_data_dir = make_path(base_data_dir, 'init_set_data')
        make_dir(base_data_dir)

        self._output_files = dict()
        self.dir = base_data_dir

    def close(self):
        for option, f in self._output_files.items():
            f.close()
            path = make_path(self.dir, 'option_' + str(option) + '.dat')
            name = make_path(self.dir, 'option_' + str(option) + '.zip')
            z = zipfile.ZipFile(name, "w")
            z.write(path, arcname='option_' + str(option) + '.dat')
            z.close()
            os.remove(path)

    def log_sample(self,
                   episode: int,
                   state: np.ndarray,
                   observation: np.ndarray,
                   option,
                   object_id,
                   can_exec: bool):
        """
        Log a single state-option sample to file.
        :param observation:
        :param episode:
        :param state: the state under consideration, represented as a feature vector
        :param option: the number of the option under consideration
        :param can_exec: whether the option could be executed in the given state
        """

        state = ', '.join(map(str, state)).replace("\n", "")
        observation = ', '.join(map(str, observation)).replace("\n", "")
        line = '{};{};{};{};{};{}\n'.format(episode, option, object_id, state, observation, can_exec)

        if option not in self._output_files:
            f_name = make_path(self.dir, 'option_' + str(option) + '.dat')
            fl = open(f_name, 'w')
            self._output_files[option] = fl

        self._output_files[option].write(line)  # option is a tuple ; option[0] is the act