computor

#!/usr/bin/env python3
import re
import sys


def sqrt(n):
    """Manual square root calculation."""
    if n < 0:
        return None
    if n == 0:
        return 0
    
    x = n
    while True:
        root = 0.5 * (x + n / x)
        if abs(root - x) < 1e-10:
            return root
        x = root


def normalize_number(number):
    """Normalize number for clean output."""
    if number.is_integer():
        return int(number)
    return round(number, 6)


def parse_side(side):
    """
    Parses one side of the equation and returns a dict {degree: coefficient}.\n
    e.g., for "5 * X^2 + 3 * X^1", it returns {2: 5, 1: 3}
    """
    if side.strip() == "":
        raise ValueError("Equation side is empty.")

    # https://regexlearn.com/learn
    # Matching pattern example: '+ 5 * X^2'
    TERM_PATTERN = re.compile(r'([+-]?\s*\d*\.?\d*)\s*\*?\s*X\^(\d+)')

    # Find all valid terms
    terms = TERM_PATTERN.findall(side)

    # Remove valid terms from the side to check if there are any invalid elements left
    cleaned_side = TERM_PATTERN.sub('', side).strip()

    # Check for any invalid terms
    if cleaned_side and cleaned_side != "0":
        raise ValueError(f"Invalid term(s) in equation side: '{cleaned_side}'")

    # Build the coefficient dictionary
    coeffs = {}
    for coeff, degree in terms:
        coeff = coeff.replace(" ", "")
        coeff = float(coeff) if coeff not in ["", "+", "-"] else float(coeff + "1")
        degree = int(degree)
        coeffs[degree] = coeffs.get(degree, 0) + coeff
    # print(f"Parsed side '{side.strip()}' to coefficients: {coeffs}")
    return coeffs


def parse_equation(equation):
    """
    Parse the equation, moving all terms to the left side and returning
    a dict {degree: coefficient}.\n
    e.g., for "5 * X^2 + 3 * X^1 = 2 * X^0", it returns {2: 5, 1: 3, 0: -2}
    """
    left, right = equation.split('=')

    left_coeffs = parse_side(left)
    right_coeffs = parse_side(right)
    
    # Move all terms to the left side (subtract the right side)
    coeffs = {}
    for degree, coeff in left_coeffs.items():
        coeffs[degree] = coeffs.get(degree, 0) + coeff
        print(f"After adding left term: degree {degree}, coeff {normalize_number(coeff)}, equation now: {get_equation_string(coeffs)}")
    for degree, coeff in right_coeffs.items():
        coeffs[degree] = coeffs.get(degree, 0) - coeff
        print(f"After subtracting right term: degree {degree}, coeff {normalize_number(coeff)}, equation now: {get_equation_string(coeffs)}")
    return coeffs


def get_equation_string(coeffs):
    """
    Get the equation in standard form.\n
    e.g., for {2: 5, 1: 3, 0: -2}, it return:\n
    5 * X^2 + 3 * X^1 - 2 * X^0
    """
    terms = []
    for degree in sorted(coeffs.keys()):
        coeff = normalize_number(coeffs[degree])
        sign = "+ " if terms and coeff >= 0 else "- " if coeff < 0 else ""
        terms.append(f"{sign}{abs(coeff)} * X^{degree}")
    equation = ' '.join(terms)
    return equation


def get_degree(coeffs):
    """Returns the degree of the polynomial."""
    if not coeffs:
        return 0
    
    for degree in sorted(coeffs.keys(), reverse=True):
        if coeffs[degree] != 0:
            return degree
    
    return 0


def solve_degree_0(coeffs):
    """Solves degree 0 equations."""
    c = coeffs.get(0, 0)
    if c == 0:
        print("Any real number is a solution.")
    else:
        print("No solution.")


def solve_degree_1(coeffs):
    """Solves degree 1 equations: ax + b = 0."""
    a = coeffs.get(1, 0)
    b = coeffs.get(0, 0)
    
    if a == 0:
        solve_degree_0(coeffs)
    else:
        solution = -b / a
        print(f"The solution is:\n{normalize_number(solution)}")


def solve_degree_2(coeffs):
    """Solves a degree 2 equation: ax² + bx + c = 0."""
    a = coeffs.get(2, 0)
    b = coeffs.get(1, 0)
    c = coeffs.get(0, 0)
    
    if a == 0:
        solve_degree_1(coeffs)
        return
    
    # Discriminant calculation: Δ = b² - 4ac
    discriminant = b * b - 4 * a * c
    
    if discriminant > 0:
        print("Discriminant is strictly positive, the two solutions are:")
        sqrt_d = sqrt(discriminant)
        x1 = (-b - sqrt_d) / (2 * a)
        x2 = (-b + sqrt_d) / (2 * a)
        print(normalize_number(x1))
        print(normalize_number(x2))
    elif discriminant == 0:
        print("Discriminant is zero, the solution is:")
        x = -b / (2 * a)
        print(normalize_number(x))
    else:
        print("Discriminant is strictly negative, the two complex solutions are:")
        real_part = -b / (2 * a)
        imaginary_part = sqrt(-discriminant) / (2 * a)
        print(f"{normalize_number(real_part)} + {normalize_number(imaginary_part)}i")
        print(f"{normalize_number(real_part)} - {normalize_number(imaginary_part)}i")


def solve_equation(equation):
    """
    Solves the given polynomial equation.
    Manages degrees 0, 1, and 2.
    e.g., for "5 * X^2 + 3 * X^1 - 2 * X^0 = 0", it computes the solutions.
    """

    # Convert the equation into coefficients, moving all terms to the left side
    coeffs = parse_equation(equation)
    
    # Display the reduced form
    print("Reduced form: " + get_equation_string(coeffs) + " = 0")

    # Determine the degree of the polynomial
    degree = get_degree(coeffs)
    print(f"Polynomial degree: {degree}")

    if degree > 2:
        print("The polynomial degree is strictly greater than 2, I can't solve.")
        sys.exit(0)

    # Solve based on the degree
    if degree == 0:
        solve_degree_0(coeffs)
    elif degree == 1:
        solve_degree_1(coeffs)
    elif degree == 2:
        solve_degree_2(coeffs)


def main():
    try:
        if len(sys.argv) == 2:
            equation = sys.argv[1]
        elif len(sys.argv) == 1:
            print("Please enter the equation:")
            equation = input()
        else:
            print("Usage: ./computor \"5 * X^2 + 3 * X^1 - 2 * X^0 = 0\"")
            sys.exit(1)

        solve_equation(equation)

    except Exception as e:
        print(f"Verify your equation syntax (e.g., '5 * X^2 + 3 * X^1 - 2 * X^0 = 0').\nDetails: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()