#!/usr/bin/env python
"""
factor.py -- Give the prime factorization of a number.

This, and the primes module, both need to be fixed to
withstand having ^C hit.

Copyright (c) 2020 Steve Witham ess_double_you_at_tiac_notthis_dot_net
BSD license https://opensource.org/licenses/BSD-2-Clause
"""

from __future__ import print_function
from numbers import Integral as _Integral
from sys import stdout as _stdout

import primes
import sys

def factor_init():
    """
    Reset the prime number information kept by the factor module.
    Recovers from the way gen_factors_powers() and prime_gen()
    respond to ^C.
    """
    global _pg, _plist, _pset
    
    _pg = primes.prime_gen() # Start a new generator if one was in progress.
    _plist = [next(_pg)]
    _pset = set(_plist)
    
factor_init()


def _gen_a_prime():
    """
    Add a prime to the table used by gen_factors_powers().
    """
    p = next(_pg)
    _plist.append(p)
    _pset.add(p)

_gen_a_prime()


def get_prime(j):
    """ 
    Get the jth prime where the zeroth prime is 2.
    Just generates primes and fills its internal list
    until the jth prime is there.
    """
    while j >= len(_plist):
        _gen_a_prime()
    return _plist[j]

def is_prime(n):
    for (p, power) in gen_factors_powers(n):
        assert p <= n
        return (p == n)

def gen_factors_powers(n):
    """ 
    Look for power-of-prime factors of n starting at 2.
    For each prime p where n is divisible by p**power, yield (p, power).
    """
    i = 0
    while n > 1:
        while i >= len(_plist):
            _gen_a_prime()
        if n > _plist[-1]:
            # Generate primes to increase the chance of ending the factor-
            # trying process by catching a big prime.
            #
            # This extends the list of primes
            # at the same rate as we're testing n for divisibility
            # by primes starting at the beginning.  At each step it tests
            # whether n *is* the new prime, but not whether the new prime
            # divides n.  (Which might be good but would be complicated because
            # we want to generate the factors in order.)
            #
            # The divisibility tester is always behind the list-extender,
            # but eventually catches up to anything the list-extender
            # adds.  n is a prime when the tester finds a prime > sqrt(n).
            #
            # If someone is (e.g.) testing ascending series of numbers
            # for primality, this prevents the process from becoming less
            # efficient than the prime generator.  Not sure if there's a
            # time/time/space tradeoff that's reasonable for all
            # applications.
            _gen_a_prime()
        if n in _pset:
            yield (n, 1)
            return

        p = _plist[i]
        if p * p > n:
            _pset.add(n)
            yield (n, 1)
            return

        power = 0
        while n % p == 0:
            power += 1
            n //= p
        if power > 0:
            yield (p, power)

        i += 1


def factors(n):
    """
    Given n, an int to factor,
    return a list of zero or more (prime, power) pairs.
    """
    return list(gen_factors_powers(n))

        
def count_divisors(n_or_bps):
    """
    Given n_or_bps -- either an int >= 1,
                      or a list or generator of (prime, power) 
                      factors of a number n,
    return the number of divisors of n, including 1 and itself.
    """
    if isinstance(n_or_bps, _Integral):
        assert n_or_bps >= 1
        bps = gen_factors_powers(n_or_bps)
    else:
        bps = n_or_bps
    
    n_divisors = 1
    for (p, power) in bps:
        n_divisors *= (power + 1)    # A choice of p**0, p**1...p**power.
    return n_divisors


def totient(n_or_bps):
    """
    Given n_or_bps -- either an int >= 1,
                      or a list or generator of (prime, power) 
                      factors of a number n,
    return the Euler totient: how many k, 1 <= k <= n, for which
        gcd(n, k) == 1.
    """
    if isinstance(n_or_bps, _Integral):
        assert n_or_bps >= 1
        bps = gen_factors_powers(n_or_bps)
    else:
        bps = n_or_bps
    
    result = 1  # Used to start with n; now we have to reconstruct
                # n from its factors.
    for (p, power) in bps:
        result *= p ** (power - 1) * (p - 1)
    return result


def print_factors(n_or_bps, as_powers=True, sep="\n", one="1", end="\n",
                      file=None, flush=False):
    """
    Given
        n_or_bps -- either an int to factor 
                    or a list or iterator of (base, power) pairs;
        as_powers -- if True and power > 1, use base**power format,
                     if True and power == 1, us just base,
                     if False, use base repeated power times,
                     values of as_powers > 1 are reserved for future use;
        sep -- what to print between successive base or base**power entries
               (default is newline);
        flush -- force flush after every (base, power) pair (default False);
        file -- file object or file-like object to output to, 
                (default sys.stdout);
        one -- what to print if no factors (default "1");
        end -- what to print when finally done (default is newline);
    format and output the factoring of n, or other (base, power) pairs.
    If there were no factors, output the "one" parameter.
    Return nothing.
    """
    assert 0 <= as_powers <= 1, "Illegal value of as_powers (%d)." % as_powers

    if isinstance(n_or_bps, _Integral):
        bps = gen_factors_powers(n_or_bps)
    else:
        bps = n_or_bps

    file = file or _stdout

    prefix = ""
    for bp in bps:
        print(prefix + format_bp(bp, as_powers, sep), end="", file=file)
        prefix = sep
        if flush:
            file.flush()
    if prefix == ""  and one: # No factors were printed.
        print(one, end="", file=file)
    print(end=end, file=file, flush=flush)


