import gurobipy as gu
import gurobipy as gu
from icecream import ic
import random
import pandas as pd
import numpy as np

class Subproblem_a:
    def __init__(self, duals_p, duals_td, duals_g, data, p, iteration, R_p, Entry_p, W, W_min):
        itr = iteration + 1
        self.P = p
        self.P_Full = data['P'].dropna().astype(int).unique().tolist()
        self.T = data['T'].dropna().astype(int).unique().tolist()
        self.D = data['D'].dropna().astype(int).unique().tolist()
        self.duals_p = duals_p
        self.duals_td = duals_td
        self.duals_g = duals_g
        self.Model = gu.Model("Subproblem")
        self.itr = itr
        self.R_p = R_p
        self.Entry_p = Entry_p
        self.W = W
        self.W_min = W_min

        self.G_d = set(self.R_p.values())
        self.R_g = {g: [i for i in self.P_Full if self.R_p[i] == g] for g in self.G_d}
        self.G = list(self.G_d)
        self.key_group = [key for key, value in self.R_g.items() if self.P in value][0]
    def buildModel(self):
        self.genVars()
        self.genCons()
        self.genFarkasObj()
        self.Model.update()
    def genVars(self):
        self.l = self.Model.addVars([self.P], self.D, vtype=gu.GRB.BINARY, name="l")
        self.LOS = self.Model.addVars([self.P], [self.itr], vtype=gu.GRB.INTEGER, name="LOS")
        self.x = self.Model.addVars([self.P], self.T, self.D, [self.itr], vtype=gu.GRB.BINARY, name="x")
        self.y = self.Model.addVars([self.P], self.D, vtype=gu.GRB.BINARY, name="y")
        self.z = self.Model.addVars([self.P], self.T, vtype=gu.GRB.BINARY, name="z")
        self.w = self.Model.addVars([self.P],self. D, vtype=gu.GRB.BINARY, name="w")
        self.S = self.Model.addVars([self.P], self.D, vtype=gu.GRB.INTEGER, name="S")
        self.App = self.Model.addVars([self.P], self.D, vtype=gu.GRB.CONTINUOUS, name="App")
        self.Model.update()

    def genCons(self):
        for p in [self.P]:
            self.Model.addLConstr(1 == gu.quicksum(self.l[p, d] for d in self.D))
            self.Model.addLConstr(self.LOS[p, self.itr] == gu.quicksum(self.l[p, d] * d for d in self.D) - self.Entry_p[p] + 1)
            self.Model.addLConstr(1 == gu.quicksum(self.z[p, t] for t in self.T))
            self.Model.addLConstr(1 == gu.quicksum(self.x[p, t, self.Entry_p[p], self.itr] for t in self.T))
            for d in range(2, len(self.D) + 1):
                self.Model.addLConstr(self.w[p, d] <= 1 - gu.quicksum(self.l[p, k] for k in range(2, d)))
            for d in self.D:
                if d < self.Entry_p[p]:
                    self.Model.addLConstr(0 == self.w[p, d])
                elif d == self.Entry_p[p]:
                    self.Model.addLConstr(1 == self.w[p, d])
                self.Model.addLConstr(self.l[p, d] <= gu.quicksum(self.x[p, t, d, self.itr] for t in self.T))
                self.Model.addLConstr(self.w[p, d] == gu.quicksum(self.x[p, t, d, self.itr] for t in self.T) + self.y[p, d])
                if d >= self.Entry_p[p] and d < len(self.D) - self.W + 1:
                    self.Model.addLConstr(gu.quicksum(self.y[p, j] for j in range(d, d + self.W)) <= self.W_min)
                for t in self.T:
                    self.Model.addLConstr(self.x[p, t, d, self.itr] <= self.z[p, t])
            for d in self.D:
                self.Model.addLConstr(self.l[p, d] * self.R_p[p] <= gu.quicksum(
                    gu.quicksum(self.x[p, t, j, self.itr] for j in range(1, d + 1)) for t in self.T) + gu.quicksum(
                    self.y[p, j] for j in range(1, d + 1)) * 0.5)
        self.Model.update()

    def genFarkasObj(self):
        self.Model.setObjective(- gu.quicksum(self.x[self.P, t, d, self.itr] * self.duals_td[t, d] for t in self.T for d in self.D) - self.duals_p[self.P], sense=gu.GRB.MINIMIZE)

    def getOptVals(self, var_name):
        variable = getattr(self, var_name, None)
        if variable is None:
            raise AttributeError(f"Variable '{var_name}' not found in the class.")

        values_dict = self.Model.getAttr("X", variable)
        values_list = list(values_dict.values())

        return values_dict, values_list

    def create_lambda_list(self, index):
        lst = [0] * len(self.P_Full)
        if 1 <= index <= len(self.P_Full):
            lst[index - 1] = 1
        return lst


