import sys
from pathlib import Path
from random import random, randint, choice
from typing import List

# Add framework src to path for imports
framework_src_path = Path(__file__).parent.parent.parent.parent / "src"
sys.path.insert(0, str(framework_src_path))

from demos.school_demo.src.individual import Student, Teacher, Counselor
from demos.school_demo.src.collectives import Neighborhood
from layer1_core.managers.population import Population


class SchoolPopulation(Population):
    """Population of students, teachers, and counselors in the school."""

    def __init__(self, param_dict: dict, neighborhoods: List[Neighborhood]):
        """Initialize the school population.

        Args:
            param_dict: Dictionary with studentCount, teacherCount, counselorCount
            neighborhoods: List of neighborhoods where students can live
        """
        super().__init__("school population")

        self.param_dict = param_dict
        self.neighborhoods = neighborhoods

        race_options = ["NHwhite", "NHBlack", "Hispanic", "Asian", "NAAHPI", "Multiple", "Other"]
        studentCount = param_dict["studentCount"]
        teacherCount = param_dict["teacherCount"]
        counselorCount = param_dict["counselorCount"]

        self.students = self.make_new_students(studentCount, race_options)
        self.teachers = self.make_new_teachers(teacherCount, race_options)
        self.counselors = self.make_new_counselors(counselorCount, race_options)
        self.agents = self.students + self.teachers + self.counselors

    def make_new_students(self, studentCount: int, race_options: List[str]) -> List[Student]:
        """Create a list of new students.

        Args:
            studentCount: Number of students to create
            race_options: List of possible race/ethnicity values

        Returns:
            List of Student instances
        """
        studentList: List[Student] = []
        for i in range(studentCount):
            age = randint(13, 18)
            race = choice(race_options)
            neighborhood = choice(self.neighborhoods) if self.neighborhoods else None
            new_student = Student(str(i), age, race, neighborhood)
            studentList.append(new_student)
        return studentList

    def make_new_teachers(self, teacherCount: int, race_options: List[str]) -> List[Teacher]:
        """Create a list of new teachers.

        Args:
            teacherCount: Number of teachers to create
            race_options: List of possible race/ethnicity values

        Returns:
            List of Teacher instances
        """
        teacherList: List[Teacher] = []
        for i in range(teacherCount):
            race = choice(race_options)
            new_teacher = Teacher(str(i), race)
            teacherList.append(new_teacher)
        return teacherList

    def make_new_counselors(self, counselorCount: int, race_options: List[str]) -> List[Counselor]:
        """Create a list of new counselors.

        Args:
            counselorCount: Number of counselors to create
            race_options: List of possible race/ethnicity values

        Returns:
            List of Counselor instances
        """
        counselorList: List[Counselor] = []
        for i in range(counselorCount):
            race = choice(race_options)
            new_counselor = Counselor(str(i), race)
            counselorList.append(new_counselor)
        return counselorList