def format_bp(bp, as_powers, sep):
    """
    Given
        bp -- a (base, power) pair,
        as_powers -- if True and power > 1, use base**power format,
                     if True and power == 1, use just base,
                     if False, use base sep base sep ...base,
                     values of as_powers > 1 are reserved for future use;
        sep -- what to insert between multiple base entries;
    return a single "base**pwr" or "base sep base sep... base" string.
    """
    base, power = bp
    if as_powers and power > 1:
        return "%d**%d" % (base, power)
    else:
        return sep.join(str(base) for i in range(power))

        
def factors_str(n_or_bps, as_powers=True, sep=" * ", one="1", end=""):
    """
    Given
        n_or_bps -- either an int to factor 
                    or a list or iterator of (base, power) pairs;
        as_powers -- if True and power > 1, use base**power format,
                     if True and power == 1, use just base,
                     if False, use base repeated power times,
                     values of as_powers > 1 are reserved for future use;
        sep -- what to insert between successive base or base**power entries
               (default is " * ");
        one -- string to return if no factors (default "1");
        end -- string to add to the end (default none);
    format and return the factoring of n, or other (base, power) pairs.
    If there were no factors, return the "one" parameter.
    """
    assert 0 <= as_powers <= 1, "Illegal value of as_powers (%d)." % as_powers

    if isinstance(n_or_bps, _Integral):
        bps = gen_factors_powers(n_or_bps)
    else:
        bps = n_or_bps

    parts = [format_bp(bp, as_powers, sep) for bp in bps]
    if parts:
        return sep.join(parts) + end
    else:
        return one + end


def _test_print_factors():
    from io import StringIO as _StringIO
    
    def print_factors_to_str(n_or_bps, as_powers=True, sep=" * ", one="1", end=""):
        out = _StringIO()
        print_factors(n_or_bps, as_powers=as_powers, sep=sep, one=one, end=end,
                            file=out)
        contents = out.getvalue()
        out.close()
        return contents

    for f in factors_str, print_factors_to_str:
        for pre_gen in False, True:
            argses = [ # Wow, a CSV constant.
                ("n_or_bps", "as_powers", "sep",  "one",  "end", "correct"),
                (1,                False,   "*",    "1",   "\n", "1\n"),
                (1,                False,   "*",    "1",     "", "1"),
                (12,               False,   "*",    "1",   "\n", "2*2*3\n"),
                (12,               False,   "*",    "1",     "", "2*2*3"),
                (1,                True,   "*",    "1",   "\n", "1\n"),
                (1,                True,   "*",    "1",     "", "1"),
                (12,               True,   "*",    "1",   "\n", "2**2*3\n"),
                (12,               True,   "*",    "1",     "", "2**2*3"),
                (1,                False,   "x",    "1",   "\n", "1\n"),
                (1,                False,   "x",    "1",     "", "1"),
                (12,               False,   "x",    "1",   "\n", "2x2x3\n"),
                (12,               False,   "x",    "1",     "", "2x2x3"),
                (1,                True,   "x",    "1",   "\n", "1\n"),
                (1,                True,   "x",    "1",     "", "1"),
                (12,               True,   "x",    "1",   "\n", "2**2x3\n"),
                (12,               True,   "x",    "1",     "", "2**2x3"),
             ]
            n_errors = 0
            for row in argses[1:]:
                pairs = zip(argses[0], row)
                kwargs = dict(pairs)
                result = kwargs["correct"]
                del kwargs["correct"]
                if pre_gen:
                    n = kwargs["n_or_bps"]
                    kwargs["n_or_bps"] = gen_factors_powers(n)
                x = f(**kwargs)
                if x != result:
                    print(kwargs)
                    print(f.__name__, repr(x))
                    print("correct", repr(result))
                    print()
                    n_errors += 1
            if n_errors > 0:
                print(n_errors, "errors.", file=_stderr)
                return False
            else:
                return True


def main(argv):
    r"""
    factor.py [-h | -p | -l | -t] [-s sep] number
    Print the factors of number.
        -h -? --help -- Print this message and exit.
        -p --powers  -- Default -- print entries as prime**power
                        or just prime if power==1.
        -l --lines   -- Opposite of --powers --
                        print prime repeated one entry per power.
        -s --sep sep -- Put sep between entries.  Default is '\n'.
                        \n, \t, \x07 etc. character escapes work.
        -t --test    -- Test print_factors() and factors_str().
    """
    # r""" ... """ to let the backslashes through as-is in docstring.
    from sys import stderr as _stderr, exit as _exit
    from char_esc import decode_char_escapes as _decode_char_escapes
    
    def usage(*print_args):
        if print_args:
            print("**** ", *print_args, file=_stderr)
            print()
        doc = main.__doc__.replace("factor.py", argv[0], 1)
        print(doc, file=_stderr)
        _exit(1)

    args = argv[1:]
    as_powers = True
    sep = "\n"
    while args and args[0].startswith("-"):
        opt = args.pop(0)
        if opt in ["-?", "-h", "--help"]:
            usage()
        elif opt in ["-p", "--powers"]:
            as_powers = True
        elif opt in ["-l", "--lines"]:
            as_powers = False
        elif opt in ["-s", "--sep"]:
            if args:
                sep = _decode_char_escapes(args.pop(0))
            else:
                usage("Missing separator after --sep.")
        elif opt in ["-t", "--test"]:
            result = _test_print_factors()
            if result == True:
                print("Test passed.", file=_stderr)
                _exit(0)
            else:
                print("Test failed.", file=_stderr)
                _exit(1)
        else:
            usage("Unrecognized option (?):", opt)
            
    if len(args) == 1:
        print_factors(int(args[0]), sep=sep, as_powers=as_powers, flush=True)
    elif len(args) > 1:
        usage("Too many arguments:", args)
    else:
        usage("Missing the number to factor.")

        
if __name__ == "__main__":
    from sys import argv

    main(argv)