import gurobipy as gp
from time import perf_counter
from gurobipy import GRB
from src.busopt.config import ModelConfig
from src.entity.network_entity import OD_demand
from src.entity.solution_entity import MidSolution, UpSolution, LowSolution
from src.algos.transit_assignment import transit_assignment_gp, update_cost

# noinspection PyShadowingNames
class Inner_relaxed_mp:

    def __init__(self, hypergraph, line_list, demand_list, fixed_up_solution, config: ModelConfig):
        self.config = config
        fixed_x = fixed_up_solution.x_sol
        demand_dict = OD_demand.destination_based_OD(demand_list)
        model = gp.Model("Inner_relaxed_mp_model")
        model.setParam("OutputFlag", 0)


        """
        Create Gurobi model.
        """

        # Variables.
        l_indices = [l.line_id for l in line_list]
        f_l = model.addVars(l_indices, vtype=GRB.CONTINUOUS, name='f_l')
        y_ln = model.addVars(l_indices, config.freq_set, vtype=GRB.BINARY, name='y_ln')
        z_ip = model.addVar(vtype=GRB.CONTINUOUS, name='z_ip', ub=1e7)  # Benders auxiliary variable

        # Constraints.
        # constraints that links f_l and y_ln:
        model.addConstrs(
            f_l[l.line_id] == gp.quicksum(n * y_ln[l.line_id, n] for n in config.freq_set) for l in line_list)
        model.addConstrs(gp.quicksum(y_ln[l.line_id, n] for n in config.freq_set) == 1.0 for l in line_list)

        # f_l - pN * x_l <= 0
        cons_1 = model.addConstrs(f_l[l_index] - config.freq_set[-1] * fixed_x[l_index] <= 0 for l_index in l_indices)

        # f_l - p1 * x_l >= 0
        cons_2 = model.addConstrs(f_l[l_index] - config.freq_set[1] * fixed_x[l_index] >= 0 for l_index in l_indices)


        # Objectives
        obj_mp_part = gp.quicksum(2 * config.c_bar * l.length * f_l[l.line_id] for l in line_list)
        obj_func = z_ip + obj_mp_part
        model.setObjective(obj_func, GRB.MINIMIZE)
        model.update()

        # register class variable
        self.model = model
        self.line_list = line_list
        self.fixed_up_solution = fixed_up_solution
        self.hypergraph = hypergraph
        self.demand_list = demand_list
        self.demand_dict = demand_dict
        self.f_l = f_l
        self.y_ln = y_ln
        self.obj_mp_part = obj_mp_part
        self.z_ip = z_ip
        self.cons_1 = cons_1
        self.cons_2 = cons_2
        self.Benders_opt_cut = []
        self.Benders_fea_cut = []
        self.subgradient_cut = []
        self.output_info = {}


    def update_fixed_up_solution(self, new_up_solution):
        self.fixed_up_solution = new_up_solution
        fixed_x = self.fixed_up_solution.x_sol
        for l in self.line_list:
            self.cons_1[l.line_id].rhs = self.config.freq_set[-1] * fixed_x[l.line_id]
            self.cons_2[l.line_id].rhs = self.config.freq_set[1] * fixed_x[l.line_id] - 1e-4
        self.model.update()


    def update_freq(self, mid_sol):
        for l in self.line_list:
            l.freq = mid_sol.f_sol[l.line_id] + l.org_freq
            for route in l.routes.values():
                for board_link in route.board_links:
                    board_link.freq = mid_sol.f_sol[l.line_id] + l.org_freq

    def clear_benders_cut(self):
        self.model.remove(self.Benders_opt_cut)
        self.model.remove(self.Benders_fea_cut)
        self.model.update()
        self.Benders_opt_cut = []
        self.Benders_fea_cut = []

    def clear_subgradient_cut(self):
        self.model.remove(self.subgradient_cut)
        self.model.update()
        self.subgradient_cut = []

    def _compute_subprob_obj(self, cur_low_sol):
        """
        Calculates the subproblem's objective function value. Q_ip(f)
        """
        user_cost = self.config.vot * (sum([value[0] * value[1] for value in cur_low_sol.v_sol.values()]) + sum(cur_low_sol.t_sol.values()))
        subsidies = 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
            subsidies += total_flow_mile * self.config.s_bar * private_portion

        return user_cost + subsidies * self.config.mu

    def solve(self):
        UB = float("inf")  # upper bound
        LB = -float("inf")  # lower bound
        iter_step = 0  # iteration step counter
        best_mid_sol, best_low_sol, best_mid_revenue = None, None, None
        epsilon = 1e-3  # optimality gap
        L_ip = 0

        start_time = perf_counter()
        is_feasible = True
        while True:
            iter_step += 1
            self.model.optimize()

            if self.model.status == GRB.Status.INFEASIBLE or self.model.status == GRB.Status.UNBOUNDED:
                print(f"Inner relaxed MP: Iteration {iter_step}, Status: {self.model.status}")
                is_feasible = False
                break

            # update lower bound and solution
            LB = max(LB, self.model.objVal)
            this_mid_sol = MidSolution(f_dict={key: value.X for key, value in self.f_l.items()})
            self.update_freq(this_mid_sol)

            # solve the resulting transit assignment problem.
            transit_assignment_gp(self.hypergraph, self.demand_list, enable_logging=False)
            this_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()})

            this_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
                this_mid_revenue += total_flow_mile * (self.config.tau + self.config.s_bar) * private_portion
                this_mid_revenue -= 2 * self.config.c_bar * this_mid_sol.f_sol[l.line_id]

            is_subprob_feasible = this_mid_revenue >= self.config.Psi0
            """
            add benders cut to the master problem
            if the subproblem returns an optimal solution, optimality cut is added:
                z_ip <= - (ub_inner - gamma(f)) * expr + gamma(f)
                expr = sum_{1-set}y_ln - sum_{0-set}y_ln - len(l_set)
            elseif subproblem is infeasible, feasibility cut is added:
                expr + 1 <= 0
            """
            expr = 0
            for l in self.line_list:
                for n in self.config.freq_set:
                    if abs(self.y_ln[l.line_id, n].X - 1.0) < 1e-5:
                        expr += self.y_ln[l.line_id, n]
                    elif abs(self.y_ln[l.line_id, n].X) < 1e-5:
                        expr += - self.y_ln[l.line_id, n]
            expr -= len(self.line_list)

            if is_subprob_feasible:
                cur_Gamma = self._compute_subprob_obj(this_low_sol)
                cur_obj_mp_part = self.obj_mp_part.getValue()
                if cur_Gamma + cur_obj_mp_part < UB:
                    UB = cur_Gamma + cur_obj_mp_part
                    best_mid_sol = this_mid_sol
                    best_low_sol = this_low_sol
                    best_mid_revenue = this_mid_revenue

                self.Benders_opt_cut.append(self.model.addConstr(self.z_ip >= (cur_Gamma - L_ip) * expr + cur_Gamma))

            else:
                self.Benders_fea_cut.append(self.model.addConstr(expr + 1.0 <= 0))
                self.model.update()

            # Print iterative information
            if iter_step == 2 or iter_step % 50 == 0:
                print(f"Iteration {iter_step}, "
                      f"LB = {LB:.2f}, UB = {UB:.2f}, Gap = {(UB - LB) * 100 / abs(UB) if UB != 0 else 100:.2f}%, "
                      f"Elapsed time: {perf_counter() - start_time:.2f} secs")

            if UB - LB < epsilon:
                print(f"Inner relaxed MP: Iteration {iter_step}, Converged! obj = {LB}")
                break

        end_time = perf_counter()
        flex_x_num = int(sum(self.fixed_up_solution.x_sol.values()))
        maximum_step = (len(self.config.freq_set) - 1) ** flex_x_num
        print(f"Inner Problem time elapsed: {end_time - start_time:.2f} s")
        print(f"Inner Problem convergence iteration: {iter_step - 1}/{maximum_step}, efficiency: {(iter_step - 1) / maximum_step}")


        self.output_info.update({
            "iteration_step": iter_step - 1,
            "max_iteration_step": maximum_step,
            "CPU (s)": end_time - start_time,
            "optimal_obj_value": LB,
            "optimal_mid_solution": best_mid_sol.f_sol if is_feasible else None,
            "optimal_low_solution": best_low_sol.v_sol if is_feasible else None,
            "Gap": ((UB - LB) * 100 / abs(UB) if UB != 0 else 100) if is_feasible else 100,
        })

        self.model.update()

        return is_feasible, best_mid_sol, best_low_sol, best_mid_revenue


