import sys
from pathlib import Path
from collections import Counter

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

from demos.school_demo.src.schoolWorld import SchoolWorld


def visualize_racial_makeup(students):
    """Create text-based visualization of racial makeup."""
    print("\n" + "="*60)
    print("RACIAL MAKEUP OF STUDENT BODY")
    print("="*60)

    # Count students by race
    race_counts = Counter(s.race for s in students)
    total = len(students)

    # Sort by count descending
    sorted_races = sorted(race_counts.items(), key=lambda x: x[1], reverse=True)

    # Find max count for scaling bars
    max_count = max(race_counts.values())
    bar_width = 40

    for race, count in sorted_races:
        percentage = (count / total) * 100
        bar_length = int((count / max_count) * bar_width)
        bar = "█" * bar_length
        print(f"{race:12} | {bar} {count:2d} ({percentage:5.1f}%)")

    print(f"\nTotal Students: {total}")


def visualize_disciplinary_records(students):
    """Create text-based visualization of disciplinary records."""
    print("\n" + "="*60)
    print("DISCIPLINARY RECORDS")
    print("="*60)

    # Collect stats
    students_with_records = [s for s in students if len(s.disciplinaryRecord) > 0]
    total_incidents = sum(len(s.disciplinaryRecord) for s in students)

    if total_incidents == 0:
        print("No disciplinary incidents recorded.")
        return

    print(f"Students with incidents: {len(students_with_records)}/{len(students)} ({100*len(students_with_records)/len(students):.1f}%)")
    print(f"Total incidents: {total_incidents}")
    print(f"Average incidents per student: {total_incidents/len(students):.2f}")

    # Distribution by number of incidents
    incident_counts = Counter(len(s.disciplinaryRecord) for s in students)

    print("\nIncident Distribution:")
    max_incidents = max(incident_counts.keys())
    for i in range(max_incidents + 1):
        count = incident_counts.get(i, 0)
        bar_length = int((count / len(students)) * 40)
        bar = "█" * bar_length
        print(f"  {i} incidents: {bar} {count:2d} students")

    # Severity analysis
    all_severities = [severity for s in students for severity in s.disciplinaryRecord]
    if all_severities:
        avg_severity = sum(all_severities) / len(all_severities)
        min_severity = min(all_severities)
        max_severity = max(all_severities)

        print(f"\nSeverity Statistics:")
        print(f"  Average: {avg_severity:.3f}")
        print(f"  Min: {min_severity:.3f}")
        print(f"  Max: {max_severity:.3f}")


def visualize_mood_changes(students, initial_moods):
    """Create text-based visualization of mood changes."""
    print("\n" + "="*60)
    print("STUDENT MOOD ANALYSIS")
    print("="*60)

    final_moods = [s.mood for s in students]
    mood_changes = [final - initial for initial, final in zip(initial_moods, final_moods)]

    # Overall stats
    avg_initial = sum(initial_moods) / len(initial_moods)
    avg_final = sum(final_moods) / len(final_moods)
    avg_change = sum(mood_changes) / len(mood_changes)

    print(f"Average mood: {avg_initial:.3f} → {avg_final:.3f} (change: {avg_change:+.3f})")

    # Count by mood change direction
    improved = sum(1 for c in mood_changes if c > 0)
    declined = sum(1 for c in mood_changes if c < 0)
    unchanged = sum(1 for c in mood_changes if c == 0)

    print(f"\nMood Changes:")
    print(f"  Improved:  {improved:2d} students ({100*improved/len(students):.1f}%)")
    print(f"  Declined:  {declined:2d} students ({100*declined/len(students):.1f}%)")
    print(f"  Unchanged: {unchanged:2d} students ({100*unchanged/len(students):.1f}%)")

    # Show distribution
    if declined > 0:
        print(f"\nMost declined: {min(mood_changes):+.3f}")
    if improved > 0:
        print(f"Most improved: {max(mood_changes):+.3f}")


