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

namespace HSchool.People;

internal static class PersonSampler
{
    public static (Dictionary<string, int> Numbers, Dictionary<string, string> Choices) Body(
        DefCatalog catalog,
        Random rng,
        bool female,
        int age)
    {
        var numbers = new Dictionary<string, int>(StringComparer.Ordinal);
        var choices = new Dictionary<string, string>(StringComparer.Ordinal);
        foreach (var def in catalog.BodyAttributes.Values.OrderBy(candidate => candidate.DefName, StringComparer.Ordinal))
        {
            if (def.Abstract)
            {
                continue;
            }

            if (def.Kind == BodyAttributeKind.Number)
            {
                numbers[def.DefName] = SampleNumber(def, rng, female, age);
            }
            else
            {
                choices[def.DefName] = SampleChoice(def, rng, female, age);
            }
        }

        var height = numbers.GetValueOrDefault("Height", 170);
        var weight = numbers.GetValueOrDefault("Weight", 65);
        choices[BodyBuilds.Attribute] = BodyBuilds.FromHeightAndWeight(height, weight);
        return (numbers, choices);
    }

    public static Dictionary<string, int> Skills(
        DefCatalog catalog,
        NameSetDef? names,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices,
        IReadOnlyList<string> traits,
        int? pupilYear,
        bool adult,
        string? nativeLanguage = null)
    {
        var values = new Dictionary<string, int>(StringComparer.Ordinal);
        foreach (var skill in catalog.Skills.Values.OrderBy(candidate => candidate.DefName, StringComparer.Ordinal))
        {
            if (!skill.Abstract && skill.Always)
            {
                RollInto(values, rng, age, choices, skill);
            }
        }

        GrantNative(values, catalog, NativeOf(names, nativeLanguage), rng, age, choices);
        if (pupilYear is { } year)
        {
            GrantYearSubjects(values, catalog, rng, age, choices, year);
        }

        if (adult)
        {
            GrantAdultExtras(values, catalog, rng, age, choices);
        }

        ApplyTraitsAndBody(values, catalog, choices, traits, values.Keys.ToArray());
        // After traits so Diligent cannot push a related tongue past the "knows a bit" cap.
        GrantRelated(values, catalog, names, NativeOf(names, nativeLanguage), rng, choices);
        return values;
    }

    /// <summary>
    /// Adds subject skills a pupil just became old enough for. Existing values stay put — yearly
    /// intake must not re-roll mathematics every September.
    /// </summary>
    public static Dictionary<string, int> EnsurePupilYear(
        DefCatalog catalog,
        NameSetDef? names,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices,
        IReadOnlyList<string> traits,
        IReadOnlyDictionary<string, int> existing,
        int pupilYear,
        string? nativeLanguage = null)
    {
        var values = new Dictionary<string, int>(existing, StringComparer.Ordinal);
        foreach (var skill in catalog.Skills.Values.OrderBy(candidate => candidate.DefName, StringComparer.Ordinal))
        {
            if (skill.Abstract || !skill.Always || values.ContainsKey(skill.DefName))
            {
                continue;
            }

            RollInto(values, rng, age, choices, skill);
        }

        GrantNative(values, catalog, NativeOf(names, nativeLanguage), rng, age, choices);
        GrantYearSubjects(values, catalog, rng, age, choices, pupilYear);

        var added = values.Keys.Where(name => !existing.ContainsKey(name)).ToArray();
        ApplyTraitsAndBody(values, catalog, choices, traits, added);
        return values;
    }

    private static string? NativeOf(NameSetDef? names, string? nativeLanguage)
    {
        if (!string.IsNullOrWhiteSpace(nativeLanguage))
        {
            return nativeLanguage;
        }

        var spoken = names?.Spoken;
        return spoken is { Count: > 0 } ? spoken[0] : null;
    }

    private static void GrantNative(
        Dictionary<string, int> values,
        DefCatalog catalog,
        string? native,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices)
    {
        if (string.IsNullOrWhiteSpace(native)
            || !catalog.Skills.TryGetValue(native, out var skill)
            || skill.Abstract)
        {
            return;
        }

        RollInto(values, rng, age, choices, skill);
    }

