package factors

import (
	"fmt"
	"slices"
	"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 // number with all found factors removed (divided)
	exponents := make(map[uint]uint)

	// try each factor starting at 2
	// if unfactored is divisible divide out until it isn't
	for f := uint(2); f*f <= unfactored; {
		if unfactored%f == 0 {
			unfactored /= f
			exponents[f] += 1
		} else {
			f += 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 {
	parts := make([]string, 0, factors.Size())
	primes := make([]int, 0, factors.Size())
	for p := range factors.exponents {
		primes = append(primes, int(p))
	}
	slices.Sort(primes)

	for _, p := range primes {
		if factors.Exponent(uint(p)) == 1 {
			parts = append(parts, fmt.Sprintf("%d", p))
		} else {
			parts = append(parts, fmt.Sprintf("%d^%d", p,
				factors.Exponent(uint(p))))
		}
	}
	return strings.Join(parts, " × ")
}