# Product of two generating functions
#
# Copyright (C) 2021--2023 Simon Dobson
#
# This file is part of epydemic, epidemic network simulations in Python.
#
# epydemic is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# epydemic 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 epydemic. If not, see <http://www.gnu.org/licenses/gpl.html>.

from itertools import combinations_with_replacement
import epydemic.gf
from epydemic.gf import GF


class ProductGF(GF):
    '''The product of two underlying generating functions.

    This class computes the derivatives etc of a product of two
    other generating functions, generally generated by multiplying
    two functions together.
    '''

    def __init__(self, gf1: GF, gf2: GF):
        super().__init__()
        self._gf1 = gf1
        self._gf2 = gf2


    # ---------- Sub-class interface ----------

    def getCoefficient(self, i: int) -> float:
        '''Extract the coefficient of i.

        To do this we extract all the numbers in the range
        (0, i), determine which of them sum to i, and add
        the coefficients of those sets that do.

        :param i: the index
        :returns: the coefficient'''

        # compute the pairs of indices
        # compute the pairs that add up to i
        forwards = [ns for ns in combinations_with_replacement(range(i + 1), 2) if sum(ns) == i]
        # reverse the pairs that aren't the same value
        backwards = [(j, i) for (i, j) in forwards if i != j]
        # concatenate the two lists
        ns = forwards + backwards

        # compute the coefficient
        c = 0
        for (i, j) in ns:
            c += self._gf1[i] * self._gf2[j]

        return c

    def evaluate(self, x: float) -> float:
        '''Evaluate the product at x. This simply multiplies the
        results of evaluating the two factors individually.

        :param x: the argument
        :returns: the value of the product of generating functions'''
        return self._gf1(x) * self._gf2(x)

    def derivative(self, order: int = 1) -> GF:
        '''Return a new generating function representing the derivative.
        This applies the product rule of calculus to generate a product
        of the derivatives.

        :param order: (optional) the order of derivative (defaults to 1)
        :returns: the derivative of the generating function'''
        if order >= 1:
            d = epydemic.gf.SumGF(ProductGF(self._gf1.dx(), self._gf2),
                                  ProductGF(self._gf1, self._gf2.dx()))
            return d.dx(order - 1)
        else:
            return self

    def scale(self, n: float) -> GF:
        '''Multiply the product by a constant. The returns of a product
        of the first factor scaled by the constant with the second. (The
        other way round would work too, of course.)

        :param n: the number
        :returns: the new generating function'''
        return ProductGF(self._gf1 * n, self._gf2)

    def product(self, gf: GF) -> GF:
        '''Multiply the generating function by another. This creates
        a nested product.

        :param gf the other generating function
        :returns: the new generating function'''
        return ProductGF(self, gf)

    def sum(self, gf: GF) -> GF:
        '''Multiply the generating function by another. This creates
        a nested sum.

        :param gf: the other generating function
        :returns: the new generating function'''
        return epydemic.gf.SumGF(self, gf)
