import re
from typing import List, Dict, Tuple
from sympy import Matrix, lcm
from .parser import parse_chemical_formula

def parse_equation(equation: str) -> Tuple[List[str], List[str]]:
    """
    Parse the equation into reactants and products.
    """
    if '->' not in equation:
        raise ValueError("Equation must contain '->'")
    parts = equation.split('->')
    if len(parts) != 2:
        raise ValueError("Invalid equation format")
    reactants = [comp.strip() for comp in parts[0].split('+')]
    products = [comp.strip() for comp in parts[1].split('+')]
    return reactants, products

def balance_equation(equation: str) -> Tuple[str, List[int]]:
    """
    Balance the chemical equation and return balanced equation and coefficients.
    """
    reactants, products = parse_equation(equation)
    compounds = reactants + products
    num_reactants = len(reactants)
    num_products = len(products)
    total_compounds = len(compounds)

    # Get all unique elements
    all_elements = set()
    atom_counts = []
    for comp in compounds:
        counts = parse_chemical_formula(comp)
        atom_counts.append(counts)
        all_elements.update(counts.keys())

    all_elements = sorted(all_elements)
    num_elements = len(all_elements)

    # Build matrix: rows = elements, columns = compounds
    matrix = []
    for el in all_elements:
        row = []
        for i, counts in enumerate(atom_counts):
            coeff = 1 if i < num_reactants else -1
            row.append(coeff * counts.get(el, 0))
        matrix.append(row)

    A = Matrix(matrix)
    # Find null space
    null_space = A.nullspace()
    if not null_space:
        raise ValueError("Equation cannot be balanced")

    # Get the solution vector (coefficients for compounds)
    sol_vec = null_space[0]
    solution = [sol_vec[i, 0] for i in range(sol_vec.rows)]
    # Make all positive by multiplying by -1 if needed
    if any(x < 0 for x in solution):
        solution = [-x for x in solution]
    # Make integer by multiplying by common denominator
    from sympy import denom, lcm
    denoms = [denom(x) for x in solution]
    common_denom = lcm(denoms)
    coefficients = [int(common_denom * x) for x in solution]
    # Minimize by dividing by gcd
    from math import gcd
    from functools import reduce
    coeff_gcd = reduce(gcd, coefficients)
    coefficients = [c // coeff_gcd for c in coefficients]

    # Build balanced equation
    balanced_reactants = []
    for i, coeff in enumerate(coefficients[:num_reactants]):
        if coeff > 1:
            balanced_reactants.append(f"{coeff} {reactants[i]}")
        else:
            balanced_reactants.append(reactants[i])

    balanced_products = []
    for i, coeff in enumerate(coefficients[num_reactants:]):
        if coeff > 1:
            balanced_products.append(f"{coeff} {products[i]}")
        else:
            balanced_products.append(products[i])

    balanced_eq = " + ".join(balanced_reactants) + " -> " + " + ".join(balanced_products)

    return balanced_eq, coefficients