/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.commons.rng.examples.jmh.core;

import java.math.BigInteger;
import java.util.function.LongSupplier;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.examples.jmh.core.LXMBenchmark.LCG128Source;
import org.apache.commons.rng.examples.jmh.core.LXMBenchmark.LXM128Source;
import org.apache.commons.rng.examples.jmh.core.LXMBenchmark.UnsignedMultiply128Source;
import org.apache.commons.rng.examples.jmh.core.LXMBenchmark.UnsignedMultiplyHighSource;
import org.apache.commons.rng.simple.RandomSource;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.RepeatedTest;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledForJreRange;
import org.junit.jupiter.api.condition.JRE;
import org.opentest4j.TestAbortedException;

/**
 * Test for methods used in {@link LXMBenchmark}. This checks the different versions of
 * 128-bit multiplication compute the correct result (verified using BigInteger).
 */
class LXMBenchmarkTest {
    /** 2^63. */
    private static final BigInteger TWO_POW_63 = BigInteger.ONE.shiftLeft(63);

    /**
     * A function operating on 2 long values.
     */
    interface LongLongFunction {
        /**
         * Apply the function.
         *
         * @param a Argument a
         * @param b Argument b
         * @return the result
         */
        long apply(long a, long b);
    }

    /**
     * A function operating on 4 long values.
     */
    interface LongLongLongLongFunction {
        /**
         * Apply the function.
         *
         * @param a Argument a
         * @param b Argument b
         * @param c Argument c
         * @param d Argument d
         * @return the result
         */
        long apply(long a, long b, long c, long d);
    }

    @Test
    @EnabledForJreRange(min = JRE.JAVA_9)
    void testMathMultiplyHigh() {
        try {
            assertUnsignedMultiply(UnsignedMultiplyHighSource::mathMultiplyHigh, null);
        } catch (NoSuchMethodError e) {
            throw new TestAbortedException("Update mathMultiplyHigh to call Math.multiplyHigh");
        }
    }

    @Test
    @EnabledForJreRange(min = JRE.JAVA_18)
    void testMathUnsignedMultiplyHigh() {
        try {
            assertUnsignedMultiply(UnsignedMultiplyHighSource::mathUnsignedMultiplyHigh, null);
        } catch (NoSuchMethodError e) {
            throw new TestAbortedException("Update mathUnsignedMultiplyHigh to call Math.unsignedMultiplyHigh");
        }
    }

    @Test
    void testUnsignedMultiplyHigh() {
        assertUnsignedMultiply(UnsignedMultiplyHighSource::unsignedMultiplyHigh, null);
    }

    @Test
    void testUnsignedMultiplyHighWithLow() {
        final long[] lo = {0};
        assertUnsignedMultiply((a, b) -> UnsignedMultiplyHighSource.unsignedMultiplyHigh(a, b, lo), lo);
    }

    @Test
    void testMhMultiplyHigh() {
        assertUnsignedMultiply(UnsignedMultiplyHighSource::mhMultiplyHigh, null);
    }

    @Test
    void testMhUnsignedMultiplyHigh() {
        assertUnsignedMultiply(UnsignedMultiplyHighSource::mhUnsignedMultiplyHigh, null);
    }

    private static void assertUnsignedMultiply(LongLongFunction fun, long[] lo) {
        final UniformRandomProvider rng = RandomSource.JSF_64.create();
        for (int i = 0; i < 100; i++) {
            final long a = rng.nextLong();
            final long b = rng.nextLong();
            Assertions.assertEquals(unsignedMultiplyHigh(a, b),
                                    fun.apply(a, b));
            if (lo != null) {
                Assertions.assertEquals(a * b, lo[0]);
            }
        }
    }

    @Test
    void testUnsignedMultiplyHighML() {
        final UniformRandomProvider rng = RandomSource.JSF_64.create();
        final long b = LXMBenchmark.ML;
        for (int i = 0; i < 100; i++) {
            final long a = rng.nextLong();
            Assertions.assertEquals(unsignedMultiplyHigh(a, b),
                UnsignedMultiplyHighSource.unsignedMultiplyHighML(a));
        }
    }

    @Test
    void testUnsignedMultiplyHighWithProducts() {
        assertUnsignedMultiply128((a, b, c, d) ->
            UnsignedMultiplyHighSource.unsignedMultiplyHigh(b, d) +
            a * d + b * c, null);
    }

    @Test
    void testUnsignedMultiply128() {
        final long[] lo = {0};
        assertUnsignedMultiply128((a, b, c, d) -> UnsignedMultiply128Source.unsignedMultiply128(a, b, c, d, lo), lo);
    }

    private static void assertUnsignedMultiply128(LongLongLongLongFunction fun, long[] lo) {
        final UniformRandomProvider rng = RandomSource.JSF_64.create();
        for (int i = 0; i < 100; i++) {
            final long a = rng.nextLong();
            final long b = rng.nextLong();
            final long c = rng.nextLong();
            final long d = rng.nextLong();
            Assertions.assertEquals(unsignedMultiplyHigh(a, b, c, d),
                                    fun.apply(a, b, c, d));
            if (lo != null) {
                Assertions.assertEquals(b * d, lo[0]);
            }
        }
    }

