#pragma once

#include <iostream>
#include <exception>
#include <functional>
#include <vector>
#include <cassert>
#include <cstdint>
#include <type_traits>

using namespace std;

// Used for decltype(lambda(args ...)).
template<class A> A undefined() {
  throw exception();
}

template<template<class T> class F, class T> class remove_outer { };
template<template<class T> class F, class A> class remove_outer<F, F<A>> {
public:
  typedef A type;
};

class Source {
public:
  virtual uint64_t readBits(int bits) = 0;
};

template<class A> class Gen {
public:
  function<A(Source&)> run;

  template<class F> Gen(const F& f)
    : run(std::move(f)) { }

  template<class F> auto map (const F& f) const
  -> Gen<decltype(f(undefined<A>()))> {
    Gen<A> lhs = *this;
    return Gen<decltype(f(undefined<A>()))>([lhs, f] (Source& s) {
      const auto a = lhs.run(s);
      return f(a);
    });
  }

  // LOL, this signature.
  template<class F> auto flatMap (const F& f) const
  -> Gen<typename remove_outer<Gen, decltype(f(undefined<A>()))>::type> {
    Gen<A> lhs = *this;
    typedef typename remove_outer<Gen, decltype(f(undefined<A>()))>::type Result;

    return Gen<Result>([lhs, f] (Source& s) {
      const auto a = lhs.run(s);
      const auto gb = f(a);
      return gb.run(s);
    });
  }

  template<class B> Gen<B*> widen() const {
    return map([] (A a) -> B* {
      return a;
    });
  }
};

template<class A, class B> Gen<pair<A, B>> operator * (const Gen<A>& lhs, const Gen<B>& rhs) {
  return Gen<pair<A, B>>([lhs, rhs] (Source& s) -> pair<A, B> {
    const auto a = lhs.run(s);
    const auto b = rhs.run(s);
    return make_pair(a, b);
  });
}

template<class F, class A, class B>
auto mapN (const F& f, const Gen<A>& genA, const Gen<B>& genB) -> Gen<typename result_of<decltype(f)>::type> {
  return (genA * genB).map([f] (const pair<A, B>& p) {
    return f(p.first, p.second);
  });
}
template<class F, class A, class B, class C>
auto mapN (const F& f, const Gen<A>& genA, const Gen<B>& genB, const Gen<C>& genC) -> Gen<typename result_of<decltype(f)>::type> {
  return (genA * genB * genC).map([f] (const pair<pair<A, B>, C>& p) {
    return f(p.first.first, p.first.second, p.second);
  });
}

namespace gen {
  template<class A> Gen<A> pure(const A& a) {
    return Gen<A>([a] (Source& s) { return a; });
  }

  template<class F> auto delay(const F& suspension)
  -> Gen<typename remove_outer<Gen, decltype(suspension())>::type> {
    typedef typename remove_outer<Gen, decltype(suspension())>::type Result;
    return Gen<Result>([suspension] (Source& s) {
      return suspension().run(s);
    });
  }

  Gen<int> int_ = Gen<int>([] (Source& s) { return (int) s.readBits(32); });

  Gen<int> intBits(int bits) {
    assert(bits >= 1);
    assert(bits <= 32);

    return Gen<int>([bits] (Source& s) { return (int) s.readBits(bits); });
  }

  template<class A> Gen<vector<A>> vecN(const Gen<A>& ga, size_t size) {
    return Gen<vector<A>>([ga, size] (Source& s) {
      std::cerr << size << std::endl;
      vector<A> result;
      result.reserve(size);
      for (size_t i = 0; i < size; i++) {
        result.push_back(ga.run(s));
      }
      return result;
    });
  }

  template<class A> Gen<vector<A>> vec(const Gen<A>& ga) {
    // For illustrative purposes limit the size to 255.
    return intBits(8).flatMap([ga] (int size) -> Gen<vector<A>> { return vecN(ga, size); });
  }

  Gen<bool> bool_ = Gen<bool>([] (Source& s) { return s.readBits(1) == 1; });

  template<class T> class type {
  public:
    template<class A, class B> static Gen<T> of(const Gen<A>& genA, const Gen<B>& genB) {
      return (genA * genB).map([] (const pair<A, B>& p) {
        return T(p.first, p.second);
      });
    }
    template<class A, class B, class C> static Gen<T> of(const Gen<A>& genA, const Gen<B>& genB, const Gen<C>& genC) {
      return (genA * genB * genC).map([] (const pair<pair<A, B>, C>& p) {
        return T(p.first.first, p.first.second, p.second);
      });
    }
  };

  template<class T> class new_ {
  public:
    template<class A, class B> static Gen<T*> of(const Gen<A>& genA, const Gen<B>& genB) {
      return (genA * genB).map([] (const pair<A, B>& p) {
        return new T(p.first, p.second);
      });
    }
    template<class A, class B, class C> static Gen<T*> of(const Gen<A>& genA, const Gen<B>& genB, const Gen<C>& genC) {
      return (genA * genB * genC).map([] (const pair<pair<A, B>, C>& p) {
        return new T(p.first.first, p.first.second, p.second);
      });
    }
  };
};

// For illustration purposes.
class RandomSource : public Source {
public:
  virtual uint64_t readBits(int bits) {
    uint64_t mask = (((uint64_t) 1) << bits) - 1;
    uint64_t value = (((uint64_t) rand()) << 32) | ((uint64_t) rand());
    return value & mask;
  }
};