from solverforge_legacy.solver.score import (
    constraint_provider,
    HardMediumSoftScore,
    Joiners,
    ConstraintFactory,
    Constraint,
)

from .domain import (
    Attendance,
    MeetingAssignment,
    PreferredAttendance,
    RequiredAttendance,
    Room,
    TimeGrain,
)


@constraint_provider
def define_constraints(constraint_factory: ConstraintFactory):
    """
    Defines all constraints for the meeting scheduling problem, organized by priority (hard, medium, soft).

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        List[Constraint]: All defined constraints.
    """
    return [
        # Hard constraints
        room_conflict(constraint_factory),
        avoid_overtime(constraint_factory),
        required_attendance_conflict(constraint_factory),
        required_room_capacity(constraint_factory),
        start_and_end_on_same_day(constraint_factory),
        # Medium constraints
        required_and_preferred_attendance_conflict(constraint_factory),
        preferred_attendance_conflict(constraint_factory),
        # Soft constraints
        do_meetings_as_soon_as_possible(constraint_factory),
        one_break_between_consecutive_meetings(constraint_factory),
        overlapping_meetings(constraint_factory),
        assign_larger_rooms_first(constraint_factory),
        room_stability(constraint_factory),
    ]


# ************************************************************************
# Hard constraints
# ************************************************************************


