from Palette.algos.charge import Charge


class MultiCharge(object):
    def __init__(self, n_ships: int, filling_time: int):
        super().__init__()
        self._n_ships = n_ships
        self._filling_time = filling_time
        self._charges = [Charge(self._filling_time)
                         for _ in range(self._n_ships)]

    def __len__(self):
        return self._n_ships

    def __getitem__(self, idx):
        return self._charges[idx]

    def step(self, state, info, idx):
        agent_info = {k: info[k]
                      for k in (info.keys()-{'fields', 'redecide_direction'})}
        agent_info['fields'] = info['fields'][idx]
        agent_info['redecide_direction'] = info['redecide_direction'][idx]
        return self._charges[idx].step(state[idx], agent_info)