#!/usr/bin/env python3

import numpy as np
import math
import argparse
import pathlib
import os

def get_primes_from_2_to(n_):
    """ Input n>=6, Returns a array of primes, 2 <= p <= n """
    n_ += 1
    sieve = np.ones(n_//3 + (n_%6==2), dtype=np.bool)
    for i in range(1,int(n_**0.5)//3+1):
        if sieve[i]:
            k=3*i+1|1
            sieve[       k*k//3     ::2*k] = False
            sieve[k*(k-2*(i&1)+4)//3::2*k] = False
    return np.r_[2,3,((3*np.nonzero(sieve)[0][1:]+1)|1)]

def get_primes(min_, max_, step_):
    primes = get_primes_from_2_to(max_)

    if len(primes) == 0 or primes[-1] < min_:
        return []

    i = 0
    while i < len(primes)-1 and primes[i] < min_:
        i += 1

    filtered_primes = [primes[i]]

    if i < len(primes)-1:
        j = i + 1
        while j < len(primes) and primes[j] <= max_:
            if primes[j] >= primes[i] + step_:
                filtered_primes.append(primes[j])
                i = j
                j = i + 1
            else:
                j += 1

    filtered_primes = np.array(filtered_primes, dtype=np.int64)
    return filtered_primes

def is_prime(n_):
    if n_ == 2:
        return True
    if n_ % 2 == 0 or n_ <= 1:
        return False

    sqr = int(math.sqrt(n_)) + 1

    for divisor in range(3, sqr, 2):
        if n_ % divisor == 0:
            return False
    return True

def check(primes_, args_):
    for p in primes_:
        assert(is_prime(p))
        assert(p >= args_.min)
        assert(p <= args_.max)

def write_hpp(fname_, string_):
    with open(fname_, 'w+') as f:
        f.write(string_)
        print("file generated at: " + fname_)

# default parameters
class defaults:
    step = 131072
    max = 17177772037
    min = step

# parse args
parser = argparse.ArgumentParser(description='generate include/warpcore/primes.hpp', formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('--min', type=int, nargs='?', const=defaults.min, default=defaults.min, help="lower bound of primes to consider")
parser.add_argument('--max', type=int, nargs='?', const=defaults.max, default=defaults.max, help="upper bound of primes to consider")
parser.add_argument('-s', '--step', type=int, nargs='?', const=defaults.step, default=defaults.step, help="minimum offset between consecutive primes")
parser.add_argument('-r', '--replace', action='store_true', help="replace existing .hpp file")
parser.add_argument('-f', '--force_replace', action='store_true', help="replace existing .hpp without asking")
parser.add_argument('--no_check', action='store_true', help="do not check for validity")

args = parser.parse_args()

# constraints
assert(args.step <= args.min <= args.max)
assert(args.min >= 6)

# file header
string  = "#ifndef WARPCORE_PRIMES_HPP\n"
string += "#define WARPCORE_PRIMES_HPP\n\n"
string += "// This file has been generated by utils/generate_primes_hpp.py\n// with the following set of parameters: "
string += "(min=" + str(args.min) + ", max=" + str(args.max) + ", step=" + str(args.step) + ")\n\n"
string += "#include <array>\n"
string += "#include <cstdint>\n\n"
string += "namespace warpcore\n{\n\n"
string += "namespace detail\n{\n\n"
string += "// prime numbers with ~" + str(args.step) + " offset\n"

# static array of primes
primes = get_primes(args.min, args.max, args.step)
assert(len(primes) != 0)

# check if generated are primes an are in valid range
if not args.no_check:
    check(primes, args)

string += "constexpr std::array<std::uint64_t, " + str(len(primes)) + "> primes = \n{\n"

# insert primes into std::array
for p in primes[:-1]:
    string += "    " + str(p) + ",\n"
string += "    " + str(primes[-1]) + "\n};\n\n"

# file footer
string += "} // namespace detail\n\n"
string += "} // namespace warpcore\n\n"
string += "#endif /* WARPCORE_PRIMES_HPP */"

# pass output to stdio or replace .hpp file
if args.replace or args.force_replace:
    fname = str(pathlib.Path(__file__).absolute().parents[1]) + "/include/warpcore/primes.hpp"

    if os.path.isfile(fname):
        if not args.force_replace:
            while True:
                choice = input("file already exists. replace? [y,n]: ")
                if choice is 'y':
                    write_hpp(fname, string)
                    break
                elif choice is 'n':
                    print(string)
                    break
                else:
                    print("please answer with y or n")
        else:
            write_hpp(fname, string)
    else:
        write_hpp(fname, string)
else:
    print(string)