# This file is a part of Julia. License is MIT: https://julialang.org/license

module DSFMT

import Base: copy, copy!, ==, hash
using Base.GMP.MPZ

export DSFMT_state, dsfmt_get_min_array_size, dsfmt_get_idstring,
       dsfmt_init_gen_rand, dsfmt_init_by_array, dsfmt_gv_init_by_array,
       dsfmt_fill_array_close_open!, dsfmt_fill_array_close1_open2!

"Mersenne Exponent"
const MEXP = 19937

"DSFMT internal state array size of N 128-bit integers."
const N = floor(Int, ((MEXP - 128) / 104 + 1))

"""
Julia DSFMT state representation size counted in 32-bit integers.

Size: (DSFMT state array of Int128 + 1)*4 + Int32 index + Int32 padding
"""
const JN32 = (N+1)*4+1+1


struct DSFMT_state
    val::Vector{Int32}

    function DSFMT_state(val::Vector{Int32} = zeros(Int32, JN32))
        length(val) == JN32 ||
            throw(DomainError(length(val), "Expected length: $JN32."))
        new(val)
    end
end

copy!(dst::DSFMT_state, src::DSFMT_state) = (copyto!(dst.val, src.val); dst)
copy(src::DSFMT_state) = DSFMT_state(copy(src.val))

==(s1::DSFMT_state, s2::DSFMT_state) = s1.val == s2.val

hash(s::DSFMT_state, h::UInt) = hash(s.val, h)


## wrapper functions

function dsfmt_get_idstring()
    idstring = ccall((:dsfmt_get_idstring,:libdSFMT),
                     Ptr{UInt8},
                     ())
    return unsafe_string(idstring)
end

function dsfmt_get_min_array_size()
    min_array_size = ccall((:dsfmt_get_min_array_size,:libdSFMT),
                           Int32,
                           ())
end

const dsfmt_min_array_size = dsfmt_get_min_array_size()

function dsfmt_init_gen_rand(s::DSFMT_state, seed::UInt32)
    ccall((:dsfmt_init_gen_rand,:libdSFMT),
          Cvoid,
          (Ptr{Cvoid}, UInt32,),
          s.val, seed)
end

function dsfmt_init_by_array(s::DSFMT_state, seed::Vector{UInt32})
    ccall((:dsfmt_init_by_array,:libdSFMT),
          Cvoid,
          (Ptr{Cvoid}, Ptr{UInt32}, Int32),
          s.val, seed, length(seed))
end

function dsfmt_gv_init_by_array(seed::Vector{UInt32})
    ccall((:dsfmt_gv_init_by_array,:libdSFMT),
          Cvoid,
          (Ptr{UInt32}, Int32),
          seed, length(seed))
end

function dsfmt_fill_array_close1_open2!(s::DSFMT_state, A::Ptr{Float64}, n::Int)
    @assert Csize_t(A) % 16 == 0 # the underlying C array must be 16-byte aligned
    @assert dsfmt_min_array_size <= n && iseven(n)
    ccall((:dsfmt_fill_array_close1_open2,:libdSFMT),
          Cvoid,
          (Ptr{Cvoid}, Ptr{Float64}, Int),
          s.val, A, n)
end

function dsfmt_fill_array_close_open!(s::DSFMT_state, A::Ptr{Float64}, n::Int)
    @assert Csize_t(A) % 16 == 0 # the underlying C array must be 16-byte aligned
    @assert dsfmt_min_array_size <= n && iseven(n)
    ccall((:dsfmt_fill_array_close_open,:libdSFMT),
          Cvoid,
          (Ptr{Cvoid}, Ptr{Float64}, Int),
          s.val, A, n)
end

# dSFMT jump poly calculation

"Represents a polynomial in GF(2)[X]"
struct GF2X
    z::BigInt
end

GF2X(s::AbstractString) = GF2X(parse(BigInt, reverse(s), base = 16))
Base.string(f::GF2X) = reverse(string(f.z, base = 16))
Base.:(==)(f::GF2X, g::GF2X) = f.z == g.z
Base.copy(f::GF2X) = GF2X(MPZ.set(f.z))

degree(f::GF2X) = Base.ndigits0z(f.z, 2) - 1
coeff(f::GF2X, i) = MPZ.tstbit(f.z, i)
setcoeff!(f::GF2X, i) = (MPZ.setbit!(f.z, i); f)

mulx!(f::GF2X, c::Int=1) = (MPZ.mul_2exp!(f.z, c), f)
xor!(f::GF2X, g::GF2X) = (MPZ.xor!(f.z, g.z); f)

# compute f*X mod m
function mulxmod!(f::GF2X, m::GF2X, deg=degree(m))::GF2X
    # precondition: degree(f) < degree(m)
    mulx!(f)
    degree(f) == deg && xor!(f, m)
    f
end

# cache for X^(2i) mod m
const _squares = Dict{GF2X, Vector{GF2X}}()

