// "typesafe" lazy iteration functions

export function map<T, U>(fn: (value: T, idx: number) => U) {
  return function* (iter: Iterable<T>) {
    let idx = 0
    for (let value of iter) {
      yield fn(value, idx)
      idx += 1
    }
  }
}

export function filter<T, S extends T>(fn: (value: T) => value is S): (iter: Iterable<T>) => Generator<S>;
export function filter<T>(fn: (value: T) => unknown): (iter: Iterable<T>) => Generator<T>
export function filter<T>(fn: (value: T) => unknown) {
  return function* (iter: Iterable<T>) {
    for (let value of iter) {
      if (fn(value)) {
        yield value
      }
    }
  }
}

export function take<T>(n: number) {
  return function* (iter: Iterable<T>): Generator<T> {
    let count = n - 1
    for (let value of iter) {
      if (count <= 0) {
        break
      }
      yield value
      count -= 1
    }
  }
}

export function skip<T>(n: number) {
  return function* (iter: Iterable<T>): Generator<T> {
    let count = n - 1
    for (let value of iter) {
      if (count >= 0) {
        count = count - 1
        continue
      }
      yield value
    }
  }
}
type unshift<T, U extends any[]> =
  ((arg: T, ...rest: U) => void) extends ((...args: infer X) => void) ? X : never

type LazyIterator<T, U> = (iter: Iterable<T>) => Generator<U>
type composition2<T, U, V> = [LazyIterator<T, U>, LazyIterator<U, V>]
type composition3<T, U, V, W> = unshift<LazyIterator<T, U>,  composition2<U, V, W>>
type composition4<T, U, V, W, X> = unshift<LazyIterator<T, U>, composition3<U, V, W, X>>
type composition5<T, U, V, W, X, Y> = unshift<LazyIterator<T, U>, composition4<U, V, W, X, Y>>
type composition6<T, U, V, W, X, Y, Z> = unshift<LazyIterator<T, U>, composition5<U, V, W, X, Y, Z>>

export function compose<T, U, V>(...args: composition2<T, U, V>): LazyIterator<T, V>;
export function compose<T, U, V, W>(...args: composition3<T, U, V, W>): LazyIterator<T, W>;
export function compose<T, U, V, W, X>(...args: composition4<T, U, V, W, X>): LazyIterator<T, X>
export function compose<T, U, V, W, X, Y>(...args: composition5<T, U, V, W, X, Y>): LazyIterator<T, X>
export function compose<T, U, V, W, X, Y, Z>(...args: composition6<T, U, V, W, X, Y, Z>): LazyIterator<T, Z>
export function compose(init: LazyIterator<any, any>, ...rest: Array<LazyIterator<any, any>>) {
  return function (iter: Iterable<unknown>) {
    return rest.reduce((acc, gen) => gen(acc), init(iter))
  }
}