"""Implement base visitor class for processing individual terms in expression"""

from abc import ABC, abstractmethod

from scipy.special import binom

from autowisp.fit_expression.FitTermsParser import FitTermsParser
from autowisp.fit_expression.FitTermsParserVisitor import FitTermsParserVisitor


class ProcessTermsVisitor(FitTermsParserVisitor, ABC):
    """
    Base visitor class for processing individual terms in fit terms expressions.

    Children should overwrite visitFit_term() to process individual terms, as
    well as implement the abstract methods.
    """

    @abstractmethod
    def _start_polynomial_expansion(self, num_output_terms, input_terms_list):
        """
        Called at the beginning of processing of a new polynomial expansion.

        Args:
            number_terms(int):    The number of terms the polynomial expansion
                will produce.

        Returns:
            None
        """

    @abstractmethod
    def _process_polynomial_term(self, input_terms, term_powers):
        """
        Called for each term in a polynomial expansion.

        Args:
            input_terms([str]):    The list of terms the polynomial expansion is
                being performed on.

            term_Powers([int]):    The powerlaw indices for each term
                in ``input_terms``.

        Returns:
            None
        """

    @abstractmethod
    def _end_polynomial_expansion(self):
        """
        Called at the end of processing of a new polynomial expansion.

        Args:
            None

        Returns:
            The terms produced by the polynomial expansion, as generated by the
            repeated calls to process_polynomial_term().
        """

    @abstractmethod
    def _start_cross_product_expansion(self, input_term_sets):
        """
        Called at the beginning of processing of a new cross product expansion.

        Args:
            input_term_sets([]):    A list of the sets of terms taking part in
                the cross product.

        Returns:
            None
        """

    @abstractmethod
    def _process_cross_product_term(self, sub_terms):
        """
        Called for each term in a cross product expansion.

        Args:
            sub_terms([]):    A list of the individual terms which must be
                multiplied together to form a single term in the resulting cross
                product expansion.

        Returns:
            None
        """

    @abstractmethod
    def _end_cross_product_expansion(self):
        """
        Called at the end of processing of a new cross product expansion.

        Args:
            None

        Returns:
            The terms produced by the cross producc expansion, as generated by
            the repeated calls to process_cross_product_term().
        """

    # Visit a parse tree produced by FitTermsParser#fit_terms_list.
    def visitFit_terms_list(self, ctx: FitTermsParser.Fit_terms_listContext):
        """Return a list of the terms in the list."""

        return [self.visit(term) for term in ctx.fit_term()]

    # Visit a parse tree produced by FitTermsParser#fit_terms_set.
    def visitFit_terms_set(self, ctx: FitTermsParser.Fit_terms_setContext):
        """Return a list of the terms in the set."""

        return self.visit(ctx.getChild(0))

    # Visit a parse tree produced by FitTermsParser#fit_polynomial.
    def visitFit_polynomial(self, ctx: FitTermsParser.Fit_polynomialContext):
        """Process the list of terms the polynomial expression expands to."""

        def next_power_set(powerlaw_indices):
            """
            Advance to the next term powerlaw indices preserting total order.

            Args:
                powerlaw_indices:    The powerlaw indices of the current term.
                    One for each term in the list from which the polynomial is
                    buing built.

            Returns:
                bool:
                    False iff the input powerlaw indices correspond to the last
                    term of the current maximum order.
            """

            for i in range(len(powerlaw_indices) - 1):
                if powerlaw_indices[i]:
                    powerlaw_indices[i + 1] += 1
                    powerlaw_indices[0] = powerlaw_indices[i] - 1
                    if i:
                        powerlaw_indices[i] = 0
                    return True
            return False

        max_order = int(ctx.order.text)
        terms_list = self.visit(ctx.fit_terms_list())
        self._start_polynomial_expansion(
            int(binom(len(terms_list) + max_order, max_order)), terms_list
        )
        term_powers = [0 for term in terms_list]
        for total_order in range(max_order):
            term_powers[-1] = 0
            term_powers[0] = total_order + 1
            self._process_polynomial_term(terms_list, term_powers)
            while next_power_set(term_powers):
                self._process_polynomial_term(terms_list, term_powers)
        return self._end_polynomial_expansion()

    # Visit a parse tree produced by FitTermsParser#fit_terms_set_cross_product.
    def visitFit_terms_set_cross_product(
        self, ctx: FitTermsParser.Fit_terms_set_cross_productContext
    ):
        """Return all possible terms combining one term from each input set."""

        term_sets = [self.visit(tset) for tset in ctx.fit_terms_set()]
        self._start_cross_product_expansion(term_sets)
        term_indices = [0 for s in term_sets]
        more_terms = True
        while more_terms:
            self._process_cross_product_term(
                [
                    term_sets[set_ind][term_indices[set_ind]]
                    for set_ind in range(len(term_sets))
                ]
            )
            more_terms = False
            for set_ind, term_ind in enumerate(term_indices):
                if term_ind < len(term_sets[set_ind]) - 1:
                    term_indices[set_ind] += 1
                    for i in range(0, set_ind):
                        term_indices[i] = 0
                    more_terms = True
                    break
        return self._end_cross_product_expansion()
