h-school/tests/HSchool.Simulation.Tests/ContagionSimulationTests.cs

using Arch.Core;
using HSchool.Ai;
using HSchool.Content;
using HSchool.People;

namespace HSchool.Simulation.Tests;

public class ContagionSimulationTests
{
    private static readonly DateTime TuesdayMorning = new(2012, 4, 3, 10, 0, 0, DateTimeKind.Utc);

    [Fact]
    public void ContagiousCarrier_InSameNode_InfectsMoreThanIsolatedControl()
    {
        using var exposed = OpenSchool(seed: 7);
        using var isolated = OpenSchool(seed: 7);

        var (carrierId, exposedId) = PickTwoPupils(exposed);
        var (_, controlId) = PickTwoPupils(isolated);
        SeedCarrier(exposed, carrierId, "TestPlague");
        SeedCarrier(isolated, carrierId, "TestPlague");

        SetPlace(exposed, carrierId, "classroom-101");
        SetPlace(exposed, exposedId, "classroom-101");
        SetPlace(isolated, carrierId, "classroom-101");
        SetPlace(isolated, controlId, "classroom-102");

        RunContagionDay(exposed);
        RunContagionDay(isolated);

        Assert.True(HasContagion(exposed, exposedId, "TestPlague"));
        Assert.False(HasContagion(isolated, controlId, "TestPlague"));
    }

    [Fact]
    public void SameSeed_SameContagionTransfers()
    {
        var first = RunGoldTransfers(seed: 11);
        var second = RunGoldTransfers(seed: 11);
        Assert.Equal(first, second);
        Assert.NotEmpty(first);
    }

    [Fact]
    public void NonContagiousDef_DoesNotJumpByClass()
    {
        using var school = OpenSchool(seed: 19);

        var schoolClass = school.Roster!.Classes.First(row => row.PupilIds.Count >= 3);
        var carrierId = schoolClass.PupilIds[0];
        var classmateIds = schoolClass.PupilIds.Skip(1).Take(2).ToArray();

        SeedCarrier(school, carrierId, "Otitis");
        foreach (var id in classmateIds.Append(carrierId))
        {
            SetPlace(school, id, schoolClass.RoomId);
        }

        RunContagionDay(school);

        foreach (var id in classmateIds)
        {
            Assert.False(HealthConditions.Has(Person(school, id), "Otitis"));
        }
    }

    [Fact]
    public void OutbreakNotice_EmitsAtThreshold()
    {
        using var school = OpenSchool(seed: 3);
        Assert.Equal(3, school.Catalog!.BehaviorRules!.DiseaseOutbreakMinNewCases);

        var room = "classroom-101";
        var pupils = school.Roster!.People
            .Where(person => person.IsStudent)
            .Take(5)
            .ToArray();
        Assert.True(pupils.Length >= 5);

        SeedCarrier(school, pupils[0].Id, "TestPlague");
        foreach (var pupil in pupils)
        {
            SetPlace(school, pupil.Id, room);
        }

        school.DrainWorldEvents();
        RunContagionDay(school);

        var contagionHits = pupils
            .Skip(1)
            .Count(person => HasContagion(school, person.Id, "TestPlague"));
        Assert.True(contagionHits >= 3);

        Assert.Contains(
            school.DrainWorldEvents(),
            row => row.Trigger.Equals(EventTriggers.DiseaseOutbreak, StringComparison.Ordinal));
    }