def room_conflict(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Hard constraint: Prevents overlapping meetings in the same room.

    Penalizes pairs of meetings scheduled in the same room whose time slots overlap, with penalty proportional to the overlap duration.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_unique_pair(
            MeetingAssignment,
            Joiners.equal(lambda assignment: assignment.room),
            Joiners.overlapping(
                lambda assignment: assignment.get_grain_index(),
                lambda assignment: assignment.get_last_time_grain_index() + 1,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_HARD,
            lambda left_assignment,
            right_assignment: right_assignment.calculate_overlap(left_assignment),
        )
        .as_constraint("Room conflict")
    )


def avoid_overtime(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Hard constraint: Prevents meetings from extending beyond available time slots.

    Penalizes meetings that end after the last available time grain, based on how far they extend beyond the schedule.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.starting_time_grain
            is not None
        )
        .if_not_exists(
            TimeGrain,
            Joiners.equal(
                lambda assignment: assignment.get_last_time_grain_index(),
                lambda time_grain: time_grain.grain_index,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_HARD,
            lambda meeting_assignment: meeting_assignment.get_last_time_grain_index(),
        )
        .as_constraint("Don't go in overtime")
    )


def required_attendance_conflict(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Hard constraint: Prevents required attendees from having overlapping meetings.

    Penalizes when a person required at multiple meetings is scheduled for overlapping meetings, proportional to the overlap duration.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_unique_pair(
            RequiredAttendance, Joiners.equal(lambda attendance: attendance.person)
        )
        .join(
            MeetingAssignment,
            Joiners.equal(
                lambda left_required, right_required: left_required.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
        )
        .join(
            MeetingAssignment,
            Joiners.equal(
                lambda left_required,
                right_required,
                left_assignment: right_required.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
            Joiners.overlapping(
                lambda attendee1, attendee2, assignment: assignment.get_grain_index(),
                lambda attendee1,
                attendee2,
                assignment: assignment.get_last_time_grain_index() + 1,
                lambda assignment: assignment.get_grain_index(),
                lambda assignment: assignment.get_last_time_grain_index() + 1,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_HARD,
            lambda left_required,
            right_required,
            left_assignment,
            right_assignment: right_assignment.calculate_overlap(left_assignment),
        )
        .as_constraint("Required attendance conflict")
    )


def required_room_capacity(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Hard constraint: Ensures rooms have enough capacity for required attendees.

    Penalizes meetings assigned to rooms with insufficient capacity, proportional to the shortfall.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.get_required_capacity()
            > meeting_assignment.get_room_capacity()
        )
        .penalize(
            HardMediumSoftScore.ONE_HARD,
            lambda meeting_assignment: meeting_assignment.get_required_capacity()
            - meeting_assignment.get_room_capacity(),
        )
        .as_constraint("Required room capacity")
    )


def start_and_end_on_same_day(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Hard constraint: Ensures meetings start and end on the same day.

    Penalizes meetings that span multiple days.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.starting_time_grain
            is not None
        )
        .join(
            TimeGrain,
            Joiners.equal(
                lambda meeting_assignment: meeting_assignment.get_last_time_grain_index(),
                lambda time_grain: time_grain.grain_index,
            ),
            Joiners.filtering(
                lambda meeting_assignment,
                time_grain: meeting_assignment.starting_time_grain.day_of_year
                != time_grain.day_of_year
            ),
        )
        .penalize(HardMediumSoftScore.ONE_HARD)
        .as_constraint("Start and end on same day")
    )


# ************************************************************************
# Medium constraints
# ************************************************************************


def required_and_preferred_attendance_conflict(
    constraint_factory: ConstraintFactory,
) -> Constraint:
    """
    Medium constraint: Discourages conflicts between required and preferred attendance for the same person.

    Penalizes when a person required at one meeting and preferred at another is scheduled for overlapping meetings, proportional to the overlap duration.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each(RequiredAttendance)
        .join(
            PreferredAttendance,
            Joiners.equal(
                lambda required: required.person, lambda preferred: preferred.person
            ),
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda assignment: assignment.starting_time_grain is not None
            ),
            Joiners.equal(
                lambda required, preferred: required.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda assignment: assignment.starting_time_grain is not None
            ),
            Joiners.equal(
                lambda required, preferred, left_assignment: preferred.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
            Joiners.overlapping(
                lambda required, preferred, assignment: assignment.get_grain_index(),
                lambda required,
                preferred,
                assignment: assignment.get_last_time_grain_index() + 1,
                lambda assignment: assignment.get_grain_index(),
                lambda assignment: assignment.get_last_time_grain_index() + 1,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_MEDIUM,
            lambda required,
            preferred,
            left_assignment,
            right_assignment: right_assignment.calculate_overlap(left_assignment),
        )
        .as_constraint("Required and preferred attendance conflict")
    )


def preferred_attendance_conflict(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Medium constraint: Discourages conflicts between preferred attendees.

    Penalizes when a person preferred at multiple meetings is scheduled for overlapping meetings, proportional to the overlap duration.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_unique_pair(
            PreferredAttendance, Joiners.equal(lambda attendance: attendance.person)
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda assignment: assignment.starting_time_grain is not None
            ),
            Joiners.equal(
                lambda left_attendance, right_attendance: left_attendance.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda assignment: assignment.starting_time_grain is not None
            ),
            Joiners.equal(
                lambda left_attendance,
                right_attendance,
                left_assignment: right_attendance.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
            Joiners.overlapping(
                lambda attendee1, attendee2, assignment: assignment.get_grain_index(),
                lambda attendee1,
                attendee2,
                assignment: assignment.get_last_time_grain_index() + 1,
                lambda assignment: assignment.get_grain_index(),
                lambda assignment: assignment.get_last_time_grain_index() + 1,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_MEDIUM,
            lambda left_attendance,
            right_attendance,
            left_assignment,
            right_assignment: right_assignment.calculate_overlap(left_assignment),
        )
        .as_constraint("Preferred attendance conflict")
    )


# ************************************************************************
# Soft constraints
# ************************************************************************


def do_meetings_as_soon_as_possible(
    constraint_factory: ConstraintFactory,
) -> Constraint:
    """
    Soft constraint: Encourages scheduling meetings earlier in the available time slots.

    Penalizes meetings scheduled later in the available time grains, proportional to their end time.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.starting_time_grain
            is not None
        )
        .penalize(
            HardMediumSoftScore.ONE_SOFT,
            lambda meeting_assignment: meeting_assignment.get_last_time_grain_index(),
        )
        .as_constraint("Do all meetings as soon as possible")
    )


