//! # Elementary Functions and Expression Manipulation
//!
//! This module provides constructor functions for creating elementary mathematical expressions
//! (like trigonometric, exponential, and power functions) and tools for manipulating
//! these expressions, such as `expand` using algebraic and trigonometric identities.
use crate::symbolic::core::Expr;
use crate::symbolic::simplify_dag::simplify;
use num_bigint::BigInt;
use num_traits::{One, ToPrimitive, Zero};
use std::sync::Arc;
/// Creates a sine expression: `sin(expr)`.
pub fn sin(expr: Expr) -> Expr {
    Expr::new_sin(expr)
}
/// Creates a cosine expression: `cos(expr)`.
pub fn cos(expr: Expr) -> Expr {
    Expr::new_cos(expr)
}
/// Creates a tangent expression: `tan(expr)`.
pub fn tan(expr: Expr) -> Expr {
    Expr::new_tan(expr)
}
/// Creates a hyperbolic sine expression: `sinh(expr)`.
pub fn sinh(expr: Expr) -> Expr {
    Expr::new_sinh(expr)
}
/// Creates a hyperbolic cosine expression: `cosh(expr)`.
pub fn cosh(expr: Expr) -> Expr {
    Expr::new_cosh(expr)
}
/// Creates a hyperbolic tangent expression: `tanh(expr)`.
pub fn tanh(expr: Expr) -> Expr {
    Expr::new_tanh(expr)
}
/// Creates a natural logarithm expression: `ln(expr)`.
pub fn ln(expr: Expr) -> Expr {
    Expr::new_log(expr)
}
/// Creates an exponential expression: `e^(expr)`.
pub fn exp(expr: Expr) -> Expr {
    Expr::new_exp(expr)
}
/// Creates a square root expression: `sqrt(expr)`.
pub fn sqrt(expr: Expr) -> Expr {
    Expr::new_sqrt(expr)
}
/// Creates a power expression: `base^exp`.
pub fn pow(base: Expr, exp: Expr) -> Expr {
    Expr::new_pow(base, exp)
}
/// Returns the symbolic representation of positive infinity.
pub const fn infinity() -> Expr {
    Expr::Infinity
}
/// Returns the symbolic representation of negative infinity.
pub const fn negative_infinity() -> Expr {
    Expr::NegativeInfinity
}
/// Creates a logarithm expression with a specified base: `log_base(expr)`.
pub fn log_base(base: Expr, expr: Expr) -> Expr {
    Expr::new_log_base(base, expr)
}
/// Creates a cotangent expression: `cot(expr)`.
pub fn cot(expr: Expr) -> Expr {
    Expr::new_cot(expr)
}
/// Creates a secant expression: `sec(expr)`.
pub fn sec(expr: Expr) -> Expr {
    Expr::new_sec(expr)
}
/// Creates a cosecant expression: `csc(expr)`.
pub fn csc(expr: Expr) -> Expr {
    Expr::new_csc(expr)
}
/// Creates an inverse cotangent expression: `acot(expr)`.
pub fn acot(expr: Expr) -> Expr {
    Expr::new_arccot(expr)
}
/// Creates an inverse secant expression: `asec(expr)`.
pub fn asec(expr: Expr) -> Expr {
    Expr::new_arcsec(expr)
}
/// Creates an inverse cosecant expression: `acsc(expr)`.
pub fn acsc(expr: Expr) -> Expr {
    Expr::new_arccsc(expr)
}
/// Creates a hyperbolic cotangent expression: `coth(expr)`.
pub fn coth(expr: Expr) -> Expr {
    Expr::new_coth(expr)
}
/// Creates a hyperbolic secant expression: `sech(expr)`.
pub fn sech(expr: Expr) -> Expr {
    Expr::new_sech(expr)
}
/// Creates a hyperbolic cosecant expression: `csch(expr)`.
pub fn csch(expr: Expr) -> Expr {
    Expr::new_csch(expr)
}
/// Creates an inverse hyperbolic sine expression: `asinh(expr)`.
pub fn asinh(expr: Expr) -> Expr {
    Expr::new_arcsinh(expr)
}
/// Creates an inverse hyperbolic cosine expression: `acosh(expr)`.
pub fn acosh(expr: Expr) -> Expr {
    Expr::new_arccosh(expr)
}
/// Creates an inverse hyperbolic tangent expression: `atanh(expr)`.
pub fn atanh(expr: Expr) -> Expr {
    Expr::new_arctanh(expr)
}
/// Creates an inverse hyperbolic cotangent expression: `acoth(expr)`.
pub fn acoth(expr: Expr) -> Expr {
    Expr::new_arccoth(expr)
}
/// Creates an inverse hyperbolic secant expression: `asech(expr)`.
pub fn asech(expr: Expr) -> Expr {
    Expr::new_arcsech(expr)
}
/// Creates an inverse hyperbolic cosecant expression: `acsch(expr)`.
pub fn acsch(expr: Expr) -> Expr {
    Expr::new_arccsch(expr)
}
/// Creates a 2-argument inverse tangent expression: `atan2(y, x)`.
pub fn atan2(y: Expr, x: Expr) -> Expr {
    Expr::new_atan2(y, x)
}
/// Returns the symbolic representation of Pi.
pub const fn pi() -> Expr {
    Expr::Pi
}
/// Returns the symbolic representation of Euler's number (e).
pub const fn e() -> Expr {
    Expr::E
}
/// Expands a symbolic expression by applying distributive, power, and trigonometric identities.
///
/// This is often the reverse of simplification and can reveal the underlying structure of an expression.
/// The expansion is applied recursively to all parts of the expression.
///
/// # Arguments
/// * `expr` - The expression to expand.
///
/// # Returns
/// A new, expanded `Expr`.
pub fn expand(expr: Expr) -> Expr {
    simplify(&expand_internal(expr))
}

