facility_location/env/pmp.py

import io
import warnings
from typing import Tuple, Dict, Optional, List, Text

import gym
import math
import numpy as np
import matplotlib.pyplot as plt
import pickle, os

from numpy import ndarray

from facility_location.utils.config import Config
from facility_location.env.facility_location_client import FacilityLocationClient
from facility_location.env.obs_extractor import ObsExtractor
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from facility_location.agent import MaskedFacilityLocationActorCriticPolicy
from facility_location.utils.policy import get_policy_kwargs


class PMPEnv(gym.Env):
    EPSILON = 1e-6

    def __init__(self,
                 cfg: Config):
        self.cfg = cfg
        self._train_region = cfg.env_specs['region']
        self._eval_region = cfg.eval_specs['region']
        self._min_n = cfg.env_specs['min_n']
        self._max_n = cfg.env_specs['max_n']
        self._min_p_ratio = cfg.env_specs['min_p_ratio']
        self._max_p_ratio = cfg.env_specs['max_p_ratio']
        self._max_steps_scale = cfg.env_specs['max_steps_scale']
        self._tabu_stable_steps_scale = cfg.env_specs['tabu_stable_steps_scale']
        self._popstar = cfg.env_specs['popstar']

        self._seed(cfg.seed)

        self._done = False

        self._set_node_edge_range()

        self._flc = FacilityLocationClient(cfg, self._np_random)
        self._obs_extractor = ObsExtractor(cfg, self._flc, self._node_range, self._edge_range)

        self._declare_spaces()

    def _declare_spaces(self) -> None:
        self.observation_space = gym.spaces.Dict({
            'node_features': gym.spaces.Box(low=0, high=1, shape=(self._node_range, self.get_node_feature_dim())),
            'static_adjacency_list': gym.spaces.Box(low=0, high=self._node_range, shape=(self._edge_range, 2), dtype=np.int64),
            'dynamic_adjacency_list': gym.spaces.Box(low=0, high=self._node_range, shape=(self._edge_range, 2), dtype=np.int64),
            'node_mask': gym.spaces.Box(low=0, high=1, shape=(self._node_range,), dtype=np.bool),
            'static_edge_mask': gym.spaces.Box(low=0, high=1, shape=(self._edge_range,), dtype=np.bool),
            'dynamic_edge_mask': gym.spaces.Box(low=0, high=1, shape=(self._edge_range,), dtype=np.bool),
        })
        if not self._popstar:
            self.action_space = gym.spaces.Discrete(self._node_range ** 2)
        else:
            self.action_space = gym.spaces.Discrete(self._node_range ** 2)

    def _set_node_edge_range(self) -> None:
        self._node_range = self._max_n + 2
        self._edge_range = int(self._max_n ** 2 * self._max_p_ratio)

    def get_node_feature_dim(self) -> int:
        return self._obs_extractor.get_node_dim()

    def _seed(self, seed: int) -> None:
        self._np_random = np.random.default_rng(seed)

    def get_reward(self) -> float:
        reward = self._obj_value[self._t - 1] - self._obj_value[self._t]
        return reward

    def _transform_action(self, action: np.ndarray) -> np.ndarray:
        if self._popstar:
            action = np.array(np.unravel_index(action, (self._node_range, self._node_range)))
        action = action - 1
        return action

    def step(self, action: np.ndarray):
        if self._done:
            raise RuntimeError('Action taken after episode is done.')
        obj_value, solution, info = self._flc.swap(action, self._t)
        self._t += 1
        self._done = (self._t == self._max_steps)
        self._obj_value[self._t] = obj_value
        self._solution[self._t] = solution
        reward = self.get_reward()
        if obj_value < self._best_obj_value - self.EPSILON:
            self._best_obj_value = obj_value
            self._best_solution = solution
            self._last_best_t = self._t
        elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0:
            self._flc.reset_tabu_time()
    
        # if self._done:
        #     print('done')
        #     for i in range(self._t):
        #         print(f'{i}:',np.where(self._solution[i]))
            
        return self._get_obs(self._t), reward, self._done, False, info

    def reset(self, seed = 0) -> Optional[Dict]:
        if self._train_region is None:
            points, demands, n, p = self._generate_new_instance()
            self._flc.set_instance(points, demands, n, p, False)
        else:
            points, demands, n, p = self._use_real_instance()
            self._flc.set_instance(points, demands, n, p, True)
        
        return self.prepare(n, p), {}

    def prepare(self, n: int, p: int) -> Dict:
        initial_obj_value, initial_solution = self._flc.compute_initial_solution()
        self._obs_extractor.reset()
        self._done = False
        self._t = 0
        self._max_steps = max(int(p * self._max_steps_scale), 5)
        self._obj_value = np.zeros(self._max_steps + 1)
        self._obj_value[0] = initial_obj_value
        self._solution = np.zeros((self._max_steps + 1, n), dtype=bool)
        self._solution[0] = initial_solution
        self._best_solution = initial_solution
        self._best_obj_value = initial_obj_value
        self._last_best_t = 0
        self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale))
        return self._get_obs(self._t)

    def render(self, mode='human', dpi=300) -> Optional[np.ndarray]:
        gdf, facilities = self._flc.get_gdf_facilities()
        if len(facilities) > 10:
            warnings.warn('Too many facilities to render. Only rendering the first 10.')
            facilities = facilities[:10]

        cm = plt.get_cmap('tab10')
        fig, axs = plt.subplots(1, 2, figsize=(12, 6), dpi=dpi)
        for i, f in enumerate(facilities):
            gdf.loc[gdf['assignment'] == f].plot(ax=axs[0],
                                                 zorder=2,
                                                 alpha=0.7,
                                                 edgecolor="k",
                                                 color=cm(i))
            gdf.loc[[f]].plot(ax=axs[0],
                              marker='*',
                              markersize=300,
                              zorder=3,
                              alpha=0.7,
                              edgecolor="k",
                              color=cm(i))
        axs[0].set_title("Facility Location", fontweight="bold")
        plot_obj_value = self._obj_value[:self._t + 1]
        axs[1].plot(plot_obj_value, marker='.', markersize=10, color='k')
        axs[1].set_title("Objective Value", fontweight="bold")
        axs[1].set_xticks(np.arange(self._max_steps + 1, step=math.ceil((self._max_steps + 1) / 10)))
        fig.tight_layout()

        if mode == 'human':
            plt.show()

        else:
            io_buf = io.BytesIO()
            fig.savefig(io_buf, format='raw', dpi=dpi)
            io_buf.seek(0)
            img_arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
                                 newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
            io_buf.close()
            return img_arr

    def close(self):
        plt.close()

    def _generate_new_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]:
        n = self._np_random.integers(self._min_n, self._max_n, endpoint=True)
        p_ratio = self._np_random.uniform(self._min_p_ratio, self._max_p_ratio)
        p = int(max(n * p_ratio, 4))

        points = self._np_random.uniform(size=(n, 2))
        while np.unique(points, axis=0).shape[0] != n:
            points = self._np_random.uniform(size=(n, 2))
        demands = self._np_random.random(size=(n,))
        return points, demands, n, p
    
    def _use_real_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]:
        data_path = './data/{}/pkl'.format(self.cfg.eval_specs['region'])
        files = os.listdir(data_path)
        files = [f for f in files if f.endswith('.pkl')]
        sample_data_path = os.path.join(data_path, files[self._np_random.integers(len(files))])
        with open(sample_data_path, 'rb') as f:
            np_data = pickle.load(f)
        
        n = self._np_random.integers(self._min_n, self._max_n, endpoint=True)
        p = max(int(n * self._np_random.uniform(self._min_p_ratio, self._max_p_ratio)), 4)
        sample_cbgs = self._np_random.choice(list(np_data[1].keys()), n, replace=False)
        points = []
        demands = []
        for cbg in sample_cbgs:
            points.append(np_data[1][cbg]['pos'])
            demands.append(np_data[1][cbg]['demand'])
        points = np.array(points)
        demands = np.array(demands)
        
        return points, demands, n, p

    def  _get_obs(self, t: int) -> Dict:
        return self._obs_extractor.get_obs(t)

    def get_initial_solution(self) -> np.ndarray:
        return self._solution[0]