def visualize_discipline_by_race(students):
    """Visualize disciplinary actions broken down by race."""
    print("\n" + "="*60)
    print("DISCIPLINARY ACTIONS BY RACE")
    print("="*60)

    # Group students by race
    students_by_race = {}
    for student in students:
        if student.race not in students_by_race:
            students_by_race[student.race] = []
        students_by_race[student.race].append(student)

    # Calculate stats for each race
    race_stats = []
    for race, race_students in students_by_race.items():
        total = len(race_students)
        with_incidents = sum(1 for s in race_students if len(s.disciplinaryRecord) > 0)
        total_incidents = sum(len(s.disciplinaryRecord) for s in race_students)
        avg_incidents = total_incidents / total if total > 0 else 0

        # Calculate average severity
        all_severities = [sev for s in race_students for sev in s.disciplinaryRecord]
        avg_severity = sum(all_severities) / len(all_severities) if all_severities else 0

        race_stats.append({
            'race': race,
            'total': total,
            'with_incidents': with_incidents,
            'pct_with_incidents': (with_incidents / total * 100) if total > 0 else 0,
            'avg_incidents': avg_incidents,
            'avg_severity': avg_severity
        })

    # Sort by average incidents descending
    race_stats.sort(key=lambda x: x['avg_incidents'], reverse=True)

    # Find max for scaling
    max_avg = max(s['avg_incidents'] for s in race_stats) if race_stats else 1
    bar_width = 30

    print("\nAverage Incidents per Student by Race:")
    for stat in race_stats:
        bar_length = int((stat['avg_incidents'] / max_avg) * bar_width) if max_avg > 0 else 0
        bar = "█" * bar_length
        print(f"{stat['race']:12} | {bar} {stat['avg_incidents']:.2f} ({stat['pct_with_incidents']:.1f}% had incidents)")

    print("\nAverage Severity by Race:")
    max_severity = max(s['avg_severity'] for s in race_stats) if race_stats else 1
    race_stats_by_severity = sorted(race_stats, key=lambda x: x['avg_severity'], reverse=True)

    for stat in race_stats_by_severity:
        if stat['avg_severity'] > 0:
            bar_length = int((stat['avg_severity'] / max_severity) * bar_width) if max_severity > 0 else 0
            bar = "█" * bar_length
            print(f"{stat['race']:12} | {bar} {stat['avg_severity']:.3f}")
        else:
            print(f"{stat['race']:12} | (no incidents)")


def visualize_mood_by_race(students, initial_moods):
    """Visualize mood changes broken down by race."""
    print("\n" + "="*60)
    print("MOOD CHANGES BY RACE")
    print("="*60)

    # Create mapping of initial moods
    initial_mood_map = {id(s): initial for s, initial in zip(students, initial_moods)}

    # Group students by race
    students_by_race = {}
    for student in students:
        if student.race not in students_by_race:
            students_by_race[student.race] = []
        students_by_race[student.race].append(student)

    # Calculate stats for each race
    race_mood_stats = []
    for race, race_students in students_by_race.items():
        initial_avg = sum(initial_mood_map[id(s)] for s in race_students) / len(race_students)
        final_avg = sum(s.mood for s in race_students) / len(race_students)
        mood_change = final_avg - initial_avg

        race_mood_stats.append({
            'race': race,
            'initial': initial_avg,
            'final': final_avg,
            'change': mood_change,
            'count': len(race_students)
        })

    # Sort by mood change (most negative first)
    race_mood_stats.sort(key=lambda x: x['change'])

    print("\nAverage Mood Change by Race:")
    print(f"{'Race':12} | {'Count':>5} | {'Initial':>7} → {'Final':>7} | {'Change':>8}")
    print("-" * 60)

    for stat in race_mood_stats:
        change_str = f"{stat['change']:+.4f}"
        print(f"{stat['race']:12} | {stat['count']:5d} | {stat['initial']:7.4f} → {stat['final']:7.4f} | {change_str:>8}")

    # Visual bar chart of final mood
    print("\nFinal Average Mood by Race:")
    max_mood = max(s['final'] for s in race_mood_stats)
    race_mood_stats_sorted = sorted(race_mood_stats, key=lambda x: x['final'], reverse=True)
    bar_width = 40

    for stat in race_mood_stats_sorted:
        bar_length = int((stat['final'] / max_mood) * bar_width) if max_mood > 0 else 0
        bar = "█" * bar_length
        print(f"{stat['race']:12} | {bar} {stat['final']:.4f}")


