/*
 * Copyright 2021 Red Hat, Inc. and/or its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.acme.vaccinationscheduler.solver;

import static java.time.temporal.ChronoUnit.DAYS;
import static java.time.temporal.ChronoUnit.MINUTES;

import java.time.LocalDateTime;
import java.util.function.Predicate;

import org.acme.vaccinationscheduler.domain.solver.PersonAssignment;
import org.optaplanner.core.api.score.buildin.bendablelong.BendableLongScore;
import org.optaplanner.core.api.score.stream.Constraint;
import org.optaplanner.core.api.score.stream.ConstraintCollectors;
import org.optaplanner.core.api.score.stream.ConstraintFactory;
import org.optaplanner.core.api.score.stream.ConstraintProvider;

public class VaccinationScheduleConstraintProvider implements ConstraintProvider {

    public static final int HARD_LEVELS_SIZE = 1;
    public static final int SOFT_LEVELS_SIZE = 5;

    private static final LocalDateTime COVID_EPOCH = LocalDateTime.of(2021, 1, 1, 0, 0);

    private BendableLongScore ofHard(long hardScore) {
        return BendableLongScore.ofHard(HARD_LEVELS_SIZE, SOFT_LEVELS_SIZE, 0, hardScore);
    }

    private BendableLongScore ofSoft(int softLevel, long softScore) {
        return BendableLongScore.ofSoft(HARD_LEVELS_SIZE, SOFT_LEVELS_SIZE, softLevel, softScore);
    }

    // Because the @PlanningVariable is nullable=true, the from() classes needed to be filtered
    private Predicate<PersonAssignment> personAssignedFilter = (personAssignment -> personAssignment.getVaccinationSlot() != null);

    @Override
    public Constraint[] defineConstraints(ConstraintFactory constraintFactory) {
        return new Constraint[]{
                // Hard constraints
                vaccinationSlotCapacity(constraintFactory),
                requiredVaccineType(constraintFactory),
                requiredVaccinationCenter(constraintFactory),
                minimumAgeVaccineType(constraintFactory),
                maximumAgeVaccineType(constraintFactory),
                readyDate(constraintFactory),
                dueDate(constraintFactory),
                // TODO restrict maximum distance
                // Medium constraints
                scheduleSecondOrLaterDosePeople(constraintFactory),
                scheduleHigherPriorityRatingPeople(constraintFactory),
                // Soft constraints
                preferredVaccineType(constraintFactory),
                preferredVaccinationCenter(constraintFactory),
                regretDistance(constraintFactory),
                idealDate(constraintFactory),
                higherPriorityRatingEarlier(constraintFactory)
        };
    }

    // ************************************************************************
    // Hard constraints
    // ************************************************************************

    Constraint vaccinationSlotCapacity(ConstraintFactory constraintFactory) {
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .groupBy(PersonAssignment::getVaccinationSlot, ConstraintCollectors.count())
                .filter((vaccinationSlot, personCount) -> personCount > vaccinationSlot.getCapacity())
                .penalizeLong("Vaccination slot capacity", ofHard(1_000),
                        (vaccinationSlot, personCount) -> personCount - vaccinationSlot.getCapacity());
    }

    Constraint requiredVaccineType(ConstraintFactory constraintFactory) {
        // Typical usage: if a person is coming for their 2nd dose, use the same vaccine type as their 1st dose.
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter((personAssignment -> personAssignment.getRequiredVaccineType() != null
                        && personAssignment.getVaccinationSlot().getVaccineType() != personAssignment.getRequiredVaccineType()))
                .penalize("Required vaccine type", ofHard(10_000_000));
    }

    Constraint requiredVaccinationCenter(ConstraintFactory constraintFactory) {
        // Typical usage: if a person is coming for their 2nd dose, enforce the same vaccination center as their 1st dose.
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter((personAssignment -> personAssignment.getRequiredVaccinationCenter() != null
                        && personAssignment.getVaccinationSlot().getVaccinationCenter() != personAssignment.getRequiredVaccinationCenter()))
                .penalize("Required vaccination center", ofHard(1_000_000));
    }

    Constraint minimumAgeVaccineType(ConstraintFactory constraintFactory) {
        // Don't inject too young people with a vaccine that has minimum age
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter(personAssignment -> personAssignment.getVaccinationSlot().getVaccineType().getMaximumAge() != null
                        && personAssignment.getAgeOnVaccinationDate()
                        < personAssignment.getVaccinationSlot().getVaccineType().getMinimumAge()
                        && personAssignment.getRequiredVaccineType() == null)
                .penalizeLong("Minimum age of vaccination type", ofHard(1),
                        personAssignment -> personAssignment.getVaccinationSlot().getVaccineType().getMinimumAge()
                                - personAssignment.getAgeOnVaccinationDate());
    }

    Constraint maximumAgeVaccineType(ConstraintFactory constraintFactory) {
        // Don't inject too oldr people with a vaccine that has maximum age
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter(personAssignment -> personAssignment.getVaccinationSlot().getVaccineType().getMaximumAge() != null
                        && personAssignment.getAgeOnVaccinationDate()
                        > personAssignment.getVaccinationSlot().getVaccineType().getMaximumAge()
                        // If the 1th dose was a max 55 year vaccine, then it's ok to inject someone who only turned 56 last week with it
                        && personAssignment.getRequiredVaccineType() == null)
                .penalizeLong("Maximum age of vaccination type", ofHard(1),
                        personAssignment -> personAssignment.getAgeOnVaccinationDate()
                                - personAssignment.getVaccinationSlot().getVaccineType().getMaximumAge());
    }

    Constraint readyDate(ConstraintFactory constraintFactory) {
        // Typical usage 1: If a person is coming for their 2nd dose, don't inject it before the ready day.
        // For example, Pfizer is ready to injected 19 days after the first dose. Moderna after 26 days.
        // Typical usage 2: If a person wants to reschedule an invited/accepted appointment,
        // set the readyDate one day after the appointment date to avoid inviting the same day (especially for multiple reschedules)
        // and also prohibit gamification (to get an earlier appointment).
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter(personAssignment -> personAssignment.getReadyDate() != null
                        && personAssignment.getVaccinationSlot().getDate().compareTo(personAssignment.getReadyDate()) < 0)
                .penalizeLong("Ready date", ofHard(1),
                        personAssignment -> DAYS.between(personAssignment.getVaccinationSlot().getDate(),
                                personAssignment.getReadyDate()));
    }

    Constraint dueDate(ConstraintFactory constraintFactory) {
        // Typical usage 1: If a person is coming for their 2nd dose, don't inject it after the due day.
        // For example, Pfizer is due to be injected 3 months after the first dose.
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter(personAssignment -> personAssignment.getDueDate() != null
                        && personAssignment.getVaccinationSlot().getDate().compareTo(personAssignment.getDueDate()) > 0)
                .penalizeLong("Due date", ofHard(1),
                        personAssignment -> DAYS.between(personAssignment.getDueDate(),
                                personAssignment.getVaccinationSlot().getDate()));
    }

    // ************************************************************************
    // Medium constraints
    // ************************************************************************

    Constraint scheduleSecondOrLaterDosePeople(ConstraintFactory constraintFactory) {
        // If a person is coming for their 2nd dose, assign them to an appointment,
        // even before 1st dose healthcare workers and older people (although 2nd dosers will normally be that too).
        // This is to avoid a snowball effect on the backlog.
        return constraintFactory
                .from(PersonAssignment.class)
                // TODO filter for ideal date is earlier or equal to planning window last day
                .filter(personAssignment -> personAssignment.getDoseNumber() > 1 && personAssignment.getVaccinationSlot() == null)
                .penalizeLong("Schedule second (or later) dose people", ofSoft(0, 1),
                        personAssignment -> personAssignment.getDoseNumber() - 1);
    }

    Constraint scheduleHigherPriorityRatingPeople(ConstraintFactory constraintFactory) {
        // Assign healthcare workers and older people for an appointment.
        // Priority rating is a person's age augmented by a few hundred points if they're a healthcare worker.
        return constraintFactory
                .from(PersonAssignment.class)
                .filter(personAssignment -> personAssignment.getVaccinationSlot() == null)
                // This is softer than scheduleSecondOrLaterDosePeople()
                // to avoid creating a backlog of 2nd dose persons, that would grow too big to respect due dates.
                .penalizeLong("Schedule higher priority rating people", ofSoft(1, 1),
                        PersonAssignment::getPriorityRating);
    }

    // ************************************************************************
    // Soft constraints
    // ************************************************************************

    Constraint preferredVaccineType(ConstraintFactory constraintFactory) {
        // Typical usage: if a person can pick a favorite vaccine type
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter((personAssignment -> personAssignment.getPreferredVaccineType() != null
                        && personAssignment.getVaccinationSlot().getVaccineType() != personAssignment.getPreferredVaccineType()))
                .penalize("Preferred vaccine type", ofSoft(2, 1_000_000_000));
    }

    Constraint preferredVaccinationCenter(ConstraintFactory constraintFactory) {
        // Typical usage: if a person is coming for their 2nd dose, prefer the same vaccination center as their 1st dose.
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter((personAssignment -> personAssignment.getPreferredVaccinationCenter() != null
                        && personAssignment.getVaccinationSlot().getVaccinationCenter() != personAssignment.getPreferredVaccinationCenter()))
                // TODO ignore the distance cost instead
                .penalize("Preferred vaccination center", ofSoft(2, 1_000_000_000));
    }

    Constraint regretDistance(ConstraintFactory constraintFactory) {
        // Minimize the distance from each person's home location to their assigned vaccination center
        // subtracted by the distance to the nearest vaccination center
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .penalizeLong("Regret distance cost", ofSoft(2, 1),
                        personAssignment -> {
                            long regretDistance = personAssignment.getRegretDistanceTo(
                                    personAssignment.getVaccinationSlot().getVaccinationCenter());
                            // Penalize outliers more for fairness
                            return regretDistance * regretDistance;
                        });
    }

    Constraint idealDate(ConstraintFactory constraintFactory) {
        // Typical usage: If a person is coming for their 2nd dose, inject it on the ideal day.
        // For example, Pfizer is ideally injected 21 days after the first dose. Moderna after 28 days.
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                .filter(personAssignment -> personAssignment.getIdealDate() != null
                        && !personAssignment.getIdealDate().equals(personAssignment.getVaccinationSlot().getDate()))
                // This constraint is softer than distanceCost() to avoid sending people
                // half-way across the country just to be one day closer to their ideal date.
                .penalizeLong("Ideal date", ofSoft(3, 1),
                        personAssignment -> {
                            long daysDiff = DAYS.between(personAssignment.getIdealDate(),
                                    personAssignment.getVaccinationSlot().getDate());
                            // Penalize outliers more for fairness
                            return daysDiff * daysDiff;
                        });
    }

    Constraint higherPriorityRatingEarlier(ConstraintFactory constraintFactory) {
        // Assign healthcare workers and older people earlier in the planning window.
        // Priority rating is a person's age augmented by a few hundred points if they're a healthcare worker.
        // Differs from scheduleHigherPriorityRatingPeople(), which requires they be assigned
        return constraintFactory
                .from(PersonAssignment.class).filter(personAssignedFilter)
                // This constraint is softer than distanceCost() to avoid sending people
                // half-way across the country just to get their vaccine one day earlier.
                .penalizeLong("Higher priority rating earlier", ofSoft(4, 1),
                        personAssignment -> personAssignment.getPriorityRating()
                                * MINUTES.between(COVID_EPOCH, personAssignment.getVaccinationSlot().getStartDateTime()));
    }

}
