from fractions import Fraction

import numpy as np
import numpy.typing as npt

from .balancing_algos import BalancingAlgorithms
from .chem_errors import BalancingError
from .utils import find_gcd, find_lcm


class Balancer(BalancingAlgorithms):
    """
    A class for balancing chemical equations automatically by different matrix methods.

    Parameters:
        matrix (npt.NDArray[np.float64]): Reaction matrix
        separator_pos (int): Position of the reaction separator (usually the separator is "=")
        round_precision (int): Coefficients rounding precision
        intify (bool): Determines whether the coefficients should be integers

    Attributes:
        coef_limit (int): max integer coefficient for \
        [_intify_coefficients][chemsynthcalc.balancer.Balancer._intify_coefficients] method
    """

    def __init__(
        self,
        matrix: npt.NDArray[np.float64],
        separator_pos: int,
        round_precision: int,
        intify: bool = True,
    ) -> None:
        super().__init__(matrix, separator_pos)

        if round_precision > 0:
            self.round_precision: int = round_precision
        else:
            raise ValueError("precision <= 0")

        self.intify: bool = intify
        self.coef_limit: int = 1_000_000

    def __str__(self) -> str:
        return f"Balancer object for matrix \n {self.reaction_matrix}"

    def __repr__(self) -> str:
        return f"Balancer({self.reaction_matrix}, {self.separator_pos}, {self.round_precision}, {self.intify})"

    def _intify_coefficients(
        self, coefficients: list[float], limit: int
    ) -> list[float | int] | list[int]:
        """
        Reduce the coefficients to integers by finding the greatest common divider.

        Parameters:
            coefficients (list): List of coefficients to intify
            limit (int): Upper limit (max int coef)

        Returns:
            A list of intified coefficients
        """
        initial_coefficients = coefficients
        frac = [Fraction(x).limit_denominator() for x in coefficients]
        vals = [
            int(
                fr.numerator
                * find_lcm([fr.denominator for fr in frac])
                / fr.denominator
            )
            for fr in frac
        ]
        coefficients = [int(val / find_gcd(vals)) for val in vals]
        if any(x > limit for x in coefficients):
            return initial_coefficients
        return coefficients

    @staticmethod
    def is_reaction_balanced(
        reactant_matrix: npt.NDArray[np.float64],
        product_matrix: npt.NDArray[np.float64],
        coefficients: list[float] | list[int],
        tolerance: float = 1e-8,
    ) -> bool:
        """
        Checks if reaction is balanced by multiplying reactant matrix and product matrix
        by the respective coefficient vector. Method is static to call it outside of balancer
        instance.

        Parameters:
            reactant_matrix (npt.NDArray[np.float64]): Matrix of reactants property generated by [ChemicalReaction][chemsynthcalc.chemical_reaction.ChemicalReaction] class
            product_matrix (npt.NDArray[np.float64]): Matrix of products property generated by [ChemicalReaction][chemsynthcalc.chemical_reaction.ChemicalReaction] class
            coefficients (list[float] | list[int]): Coefficients
            tolerance (float): tolerance limit for the *np.allclose* function

        Returns:
            True if balanced within tolerance

        Examples:
            >>> reaction = ChemicalReaction("NH4ClO4+HNO3+HCl=HClO4+NOCl+N2O+N2O3+H2O+Cl2")
            >>> Balancer.is_reaction_balanced(reaction.reactant_matrix, reaction.product_matrix, [64, 167, 137, 80, 43, 64, 30, 240, 39])
            True
            >>> reaction = ChemicalReaction("H2+O2=H2O")
            >>> Balancer.is_reaction_balanced(reaction.reactant_matrix, reaction.product_matrix, [2,2,2])
            False
        """
        try:
            reactants = np.multiply(
                reactant_matrix.T,
                np.array(coefficients)[: reactant_matrix.shape[1], None],
            )
            products = np.multiply(
                product_matrix.T,
                np.array(coefficients)[reactant_matrix.shape[1] :, None],
            )
            return np.allclose(
                reactants.sum(axis=0), products.sum(axis=0), rtol=tolerance
            )

        except Exception:
            return False

    def _calculate_by_method(self, method: str) -> list[float | int] | list[int]:
        """
        Compute the coefficients list by a specific method.

        Parameters:
            method (str): One of 4 currently implemented methods (inv, gpinv, ppinv, comb)

        Returns:
            A list of coefficients

        Raise:
            ValueError if method is not found. <br />
            [BalancingError][chemsynthcalc.chem_errors.BalancingError] if can't balance reaction by specified method.
        """
        match method:
            case "inv":
                coefficients: list[float] = np.round(
                    self._inv_algorithm(), decimals=self.round_precision
                ).tolist()  # type: ignore

            case "gpinv":
                coefficients: list[float] = np.round(
                    self._gpinv_algorithm(), decimals=self.round_precision + 2
                ).tolist()  # type: ignore

            case "ppinv":
                coefficients: list[float] = np.round(
                    self._ppinv_algorithm(), decimals=self.round_precision + 2
                ).tolist()  # type: ignore

            case "comb":
                res: npt.NDArray[np.int32] | None = self._comb_algorithm()
                if res is not None:
                    return res.tolist()  # type: ignore
                else:
                    raise BalancingError(f"Can't balance reaction by {method} method")

            case _:
                raise ValueError(f"No method {method}")

        if (
            Balancer.is_reaction_balanced(
                self.reactant_matrix, self.product_matrix, coefficients
            )
            and all(x > 0 for x in coefficients)
            and len(coefficients) == self.reaction_matrix.shape[1]
        ):
            if self.intify:
                intified = self._intify_coefficients(coefficients, self.coef_limit)
                if all(x < self.coef_limit for x in intified):
                    return intified
                else:
                    return coefficients
            else:
                return coefficients
        else:
            raise BalancingError(f"Can't balance reaction by {method} method")

    def inv(self) -> list[float | int] | list[int]:
        """
        A high-level function call to compute coefficients by Thorne method.

        Returns:
            A list of coefficients
        """
        return self._calculate_by_method("inv")

    def gpinv(self) -> list[float | int] | list[int]:
        """
        A high-level function call to compute coefficients by
        Risteski general pseudoinverse method.

        Returns:
            A list of coefficients
        """
        return self._calculate_by_method("gpinv")

    def ppinv(self) -> list[float | int] | list[int]:
        """
        A high-level function call to compute coefficients by
        Risteski partial pseudoinverse method.

        Returns:
            A list of coefficients
        """
        return self._calculate_by_method("ppinv")

    def comb(self) -> list[float | int] | list[int]:
        """
        A high-level function call to compute coefficients by
        combinatorial method.

        Returns:
            A list of coefficients
        """
        return self._calculate_by_method("comb")

    def auto(self) -> tuple[list[float | int] | list[int], str]:
        """
        A high-level function call to automatically compute coefficients
        by sequentially calling inv, gpinv, ppinv methods.

        Returns:
            A list of coefficients

        Raise:
            [BalancingError][chemsynthcalc.chem_errors.BalancingError] if can't balance reaction by any method.
        """
        try:
            return (self.inv(), "inverse")
        except Exception:
            pass
        try:
            return (self.gpinv(), "general pseudoinverse")
        except Exception:
            pass
        try:
            return (self.ppinv(), "partial pseudoinverse")
        except Exception:
            raise BalancingError("Can't balance this reaction by any method")