def visualize_friendships_by_race(friendship_manager, students):
    """Visualize friendship network patterns by race."""
    print("\n" + "="*60)
    print("FRIENDSHIP NETWORK BY RACE")
    print("="*60)

    # Get all friendships
    friendships = friendship_manager.friendships

    if not friendships:
        print("No friendships found.")
        return

    # Count total friendships
    print(f"Total Friendships: {len(friendships)}")

    # Analyze same-race vs cross-race friendships
    same_race = 0
    cross_race = 0

    # Track friendships by race pair
    race_pair_counts = {}

    for friendship in friendships:
        student1 = friendship.agent_1
        student2 = friendship.agent_2

        race1 = student1.race
        race2 = student2.race

        if race1 == race2:
            same_race += 1
            pair_key = f"{race1} ↔ {race1}"
        else:
            cross_race += 1
            # Normalize pair key so (A,B) and (B,A) are counted together
            pair_key = f"{min(race1, race2)} ↔ {max(race1, race2)}"

        race_pair_counts[pair_key] = race_pair_counts.get(pair_key, 0) + 1

    total = same_race + cross_race
    same_race_pct = (same_race / total * 100) if total > 0 else 0
    cross_race_pct = (cross_race / total * 100) if total > 0 else 0

    print(f"\nSame-race friendships: {same_race} ({same_race_pct:.1f}%)")
    print(f"Cross-race friendships: {cross_race} ({cross_race_pct:.1f}%)")

    # Show top race pairs
    print("\nTop Friendship Pairs:")
    sorted_pairs = sorted(race_pair_counts.items(), key=lambda x: x[1], reverse=True)
    max_count = max(race_pair_counts.values()) if race_pair_counts else 1
    bar_width = 30

    for i, (pair, count) in enumerate(sorted_pairs[:10]):  # Show top 10
        bar_length = int((count / max_count) * bar_width)
        bar = "█" * bar_length
        pct = (count / total * 100) if total > 0 else 0
        is_same_race = "↔" in pair and pair.split(" ↔ ")[0] == pair.split(" ↔ ")[1]
        marker = "●" if is_same_race else "○"
        print(f"{marker} {pair:30} | {bar} {count:2d} ({pct:5.1f}%)")

    print("\n● = same-race, ○ = cross-race")

    # Average friendship strength
    avg_strength = sum(f.strength for f in friendships) / len(friendships)
    print(f"\nAverage friendship strength: {avg_strength:.3f}")

    # Strength by same-race vs cross-race
    same_race_friendships = [f for f in friendships if f.agent_1.race == f.agent_2.race]
    cross_race_friendships = [f for f in friendships if f.agent_1.race != f.agent_2.race]

    if same_race_friendships:
        avg_same_race_strength = sum(f.strength for f in same_race_friendships) / len(same_race_friendships)
        print(f"  Same-race: {avg_same_race_strength:.3f}")

    if cross_race_friendships:
        avg_cross_race_strength = sum(f.strength for f in cross_race_friendships) / len(cross_race_friendships)
        print(f"  Cross-race: {avg_cross_race_strength:.3f}")

    # Count friendships per student
    print("\n" + "-"*60)
    print("FRIENDSHIP COUNTS BY STUDENT")
    print("-"*60)

    friendship_counts = {}
    for student in students:
        friendship_counts[student] = 0

    # Count each student's friendships
    for friendship in friendships:
        friendship_counts[friendship.agent_1] += 1
        friendship_counts[friendship.agent_2] += 1

    # Find students with most and least friendships
    sorted_students = sorted(friendship_counts.items(), key=lambda x: x[1], reverse=True)

    # Show top 10 most connected students
    print("\nMost Connected Students:")
    for i, (student, count) in enumerate(sorted_students[:10]):
        bar_length = int((count / max(friendship_counts.values())) * 20) if friendship_counts else 0
        bar = "█" * bar_length
        print(f"  {i+1:2d}. Student #{student.agent_id} ({student.race:12}): {bar} {count} friends")

    # Calculate average friendships by race
    print("\nAverage Friendships by Race:")
    race_friendship_totals = {}
    race_student_counts = {}

    for student, count in friendship_counts.items():
        race = student.race
        if race not in race_friendship_totals:
            race_friendship_totals[race] = 0
            race_student_counts[race] = 0
        race_friendship_totals[race] += count
        race_student_counts[race] += 1

    race_avg_friendships = []
    for race in race_friendship_totals:
        avg = race_friendship_totals[race] / race_student_counts[race]
        race_avg_friendships.append((race, avg, race_student_counts[race]))

    # Sort by average friendships
    race_avg_friendships.sort(key=lambda x: x[1], reverse=True)

    max_avg = max(avg for _, avg, _ in race_avg_friendships) if race_avg_friendships else 1
    for race, avg, count in race_avg_friendships:
        bar_length = int((avg / max_avg) * 30) if max_avg > 0 else 0
        bar = "█" * bar_length
        print(f"{race:12} | {bar} {avg:.2f} friends/student (n={count})")

    # Show distribution of friendship counts
    print("\nFriendship Distribution:")
    friend_count_distribution = Counter(friendship_counts.values())
    max_friends = max(friendship_counts.values()) if friendship_counts else 0

    for i in range(max_friends + 1):
        student_count = friend_count_distribution.get(i, 0)
        if student_count > 0:
            bar_length = int((student_count / len(students)) * 30)
            bar = "█" * bar_length
            pct = (student_count / len(students)) * 100
            print(f"  {i} friends: {bar} {student_count:2d} students ({pct:.1f}%)")


