import logging
from sage.all import sage_eval, PolynomialRing, QQ, SR, log

class PFOperator:
    """
    A class representing a Picard-Fuchs operator in a single variable z.
    Independent of the given order, the variables are assumed to be left of
    the derivatives.
    """
    def _operator_from_string(self, operator_string: str) -> "sage.rings.polynomial.polynomial_ring.Polynomial":
        z_names = ['z%d' % i for i in range(self.no_variables)]
        theta_names = ['theta%d' % i for i in range(self.no_variables)]
        self.ring = PolynomialRing(QQ, z_names + theta_names)
        self.z_gens = self.ring.gens()[:self.no_variables]
        self.theta_gens = self.ring.gens()[self.no_variables:]

        try:
            operator = sage_eval(operator_string, locals=self.ring.gens_dict())
        except Exception as e:
            raise ValueError("Invalid operator string: %s" % e)
        return self.ring(operator)

    def __init__(self, operator_string: str, no_variables: int = 1):
        self.no_variables = no_variables
        self.operator = self._operator_from_string(operator_string)
        logging.info("Initialised PFOperator: %s", self.operator)

class Period:
    """
    A class representing a period as a formal power series in z-variables and their logs.

    The coefficients are stored in a dictionary of dictionaries, where the first key is the multi-index
    of the logarithmic part and the second key is the multi-index of the z-variables. 
    So, for example, the coefficient of

        log(z0)^2 * log(z1) * z0^3 * z1^2 would be stored as coefficients[(2, 1)][(3, 2)]

    """

    def _initialise_ring(self):
        z_names = ['z%d' % i for i in range(self.no_variables)]
        log_names = ['L%d' % i for i in range(self.no_variables)]
        ring = PolynomialRing(QQ, z_names + log_names)
        z_gens = ring.gens()[:self.no_variables]
        log_gens = ring.gens()[self.no_variables:]

        return ring, z_gens, log_gens  # type: (PolynomialRing, list, list)

    def _period_from_string(self, period_string: str) -> dict:
        ring, z_gens, log_gens = self._initialise_ring()
        log_of_z = dict(zip(z_gens, log_gens))

        def log(zi):
            try:
                return log_of_z[zi]
            except (KeyError, TypeError):
                raise ValueError(
                    "log(...) may only be applied to one of the z-variables z0, ..., z%d"
                    % (self.no_variables - 1)
                )

        parse_locals = dict(ring.gens_dict())
        parse_locals['log'] = log

        try:
            expression = ring(sage_eval(period_string, locals=parse_locals))
        except Exception as e:
            raise ValueError("Invalid period string: %s" % e)

        coefficients = {}
        for coeff, monomial in expression:
            exponents = monomial.exponents()[0]
            z_index = tuple(exponents[:self.no_variables])
            log_index = tuple(exponents[self.no_variables:])
            coefficients.setdefault(log_index, {})[z_index] = coeff

        return coefficients
    
    def _period_to_string(self) -> str:
        ring, z_gens, log_gens = self._initialise_ring()
        expression = ring(0)
        for log_index, z_dict in self.coefficients.items():
            for z_index, coeff in z_dict.items():
                monomial = coeff
                for i in range(self.no_variables):
                    monomial *= (log_gens[i] ** log_index[i]) * (z_gens[i] ** z_index[i])
                expression += monomial

        log_substitutions = {log_gens[i]: log(SR(z_gens[i])) for i in range(self.no_variables)}
        return str(SR(expression).subs(log_substitutions))

    def _max_z_degree(self) -> int:
        # Highest total z-degree (sum of the z-multi-index) among all coefficients.
        if not self.coefficients:
            return 0
        return max(
            sum(z_index)
            for z_dict in self.coefficients.values()
            for z_index in z_dict
        )

    def _truncate_coefficients(self, order: int) -> dict:
        # Drop every (log_index, z_index) entry whose total z-degree exceeds order,
        # removing it from the dictionary rather than merely zeroing it out.
        truncated = {}
        for log_index, z_dict in self.coefficients.items():
            kept = {z_index: coeff for z_index, coeff in z_dict.items() if sum(z_index) <= order}
            if kept:
                truncated[log_index] = kept
        return truncated

    def _apply_theta(self, coefficients: dict, index: int) -> dict:
        """
        Applies the logarithmic derivative theta_i = z_i * d/dz_i once to a coefficients dict of the same
        shape as self.coefficients. It uses the product rule

            theta_i(z^a log(z)^k) = a_i * z^a log(z)^k + k_i * z^a log(z)^(k - e_i).
        """
        result = {}
        for log_index, z_dict in coefficients.items():
            for z_index, coeff in z_dict.items():
                a_i = z_index[index]
                if a_i != 0:
                    inner = result.setdefault(log_index, {})
                    inner[z_index] = inner.get(z_index, 0) + a_i * coeff

                k_i = log_index[index]
                if k_i != 0:
                    lowered_log_index = log_index[:index] + (k_i - 1,) + log_index[index + 1:]
                    inner = result.setdefault(lowered_log_index, {})
                    inner[z_index] = inner.get(z_index, 0) + k_i * coeff
        return result

    def apply_operator(self, pf_operator: PFOperator) -> "Period":
        """
        Applies a PFOperator to this period and returns the result as a new Period,
        truncated to self.order. Each monomial of the operator is normalised as
        coeff * z^p * theta^q (see PFOperator's docstring), so theta^q is applied
        to the period first and the result is then multiplied by coeff * z^p.
        """
        if pf_operator.no_variables != self.no_variables:
            raise ValueError(
                "Variable count mismatch: period has %d variable(s), operator has %d."
                % (self.no_variables, pf_operator.no_variables)
            )

        result_coefficients = {}
        for monomial_coeff, monomial in pf_operator.operator:
            exponents = monomial.exponents()[0]
            z_exponents = exponents[:self.no_variables]
            theta_exponents = exponents[self.no_variables:]

            term = self.coefficients
            for i, power in enumerate(theta_exponents):
                for _ in range(power):
                    term = self._apply_theta(term, i)

            for log_index, z_dict in term.items():
                inner = result_coefficients.setdefault(log_index, {})
                for z_index, coeff in z_dict.items():
                    shifted_z_index = tuple(z_index[i] + z_exponents[i] for i in range(self.no_variables))
                    inner[shifted_z_index] = inner.get(shifted_z_index, 0) + monomial_coeff * coeff

        result_coefficients = {
            log_index: {z_index: c for z_index, c in z_dict.items() if c != 0}
            for log_index, z_dict in result_coefficients.items()
        }
        result_coefficients = {log_index: z_dict for log_index, z_dict in result_coefficients.items() if z_dict}

        return Period(
            no_variables=self.no_variables,
            coefficients=result_coefficients,
            order=self.order,
        )

    def __init__(self, no_variables: int = 1, coefficients: dict = None, period_string: str = None, order: int = None):
        self.no_variables = no_variables

        # Use coefficients or period_string to initialize the period
        if period_string is not None:
            if coefficients is not None:
                raise ValueError("Provide either coefficients or period_string, not both.")
            self.coefficients = self._period_from_string(period_string)
            logging.debug("Using string for initialisation.")
        elif coefficients is None:
            self.coefficients = {}
        else:
            self.coefficients = coefficients
            self.period_string = self._period_to_string()
            logging.debug("Using coefficients for initialisation.")

        if order is None:
            self.order = self._max_z_degree()
            logging.debug("Order not provided, using maximum z-degree: %d", self.order)
        else:
            self.coefficients = self._truncate_coefficients(order)
            self.order = order
            self.period_string = self._period_to_string()
            logging.debug("Truncated coefficients to order %d.", self.order)

        logging.info("Initialised Period in %d variables at order %d.", self.no_variables, self.order)