# %%
PROJECT_PATH = '/home/md/Work/ligalytics/leagues_stable/'
import os, sys
sys.path.insert(0, PROJECT_PATH)
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "leagues.settings")
os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"
import random

# XPRESS ENVIRONMENT
os.environ['XPRESSDIR'] = "/opt/xpressmp_8.4"
os.environ['XPRESS'] = "/opt/xpressmp_8.4/bin"
os.environ['LD_LIBRARY_PATH'] = os.environ['XPRESSDIR']+"/lib:"+os.environ['LD_LIBRARY_PATH']
os.environ['DYLD_LIBRARY_PATH'] = os.environ['XPRESSDIR']+"/lib:"
os.environ['SHLIB_PATH'] = os.environ['XPRESSDIR']+"/lib:"
os.environ['LIBPATH'] = os.environ['XPRESSDIR']+"/lib:"
# os.environ['PYTHONPATH'] =os.environ['XPRESSDIR']+"/lib:"+os.environ['PYTHONPATH']
os.environ['PYTHONPATH'] =os.environ['XPRESSDIR']+"/lib:"
os.environ['CLASSPATH'] =os.environ['XPRESSDIR']+"/lib/xprs.jar:"
os.environ['CLASSPATH'] =os.environ['XPRESSDIR']+"/lib/xprb.jar:"+os.environ['CLASSPATH']
os.environ['CLASSPATH'] =os.environ['XPRESSDIR']+"/lib/xprm.jar:"+os.environ['CLASSPATH']
os.environ['PATH'] =os.environ['XPRESSDIR']+"/bin:"+os.environ['PATH']


from leagues import settings
settings.DATABASES['default']['NAME'] = PROJECT_PATH+'/db.sqlite3'

import django
django.setup()

from scheduler.models import *
from pulp import *


# import xpress as xp
# xp.controls.outputlog = 1



scenario = Scenario.objects.get(id=52)
solver = "xpress"


# %%

teams = Team.objects.filter(season=scenario.season,active=True)
gameRequirements = GameRequirement.objects.filter(scenario=scenario)
pairings = Pairing.objects.filter(scenario=scenario,type="Home",prio="Hard")
days = Day.objects.filter(season=scenario.season)

daysPerRound = defaultdict(lambda:[])
for d in days:
    daysPerRound[d.round].append(d)

print ('#####################################################')
print ('# SOLVING MODEL (UCL24 - CREATE FEASIBLE SCHEDULE)  #')
print ('#####################################################')

rounds = daysPerRound.keys()
nRounds = len(rounds)
single_day_rounds = [r for r in daysPerRound if len(daysPerRound[r]) == 1]
two_day_rounds = [r for r in daysPerRound if len(daysPerRound[r]) == 2]
team_ids = list(teams.values_list('id',flat=True))
getTeamByID = {
    t.id:t for t in teams
}
gamereqs = []
opponents_from_pot = defaultdict(lambda:defaultdict(lambda:[]))
for req in gameRequirements:
    gamereqs.append((req.team1.id,req.team2.id))
    opponents_from_pot[req.team1.id][req.team2.pot].append(req.team2.id)
    opponents_from_pot[req.team2.id][req.team1.pot].append(req.team1.id)
pot = {}
teams_in_pot = {}
for i in teams.values_list('pot',flat=True).distinct():
    pot[i] = teams.filter(pot=i)
    teams_in_pot[i] = list(teams.filter(pot=i).values_list('id',flat=True))

model = pulp.LpProblem("", pulp.LpMinimize)            
x = { (r,t1,t2) : pulp.LpVariable('x_'+str(r)+'_'+str(t1)+'_'+str(t2) , lowBound = 0, upBound = 1, cat = pulp.LpInteger) for r in rounds for t1 in team_ids for t2 in team_ids if (t1,t2) in gamereqs or (t2,t1) in gamereqs}
home = {(r,t) : pulp.LpVariable('home_'+str(r)+'_'+str(t) , lowBound = 0, upBound = 1, cat = pulp.LpInteger) for r in rounds for t in team_ids}
penalty = {t : pulp.LpVariable('penalty_'+str(t) , lowBound = 0, cat = pulp.LpContinuous) for t in team_ids}

