{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp rl.envs.recsys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RecSys RL Env\n",
    "> OpenAI Gym's Box environment that simulates a recommendation system by picking item and giving feedback signals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import os\n",
    "import gym\n",
    "import numpy as np\n",
    "from sklearn.decomposition import NMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Env(gym.Env):\n",
    "    \"\"\"\n",
    "    Environment for the recommender system\n",
    "    https://github.com/openai/gym/blob/master/gym/core.py\n",
    "    \"\"\"\n",
    "    metadata = {'render.modes': ['human']}\n",
    "    reward_range = (0, 1)\n",
    "\n",
    "    def __init__(self, config, rating_matrix, dataset_name: str):\n",
    "        \"\"\"\n",
    "        Initialize Env\n",
    "        :param config: configurations\n",
    "        :param rating_matrix: rating matrix\n",
    "        :param dataset_name: dataset name\n",
    "        \"\"\"\n",
    "        assert dataset_name in ['train', 'val', 'test']\n",
    "        self.config = config\n",
    "        self.action_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(config.action_size,))\n",
    "        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(config.state_size,))\n",
    "\n",
    "        self.rating_matrix = rating_matrix\n",
    "        rating_matrix_coo = rating_matrix.tocoo()\n",
    "        rating_matrix_rows = rating_matrix_coo.row\n",
    "        rating_matrix_columns = rating_matrix_coo.col\n",
    "        self.rating_matrix_index_set = set(zip(*(rating_matrix_rows, rating_matrix_columns)))\n",
    "        self.env_name = 'env_' + dataset_name + '_' + str(self.config.env_n_components) + '.npy'\n",
    "        self.env_path = os.path.join(config.saves_folder_path, self.env_name)\n",
    "\n",
    "        self.rating_matrix_pred = None\n",
    "        self.load_env()\n",
    "\n",
    "        self.state = None\n",
    "        self.reset()\n",
    "\n",
    "    def load_env(self):\n",
    "        \"\"\"\n",
    "        Load environment\n",
    "        \"\"\"\n",
    "        if not os.path.exists(self.env_path):\n",
    "            env_model = NMF(n_components=self.config.env_n_components, init='random', tol=self.config.env_tol,\n",
    "                            max_iter=self.config.env_max_iter, alpha_H=self.config.env_alpha, \n",
    "                            alpha_W=self.config.env_alpha, verbose=True, random_state=0)\n",
    "            print('-' * 50)\n",
    "            print('Train environment:')\n",
    "            W = env_model.fit_transform(X=self.rating_matrix)\n",
    "            H = env_model.components_\n",
    "            self.rating_matrix_pred = W @ H\n",
    "            print('-' * 50)\n",
    "            np.save(self.env_path, self.rating_matrix_pred)\n",
    "            print('Save environment:', self.env_path)\n",
    "        else:\n",
    "            self.rating_matrix_pred = np.load(self.env_path)\n",
    "            print('Load environment:', self.env_path)\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Reset the environment\n",
    "        :return: state\n",
    "        \"\"\"\n",
    "        while True:\n",
    "            group_id = np.random.choice(range(1, self.config.total_group_num + 1))\n",
    "            nonzero_row, nonzero_col = self.rating_matrix[group_id, :].nonzero()\n",
    "            if len(nonzero_col) >= self.config.history_length:\n",
    "                break\n",
    "        history = np.random.choice(nonzero_col, size=self.config.history_length, replace=False).tolist()\n",
    "        self.state = [group_id] + history\n",
    "        return self.state\n",
    "\n",
    "    def step(self, action: int):\n",
    "        \"\"\"\n",
    "        Take one action to the environment\n",
    "        :param action: action\n",
    "        :return: new_state, reward, done, info\n",
    "        \"\"\"\n",
    "        group_id = self.state[0]\n",
    "        history = self.state[1:]\n",
    "\n",
    "        if (group_id, action) in self.rating_matrix_index_set:\n",
    "            reward = self.rating_matrix[group_id, action]\n",
    "        else:\n",
    "            reward_probability = self.rating_matrix_pred[group_id, action]\n",
    "            reward = np.random.choice(self.config.rewards, p=[1 - reward_probability, reward_probability])\n",
    "\n",
    "        if reward > 0:\n",
    "            history = history[1:] + [action]\n",
    "\n",
    "        new_state = [group_id] + history\n",
    "        self.state = new_state\n",
    "        done = False\n",
    "        info = {}\n",
    "\n",
    "        return new_state, reward, done, info\n",
    "\n",
    "    def render(self, mode='human'):\n",
    "        \"\"\"\n",
    "        Render the environment\n",
    "        :param mode: mode\n",
    "        \"\"\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config(object):\n",
    "    action_size = 1\n",
    "    state_size = 6\n",
    "    env_n_components = 32\n",
    "    saves_folder_path = './'\n",
    "    env_tol = 1e-4\n",
    "    env_max_iter = 1000\n",
    "    env_alpha = 0.001\n",
    "    total_group_num = 3\n",
    "    history_length = 2\n",
    "    rewards = [0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load environment: ./env_val_32.npy\n"
     ]
    }
   ],
   "source": [
    "from scipy.sparse import coo_matrix\n",
    "\n",
    "group_ids = np.array([1,1,1,2,2,3])\n",
    "item_ids = np.array([1,2,3,2,3,1])\n",
    "ratings = np.array([1,0,0,1,1,0])\n",
    "\n",
    "rating_matrix = coo_matrix((ratings, (group_ids, item_ids)), shape=(4, 4)).tocsr()\n",
    "\n",
    "env = Env(config=config, rating_matrix=rating_matrix, dataset_name='val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 2, 3]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state = env.reset()\n",
    "state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([2, 2, 3], 0)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_state, reward, _, _ = env.step(0)\n",
    "new_state, reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-19 13:11:10\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "numpy  : 1.19.5\n",
      "IPython: 5.5.0\n",
      "gym    : 0.17.3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "!pip install -q watermark\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
