from pettingzoo import AECEnv
from pettingzoo.utils.agent_selector import agent_selector
from gym import spaces
import rlcard
import random
import numpy as np
from pettingzoo.utils import wrappers


def env(**kwargs):
    env = raw_env(**kwargs)
    env = wrappers.CaptureStdoutWrapper(env)
    env = wrappers.TerminateIllegalWrapper(env, illegal_reward=-1)
    env = wrappers.AssertOutOfBoundsWrapper(env)
    env = wrappers.OrderEnforcingWrapper(env)
    return env


class raw_env(AECEnv):

    metadata = {'render.modes': ['human'], "name": "liars_dice_v0"}

    def __init__(self, num_players=4, num_dice=6, num_sides=6, common_hand=False, opponents_hand_visible=False):
        super().__init__()
        self.opponents_hand_visible = opponents_hand_visible
        self.num_dice = num_dice
        self.num_sides = num_sides
        self.common_hand = common_hand
        self.num_players = num_players
        self.num_locations = num_players if not common_hand else num_players + 1 
        self.previous_agent = None

        agents = []
        for i in range(num_players):
            agents.append("player_" + str(i))
        self.agents = agents
        self.possible_agents = agents

        self.num_bets = num_dice * num_sides * num_players if not common_hand else num_dice * num_sides * (1 + num_players)
        self.num_hands_visible = 1 + common_hand if not opponents_hand_visible else 1 + num_players + common_hand

        self.act_shape = act_shape = 1 + self.num_bets
        self.obs_shape = obs_shape = (num_dice * num_sides * self.num_hands_visible) + act_shape # previous bet is in observation 

        self.action_spaces = {name: spaces.Discrete(1 + self.num_bets) for name in self.agents}
        self.observation_spaces = {name: spaces.Dict({
            'observation': spaces.Box(low=0, high=1, shape=(obs_shape,), dtype=np.bool),
            'action_mask': spaces.Box(low=0, high=1, shape=(act_shape,), dtype=np.bool)
        }) for name in self.agents}

        self.board = np.zeros((num_players, num_dice, num_sides)) if not common_hand else np.zeros((num_players + 1, num_dice, num_sides))
        self.agent_selection = self.agents[np.random.choice(num_players)]
        self.current_bid = None
        self.rewards = {name: 0 for name in self.agents}
        self.dones = {name: True for name in self.agents}
        self.infos = {name: {} for name in self.agents}
        self.obs_n = None
        self._cumulative_rewards = np.zeros(num_players) # must be here for petting zoo 
        self.reward_range = (0,1)

        possible_bids = []
        for i in range(num_sides):
            hand_size = (num_players + common_hand) * num_dice
            while hand_size > 0:
                possible_bid = []
                for j in range(hand_size):
                    possible_bid.append(i)
                possible_bids.append(possible_bid)
                hand_size -= 1
        
        self.possible_bids = list(sorted(possible_bids, key = len)) # sorted from worst to best hand for simplicity 

    def reset(self):
        agent_index = np.random.choice(self.num_players)
        self.agent_selection = self.agents[agent_index]
        self.current_bid = None

        self.board = []
        for _ in range(self.num_players + self.common_hand):
            hand = []
            for _ in range(self.num_dice):
                die_roll = np.random.choice(self.num_sides)
                blank_die = np.zeros(self.num_sides, dtype=int)
                blank_die[die_roll] = 1
                hand.append(blank_die)
            self.board.append(hand)

        self.board = np.array(self.board)
        obs = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])

        self.previous_agent = None 
        self.action_mask = np.ones(len(self.possible_bids))
        self.action_mask[0] = 0
        self.rewards = {name: 0 for name in self.agents}

        obs = self.observe(self.agent_selection)

        return obs, self.agent_selection
    
    def current_bid_true(self):
        bid_value = self.current_bid[0]
        bid_size = len(self.current_bid)

        die_counts = np.sum(np.sum(self.board, axis=1), axis=0)
        if(die_counts[bid_value] < bid_size):
            return False 
        else:
            return True 

    def step(self, action):
        done = False
        agent_index = self.agents.index(self.agent_selection)
        if(action == 0): # Bluff has been called 
            # only action that ends the game 
            self.dones = {name: True for name in self.agents}
            # Check if previous bid was true 
            self.rewards = {name: 0 for name in self.agents}
            if(self.current_bid_true()):
                self.rewards[self.previous_agent] = 1
                self.rewards[self.agent_selection] = -1
            else:
                self.rewards[self.agent_selection] = 1
                self.rewards[self.previous_agent] = -1
            done = True
        else:
            self.current_bid = self.possible_bids[action - 1]
            self.dones = {name: False for name in self.agents}
            self.rewards = {name: 0 for name in self.agents}

        # next obs for agent that just performed action 
        #self.obs_n = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])
        self.obs_n = [self.observe(agent) for agent in self.agents] if not done else [None for _ in self.agents]
        
        self.previous_agent = self.agent_selection
        agent_index += 1
        if(agent_index >= len(self.agents)):
            agent_index = 0
        self.agent_selection = self.agents[agent_index]

    def observe(self, agent):
        agent_index = self.agents.index(agent)
        observation = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])
        action_mask = np.ones(len(self.possible_bids))
        if(self.previous_agent is None):
            action_mask[0] = 0 # can't challange first bet 

        if (len(observation.shape)) < 3:
            observation = np.expand_dims(observation, axis=2)  

        if(self.current_bid is not None):
            bid_index = self.possible_bids.index(self.current_bid)
            action_mask[1:bid_index+2] = 0 # First action is challange

        bid_ohe = np.zeros(self.num_bets)
        if(self.current_bid is not None):
            current_bid_index = self.possible_bids.index(self.current_bid)
            bid_ohe[current_bid_index] = 1
        observation = observation.flatten()
        observation = np.append(observation, bid_ohe)

        return {'observation': observation, 'action_mask': action_mask}

    def render(self, mode='human'):
        print("no rendering")
        

