import Prando from 'prando';
import { none } from 'fp-ts/lib/Option.js';
import G from './gen.js';
import { test, expect } from 'vitest';
import { chiSquare, computeChiSquaredPValue, computeCramersV } from './stats.js';

test('can generate and shrink bytes', async () => {
  const prng = G();
  const gen = prng.genBytes();
  const bytes: number[][] = [];
  for (let i = 1; i < 10; i++) {
    bytes.push(gen.generate(i));
  }
  console.log('Generated bytes:', bytes);
  const replay = prng.replay().genBytes();
  for (let i = 1; i < 10; i++) {
    const replayed = replay.generate(i);
    expect(replayed, `Replaying at ${i}`).toEqual(bytes[i - 1]);
  }
});

test('can generate an arbitrary integer', () => {
  const gen = G().genInt();
  const value = gen.generate();
  expect(typeof value).toBe('number');
});

test('can generate N arbitrary integers', () => {
  const gen = G().genInt();
  const value = gen.generateN(5);
  expect(value.length).toBe(5);
});

test('can generate arbitrary integer within bounds', () => {
  const gen = G().genIntWithin(1000, 200000);
  for (let i = 0; i < 100; i++) {
    const value = gen.generate();
    expect(value).toBeGreaterThanOrEqual(1000);
    expect(value).toBeLessThanOrEqual(200000);
  }
});

test('generated integers are uniformly distributed', () => {
  const lb = 0;
  const ub = 100000;
  const gen = G().genIntWithin(lb, ub);
  const range = ub - lb;
  const buckets = 100;
  const numSamples = 100000;
  // indexed by lower bound
  const counts: Record<number, number> = {};
  for (let i = 1; i <= buckets; i++) {
    counts[Math.ceil(range / buckets) * i + lb] = 0;
  }
  const values = [];
  for (let i = 0; i < numSamples; i++) {
    const value = gen.generate();
    values.push(value);
    const bucket = Object.keys(counts).find((b) => value < parseInt(b));
    if (!bucket) {
      console.log('ignoring bucket for value', value);
    } else {
      const b = parseInt(bucket);
      counts[b] = (counts[b] || 0) + 1;
    }
  }

  // compute χ² statistic
  const expectedCount = numSamples / buckets;
  const chiSquared = chiSquare(Object.values(counts), new Array(buckets).fill(expectedCount));
  // compute p-value and Cramér's V.
  const cramerv = computeCramersV(chiSquared, numSamples, buckets);
  const pValue = computeChiSquaredPValue(chiSquared, buckets);
  expect(cramerv, `Chi-squared: ${chiSquared}, Cramer V: ${cramerv}`).toBeLessThan(0.015);
  expect(pValue, `Chi-squared: ${chiSquared}, p-value: ${pValue}`).toBeGreaterThan(0.05);
});

test('can generate uniformly distributed arbitrary strings shorter than 256 characters', () => {
  const gen = G().genString();
  const buckets = 16;
  const numSamples = (1 << 16);
  const lengths = new Array(buckets).fill(0);
  for (let i = 0; i < numSamples; i++) {
    const value = gen.generate();
    lengths[value.length % buckets]++;
    expect(value.length).toBeLessThan(256);
  }

  const chiSquared = chiSquare(lengths, new Array(buckets).fill(numSamples / buckets));
  const cramerv = computeCramersV(chiSquared, numSamples, buckets);
  const pValue = computeChiSquaredPValue(chiSquared, buckets);
  expect(cramerv, `Chi-squared: ${chiSquared}, Cramer V: ${cramerv}`).toBeLessThan(0.015);
  expect(pValue, `Chi-squared: ${chiSquared}, p-value: ${pValue}`).toBeGreaterThan(0.05);
});