import os
import random

import numpy as np
from gym_minigrid.minigrid import *
from gym_minigrid.register import register


def open_file(file):
    file_grid = []
    # if os.path.exists("/aubret/These/stalk/tools/resources/" + file):
    #     f = open("/aubret/These/stalk/tools/resources/" + file, "r")
    # else:
    f = open("resources/" + file, "r")
    for l in f:
        wd = []
        for c in l:
            if c == "_":
                wd.append(0)
            else:
                wd.append(1)
        file_grid.append(wd)
    np_file_grid = np.array(file_grid).transpose(1, 0)
    f.close()
    return np_file_grid


class EmptyEnv(MiniGridEnv):
    """
    Empty grid environment, no obstacles, sparse reward
    """

    def __init__(self, random_pos=False,manual_start_pos=None,**kwargs):
        self.random_pos=random_pos
        self.manuel_start_pos=manual_start_pos
        super().__init__(
            # Set this to True for maximum speed
            see_through_walls=True,
            **kwargs
        )

    def _gen_grid(self, width, height,total_sum=None,**kwargs):

        # Create an empty grid
        self.grid = Grid(width, height,total_sum=None,**kwargs)#**kwargs)

        # Generate the surrounding walls
        self.grid.wall_rect(0, 0, width, height)

        # Place the agent in the top-left corner
        if not self.random_pos:
            self.start_pos = (width // 2, height // 2)
            # self.start_pos=(1,1)
            self.start_dir = 0
        else:
            self.place_agent()



        # Place a goal square in the bottom-right corner

        #if self.goal is None :
        ###############The goal definition
        #self.pos_goal = (width-2, height-2)
        #self.grid.set(self.pos_goal[0], self.pos_goal[1], Goal())



        #else :
        #    index = np.argmax(self.goal,axis=0)
        #    self.pos_goal = (width-2,1) if index ==0 else (1,height-2)
        #    self.grid.set(self.pos_goal[0], self.pos_goal[1], Goal())
        """
        while True :
            self.start_pos = (random.randint(1,width-2),random.randint(1,height-2))
            if self.start_pos != self.pos_goal:
                break
        self.start_dir = random.randint(0,3)
        """

        self.mission = "get to the green goal square"

    def _reward(self):

        """
        Compute the reward to be given upon success
        """
        #x,y = self.agent_pos
        #goal_x, goal_y = self.pos_goal

        #reward = -np.sqrt(math.pow(x-goal_x,2)+math.pow(y-goal_y,2)) / self.width
        #return reward

        return 1
        #return 1 - 0.9 * (self.step_count / self.max_steps)

class EmptyEnv6x6(EmptyEnv):
    def __init__(self):
        super().__init__(size=6)

class EmptyEnv16x16(EmptyEnv):
    def __init__(self):
        super().__init__(size=16)

class EmptyWallEnv(EmptyEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _gen_grid(self, width, height,*args,**kwargs):
        super()._gen_grid( width, height,*args,**kwargs)
        self.grid.vert_wall(width//2 +2 , 1,(height//2)+6)
        if DENSE_REWARD and not self.render_stats:
            for i in range(width//2+2+1,width-1):
                for j in range(2,(height//2)+6):
                    self.grid.set(i, j, Goal())

class EmptyWallHorzEnv(EmptyEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _gen_grid(self, width, height,*args,**kwargs):
        super()._gen_grid( width, height,*args,**kwargs)
        self.grid.horz_wall(width//2-6, (height//2), 5)




class EmptyFile(EmptyEnv):
    FILE="hard_env.txt"
    def __init__(self,**kwargs):
        self.change_start_pos=False
        self.file_grid=open_file(self.FILE)
        super().__init__(**kwargs)


    def _gen_grid(self, width, height,*args,**kwargs):
        super()._gen_grid( width, height,*args,**kwargs)
        for i in range(self.width):
            for j in range(self.height):
                if self.file_grid[i,j] == 1:
                    self.grid.set(i, j, Wall())

        # self.place_agent()
        if not self.random_pos:
            if not self.manuel_start_pos:
                self.start_pos = (1, height -2)
                self.start_dir = 0
            else:
                self.start_pos = self.manuel_start_pos
                self.start_dir = 0
        else:
            self.place_agent()


class HardWall(EmptyFile):
    FILE="hard_env.txt"

class Loop(EmptyFile):
    FILE="loop.txt"

class LoopFail(EmptyFile):
    FILE="loop_fail.txt"

class LoopWay(EmptyFile):
    FILE="loop_fail.txt"

    def step(self, action):
        if self.agent_pos[0] == 2 and self.agent_pos[1]==17 and (action == 2):
            self.agent_pos=(26, 18)
            obs=self.gen_obs()
            infos={}
            infos["agent_pos"] = self.agent_pos
            if self.stats is not None:
                infos["stats"] = self.stats
            return obs,0,False,infos
        elif self.agent_pos[0] == 26 and self.agent_pos[1]==18 and (action == 0):
            self.agent_pos=(2, 17)
            obs=self.gen_obs()
            infos={}
            infos["agent_pos"] = self.agent_pos
            if self.stats is not None:
                infos["stats"] = self.stats
            return obs,0,False,infos
        else:
            return super().step(action)




