# refer to MoutainCar-v0
# Game: https://zh.stardewvalleywiki.com/%E9%B1%BC
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import random

class FishingEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 30
    }
    def __init__(self):
        self.min_position = 0
        self.max_position = 5
        self.max_speed = 0.2
        self.force = 0.02
        self.gravity = 0.01

        self.garb_size = 1.0
        self.catch = True
        self.keep_time = 50
        self.success_time = 400

        self.max_fish_position = self.max_position + self.garb_size - 0.1
        self.min_fish_position = self.min_position + 0.1
        self.fish_position = 0.5
        self.fish_max_speed = 0.1 #
        self.fish_max_acc = 0.01 #

        self.viewer = None

        self.action_space = spaces.Discrete(2)
        self.low = np.array([self.min_position, -self.max_speed, self.min_fish_position, -self.fish_max_speed])
        self.high = np.array([self.max_position, self.max_speed, self.max_fish_position, self.fish_max_speed])
        self.observation_space = spaces.Box(self.low, self.high, dtype=np.float32)

        self.seed()
        
    
    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]
    
    def step(self, action):
        assert self.action_space.contains(action), "%r (%s) invalid" % (action, type(action))

        pp, vp, pf, vf = self.state
        # pod
        vp += (action*self.force - self.gravity)
        vp = np.clip(vp, -self.max_speed, self.max_speed)
        pp += vp
        if pp > self.max_position:
            pp = self.max_position
            vp = 0
        elif pp < self.min_position:
            pp = self.min_position
            if vp < -0.05:
                vp = -vp*0.8
            else: vp = 0
        # fish
        af = (float)(random.randint(-10,10))* self.fish_max_acc/10.0 
        vp = np.clip(vp, -self.fish_max_speed, self.fish_max_speed)
        vf += af
        pf += vf
        if pf < self.min_position or pf > self.max_fish_position:
            vf = 0
        pf = np.clip(pf, self.min_position, self.max_fish_position)
        # catch?
        if pf > pp and pf < pp + self.garb_size:
            self.catch = True
        else: self.catch = False
        reward = -1.0
        if self.catch: 
            self.keep_time += 1
            reward = 1.0
        else : 
            self.keep_time -= 1
        done = bool(self.keep_time >= self.success_time or self.keep_time <= 0)
        
        self.state = (pp,vp,pf,vf)
        return np.array(self.state), reward, done, {}
    
    def reset(self):
        self.state = np.array([0,0,0.5,0])
        self.keep_time = 20
        self.catch = True
        return self.state
    
    def render(self, mode='human'):
        screen_width = 200
        screen_height = 400

        fish_height = 0.1
        world_height = self.max_fish_position - self.min_position + fish_height # fish length for drawing
        scale = screen_height/world_height
        garb_width = scale*self.garb_size/5

        if self.viewer is None:
            from gym.envs.classic_control import rendering
            self.viewer = rendering.Viewer(screen_width, screen_height)
            
            clearance = 0
            water = rendering.FilledPolygon(
                [(screen_width/2-garb_width/2,0),
                 (screen_width/2+garb_width/2,0),
                 (screen_width/2+garb_width/2,screen_height),
                 (screen_width/2-garb_width/2,screen_height)])
            water.set_color(0,0,.8)
            self.viewer.add_geom(water)

            l, r, t, b = -garb_width+2, garb_width-2, self.garb_size*scale, 0
            garb = rendering.FilledPolygon([(l,b), (l,t), (r,t), (r,b)])
            garb.add_attr(rendering.Transform(translation = (screen_width/2, clearance)))
            self.Cg = rendering.Transform()
            garb.add_attr(self.Cg)
            garb.set_color(0,.8,0)
            self.viewer.add_geom(garb)
 
            fish = rendering.FilledPolygon([
                (-10,0),(-1,-10),(-5,-15),(5,-15),(1,-10),(10,0),(0,10)
            ])
            fish.set_color(.8,0,0)
            fish.add_attr(rendering.Transform(translation = (screen_width/2, clearance)))
            self.Cf = rendering.Transform()
            fish.add_attr(self.Cf)
            self.viewer.add_geom(fish)

            ll, rr, tt, bb = 2, 32, 10, 0
            self.current = rendering.FilledPolygon([(ll,bb), (ll,tt), (rr,tt), (rr,bb)])
            self.current.add_attr(rendering.Transform(translation = (screen_width/2 + garb_width, clearance)))
            self.Cc = rendering.Transform()
            self.current.add_attr(self.Cc)
            self.viewer.add_geom(self.current)
            
        
        self.Cg.set_translation(0,self.state[0]*scale)
        self.Cf.set_translation(0,self.state[2]*scale)
        p = self.keep_time/self.success_time
        self.current.set_color(1-p,p,0)
        self.Cc.set_translation(0,p*self.max_position*scale)
        
        return self.viewer.render(return_rgb_array = mode=='rgb_array')

    def get_keys_to_action(self):
        return {():0,(ord(" "),):1} #control with left and right arrow keys 
    
    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None