fn expand_internal(expr: Expr) -> Expr {
    match expr {
        Expr::Dag(node) => expand_internal(node.to_expr().expect("Expand")),
        Expr::Complex(re, im) => {
            // Convert Complex(re, im) to Add(re, Mul(im, i)) for expansion
            let i = Expr::Variable("i".to_string());
            let expanded_re = expand_internal((*re).clone());
            let expanded_im = expand_internal((*im).clone());
            Expr::Add(
                Arc::new(expanded_re),
                Arc::new(Expr::Mul(Arc::new(expanded_im), Arc::new(i))),
            )
        }
        Expr::Add(a, b) => Expr::Add(
            Arc::new(expand_internal((*a).clone())),
            Arc::new(expand_internal((*b).clone())),
        ),
        Expr::Sub(a, b) => Expr::Sub(
            Arc::new(expand_internal((*a).clone())),
            Arc::new(expand_internal((*b).clone())),
        ),
        Expr::Mul(a, b) => expand_mul((*a).clone(), (*b).clone()),
        Expr::Div(a, b) => Expr::Div(
            Arc::new(expand_internal((*a).clone())),
            Arc::new(expand_internal((*b).clone())),
        ),
        Expr::Power(b, e) => expand_power(&b, &e),
        Expr::Log(arg) => expand_log(&arg),
        Expr::Sin(arg) => expand_sin(&arg),
        Expr::Cos(arg) => expand_cos(&arg),
        Expr::Exp(arg) => expand_exp(&arg),
        _ => expr,
    }
}
/// Expands multiplication over addition: `a*(b+c) -> a*b + a*c`.
pub(crate) fn expand_mul(a: Expr, b: Expr) -> Expr {
    let a_exp = expand_internal(a);
    let b_exp = expand_internal(b);
    match (a_exp, b_exp) {
        (l, Expr::Add(m, n)) => Expr::Add(
            Arc::new(expand_internal(Expr::Mul(Arc::new(l.clone()), m))),
            Arc::new(expand_internal(Expr::Mul(Arc::new(l), n))),
        ),
        (Expr::Add(m, n), r) => Expr::Add(
            Arc::new(expand_internal(Expr::Mul(m, Arc::new(r.clone())))),
            Arc::new(expand_internal(Expr::Mul(n, Arc::new(r)))),
        ),
        (l, r) => Expr::Mul(Arc::new(l), Arc::new(r)),
    }
}
/// Expands powers, e.g., `(a*b)^c -> a^c * b^c` and `(a+b)^n -> a^n + ...` (binomial expansion).
pub(crate) fn expand_power(base: &Arc<Expr>, exp: &Arc<Expr>) -> Expr {
    let b_exp = expand_internal(base.as_ref().clone());
    let e_exp = expand_internal(exp.as_ref().clone());
    match (b_exp, e_exp) {
        (Expr::Mul(f, g), e) => Expr::Mul(
            Arc::new(expand_internal(Expr::Power(f, Arc::new(e.clone())))),
            Arc::new(expand_internal(Expr::Power(g, Arc::new(e)))),
        ),
        (Expr::Add(a, b), Expr::BigInt(n)) => {
            if let Some(n_usize) = n.to_usize() {
                expand_binomial(a, b, n_usize)
            } else {
                Expr::Power(Arc::new(Expr::Add(a, b)), Arc::new(Expr::BigInt(n)))
            }
        }
        (Expr::Add(a, b), Expr::Constant(c)) if c.fract() == 0.0 && c >= 0.0 => {
            let n_usize = c as usize;
            expand_binomial(a, b, n_usize)
        }
        (b, e) => Expr::Power(Arc::new(b), Arc::new(e)),
    }
}

