"""
Gale-Shapley Algorithm Implementation

Python port of the C implementation from the original project.
Implements the Stable Marriage Problem for roommate matching and room allocation.

Original C version: https://github.com/Harshwardhan-Deshmukh03/Roommate-allocation-using-Gale-Shapley-Algorithm.git
"""

from typing import Dict, List, Tuple, Optional


def gale_shapley(pref_matrix: List[List[int]], n: int) -> Dict[int, int]:
    """
    Run the Gale-Shapley algorithm for stable matching.

    This is a direct Python translation of the C galeShapley() function.
    Students 0..n-1 are "proposers" and students n..2n-1 are "acceptors".

    Args:
        pref_matrix: 2D preference list. pref_matrix[i] is the ordered preference 
                     list for person i. For proposers (0..n-1), preferences are 
                     indices in the acceptor range (n..2n-1). For acceptors, 
                     preferences are proposer indices (0..n-1).
        n: Number of pairs (half the total participants).

    Returns:
        Dictionary mapping each person to their matched partner.
    """
    # match[i] = partner of person i, -1 means unmatched
    match = [-1] * (2 * n)
    free_count = n

    # Track which preference index each proposer is up to
    next_proposal = [0] * n

    while free_count > 0:
        # Find the first free proposer
        m = -1
        for i in range(n):
            if match[i] == -1:
                m = i
                break

        if m == -1:
            break

        # Proposer m proposes to their next preferred acceptor
        while next_proposal[m] < n:
            w = pref_matrix[m][next_proposal[m]]
            next_proposal[m] += 1

            if match[w] == -1:
                # w is free, accept the proposal
                match[w] = m
                match[m] = w
                free_count -= 1
                break
            else:
                # w is already matched, check if w prefers m over current partner
                m1 = match[w]
                if _prefers_new_over_current(pref_matrix, w, m, m1):
                    # w prefers m over m1 — switch
                    match[w] = m
                    match[m] = w
                    match[m1] = -1  # m1 becomes free
                    break
                # else w rejects m, m tries next preference

    return match


def _prefers_new_over_current(
    pref_matrix: List[List[int]], w: int, m_new: int, m_current: int
) -> bool:
    """
    Check if acceptor w prefers m_new over m_current.
    
    Mirrors the C function wPrefersm1Overm() but with clearer naming.
    Returns True if w prefers m_new over m_current.
    """
    for pref in pref_matrix[w]:
        if pref == m_new:
            return True
        if pref == m_current:
            return False
    return False


def run_roommate_matching(students: List[dict]) -> List[Tuple[int, int]]:
    """
    Stage 1: Match students into roommate pairs using Gale-Shapley.

    Args:
        students: List of student dicts with 'id', 'name', 'cgpa', 'pref_roommate'.

    Returns:
        List of (student_id_1, student_id_2) roommate pairs.
    """
    n_students = len(students)
    if n_students < 2 or n_students % 2 != 0:
        raise ValueError(
            f"Need an even number of students (≥ 2). Got {n_students}."
        )

    n = n_students // 2  # Number of pairs

    # Build preference matrix: proposers are 0..n-1, acceptors are n..2n-1
    # Each student's pref_roommate contains IDs of other students they prefer
    pref_matrix = []
    for student in students:
        prefs = student["pref_roommate"]
        # Filter out only valid preferences (should be IDs of other students)
        # For proposers (0..n-1): their prefs should reference acceptors (n..2n-1)
        # For acceptors (n..2n-1): their prefs should reference proposers (0..n-1)
        sid = student["id"]
        if sid < n:
            # Proposer: filter prefs to only include acceptor IDs (n..2n-1)
            filtered = [p for p in prefs if n <= p < 2 * n]
        else:
            # Acceptor: filter prefs to only include proposer IDs (0..n-1)
            filtered = [p for p in prefs if 0 <= p < n]
        pref_matrix.append(filtered)

    # Run Gale-Shapley
    match = gale_shapley(pref_matrix, n)

    # Extract unique pairs (only from proposer side to avoid duplicates)
    pairs = []
    seen = set()
    for i in range(n):
        partner = match[i]
        if partner != -1 and i not in seen and partner not in seen:
            pairs.append((i, partner))
            seen.add(i)
            seen.add(partner)

    return pairs


