namespace Just.PreciseMath.Tests;

public class PreciseMathSqrtTests
{
    [Theory]
    [InlineData(1.0, 1.0)]
    [InlineData(4.0, 2.0)]
    [InlineData(9.0, 3.0)]
    [InlineData(0.25, 0.5)]
    [InlineData(2.25, 1.5)]
    public void ExactBinarySquaresHaveExactRoots(double square, double root)
    {
        DoubleDouble actual = DDMath.Sqrt(new DoubleDouble(square));
        actual.High.ShouldBe(root);
        BitConverter.DoubleToInt64Bits(actual.Low).ShouldBe(0L);
    }

    [Fact]
    public void PowersOfFourHaveExactRootsAcrossTheFiniteRange()
    {
        // sqrt(2^(2k)) = 2^k exactly, including the minimum subnormal input.
        for (int exponent = -1074; exponent <= 1022; exponent += 2)
        {
            DoubleDouble actual = DDMath.Sqrt(new DoubleDouble(Math.ScaleB(1.0, exponent)));
            actual.High.ShouldBe(Math.ScaleB(1.0, exponent / 2));
            BitConverter.DoubleToInt64Bits(actual.Low).ShouldBe(0L);
        }
    }

    [Theory]
    [InlineData(2.0)]
    [InlineData(3.0)]
    [InlineData(5.0)]
    [InlineData(1e-308)]
    [InlineData(1e308)]
    public void IrrationalRootsRetainMoreThanBinary64Precision(double value)
    {
        DoubleDouble input = new(value);
        DoubleDouble actual = DDMath.Sqrt(input);
        actual.Low.ShouldNotBe(0.0);
        AssertSqrtBound(input, actual);
    }

    [Fact]
    public void LowComponentsOnEitherSideOfAnExactSquareAffectTheRoot()
    {
        foreach (double low in new[] { Math.ScaleB(1.0, -54), -Math.ScaleB(1.0, -54) })
        {
            DoubleDouble input = DoubleDouble.FromComponents(1.0, low);
            DoubleDouble actual = DDMath.Sqrt(input);
            actual.High.ShouldBe(1.0);
            (Math.Sign(actual.Low) == Math.Sign(low)).ShouldBeTrue();
            AssertSqrtBound(input, actual);
        }
    }

    [Fact]
    public void SpecialValuesMatchBinary64AndRemainCanonical()
    {
        double[] values = [0.0, -0.0, double.PositiveInfinity, double.NegativeInfinity,
            double.NaN, -1.0, -double.Epsilon, double.MinValue];
        foreach (double value in values)
        {
            DoubleDouble actual = DDMath.Sqrt(new DoubleDouble(value));
            double expected = Math.Sqrt(value);
            // Construction canonicalizes all NaNs to double.NaN, including
            // domain errors; zeros must instead be compared by their sign bits.
            BitConverter.DoubleToInt64Bits(actual.High).ShouldBe(
                BitConverter.DoubleToInt64Bits(double.IsNaN(expected) ? double.NaN : expected));
            BitConverter.DoubleToInt64Bits(actual.Low).ShouldBe(0L);
            DoubleDouble.IsCanonical(actual).ShouldBeTrue();
        }

        DoubleDouble negative = DoubleDouble.FromComponents(-1.0, Math.ScaleB(1.0, -54));
        DoubleDouble.IsNaN(DDMath.Sqrt(negative)).ShouldBeTrue();
    }

    [Fact]
    public void PositiveFiniteInputsMeetExactRelativeErrorBound()
    {
        // Every binary64 exponent, including subnormal input exponents. Test zero,
        // dense, and sparse lows of either sign, without relying on DD arithmetic
        // to compute the expected square root or its square.
        Random random = new(314159);
        for (int exponent = -1074; exponent <= 1023; ++exponent)
        {
            double high = Math.ScaleB(1.0 + random.NextDouble(), exponent);
            double low = Math.ScaleB(random.NextDouble(), exponent - 54);
            double power = Math.ScaleB(1.0, exponent);
            foreach (double boundary in new[] { Math.BitDecrement(power), power, Math.BitIncrement(power) })
            {
                if (boundary > 0.0)
                {
                    DoubleDouble input = new(boundary);
                    AssertSqrtBound(input, DDMath.Sqrt(input));
                }
            }
            foreach (double residual in new[] { 0.0, low, -low, double.Epsilon, -double.Epsilon })
            {
                DoubleDouble input = DoubleDouble.FromComponents(high, residual);
                if (input.High > 0.0)
                {
                    AssertSqrtBound(input, DDMath.Sqrt(input));
                }
            }
        }

        // Pin the tiny-input precision loss described in the optional legacy review,
        // and guard against intermediate squaring overflow near the finite limit.
        DoubleDouble[] boundaries = [new(1e-308), new(double.Epsilon),
            new(Math.BitDecrement(Math.ScaleB(1.0, -1022))), new(Math.ScaleB(1.0, -1022)),
            new(Math.BitIncrement(Math.ScaleB(1.0, -1022))), new(double.MaxValue),
            DoubleDouble.FromComponents(double.MaxValue, Math.ScaleB(1.0, 969)),
            DoubleDouble.FromComponents(double.MaxValue, -Math.ScaleB(1.0, 969)),
            DoubleDouble.FromComponents(double.MaxValue, Math.BitDecrement(Math.ScaleB(1.0, 970))),
            DoubleDouble.FromComponents(double.MaxValue, -Math.ScaleB(1.0, 970)),
            DoubleDouble.FromComponents(1.0, Math.ScaleB(1.0, -53)),
            DoubleDouble.FromComponents(1.0, -Math.ScaleB(1.0, -54)),
            DoubleDouble.FromComponents(4.0, Math.ScaleB(1.0, -51)),
            DoubleDouble.FromComponents(4.0, -Math.ScaleB(1.0, -52))];
        foreach (DoubleDouble input in boundaries)
        {
            AssertSqrtBound(input, DDMath.Sqrt(input));
        }
    }

