package org.sample;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.reflect.Method;
import java.math.BigInteger;
import java.util.concurrent.TimeUnit;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;

@Warmup(iterations = 2)
@Measurement(iterations = 5)
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@State(Scope.Benchmark)
public class BigIntegerBench2 {
    static final int K = 0xCAFEBABE; // any int for multiplier value
    static final int OFFSET = 0; // we'll go from the start of array

    static final MethodHandles.Lookup lookup = MethodHandles.lookup();

    static final MethodHandle squareMethodHandle, mulAddMethodHandle,
        implMultiplyToLenHandle, implSquareToLenHandle, multiplyMethodHandle;

    private static MethodHandle handleFor(Class klass, String name, Class<?>... parameterTypes)
        throws NoSuchMethodException, SecurityException, IllegalAccessException {
        Method m = klass.getDeclaredMethod(name, parameterTypes);
        m.setAccessible(true);
        return lookup.unreflect(m);
    }

    static {

        try {
            squareMethodHandle = handleFor(BigInteger.class, "square");
            multiplyMethodHandle = handleFor(BigInteger.class, "multiply",
                    BigInteger.class);
            mulAddMethodHandle = handleFor(BigInteger.class, "mulAdd",
                    int[].class, int[].class, int.class, int.class, int.class);
            implMultiplyToLenHandle = handleFor(BigInteger.class, "implMultiplyToLen",
                    int[].class, int.class, int[].class, int.class, int[].class);
            implSquareToLenHandle = handleFor(BigInteger.class, "implSquareToLen",
                    int[].class, int.class, int[].class, int.class);
        } catch (ReflectiveOperationException roe) {
            throw new AssertionError("Reflection exception: " + roe, roe);
        }
    }

    private static final Sink s = new Sink();

    @Param({"1", "2", "3", "10", "50", "90", "127"})
    int size;

    BigInteger bi;
    BigInteger bi_copy; // for multiplication

    int[] out;
    int[] in;

    @Setup(Level.Trial)
    public void init() {
        byte[] src = new byte[4 * size];
        for(int i = 0; i < src.length; i++) {
            src[i] = (byte)(0xFF);
        }
        bi = new BigInteger(1, src);
        bi_copy = new BigInteger(1, src);
        out = new int[2 * size];
        for (int i = 0; i < 2 * size; i++) {
            out[i] = i; // anything is fine
        }
        in = new int[size];
        for (int i = 0; i < size; i++) {
            in[i] = i; // anything is fine
        }
    }

    // @Benchmark
    // public BigInteger square() throws Throwable {
    //     return (BigInteger)squareMethodHandle.invokeExact(bi);
    // }

    // @Benchmark
    // public int mulAdd() throws Throwable {
    //     return (int)mulAddMethodHandle.invokeExact(out, in, OFFSET, size, K);
    // }

    // @Benchmark
    // public int[] implSquareToLen() throws Throwable {
    //     return (int[])implSquareToLenHandle.invokeExact(in, in.length, out, out.length);
    // }

    // @Benchmark
    // public int[] implMutliplyToLen() throws Throwable {
    //     return (int[])implMultiplyToLenHandle.invokeExact(in, in.length, in, in.length, out);
    // }

    @Benchmark
    public void square() throws Throwable {
        s.drain((BigInteger)squareMethodHandle.invokeExact(bi));
    }

    @Benchmark
    public void multiply() throws Throwable {
        s.drain((BigInteger)multiplyMethodHandle.invokeExact(bi, bi_copy));
    }

    @Benchmark
    public void mulAdd() throws Throwable {
        s.drain((int)mulAddMethodHandle.invokeExact(out, in, OFFSET, size, K));
    }

    @Benchmark
    public void implSquareToLen() throws Throwable {
        s.drain((int[])implSquareToLenHandle.invokeExact(in, in.length, out, out.length));
    }

    @Benchmark
    public void implMutliplyToLen() throws Throwable {
        s.drain((int[])implMultiplyToLenHandle.invokeExact(in, in.length, in, in.length, out));
    }
}