from gurobipy import GRB
import gurobipy as gp

import numpy as np

from src.algos.transit_assignment import transit_assignment_gp
from src.busopt.config import ModelConfig


class Outer_bounding_prob_base:
    """
    subproblem to compute a lower bound for the tri-level objective function,
    given a determined upper-level solution x.
    it has three separate problems:
    (1) minimize passengers' cost (in SO form)
    (2) minimize company's operating cost
    (3) minimize social funding cost induced by subsiding
    """

    def __init__(self, hypergraph, line_list, demand_list, config: ModelConfig, fixed_up_solution=None):
        self.hypergraph = hypergraph
        self.line_list = line_list
        self.demand_list = demand_list
        self.fixed_x_sol = fixed_up_solution.x_sol
        self.config = config

        # initialize Gurobi model
        self.model = gp.Model("Outer_bounding_prob_base")
        self.model.setParam("OutputFlag", 0)

    def create_f_space(self):
        """
        create the feasible region for variable f, M(x)
        (1) f_l - BigM * x_l <= 0
        (2) f_l - minimum_feq * x_l >= 0
        (3) fleet size: \sum {2 T_l f} <= 0
        (4) f_l to y_ln
        """
        model = self.model
        CONSTRS_LINKING_BIGM = 100

        l_indices = [l.line_id for l in self.line_list]
        fm_l = model.addVars(l_indices, vtype=GRB.CONTINUOUS, name='f_l')
        y_ln = model.addVars(l_indices, self.config.freq_set, vtype=GRB.BINARY, name='y_ln')

        # (1) f_l - BigM * x_l <= 0
        model.addConstrs(fm_l[l_index] - CONSTRS_LINKING_BIGM * self.fixed_x_sol[l_index] <= 0 for l_index in l_indices)

        # (2) f_l - minimum_feq * x_l >= 0
        model.addConstrs(fm_l[l_index] - self.config.freq_set[1] * self.fixed_x_sol[l_index] >= 0 for l_index in l_indices)

        # (3) fleet size: \sum {2 T_l f} <= 0
        model.addConstr(gp.quicksum(2 * l.time * fm_l[l.line_id] for l in self.line_list) <= self.config.private_fleet_size)

        # (4) f_l to y_ln
        model.addConstrs(fm_l[l.line_id] == gp.quicksum(n * y_ln[l.line_id, n] for n in self.config.freq_set) for l in self.line_list)
        model.addConstrs(gp.quicksum(y_ln[l.line_id, n] for n in self.config.freq_set) == 1.0 for l in self.line_list)

        # register class variable
        self.fm_l = fm_l
        self.y_ln = y_ln

    def create_vt_space(self):
        """
        create the feasible region for variable v and t, Omega(f).
        Since the Omega(f) is originally defined in the space of hyperpath flow,
        to avoid path enumeration, an OD-based reformulation is needed.
        This is in reference to Xu et al. (2022)

        flow balance constraint:
        sum{j in N_i-} sum{s' in S_j, (j,i) in M_s'} u_js'^w * e_jis'
        - sum{s in S_i} u_is^w
        = q_iw    forall i in N, w in W

        non-negative constraint:
        u_is^w >= 0

        connection between u_is^w and v_ij
        v_ij = sum{w in W, s in S_i, (i,j) in M_s} u_is^w * e_ijs

        connection between u_is^w and t_i
        t_i = sum{w in W, s in S_i} u_is^w / sum{(i,j) in M_s} f_ij

        u_is^w -- flow leaving node i that belongs to OD-pair w via hyperlink s
        e_ijs -- weight of a link (i,j) including in hyperlink s
        S_i -- set of hyperlinks that take node i as head node
        M_s -- set of regular links in hyperlink s
        """
        model = self.model
        hypergraph = self.hypergraph
        demand_list = self.demand_list

        v_a = model.addVars([a.link_id for a in hypergraph.link_dict.values()],
                            vtype=GRB.CONTINUOUS, name='v_a')

        t_i = model.addVars([i.node_id for i in hypergraph.node_dict.values() if i.node_type == 'stop'],
                            vtype=GRB.CONTINUOUS, name='t_i')

        u_isw = model.addVars([(i.node_id, s.id, w.demand_id)
                               for i in hypergraph.node_dict.values() for s in i.out_hyperlinks for w in demand_list],
                              vtype=GRB.CONTINUOUS, name='u_isw')

        # flow balance constraint (in terms of hyperlink flow)
        for i in hypergraph.node_dict.values():
            for w in demand_list:
                sum_in_flow = gp.quicksum(u_isw[s.from_node.node_id, s.id, w.demand_id] * s.link_prob[s.link_dict[i.node_id]]
                                          for s in i.in_hyperlinks)
                sum_out_flow = gp.quicksum(u_isw[i.node_id, s.id, w.demand_id]
                                           for s in i.out_hyperlinks)
                if i.node_id == w.ori_node.node_id:
                    rhs = -w.demand
                elif i.node_id == w.des_node.node_id:
                    rhs = w.demand
                else:
                    rhs = 0
                model.addConstr(sum_in_flow - sum_out_flow == rhs)

        # connection between u_is^w and v_ij
        for a in hypergraph.link_dict.values():
            i = a.from_node
            link_flow = gp.quicksum(u_isw[i.node_id, s.id, w.demand_id] * s.link_prob[a]
                                    for s in i.out_hyperlinks for w in demand_list if a in s.link_prob.keys())
            model.addConstr(v_a[a.link_id] == link_flow)

        # connection between u_is^w and t_i
        for i in [_ for _ in hypergraph.node_dict.values() if _.node_type == 'stop']:
            node_waiting = gp.quicksum(u_isw[i.node_id, s.id, w.demand_id] / sum([a.freq for a in s.link_dict.values()])
                                       for s in i.out_hyperlinks for w in demand_list if s.has_waiting)
            model.addConstr(t_i[i.node_id] == node_waiting)

        # register class variable
        self.v_a = v_a
        self.t_i = t_i
        self.u_isw = u_isw