    [Fact]
    public void ClassContagion_ChainsAcrossDays()
    {
        using var school = OpenSchool(seed: 23);
        var schoolClass = school.Roster!.Classes.First(row => row.PupilIds.Count >= 3);
        var a = schoolClass.PupilIds[0];
        var b = schoolClass.PupilIds[1];
        var c = schoolClass.PupilIds[2];

        SeedCarrier(school, a, "TestClassPlague");
        foreach (var id in new[] { a, b, c })
        {
            SetPlace(school, id, schoolClass.RoomId);
        }

        RunContagionDay(school);
        Assert.True(HasContagion(school, b, "TestClassPlague") || HasContagion(school, c, "TestClassPlague"));

        // Next calendar day: newly infected classmates can pass it on.
        school.Clock.JumpTo(school.Clock.Time.AddDays(1));
        foreach (var id in new[] { a, b, c })
        {
            SetPlace(school, id, schoolClass.RoomId);
        }

        RunContagionDay(school);
        var infected = new[] { b, c }.Count(id =>
            HealthConditions.Has(Person(school, id), "TestClassPlague"));
        Assert.Equal(2, infected);
    }

    [Fact]
    public void ImmuneAndAlreadySick_AreNotInfected()
    {
        using var school = OpenSchool(seed: 5);
        var (carrierId, targetId) = PickTwoPupils(school);
        SeedCarrier(school, carrierId, "TestPlague");
        SeedCarrier(school, targetId, "TestPlague");
        DiseaseImmunities.Grant(Person(school, targetId), "TestPlague", school.Clock.Time.AddDays(10));

        // Second healthy classmate for the immunity case after clearing the already-sick one.
        var healthy = school.Roster!.People
            .Where(person => person.IsStudent && person.Id != carrierId && person.Id != targetId)
            .First();
        DiseaseImmunities.Grant(healthy, "TestPlague", school.Clock.Time.AddDays(10));

        SetPlace(school, carrierId, "classroom-101");
        SetPlace(school, targetId, "classroom-101");
        SetPlace(school, healthy.Id, "classroom-101");

        var beforeSick = Person(school, targetId).Conditions!.Count;
        RunContagionDay(school);

        Assert.Equal(beforeSick, Person(school, targetId).Conditions!.Count);
        Assert.False(HasContagion(school, targetId, "TestPlague"));
        Assert.False(HasContagion(school, healthy.Id, "TestPlague"));
        Assert.False(HealthConditions.Has(healthy, "TestPlague"));
    }

    private static HashSet<string> RunGoldTransfers(int seed)
    {
        using var school = OpenSchool(seed);
        var (carrierId, exposedId) = PickTwoPupils(school);
        SeedCarrier(school, carrierId, "TestPlague");
        SetPlace(school, carrierId, "classroom-101");
        SetPlace(school, exposedId, "classroom-101");
        RunContagionDay(school);

        return school.Roster!.People
            .Where(person => HasContagion(school, person.Id, "TestPlague"))
            .Select(person => person.Id)
            .OrderBy(id => id, StringComparer.Ordinal)
            .ToHashSet(StringComparer.Ordinal);
    }

    private static void RunContagionDay(School school)
    {
        school.LastDiseaseDay = int.MinValue;
        DiseaseSystem.Apply(school, gameMinutes: 0);
    }

    private static void SeedCarrier(School school, string personId, string defName)
    {
        HealthConditions.Add(
            Person(school, personId),
            new HealthCondition
            {
                DefName = defName,
                Severity = 0.2f,
                Progress = 0.05f,
                Source = "seed",
                StartedAt = school.Clock.Time.AddHours(-6),
            });
    }

    private static bool HasContagion(School school, string personId, string defName)
    {
        var list = Person(school, personId).Conditions;
        return list is not null
            && list.Any(row =>
                row.DefName.Equals(defName, StringComparison.Ordinal)
                && string.Equals(row.Source, "contagion", StringComparison.Ordinal));
    }

    private static (string CarrierId, string OtherId) PickTwoPupils(School school)
    {
        var pupils = school.Roster!.People.Where(person => person.IsStudent).Take(2).ToArray();
        Assert.Equal(2, pupils.Length);
        return (pupils[0].Id, pupils[1].Id);
    }

    private static Person Person(School school, string id) =>
        school.Roster!.People.First(row => row.Id.Equals(id, StringComparison.Ordinal));

