# -*- coding: utf-8 -*-
"""
Implementation of Series handling functions.
"""
from mathics.core.atoms import Integer, Integer0, Rational
from mathics.core.convert.expression import to_mathics_list
from mathics.core.element import BaseElement
from mathics.core.evaluation import Evaluation
from mathics.core.expression import Expression
from mathics.core.list import ListExpression
from mathics.core.rules import BasePattern
from mathics.core.symbols import Atom, Symbol, SymbolPlus, SymbolPower, SymbolTimes
from mathics.core.systemsymbols import (
    SymbolComplexInfinity,
    SymbolD,
    SymbolDirectedInfinity,
    SymbolFactorial,
    SymbolIndeterminate,
    SymbolInfinity,
)

IntegerMinusOne = Integer(-1)
SymbolSeries = Symbol("Series")
SymbolSeriesData = Symbol("SeriesData")


def same_monomial(expr, x, x0):
    """
    Checks if expr == (x-x0)
    """
    if x0.is_zero and expr.sameQ(x):
        return True
    if expr.get_head() is not SymbolPlus:
        return False
    y, y0 = expr.elements
    if y0.sameQ(x):
        if x0.sameQ(-y) or y.sameQ(-x0):
            return True
    if y.sameQ(x):
        if x.sameQ(-y) or y.sameQ(-x):
            return True
    return False


# def to_series_term(series, term, x, x0):
#     return None, None
#     coeff = None
#     power = None
#     reminder = None
#     if isinstance(term, Atom):
#         if term.sameQ(x):
#             coeff = Integer1
#             power = 1
#             reminder = -x0
#         else:
#             coeff = term
#             power = 0
#             reminder = None
#     else:
#         head, elements = term.head, term.elements
#         if head is SymbolPower:
#             base, exponent = elements
#             head_exponent = exponent.get_head()
#             if head_exponent is SymbolInteger:
#                 power = exponent.to_float()
#             elif head_exponent is SymbolRational:
#                 num, den = exponent.value.as_numer_denom()
#                 if den != 1 and den != series[-1]:
#                     return None, None
#                 power = exponent.to_float()
#             else:
#                 return None, None
#             if x.sameQ(base):
#                 reminder = -term + Expression(SymbolPower, x - x0, exponent)
#                 coeff = 1
#             elif same_monomial(base, x, x0):
#                 coeff = 1
#                 reminder = None
#             else:
#                 return None, None
#         elif head is SymbolTimes:
#             coeffs_free = []
#             coeffs_powers = []
#             coeffs_x = []
#             for element in elements:
#                 if x.sameQ(element):
#                     coeffs_x.append(x)
#                 elif isinstance(element, Atom):
#                     coeffs_free.append(element)
#                 elif element.get_head() is SymbolPower:
#                     coeffs_powers.append(element)
#                 else:
#                     return None, None
#             # All the factors are of the form x^k
#             if all(x.sameQ(element.elements[0]) for element in coeffs_powers):
#                 coeff = Expression(SymbolTimes, *coeffs_free)
#                 power = len(coeffs_x)
#                 exponents = [element.elements[1] for element in coeffs_powers]
#                 if not all(
#                     isinstance(exponent, (Integer, Rational)) for exponent in exponents
#                 ):
#                     return None, None
#                 power = power + sum([exp.to_python() for exp in exponents])
#                 reminder = -term + Expression(
#                     SymbolTimes, coeff, Expression(SymbolPower, x - x0, power)
#                 )
#             # All the factors are of the form (x-x0)^k
#             elif (
#                 all(same_monomial(element.elements[0], x, x0) for element in coeffs_powers)
#                 and len(coeffs_x) == 0
#             ):
#                 coeff = Expression(SymbolTimes, *coeffs_free)
#                 exponents = [element.elements[1] for element in coeffs_powers]
#                 if not all(
#                     isinstance(exponent, (Integer, Rational)) for exponent in exponents
#                 ):
#                     return None, None
#                 power = power + sum([exp.to_python() for exp in exponents])
#                 reminder = None
#             else:
#                 return None, None
#     if coeff:
#         power = power * series[-1]
#         if power > series[-2] or (power - int(power)) != 0:
#             return None, None
#         nmin = series[1]
#         power = int(power) - nmin
#         if power < 0:
#             nmin = nmin + power
#             newdata = [coef] + [Integer0 for i in range(1 - power)] + series[0].elements
#         else:
#             newdata = [
#                 c + coeff if p == power else c for p, c in enumerate(series[0].elements)
#             ]
#         return (
#             ListExpression(*newdata),
#             nmin,
#             series[-2],
#             series[-1],
#         ), reminder
#     return None, None


