import itertools
import time

from gurobipy import GRB
import gurobipy as gp

import numpy as np
from src.algos.transit_assignment import transit_assignment_gp, update_cost
from src.busopt.config import ModelConfig
from src.entity.solution_entity import MidSolution, LowSolution


class HPR:
    """
    high-point relaxation of tri-level model
    """

    def __init__(self, hypergraph, line_list, demand_list, mid_sol, config: ModelConfig):
        self.hypergraph = hypergraph
        self.line_list = line_list
        self.demand_list = demand_list
        self.mid_sol = mid_sol
        self.config = config

        # initialize Gurobi model
        self.model = gp.Model("HPR")
        self.model.setParam("OutputFlag", 0)
        self._update_link_frequencies()
        self.create_vt_space()
        self.create_convex_approximation()
        self.create_objective()


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


    def create_objective(self):
        obj_func_1 = self.config.vot * (gp.quicksum(self.apx_a.values()) + gp.quicksum(self.t_i.values()))
        obj_func_2 = 2 * self.config.c_bar * gp.quicksum(l.length * self.mid_sol.f_sol[l.line_id] for l in self.line_list)
        obj_func_3 = 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_3 += self.config.mu * self.config.s_bar * (l.freq - l.org_freq) / l.freq * expression

        self.model.setObjective(obj_func_1 + obj_func_2 + obj_func_3, GRB.MINIMIZE)
        self.obj_func_1, self.obj_func_2, self.obj_func_3 = obj_func_1, obj_func_2, obj_func_3

    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


    def create_convex_approximation(self):
        # update link/hyperlink frequencies
        self._update_link_frequencies()

        # auxiliary variables for approximation operation
        # apx_a = C_a(v_a) * v_a
        apx_a = self.model.addVars([a.link_id for a in self.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()

    def _update_link_frequencies(self):
        for l in self.line_list:
            new_freq = self.mid_sol.f_sol[l.line_id]
            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 = 10, 5
        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()

        for a in self.hypergraph.link_dict.values():
            a.flow = float(self.v_a[a.link_id].X)
        for i in self.hypergraph.node_dict.values():
            if i.node_type == 'stop':
                i.wait = float(self.t_i[i.node_id].X)
        update_cost(self.hypergraph)

        mid_revenue = 0
        for l in self.line_list:
            total_flow_mile = sum([tl.flow * tl.length for r in l.routes.values() for tl in r.transit_links])
            private_portion = (l.freq - l.org_freq) / l.freq
            mid_revenue += total_flow_mile * (self.config.tau + self.config.s_bar) * private_portion
            mid_revenue -= 2 * self.config.c_bar * self.mid_sol.f_sol[l.line_id]

        low_sol = LowSolution(
            v_dict={a.link_id: (float(a.flow), float(a.travel_cost)) for a in self.hypergraph.link_dict.values()},
            t_dict={i.node_id: i.wait for i in self.hypergraph.node_dict.values()})

        return self.model.objVal, mid_revenue, low_sol, self.obj_func_1.getValue(), self.obj_func_2.getValue(), self.obj_func_3.getValue()


def solve_up_HPR(hypergraph, line_list, demand_list, up_sol, config: ModelConfig):
    best_mid_sol, best_low_sol, best_mid_revenue = None, None, None
    best_user_cost, best_oper_cost, best_funding_cost = None, None, None
    best_obj_value = float("inf")
    is_inner_feasible = True

    all_feasible_f = {}
    for k, v in up_sol.x_sol.items():
        if abs(v) <= 1e-3:
            all_feasible_f[k] = [config.freq_set[0]]
        elif abs(v - 1.0) <= 1e-3:
            all_feasible_f[k] = config.freq_set[1:]
        else:
            raise Exception

    combinations = [dict(zip(all_feasible_f.keys(), combo)) for combo in itertools.product(*all_feasible_f.values())]
    start_time = time.time()
    for i, v in enumerate(combinations):
        this_mid_sol = MidSolution(f_dict=v)
        hpr = HPR(hypergraph, line_list, demand_list, this_mid_sol, config)
        this_obj_value, this_mid_revenue, this_low_sol, this_user_cost, this_oper_cost, this_funding_cost = hpr.solve()
        if this_mid_revenue >= config.Psi0 and this_obj_value < best_obj_value:
            best_mid_sol = this_mid_sol
            best_low_sol = this_low_sol
            best_obj_value = this_obj_value
            best_mid_revenue = this_mid_revenue
            best_user_cost = this_user_cost
            best_oper_cost = this_oper_cost
            best_funding_cost = this_funding_cost

        if i % 50 == 0:
            print(f"HPR iteration {i}/{len(combinations)}, best bound {best_obj_value:.2f}, time: {time.time() - start_time:.2f} secs")

    if best_mid_sol is None:
        is_inner_feasible = False

    print(f"HPR terminate")


    return {"is_inner_feasible": is_inner_feasible,
            "best_mid_sol": best_mid_sol,
            "best_low_sol": best_low_sol,
            "best_mid_revenue": best_mid_revenue,
            "best_user_cost": best_user_cost,
            "best_oper_cost": best_oper_cost,
            "best_funding_cost": best_funding_cost}