    /// <summary>
    /// Other tongues of the same name set, at a low level. Skips keys already granted as native
    /// or by a subject — a Belarusian pupil who studies Russian keeps the school value.
    /// </summary>
    private static void GrantRelated(
        Dictionary<string, int> values,
        DefCatalog catalog,
        NameSetDef? names,
        string? native,
        Random rng,
        IReadOnlyDictionary<string, string> choices)
    {
        if (names is null || names.RelatedLanguageChance <= 0)
        {
            return;
        }

        foreach (var name in names.Spoken.OrderBy(candidate => candidate, StringComparer.Ordinal))
        {
            if (name.Equals(native, StringComparison.Ordinal)
                || values.ContainsKey(name)
                || !catalog.Skills.TryGetValue(name, out var skill)
                || skill.Abstract)
            {
                continue;
            }

            if (rng.NextDouble() >= names.RelatedLanguageChance)
            {
                continue;
            }

            var rolled = (int)Math.Round(names.RelatedLanguageMean + (names.RelatedLanguageStdDev * NextGaussian(rng)));
            rolled = Clamp(rolled, skill.Range);
            rolled = Math.Min(rolled, names.RelatedLanguageMax);
            rolled = ApplyBodyLimits(rolled, skill, choices);
            values[skill.DefName] = Clamp(rolled, skill.Range);
        }
    }

    private static void GrantYearSubjects(
        Dictionary<string, int> values,
        DefCatalog catalog,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices,
        int pupilYear)
    {
        foreach (var subject in catalog.Subjects.Values.OrderBy(def => def.DefName, StringComparer.Ordinal))
        {
            if (subject.Abstract || pupilYear < subject.Grades.Min || pupilYear > subject.Grades.Max)
            {
                continue;
            }

            foreach (var share in subject.Skills.OrderBy(row => row.Skill, StringComparer.Ordinal))
            {
                if (!catalog.Skills.TryGetValue(share.Skill, out var skill) || skill.Abstract)
                {
                    continue;
                }

                RollInto(values, rng, age, choices, skill);
            }
        }
    }

    private static void GrantAdultExtras(
        Dictionary<string, int> values,
        DefCatalog catalog,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices)
    {
        var academic = catalog.Skills.Values
            .Where(skill =>
                !skill.Abstract
                && !skill.Always
                && !skill.Work
                && skill.AdultChance <= 0
                && SubjectUses(catalog, skill.DefName))
            .OrderBy(skill => skill.DefName, StringComparer.Ordinal)
            .ToList();
        PickInto(values, rng, age, choices, academic, rng.Next(3));

        var work = catalog.Skills.Values
            .Where(skill => !skill.Abstract && skill.Work)
            .OrderBy(skill => skill.DefName, StringComparer.Ordinal)
            .ToList();
        PickInto(values, rng, age, choices, work, rng.Next(3));

        foreach (var skill in catalog.Skills.Values.OrderBy(candidate => candidate.DefName, StringComparer.Ordinal))
        {
            if (skill.Abstract || skill.AdultChance <= 0 || values.ContainsKey(skill.DefName))
            {
                continue;
            }

            if (rng.NextDouble() < skill.AdultChance)
            {
                RollInto(values, rng, age, choices, skill);
            }
        }
    }

    private static bool SubjectUses(DefCatalog catalog, string skillName)
    {
        foreach (var subject in catalog.Subjects.Values)
        {
            if (subject.Abstract)
            {
                continue;
            }

            foreach (var share in subject.Skills)
            {
                if (share.Skill.Equals(skillName, StringComparison.Ordinal))
                {
                    return true;
                }
            }
        }

        return false;
    }

    private static void PickInto(
        Dictionary<string, int> values,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices,
        List<SkillDef> pool,
        int count)
    {
        var remaining = pool.Where(skill => !values.ContainsKey(skill.DefName)).ToList();
        for (var n = 0; n < count && remaining.Count > 0; n++)
        {
            var index = rng.Next(remaining.Count);
            RollInto(values, rng, age, choices, remaining[index]);
            remaining.RemoveAt(index);
        }
    }

