# Copyright (C) 2018 by the RROMPy authors
#
# This file is part of RROMPy.
#
# RROMPy is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# RROMPy 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with RROMPy. If not, see <http://www.gnu.org/licenses/>.
#

import numpy as np

__all__ = ['squareResonances']

prime_v = []

def squareResonances(a:int, b:int, zero_terms : bool = True):
    spectrum = []
    a = max(int(np.floor(a)), 0)
    b = max(int(np.ceil(b)), 0)
    global prime_v
    if len(prime_v) == 0:
        prime_v = [2, 3]
    if a > prime_v[-1]:
        for i in range(prime_v[-1], a, 2):
            getLowestPrimeFactor(i)
    for i in range(a, b + 1):
        spectrum = spectrum + [i] * countSquareSums(i, zero_terms)
    return np.array(spectrum)

def getLowestPrimeFactor(n:int):
    global prime_v
    for x in prime_v:
        if n % x == 0:
            return x
    if prime_v[-1] < n:
        prime_v = prime_v + [n]
    return n

def primeFactorize(n:int):
    factors = []
    number = n
    while number > 1:
        factor = getLowestPrimeFactor(number)
        factors.append(factor)
        number = int(number / factor)
    if n < -1:
        factors[0] = -factors[0]
    return list(factors)

def countSquareSums(n:int, zero_terms : bool = True):
    if n < 2: return (n + 1) * zero_terms
    factors = primeFactorize(n)
    funique, fcounts = np.unique(factors, return_counts = True)
    count = 1
    for fac, con in zip(funique, fcounts): #using number theory magic
        if fac % 4 == 1:
            count = count * (con + 1)
        elif fac % 4 == 3 and con % 2 == 1:
            return 0
    return count + (2 * zero_terms - 1) * (int(n ** .5) ** 2 == n)