# Sum 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>.

import epydemic.gf
from epydemic.gf import GF


class SumGF(GF):
    '''The sum of two underlying generating functions.

    This class computes the derivatives etc of a sum of two
    other generating functions, generally generated by adding
    their corresponding coefficients 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.

        :param i: the index
        :returns: the coefficient'''
        return self._gf1[i] + self._gf2[i]

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

        :param x: the argumentr
        :returns: the value of the sum 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 creates a sum of the derivatives of the terms.

        :param order: (optional) the order of derivative (defaults to 1)
        :returns: the deerivative of the generating function'''
        return SumGF(self._gf1.dx(order), self._gf2.dx(order))

    def scale(self, n: float) -> GF:
        '''Multiply the product by a constant. This distributes
        the constant multiplication across the product.

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

    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 epydemic.gf.ProductGF(self, gf)

    def sum(self, gf: GF) -> GF:
        '''Form the sum of this generating function with another. This
        creates a nested sum.

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