from pettingzoo import AECEnv
from pettingzoo.utils.agent_selector import agent_selector
from gym import spaces
import rlcard
import random
import numpy as np
from pettingzoo.utils import wrappers


def env(**kwargs):
    env = raw_env(**kwargs)
    env = wrappers.CaptureStdoutWrapper(env)
    env = wrappers.TerminateIllegalWrapper(env, illegal_reward=-1)
    env = wrappers.AssertOutOfBoundsWrapper(env)
    env = wrappers.OrderEnforcingWrapper(env)
    return env


class raw_env(AECEnv):

    metadata = {'render.modes': ['human'], "name": "liars_dice_v0"}

    def __init__(self, num_players=4, num_dice=6, num_sides=6, common_hand=False, opponents_hand_visible=False):
        super().__init__()
        self.opponents_hand_visible = opponents_hand_visible
        self.num_dice = num_dice
        self.num_sides = num_sides
        self.common_hand = common_hand
        self.num_players = num_players
        self.num_locations = num_players if not common_hand else num_players + 1 
        self.previous_agent = None

        agents = []
        for i in range(num_players):
            agents.append("player_" + str(i))
        self.agents = agents
        self.possible_agents = agents

        self.num_bets = num_dice * num_sides * num_players if not common_hand else num_dice * num_sides * (1 + num_players)
        self.num_hands_visible = 1 + common_hand if not opponents_hand_visible else 1 + num_players + common_hand

        self.act_shape = act_shape = 1 + self.num_bets
        self.obs_shape = obs_shape = (num_dice * num_sides * self.num_hands_visible) + act_shape # previous bet is in observation 

        self.action_spaces = {name: spaces.Discrete(1 + self.num_bets) for name in self.agents}
        self.observation_spaces = {name: spaces.Dict({
            'observation': spaces.Box(low=0, high=1, shape=(obs_shape,), dtype=np.bool),
            'action_mask': spaces.Box(low=0, high=1, shape=(act_shape,), dtype=np.bool)
        }) for name in self.agents}

        self.board = np.zeros((num_players, num_dice, num_sides)) if not common_hand else np.zeros((num_players + 1, num_dice, num_sides))
        self.agent_selection = self.agents[np.random.choice(num_players)]
        self.current_bid = None
        self.rewards = {name: 0 for name in self.agents}
        self.dones = {name: True for name in self.agents}
        self.infos = {name: {} for name in self.agents}
        self.obs_n = None
        self._cumulative_rewards = np.zeros(num_players) # must be here for petting zoo 
        self.reward_range = (0,1)

        possible_bids = []
        for i in range(num_sides):
            hand_size = (num_players + common_hand) * num_dice
            while hand_size > 0:
                possible_bid = []
                for j in range(hand_size):
                    possible_bid.append(i)
                possible_bids.append(possible_bid)
                hand_size -= 1
        
        self.possible_bids = list(sorted(possible_bids, key = len)) # sorted from worst to best hand for simplicity 

    def reset(self):
        agent_index = np.random.choice(self.num_players)
        self.agent_selection = self.agents[agent_index]
        self.current_bid = None

        self.board = []
        for _ in range(self.num_players + self.common_hand):
            hand = []
            for _ in range(self.num_dice):
                die_roll = np.random.choice(self.num_sides)
                blank_die = np.zeros(self.num_sides, dtype=int)
                blank_die[die_roll] = 1
                hand.append(blank_die)
            self.board.append(hand)

        self.board = np.array(self.board)
        obs = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])

        self.previous_agent = None 
        self.action_mask = np.ones(len(self.possible_bids))
        self.action_mask[0] = 0
        self.rewards = {name: 0 for name in self.agents}

        obs = self.observe(self.agent_selection)

        return obs, self.agent_selection
    
    def current_bid_true(self):
        bid_value = self.current_bid[0]
        bid_size = len(self.current_bid)

        die_counts = np.sum(np.sum(self.board, axis=1), axis=0)
        if(die_counts[bid_value] < bid_size):
            return False 
        else:
            return True 

    def step(self, action):
        done = False
        agent_index = self.agents.index(self.agent_selection)
        if(action == 0): # Bluff has been called 
            # only action that ends the game 
            self.dones = {name: True for name in self.agents}
            # Check if previous bid was true 
            self.rewards = {name: 0 for name in self.agents}
            if(self.current_bid_true()):
                self.rewards[self.previous_agent] = 1
                self.rewards[self.agent_selection] = -1
            else:
                self.rewards[self.agent_selection] = 1
                self.rewards[self.previous_agent] = -1
            done = True
        else:
            self.current_bid = self.possible_bids[action - 1]
            self.dones = {name: False for name in self.agents}
            self.rewards = {name: 0 for name in self.agents}

        # next obs for agent that just performed action 
        #self.obs_n = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])
        self.obs_n = [self.observe(agent) for agent in self.agents] if not done else [None for _ in self.agents]
        
        self.previous_agent = self.agent_selection
        agent_index += 1
        if(agent_index >= len(self.agents)):
            agent_index = 0
        self.agent_selection = self.agents[agent_index]

    def observe(self, agent):
        agent_index = self.agents.index(agent)
        observation = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:])
        action_mask = np.ones(len(self.possible_bids))
        if(self.previous_agent is None):
            action_mask[0] = 0 # can't challange first bet 

        if (len(observation.shape)) < 3:
            observation = np.expand_dims(observation, axis=2)  

        if(self.current_bid is not None):
            bid_index = self.possible_bids.index(self.current_bid)
            action_mask[1:bid_index+2] = 0 # First action is challange

        bid_ohe = np.zeros(self.num_bets)
        if(self.current_bid is not None):
            current_bid_index = self.possible_bids.index(self.current_bid)
            bid_ohe[current_bid_index] = 1
        observation = observation.flatten()
        observation = np.append(observation, bid_ohe)

        return {'observation': observation, 'action_mask': action_mask}

    def render(self, mode='human'):
        print("no rendering")
        

# previous bet is in observation # must be here for petting zoo # sorted from worst to best hand for simplicity # Bluff has been called # only action that ends the game # Check if previous bid was true # next obs for agent that just performed action #self.obs_n = self.board[agent_index,:] if not self.common_hand else np.concatenate(self.board[1,:], self.board[agent_index,:]) # can't challange first bet # First action is challange