class Outer_bounding_prob_1(Outer_bounding_prob_base):
    """
    subproblem that compute the minimum user cost (given x and f)
        min sum{a in A} C_a(v_a, f_a) * v_a + sum{i in N_s} t_i
    s.t.
    (v,t) in Omega(f)

    Could prove that the objective is monotone to f
    So that take maximum possible value of f (ignore the fleet size constraint),
    result in the minimum user costs.

    nonlinear term is linearized through PWL techniques.
    """

    def __init__(self, hypergraph, line_list, demand_list, config: ModelConfig, fixed_up_solution=None):
        super().__init__(hypergraph, line_list, demand_list, config, fixed_up_solution)

        # update link/hyperlink frequencies
        self._update_link_frequencies()

        # build model
        self.create_vt_space()

        # auxiliary variables for approximation operation
        # apx_a = C_a(v_a) * v_a
        apx_a = self.model.addVars([a.link_id for a in hypergraph.link_dict.values()],
                                   vtype=GRB.CONTINUOUS, name='apx_a')
        self.apx_a = apx_a
        transit_assignment_gp(self.hypergraph, self.demand_list, enable_logging=False)  # do assignment once, to set the variables' upper bound
        self._approximation()

        # objective function
        obj_func = self.config.vot * (gp.quicksum(self.apx_a.values()) + gp.quicksum(self.t_i.values()))
        self.model.setObjective(obj_func, GRB.MINIMIZE)

    def _update_link_frequencies(self):
        """update link frequencies with maximum possible value"""
        for l in self.line_list:
            new_freq = self.config.freq_set[-1] if self.fixed_x_sol[l.line_id] != 0 else 0
            l.freq = new_freq + l.org_freq
            for route in l.routes.values():
                for board_link in route.board_links:
                    board_link.freq = new_freq + l.org_freq
        for s in self.hypergraph.hyperlink_dict.values():
            s.update_link_prob()

    def _approximation(self):
        """
        Using PWL to linearize the nonlinear term
        """
        model = self.model

        nonlinear_links = [a for a in self.hypergraph.link_dict.values() if a.link_type in ["transit", "board"]]
        linear_links = [a for a in self.hypergraph.link_dict.values() if a.link_type in ["alight", "walk"]]
        EXPAND_COEF = 1.5  # upper limit = expand_coef * link_flow
        X1_NUM, X2_NUM = 5, 3
        x1_pts = {a: list(np.linspace(0, EXPAND_COEF * a.flow, X1_NUM)) for a in nonlinear_links}
        x2_pts = {a: list(np.linspace(0, EXPAND_COEF * a.dept_link.flow, X2_NUM)) for a in nonlinear_links}

        # auxiliary variables:
        scatter_index = [(a.link_id, i, j) for a in nonlinear_links for i in range(X1_NUM - 1) for j in range(X2_NUM - 1)]
        K_U_1 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_U_1')
        K_U_2 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_U_2')
        K_U_3 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_U_3')
        K_L_1 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_L_1')
        K_L_2 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_L_2')
        K_L_3 = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='K_L_3')
        H_U = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='H_U', ub=[1.0] * len(scatter_index))
        H_L = model.addVars(scatter_index, vtype=GRB.CONTINUOUS, name='H_L', ub=[1.0] * len(scatter_index))

        # Constraints:
        for a in nonlinear_links:

            for i in range(X1_NUM - 1):
                for j in range(X2_NUM - 1):
                    model.addConstr(K_U_1[a.link_id, i, j] + K_U_2[a.link_id, i, j] + K_U_3[a.link_id, i, j] == H_U[a.link_id, i, j])
                    model.addConstr(K_L_1[a.link_id, i, j] + K_L_2[a.link_id, i, j] + K_L_3[a.link_id, i, j] == H_L[a.link_id, i, j])

            model.addConstr(
                gp.quicksum((H_U[a.link_id, i, j] + H_L[a.link_id, i, j]) for i in range(X1_NUM - 1) for j in range(X2_NUM - 1)) == 1)

            model.addConstr(
                self.v_a[a.link_id] ==
                gp.quicksum(K_U_1[a.link_id, i, j] * x1_pts[a][i] + K_U_2[a.link_id, i, j] * x1_pts[a][i] + K_U_3[a.link_id, i, j] * x1_pts[a][i + 1] +
                            K_L_1[a.link_id, i, j] * x1_pts[a][i] + K_L_2[a.link_id, i, j] * x1_pts[a][i + 1] + K_L_3[a.link_id, i, j] * x1_pts[a][i + 1]
                            for i in range(X1_NUM - 1) for j in range(X2_NUM - 1)))

            model.addConstr(
                self.v_a[a.dept_link.link_id] ==
                gp.quicksum(K_U_1[a.link_id, i, j] * x2_pts[a][j] + K_U_2[a.link_id, i, j] * x2_pts[a][j + 1] + K_U_3[a.link_id, i, j] * x2_pts[a][j + 1] +
                            K_L_1[a.link_id, i, j] * x2_pts[a][j] + K_L_2[a.link_id, i, j] * x2_pts[a][j] + K_L_3[a.link_id, i, j] * x2_pts[a][j + 1]
                            for i in range(X1_NUM - 1) for j in range(X2_NUM - 1)))

            approx_func = a.second_func()
            model.addConstr(
                self.apx_a[a.link_id] ==
                gp.quicksum(K_U_1[a.link_id, i, j] * approx_func(x1_pts[a][i], x2_pts[a][j]) + K_U_2[a.link_id, i, j] * approx_func(x1_pts[a][i], x2_pts[a][j + 1]) +
                            K_U_3[a.link_id, i, j] * approx_func(x1_pts[a][i + 1], x2_pts[a][j + 1]) +
                            K_L_1[a.link_id, i, j] * approx_func(x1_pts[a][i], x2_pts[a][j]) + K_L_2[a.link_id, i, j] * approx_func(x1_pts[a][i + 1], x2_pts[a][j]) +
                            K_L_3[a.link_id, i, j] * approx_func(x1_pts[a][i + 1], x2_pts[a][j + 1])
                            for i in range(X1_NUM - 1) for j in range(X2_NUM - 1)))

        for a in linear_links:
            self.model.addConstr(self.apx_a[a.link_id] == a.travel_cost * self.v_a[a.link_id])

    def solve(self):
        self.model.optimize()
        self.model.update()
        return self.model.objVal


