/* This script is part of radix_info.
   Copyright (C) 2023  Adrien Hopkins

   This program is free software: you can redistribute it and/or modify
   it under the terms of version 3 of the GNU General Public License
   as published by the Free Software Foundation.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package factors

import (
	"fmt"
	"strings"
)

// PrimeFactorization represents how a number can be represented as a
// product of prime numbers.  This is surprisingly useful for learning
// about number radices.
type PrimeFactorization struct {
	exponents map[uint]uint
}

// PrimeFactorize creates a PrimeFactorization by factorizing a number.
func PrimeFactorize(n uint) PrimeFactorization {
	if n == 0 {
		return PrimeFactorization{map[uint]uint{0: 1}}
	}

	unfactored := n
	exponents := make(map[uint]uint)

	for possibleFactor := uint(2); possibleFactor*possibleFactor <= unfactored; possibleFactor++ {
		for unfactored%possibleFactor == 0 {
			unfactored /= possibleFactor
			exponents[possibleFactor] += 1
		}
	}

	// by this point, unfactored is guaranteed to be 1 or prime
	if unfactored > 1 {
		exponents[unfactored] += 1
	}
	return PrimeFactorization{exponents}
}

// Exponent returns the exponent for the prime p.
func (factors PrimeFactorization) Exponent(p uint) uint {
	return factors.exponents[p]
}

// ExponentMap returns a map mapping primes to their exponents.
func (factors PrimeFactorization) ExponentMap() map[uint]uint {
	exponentMap := make(map[uint]uint, len(factors.exponents))
	for p, e := range factors.exponents {
		exponentMap[p] = e
	}
	return exponentMap
}

// Primes returns a slice containing all of the primes with nonzero exponents
func (factors PrimeFactorization) Primes() []uint {
	primes := make([]uint, 0, len(factors.exponents))
	for p := range factors.exponents {
		if factors.exponents[p] > 0 {
			primes = append(primes, p)
		}
	}
	return primes
}

// Size returns how many primes in this factorization have nonzero exponents.
func (factors PrimeFactorization) Size() int {
	return len(factors.exponents)
}

// String returns a string representation of the prime factorization,
// which looks like "2^2 × 3".
func (factors PrimeFactorization) String() string {
	primesSorted := sortUints(factors.Primes())

	parts := make([]string, 0, factors.Size())
	for _, p := range primesSorted {
		if factors.Exponent(p) == 1 {
			parts = append(parts, fmt.Sprintf("%d", p))
		} else {
			parts = append(parts, fmt.Sprintf("%d^%d", p, factors.Exponent(p)))
		}
	}
	return strings.Join(parts, " × ")
}

func isPrime(p uint) bool {
	return p > 1 && PrimeFactorize(p).Exponent(p) == 1
}