    [Fact]
    public void RootRoundingMidpointsAndAdjacentLowsMeetTheBound()
    {
        // Exact squares of root midpoints 1 + 2^-53 and 2 - 2^-53:
        // (1 + 2^-52) + 2^-106 and (4 - 2^-51) + 2^-106.
        // Bracket each with adjacent low doubles, then exercise exponent parity
        // and rescaling. At tiny exponents the input itself loses low precision;
        // the oracle always checks the exact stored input, not the unscaled square.
        double midpointLow = Math.ScaleB(1.0, -106);
        for (int exponent = -1074; exponent <= 1022; ++exponent)
        {
            foreach (double high in new[] { Math.BitIncrement(1.0), Math.BitDecrement(4.0) })
            {
                foreach (double low in new[] { Math.BitDecrement(midpointLow), midpointLow, Math.BitIncrement(midpointLow) })
                {
                    DoubleDouble input = DoubleDouble.FromComponents(Math.ScaleB(high, exponent), Math.ScaleB(low, exponent));
                    if (DoubleDouble.IsFinite(input) && input.High > 0.0)
                    {
                        AssertSqrtBound(input, DDMath.Sqrt(input));
                    }
                }
            }
        }
    }

    [Fact]
    public void PublicFactoryNormalizesLegacyZeroHighInputsBeforeTakingTheRoot()
    {
        // Legacy raw (0, low) examples are normalized at the public boundary.
        foreach (double low in new[] { double.Epsilon, 1e-308, 2.0, double.MaxValue })
        {
            DoubleDouble input = DoubleDouble.FromComponents(0.0, low);
            AssertSqrtBound(input, DDMath.Sqrt(input));
        }
        DoubleDouble.IsNaN(DDMath.Sqrt(DoubleDouble.FromComponents(0.0, -1.0))).ShouldBeTrue();
    }

    [Fact]
    public void DenseLowNormalizationBoundariesMeetTheBound()
    {
        // Half-ulp lows and their neighbors test both sides of input normalization.
        Random random = new(161803);
        for (int exponent = -1074; exponent <= 1023; ++exponent)
        {
            double high = Math.ScaleB(1.0 + random.NextDouble(), exponent);
            double halfUlp = Math.ScaleB(1.0, exponent - 53);
            foreach (double magnitude in new[] { Math.BitDecrement(halfUlp), halfUlp, Math.BitIncrement(halfUlp) })
            {
                foreach (double low in new[] { magnitude, -magnitude })
                {
                    DoubleDouble input = DoubleDouble.FromComponents(high, low);
                    if (DoubleDouble.IsFinite(input) && input.High > 0.0)
                    {
                        AssertSqrtBound(input, DDMath.Sqrt(input));
                    }
                }
            }
        }
    }

    [Fact]
    public void RepresentableSparseCorrectionsAreNotDiscarded()
    {
        // sqrt(1+d) = 1+d/2+O(d^2). For these dyadic d, the quadratic term
        // is below half an ulp of d/2, including when d/2 is the minimum subnormal.
        foreach (int exponent in new[] { -100, -500, -1000, -1073 })
        {
            foreach (double sign in new[] { -1.0, 1.0 })
            {
                double low = Math.ScaleB(sign, exponent);
                DoubleDouble actual = DDMath.Sqrt(DoubleDouble.FromComponents(1.0, low));
                actual.High.ShouldBe(1.0);
                actual.Low.ShouldBe(Math.ScaleB(sign, exponent - 1));
                DoubleDouble.IsCanonical(actual).ShouldBeTrue();
            }
        }
    }

    [Theory]
    [InlineData(-1.0)]
    [InlineData(1.0)]
    public void ScalingPreservesRepresentableSparseRootCorrections(double sign)
    {
        // sqrt(2^1000 + d) = 2^500 + d/2^501 + O(d^2/2^1500).
        // Scaling d=+/-2^-174 down by 2^1000 erases it, but the root's
        // +/-2^-675 low is representable. Prove its rounding interval exactly.
        DoubleDouble input = DoubleDouble.FromComponents(Math.ScaleB(1.0, 1000), Math.ScaleB(sign, -174));
        AssertSparseRoot(input, Math.ScaleB(1.0, 500), Math.ScaleB(sign, -675));
    }

