#All Rights Reserved John Salguero
#Steps that are generated

from sympy import *
from sympy.parsing.sympy_parser import (
    parse_expr,
    standard_transformations,
    implicit_multiplication_application
)
transformations = standard_transformations + (implicit_multiplication_application,)

def move_all_to_one_side(equation):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    new_expr = left_expr - right_expr
    step["after"] = f"{sstr(new_expr)} = 0"
    
    step["step"] = f"Subtract both sides by {sstr(right_expr)}"
    step["rule"] = "Subtraction Property of Equality"
    return step
   
def add_both_sides(equation, value):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = clean(left_expr + value)
    new_right_expr = clean(right_expr + value)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Add both sides by {sstr(value)}"
    step["rule"] = "Addition Property of Equality"
    return step
   
def subtract_both_sides(equation, value):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = clean(left_expr - value)
    new_right_expr = clean(right_expr - value)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Subtract both sides by {sstr(value)}"
    step["rule"] = "Subtraction Property of Equality"
    return step
    
def divide_both_sides(equation, value):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = clean(left_expr / value)
    new_right_expr = clean(right_expr / value)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Divide both sides by {sstr(value)}"
    step["rule"] = "Division Property of Equality"
    return step
    
def multiply_both_sides(equation, value):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = left_expr * value
    new_right_expr = right_expr * value
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Multiply both sides by {sstr(value)}"
    step["rule"] = "Multiplication Property of Equality"
    return step
    
def square_root_both_sides(equation):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x_generic = symbols('x')
    x_pos = symbols('x', positive=True)
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = sqrt(left_expr.subs(x_generic, x_pos))
    new_right_expr = sqrt(right_expr)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}, {sstr(new_left_expr)} = -{sstr(new_right_expr)}"
    
    step["step"] = f"Take the square root of both sides"
    step["rule"] = "Square Root Property of Equality"
    return step
    
def square_both_sides(equation):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = clean(left_expr * left_expr)
    new_right_expr = clean(right_expr * right_expr)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Square both sides"
    step["rule"] = "Multiplication property of equality"
    return step
  
def factor_collect(equation):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = factor(left_expr)
    new_right_expr = factor(right_expr)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Collect the factors"
    step["rule"] = "Factoring by grouping"
    return step
  
def factor_form_collection(equation, factor):
    # Collect factors of factor
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = collect(left_expr, factor)
    new_right_expr = collect(right_expr, factor)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    
    step["step"] = f"Collect the factors using factor {sstr(factor)}"
    step["rule"] = "Factor by grouping"
    return step
    
def factor_out(equation, factor):
    step = {}
   
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    new_left_expr = clean(cancel(left_expr / factor))
    new_left_expr = Mul(factor, new_left_expr, evaluate = False)
    step["after"] = f"{sstr(new_left_expr)} = {sstr(right_expr)}"
    
    step["step"] = f"Factor out the Greatest Common Factor, {sstr(factor)}"
    step["rule"] = "Reverse Distributive Property"
    return step
    
