package main

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

type Config struct {
	PrimeList bool
	DontLoad bool
	NumRoutines int
	UseTensors bool
	UseTxtRead bool
	UseTxtWrite bool
}

func calculateProgress(start, end, index int64) float32 {
	steps := (end - start)
	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, config Config) []int64 {
	var primes []int64

	bootTime := time.Now()

	if config.DontLoad {
		primes = []int64{2, 3, 5}
	} else {
		primes = loadPrimes(config.UseTxtRead)
	}

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

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

	startTime := time.Now()

	if config.UseTensors {
		primes = generatePrimesGPU(upperLimit, primes, progress)
	} else {
		primes = generatePrimesCPU(upperLimit, primes, config.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)

	continueFrom := primes[len(primes) - 1] + 2
	for i := continueFrom; i <= upperLimit; i += 2 {
		select {
		case progress <- calculateProgress(continueFrom, upperLimit, 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), 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 config Config

	flag.BoolVar(&config.PrimeList, "p", false, "Only calculate and print prime list")
	flag.BoolVar(&config.DontLoad, "d", false, "Don't load precalculated primes, calculate from 0")
	flag.IntVar(&config.NumRoutines, "r", 0, "How many routines to use for calculation. 0 = number of available CPU cores")
	flag.BoolVar(&config.UseTensors, "t", false, "Use tensorflow")
	flag.BoolVar(&config.UseTxtRead, "read-legacy", false, "Read legacy prime.txt file instead of prime.msgpack")
	flag.BoolVar(&config.UseTxtWrite, "write-legacy", false, "Write legacy prime.txt file instead of prime.msgpack")
	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 config.PrimeList {
		onlyGenerate(number, config)
	} else {
		calculate(number, config)
	}
}

func onlyGenerate(number int64, config Config) {
	primes := generatePrimes(number, config)
	writePrimes(primes, config.UseTxtWrite)
}

func calculate(number int64, config Config) {
	primes := generatePrimes(number, config)

	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)
}