import random
import numpy as np
from scipy.ndimage import rotate
from scipy.misc import imresize

class CatchEnv4:
    def __init__(self):
        self.size = 21
        self.image = np.zeros((self.size, self.size))
        self.state = []
        self.fps = 4
        self.output_shape = (84, 84)

    def reset_random(self):
        self.image.fill(0)
        self.pos = np.random.randint(2, self.size-2)
        self.vx = np.random.randint(5) - 2
        self.vy = 1
        self.ballx, self.bally = np.random.randint(self.size), 4
            
        self.image[self.bally, self.ballx] = 1
        self.image[-5, self.pos - 1:self.pos+1] = np.ones(1)
 
        return self.step(2)[0]


    def step(self, action):
        def left():
            if self.pos > 3:
                self.pos -= 2
        def right():
            if self.pos < 17:
                self.pos += 2
        def noop():
            pass
        {0: left, 1: right, 2: noop}[action]()

        
        self.image[self.bally, self.ballx] = 0
        self.ballx += self.vx
        self.bally += self.vy
        if self.ballx > self.size - 1:
            self.ballx -= 2 * (self.ballx - (self.size-1))
            self.vx *= -1
        elif self.ballx < 0:
            self.ballx += 2 * (0 - self.ballx)
            self.vx *= -1
        self.image[self.bally, self.ballx] = 1

        self.image[-5].fill(0)
        self.image[-5, self.pos-1:self.pos+1] = np.ones(1)
        
        terminal = self.bally == self.size - 2 - 4
        reward = int(self.pos - 1 <= self.ballx <= self.pos + 1) if terminal else 0

        [self.state.append(imresize(self.image, (84, 84))) for _ in range(self.fps - len(self.state) + 1)]
        self.state = self.state[-self.fps:]
        
        return np.transpose(self.state, [1, 2, 0]), reward, terminal

    def get_num_actions(self):
        return 3

    def reset(self):
        return self.reset_random()

    def state_shape(self):
        return (self.fps,) + self.output_shape
    
    def show_state(self, i):
        plt.imshow(self.image)
        plt.imsave('image_'+str(i)+'.jpg', self.image)

def test():
    env = CatchEnv4()
    i = 0
    for ep in range(1):
        env.reset()
        state, reward, terminal = env.step(1) 
        while not terminal:
            state, reward, terminal = env.step(np.random.randint(0,2))
            state = np.squeeze(state)
            env.show_state(i)
            i += 1
        print(reward)

if __name__ == "main":
    test()