import {Option, getOrElse, isSome, none, some} from 'fp-ts/lib/Option.js';
import {LazyArg, pipe} from 'fp-ts/lib/function.js';
import Prando from 'prando';

type Result = 'OK' | 'Falsified' | 'Exception';

type ShrinkResult<A> = {counterexample: A; shrinks: number};

type TestResult<A> = {
  seed: number | null;
  result: Result;
  labels: string[];
  counterexample: ShrinkResult<A> | null;
};

type Predicate<A> = (a: A) => Promise<boolean>;

export type Gen<A> = (rng: Prando) => (size: number) => A;

type Shrinker<A> = (a: A) => A[];

/// Label test input with some arbitrary string for classification purpose
type Label<A> = (a: A) => string;

/// classify test result
export function classify<A>(results: TestResult<A>): [number, string][] {
  if (results.labels.length == 0) return [];

  const distribution = new Map<string, number>();
  const result: [number, string][] = [];
  results.labels.sort();
  results.labels.forEach(s => {
    const num = distribution.get(s) || 0;
    distribution.set(s, num + 1);
  });

  distribution.forEach((v, k) => {
    result.push([(v / results.labels.length) * 100.0, k]);
  });

  return result;
}

type Property<A> = (seed: number, size: number) => Promise<TestResult<A>>;

const MAX_SUCCESS = 100;

export async function findMinimalCounterExample<A>(
  x: A,
  predicate: Predicate<A>,
  shrinker: Shrinker<A>,
  depth = 0,
): Promise<ShrinkResult<A>> {
  const xs = shrinker(x);
  let counterexample = x;
  let shrinks = depth;
  for (const y of xs) {
    if (!(await predicate(y))) {
      const shrink = await findMinimalCounterExample(
        y,
        predicate,
        shrinker,
        depth + 1,
      );
      if (shrink.shrinks > depth) {
        counterexample = shrink.counterexample;
        shrinks = shrink.shrinks;
      }
    }
  }
  return {counterexample, shrinks};
}

export function generate<A>(gen: Gen<A>): A {
  const rng = new Prando(Math.random() * 1000);
  return gen(rng)(100);
}

export async function property<A>(
  generator: Gen<A>,
  predicate: Predicate<A>,
  shrinker?: Shrinker<A>,
  label?: Label<A>,
): Promise<Property<A>> {
  return async (seed: number, size: number) => {
    const rng = new Prando(seed);
    const gen = generator(rng);
    const labels: string[] = [];

    let i = 0;
    for (; i < MAX_SUCCESS; i++) {
      const x = gen(size);
      if (!(await predicate(x))) {
        let counterexample = {counterexample: x, shrinks: 0};
        if (shrinker !== undefined) {
          counterexample = await findMinimalCounterExample(
            x,
            predicate,
            shrinker,
          );
        }
        return {
          seed,
          result: 'Falsified',
          tests: i + 1,
          labels: [],
          counterexample,
        };
      }
      if (label !== undefined) {
        labels.push(label(x));
      }
      size++;
    }
    return {seed, result: 'OK', tests: i, labels, counterexample: null};
  };
}

// basic generators

export function listOf<A>(gen: Gen<A>): Gen<A[]> {
  return (rng: Prando) => {
    const g = gen(rng);
    return (size: number) => {
      const result = [];
      for (let i = 0; i < size; i++) {
        result.push(g(size));
      }
      return result;
    };
  };
}

/// generates a map with exactly size elements
export function mapOf<K, V>(genK: Gen<K>, genV: Gen<V>): Gen<Map<K, V>> {
  return (rng: Prando) => {
    const gk = genK(rng);
    const gv = genV(rng);
    return (size: number) => {
      const result = new Map<K, V>();
      for (let i = 0; i < size; i++) {
        result.set(gk(size), gv(size));
      }
      return result;
    };
  };
}

/// generates a map with at least size+1 elements. As size is expected
/// to be positive or null this guarantees at least one element hence
/// the name
export function mapOf1<K, V>(genK: Gen<K>, genV: Gen<V>): Gen<Map<K, V>> {
  return (rng: Prando) => {
    const gk = genK(rng);
    const gv = genV(rng);
    return (size: number) => {
      const result = new Map<K, V>();
      for (let i = 0; i <= size; i++) {
        result.set(gk(size), gv(size));
      }
      return result;
    };
  };
}

