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

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

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

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

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

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

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

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

const MAX_SUCCESS = 100;

export function findMinimalCounterExample<A>(
  x: A,
  predicate: Predicate<A>,
  shrinker: Shrinker<A>,
  depth: number = 0,
): ShrinkResult<A> {
  let xs = shrinker(x);
  let counterexample = x;
  let shrinks = depth;
  for (let y of xs) {
    if (!predicate(y)) {
      let shrink = 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 {
  let rng = new Prando(Math.random() * 1000);
  return gen(rng)(100);
}

export function property<A>(
  generator: Gen<A>,
  predicate: Predicate<A>,
  shrinker?: Shrinker<A>,
): Property<A> {
  return (seed: number, size: number) => {
    let rng = new Prando(seed);
    let gen = generator(rng);
    let i = 0;
    for (; i < MAX_SUCCESS; i++) {
      let x = gen(size);
      if (!predicate(x)) {
        let counterexample = {counterexample: x, shrinks: 0};
        if (shrinker !== undefined) {
          counterexample = findMinimalCounterExample(x, predicate, shrinker);
        }
        return {seed, result: 'Falsified', tests: i + 1, counterexample};
      }
      size++;
    }
    return {seed, result: 'OK', tests: i, counterexample: null};
  };
}

// basic generators

export function listOf<A>(gen: Gen<A>): Gen<A[]> {
  return (rng: Prando) => {
    let g = gen(rng);
    return (size: number) => {
      let 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) => {
    let gk = genK(rng);
    let gv = genV(rng);
    return (size: number) => {
      let 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) => {
    let gk = genK(rng);
    let gv = genV(rng);
    return (size: number) => {
      let 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) => {
    let 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]> {
  let keys = Array.from(map.keys());
  return (rng: Prando) => {
    let gkey = elements(keys)(rng);
    return (size: number) => {
      let key = gkey(size);
      let 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 let someInt: Gen<number> = (rng: Prando) => (size: number) => {
  const bound = Math.pow(2, size);
  return rng.nextInt(-bound, bound);
};

/// Choose some random positive integer.
export let 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[][] {
  let result = [];

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

  if (arr.length > 3) {
    let half = Math.floor(arr.length / 2);
    let copy = arr.slice();
    let 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 = {
  run<A>(prop: Property<A>): TestResult<A> {
    let seed = Math.floor(Math.random() * 1000000);
    return prop(seed, 0);
  },
};

// helpers

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