def run_room_allocation(
    students: List[dict],
    roommate_pairs: List[Tuple[int, int]],
    rooms: List[dict],
) -> List[dict]:
    """
    Stage 2: Allocate rooms to roommate pairs based on CGPA ranking.

    Higher CGPA pairs get priority in room selection (Gale-Shapley on rooms).

    Args:
        students: List of student dicts with 'id', 'name', 'cgpa', 'pref_room'.
        roommate_pairs: List of (id1, id2) roommate pairs from Stage 1.
        rooms: List of room dicts with 'room_id' and 'room_number'.

    Returns:
        List of allocation dicts: {roommate1, roommate2, room_number, room_id, pair_cgpa}
    """
    n = len(roommate_pairs)
    n_rooms = len(rooms)

    if n_rooms < n:
        raise ValueError(
            f"Not enough rooms ({n_rooms}) for {n} pairs."
        )

    # Build student lookup
    student_map = {s["id"]: s for s in students}

    # Rank pairs by the higher CGPA in each pair (descending)
    pair_cgpas = []
    for id1, id2 in roommate_pairs:
        cgpa1 = student_map[id1]["cgpa"]
        cgpa2 = student_map[id2]["cgpa"]
        max_cgpa = max(cgpa1, cgpa2)
        pair_cgpas.append((max_cgpa, id1, id2))

    # Sort descending by max CGPA
    pair_cgpas.sort(key=lambda x: x[0], reverse=True)

    # Build room preference matrix for Gale-Shapley
    # Proposers: ranked pairs (0..n-1) → these map to pair_cgpas indices
    # Acceptors: rooms (n..2n-1) → these map to rooms indices (offset by n)
    room_id_to_index = {rooms[j]["room_id"]: j + n for j in range(n)}

    pref_matrix = [[] for _ in range(2 * n)]

    # For each ranked pair, get the higher-CGPA student's room preferences
    for rank_idx, (max_cgpa, id1, id2) in enumerate(pair_cgpas):
        # Use the higher CGPA student's room preferences
        if student_map[id1]["cgpa"] >= student_map[id2]["cgpa"]:
            prefs = student_map[id1].get("pref_room", [])
        else:
            prefs = student_map[id2].get("pref_room", [])

        # Map room IDs to room indices (offset by n for acceptor range)
        mapped_prefs = []
        for room_id in prefs:
            if room_id in room_id_to_index:
                mapped_prefs.append(room_id_to_index[room_id])
        pref_matrix[rank_idx] = mapped_prefs

    # Rooms accept any student in order (no real preference)
    for j in range(n):
        pref_matrix[n + j] = list(range(n))

    # Run Gale-Shapley for room allocation
    match = gale_shapley(pref_matrix, n)

    # Build index-to-room mapping
    index_to_room = {j + n: rooms[j] for j in range(n)}

    # Build final allocation
    allocations = []
    for rank_idx, (max_cgpa, id1, id2) in enumerate(pair_cgpas):
        room_index = match[rank_idx]
        room = index_to_room.get(room_index, {"room_number": "N/A", "room_id": -1})

        allocations.append({
            "roommate1_id": id1,
            "roommate1_name": student_map[id1]["name"],
            "roommate1_cgpa": student_map[id1]["cgpa"],
            "roommate2_id": id2,
            "roommate2_name": student_map[id2]["name"],
            "roommate2_cgpa": student_map[id2]["cgpa"],
            "room_number": room["room_number"],
            "room_id": room["room_id"],
            "pair_max_cgpa": max_cgpa,
        })

    return allocations


def run_full_allocation(students: List[dict], rooms: List[dict]) -> List[dict]:
    """
    Run the complete two-stage allocation pipeline.

    Stage 1: Gale-Shapley for roommate matching.
    Stage 2: CGPA-ranked Gale-Shapley for room allocation.

    Args:
        students: List of student dicts.
        rooms: List of room dicts.

    Returns:
        List of allocation result dicts.
    """
    # Stage 1: Roommate matching
    roommate_pairs = run_roommate_matching(students)

    # Stage 2: Room allocation
    allocations = run_room_allocation(students, roommate_pairs, rooms)

    return allocations