def visualize_friendship_mood_relationship(friendship_manager, students):
    """Visualize the relationship between friendships and mood."""
    print("\n" + "="*60)
    print("FRIENDSHIP-MOOD RELATIONSHIP")
    print("="*60)

    friendships = friendship_manager.friendships

    # Count friendships and calculate total strength per student
    friendship_counts = {}
    friendship_strengths = {}

    for student in students:
        friendship_counts[student] = 0
        friendship_strengths[student] = 0.0

    for friendship in friendships:
        friendship_counts[friendship.agent_1] += 1
        friendship_counts[friendship.agent_2] += 1
        friendship_strengths[friendship.agent_1] += friendship.strength
        friendship_strengths[friendship.agent_2] += friendship.strength

    # Group students by number of friends and calculate average mood
    print("\nAverage Mood by Number of Friends:")
    friends_to_moods = {}
    max_friends = max(friendship_counts.values()) if friendship_counts else 0

    for student in students:
        num_friends = friendship_counts[student]
        if num_friends not in friends_to_moods:
            friends_to_moods[num_friends] = []
        friends_to_moods[num_friends].append(student.mood)

    for i in range(max_friends + 1):
        if i in friends_to_moods:
            moods = friends_to_moods[i]
            avg_mood = sum(moods) / len(moods)
            count = len(moods)
            bar_length = int((avg_mood / 0.51) * 40)  # Scale assuming moods near 0.5
            bar = "█" * bar_length
            print(f"  {i} friends: {bar} {avg_mood:.4f} (n={count:2d})")

    # Correlation analysis - simple approach
    print("\nCorrelation Analysis:")

    # Calculate correlation between friend count and mood
    friend_counts_list = [friendship_counts[s] for s in students]
    moods_list = [s.mood for s in students]

    # Simple correlation coefficient
    n = len(students)
    mean_friends = sum(friend_counts_list) / n
    mean_mood = sum(moods_list) / n

    covariance = sum((friend_counts_list[i] - mean_friends) * (moods_list[i] - mean_mood) for i in range(n)) / n
    std_friends = (sum((f - mean_friends) ** 2 for f in friend_counts_list) / n) ** 0.5
    std_mood = (sum((m - mean_mood) ** 2 for m in moods_list) / n) ** 0.5

    correlation = covariance / (std_friends * std_mood) if (std_friends * std_mood) > 0 else 0

    print(f"  Friend count ↔ Mood: {correlation:+.3f}")

    if abs(correlation) < 0.1:
        strength_desc = "negligible"
    elif abs(correlation) < 0.3:
        strength_desc = "weak"
    elif abs(correlation) < 0.5:
        strength_desc = "moderate"
    else:
        strength_desc = "strong"

    direction = "positive" if correlation > 0 else "negative"
    print(f"  ({strength_desc} {direction} correlation)")

    # Analyze friendship strength vs mood
    print("\nAverage Mood by Total Friendship Strength:")

    # Group students by friendship strength quartiles
    students_with_friends = [s for s in students if friendship_strengths[s] > 0]

    if students_with_friends:
        # Divide into quartiles
        sorted_by_strength = sorted(students_with_friends, key=lambda s: friendship_strengths[s])
        quartile_size = len(sorted_by_strength) // 4

        if quartile_size > 0:
            quartiles = [
                ("Low strength", sorted_by_strength[:quartile_size]),
                ("Medium-low", sorted_by_strength[quartile_size:2*quartile_size]),
                ("Medium-high", sorted_by_strength[2*quartile_size:3*quartile_size]),
                ("High strength", sorted_by_strength[3*quartile_size:])
            ]

            for label, quartile_students in quartiles:
                avg_mood = sum(s.mood for s in quartile_students) / len(quartile_students)
                avg_strength = sum(friendship_strengths[s] for s in quartile_students) / len(quartile_students)
                bar_length = int((avg_mood / 0.51) * 30)
                bar = "█" * bar_length
                print(f"  {label:13}: {bar} {avg_mood:.4f} (avg strength: {avg_strength:.2f})")

    # Students with no friends
    students_no_friends = [s for s in students if friendship_counts[s] == 0]
    if students_no_friends:
        avg_mood_no_friends = sum(s.mood for s in students_no_friends) / len(students_no_friends)
        print(f"\n  No friends ({len(students_no_friends)} students): avg mood = {avg_mood_no_friends:.4f}")

    students_with_any_friends = [s for s in students if friendship_counts[s] > 0]
    if students_with_any_friends:
        avg_mood_with_friends = sum(s.mood for s in students_with_any_friends) / len(students_with_any_friends)
        print(f"  With friends ({len(students_with_any_friends)} students): avg mood = {avg_mood_with_friends:.4f}")

        mood_difference = avg_mood_with_friends - avg_mood_no_friends
        print(f"  Difference: {mood_difference:+.4f}")


