"""Utility functions and classes for the school simulation."""

import sys
from collections import Counter
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from demos.school_demo.src.schoolWorld import SchoolWorld


class TeeOutput:
    """Write output to both console and file simultaneously.

    Usage:
        with TeeOutput('output.txt'):
            print("This goes to both console and file")
    """

    def __init__(self, filepath):
        """Initialize TeeOutput.

        Args:
            filepath: Path to the output file
        """
        self.file = open(filepath, 'w')
        self.stdout = sys.stdout

    def write(self, message):
        """Write message to both stdout and file."""
        self.stdout.write(message)
        self.file.write(message)

    def flush(self):
        """Flush both stdout and file buffers."""
        self.stdout.flush()
        self.file.flush()

    def close(self):
        """Close the output file."""
        self.file.close()

    def __enter__(self):
        """Enter context manager - redirect stdout."""
        sys.stdout = self
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager - restore stdout and close file."""
        sys.stdout = self.stdout
        self.close()


def print_header(text: str):
    """Print a formatted header.

    Args:
        text: Header text to display
    """
    print("\n" + "=" * 70)
    print(text)
    print("=" * 70)


def print_initial_state(world: 'SchoolWorld'):
    """Print the initial state of the simulation.

    Args:
        world: The SchoolWorld instance
    """
    print_header("SCHOOL SIMULATION - INITIAL STATE")

    print(f"\nPopulation:")
    print(f"  Students:   {len(world.population.students)}")
    print(f"  Teachers:   {len(world.population.teachers)}")
    print(f"  Counselors: {len(world.population.counselors)}")

    print(f"\nNeighborhoods: {len(world.neighborhoods)}")

    print(f"\nServices registered:")
    print(f"  - FriendshipManager")
    print(f"  - DisciplineService")
    print(f"  - CounselingService")
    print(f"  - HousingService")

    # Show racial makeup
    race_counts = Counter(s.race for s in world.population.students)
    print(f"\nStudent demographics:")
    for race, count in sorted(race_counts.items(), key=lambda x: x[1], reverse=True):
        pct = (count / len(world.population.students)) * 100
        print(f"  {race:12}: {count:3d} ({pct:5.1f}%)")

    # Show initial friendships
    print(f"\nInitial friendships: {len(world.friendshipManager.friendships)}")

    print(f"\nScheduled events: {len(world.schedule.get_events())}")


def print_final_state(world: 'SchoolWorld'):
    """Print the final state after simulation.

    Args:
        world: The SchoolWorld instance after simulation
    """
    print_header("SIMULATION COMPLETE - FINAL STATE")

    students = world.population.students

    # Mood statistics
    moods = [s.mood for s in students]
    avg_mood = sum(moods) / len(moods)
    min_mood = min(moods)
    max_mood = max(moods)

    print(f"\nStudent Mood Statistics:")
    print(f"  Average: {avg_mood:.3f}")
    print(f"  Min:     {min_mood:.3f}")
    print(f"  Max:     {max_mood:.3f}")

    # Disciplinary records
    total_incidents = sum(len(s.disciplinaryRecord) for s in students)
    students_with_incidents = sum(1 for s in students if len(s.disciplinaryRecord) > 0)

    print(f"\nDisciplinary Statistics:")
    print(f"  Total incidents:           {total_incidents}")
    print(f"  Students with incidents:   {students_with_incidents}/{len(students)}")
    print(f"  Average per student:       {total_incidents/len(students):.2f}")

    # Friendship statistics
    friendships = world.friendshipManager.friendships
    if friendships:
        avg_strength = sum(f.strength for f in friendships) / len(friendships)
        print(f"\nFriendship Statistics:")
        print(f"  Total friendships:     {len(friendships)}")
        print(f"  Average strength:      {avg_strength:.3f}")

        # Count friends per student
        friend_counts = {}
        for s in students:
            friend_counts[s] = 0
        for f in friendships:
            friend_counts[f.agent_1] += 1
            friend_counts[f.agent_2] += 1

        avg_friends = sum(friend_counts.values()) / len(students)
        max_friends = max(friend_counts.values())
        no_friends = sum(1 for c in friend_counts.values() if c == 0)

        print(f"  Average friends/student: {avg_friends:.2f}")
        print(f"  Max friends (any student): {max_friends}")
        print(f"  Students with no friends: {no_friends}")

    print(f"\nFinal step: {world.schedule.get_current_step()}")
    print(f"Remaining events: {len(world.schedule.get_events())}")
