PKHeX/PKHeX.Core/Editing/HiddenPower.cs

using System;
using System.Runtime.CompilerServices;

namespace PKHeX.Core
{
    /// <summary>
    /// Logic for calculating a Hidden Power Type based on IVs and generation-format.
    /// </summary>
    public static class HiddenPower
    {
        /// <summary>
        /// Gets the current Hidden Power Type of the input <see cref="IVs"/> for the requested format generation.
        /// </summary>
        /// <param name="IVs">Current IVs</param>
        /// <returns>Hidden Power Type of the <see cref="IVs"/></returns>
        /// <param name="format">Generation format</param>
        public static int GetType(ReadOnlySpan<int> IVs, int format)
        {
            if (format <= 2)
                return GetTypeGB(IVs);
            return GetType(IVs);
        }

        /// <summary>
        /// Gets the current Hidden Power Type of the input <see cref="IVs"/> for Generations 3+
        /// </summary>
        /// <param name="IVs">Current IVs</param>
        /// <returns>Hidden Power Type of the <see cref="IVs"/></returns>
        public static int GetType(ReadOnlySpan<int> IVs)
        {
            int hp = 0;
            for (int i = 0; i < 6; i++)
                hp |= (IVs[i] & 1) << i;
            hp *= 0xF;
            hp /= 0x3F;
            return hp;
        }

        /// <summary>
        /// Gets the current Hidden Power Type of the input <see cref="IVs"/> for Generations 1 &amp; 2
        /// </summary>
        /// <param name="IVs">Current IVs</param>
        /// <returns>Hidden Power Type of the <see cref="IVs"/></returns>
        public static int GetTypeGB(ReadOnlySpan<int> IVs)
        {
            var atk = IVs[1];
            var def = IVs[2];
            return ((atk & 3) << 2) | (def & 3);
        }

        /// <summary>
        /// Modifies the provided <see cref="IVs"/> to have the requested <see cref="hiddenPowerType"/> for Generations 1 &amp; 2
        /// </summary>
        /// <param name="hiddenPowerType">Hidden Power Type</param>
        /// <param name="IVs">Current IVs</param>
        /// <returns>True if the Hidden Power of the <see cref="IVs"/> is obtained, with or without modifications</returns>
        public static bool SetTypeGB(int hiddenPowerType, Span<int> IVs)
        {
            IVs[1] = (IVs[1] & ~3) | (hiddenPowerType >> 2);
            IVs[2] = (IVs[2] & ~3) | (hiddenPowerType & 3);
            return true;
        }

        /// <summary>
        /// Modifies the provided <see cref="IVs"/> to have the requested <see cref="hiddenPowerType"/>.
        /// </summary>
        /// <param name="hiddenPowerType">Hidden Power Type</param>
        /// <param name="IVs">Current IVs (6 total)</param>
        /// <param name="format">Generation format</param>
        /// <returns>True if the Hidden Power of the <see cref="IVs"/> is obtained, with or without modifications</returns>
        public static bool SetIVsForType(int hiddenPowerType, Span<int> IVs, int format)
        {
            if (format <= 2)
                return SetTypeGB(hiddenPowerType, IVs);
            return SetIVsForType(hiddenPowerType, IVs);
        }

        /// <summary>
        /// Sets the <see cref="IVs"/> to the requested <see cref="hpVal"/> for Generation 3+ game formats.
        /// </summary>
        /// <param name="hpVal">Hidden Power Type</param>
        /// <param name="IVs">Current IVs (6 total)</param>
        /// <returns>True if the Hidden Power of the <see cref="IVs"/> is obtained, with or without modifications</returns>
        public static bool SetIVsForType(int hpVal, Span<int> IVs)
        {
            int flawlessCount = IVs.Count(31);
            if (flawlessCount == 0)
                return false;

            if (flawlessCount == IVs.Length)
            {
                SetIVs(hpVal, IVs); // Get IVs
                return true;
            }

            int current = GetType(IVs);
            if (current == hpVal)
                return true; // no mods necessary

            // Required HP type doesn't match IVs. Make currently-flawless IVs flawed.
            Span<int> scratch = stackalloc int[IVs.Length];
            Span<int> result = stackalloc int[IVs.Length];
            var success = GetSuggestedHiddenPowerIVs(hpVal, IVs, scratch, result);
            if (!success)
                return false; // can't force hidden power?

            // set IVs back to array
            result.CopyTo(IVs);
            return true;
        }