    private static void RollInto(
        Dictionary<string, int> values,
        Random rng,
        int age,
        IReadOnlyDictionary<string, string> choices,
        SkillDef skill)
    {
        if (values.ContainsKey(skill.DefName))
        {
            return;
        }

        var mean = MeanForAge(skill, age);
        var stdDev = skill.Distribution?.StdDev ?? 10f;
        var rolled = (int)Math.Round(mean + (stdDev * NextGaussian(rng)));
        rolled = Clamp(rolled, skill.Range);
        rolled = ApplyBodyLimits(rolled, skill, choices);
        values[skill.DefName] = Clamp(rolled, skill.Range);
    }

    private static void ApplyTraitsAndBody(
        Dictionary<string, int> values,
        DefCatalog catalog,
        IReadOnlyDictionary<string, string> choices,
        IReadOnlyList<string> traits,
        IEnumerable<string> names)
    {
        var allowed = names as ISet<string> ?? new HashSet<string>(names, StringComparer.Ordinal);
        var touched = new HashSet<string>(StringComparer.Ordinal);
        foreach (var traitName in traits)
        {
            if (!catalog.Traits.TryGetValue(traitName, out var trait))
            {
                continue;
            }

            foreach (var modifier in trait.SkillModifiers)
            {
                if (!allowed.Contains(modifier.Skill)
                    || !values.TryGetValue(modifier.Skill, out var current)
                    || !catalog.Skills.TryGetValue(modifier.Skill, out var skill))
                {
                    continue;
                }

                values[modifier.Skill] = Clamp(current + modifier.Offset, skill.Range);
                touched.Add(modifier.Skill);
            }
        }

        // The body has the last word. Without this pass a trait modifier reopened what the body
        // closed: a skinny Bully rolled Strength 45, the trait added 8, and the "Skinny caps
        // Strength at 45" limit was silently gone.
        foreach (var skillName in touched)
        {
            var skill = catalog.Skills[skillName];
            values[skillName] = Clamp(ApplyBodyLimits(values[skillName], skill, choices), skill.Range);
        }
    }

    public static List<string> Traits(DefCatalog catalog, Random rng, IReadOnlyList<string> roles, int age)
    {
        var picked = new List<string>();
        var count = rng.Next(3);
        if (count == 0 || catalog.Traits.Count == 0)
        {
            return picked;
        }

        var pool = catalog.Traits.Values
            .Where(trait => !trait.Abstract && TraitFits(trait, roles, age))
            .OrderBy(trait => trait.DefName, StringComparer.Ordinal)
            .ToList();

        for (var n = 0; n < count && pool.Count > 0; n++)
        {
            var chosen = WeightedTrait(pool, rng);
            if (chosen is null)
            {
                break;
            }

            picked.Add(chosen.DefName);
            var banned = catalog.TraitIncompatibilities(chosen.DefName);
            pool.RemoveAll(trait =>
                trait.DefName == chosen.DefName
                || banned.Contains(trait.DefName));
        }

        picked.Sort(StringComparer.Ordinal);
        return picked;
    }

    public static Dictionary<string, float> Needs(DefCatalog catalog)
    {
        var values = new Dictionary<string, float>(StringComparer.Ordinal);
        foreach (var need in catalog.Needs.Values.OrderBy(candidate => candidate.DefName, StringComparer.Ordinal))
        {
            if (!need.Abstract)
            {
                values[need.DefName] = need.Initial;
            }
        }

        return values;
    }

    internal static int ApplyBodyLimits(int value, SkillDef skill, IReadOnlyDictionary<string, string> choices)
    {
        foreach (var limit in skill.BodyLimits)
        {
            if (!choices.TryGetValue(limit.Attribute, out var actual))
            {
                continue;
            }

            if (limit.Value is not null && !actual.Equals(limit.Value, StringComparison.Ordinal))
            {
                continue;
            }

            if (limit.Min is { } min)
            {
                value = Math.Max(value, min);
            }

            if (limit.Max is { } max)
            {
                value = Math.Min(value, max);
            }
        }

        return value;
    }

    private static bool TraitFits(TraitDef trait, IReadOnlyList<string> roles, int age)
    {
        if (trait.Age is { } range && (age < range.Min || age > range.Max))
        {
            return false;
        }

        if (trait.Roles.Count == 0)
        {
            return true;
        }

        foreach (var role in trait.Roles)
        {
            foreach (var have in roles)
            {
                if (role.Equals(have, StringComparison.OrdinalIgnoreCase))
                {
                    return true;
                }
            }
        }

        return false;
    }