    @Test
    void testUnsignedSquareHigh() {
        final UniformRandomProvider rng = RandomSource.JSF_64.create();
        for (int i = 0; i < 100; i++) {
            final long a = rng.nextLong();
            Assertions.assertEquals(unsignedMultiplyHigh(a, a),
                                    UnsignedMultiply128Source.unsignedSquareHigh(a));
        }
    }

    @Test
    void testUnsignedSquare128() {
        final long[] lo = {0};
        final UniformRandomProvider rng = RandomSource.JSF_64.create();
        for (int i = 0; i < 100; i++) {
            final long a = rng.nextLong();
            final long b = rng.nextLong();
            Assertions.assertEquals(unsignedMultiplyHigh(a, b, a, b),
                                    UnsignedMultiply128Source.unsignedSquare128(a, b, lo));
            Assertions.assertEquals(b * b, lo[0]);
        }
    }

    /**
     * Compute the unsigned multiply of two values using BigInteger.
     *
     * @param a First value
     * @param b Second value
     * @return the upper 64-bits of the 128-bit result
     */
    private static long unsignedMultiplyHigh(long a, long b) {
        final BigInteger ba = toUnsignedBigInteger(a);
        final BigInteger bb = toUnsignedBigInteger(b);
        return ba.multiply(bb).shiftRight(64).longValue();
    }

    /**
     * Compute the unsigned multiply of two 128-bit values using BigInteger.
     *
     * <p>This computes the upper 64-bits of a 128-bit result that would
     * be generated by multiplication to a native 128-bit integer type.
     *
     * @param a First value
     * @param b Second value
     * @return the upper 64-bits of the truncated 128-bit result
     */
    private static long unsignedMultiplyHigh(long ah, long al, long bh, long bl) {
        final BigInteger a = toUnsignedBigInteger(ah, al);
        final BigInteger b = toUnsignedBigInteger(bh, bl);
        return a.multiply(b).shiftRight(64).longValue();
    }

    /**
     * Create a big integer treating the value as unsigned.
     *
     * @param v Value
     * @return the big integer
     */
    private static BigInteger toUnsignedBigInteger(long v) {
        return v < 0 ?
            TWO_POW_63.or(BigInteger.valueOf(v & Long.MAX_VALUE)) :
            BigInteger.valueOf(v);
    }

    /**
     * Create a 128-bit big integer treating the value as unsigned.
     *
     * @param hi High part of value
     * @param lo High part of value
     * @return the big integer
     */
    private static BigInteger toUnsignedBigInteger(long hi, long lo) {
        return toUnsignedBigInteger(hi).shiftLeft(64).or(toUnsignedBigInteger(lo));
    }

    /**
     * Test all 128-bit LCG implementations compute the same state update.
     */
    @RepeatedTest(value = 50)
    void testLcg128() {
        final long ah = RandomSource.createLong();
        final long al = RandomSource.createLong() | 1;
        final LongSupplier a = new LCG128Source.ReferenceLcg128(ah, al)::getAsLong;
        final LongSupplier[] b = new LongSupplier[] {
            new LCG128Source.ReferenceLcg128(ah, al)::getAsLong,
            new LCG128Source.ReferenceLcg128Final(ah, al)::getAsLong,
            new LCG128Source.CompareUnsignedLcg128(ah, al)::getAsLong,
            new LCG128Source.ConditionalLcg128(ah, al)::getAsLong,
            new LCG128Source.ConditionalLcg128Final(ah, al)::getAsLong,
            new LCG128Source.BranchlessLcg128(ah, al)::getAsLong,
            new LCG128Source.BranchlessFullLcg128(ah, al)::getAsLong,
            new LCG128Source.BranchlessFullComposedLcg128(ah, al)::getAsLong,
        };
        for (int i = 0; i < 10; i++) {
            final long expected = a.getAsLong();
            for (int j = 0; j < b.length; j++) {
                Assertions.assertEquals(expected, b[j].getAsLong());
            }
        }
    }

    /**
     * Test all 128-bit LCG based implementation of a LXM generator compute the same state update.
     */
    @RepeatedTest(value = 50)
    void testLxm128() {
        final long ah = RandomSource.createLong();
        final long al = RandomSource.createLong();
        // Native seed: LCG add, LCG state, XBG
        // This requires the initial state of the XBG and LCG.
        final long[] seed = {ah, al,
                             LXM128Source.S0, LXM128Source.S1,
                             LXM128Source.X0, LXM128Source.X1};
        final UniformRandomProvider rng = RandomSource.L128_X128_MIX.create(seed);
        final LongSupplier a = rng::nextLong;
        final LongSupplier[] b = new LongSupplier[] {
            new LXM128Source.ReferenceLxm128(ah, al)::getAsLong,
            new LXM128Source.BranchlessLxm128(ah, al)::getAsLong,
        };
        for (int i = 0; i < 10; i++) {
            final long expected = a.getAsLong();
            for (int j = 0; j < b.length; j++) {
                Assertions.assertEquals(expected, b[j].getAsLong());
            }
        }
    }
}