def series_plus_series(series1, series2):
    """
    Tries to reduce the sum of series1 and series2 on the same variable and neighbourhood
    to a single series. None if it is not possible.
    """
    # First, check that series1 has the smaller power.
    if series1[1] / series1[3] > series2[1] / series2[3]:
        series1, series2 = series2, series1

    data1, nmin1, nmax1, den1 = series1
    data2, nmin2, nmax2, den2 = series2
    data1 = data1.elements
    data2 = data2.elements

    den = den1 * den2
    nmin1_ = nmin1 * int(den1 / den)
    nmin2_ = nmin2 * int(den2 / den)
    nmin = nmin1_
    offset2 = nmin2_ - nmin1
    nmax1_ = nmax1 * int(den1 / den) - nmin
    nmax2_ = nmax2 * int(den2 / den) - nmin
    nmax = min(nmax1_, nmax2_)  # relative to nmin
    len_newdata = nmax
    nmax = nmax + nmin
    data = [Integer0 for x in range(len_newdata)]

    for k, coeff in enumerate(data1):
        p = k * int(den1 / den)
        if p < len_newdata:
            data[p] = coeff

    for k, coeff in enumerate(data2):
        p = k * int(den2 / den) + offset2
        if p < len_newdata:
            if data[p].is_zero:
                data[p] = coeff
            else:
                data[p] = Expression(SymbolPlus, data[p], coeff)

    data = ListExpression(*data)
    result = reduce_series_trailing_zeros((data, nmin, nmax, den))
    return result


def series_times_series(series1, series2):
    """
    Tries to reduce the product of series1 and series2 on the same variable and neighbourhood
    to a single series. None if it is not possible
    """
    data1, nmin1, nmax1, den1 = series1
    data2, nmin2, nmax2, den2 = series2
    data1 = data1.elements
    data2 = data2.elements
    # maybe we should use the MCD
    den = den1 * den2
    offset1 = int(den1 / den)
    offset2 = int(den2 / den)
    nmin = nmin1 * offset1 + nmin2 * offset2
    nmax = min(nmax1 * den2, nmax2 * den1)
    len_newdata = min(len(data1) * len(data2), nmax)
    data = [Integer0 for k in range(len_newdata)]

    for k1, c1 in enumerate(data1):
        for k2, c2 in enumerate(data2):
            pos = k1 * offset1 + k2 * offset2
            if pos >= len_newdata:
                continue
            if data[pos].is_zero:
                data[pos] = Expression(SymbolTimes, c1, c2)
            elif data[pos].get_head() is SymbolPlus:
                data[pos] = Expression(
                    SymbolPlus, Expression(SymbolTimes, c1, c2), *(data[pos].elements)
                )
            else:
                data[pos] = Expression(
                    SymbolPlus, Expression(SymbolTimes, c1, c2), data[pos]
                )

    data = ListExpression(*data)
    return reduce_series_trailing_zeros((data, nmin, nmax, den))


def _series_times_rational_power(series, num_power, den_power):
    """
    Tries to reduce the producto of a series with a power of the basic monomial
    to a single series. None if it is not possible
    """
    data, nmin, nmax, den = series
    data = data.elements
    # TODO: use the MCD
    den_ = den * den_power
    nmin_ = int(nmin * den_ / den + num_power * den_ / den_power)
    nmax_ = int(nmax * den_ / den + num_power * den_ / den_power)
    granularity = int(den_ / den_power)
    data_ = [Integer0 for u in range(len(data) * granularity)]
    for k, val in enumerate(data):
        data_[k * granularity] = val
    return reduce_series_trailing_zeros(
        (ListExpression(*data_), nmin_, nmax_, den * den_power)
    )


def reduce_series_trailing_zeros(series):
    """
    Reduce the series representation to start with a non-zero
    coefficient, and to remove coefficients of order larger that
    nmax
    """
    data, nmin, nmax, den = series
    data = data.elements
    if len(data) == 0:
        return series
    i = 0
    while i < len(data) and data[i].is_zero:
        i = i + 1
    nmin = nmin + i
    data = data[i:]
    useful_len = nmax - nmin
    if len(data) > useful_len:
        data = data[:useful_len]
    if len(data) == 0:
        nmin = nmax
    result = (ListExpression(*data), nmin, nmax, den)
    return result


def reduce_series(series):
    series = reduce_series_trailing_zeros(series)

    def factors(den):
        # TODO: extend the list of primes?
        for factor in [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]:
            if den < factor:
                break
            if den % factor == 0:
                yield (factor)

    def reduce_dataseries(series, factor):
        data, nmin, nmax, den = series
        data = data.elements
        notdone = True
        while notdone:
            if (den % factor == 0) and (nmin % factor == 0) and (nmax % factor == 0):
                if all(
                    q.is_zero for q in data[1 + factor :: 2] for r in range(factor - 1)
                ):
                    data = data[0::factor]
                    nmin, nmax, den = (
                        int(nmin / factor),
                        int(nmax / factor),
                        int(den / factor),
                    )
                    continue
            notdone = False

        if series[-1] == den:
            return series
        return (ListExpression(*data), nmin, nmax, den)

    for factor in factors(series[-1]):
        series = reduce_dataseries(series, factor)

    return series