class MasterProblem_a:
    def __init__(self, df, current_iteration, nr, T_Max, Nr_c, R_p, farkas = None, farkas_schedule=None, farkas_los = None):
        self.iteration = current_iteration
        self.P = df['P'].dropna().astype(int).unique().tolist()
        self.D = df['D'].dropna().astype(int).unique().tolist()
        self.T = df['T'].dropna().astype(int).unique().tolist()
        self.A = [1]
        self.A_Farkas = [i for i in range(1, current_iteration + 1)]
        self.Model = gu.Model("MasterProblem")
        self.cons_p_max = {}
        self.cons_max_group = {}
        self.Nr = Nr_c
        self.cons_lmbda = {}
        self.output_len = nr
        self.T_max = T_Max
        self.all_schedules = {}
        self.schedules_farkas_x = farkas_schedule
        self.schedules_farkas_los = farkas_los
        self.all_los = {}
        self.beta = 1
        self.R_p = R_p
        self.farkas = farkas

        self.G_d = set(self.R_p.values())
        self.R_g = {g: [i for i in self.P if self.R_p[i] == g] for g in self.G_d}
        self.G = list(self.G_d)
        ic(self.G, self.R_g, self.R_g, self.R_p)

    def buildModel(self):
        self.genVars()
        self.genConsFarkas()
        self.Model.Params.InfUnbdInfo = 1
        self.genFarkasObj()
        self.Model.update()

    def genVars(self):
        self.lmbda = self.Model.addVars(self.P, self.A, lb = 0,vtype=gu.GRB.INTEGER, name='lmbda')
        self.max = self.Model.addVars(self.G, lb = 0, vtype=gu.GRB.INTEGER, name=f"max")

    def genConsFarkas(self):
        for p in self.P:
            print(f'Nr: {self.Nr[p]}')
            self.cons_lmbda[p] = self.Model.addConstr(self.Nr[p] == 0, name=f"lmbda({p})")
        for t in self.T:
            for d in self.D:
                self.cons_p_max[t, d] = self.Model.addConstr(0 <= self.T_max[t, d], name=f"p_max({t},{d})")
        self.Model.update()

    def genFarkasObj(self):
        self.Model.setObjective(0, sense=gu.GRB.MINIMIZE)

    def solRelModel(self):
        self.Model.Params.OutputFlag = 1
        self.Model.Params.Method = 1
        for v in self.Model.getVars():
            v.setAttr('vtype', 'C')
            v.setAttr('lb', 0.0)
        self.Model.update()
        self.Model.optimize()
        if self.Model.status == 3:
            print('\nThe following constraints and variables are in the IIS:')
            self.Model.computeIIS()
            self.Model.write(f"iis{self.iteration}.ilp")
            for c in self.Model.getConstrs():
                if c.IISConstr: print(f'\t{c.constrname}: {self.Model.getRow(c)} {c.Sense} {c.RHS}')
            for v in self.Model.getVars():
                if v.IISLB: print(f'\t{v.varname} ≥ {v.LB}')
                if v.IISUB: print(f'\t{v.varname} ≤ {v.UB}')

    def addSchedule(self, schedule):
        self.all_schedules.update(schedule)

    def addLOS(self, schedule):
        self.all_los.update(schedule)

    def getFarkasDuals(self):
        return {(t, d): self.cons_p_max[t, d].FarkasDual for t in self.T for d in self.D}, {p: self.cons_lmbda[p].FarkasDual for p in self.P}, {(g, i): 0 for g in self.G for i in self.R_g[g]}

    def generate_start_solution(self):
        from collections import defaultdict

        start_solution_x = {(p, t, d): 0.0 for p in self.P for t in self.T for d in self.D}

        for (t, d), constraint in self.cons_p_max.items():
            involved_lambdas = set()

            lhs = self.Model.getRow(constraint)

            for i in range(lhs.size()):
                var = lhs.getVar(i)
                if var.varName.startswith("lmbda["):
                    key = tuple(map(int, var.varName[6:-1].split(',')))
                    involved_lambdas.add(key)

            for (p, a) in involved_lambdas:
                start_solution_x[(p, t, d)] = 1.0

        start_solution_los = {}
        obj_expr = self.Model.getObjective()

        for i in range(obj_expr.size()):
            var = obj_expr.getVar(i)
            coeff = obj_expr.getCoeff(i)
            if var.varName.startswith("lmbda["):
                key = tuple(map(int, var.varName[6:-1].split(',')))
                start_solution_los[key] = coeff

        new_data = defaultdict(list)
        for (i, j), value in start_solution_los.items():
            new_data[i].append(value)
        start_solution_los_f = {k: v[0] for k, v in new_data.items()}

        return start_solution_x, start_solution_los_f

    def addLambdaVar(self, p, a, col, coef):
        new_col = gu.Column(col, self.Model.getConstrs())
        var_name = f"lambda[{p},{a}]"
        self.lmbda[p, a] = self.Model.addVar(obj=coef[0], vtype=gu.GRB.INTEGER, column=new_col, name=var_name)
        self.Model.update()


