h-school/src/HSchool.People/YearlyIntake.cs

namespace HSchool.People;

/// <summary>
/// 1 September: promote every class, graduate the oldest year present, fill the vacated rooms
/// with a new first year. Capacity is the map — intake replaces graduates, it does not grow.
/// </summary>
public static class YearlyIntake
{
    public const int MaxChildrenInFamily = 3;

    /// <summary>
    /// September 1sts strictly after <paramref name="before"/> and at or before <paramref name="after"/>.
    /// A school created on 1 September does not intake that morning — its roster is already that year.
    /// </summary>
    public static IEnumerable<DateTime> DatesBetween(DateTime before, DateTime after)
    {
        var utcBefore = DateTime.SpecifyKind(before, DateTimeKind.Utc);
        var utcAfter = DateTime.SpecifyKind(after, DateTimeKind.Utc);
        if (utcAfter <= utcBefore)
        {
            yield break;
        }

        var date = new DateTime(utcBefore.Year, 9, 1, 0, 0, 0, DateTimeKind.Utc);
        if (date <= utcBefore)
        {
            date = date.AddYears(1);
        }

        while (date <= utcAfter)
        {
            yield return date;
            date = date.AddYears(1);
        }
    }

    public static Roster Apply(
        DefCatalog catalog,
        Roster roster,
        int schoolSeed,
        string nameSetId,
        DateTime asOf,
        string? nativeLanguage = null)
    {
        ArgumentNullException.ThrowIfNull(catalog);
        ArgumentNullException.ThrowIfNull(roster);
        ArgumentException.ThrowIfNullOrWhiteSpace(nameSetId);

        if (!catalog.NameSets.TryGetValue(nameSetId, out var names))
        {
            throw new ArgumentException($"Unknown name set '{nameSetId}'.", nameof(nameSetId));
        }

        var native = NativeLanguages.Pick(names, schoolSeed, nativeLanguage, rollIfOmitted: false);

        if (roster.Classes.Count == 0)
        {
            return roster;
        }

        var when = DateTime.SpecifyKind(asOf, DateTimeKind.Utc);
        var yearStart = SchoolYears.StartOn(when);
        var maxYear = roster.Classes.Max(schoolClass => schoolClass.Year);
        var graduatingIds = roster.Classes
            .Where(schoolClass => schoolClass.Year == maxYear)
            .SelectMany(schoolClass => schoolClass.PupilIds)
            .ToHashSet(StringComparer.Ordinal);

        var remainingPeople = roster.People
            .Where(person => !graduatingIds.Contains(person.Id))
            .ToDictionary(person => person.Id, StringComparer.Ordinal);

        var families = new List<Family>();
        foreach (var family in roster.Families.OrderBy(candidate => IndexOf(candidate.Id)))
        {
            var childIds = family.ChildIds.Where(id => !graduatingIds.Contains(id)).ToArray();
            var parentIds = family.ParentIds.Where(id => remainingPeople.ContainsKey(id)).ToArray();
            var keepParents = childIds.Length > 0
                ? parentIds
                : parentIds.Where(id => remainingPeople[id].IsStaff).ToArray();

            foreach (var parentId in parentIds)
            {
                if (!keepParents.Contains(parentId, StringComparer.Ordinal))
                {
                    remainingPeople.Remove(parentId);
                }
            }

            if (keepParents.Length == 0 && childIds.Length == 0)
            {
                continue;
            }

            // The counter travels with the family: a graduate's id must never come back. So do the
            // father's name and the surname — a single-mother household still names its children
            // after the father, and he is not in the roster to be asked.
            families.Add(family with
            {
                ParentIds = keepParents,
                ChildIds = childIds,
                NextChild = family.NextChildIndex,
            });
            foreach (var parentId in keepParents)
            {
                var parent = remainingPeople[parentId];
                remainingPeople[parentId] = parent with { IsParent = childIds.Length > 0 };
            }
        }

        var classes = new SchoolClass[roster.Classes.Count];
        var newSeats = new List<PupilSeat>();
        var incumbents = remainingPeople
            .Where(pair => pair.Value.IsStudent)
            .Select(pair => pair.Key)
            .ToHashSet(StringComparer.Ordinal);
        for (var i = 0; i < roster.Classes.Count; i++)
        {
            var schoolClass = roster.Classes[i];
            if (schoolClass.Year == maxYear)
            {
                classes[i] = schoolClass with { Year = 1, PupilIds = [] };
                for (var seat = 0; seat < schoolClass.Capacity; seat++)
                {
                    newSeats.Add(new PupilSeat(schoolClass.Id, schoolClass.RoomId, Year: 1, schoolClass.Letter));
                }
            }
            else
            {
                var kept = schoolClass.PupilIds.Where(id => remainingPeople.ContainsKey(id)).ToArray();
                classes[i] = schoolClass with { Year = schoolClass.Year + 1, PupilIds = kept };
            }
        }

        FillFirstYear(
            catalog,
            names,
            schoolSeed,
            yearStart,
            when,
            families,
            remainingPeople,
            newSeats,
            nextFamilyIndex: roster.Families.Select(family => IndexOf(family.Id)).DefaultIfEmpty(-1).Max() + 1,
            native);

        GrantPromotedSkills(catalog, names, schoolSeed, when, remainingPeople, classes, incumbents, native);

        var people = remainingPeople.Values.ToList();
        var filled = AttachPupils(classes, people);
        return new Roster(people, families, filled);
    }

