import gym
import numpy as np

from gym import spaces
from scipy import stats
from scipy.special import expit

class ToyPricing(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 4}

    def __init__(self, render_mode=None, num_states = 5, p = np.array([-0.3, 0.1], dtype = float)):
        self.states = list(range(num_states))
        self.actions = list(range(2))

        self.p = p
        self.window_size = 512  # The size of the PyGame window

        # Observations are dictionaries with the agent's and the target's location.
        # Each location is encoded as an element of {0, ..., `size`}^2, i.e. MultiDiscrete([size, size]).
        self.observation_space = spaces.Box(0, num_states - 1, shape=(1,), dtype=int)

        # We have 2 actions, corresponding "high", "low"
        self.action_space = spaces.Discrete(2)

    def _get_obs(self):
        return self._current_state
    
    def _get_info(self):
        return {"state": self._current_state}

    def _get_next_state(self, state, action):
        s_a = np.concatenate((state.reshape(-1, 1), action.reshape(-1, 1)), axis = 1)
        return np.array([stats.binom.rvs(len(self.states) - 1, expit(s_a @ self.p))], dtype = int)

    def transition_pmf(self, state, action, next_state):
        s_a = np.concatenate((state.reshape(-1, 1), action.reshape(-1, 1)), axis = 1)
        return np.array([stats.binom.pmf(next_state, len(self.states) - 1, expit(s_a @ self.p))], dtype = float)

    def init_pmf(self, state):
        return np.array([stats.binom.pmf(state, len(self.states) - 1, 0)], dtype = float)

    def reset(self, seed=None):
        # We need the following line to seed self.np_random
        super().reset(seed=seed)

        # Always start at first state
        self._current_state = np.array([0], dtype=int)

        observation = self._get_obs()
        info = self._get_info()

        return observation, info

    def step(self, action):
        """will use TimeLimit wrapper for truncation..."""
        observation = self._get_obs()
        self._current_state = self._get_next_state(observation, action)

        # Reward
        reward = (action * self._current_state).item()

        next_observation = self._get_obs()
        info = self._get_info()
        terminated = False

        return next_observation, reward, terminated, False, info

    def render(self):
        pass

    def _render_frame(self):
        pass

    def close(self):
        pass