    [Fact]
    public void SparseRootCorrectionsSurviveScalingAndCorrectionDivisionBoundaries()
    {
        // Exact squares with both even and odd high exponents. Include scaled
        // lows on both sides of the separate-correction dispatch, subnormal
        // scaled lows, fully erased scaled lows, and subnormal output lows.
        foreach (int exponent in new[] { 0, 2, 100, 1000, 1022 })
        {
            foreach (double rootMantissa in new[] { 1.0, 1.5 })
            {
                double high = Math.ScaleB(rootMantissa * rootMantissa, exponent);
                double rootHigh = Math.ScaleB(rootMantissa, exponent / 2);
                foreach (int gap in new[] { -1019, -1020, -1021, -1074, -1075, -1174 })
                {
                    double low = Math.ScaleB(1.0, exponent + gap);
                    double rootLow = Math.ScaleB(low, -(exponent / 2)) / (2.0 * rootMantissa);
                    if (rootLow == 0.0)
                    {
                        continue; // This matrix targets representable nonzero corrections.
                    }
                    foreach (double sign in new[] { -1.0, 1.0 })
                    {
                        // AssertSparseRoot independently proves this candidate's
                        // rounding interval; the first-order formula is not an oracle.
                        AssertSparseRoot(DoubleDouble.FromComponents(high, sign * low), rootHigh, sign * rootLow);
                    }
                }
            }
        }
    }

    [Fact]
    public void SparseRootDispatchNeighborsRetainBothSignsWithoutDoubleCounting()
    {
        // With high=2^1000, lows around 2^-20 scale to the dispatch at 2^-1020.
        // Keep the predicted output normal so candidate scaling is exact.
        double threshold = Math.ScaleB(1.0, -20);
        foreach (double low in new[] { Math.BitDecrement(threshold), threshold, Math.BitIncrement(threshold) })
        {
            foreach (double sign in new[] { -1.0, 1.0 })
            {
                DoubleDouble input = DoubleDouble.FromComponents(Math.ScaleB(1.0, 1000), sign * low);
                AssertSparseRoot(input, Math.ScaleB(1.0, 500), Math.ScaleB(sign * low, -501));
            }
        }
    }

    private static void AssertSparseRoot(DoubleDouble input, double high, double low)
    {
        BigInteger x = (Units(input.High) + Units(input.Low)) << 1074;
        BigInteger center = Units(high) + Units(low);
        BigInteger previous = Units(high) + Units(Math.BitDecrement(low));
        BigInteger following = Units(high) + Units(Math.BitIncrement(low));
        // Strict bounds avoid relying on a tie rule or a rounded sqrt oracle.
        ((previous + center) * (previous + center) < (x << 2)).ShouldBeTrue();
        ((center + following) * (center + following) > (x << 2)).ShouldBeTrue();
        DoubleDouble actual = DoubleDouble.Sqrt(input);
        actual.High.ShouldBe(high);
        BitConverter.DoubleToInt64Bits(actual.Low).ShouldBe(BitConverter.DoubleToInt64Bits(low));
        DoubleDouble facade = DDMath.Sqrt(input);
        BitConverter.DoubleToInt64Bits(facade.High).ShouldBe(BitConverter.DoubleToInt64Bits(actual.High));
        BitConverter.DoubleToInt64Bits(facade.Low).ShouldBe(BitConverter.DoubleToInt64Bits(actual.Low));
        AssertSqrtBound(input, actual);
    }

    private static void AssertSqrtBound(DoubleDouble input, DoubleDouble actual)
    {
        DoubleDouble.IsFinite(actual).ShouldBeTrue();
        DoubleDouble.IsCanonical(actual).ShouldBeTrue();
        (actual.High > 0.0).ShouldBeTrue();

        // Independent exact rational oracle: x = X*2^-1074, y = Y*2^-1074.
        // For positive x and y, |y/sqrt(x) - 1| <= t iff
        // x*(1-t)^2 <= y^2 <= x*(1+t)^2. With t=2^-100 we
        // cross-multiply to integers: no rounded sqrt, product, or DD conversion.
        BigInteger x = (Units(input.High) + Units(input.Low)) << 1074;
        BigInteger y = Units(actual.High) + Units(actual.Low);
        BigInteger scale = BigInteger.One << 100;
        BigInteger squared = (y * y) << 200;
        BigInteger lower = x * (scale - 1) * (scale - 1);
        BigInteger upper = x * (scale + 1) * (scale + 1);
        (squared >= lower && squared <= upper).ShouldBeTrue(
            $"Sqrt bound failed for ({input.High:R}, {input.Low:R}): ({actual.High:R}, {actual.Low:R})");
    }

    private static BigInteger Units(double value)
    {
        long bits = BitConverter.DoubleToInt64Bits(value);
        int exponent = (int)((bits >> 52) & 0x7ff);
        BigInteger significand = bits & 0xfffffffffffffL;
        if (exponent != 0)
        {
            significand += BigInteger.One << 52;
            significand <<= exponent - 1;
        }
        return bits < 0 ? -significand : significand;
    }
}