data_loader.py

# data_loader.py

"""
This module is responsible for converting raw input data into a list of
atomic 'Task' objects that the solver can schedule.

Responsibilities:
- Ingest lists of faculties, subjects, sections, and their allocations.
- Convert subject credits into the correct number of weekly Task instances.
- Ensure THEORY subjects create 1-hour tasks per credit.
- Ensure LAB, SOFTSKILL, and FORUM subjects create a single 2-hour task.
- Assign a unique ID to each task.
- Populate each task with its corresponding faculty, subject, and section.
- Correctly handle elective groups for multi-hour subjects.

This module performs NO scheduling logic, constraint creation, or optimization.
It is purely a data transformation and preparation layer.
"""

from dataclasses import dataclass
from typing import List, Optional

# Import the core data models
from models import Faculty, Subject, Section, Task, SubjectType

@dataclass(frozen=True)
class Allocation:
    """
    A simple dataclass to represent the raw input mapping of who teaches what
    to whom. This is a cleaner alternative to using tuples or dicts.
    """
    faculty_id: str
    subject_code: str
    section_id: str
    elective_group_id: Optional[str] = None


def prepare_scheduling_tasks(
    allocations: List[Allocation],
    faculties: List[Faculty],
    subjects: List[Subject],
    sections: List[Section]
) -> List[Task]:
    """
    Processes raw allocation data and generates a flat list of atomic Task
    objects ready for the solver.

    Args:
        allocations: A list of Allocation objects defining the teaching load.
        faculties: A list of all available Faculty objects.
        subjects: A list of all available Subject objects.
        sections: A list of all available Section objects.

    Returns:
        A list of Task objects, where each task is an atomic unit to be scheduled.
    """
    # Create lookup dictionaries for efficient access
    faculties_by_id = {f.id: f for f in faculties}
    subjects_by_code = {s.subject_code: s for s in subjects}
    sections_by_id = {s.section_id: s for s in sections}

    all_tasks: List[Task] = []

    for alloc in allocations:
        # Retrieve the full objects using the IDs from the allocation
        try:
            faculty = faculties_by_id[alloc.faculty_id]
            subject = subjects_by_code[alloc.subject_code]
            section = sections_by_id[alloc.section_id]
        except KeyError as e:
            print(f"Error: Invalid ID in allocation {alloc}. Missing key: {e}")
            continue

        # --- Task Generation Logic ---

        if subject.subject_type == SubjectType.THEORY:
            # For a theory subject, create one 1-hour task for each credit.
            for i in range(subject.credits):
                # *** CRITICAL FIX FOR ELECTIVES ***
                # For a 3-credit elective, we need groups like 'group_0', 'group_1', 'group_2'
                # to pair the correct hours across different sections.
                
                group_id = None
                if alloc.elective_group_id:
                    group_id = f"{alloc.elective_group_id}_{i}"
                elif subject.is_core:
                    # Treat Core Theory as a "Joint Class" for all batches in the section (e.g. 6a-E1 + 6a-E2)
                    # We generate a synthetic group ID: CORE_CS101_6A_0
                    parent_sec = section.section_id.split('-')[0]
                    group_id = f"CORE_{subject.subject_code}_{parent_sec}_{i}"

                task = Task(
                    task_id=f"{subject.subject_code}-{section.section_id}-{i}",
                    faculty=faculty,
                    subject=subject,
                    section=section,
                    duration=1,
                    elective_group_id=group_id
                )
                all_tasks.append(task)
        
        elif subject.subject_type in [SubjectType.LAB, SubjectType.SOFTSKILL, SubjectType.FORUM]:
            # For labs and other block sessions, create exactly ONE 2-hour task.
            # The elective group applies to the entire 2-hour block as a single unit.
            task = Task(
                task_id=f"{subject.subject_code}-{section.section_id}-BLOCK",
                faculty=faculty,
                subject=subject,
                section=section,
                duration=2,
                elective_group_id=alloc.elective_group_id
            )
            all_tasks.append(task)

    return all_tasks

# --- Example Usage ---
if __name__ == '__main__':
    # This block demonstrates how to use the prepare_scheduling_tasks function.
    # It will only run when this file is executed directly.

    # 1. Define sample data using the core models
    faculty1 = Faculty(id="F001", name="Dr. Smith", designation="Professor", max_hours_per_week=10)
    faculty2 = Faculty(id="F002", name="Dr. Jones", designation="Asst. Professor", max_hours_per_week=12)

    subject_theory = Subject(subject_code="CS101", name="Intro to CS", credits=4, subject_type=SubjectType.THEORY)
    subject_lab = Subject(subject_code="CS101L", name="CS Lab", credits=1, subject_type=SubjectType.LAB)
    subject_elective = Subject(subject_code="CS555", name="Advanced AI", credits=3, subject_type=SubjectType.THEORY)


    section_a = Section(section_id="5A", semester=5, student_strength=60)
    section_b = Section(section_id="5B", semester=5, student_strength=62)

    # 2. Define the teaching allocations
    sample_allocations = [
        Allocation(faculty_id="F001", subject_code="CS101", section_id="5A"),
        Allocation(faculty_id="F002", subject_code="CS101L", section_id="5A"),
        # Elective subject taught by the same faculty to both sections
        Allocation(faculty_id="F001", subject_code="CS555", section_id="5A", elective_group_id="ELEC01"),
        Allocation(faculty_id="F001", subject_code="CS555", section_id="5B", elective_group_id="ELEC01"),
    ]

    # 3. Call the data loader function
    generated_tasks = prepare_scheduling_tasks(
        allocations=sample_allocations,
        faculties=[faculty1, faculty2],
        subjects=[subject_theory, subject_lab, subject_elective],
        sections=[section_a, section_b]
    )

    # 4. Print the results to verify
    print(f"--- Generated {len(generated_tasks)} Tasks ---\n")
    for task in generated_tasks:
        print(f"Task ID: {task.task_id}")
        print(f"  Subject: {task.subject.name} ({task.subject.subject_type.name})")
        print(f"  Faculty: {task.faculty.name}")
        print(f"  Section: {task.section.section_id}")
        print(f"  Duration: {task.duration} hour(s)")
        if task.elective_group_id:
            print(f"  Elective Group: {task.elective_group_id}")
        print("-" * 20)

    # Expected Output:
    # Elective tasks for CS555 will now have group IDs like "ELEC01_0", "ELEC01_1", "ELEC01_2"