class EvalPMPEnv(PMPEnv):
    def __init__(self,
                 cfg: Config,
                positions, demands, n, p, boost=False):
        self._eval_np = (n,p)
        self._eval_seed = cfg.eval_specs['seed']
        self._boost = boost
        self.points = positions
        self.demands = demands
        self._n = n
        self._p = p
        
        super().__init__(cfg)

    def _set_node_edge_range(self) -> None:
        n, p = self._eval_np

        self._node_range = n + 2
        self._edge_range = n * p

    def get_eval_num_cases(self) -> int:
        return self._eval_num_cases

    def get_eval_np(self) -> Tuple[int, int]:
        return self._eval_np

    def reset_instance_id(self) -> None:
        self._instance_id = 0

    def step(self, action: np.ndarray):
        if self._done:
            raise RuntimeError('Action taken after episode is done.')
        obj_value, solution, info = self._flc.swap(action, self._t)
        self._t += 1
        self._done = (self._t == self._max_steps)
        self._obj_value[self._t] = obj_value
        self._solution[self._t] = solution
        reward = self.get_reward()
        if obj_value < self._best_obj_value - self.EPSILON:
            self._best_obj_value = obj_value
            self._best_solution = solution
            self._last_best_t = self._t
        elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0:
            self._flc.reset_tabu_time() 
        print(self._t, self._max_steps)
                
        return self._get_obs(self._t), reward, self._done, False, info

    def get_reward(self) -> float:
        if self._done:
            reward = -np.min(self._obj_value)
        else:
            reward = 0.0
        
        return reward
    
    def get_best_solution(self) -> np.ndarray:
        return self._best_solution
    
    def reset(self, seed = 0) -> Dict:
        self._flc.set_instance(self.points, self.demands, self._n, self._p, False)
        return self.prepare(self._n, self._p, self._boost), {}
    
    def prepare(self, n: int, p: int, boost: bool) -> Dict:
        initial_obj_value, initial_solution = self._flc.compute_initial_solution()
        self._obs_extractor.reset()
        self._done = False
        self._t = 0
        self._max_steps = max(int(p * self._max_steps_scale), 5)
        if boost:
            self._max_steps = max(int(self._max_steps_scale / 10), 5)
        self._obj_value = np.zeros(self._max_steps + 1)
        self._obj_value[0] = initial_obj_value
        self._solution = np.zeros((self._max_steps + 1, n), dtype=bool)
        self._solution[0] = initial_solution
        self._best_solution = initial_solution
        self._best_obj_value = initial_obj_value
        self._last_best_t = 0
        self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale))
        return self._get_obs(self._t)

    def get_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]:
        points, demands, n, p = self._flc.get_instance()
        return points, demands, n, p

    def get_distance_and_cost(self) -> Tuple[np.ndarray, np.ndarray]:
        return self._flc.get_distance_and_cost_matrix()

    def evaluate(self, solution: np.ndarray) -> float:
        self._flc.set_solution(solution)
        obj_value = self._flc.compute_obj_value()
        return obj_value

