import numpy as np

'''
class StaticFns:

    @staticmethod
    def termination_fn(obs, act, next_obs):
        assert len(obs.shape) == len(next_obs.shape) == len(act.shape) == 2

        x = next_obs[:, 2]
        not_done = np.isfinite(next_obs).all(axis=-1) \
                   * (x >= 0.2) \
                   * (x <= 1.0)

        done = ~not_done
        done = done[:, None]
        return done
'''

class StaticFns:

    @staticmethod
    def termination_fn(obs, act, next_obs, env):
        assert len(obs.shape) == len(next_obs.shape) == len(act.shape) == 2

        done = np.linalg.norm(next_obs[:, :2] - env.target_goal, axis=1) <= 0.5
        done = done[:, None]

        return done