        // Non-recursive https://en.wikipedia.org/wiki/Heap%27s_algorithm
        private static bool GetSuggestedHiddenPowerIVs(int hpVal, ReadOnlySpan<int> original, Span<int> ivs, Span<int> best)
        {
            const int max = 31;

            // Get a list of indexes that can be mutated
            Span<int> indexes = stackalloc int[original.Length];
            int flaw = 0;
            for (int i = 0; i < original.Length; i++)
            {
                if (original[i] == max)
                    indexes[flaw++] = i;
            }
            indexes = indexes[..flaw];
            Span<int> c = stackalloc int[indexes.Length];

            int mutated = c.Length + 1; // result tracking
            for (int i = 1; i < c.Length;)
            {
                if (c[i] >= i) // Reset the state and simulate popping the stack by incrementing the pointer.
                {
                    c[i++] = 0;
                    continue;
                }

                var x = (i & 1) == 1 ? c[i] : 0;
                Swap(ref indexes[i], ref indexes[x]);

                // Inlined continuance check
                original.CopyTo(ivs);
                var q = Math.Min(indexes.Length, mutated);
                for (var j = 0; j < q; j++)
                {
                    ivs[indexes[j]] ^= 1;
                    if (hpVal != GetType(ivs))
                        continue;

                    var ct = j + 1;
                    if (ct >= mutated)
                        break; // any further flaws are always worse

                    mutated = ct;
                    ivs.CopyTo(best);
                    if (j == 0) // nothing will be better than only 1 flaw
                        return true;
                    break; // any further flaws are always worse
                }

                c[i]++;
                i = 1;
            }

            return mutated <= c.Length; // did we actually find a suitable result?
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static void Swap<T>(ref T a, ref T b) => (a, b) = (b, a);

        /// <summary>Calculate the Hidden Power Type of the entered IVs.</summary>
        /// <param name="type">Hidden Power Type</param>
        /// <param name="ivs">Individual Values (H/A/B/S/C/D)</param>
        /// <param name="format">Generation specific format</param>
        public static void SetIVs(int type, Span<int> ivs, int format = PKX.Generation)
        {
            if (format <= 2)
            {
                ivs[1] = (ivs[1] & ~3) | (type >> 2);
                ivs[2] = (ivs[2] & ~3) | (type & 3);
            }
            else
            {
                var bits = DefaultLowBits[type];
                for (int i = 0; i < 6; i++)
                    ivs[i] = (ivs[i] & 0x1E) + ((bits >> i) & 1);
            }
        }

        /// <summary>
        /// Hidden Power IV values (even or odd) to achieve a specified Hidden Power Type
        /// </summary>
        /// <remarks>
        /// There are other IV combinations to achieve the same Hidden Power Type.
        /// These are just precomputed for fast modification.
        /// Individual Values (H/A/B/S/C/D)
        /// </remarks>
        public static readonly byte[] DefaultLowBits =
        {
            0b000011, // Fighting
            0b001000, // Flying
            0b001011, // Poison
            0b001111, // Ground
            0b010011, // Rock
            0b011001, // Bug
            0b011101, // Ghost
            0b011111, // Steel
            0b100101, // Fire
            0b101001, // Water
            0b101101, // Grass
            0b101111, // Electric
            0b110101, // Psychic
            0b111001, // Ice
            0b111101, // Dragon
            0b111111, // Dark
        };
    }
}