    private static TraitDef? WeightedTrait(List<TraitDef> pool, Random rng)
    {
        var total = 0;
        foreach (var trait in pool)
        {
            total += Math.Max(trait.Weight, 1);
        }

        if (total <= 0)
        {
            return null;
        }

        var pick = rng.Next(total);
        foreach (var trait in pool)
        {
            pick -= Math.Max(trait.Weight, 1);
            if (pick < 0)
            {
                return trait;
            }
        }

        return pool[^1];
    }

    private static int SampleNumber(BodyAttributeDef def, Random rng, bool female, int age)
    {
        var row = MatchDistribution(def.Distributions, female, age);
        var mean = row?.Distribution.Mean ?? 0;
        var stdDev = row?.Distribution.StdDev ?? 1;
        var value = (int)Math.Round(mean + (stdDev * NextGaussian(rng)));
        if (row?.Range is { } range)
        {
            return Clamp(value, range);
        }

        return value;
    }

    private static string SampleChoice(BodyAttributeDef def, Random rng, bool female, int age)
    {
        var options = def.Options.Where(option => OptionFits(option, female, age)).ToList();
        if (options.Count == 0)
        {
            options = [.. def.Options];
        }

        var total = 0;
        foreach (var option in options)
        {
            total += Math.Max(option.Weight, 1);
        }

        var pick = rng.Next(Math.Max(total, 1));
        foreach (var option in options)
        {
            pick -= Math.Max(option.Weight, 1);
            if (pick < 0)
            {
                return option.Value;
            }
        }

        return options[^1].Value;
    }

    private static bool OptionFits(WeightedOption option, bool female, int age)
    {
        if (option.AgeMin is { } min && age < min)
        {
            return false;
        }

        if (option.AgeMax is { } max && age > max)
        {
            return false;
        }

        if (option.Sex is null)
        {
            return true;
        }

        var want = female ? "female" : "male";
        return option.Sex.Equals(want, StringComparison.OrdinalIgnoreCase);
    }

    private static SexAgeDistribution? MatchDistribution(
        IReadOnlyList<SexAgeDistribution> rows,
        bool female,
        int age)
    {
        SexAgeDistribution? unisex = null;
        foreach (var row in rows)
        {
            if (row.AgeMin is { } min && age < min)
            {
                continue;
            }

            if (row.AgeMax is { } max && age > max)
            {
                continue;
            }

            if (row.Sex is null)
            {
                unisex ??= row;
                continue;
            }

            var want = female ? "female" : "male";
            if (row.Sex.Equals(want, StringComparison.OrdinalIgnoreCase))
            {
                return row;
            }
        }

        return unisex ?? (rows.Count == 0 ? null : rows[0]);
    }

    private static float MeanForAge(SkillDef skill, int age)
    {
        if (skill.AgeMeans.Count == 0)
        {
            return skill.Distribution?.Mean ?? (skill.Range.Min + skill.Range.Max) / 2f;
        }

        var points = skill.AgeMeans.OrderBy(point => point.Age).ToList();
        if (age <= points[0].Age)
        {
            return points[0].Mean;
        }

        if (age >= points[^1].Age)
        {
            return points[^1].Mean;
        }

        for (var i = 1; i < points.Count; i++)
        {
            if (age > points[i].Age)
            {
                continue;
            }

            var a = points[i - 1];
            var b = points[i];
            var span = b.Age - a.Age;
            var t = span == 0 ? 0f : (age - a.Age) / (float)span;
            return a.Mean + (t * (b.Mean - a.Mean));
        }

        return points[^1].Mean;
    }

    private static int Clamp(int value, IntRange range) => Math.Clamp(value, range.Min, range.Max);

    private static double NextGaussian(Random rng)
    {
        double u1;
        do
        {
            u1 = rng.NextDouble();
        }
        while (u1 <= double.Epsilon);

        var u2 = rng.NextDouble();
        return Math.Sqrt(-2d * Math.Log(u1)) * Math.Cos(2d * Math.PI * u2);
    }
}