    private static void GrantPromotedSkills(
        DefCatalog catalog,
        NameSetDef names,
        int schoolSeed,
        DateTime when,
        Dictionary<string, Person> people,
        SchoolClass[] classes,
        HashSet<string> incumbents,
        string? nativeLanguage)
    {
        var yearByClass = classes.ToDictionary(
            schoolClass => schoolClass.Id,
            schoolClass => schoolClass.Year,
            StringComparer.Ordinal);
        foreach (var id in incumbents)
        {
            if (!people.TryGetValue(id, out var person)
                || person.ClassId is not { } classId
                || !yearByClass.TryGetValue(classId, out var year))
            {
                continue;
            }

            var rng = new Random(Seed.Mix(schoolSeed, person.Id, when.Year, Seed.SkillGrantSalt));
            people[id] = person with
            {
                Skills = PersonSampler.EnsurePupilYear(
                    catalog,
                    names,
                    rng,
                    person.AgeOn(when),
                    person.Choices,
                    person.Traits,
                    person.Skills,
                    year,
                    nativeLanguage),
            };
        }
    }

    private static void FillFirstYear(
        DefCatalog catalog,
        NameSetDef names,
        int schoolSeed,
        DateTime yearStart,
        DateTime asOf,
        List<Family> families,
        Dictionary<string, Person> people,
        List<PupilSeat> seats,
        int nextFamilyIndex,
        string? nativeLanguage)
    {
        var cursor = 0;
        foreach (var family in families.ToArray())
        {
            if (cursor >= seats.Count)
            {
                break;
            }

            var index = IndexOf(family.Id);
            if (index < 0 || family.ChildIds.Count == 0 || family.ChildIds.Count >= MaxChildrenInFamily)
            {
                continue;
            }

            if (!WantsYoungerSibling(schoolSeed, index, yearStart.Year))
            {
                continue;
            }

            var members = family.ParentIds.Concat(family.ChildIds)
                .Select(id => people[id])
                .ToArray();

            // A household with no adult left cannot take in a first-year; nor can one whose
            // father's name is unknown, since the newcomer's patronymic comes from it.
            if (!members.Any(person => !person.IsStudent)
                || FamilyFactory.FatherGivenOf(family, members) is null)
            {
                continue;
            }

            var childIndex = family.NextChildIndex;
            var rng = new Random(Seed.Mix(schoolSeed, index, Seed.IntakeSalt + yearStart.Year * 10 + childIndex));
            var child = FamilyFactory.AddChild(catalog, names, rng, family, members, seats[cursor], yearStart, asOf, childIndex, nativeLanguage);
            people[child.Id] = child;
            var familyAt = families.FindIndex(candidate => candidate.Id.Equals(family.Id, StringComparison.Ordinal));
            families[familyAt] = family with
            {
                ChildIds = [.. family.ChildIds, child.Id],
                NextChild = childIndex + 1,
            };
            foreach (var parentId in family.ParentIds)
            {
                people[parentId] = people[parentId] with { IsParent = true };
            }

            cursor++;
        }

        if (cursor >= seats.Count)
        {
            return;
        }

        var leftover = seats.Skip(cursor).ToArray();
        foreach (var plan in FamilyPlanner.Singletons(leftover, Math.Max(nextFamilyIndex, 0)))
        {
            var (created, members) = FamilyFactory.Create(catalog, names, schoolSeed, plan, yearStart, asOf, nativeLanguage);
            families.Add(created);
            foreach (var member in members)
            {
                people[member.Id] = member;
            }
        }
    }

    private static IReadOnlyList<SchoolClass> AttachPupils(SchoolClass[] classes, List<Person> people)
    {
        var pupils = new Dictionary<string, List<string>>(StringComparer.Ordinal);
        foreach (var schoolClass in classes)
        {
            pupils[schoolClass.Id] = [.. schoolClass.PupilIds];
        }

        foreach (var person in people)
        {
            if (!person.IsStudent || person.ClassId is not { } classId || !pupils.TryGetValue(classId, out var list))
            {
                continue;
            }

            if (!list.Contains(person.Id, StringComparer.Ordinal))
            {
                list.Add(person.Id);
            }
        }

        for (var i = 0; i < classes.Length; i++)
        {
            var schoolClass = classes[i];
            classes[i] = schoolClass with { PupilIds = pupils[schoolClass.Id] };
        }

        return classes;
    }

    private static bool WantsYoungerSibling(int schoolSeed, int familyIndex, int intakeYear)
    {
        var rng = new Random(Seed.Mix(schoolSeed, familyIndex, Seed.IntakeSalt ^ intakeYear));
        return rng.Next(100) < 40;
    }

    private static int IndexOf(string familyId) =>
        familyId.Length > 1 && familyId[0] == 'f' && int.TryParse(familyId.AsSpan(1), out var index)
            ? index
            : -1;
}