def reduce_series_plus(series, terms, x, x0):
    """
    Tries to reduce the sum of a Series with elements in expr_list.
    """
    old_series, old_terms = series, terms
    other_terms = []
    if not terms:
        return series, None

    # Loop over terms
    for term in terms:
        if term.is_zero:
            continue
        if isinstance(term, Atom):
            other_terms.append(term)
            continue
        term_head = term.head
        term_elements = term.elements
        if term_head is SymbolSeriesData:
            y, y0, data, nummin, nummax, den = term_elements
            if not x.sameQ(y):
                data = ListExpression(term)
                y = x
                y0 = x0
                nummin = 0
                nummax = series[2]
                den = series[3]
            elif not x0.sameQ(y0):
                other_terms.append(term)
                continue
            new_series = series_plus_series(
                series,
                (
                    data,
                    nummin.get_int_value(),
                    nummax.get_int_value(),
                    den.get_int_value(),
                ),
            )
            if new_series:
                series = new_series
                continue
        # TODO: handle constant terms, and terms of the form x^n, a x^n
        other_terms.append(term)
    if len(other_terms) == 0:
        return series, None

    if series is old_series:
        if len(other_terms) == len(old_terms) and all(
            term.sameQ(old_term) for term, old_term in zip(other_terms, old_terms)
        ):
            return None, None
    series = reduce_series(series)
    return series, other_terms


def build_series(
    f: BaseElement, x: BaseElement, x0: BaseElement, n: Integer, evaluation: Evaluation
) -> BaseElement:
    """
    Builds the series expansion of f on x around x0, upto order n.
    """
    # TODO:
    # - Deal with special cases (monomials, fractional powers, singular cases)
    # - Use definitions of Series for builtin functions
    # - Asymptotic series
    x_name = x.get_name()
    vars = {
        x_name: x0,
    }
    x_pattern = BasePattern.create(x, evaluation=evaluation)

    if f.is_free(x_pattern, evaluation):
        # print(x, " not in ", f)
        return f

    data = [f.replace_vars(vars)]
    df = f
    for i in range(n.get_int_value()):
        df = Expression(SymbolD, df, x).evaluate(evaluation)
        newcoeff = df.replace_vars(vars).evaluate(evaluation)
        factorial = Expression(SymbolFactorial, Integer(i + 1))
        newcoeff = Expression(
            SymbolTimes,
            Expression(SymbolPower, factorial, IntegerMinusOne),
            newcoeff,
        ).evaluate(evaluation)
        if newcoeff in (
            SymbolInfinity,
            SymbolComplexInfinity,
            SymbolIndeterminate,
        ) or newcoeff.get_head() in (SymbolDirectedInfinity,):
            return Expression(
                f.get_head(),
                *[
                    build_series(element, x, x0, n, evaluation)
                    for element in f.elements
                ],
            )
        data.append(newcoeff)
    data = ListExpression(*data).evaluate(evaluation)
    series = reduce_series_trailing_zeros((data, 0, n.get_int_value() + 1, 1))
    return Expression(
        SymbolSeriesData,
        x,
        x0,
        series[0],
        Integer(series[1]),
        Integer(series[2]),
        Integer(series[3]),
    )


def series_derivative(series, x, x0, y, evaluation):
    """
    Evaluates the derivative of the series
    """
    data, nmin, nmax, den = series
    coeffs = list(data.elements)
    if all(
        not coeff.has_symbol(y.get_name())
        for coeff in coeffs
        if hasattr(coeff, "has_symbol")
    ):
        dcoeffs = None
    else:
        dcoeffs = [Expression(SymbolD, coeff, y) for coeff in coeffs]

    prefactor = Expression(SymbolD, x - x0, y).evaluate(evaluation)

    if Integer0.sameQ(prefactor):
        if dcoeffs:
            return reduce_series_trailing_zeros(
                (
                    to_mathics_list(*[coeff.evaluate(evaluation) for coeff in dcoeffs]),
                    nmin,
                    nmax,
                    den,
                )
            )
        else:
            return None
    nmax = nmax - den
    trailings0 = [Integer0 for k in range(den)]
    if nmin != 0:
        if den == 1:
            coeffs2 = [
                (Integer(k + nmin) * prefactor * coeff)
                for k, coeff in enumerate(coeffs)
            ] + trailings0
        else:
            coeffs2 = [
                (Rational(k + nmin, den) * prefactor * coeff)
                for k, coeff in enumerate(coeffs)
            ] + trailings0
        nmin = nmin - den
        if dcoeffs:
            dcoeffs = trailings0 + dcoeffs
    else:
        coeffs = coeffs[1:]
        if den == 1:
            coeffs2 = [
                (Integer(k + 1) * prefactor * coeff) for k, coeff in enumerate(coeffs)
            ] + trailings0
        else:
            coeffs2 = [
                (Rational(k + 1, den) * prefactor * coeff)
                for k, coeff in enumerate(coeffs)
            ] + trailings0

    if dcoeffs:
        new_coeffs = [
            (c1 + c2).evaluate(evaluation) for c1, c2 in zip(dcoeffs, coeffs2)
        ]
    else:
        new_coeffs = [c1.evaluate(evaluation) for c1 in coeffs2]

    result = reduce_series_trailing_zeros(
        (ListExpression(*new_coeffs), nmin, nmax, den)
    )
    return result