class MULTIPMP(PMPEnv):
    EPSILON = 1e-6
    def __init__(self,
                 cfg,
                 data_npy,
                 boost = False):
        self.cfg = cfg
        self.data_npy =  data_npy
        self._boost = boost
        self._all_points, self._all_demands, self._n, self._all_p = self._load_multi_facility_data(data_npy)
        self.boost = boost
        self._all_solutions = self._load_multi_facility_solutions(boost)
        print('all_solutions:', self._all_solutions)
        self._final_solutions = list(self._all_solutions)
        self._num_types = len(self._all_p)
        self._current_type = 0
        self._all_max_steps, self._old_mask, self._new_mask = self._get_max_steps()
        super().__init__(cfg)

    def _set_node_edge_range(self) -> None:
        self._node_range = self._n + 2
        self._edge_range = self._n * max(self._all_p)

    def step(self, action: np.ndarray):
        if self._num_types == 1:
            reward = self.get_reward()
            self._done = True
            pickle.dump(self._final_solutions, open('./facility_location/solutions.pkl', 'wb'))
            return self._get_obs(self._t), reward, self._done, False, {}
            
        if self._done:
            raise RuntimeError('Action taken after episode is done.')
        obj_value, solution, info = self._flc.swap(action, self._t, stage=2)
        self._t += 1
        
        self._done = (self._t == self._all_max_steps[-1] and self._current_type == len(self._all_max_steps) - 1)
        self._obj_value[self._t] = obj_value
        self._solution[self._t] = solution
        reward = self.get_reward()
        if obj_value < self._best_obj_value - self.EPSILON:
            self._best_obj_value = obj_value
            self._best_solution = solution
            self._last_best_t = self._t
        elif (self._t - self._last_best_t) % self._tabu_stable_steps == 0:
            self._flc.reset_tabu_time()
        
        print(self._t, self._all_max_steps[self._current_type])
        if self._t == self._all_max_steps[self._current_type] and not self._done:
            self._t = 0
            self._multi_obj += obj_value
            self._final_solutions[self._current_type] = solution
            self._update_type()
            
        if self._done:
            pickle.dump(self._final_solutions, open('./facility_location/solutions.pkl', 'wb'))
            
        return self._get_obs(self._t), reward, self._done, False, info

    def reset(self, seed = 0) -> Optional[Dict]:
        self._current_type = 0
        points = self._all_points
        demands = self._all_demands[:,0]
        n = self._n
        p = self._all_p[0]
        solution = self._all_solutions[0]
        self._multi_obj = 0
        
        self._flc.set_instance(points, demands, n, p, True)
        
        return self.prepare(n, p, solution), {}
    
    def _update_type(self):
        if self._current_type >= self._num_types:
            raise RuntimeError('Action taken after episode is done.')
        if self._current_type < self._num_types - 1:
            self._current_type += 1
            print(f'current type: {self._current_type}')
            print(self._num_types)
            points = self._all_points
            demands = self._all_demands[:,self._current_type]
            n = self._n
            p = self._all_p[self._current_type]
            solution = self._all_solutions[self._current_type]
            self._flc.set_instance(points, demands, n, p, True)
            self.prepare(n, p, solution)

    def prepare(self, n: int, p: int, solution: list) -> Dict:
        self._obs_extractor.reset()
        self._done = False
        self._t = 0
        if len(self._all_p) > 1:
            self._max_steps = self._all_max_steps[self._current_type]
            self._flc.init_facility_mask(self._old_mask[self._current_type], self._new_mask[self._current_type])
        else:
            self._max_steps = 0
            
        initial_solution = solution
        initial_obj_value = self._flc.compute_obj_value_from_solution(initial_solution,stage=2)
        
        self._obj_value = np.zeros(self._max_steps + 1)
        self._obj_value[0] = initial_obj_value
        self._solution = np.zeros((self._max_steps + 1, n), dtype=bool)
        self._solution[0] = initial_solution
        self._best_solution = initial_solution
        self._best_obj_value = initial_obj_value
        self._last_best_t = 0
        self._tabu_stable_steps = max(1, round(self._max_steps * self._tabu_stable_steps_scale))
        return self._get_obs(self._t)
    
    def _get_max_steps(self) -> list:
        # print(self._all_solutions)
        tmp_all_solitions = list(self._all_solutions)
        max_steps = []
        old_idx = []
        new_idx = []
        for t in range(self._num_types):
            count_true = [sum(s) for s in zip(*(tmp_all_solitions[t:]))]
            # print(count_true)
            old = [i for i in range(len(count_true)) if count_true[i] > 1 and tmp_all_solitions[t][i]]
            new = [i for i in range(len(count_true)) if count_true[i] == 0]
            if len(old):
                old_idx.append(old)
                new_idx.append(new)
                max_steps.append(len(old))
                for i in old:
                    count_true[i] = count_true[i] - 1
        # print(max_steps, old_idx, new_idx)
        return max_steps, old_idx, new_idx

    def _generate_new_instance(self) -> Tuple[np.ndarray, np.ndarray, int, int]:
        n = self._np_random.integers(self._min_n, self._max_n, endpoint=True)
        p_ratio = self._np_random.uniform(self._min_p_ratio, self._max_p_ratio)
        p = int(max(n * p_ratio, 4))

        points = self._np_random.uniform(size=(n, 2))
        while np.unique(points, axis=0).shape[0] != n:
            points = self._np_random.uniform(size=(n, 2))
        demands = self._np_random.random(size=(n,))
        return points, demands, n, p
    
    def _load_multi_facility_data(self, data_npy) -> Tuple[np.ndarray, np.ndarray]:
        data = data_npy.split('\n')
        n = len(data)
        p = int((len(data[0].split(' '))-2) / 2)
        
        positions = []
        demands = []
        actual_facilities = []
        ps = []
        for row in data:
            row = row.split(' ')
            row = [x for x in row if len(x)]
            positions.append([float(row[0]), float(row[1])])

            demand = []
            for i in range(2, 2+p):
                demand.append(float(row[i]))
            demands.append(demand)
            
            actual_facility = []
            for i in range(2+p, 2+2*p):
                actual_facility.append(bool(int(float(row[i]))))
            actual_facilities.append(actual_facility)

        positions = np.array(positions)
        positions = np.deg2rad(positions)
        demands = np.array(demands)
        actual_facilities = np.array(actual_facilities)
        ps = actual_facilities.sum(axis=0)
        
        return positions, demands, n, ps
        
    def _load_multi_facility_solutions(self, boost) -> list:
        def load_model(positions, demands, n, p, boost):
            eval_env = EvalPMPEnv(self.cfg, positions, demands, n, p, boost)
            eval_env = DummyVecEnv([lambda: eval_env])

            policy_kwargs = get_policy_kwargs(self.cfg)
            test_model = PPO(MaskedFacilityLocationActorCriticPolicy,
                        eval_env,
                        verbose=1,
                        policy_kwargs=policy_kwargs,
                        device='cuda:1')
            train_model = PPO.load(self.cfg.load_model_path)
            test_model.set_parameters(train_model.get_parameters())
            return test_model, eval_env

        def get_optimal_solution(model, eval_env):
            obs = eval_env.reset()
            done = False
            while not done:
                action, _ = model.predict(obs, deterministic=True)
                obs, _, done, info = eval_env.step(action)
            return eval_env.get_attr('_best_solution')[0]
        
        multi_solutions = []
        for i in range(len(self._all_p)):
            positions = self._all_points
            demands = self._all_demands[:,i]
            n = self._n
            p = self._all_p[i]
            model, env = load_model(positions,demands,n,p,boost)
            multi_solutions.append(get_optimal_solution(model, env))
    
        return multi_solutions
    
    def get_reward(self) -> float:
        if self._done:
            reward = np.min(self._obj_value)
        else:
            reward = 0.0
        return reward