fn expand_binomial(a: Arc<Expr>, b: Arc<Expr>, n: usize) -> Expr {
    let mut sum = Expr::BigInt(BigInt::zero());
    for k in 0..=n {
        let bin_coeff = Expr::BigInt(binomial_coefficient(n, k));
        let term1 = Expr::Power(a.clone(), Arc::new(Expr::BigInt(BigInt::from(k))));
        let term2 = Expr::Power(b.clone(), Arc::new(Expr::BigInt(BigInt::from(n - k))));
        let term = Expr::Mul(
            Arc::new(bin_coeff),
            Arc::new(Expr::Mul(Arc::new(term1), Arc::new(term2))),
        );
        sum = Expr::Add(Arc::new(sum), Arc::new(expand_internal(term)));
    }
    sum
}
/// Expands logarithms using identities like `log(a*b) -> log(a) + log(b)`.
pub(crate) fn expand_log(arg: &Arc<Expr>) -> Expr {
    let arg_exp = expand_internal(arg.as_ref().clone());
    match arg_exp {
        Expr::Mul(a, b) => Expr::Add(
            Arc::new(expand_internal(Expr::Log(a))),
            Arc::new(expand_internal(Expr::Log(b))),
        ),
        Expr::Div(a, b) => Expr::Sub(
            Arc::new(expand_internal(Expr::Log(a))),
            Arc::new(expand_internal(Expr::Log(b))),
        ),
        Expr::Power(b, e) => Expr::Mul(e, Arc::new(expand_internal(Expr::Log(b)))),
        a => Expr::Log(Arc::new(a)),
    }
}
/// Expands `sin` using sum-angle identities, e.g., `sin(a+b)`.
pub(crate) fn expand_sin(arg: &Arc<Expr>) -> Expr {
    let arg_exp = expand_internal(arg.as_ref().clone());
    match arg_exp {
        Expr::Add(a, b) => Expr::Add(
            Arc::new(Expr::Mul(
                Arc::new(sin(a.as_ref().clone())),
                Arc::new(cos(b.as_ref().clone())),
            )),
            Arc::new(Expr::Mul(
                Arc::new(cos(a.as_ref().clone())),
                Arc::new(sin(b.as_ref().clone())),
            )),
        ),
        a => Expr::Sin(Arc::new(a)),
    }
}
/// Expands `cos` using sum-angle identities, e.g., `cos(a+b)`.
pub(crate) fn expand_cos(arg: &Arc<Expr>) -> Expr {
    let arg_exp = expand_internal(arg.as_ref().clone());
    match arg_exp {
        Expr::Add(a, b) => Expr::Sub(
            Arc::new(Expr::Mul(
                Arc::new(cos(a.as_ref().clone())),
                Arc::new(cos(b.as_ref().clone())),
            )),
            Arc::new(Expr::Mul(
                Arc::new(sin(a.as_ref().clone())),
                Arc::new(sin(b.as_ref().clone())),
            )),
        ),
        a => Expr::Cos(Arc::new(a)),
    }
}

/// Expands `exp` using identities like `exp(a+b) -> exp(a) * exp(b)`.
pub(crate) fn expand_exp(arg: &Arc<Expr>) -> Expr {
    let arg_exp = expand_internal(arg.as_ref().clone());
    match arg_exp {
        Expr::Add(a, b) => Expr::Mul(
            Arc::new(expand_internal(Expr::Exp(a))),
            Arc::new(expand_internal(Expr::Exp(b))),
        ),
        Expr::Mul(a, b) => {
            let is_i = |e: &Expr| matches!(e, Expr::Variable(name) if name == "i");
            if is_i(&a) {
                // exp(i * b) = cos(b) + i * sin(b)
                let cos_b = Expr::Cos(b.clone());
                let sin_b = Expr::Sin(b.clone());
                Expr::Add(Arc::new(cos_b), Arc::new(Expr::Mul(a, Arc::new(sin_b))))
            } else if is_i(&b) {
                // exp(a * i) = cos(a) + i * sin(a)
                let cos_a = Expr::Cos(a.clone());
                let sin_a = Expr::Sin(a.clone());
                Expr::Add(Arc::new(cos_a), Arc::new(Expr::Mul(b, Arc::new(sin_a))))
            } else {
                Expr::Exp(Arc::new(Expr::Mul(a, b)))
            }
        }
        a => Expr::Exp(Arc::new(a)),
    }
}
/// Helper to compute binomial coefficients C(n, k) = n! / (k! * (n-k)!).
pub fn binomial_coefficient(n: usize, k: usize) -> BigInt {
    if k > n {
        return BigInt::zero();
    }
    let mut res = BigInt::one();
    for i in 0..k {
        res = (res * (n - i)) / (i + 1);
    }
    res
}
