"""
Kronos: A simple scheduler for graduate training programme

Entities: User, Schedule, Rotation
"""

from operator import itemgetter
from datetime import datetime, timedelta


def getRotationCapacity(rotationId, startDate, endDate, assignments):
    """ Calculate number of users assigned to a particular rotation during the specified duration
    """
    start = datetime.strptime(startDate, "%d%m%Y")
    end = datetime.strptime(endDate, "%d%m%Y")
    duration = int((end - start).days / 7.0)
    # Weeks involved during the rotation
    weeks = [(start + timedelta(weeks=x)).strftime("%W%Y") for x in range(0, duration)]
    capacity = sum(itemgetter(*weeks)(assignments[rotationId][0][0]))
    return capacity


def score_assignment(
    assignments,
    solution,
    earliestAvailableDate,
    core_rotations=["PMO", "PE", "SE", "PM"],
    rotation_duration={
        "PMO": 12,
        "PE": 12,
        "SE": 12,
        "PM": 12,
        "SYS": 12,
        "ARC": 12,
        "ANA": 12,
    },
):
    """ Calculate loss function for suggested solution (negative = better)
    Parameters:
        assignments (dict): global assignment object by rotation
        solution (dict): rotation assignment for a user
        earliestAvailableDate (date): earliest date where a user can be assigned a rotation
        core_rotations (list): rotation that should be completed first
        rotation_duration (dict): duration of each rotation
    """
    print(solution)
    # SOFT CONSTRAINT 1 - Core rotations should be completed in the first 4 rotations if possible
    core_first_loss = sum(
        [
            -3 if x[0] in core_rotations else 0
            for x in solution
            if int(x[1]) <= len(core_rotations)
        ]
    )
    # SOFT CONSTRAINT 2 - External Assignment must be assigned last
    external_assignment_loss = (
        99 if "EXT" in [x[0] for x in solution] and solution[-1][0] != "EXT" else 0
    )

    # Calculate timing of each rotation from solution
    solution = [
        (
            x[0],
            rotation_duration[x[0]]
            + (sum([rotation_duration[x[0]] for x in solution[:i]]) if i != 0 else 0),
        )
        for i, x in enumerate(solution)
    ]
    startDate = earliestAvailableDate
    schedule = []
    for x in solution:
        endDate = startDate + timedelta(weeks=x[1]) - timedelta(days=1)
        # Make sure the date falls on weekday
        if endDate.weekday() >= 5:
            endDate -= timedelta(endDate.weekday() - 4)
        schedule.append(
            (x[0], startDate.strftime("%d%m%Y"), endDate.strftime("%d%m%Y"))
        )
        startDate += timedelta(weeks=x[1])

    spread_first_loss = sum(
        [getRotationCapacity(x[0], x[1], x[2], assignments) for x in schedule]
    )
    loss = core_first_loss + external_assignment_loss + spread_first_loss
    return loss


def schedule2assignments(schedule):
    """ Convert schedule object to assignment object
    """
    rotations = {}
    for userId, userSchedule in schedule.items():
        for rotation in userSchedule:
            id = rotation["rotationId"]
            if id not in rotations:
                rotations[id] = [[{}], []]
            print(rotations[id][0][0])
            startDate, endDate = itemgetter("startDate", "endDate")(rotation)
            start = datetime.strptime(startDate, "%d%m%Y")
            end = datetime.strptime(endDate, "%d%m%Y")
            duration = int((end - start).days / 7.0)
            for i in range(duration):
                date = (start + timedelta(weeks=i)).strftime("%W%Y")
                if date not in rotations[id][0][0]:
                    rotations[id][0][0][date] = 0
                rotations[id][0][0][date] += 1
            rotations[id][1].append((userId, startDate, endDate))
        sortedDate = sorted(list(rotations[id][0][0].keys()))
        if len(rotations[id][0]) < 2:
            rotations[id][0].append(sortedDate[0])
            rotations[id][0].append(sortedDate[-1])
        elif sortedDate[0] < rotations[id][0][1]:
            rotations[id][0][1] = sortedDate[0]
        elif len(rotations[id][0]) > 2 and sortedDate[-1] > rotations[id][0][2]:
            rotations[id][0][2] = sortedDate[-1]
    print(rotations)
    return rotations


def assignments2schedule(assignments):
    """ Convert assignment object to overall schedule
    """
    users = {}
    for rotationId, rotationInfo in assignments.items():
        for userId, userAssignment in rotationInfo[1].items():
            if userId not in users:
                users[userId] = []
            users[userId].append(
                {
                    "rotationId": rotationId,
                    "startDate": userAssignment[0],
                    "endDate": userAssignment[1],
                }
            )
    print(users)
    return users


def generateUserSchedule(user, assignments, scoring_function):
    """ Generate most optimal user schedule
    Parameters:
        user (object): User
        assignments (dict): Time-bounded assignments
        scoring_function (function): scoring function to rank possible assignments

    Returns:
        schedule (list): list of rotations
    """
    return [{"rotationId": "PMO", "startDate": "012018"}]


def getOverallSchedule(users):
    """ Generate overall schedule from individual user's schedule
    Parameters:
        users (list): list of Users

    Returns:
        schedule (dict): overall assignments
    """
    return {}


def getConflictingAssignments(schedule):
    """ Get list of assignments which exceeded rotation capacity
    Parameters:
        schedule (dict): overall assignments

    Returns:
        confictingAssignmentsByRotation (dict): overall schedule with conflicting assignments
    """
    return {}


if __name__ == "__main__":
    pass

[(14, 'datetime.datetime.strptime', 'datetime.strptime', (['startDate', '"""%d%m%Y"""'], {}), False, 'from datetime import datetime, timedelta\n'), (15, 'datetime.datetime.strptime', 'datetime.strptime', (['endDate', '"""%d%m%Y"""'], {}), False, 'from datetime import datetime, timedelta\n'), (79, 'datetime.timedelta', 'timedelta', ([], {'weeks': 'x[1]'}), False, 'from datetime import datetime, timedelta\n'), (19, 'operator.itemgetter', 'itemgetter', (['*weeks'], {}), False, 'from operator import itemgetter\n'), (72, 'datetime.timedelta', 'timedelta', ([], {'days': '(1)'}), False, 'from datetime import datetime, timedelta\n'), (99, 'datetime.datetime.strptime', 'datetime.strptime', (['startDate', '"""%d%m%Y"""'], {}), False, 'from datetime import datetime, timedelta\n'), (100, 'datetime.datetime.strptime', 'datetime.strptime', (['endDate', '"""%d%m%Y"""'], {}), False, 'from datetime import datetime, timedelta\n'), (72, 'datetime.timedelta', 'timedelta', ([], {'weeks': 'x[1]'}), False, 'from datetime import datetime, timedelta\n'), (98, 'operator.itemgetter', 'itemgetter', (['"""startDate"""', '"""endDate"""'], {}), False, 'from operator import itemgetter\n'), (18, 'datetime.timedelta', 'timedelta', ([], {'weeks': 'x'}), False, 'from datetime import datetime, timedelta\n'), (103, 'datetime.timedelta', 'timedelta', ([], {'weeks': 'i'}), False, 'from datetime import datetime, timedelta\n')]