    private static void SetPlace(School school, string personId, string? nodeId)
    {
        var query = new QueryDescription().WithAll<PersonIdentity, Presence>();
        school.World.Query(
            in query,
            (ref PersonIdentity identity, ref Presence presence) =>
            {
                if (!identity.Id.Equals(personId, StringComparison.Ordinal))
                {
                    return;
                }

                presence = nodeId is null
                    ? Presence.OffCampus
                    : new Presence(nodeId, 0f, nodeId, false, []);
            });
    }

    private static School OpenSchool(int seed)
    {
        var (catalog, map) = ContagionCatalog();
        var roster = RosterGenerator.Generate(catalog, map, schoolSeed: seed, "Russia", TuesdayMorning);
        var pool = ApplicantPool.Create(catalog, roster, schoolSeed: seed, "Russia", TuesdayMorning);
        var school = School.Create(1, "Зараза", TuesdayMorning, catalog, map);
        school.InstallPeople(roster, seed, "Russia", pool);
        school.ConfigurePresence(weekDays: 5, maxDecisionsPerTick: 10_000);
        return school;
    }

    private static (DefCatalog Catalog, MapLayout Map) ContagionCatalog()
    {
        var root = Path.Combine(AppContext.BaseDirectory, "vanilla");
        var documents = new List<ContentDocument>();
        foreach (var path in Directory.EnumerateFiles(root, "*.*", SearchOption.AllDirectories))
        {
            if (!path.EndsWith(".jsonc", StringComparison.OrdinalIgnoreCase)
                && !path.EndsWith(".json", StringComparison.OrdinalIgnoreCase))
            {
                continue;
            }

            var relative = Path.GetRelativePath(root, path).Replace('\\', '/');
            documents.Add(new ContentDocument(CatalogLoader.CorePackId, relative, File.ReadAllText(path)));
        }

        documents.Add(
            new ContentDocument(
                CatalogLoader.CorePackId,
                "defs/diseases/test-plague.jsonc",
                """
                {
                  "defName": "TestPlague",
                  "family": "respiratory",
                  "incubationDays": 0,
                  "immunityDays": 1,
                  "baseChancePerDay": 0,
                  "nodeContagionChancePerDay": 1,
                  "classContagionChancePerDay": 0,
                  "contagiousDuringIncubation": true,
                  "stages": [
                    { "minSeverity": 0, "severityPerDay": 0.1, "progressPerDay": 0.1, "lessonLearningFactor": 1, "stayHomeChance": 0 }
                  ]
                }
                """));
        documents.Add(
            new ContentDocument(
                CatalogLoader.CorePackId,
                "defs/diseases/test-class-plague.jsonc",
                """
                {
                  "defName": "TestClassPlague",
                  "family": "respiratory",
                  "incubationDays": 0,
                  "immunityDays": 1,
                  "baseChancePerDay": 0,
                  "nodeContagionChancePerDay": 0,
                  "classContagionChancePerDay": 1,
                  "contagiousDuringIncubation": true,
                  "stages": [
                    { "minSeverity": 0, "severityPerDay": 0.1, "progressPerDay": 0.1, "lessonLearningFactor": 1, "stayHomeChance": 0 }
                  ]
                }
                """));
        documents.Add(
            new ContentDocument(
                CatalogLoader.CorePackId,
                "patches/disable-disease-onset.jsonc",
                """
                { "target": "Behavior", "ops": [ { "op": "replace", "path": "/diseaseVectorScale", "value": 0 } ] }
                """));

        var catalog = new CatalogLoader().Load([CatalogLoader.CorePackId], documents);
        var map = CatalogLoader.LastDefaultMap([CatalogLoader.CorePackId], documents);
        Assert.NotNull(map);
        Assert.Equal(0f, catalog.BehaviorRules!.DiseaseVectorScale);
        Assert.True(catalog.Diseases.ContainsKey("TestPlague"));
        return (catalog, map);
    }
}