# Parameter
random.seed(29)
W, W_min = 5, 3
E_app, E_min, baselevel, lin = 0.6, 0.1, 0.5, 0.1
R_p, Entry_p, Max_t, P, D, T, N_c, R_p_c, Entry_p_c, Nr_c, unique_combinations, num_unique_combinations = {1: 4, 2: 1, 3: 1, 4: 1, 5: 1, 6: 2, 7: 2, 8: 3, 9: 1, 10: 1}, {1: 2, 2: 7, 3: 5, 4: 1, 5: 2, 6: 4, 7: 3, 8: 3, 9: 1, 10: 6}, {(1, 1): 3, (1, 2): 3, (1, 3): 3, (1, 4): 0, (1, 5): 0, (1, 6): 3, (1, 7): 3, (2, 1): 0, (2, 2): 1, (2, 3): 1, (2, 4): 1, (2, 5): 1, (2, 6): 1, (2, 7): 1, (3, 1): 0, (3, 2): 1, (3, 3): 1, (3, 4): 1, (3, 5): 1, (3, 6): 1, (3, 7): 1, (4, 1): 2, (4, 2): 2, (4, 3): 2, (4, 4): 0, (4, 5): 0, (4, 6): 2, (4, 7): 2, (5, 1): 0, (5, 2): 4, (5, 3): 4, (5, 4): 4, (5, 5): 4, (5, 6): 4, (5, 7): 4}, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6, 7, 8, 9], {1: 4, 2: 1, 3: 1, 4: 1, 5: 1, 6: 2, 7: 2, 8: 3, 9: 1}, {1: 2, 2: 7, 3: 5, 4: 1, 5: 2, 6: 4, 7: 3, 8: 3, 9: 6}, {1: 1, 2: 1, 3: 1, 4: 2, 5: 1, 6: 1, 7: 1, 8: 1, 9: 1}, {(4, 2): [1], (1, 7): [2], (1, 5): [3], (1, 1): [4, 9], (1, 2): [5], (2, 4): [6], (2, 3): [7], (3, 3): [8], (1, 6): [10]}, 9


# **** Prerequisites ****
# Create Dataframes
data = pd.DataFrame({
    'P': N_c + [np.nan] * (max(len(N_c), len(T), len(D)) - len(N_c)),
    'T': T + [np.nan] * (max(len(N_c), len(T), len(D)) - len(T)),
    'D': D + [np.nan] * (max(len(N_c), len(T), len(D)) - len(D))
})

# **** Farkas ****
schedules_farkas_los = {}
schedules_farkas_x = {}
ModelInfeasible = True
infeasible_count = 0

master_farkas = MasterProblem_a(data, infeasible_count, 50, Max_t, Nr_c, R_p, True, schedules_farkas_x, schedules_farkas_los)
master_farkas.buildModel()
master_farkas.solRelModel()

while (ModelInfeasible) and infeasible_count < 200:

    master_farkas.solRelModel()
    if master_farkas.Model.Status == 3:
        farkas_ts, farkas_i, farkas_g = master_farkas.getFarkasDuals()
        ModelInfeasible = True
        infeasible_count += 1
    elif master_farkas.Model.Status != 3:
        ModelInfeasible = False
    else:
        break

    for index in N_c:
        subproblem_farkas = Subproblem_a(farkas_i, farkas_ts, farkas_g, data, index, 0, R_p_c, Entry_p_c, W, W_min)
        subproblem_farkas.buildModel()
        subproblem_farkas.Model.optimize()
        reducedCost = subproblem_farkas.Model.objval

        # Generate and add columns with reduced cost
        if reducedCost < -1e-7:
            Schedules_x_farkas, x_list_farkas = subproblem_farkas.getOptVals('x')
            master_farkas.addSchedule(Schedules_x_farkas)
            Schedules_LOS_farkas, LOS_list_farkas = subproblem_farkas.getOptVals('LOS')
            master_farkas.addLOS(Schedules_LOS_farkas)
            lambda_list_farkas = subproblem_farkas.create_lambda_list(index)
            master_farkas.addLambdaVar(index, infeasible_count + 1, [*lambda_list_farkas, *x_list_farkas], LOS_list_farkas)

        else:
            print(f'No Column found for Patient {index} for Initalization.')