def trinomial_by_grouping(equation, inner):
    # expects n (ax**2+bx+c) = rhs : inner = (ax**2+bx+c), b != 0, c != 0
    
    #4 steps
    steps = [{}]
   
    current = equation
    steps[-1]["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    n = simplify(left_expr / inner)
    poly = inner.as_poly(x)
    a = poly.coeff_monomial(x**2)
    b = poly.coeff_monomial(x)
    c = poly.coeff_monomial(1)
    ac = Abs(a * c)
    
    ## Split Coeficients
    factor1 = Integer(1)
    factor2 = ac
    while factor1 < factor2:
        if ac % factor1 == 0:
            factor2 = ac / factor1
            if c.is_negative:
                if Abs(factor1 - factor2) == Abs(b):
                    break
            else:
                if factor1 + factor2 == Abs(b):
                    break
        factor1 = factor1 + 1
    
    if factor1 > factor2:
        return []
    
    if c.is_negative:
        action = "differ by"
        if b.is_negative:
            left_expr = Add(a * x**2, factor1 * x, -factor2 * x, c, evaluate=False)
        else:
            left_expr = Add(a * x**2, -factor1 * x, factor2 * x, c, evaluate=False)
    else:
        action = "add up to"
        if b.is_negative:
            left_expr = Add(a * x**2, -factor1 * x, -factor2 * x, c, evaluate=False)
        else:
            left_expr = Add(a * x**2, factor1 * x, factor2 * x, c, evaluate=False)
    if n != 1:
        new_left_expr = Mul(n, left_expr, evaluate=False)
    else:
        new_left_expr = left_expr
    
    steps[-1]["after"] = f"{sstr(new_left_expr)} = {sstr(right_expr)}"
    steps[-1]["step"] = f"Seperate the x coeficients equal to the factors of {sstr(ac)} that {action} {sstr(Abs(b))}"
    steps[-1]["rule"] = "Split Coeficients"
    
    ## Factor Out X on left term
    steps.append({})
    steps[-1]["before"] = steps[-2]["after"]
    
    terms = left_expr.as_ordered_terms()
    t1, t2, t3, t4 = terms[0], terms[1], terms[2], terms[3]
    factored_part1 = factor(t1 + t2)
    factored_part2 = factor(t3 + t4)
    rest = sum(terms[2:])
    new_left_expr = Add(factored_part1, rest, evaluate=False)
    new_left_expr = Mul(n, new_left_expr, evaluate=False)
    
    steps[-1]["after"] = f"{sstr(new_left_expr)} = {sstr(right_expr)}"
    steps[-1]["step"] = f"Factor out the x from the left two terms of the inner polynomial"
    steps[-1]["rule"] = "Reverse Distributive Property"
    
    ## Factor Out GCD on right term
    steps.append({})
    steps[-1]["before"] = steps[-2]["after"]
    
    left_expr = Add(factored_part1, factored_part2, evaluate=False)
    new_left_expr = Mul(n, left_expr, evaluate=False)
    
    steps[-1]["after"] = f"{sstr(new_left_expr)} = {sstr(right_expr)}"
    steps[-1]["step"] = f"Factor out the GCD from the right two terms of the inner polynomial"
    steps[-1]["rule"] = "Reverse Distributive Property"
    
    ## Add Like Terms
    if steps[-1]["before"] != steps[-1]["after"]:
        steps.append({})
    steps[-1]["before"] = steps[-2]["after"]
    terms = left_expr.as_ordered_terms()
    factors = [set(t.as_ordered_factors()) for t in terms]
    common = set.intersection(*factors)
    base = list(common)[0]
    coeffs = [t.coeff(base) for t in terms]
    new_expr = sum(coeffs) * base
    new_left_expr = flatten_mul(Mul(n, new_expr, evaluate=False))
    
    steps[-1]["after"] = f"{sstr(new_left_expr)} = {sstr(right_expr)}"
    steps[-1]["step"] = f"Collect the like terms"
    steps[-1]["rule"] = "Combine the like terms"
    
    return steps
    
def solve_roots(equation):
    # expects n(ax+b)(x+c) = 0
    
    #4 steps
    steps = []
    
    left, right = equation.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    factors = left_expr.as_ordered_factors()
    x_factors = [f for f in factors if f.has(x)]
    
    solutions = ""
    for factor in x_factors:
        clean_factor = clean(factor)
        steps.append({})
        if solutions:
            solutions = solutions + ", "
        steps[-1]["before"] = equation
        steps[-1]["after"] = f"{sstr(clean_factor)} = 0"
        steps[-1]["step"] = f"Focus on a root"
        steps[-1]["rule"] = "Root of a polynomial"
        
        current = steps[-1]["after"]
        a = clean_factor.coeff(x)
        b = clean_factor.subs(x, 0)
        if b.is_zero == False:
            if b.is_negative:
                steps.append(add_both_sides(current, -b))
            elif b.is_positive:
                steps.append(subtract_both_sides(current, b))
        current = steps[-1]["after"]
        left, right = current.split("=")
        left_expr = parse_expr(left)
        right_expr = parse_expr(right)
        steps[-1]["after"] = f"{sstr(left_expr)} = {sstr(right_expr)}"
        current = steps[-1]["after"]
        if a != 1 and a != -1:
            steps.append(divide_both_sides(current, a))
        elif a == -1:
            steps.append(multiply_both_sides(current, a))
        solutions += steps[-1]["after"]
        
    steps.append({})
    steps[-1]["before"] = equation
    steps[-1]["after"] = solutions
    steps[-1]["step"] = f"Solved The Roots"
    steps[-1]["rule"] = "Root of a polynomial"
    
    return steps
    
def combine_like_terms(equation):
    step = {}
    
    current = equation
    step["before"] = current
    left, right = current.split("=")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    ## Combine Left Terms
    left_terms = left_expr.as_ordered_terms()
    # group by base
    left_groups = {}
    for t in left_terms:
        coeff, rest = t.as_coeff_Mul()
        left_groups.setdefault(rest, 0)
        left_groups[rest] += coeff
    # rebuild manually
    new_left_terms = []
    for base, coeff in left_groups.items():
        if coeff != 0:
            new_left_terms.append(coeff * base)
    new_left_expr = sum(new_left_terms)
    
    ## Comnine Right Terms
    right_terms = right_expr.as_ordered_terms()
    # group by base
    right_groups = {}
    for t in right_terms:
        coeff, rest = t.as_coeff_Mul()
        right_groups.setdefault(rest, 0)
        right_groups[rest] += coeff
    # rebuild manually
    new_right_terms = []
    for base, coeff in right_groups.items():
        if coeff != 0:
            new_right_terms.append(coeff * base)
    new_right_expr = sum(new_right_terms)
    
    step["after"] = f"{sstr(new_left_expr)} = {sstr(new_right_expr)}"
    step["step"] = "Collect Like Terms"
    step["rule"] = "Combine the like terms"
    return step
    
def distribute_step(equation):
    steps = []
    
    current = equation
    left, right = current.split("=")
    left = left.replace("-(", "-1*(")
    x = symbols('x')
    
    left_expr = parse_expr(left, transformations=transformations, evaluate=False)
    right_expr = parse_expr(right, transformations=transformations, evaluate=False)
    
    print(f"calling distribute_once with expression: {sstr(left_expr)}")
    new_left_expr, distributed = distribute_once(left_expr)
    
    if distributed != None:
        steps.append({})
        steps[-1]["before"] = current
        steps[-1]["after"] = f"{sstr(safe_format(new_left_expr))} = {sstr(right_expr)}"
        steps[-1]["step"] = f"Distribute out {sstr(distributed)}"
        steps[-1]["rule"] = "Distributive Law of Multiplication"
        
    return steps
    
def build_ordered_add(args):
    expr = args[0]
    for a in args[1:]:
        expr = Add(expr, a, evaluate=False)
    return expr    
    
def distribute_once(expr):
    expr = flatten_mul(expr)

    # ------------------------------------------------------------
    # STEP 1: ONLY HANDLE DIRECT DISTRIBUTION CASES
    # (i.e. Mul where one factor is Add)
    # ------------------------------------------------------------
    if expr.is_Mul:

        print(f"expr: {sstr(expr)}")

        add_part = None
        other_parts = []

        # extract Add factor + everything else
        for arg in expr.args:
            print(f"arg: {sstr(arg)}")

            if arg.is_Add and add_part is None:
                add_part = arg
            else:
                other_parts.append(arg)

        # --------------------------------------------------------
        # DISTRIBUTION RULE
        # --------------------------------------------------------
        if add_part is not None:
            print(f"expr used: {sstr(expr)}, add used: {sstr(add_part)}")

            distributed_value = Mul(*other_parts)

            distributed_terms = [
                Mul(*other_parts, term)
                for term in add_part.args
            ]

            new_expr = build_ordered_add(distributed_terms)

            return new_expr, distributed_value

    # ------------------------------------------------------------
    # STEP 2: PRIORITY-BASED RECURSION (IMPORTANT FIX)
    # ------------------------------------------------------------
    if expr.args:
        print(f"step2 args:{expr.args}")
        # PASS 1: ONLY distributable Mul(Add(...))
        for i, arg in enumerate(expr.args):
            if arg.is_Mul and arg.has(Add):

                new_arg, distributed = distribute_once(arg)

                if new_arg != arg:
                    new_args = list(expr.args)
                    new_args[i] = new_arg
                    return build_ordered_add(new_args), distributed

        # PASS 2: ONLY recurse into Add or structured nodes
        for i, arg in enumerate(expr.args):

            # IMPORTANT FILTER: skip pure Mul like -4*x
            if arg.is_Mul and not any(a.is_Add for a in arg.args):
                continue

            new_arg, distributed = distribute_once(arg)

            if new_arg != arg:
                new_args = list(expr.args)
                new_args[i] = new_arg
                return build_ordered_add(new_args), distributed

    # ------------------------------------------------------------
    # STEP 3: NO CHANGE
    # ------------------------------------------------------------
    return expr, None
    
def clean(expr):

    # remove explicit 1 multipliers
    expr = expr.replace(
        lambda e: isinstance(e, Mul),
        lambda e: Mul(*[arg for arg in e.args if arg != 1])
    )

    return expr
    
def safe_format(expr):
    if expr.is_Mul:
        args = []
        for a in expr.args:
            a = safe_format(a)
            if a == 1:
                continue
            args.append(a)
        return Mul(*args, evaluate=False)

    if expr.is_Add:
        return expr.func(*[safe_format(a) for a in expr.args], evaluate=False)

    return expr
    
def flatten_mul(expr):
    if expr.is_Mul:
        args = []
        for arg in expr.args:
            if arg.is_Mul:
                args.extend(arg.args)
            else:
                args.append(arg)
        return Mul(*args, evaluate=False)
    return expr