def test_school_simulation_basic():
    """Test basic school simulation setup and execution."""

    # Setup parameters
    params = {
        "studentCount": 50,
        "teacherCount": 2,
        "counselorCount": 1,

    }


    # Create the world
    world = SchoolWorld(params)

    # Verify population was created
    assert world.population is not None
    assert len(world.population.students) == 50
    assert len(world.population.teachers) == 2
    assert len(world.population.counselors) == 1

    # Setup events
    world.setup_events()

    # Verify events were scheduled
    assert len(world.schedule.get_events()) > 0

    # Define stopping condition (stop after 100 steps)
    def stop_at_1000(schedule):
        return schedule.get_current_step() >= 1000

    # Run the simulation
    world.runner.run(stopping_condition=stop_at_1000)

    # Verify simulation ran
    assert world.schedule.get_current_step() >= 1000
    # Note: May have a few events remaining at the final step
    assert len(world.schedule.get_events()) < 10  # Most events executed

    print(f"✓ Simulation completed {world.schedule.get_current_step()} steps")
    print(f"✓ Population: {len(world.population.students)} students")


def test_student_mood_changes():
    """Test that student moods change during simulation."""

    params = {
        "studentCount": 50,
        "teacherCount": 2,
        "counselorCount": 1
    }

    world = SchoolWorld(params)

    # Record initial moods
    initial_moods = [s.mood for s in world.population.students]

    # Setup and run
    world.setup_events()
    world.runner.run(stopping_condition=lambda s: s.get_current_step() >= 50)

    # Check that at least some moods changed
    final_moods = [s.mood for s in world.population.students]

    # At least one student should have different mood
    mood_changed = any(initial != final for initial, final in zip(initial_moods, final_moods))

    assert mood_changed, "Student moods should change during simulation"

    print(f"✓ Student moods changed during simulation")


def test_with_visualizations():
    """Run simulation with detailed visualizations."""

    params = {
        "studentCount": 100,
        "teacherCount": 5,
        "counselorCount": 2
    }

    print("\nInitializing school simulation...")
    world = SchoolWorld(params)

    # Show initial demographics
    visualize_racial_makeup(world.population.students)

    # Record initial moods
    initial_moods = [s.mood for s in world.population.students]

    # Setup and run simulation
    print(f"\nRunning simulation for 1000 steps...")
    world.setup_events()
    world.runner.run(stopping_condition=lambda s: s.get_current_step() >= 1000)

    print(f"\n✓ Simulation completed at step {world.schedule.get_current_step()}")

    # Show overall results
    visualize_disciplinary_records(world.population.students)
    visualize_mood_changes(world.population.students, initial_moods)

    # Show results broken down by race
    visualize_discipline_by_race(world.population.students)
    visualize_mood_by_race(world.population.students, initial_moods)
    visualize_friendships_by_race(world.friendshipManager, world.population.students)
    visualize_friendship_mood_relationship(world.friendshipManager, world.population.students)


if __name__ == "__main__":
    print("Running school demo tests...\n")

    try:
        test_school_simulation_basic()
        print()
        test_student_mood_changes()
        print("\n✓ Basic tests passed!")

        # Run visualization test
        print("\n" + "="*60)
        print("RUNNING DEMO WITH VISUALIZATIONS")
        print("="*60)
        test_with_visualizations()

        print("\n" + "="*60)
        print("✓ ALL TESTS PASSED!")
        print("="*60)
    except AssertionError as e:
        print(f"\n✗ Test failed: {e}")
    except Exception as e:
        print(f"\n✗ Error: {e}")
        import traceback
        traceback.print_exc()