import gymnasium as gym
from gymnasium import spaces
from gymnasium import utils
import numpy as np
import matplotlib.pyplot as plt

import logging
logger = logging.getLogger(__name__)


class InventoryEnv(gym.Env, utils.EzPickle):
    """Inventory control with lost sales environment
    TO BE EDITED
    This environment corresponds to the version of the inventory control
    with lost sales problem described in Example 1.1 in Algorithms for
    Reinforcement Learning by Csaba Szepesvari (2010).
    https://sites.ualberta.ca/~szepesva/RLBook.html
    """
    def __init__(self, n=100, k=5, c=2, h=2, p=3, lam=8):
        self.n = n
        self.action_space = spaces.Discrete(n)
        self.observation_space = spaces.Discrete(n)
        self.max = n
        self.state = n
        self.k = k
        self.c = c
        self.h = h
        self.p = p
        self.lam = lam
    def demand(self):
        return np.random.poisson(self.lam)
    def plot_demand(self):
        data = np.random.poisson(self.lam, 1000)
        # Plot the histogram
        plt.hist(data, bins=range(data.min(), data.max()+2), align='left')
        plt.xlabel('Value')
        plt.ylabel('Frequency')
        plt.title('Histogram of Demand Distribution')
        plt.show()
    def transition(self, x, a, d):
        m = self.max
        return max(min(x + a, m) - d, 0)
    def reward(self, x, a, y):
        k = self.k
        m = self.max
        c = self.c
        h = self.h
        p = self.p
        r = -k * (a > 0) - c * max(min(x + a, m) - x, 0) - h * x + p * max(min(x + a, m) - y, 0)
        return r
    def step(self, action):
        assert self.action_space.contains(action)
        obs = self.state
        demand = self.demand()
        obs2 = self.transition(obs, action, demand)
        self.state = obs2
        reward = self.reward(obs, action, obs2)
        terminated = True
        info = {}
        truncated = False
        return obs2, reward, terminated, truncated, info
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        observation = self.observation_space.sample()
        self.state = observation
        info = {}
        #print(observation)
        return self.state,info  # reward, done, info can't be included
if __name__ == '__main__':
    a=InventoryEnv()
    a.plot_demand()