package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"strconv"
	"time"
)

func calculateProgress(start, end, step, index int64) float32 {
	steps := (end - start) / step
	current := index - start
	if current == 0 {
		return 0
	}
	return float32(current) / float32(steps)
}

func printProgress(progress <-chan float32, message string) {
	for {
		time.Sleep(time.Second)
		p, channelOpen := <-progress
		if !channelOpen {
			break
		}
		fmt.Printf("%s %.2f%%\n", message, p*100)
	}
}

func checkIsDivisibleByPrime(number int64, offset, stride int, primes *[]int64, resultChannel chan bool, ctx context.Context) {
	threshold := number / 2
	for i := offset; i < len(*primes); i += stride {
		select {
			case <-ctx.Done():
				resultChannel <- false
				return
			default:
				// no need to continue checking if number is greater than half our number, division is already impossible
				if (*primes)[i] > threshold {
					resultChannel <- false
					return
				}
				// only have to check primes, because all other numbers are multiple of primes.
				// if a number is divisible by a multiple of a prime, it is by definition also divisible by the prime itself.
				if number % (*primes)[i] == 0 {
					resultChannel <- true
					return
				}
		}
	}
	resultChannel <- false
}

func generatePrimes(upperLimit int64, loadList bool, numRoutines int, useTensors bool) []int64 {
	var primes []int64

	bootTime := time.Now()

	if loadList {
		primes = loadPrimes()
	} else {
		primes = []int64{2, 3, 5}
	}

	fmt.Printf("Startup time: %v\n", time.Now().Sub(bootTime))

	progress := make(chan float32)
	go printProgress(progress, "Generating primes")

	startTime := time.Now()

	if useTensors {
		primes = generatePrimesGPU(upperLimit, primes, progress)
	} else {
		primes = generatePrimesCPU(upperLimit, primes, numRoutines, progress)
	}
	close(progress)

	endTime := time.Now()

	fmt.Printf("Prime generation took %v\nLargest prime found: %d\nTotal prime number count: %d\n", endTime.Sub(startTime), primes[len(primes) - 1], len(primes))

	return primes
}

func generatePrimesCPU(upperLimit int64, primes []int64, numRoutines int, progress chan float32) []int64 {
	if numRoutines == 0 {
		numRoutines = runtime.NumCPU()
	}

	fmt.Printf("Calculating with %d routines\n\n", numRoutines)

	for i := (primes[len(primes) - 1] + 2); i <= upperLimit; i += 2 {
		select {
		case progress <- calculateProgress(6, upperLimit, 1, i):
		default:
		}

		isPrime := true

		checkCtx, checkCancel := context.WithCancel(context.Background())
		resultChannel := make(chan bool)

		for r := 0; r < numRoutines; r++ {
			go checkIsDivisibleByPrime(i, r, numRoutines, &primes, resultChannel, checkCtx)
		}

		for r := 0; r < numRoutines; r++ {
			if <-resultChannel {
				isPrime = false
				checkCancel()
				// can't break here, because all channel writes have to be read to avoid deadlocks
			}
		}

		if isPrime {
			primes = append(primes, i)
		}
	}
	return primes
}

func calculatePrimeParts(number int64, primes []int64) []int64 {
	// don't calculate if number is a prime itself
	if primes[len(primes)-1] == number {
		return []int64{number}
	}

	progress := make(chan float32)

	go printProgress(progress, "Calculating")

	var primeParts []int64

	for i := int64(len(primes) - 1); i >= 0; i-- {
		select {
		case progress <- 1.0 - calculateProgress(0, int64(len(primes) - 1), 1, i):
		default:
		}
		flooredDiv := number / primes[i]
		if flooredDiv == 0 {
			continue
		}
		primeParts = append(primeParts, primes[i])
		number = number - (flooredDiv * primes[i])
		if number == 0 {
			break
		}
	}
	close(progress)
	if number != 0 {
		primeParts = append(primeParts, 1)
	}
	return primeParts
}

func main() {
	var (
		primeList bool
		dontLoad bool
		numRoutines int
		useTensors bool
	)

	flag.BoolVar(&primeList, "p", false, "Only calculate and print prime list")
	flag.BoolVar(&dontLoad, "d", false, "Don't load precalculated primes, calculate from 0")
	flag.IntVar(&numRoutines, "r", 0, "How many routines to use for calculation. 0 = number of available CPU cores")
	flag.BoolVar(&useTensors, "t", false, "Use tensorflow")
	flag.Parse()
	numStr := flag.Arg(0)

	if numStr == "" {
		fmt.Printf("Usage: %s [-p|d] <number>\n", os.Args[0])
		flag.PrintDefaults()
		return
	}

	number, err := strconv.ParseInt(numStr, 10, 64)
	if err != nil {
		log.Fatal(err)
	}

	if primeList {
		onlyGenerate(number, dontLoad, numRoutines, useTensors)
	} else {
		calculate(number, dontLoad, numRoutines, useTensors)
	}
}

func onlyGenerate(number int64, dontLoad bool, numRoutines int, useTensors bool) {
	primes := generatePrimes(number, !dontLoad, numRoutines, useTensors)

	file, err := os.Create("prime.txt")
	if err != nil {
		log.Fatal(err)
	}

	writer := bufio.NewWriter(file)

	for _, prime := range primes {
		writer.WriteString(fmt.Sprintf("%d\n", prime))
	}

	writer.Flush()
	file.Close()
}

func calculate(number int64, dontLoad bool, numRoutines int, useTensors bool) {
	primes := generatePrimes(number, !dontLoad, numRoutines, useTensors)

	primeParts := calculatePrimeParts(number, primes)

	var sum int64

	for i, prime := range primeParts {
		sum += prime
		print(prime)

		// if not last element
		if i < len(primeParts) - 1 {
			print(" + ")
		}
	}
	print(" = ")
	println(sum)
}

func loadPrimes() []int64 {
	file, err := os.Open("prime.txt")
	if err != nil {
		log.Fatal(err)
	}

	var primes []int64

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		nextPrime, err := strconv.ParseInt(scanner.Text(), 10, 64)
		if err != nil {
			log.Fatal(err)
		}
		primes = append(primes, nextPrime)
	}

	file.Close()

	return primes
}