from random import random, randint
from environment import Environment
from constants import ACTION_FORWARD, ACTION_TURN_LEFT
from constants import ACTION_PROBE_GHOST, ACTION_BUST_GHOST
from constants import ACTION_TURN_RIGHT, ACTION_NOP, ACTION_STOP
from constants import CLEAN, WALL, RED, ORANGE, YELLOW, GREEN, BUSTED
from constants import EAST

class Percept:
    def __init__(self, attributes):
        self.attributes = attributes


class VacuumEnvironment(Environment):

    """
    Create a vacuum environment with the given width, height, world-gen element biases and PRF seed
    """
    def __init__(self, env_x=5, env_y=5, dirt_bias=0.1, wall_bias=0.0, gold_bias=0.0, world_seed=None):
        super().__init__()
        self.env_x = env_x
        self.env_y = env_y
        self.world = None
        self.ghost_location = (randint(1, self.env_x-2), randint(1, self.env_y-2))
        print(f"Ghost placed at {self.ghost_location[1]}, {self.ghost_location[0]}")
        self.probe_percept = None
        self.bust_percept = None
        self.wallify()

    """
    Add thing to environment
    """
    def add_thing(self, thing, location=None):
        thing.facing = (1, 0)
        super().add_thing(thing, location)

    def percept(self, agent):
        percepts = {"bump": agent.bump}
        if self.probe_percept:
            percepts["probe"] = self.probe_percept
            self.probe_percept = None
        if self.bust_percept:
            percepts["bust"] = self.bust_percept == 'yes'
            self.bust_percept = None
        return Percept(percepts)

  
    """
    Process actions generated by agents in environment
    """
    def execute_action(self, agent, action):
        agent.bump = False
        if action == ACTION_FORWARD:
            new_location = (agent.location[0] + agent.facing[0], agent.location[1] + agent.facing[1])
            agent.bump = self.world[new_location[0]][new_location[1]] == WALL
            agent.location = agent.location if agent.bump else new_location
        elif action == ACTION_TURN_LEFT:
            """
            NORTH -> WEST   |  ( 0, -1) -> (-1,  0)
            EAST  -> NORTH  |  ( 1,  0) -> ( 0, -1)
            SOUTH -> EAST   |  ( 0,  1) -> ( 1,  0)
            WEST  -> SOUTH  |  (-1,  0) -> ( 0,  1)
            """
            agent.facing = (agent.facing[1], -agent.facing[0] if agent.facing[0] != 0 else agent.facing[0])
        elif action == ACTION_TURN_RIGHT:
            agent.facing = (-agent.facing[1] if agent.facing[1] != 0 else agent.facing[1], agent.facing[0])        
        elif action == ACTION_PROBE_GHOST:
            self.probe_percept = self.make_probe_percept((agent.location[0], agent.location[1]), self.ghost_location)
            self.world[agent.location[0]][agent.location[1]] = self.probe_percept
        elif action == ACTION_BUST_GHOST:
            aloc = (agent.location[0], agent.location[1])
            gloc = self.ghost_location
            busted = 'yes' if aloc == gloc else "no"
            if busted == 'yes':
                agent.add_bust_reward()
                self.world[agent.location[0]][agent.location[1]] = BUSTED
            self.bust_percept = busted
        elif action == ACTION_NOP:
            pass
        elif action == ACTION_STOP:
            agent.is_alive = False
        else:
            raise(Exception(f"Bad action {action}"))
            
        return True
    
    def coord_distance(self, c1, c2):
        row1, col1 = c1
        row2, col2 = c2
        return abs(row1-row2) + abs(col1-col2)
    
    def make_probe_percept(self, agent_loc, ghost_loc):
        distance = self.coord_distance(agent_loc, ghost_loc)
        distance = 7 if distance > 7 else distance
        #print(f"Probe percept, {agent_loc} {ghost_loc} effective distance is {distance}")
        probe_probs = \
            {0:	[.94, .97, .99, 1.00],
             1: [.15, .95, .97, 1.00], 
             2: [.05, .90, .96, 1.00],
             3: [.03, .10, .90, 1.00],
             4: [.01, .10, .85, 1.00],
             5: [.01, .05, .15, 1.00], 
             6: [.00, .01, .03, 1.00],
             7: [.00, .01, .02, 1.00]}
        probs = probe_probs[distance]
        probe = None
        r = random()
        #print(f"Probs is {probs} random is {r}")
        if r <= probs[0]:
            probe = RED
        elif r <= probs[1]:
            probe = ORANGE
        elif r <= probs[2]:
            probe = YELLOW
        else:
            probe = GREEN
        #print(f"Returning {probe}")
        return probe

    """
    Random-generate an environment for the vacuum with an optional seed
    """  
    def wallify(self):
        self.world =  [
                [
                WALL if
                    x == 0 or
                    x == self.env_x - 1 or
                    y == 0 or
                    y == self.env_y - 1
                else CLEAN
                for y in range(self.env_y)
            ]
            for x in range(self.env_x)
        ]
                
    # CAUTION!
    #   The position argument is possibly breaking since it's just 
    #   the default place where Things are placed (see defn of default_location)
    
    def prep_agent(self, agent, recon_type):
        if not recon_type in ['None', 'WallsOnly', 'Full', 'WallsAndGold']:
            raise(Exception(f"Bad recon type {recon_type}"))
        recon = {'width': self.env_x,
                 'height': self.env_y,
                 'position': (1,1),
                 'heading': EAST}
        if recon_type == "WallsOnly":
            recon['walls'] = self.env_positions(WALL)
        elif recon_type == "WallsAndGold":
            recon['walls'] =  self.env_positions(WALL)
            recon['gold']  = self.env_positions(GOLD)
        elif recon_type == "Full":
            recon['walls'] =  self.env_positions(WALL)
            recon['gold']  = self.env_positions(GOLD)
            recon['dirt']  =  self.env_positions(DIRT)
        agent.prep(recon)
    
    def env_positions(self, env):
        return [(r,c) for r in range(self.env_y) for c in range(self.env_x) if self.world[c][r] == env]

        
        
   