# compute f^2 mod m
function sqrmod!(f::GF2X, m::GF2X)::GF2X
    d = degree(m)-1
    0 <= degree(f) <= d || throw(DomainError("f must satisfy 0 <= degree(f) <= degree(m)-1"))
    sqrs = get!(_squares, m) do
        x2i = GF2X(1)
        GF2X[copy(mulxmod!(mulxmod!(x2i, m, d+1), m, d+1)) for i=1:d]
    end
    foldl(filter(i->coeff(f, i), 0:degree(f)); init=GF2X(0)) do g, i
        i <= d÷2 ? # optimization for "simple" squares
            setcoeff!(g, 2i) :
            xor!(g, sqrs[i])
    end
end

# compute X^e mod m
function powxmod(e::BigInt, m::GF2X)::GF2X
    e < 0 && throw(DomainError("e must be >= 0"))
    foldl(Base.ndigits0z(e, 2)-1:-1:0; init=GF2X(1)) do f, i
        MPZ.tstbit(e, i) ?
            mulxmod!(sqrmod!(f, m), m) :
            sqrmod!(f, m)
    end
end

"Cached jump polynomials for `MersenneTwister`."
const JumpPolys = Dict{BigInt,GF2X}()

const CharPoly_ref = Ref{GF2X}()
# Ref because it can not be initialized at load time
function CharPoly()
    if !isassigned(CharPoly_ref)
        CharPoly_ref[] = GF2X(Poly19937)
    end
    return CharPoly_ref[]
end

"""
    calc_jump(steps::Integer)

Compute the dSFMT jump polynomial for `steps` steps (by default for
the Mersenne-Twister with exponent 19937). Note that `steps` should be
less than the period (e.g. ``steps ≪ 2^19937-1``).
"""
function calc_jump(steps::Integer,
                   charpoly::GF2X=CharPoly())::GF2X
    steps < 0 && throw(DomainError("jump steps must be >= 0 (got $steps)"))
    if isempty(JumpPolys)
        JumpPolys[big(10)^20] = GF2X(JPOLY1e20)
    end
    get!(JumpPolys, steps) do
        powxmod(big(steps), charpoly)
    end
end


## dSFMT jump

function dsfmt_jump(s::DSFMT_state, jp::GF2X)
    degree(jp) < degree(CharPoly()) || throw(DomainError("degree of jump polynomial must be < $(degree(CharPoly())), got $(degree(jp))"))

    val = s.val
    nval = length(val)
    index = val[nval - 1]
    work = zeros(Int32, JN32)
    rwork = reinterpret(UInt64, work)
    dsfmt = Vector{UInt64}(undef, nval >> 1)
    ccall(:memcpy, Ptr{Cvoid}, (Ptr{UInt64}, Ptr{Int32}, Csize_t),
          dsfmt, val, (nval - 1) * sizeof(Int32))
    dsfmt[end] = UInt64(N*2)

    for i in 0:degree(jp)
        coeff(jp, i) && dsfmt_jump_add!(rwork, dsfmt)
        dsfmt_jump_next_state!(dsfmt)
    end

    rwork[end] = index
    return DSFMT_state(work)
end

function dsfmt_jump_add!(dest::AbstractVector{UInt64}, src::Vector{UInt64})
    dp = dest[end] >> 1
    sp = src[end] >> 1
    diff = ((sp - dp + N) % N)
    i = 1
    while i <= N-diff
        j = i*2-1
        p = j + diff*2
        dest[j]   ⊻= src[p]
        dest[j+1] ⊻= src[p+1]
        i += 1
    end
    while i <= N
        j = i*2-1
        p = j + (diff - N)*2
        dest[j]   ⊻= src[p]
        dest[j+1] ⊻= src[p+1]
        i += 1
    end
    dest[N*2+1] ⊻= src[N*2+1]
    dest[N*2+2] ⊻= src[N*2+2]
    return dest
end

function dsfmt_jump_next_state!(mts::Vector{UInt64})
    POS1 = 117
    SL1  = 19
    SR   = 12
    MSK1 = 0x000ffafffffffb3f
    MSK2 = 0x000ffdfffc90fffd

    idx = (mts[end] >> 1) % N

    a = idx*2+1
    b = ((idx + POS1) % N)*2+1
    u = N*2+1

    t0 = mts[a]
    t1 = mts[a+1]
    L0 = mts[u]
    L1 = mts[u+1]
    mts[u]   = xor(t0 << SL1, L1 >> 32, L1 << 32, mts[b])
    mts[u+1] = xor(t1 << SL1, L0 >> 32, L0 << 32, mts[b+1])
    mts[a]   = xor(mts[u]   >> SR, mts[u]   & MSK1, t0)
    mts[a+1] = xor(mts[u+1] >> SR, mts[u+1] & MSK2, t1)

    mts[end] = (mts[end] + 2) % (N*2)
    return mts
end

"Jump polynomial for 10^20 steps for dSFMT with exponent 19937"
const JPOLY1e20 = "[encoded data, 5046 characters]"

"Characteristic polynomial used by dSFMT with exponent 19937"
const Poly19937 = "[encoded data, 5046 characters]"

end # module
