import math
import numpy as np
import sympy as sp

# ---------------- PRIME FUNCTIONS ----------------

def primelist(n):
    """Return a list of all prime numbers up to n"""
    def is_prime(num):
        """Check if num is prime"""
        if num < 2:
            return False
        for i in range(2, int(math.isqrt(num)) + 1):
            if num % i == 0:
                return False
        return True
    
    def find_prime():
        """Generate primes up to n"""
        return [i for i in range(2, n+1) if is_prime(i)]
    
    return find_prime()

def biggestprime(n):
    """Return the largest prime ≤ n"""
    for i in range(n, 1, -1):
        if all(i % j != 0 for j in range(2, int(math.isqrt(i))+1)):
            return i
    return None

def smallestprime(n):
    """Return the smallest prime ≥ n"""
    i = max(n, 2)
    while True:
        if all(i % j != 0 for j in range(2, int(math.isqrt(i))+1)):
            return i
        i += 1

# ---------------- TRIG FUNCTIONS ----------------

def cos(n):
    """Return cosine of n (radians)"""
    return math.cos(n)

def sin(n):
    """Return sine of n (radians)"""
    return math.sin(n)

def tan(n):
    """Return tangent of n (radians)"""
    return math.tan(n)

# ---------------- LINEAR EQUATIONS ----------------

def linear_solving(a, b, result):
    """Solve linear equation a*x + b = result; returns solution x"""
    x = sp.symbols('x')
    solution = sp.solve(a*x + b - result, x)
    return solution

def solve_x(a, result):
    """Solve simple equation a*x = result"""
    if a == 0:
        return "No solution" if result != 0 else "Infinite solutions"
    return result / a

# ---------------- SIGMA NOTATION ----------------

def sigma_notaion(E, n, k=1):
    """Compute sum of E from k to n. E can be a string expression using 'k'"""
    k_sym = sp.symbols('k')
    expr = sp.sympify(E)
    total = sum(expr.subs(k_sym, i) for i in range(k, n+1))
    return total

# ---------------- KNUTH ARROW ----------------

def knuth_uparrow(a, b, arrow):
    """Implement Knuth's up-arrow notation: ↑, ↑↑, ↑↑↑"""
    if arrow == 1:
        return a ** b
    elif arrow == 2:
        result = a
        for _ in range(b-1):
            result = a ** result
        return result
    elif arrow == 3:
        result = a
        for _ in range(b-1):
            temp = a
            for _ in range(result-1):
                temp = a ** temp
            result = temp
        return result
    else:
        raise ValueError("Arrow value > 3 not implemented")

# ---------------- ARRAY OPERATIONS ----------------

def array_add(array_of_a, array_of_b):
    """Return element-wise addition of two arrays"""
    a = np.array(array_of_a)
    b = np.array(array_of_b)
    return a + b

def array_sub(array_of_a, array_of_b):
    """Return element-wise subtraction of two arrays"""
    a = np.array(array_of_a)
    b = np.array(array_of_b)
    return a - b

def array_multi(array_of_a, array_of_b):
    """Return element-wise multiplication of two arrays"""
    a = np.array(array_of_a)
    b = np.array(array_of_b)
    return a * b

def array_div(array_of_a, array_of_b):
    """Return element-wise division of two arrays"""
    a = np.array(array_of_a, dtype=float)
    b = np.array(array_of_b, dtype=float)
    return a / b

# ---------------- REPEATED POWER ----------------

def repeated_power(a, n, r):
    """Compute a^(a^(...)) r times with n as exponent if needed"""
    result = a ** n
    for _ in range(r-1):
        result = a ** result
    return result

# ---------------- MATH ----------------

def sqrt(n):
    """Return square root of n"""
    return math.sqrt(n)

def plus_minus(a, b):
    """Return tuple: (a+b, a-b)"""
    return (a + b, a - b)

def minus_plus(a, b):
    """Return tuple: (a-b, a+b)"""
    return (a - b, a + b)

def similarity(a, b):
    """Return cosine similarity if a,b are vectors"""
    a = np.array(a)
    b = np.array(b)
    if np.linalg.norm(a) == 0 or np.linalg.norm(b) == 0:
        return 0
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def div_all(a, b):
    """Return tuple: (a/b, b/a)"""
    return (a/b, b/a)

# ---------------- SET OPERATIONS ----------------

def union(set_a, set_b):
    """Return a list of unique elements from both set_a and set_b.
    Uses a set 'seen' to track added elements so each appears only once.
    Iterates through set_a + set_b and includes each element only if not seen before.
    """
    seen = set()
    return [x for x in set_a + set_b if not (x in seen or seen.add(x))]