def one_break_between_consecutive_meetings(
    constraint_factory: ConstraintFactory,
) -> Constraint:
    """
    Soft constraint: Penalizes consecutive meetings without a break.

    Penalizes pairs of meetings that are scheduled consecutively without at least one time grain break between them.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.starting_time_grain
            is not None
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda assignment: assignment.starting_time_grain is not None
            ),
            Joiners.equal(
                lambda left_assignment: left_assignment.get_last_time_grain_index(),
                lambda right_assignment: right_assignment.get_grain_index() - 1,
            ),
        )
        .penalize(HardMediumSoftScore.of_soft(100))
        .as_constraint("One TimeGrain break between two consecutive meetings")
    )


def overlapping_meetings(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Soft constraint: Discourages overlapping meetings, even in different rooms.

    Penalizes pairs of meetings that overlap in time, regardless of room, proportional to the overlap duration.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(
            lambda meeting_assignment: meeting_assignment.starting_time_grain
            is not None
        )
        .join(
            constraint_factory.for_each_including_unassigned(MeetingAssignment).filter(
                lambda meeting_assignment: meeting_assignment.starting_time_grain
                is not None
            ),
            Joiners.greater_than(
                lambda left_assignment: left_assignment.meeting.id,
                lambda right_assignment: right_assignment.meeting.id,
            ),
            Joiners.overlapping(
                lambda assignment: assignment.get_grain_index(),
                lambda assignment: assignment.get_last_time_grain_index() + 1,
            ),
        )
        .penalize(
            HardMediumSoftScore.of_soft(10),
            lambda left_assignment, right_assignment: left_assignment.calculate_overlap(
                right_assignment
            ),
        )
        .as_constraint("Overlapping meetings")
    )


def assign_larger_rooms_first(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Soft constraint: Penalizes using smaller rooms when larger rooms are available.

    Penalizes when a meeting is assigned to a room while larger rooms exist, proportional to the capacity difference.

    Args:
        constraint_factory (ConstraintFactory): The constraint factory.
    Returns:
        Constraint: The defined constraint.
    """
    return (
        constraint_factory.for_each_including_unassigned(MeetingAssignment)
        .filter(lambda meeting_assignment: meeting_assignment.room is not None)
        .join(
            Room,
            Joiners.less_than(
                lambda meeting_assignment: meeting_assignment.get_room_capacity(),
                lambda room: room.capacity,
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_SOFT,
            lambda meeting_assignment, room: room.capacity
            - meeting_assignment.get_room_capacity(),
        )
        .as_constraint("Assign larger rooms first")
    )


def room_stability(constraint_factory: ConstraintFactory) -> Constraint:
    """
    Soft constraint: Encourages room stability for people attending multiple meetings.

    Penalizes when a person attends meetings in different rooms that are close in time.
    Uses weighted penalty: back-to-back room switches cost more than switches with gaps.
    """
    return (
        constraint_factory.for_each(Attendance)
        .join(
            Attendance,
            Joiners.equal(lambda a: a.person),
            Joiners.filtering(lambda left, right: left.meeting_id != right.meeting_id),
        )
        .join(
            MeetingAssignment,
            Joiners.equal(
                lambda left, right: left.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
        )
        .join(
            MeetingAssignment,
            Joiners.equal(
                lambda left, right, left_assignment: right.meeting_id,
                lambda assignment: assignment.meeting.id,
            ),
            Joiners.less_than(
                lambda left, right, left_assignment: left_assignment.get_grain_index(),
                lambda assignment: assignment.get_grain_index(),
            ),
            Joiners.filtering(
                lambda left, right, left_assignment, right_assignment: left_assignment.room != right_assignment.room
            ),
            Joiners.filtering(
                lambda left, right, left_assignment, right_assignment: right_assignment.get_grain_index()
                - left_assignment.meeting.duration_in_grains
                - left_assignment.get_grain_index()
                <= 2
            ),
        )
        .penalize(
            HardMediumSoftScore.ONE_SOFT,
            lambda left, right, left_assignment, right_assignment: 3 - (
                right_assignment.get_grain_index()
                - left_assignment.meeting.duration_in_grains
                - left_assignment.get_grain_index()
            ),
        )
        .as_constraint("Room stability")
    )