from typing import Text

import numpy as np
import pulp
from spopt.locate import PMedian

from facility_location.env import EvalPMPEnv


class PMPSolver:
    def __init__(self, solver: Text, env: EvalPMPEnv):
        if solver == 'GUROBI':
            self._solver = pulp.GUROBI(msg=False)
        elif solver == 'GUROBI_CMD':
            self._solver = pulp.GUROBI_CMD(msg=False)
        elif solver == 'PULP_CBC_CMD':
            self._solver = pulp.PULP_CBC_CMD(msg=False)
        elif solver == 'GLPK_CMD':
            self._solver = pulp.GLPK_CMD(msg=False)
        elif solver == 'MOSEK':
            self._solver = pulp.MOSEK(msg=False)
        else:
            raise ValueError(f'Solver {solver} not supported.')

        self.env = env

    def solve(self):
        _, demands, _, p = self.env.get_instance()
        distance_matrix, _ = self.env.get_distance_and_cost()
        pmedian_from_cost_matrix = PMedian.from_cost_matrix(distance_matrix, demands, p_facilities=p)
        pmedian_from_cost_matrix = pmedian_from_cost_matrix.solve(self._solver)
        solution = np.array([len(temp) > 0 for temp in pmedian_from_cost_matrix.fac2cli], dtype=bool)
        return solution