class Outer_bounding_prob_2(Outer_bounding_prob_base):
    """
    subproblem that computes the minimum operating costs.
    operating cost = sum{l in L} 2 C^M d_l f_l
    Don't need an optimization model
    """

    def __init__(self, hypergraph, line_list, demand_list, config: ModelConfig, fixed_up_solution=None):
        super().__init__(hypergraph, line_list, demand_list, config, fixed_up_solution)

    def solve(self):
        obj_value = 0
        for l in self.line_list:
            freq = self.config.freq_set[1] if self.fixed_x_sol[l.line_id] != 0 else 0
            obj_value += 2 * self.config.c_bar * l.length * freq

        return obj_value


class Outer_bounding_prob_3(Outer_bounding_prob_base):
    """
    subproblem that computes the minimum social funding costs.
        min mu * sum{l in L} (sum{a in A_t^l} subsidy * d_a *v_a) f_l / (f_l + f_l0)
    s.t.
    (v,t) in Omega(f)
    """

    def __init__(self, hypergraph, line_list, demand_list, config: ModelConfig, fixed_up_solution=None):
        super().__init__(hypergraph, line_list, demand_list, config, fixed_up_solution)

        # update frequencies
        self._update_link_frequencies()

        # build model
        self.create_vt_space()

        # objective function
        obj_func = 0
        for l in self.line_list:
            for r in l.routes.values():
                expression = gp.quicksum(a.length * self.v_a[a.link_id] for a in r.transit_links)
                obj_func += self.config.mu * self.config.s_bar * (l.freq - l.org_freq) / l.freq * expression

        # impose punishment on walking links to avoid unwanted walking
        # this is only valid when walking links is never used
        obj_func += 1e6 * gp.quicksum(self.v_a[a.link_id] for a in self.hypergraph.link_dict.values() if a.link_type == 'walk')

        self.model.setObjective(obj_func, GRB.MINIMIZE)

    def solve(self):
        self.model.optimize()
        self.model.update()
        return self.model.objVal

    def _update_link_frequencies(self):
        """update link frequencies with minimum possible value"""
        for l in self.line_list:
            new_freq = self.config.freq_set[1] if self.fixed_x_sol[l.line_id] != 0 else 0
            l.freq = new_freq + l.org_freq
            for route in l.routes.values():
                for board_link in route.board_links:
                    board_link.freq = new_freq + l.org_freq
        for s in self.hypergraph.hyperlink_dict.values():
            s.update_link_prob()



class Outer_bounding_prob:
    def __init__(self, hypergraph, line_list, demand_list, fixed_up_solution=None):
        config = hypergraph.config
        self.p1 = Outer_bounding_prob_1(hypergraph, line_list, demand_list, config, fixed_up_solution)
        self.p2 = Outer_bounding_prob_2(hypergraph, line_list, demand_list, config, fixed_up_solution)
        self.p3 = Outer_bounding_prob_3(hypergraph, line_list, demand_list, config, fixed_up_solution)




    def solve(self, reset=True):
        v1 = self.p1.solve()
        v2 = self.p2.solve()
        v3 = self.p3.solve()
        if reset:
            self.p1.model.reset()
            self.p3.model.reset()

        return {"v1": v1, "v2": v2, "v3": v3}