# each game has to be played
for (t1,t2) in gamereqs:
    model += lpSum([x[r,t1,t2]+x[r,t2,t1] for r in rounds])  == 1

for t in team_ids:
    for r in rounds:
        # each team plays once in each round
        model += lpSum([x[r,t,t2]+x[r,t2,t] for t2 in team_ids if (r,t,t2) in x.keys()]) == 1
        # couple homes
        # model += home[r,t] == lpSum([x[r,t,t2] for t2 in team_ids if (r,t,t2) in x.keys()])
        model += home[r,t] == lpSum([x[r,t,t2] for t2 in team_ids if (r,t,t2) in x.keys()])
    
    # structural (home) requirements
    for r in range(1,nRounds-1):
        model += lpSum([home[r2,t] for r2 in range(r,r+3)]) <= 2
        model += lpSum([1-home[r2,t] for r2 in range(r,r+3)]) <= 2
    model += lpSum([home[r2,t] for r2 in range(1,3)]) == 1
    model += lpSum([home[r2,t] for r2 in range(nRounds-1,nRounds+1)]) == 1
    
    # play home against each pot
    for p in pot:
        model += lpSum([x[r,t,t2] for r in rounds for t2 in team_ids if (r,t,t2) in x.keys() and t2 in teams_in_pot[p]]) == 1

# stadium/city clashes
for r in single_day_rounds:
    for pair in pairings.filter(dist__in=[0,1]):
        model += lpSum(home[r,pair.team1.id]+home[r,pair.team2.id]) <= 1
for r in two_day_rounds:
    for pair in pairings.filter(dist=1):
        model += lpSum(home[r,pair.team1.id]+home[r,pair.team2.id]) <= 1


# penalty for two POT-1 games in a row
for t in team_ids:
    for r in range(1,nRounds):
        model += lpSum(x[r,t,t2] + x[r,t2,t] + x[r+1,t,t3] + x[r+1,t3,t] for t2 in opponents_from_pot[t][1] for t3 in opponents_from_pot[t][1] if t2 != t3) <= 1 + penalty[t]


# some objective to avoid pulp from failing    
model += lpSum(penalty[key] for key in penalty.keys())

if solver == "CBC":
    model.solve(PULP_CBC_CMD(maxSeconds = 120,msg=1))
elif solver == "Gurobi":
    model.solve(GUROBI(TimeLimit=120, msg=1))
else:
    model.solve(XPRESS(msg=1,maxSeconds=120))

try:
    objective = value(model.objective)
    games = [(key[1],key[2],key[0]) for key in x.keys() if x[key].value() > 0]
except:
    objective = -1
    games = []

# return objective, games 

# return games



# %%


if model.status == 1:
    home_dict = {}
    away_dict = {}
    for key in x.keys():
        if x[key].value() > 0:
            # print(key[0],getTeamByID[key[1]],getTeamByID[key[2]])
            home_dict[key[0],key[1]] = key[2]
            away_dict[key[0],key[2]] = key[1]

    sol = " \
        <style> \
        table, th, td { \
        border: 1px solid black; \
        border-collapse: collapse; \
        } \
        </style> \
    "
    sol += "<table style='border:1px solid black'>\n"
    sol += "<thead>\n"
    sol += "<tr><th></th>"
    for r in rounds:
        sol += f"<th>{r}</th>"
    sol += "</tr>"
    sol += "</thead>\n"
    sol += "<tbody>\n"
    for t in team_ids:
        tname = getTeamByID[t]
        sol += f"<tr><td>{tname}</td>"
        for r in rounds:
            if (r,t) in home_dict.keys():
                opponent = getTeamByID[home_dict[(r,t)]]
                sol += f"<td style='background-color:lightyellow'>{opponent.pot} - {opponent}</td>"
            elif (r,t) in away_dict.keys():
                opponent = getTeamByID[away_dict[(r,t)]]
                sol += f"<td style='background-color:lightsteelblue'>{opponent.pot} - {opponent}</td>"
            else:
                sol += "<td></td>"
        sol += "</tr>"
    sol += "</tbody>\n"
    sol += "</table>\n"
    sol += "<br><br><br>\n"

    with open('draw_schedule.html', 'w') as f:
        f.write(sol)

# %%

# games = [(key[1],key[2],key[0]) for key in x.keys() if model.getSolution(x[key]) > 0]
# %%
