//! returns: 0
//! Mersenne Twister PRNG with statistical testing.

constant N: u32 = 624;
constant M: u32 = 397;
constant MATRIX_A: u32 = 0x9908B0DF;
constant UPPER_MASK: u32 = 0x80000000;
constant LOWER_MASK: u32 = 0x7FFFFFFF;

record MtState: Copy {
    mt: *mut [u32],
    mti: u32,
}

fn mtInit(s: *mut MtState, seed: u32) {
    set s.mt[0] = seed;
    let mut i: u32 = 1;
    while i < N {
        let prev: u32 = s.mt[i - 1];
        let xored: u32 = prev ^ (prev >> 30);

        let lo: u32 = xored & 0xFFFF;
        let hi: u32 = xored >> 16;
        let c: u32 = 1812433253;
        let clo: u32 = c & 0xFFFF;
        let chi: u32 = c >> 16;

        let ll: u32 = clo * lo;
        let lh: u32 = clo * hi;
        let hl: u32 = chi * lo;

        let result: u32 = ll + ((lh + hl) << 16) + i;
        set s.mt[i] = result;
        set i += 1;
    }
    set s.mti = N;
}

fn generateNumbers(s: *mut MtState) {
    let mut i: u32 = 0;

    while i < N - M {
        let y: u32 = (s.mt[i] & UPPER_MASK) | (s.mt[i + 1] & LOWER_MASK);
        let mut mag: u32 = 0;
        if y & 1 == 1 {
            set mag = MATRIX_A;
        }
        set s.mt[i] = s.mt[i + M] ^ (y >> 1) ^ mag;
        set i += 1;
    }

    while i < N - 1 {
        let y: u32 = (s.mt[i] & UPPER_MASK) | (s.mt[i + 1] & LOWER_MASK);
        let mut mag: u32 = 0;
        if y & 1 == 1 {
            set mag = MATRIX_A;
        }
        set s.mt[i] = s.mt[i + M - N] ^ (y >> 1) ^ mag;
        set i += 1;
    }

    let y: u32 = (s.mt[N - 1] & UPPER_MASK) | (s.mt[0] & LOWER_MASK);
    let mut mag: u32 = 0;
    if y & 1 == 1 {
        set mag = MATRIX_A;
    }
    set s.mt[N - 1] = s.mt[M - 1] ^ (y >> 1) ^ mag;

    set s.mti = 0;
}

fn mtNext(s: *mut MtState) -> u32 {
    if s.mti >= N {
        generateNumbers(s);
    }

    let mut y: u32 = s.mt[s.mti];
    set s.mti += 1;

    set y ^= (y >> 11);
    set y ^= ((y << 7) & 0x9D2C5680);
    set y ^= ((y << 15) & 0xEFC60000);
    set y ^= (y >> 18);

    return y;
}

fn testKnownSequence(s: *mut MtState) -> i32 {
    mtInit(s, 1);

    let v0: u32 = mtNext(s);
    assert v0 == 1791095845;

    let v1: u32 = mtNext(s);
    assert v1 == 4282876139;

    let v2: u32 = mtNext(s);
    assert v2 == 3093770124;

    let v3: u32 = mtNext(s);
    assert v3 == 4005303368;

    let v4: u32 = mtNext(s);
    assert v4 == 491263;

    return 0;
}

fn testDeterminism(s: *mut MtState) -> i32 {
    mtInit(s, 42);

    let mut first: [u32; 10] = [0; 10];
    let mut i: u32 = 0;
    while i < 10 {
        set first[i] = mtNext(s);
        set i += 1;
    }

    mtInit(s, 42);

    set i = 0;
    while i < 10 {
        let v: u32 = mtNext(s);
        if v <> first[i] { return i as i32 + 1; }
        set i += 1;
    }

    return 0;
}

fn testDifferentSeeds(s: *mut MtState) -> i32 {
    mtInit(s, 1);
    let a: u32 = mtNext(s);

    mtInit(s, 2);
    let b: u32 = mtNext(s);

    mtInit(s, 3);
    let c: u32 = mtNext(s);

    assert a <> b;
    assert b <> c;
    assert a <> c;

    return 0;
}

fn testChiSquared(s: *mut MtState) -> i32 {
    mtInit(s, 12345);

    constant NUM_BINS: u32 = 16;
    constant NUM_SAMPLES: u32 = 1600;
    constant EXPECTED: u32 = 100;

    let mut bins: [u32; 16] = [0; 16];

    let mut i: u32 = 0;
    while i < NUM_SAMPLES {
        let val: u32 = mtNext(s);
        let bin: u32 = val >> 28;
        set bins[bin] += 1;
        set i += 1;
    }

    let mut chi2Scaled: u32 = 0;
    let mut b: u32 = 0;
    while b < NUM_BINS {
        let obs: i32 = bins[b] as i32;
        let exp: i32 = EXPECTED as i32;
        let diff: i32 = obs - exp;
        set chi2Scaled += (diff * diff) as u32;
        set b += 1;
    }

    assert chi2Scaled <= 5000;

    set b = 0;
    while b < NUM_BINS {
        assert bins[b] <> 0;
        set b += 1;
    }

    set b = 0;
    while b < NUM_BINS {
        assert bins[b] <= NUM_SAMPLES / 2;
        set b += 1;
    }

    return 0;
}

fn testRegeneration(s: *mut MtState) -> i32 {
    mtInit(s, 7);

    let mut last: u32 = 0;
    let mut i: u32 = 0;
    while i < 700 {
        set last = mtNext(s);
        set i += 1;
    }

    mtInit(s, 7);
    let mut last2: u32 = 0;
    set i = 0;
    while i < 700 {
        set last2 = mtNext(s);
        set i += 1;
    }

    assert last == last2;

    let mut more: u32 = 0;
    set i = 0;
    while i < 700 {
        set more = mtNext(s);
        set i += 1;
    }
    assert more <> 0;

    return 0;
}

fn testBitCoverage(s: *mut MtState) -> i32 {
    mtInit(s, 999);

    let mut orAll: u32 = 0;
    let mut andAll: u32 = 0xFFFFFFFF;

    let mut i: u32 = 0;
    while i < 200 {
        let v: u32 = mtNext(s);
        set orAll |= v;
        set andAll &= v;
        set i += 1;
    }

    assert orAll == 0xFFFFFFFF;
    assert andAll <> 0xFFFFFFFF;
    assert andAll == 0;

    return 0;
}

@default fn main() -> i32 {
    let mut mt: [u32; 624] = [0; 624];
    let mut s: MtState = MtState {
        mt: &mut mt[..],
        mti: 625,
    };

    let r1: i32 = testKnownSequence(&mut s);
    if r1 <> 0 { return 10 + r1; }

    let r2: i32 = testDeterminism(&mut s);
    if r2 <> 0 { return 20 + r2; }

    let r3: i32 = testDifferentSeeds(&mut s);
    if r3 <> 0 { return 30 + r3; }

    let r4: i32 = testChiSquared(&mut s);
    if r4 <> 0 { return 40 + r4; }

    let r5: i32 = testRegeneration(&mut s);
    if r5 <> 0 { return 50 + r5; }

    let r6: i32 = testBitCoverage(&mut s);
    if r6 <> 0 { return 60 + r6; }

    return 0;
}