# noinspection PyShadowingNames
class Outer_relaxed_mp:
    def __init__(self, hypergraph, line_list, demand_list, config: ModelConfig):
        self.config = config
        model = gp.Model("Outer_relaxed_mp_model")
        model.setParam("OutputFlag", 0)

        """
        Create Gurobi model.
        """

        # Variables.
        l_indices = [l.line_id for l in line_list]
        x_l = model.addVars(l_indices, vtype=GRB.BINARY, name='x_l')
        z_op = model.addVar(vtype=GRB.CONTINUOUS, name='z_op', lb=0)    # Benders auxiliary variable


        # Objectives.
        model.setObjective(z_op, GRB.MINIMIZE)
        model.update()

        # register class variable
        self.model = model
        self.hypergraph = hypergraph
        self.line_list = line_list
        self.demand_list = demand_list
        self.x_l = x_l
        self.z_op = z_op
        self.output_info = {}


    def _compute_subprob_obj(self, cur_mid_sol, cur_low_sol):
        """
        Calculates the subproblem's objective function value. Q_op(x)
        """
        user_cost = self.config.vot * (sum([value[0] * value[1] for value in cur_low_sol.v_sol.values()]) + sum(cur_low_sol.t_sol.values()))
        oper_cost = sum([2 * self.config.c_bar * l.length * cur_mid_sol.f_sol[l.line_id] for l in self.line_list])
        subsidies = 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
            subsidies += total_flow_mile * self.config.s_bar * private_portion

        return {"Q_op": user_cost + oper_cost + subsidies * self.config.mu,
                "user_cost": user_cost,
                "oper_cost": oper_cost,
                "subsidies": subsidies}


    def solve(self):
        UB = float("inf")  # upper bound
        LB = -float("inf")  # lower bound
        iter_step = 0  # iteration step counter
        best_up_sol, best_mid_sol, best_low_sol = None, None, None
        best_user_cost, best_mid_oper, best_funding = None, None, None
        best_sol_iter, best_sol_time = None, None
        inner_model = None
        best_mid_revenue = None
        solution_set = []
        self.output_info['iteration_details'] = []

        L_op = 0

        start_time = perf_counter()
        while True:
            iter_step += 1
            print(f"========================== O-BLP Master Problem: Iteration {iter_step} =================================")
            print(f"Start optimizing")
            self.model.optimize()

            if self.model.status == GRB.Status.INFEASIBLE or self.model.status == GRB.Status.UNBOUNDED:
                print(f"O-BLP Master Problem: Iteration {iter_step}, Status: {self.model.status}")
                break

            # update lower bound and solution
            LB = max(LB, self.model.objVal)
            this_up_sol = UpSolution(x_dict={key: value.X for key, value in self.x_l.items()})
            print(f"\nSolving master problem...")
            print(f"Up solution in this iteration: {this_up_sol} | Master problem Obj.value: {self.model.objVal:.2f}")

            # Termination Criteria: reach a solution once acquired
            if this_up_sol in solution_set:
                print(f"reach a solution once acquired")
                print(f"O-BLP Master Problem:  Iteration {iter_step}, Converged!")
                break
            else:
                solution_set.append(this_up_sol)

            this_mid_sol, Q_op, this_user_cost, this_oper_cost, this_subsidies = None, None, None, None, None
            relaxed_prob_obj_value = None

            # solve the subproblem to optimal
            # can prove that it is equivalent to solve the Inner Problem
            print(f"\nSolving subproblem...")
            if inner_model is None:
                inner_model = Inner_relaxed_mp(self.hypergraph, self.line_list, self.demand_list, this_up_sol, self.config)
                inner_model.update_fixed_up_solution(this_up_sol)
            else:
                inner_model.update_fixed_up_solution(this_up_sol)
                inner_model.clear_benders_cut()
                inner_model.clear_subgradient_cut()
            is_inner_feasible, this_mid_sol, this_low_sol, this_mid_revenue = inner_model.solve()

            """
            add benders cut to the master problem
            if the subproblem returns an optimal solution, optimality cut is added:
                z_ip >= (gamma(x) - L_op) * expr + gamma(x)
                expr = sum_{1-set}x_l - sum_{0-set}x_l - len({1-set}})
            elseif subproblem is infeasible, feasibility cut is added:
                expr + 1 <= 0
            """
            expr = 0
            length_1set = 0
            for l in self.line_list:
                if abs(self.x_l[l.line_id].X - 1.0) < 1e-5:
                    expr += self.x_l[l.line_id]
                    length_1set += 1
                elif abs(self.x_l[l.line_id].X) < 1e-5:
                    expr -= self.x_l[l.line_id]
            expr -= length_1set

            if not is_inner_feasible:
                # feasibility cut
                self.model.addConstr(expr + 1.0 <= 0)
                self.model.update()
                # Print iterative information
                print(f"\n[SUMMARY]")
                print(f"Subproblem is infeasible, add a feasibility cut")
                print(f"Best incumbent solution: {best_up_sol}, {best_mid_sol}")
                print(f"LB = {LB:.2f}, UB = {UB:.2f}, Gap = {(UB - LB) * 100 / abs(UB) if UB != 0 else 100:.2f}% ")

            else:
                # calculate Q_op(x) and update upper bound
                dict_subprob = self._compute_subprob_obj(this_mid_sol, this_low_sol)
                Q_op, this_user_cost, this_oper_cost, this_subsidies = (
                    dict_subprob["Q_op"], dict_subprob["user_cost"], dict_subprob["oper_cost"], dict_subprob["subsidies"])

                if UB > Q_op:
                    UB = Q_op
                    best_up_sol = this_up_sol
                    best_mid_sol = this_mid_sol
                    best_low_sol = this_low_sol
                    best_user_cost = this_user_cost
                    best_mid_oper = this_oper_cost
                    best_funding = this_subsidies
                    best_mid_revenue = this_mid_revenue
                    best_sol_iter = iter_step
                    best_sol_time = perf_counter() - start_time

                # optimality cut
                self.model.addConstr(self.z_op >= (Q_op - L_op) * expr + Q_op)

                # Print iterative information
                print(f"\n[SUMMARY]")
                print(f"New incumbent: {this_up_sol == best_up_sol}")
                print(f"This iteration solution: {this_up_sol}, {this_mid_sol}")
                print(f"Best incumbent solution: {best_up_sol}, {best_mid_sol}")
                print(f"objective value (total cost): {Q_op:.2f}")
                print(f"passenger cost: {this_user_cost:.2f}, operating cost: {this_oper_cost:.2f}, funding (not cost): {this_subsidies:.2f}")
                print(f"LB = {LB:.2f}, UB = {UB:.2f}, Gap = {(UB - LB) * 100 / abs(UB) if UB != 0 else 100:.2f}%")

            cum_time = perf_counter() - start_time
            print(f"Cumulative elapsed time: {cum_time:.2f} secs")
            print(f"====================================================================================================\n")

            self.output_info['iteration_details'].append({
                "iter_step": iter_step,
                "objective_value": Q_op, "passenger cost": this_user_cost, "operating cost": this_oper_cost, "funding": this_subsidies,
                "lb": LB, "ub": UB, "gap": (UB - LB) * 100 / abs(UB) if UB != 0 else 100,
                "relax_bound": relaxed_prob_obj_value,
                "cumulative_elapsed_time": cum_time,
                "new_incumbent": this_up_sol == best_up_sol if this_mid_sol is not None else False,
                "this_up_sol": this_up_sol.x_sol if this_mid_sol is not None else None,
                "this_mid_sol": this_mid_sol.f_sol if this_mid_sol is not None else None,
                "best_up_sol": best_up_sol.x_sol if best_mid_sol is not None else None,
                "best_mid_sol": best_mid_sol.f_sol if best_mid_sol is not None else None
            })

        elapsed_time = perf_counter() - start_time
        iterations = iter_step - 1
        maximum_step = 2 ** len(self.line_list)

        # Optimization summary
        print(f"\n\n ============================ O-BLP Optimization Summary ============================")
        print(f"Time Elapsed: {elapsed_time:.2f} secs")
        print(f"Convergence Iteration: {iterations}/{maximum_step}, Efficiency: {iterations / maximum_step}")
        print(f"Best solution: {best_up_sol}, {best_mid_sol}")
        print(f"Best objective value (total cost): {UB}")
        print(f"Best solution found at iteration {best_sol_iter}")
        print(f"Best solution found cost: {best_sol_time:.2f} secs")
        print(f"Relative gap: {(UB - LB) * 100 / abs(UB) if UB != 0 else 100:.2f}%")
        print(f"passenger cost: {best_user_cost:.2f}, operating cost: {best_mid_oper:.2f}, funding (not cost): {best_funding:.2f}")

        # Save results
        self.output_info.update({
            "iterations": iterations,
            "max_iteration_step": maximum_step,
            "max_f_iterations": len(self.config.freq_set) ** self.config.route_num,
            "CPU (s)": elapsed_time,
            "upper_bound": UB,
            "lower_bound": LB,
            "Gap": (UB - LB) * 100 / abs(UB) if UB != 0 else 100,
            "iteration_find_best": best_sol_iter,
            "time_find_best": best_sol_time,
            "optimal_obj_user_cost": best_user_cost,
            "optimal_obj_mid_oper": best_mid_oper,
            "optimal_obj_funding": best_funding,
            "optimal_mid_revenue": best_mid_revenue,
            "optimal_up_solution": best_up_sol.x_sol,
            "optimal_mid_solution": best_mid_sol.f_sol,
            "optimal_low_solution": best_low_sol.v_sol,
            "L_op": L_op,
            "config": self.config.to_dict()
        })