def intersection(set_a, set_b):
    """Return a list of elements common to both set_a and set_b.
    Converts set_b to a set for fast membership checking.
    Iterates through set_a and includes elements present in set_b.
    """
    b_set = set(set_b)
    return [x for x in set_a if x in b_set]

def difference(set_a, set_b):
    """Return a list of elements in set_a that are not in set_b.
    Converts set_b to a set for fast lookup.
    Iterates through set_a and includes only elements not in set_b.
    """
    b_set = set(set_b)
    return [x for x in set_a if x not in b_set]

def symmetric_difference(set_a, set_b):
    """Return a list of elements that are in either set_a or set_b, but not both.
    Converts both lists to sets for faster membership checks.
    Includes elements in set_a not in set_b, plus elements in set_b not in set_a.
    """
    b_set = set(set_b)
    a_set = set(set_a)
    return [x for x in set_a if x not in b_set] + [x for x in set_b if x not in a_set]



# ---------------- FACTORIAL ----------------

def factorial(n):
    """Return factorial of n"""
    return math.factorial(n)

# ---------------- INTEGRAL ----------------

def intergal(step, f, x, dx):
    """Numerical integration using simple Riemann sum"""
    total = 0
    for i in range(step):
        total += f(x + i*dx) * dx
    return total

# ---------------- OTHER MATH ----------------

def absolute_value(val):
    """Return absolute value"""
    return abs(val)

def xor(a, b):
    """Bitwise XOR"""
    return a ^ b

def imaginary(val):
    """Return imaginary part of complex number"""
    return val.imag

def real(val):
    """Return real part of complex number"""
    return val.real

def partial(expr_str, var_str, val_dict):
    """Compute partial derivative of expr_str w.r.t var_str and substitute values"""
    expr = sp.sympify(expr_str)
    var = sp.symbols(var_str)
    return sp.diff(expr, var).subs(val_dict)

# ---------------- QUANTIFIERS ----------------

def universal_quantifier(condition, N):
    """Universal quantifier ∀n ∈ N : condition(n)"""
    return all(condition(n) for n in N)

def existential_quantifier(condition, N):
    """Existential quantifier ∃n ∈ N : condition(n)"""
    return any(condition(n) for n in N)

def uniqueness_quantifier(condition, N):
    """Uniqueness quantifier ∃! n ∈ N : condition(n)"""
    return sum(1 for n in N if condition(n)) == 1

# ---------------- LOGARITHMS ----------------

def log(x, base=math.e):
    """Natural log by default; base-n log if specified"""
    if isinstance(x, (int, float)):
        return math.log(x, base)
    return sp.log(x, base)

def in_log(base, x):
    """Check if log exists for base and x; returns symbolic or numeric value"""
    if isinstance(x, (int, float)) and x <= 0:
        raise ValueError("Log undefined for x <= 0")
    return sp.log(x, base)

# ---------------- LIMIT ----------------

def lim(f, x, x0):
    """Compute the limit of f(x) as x -> x0"""
    if callable(f):
        delta = 1e-7
        return f(x0 + delta)
    elif isinstance(f, sp.Expr):
        return sp.limit(f, x, x0)
    else:
        raise TypeError("f must be a function or sympy expression")

# ---------------- SET MEMBERSHIP ----------------

def belong_to(s, n):
    """Check if n belongs to set s"""
    return n in s

def dont_belong_to(s, n):
    """Check if n does NOT belong to set s"""
    return n not in s

# ---------------- EQUALITY -----------------------
def equal_to(a,b):
    """ check if a equal to b """
    return a == b

def not_equal_to(a,b):
    """ check if a not equal to b """
    return a != b
# ---------------- MORE_GREATER ------------------------
def much_greater(a, b):
    """Check if a is more than twice b."""
    return a > 2 * b

def much_smaller(a, b):
    """Check if a is less than half of b."""
    return a < 0.5 * b

def roughly_equal(a, b, tolerance=0.1):
    """Check if a and b are approximately equal within a tolerance (default 10%)."""
    return abs(a - b) <= tolerance * max(abs(a), abs(b))
# ---------------- CUSTOM MATH --------------------------
def star(a,b):
    """ Take the first number and add twice the second number(2 * b) """
    return a + 2 * b

def plus_twist(a,b):
    """ add a and b but alway add extra 1 """
    return a + b + 1

def multiply_tweak(a,b):
    """ mutil a by b and then add a """
    return a * b + a

def swap_add(a,b):
    """ take b and add half of a """
    return b + a / 2

def double_add(a,b):
    """ add a and b,and mirror,and add both result """
    return (a + b) + (a + b)

def factorial_power(a,n):
    """ factorize n first,then rise a to n! """
    return a ** math.factorial(n)