/// Select randomly one element from given list
export function elements<A>(arr: A[]): Gen<A> {
  return (rng: Prando) => (_size: number) => {
    const index = rng.nextInt(0, arr.length - 1);
    return arr[index];
  };
}

export function frequency<A>(distribution: [number, Gen<A>][]): Gen<A> {
  const total: number = distribution.reduce((acc, [freq, _]) => acc + freq, 0);

  return (rng: Prando) => (size: number) => {
    const idx = Math.floor(rng.next(0, total));
    const [selectedGen, _] = distribution.reduce(
      ([found, sum], [freq, gen]) => {
        const current = sum + freq;
        if (isSome(found)) {
          return [found, current];
        } else if (idx <= current) {
          return [some(gen), current];
        } else {
          return [none, current];
        }
      },
      [none as Option<Gen<A>>, 0],
    );
    const generator: Gen<A> = pipe(
      selectedGen,
      getOrElse((() => {
        throw new Error(
          `invalid frequency distribution ${distribution}, idx=${idx}`,
        );
      }) as LazyArg<Gen<A>>),
    );
    return generator(rng)(size);
  };
}

/// Select one key-value pair from given map
export function keyValueOf<K, V>(map: Map<K, V>): Gen<[K, V]> {
  const keys = Array.from(map.keys());
  return (rng: Prando) => {
    const gkey = elements(keys)(rng);
    return (size: number) => {
      const key = gkey(size);
      const value = map.get(key)!;
      return [key, value];
    };
  };
}

/// Choose number between some bounds.
export function choose(min: number, max: number): Gen<number> {
  return (rng: Prando) => (_size: number) => rng.nextInt(min, max);
}

/// Choose some random integer.
export const someInt: Gen<number> = (rng: Prando) => (size: number) => {
  const bound = Math.pow(2, size);
  return rng.nextInt(-bound, bound);
};

/// Choose some random positive integer.
export const positiveInt: Gen<number> = (rng: Prando) => (size: number) => {
  const bound = Math.pow(2, Math.round(Math.sqrt(size)));
  return rng.nextInt(0, bound);
};

/// Generate a pair of generated values
export function pairOf<A>(gen: Gen<A>): Gen<[A, A]> {
  return (rng: Prando) => (size: number) => [gen(rng)(size), gen(rng)(size)];
}

// basic shrinkers
export function shrinklist<A>(arr: A[]): A[][] {
  const result = [];

  const tail = arr.slice(1);
  const keep1 = arr.slice();
  keep1.splice(1, 1);
  result.push(tail, keep1);

  if (arr.length > 3) {
    const half = Math.floor(arr.length / 2);
    const copy = arr.slice();
    const secondHalf = copy.slice(half);
    result.push(secondHalf);
  }

  return result;
}

export function shrinkPositive(i: number): number[] {
  if (i === 0) return [];

  return [i - 1, Math.round(Math.sqrt(i))];
}

// shrink pair
export function shrinkPair<A, B>(
  shrinkA: Shrinker<A>,
  shrinkB: Shrinker<B>,
): Shrinker<[A, B]> {
  return ([x, y]) => {
    const shrinkx = shrinkA(x);
    const shrinky = shrinkB(y);
    // assumes x >= y
    const resulta: [A, B][] = shrinkx.map(e => [e, y]);
    const resultb: [A, B][] = shrinky.map(e => [x, e]);
    const resultboth: [A, B][] = shrinkx.map((e, i) => [e, shrinky[i]]);
    return resulta.concat(resultb, resultboth);
  };
}

export const Property = {
  async run<A>(prop: Promise<Property<A>>): Promise<TestResult<A>> {
    const seed = Math.floor(Math.random() * 1000000);
    const p = await prop;
    return await p(seed, 0);
  },
};

// helpers

export function forall<A, B>(
  gen: Gen<A>,
  k: (a: A) => Promise<Property<B>>,
): Promise<Property<B>> {
  const p = async (seed: number, size: number) => {
    const rng = new Prando(seed);
    const a = gen(rng)(size);
    const p = await k(a);
    return await p(seed, size);
  };

  return Promise.resolve(p);
}