RustAst/src/ast/rtl/math.rs

use crate::ast::environment::Environment;
use crate::ast::types::{NativeFunction, Purity, Signature, StaticType, Value};
use std::cell::RefCell;
use std::f64::consts;
use std::rc::Rc;

pub fn register(env: &Environment) {
    // Constants
    env.register_constant("PI", StaticType::Float, Value::Float(consts::PI))
        .doc("Mathematical constant π ≈ 3.14159265.");
    env.register_constant("E", StaticType::Float, Value::Float(consts::E))
        .doc("Mathematical constant e ≈ 2.71828182 (Euler's number).");

    // Helper to get f64 from Value (Int or Float)
    fn to_f64(v: &Value) -> Option<f64> {
        match v {
            Value::Float(f) => Some(*f),
            Value::Int(i) => Some(*i as f64),
            _ => None,
        }
    }

    // Unary functions (float -> float)
    let register_unary = |name: &'static str, f: fn(f64) -> f64, doc: &'static str| {
        let ty = StaticType::Function(Box::new(Signature {
            params: StaticType::Tuple(vec![StaticType::Float]),
            ret: StaticType::Float,
        }));
        env.register_native_fn(name, ty, Purity::Pure, move |args| {
            if let Some(x) = to_f64(&args[0]) {
                Value::Float(f(x))
            } else {
                Value::Void
            }
        }).doc(doc);
    };

    // Unary functions (float -> int)
    let register_to_int = |name: &'static str, f: fn(f64) -> f64, doc: &'static str| {
        let ty = StaticType::Function(Box::new(Signature {
            params: StaticType::Tuple(vec![StaticType::Float]),
            ret: StaticType::Int,
        }));
        env.register_native_fn(name, ty, Purity::Pure, move |args| {
            if let Some(x) = to_f64(&args[0]) {
                Value::Int(f(x) as i64)
            } else {
                Value::Void
            }
        }).doc(doc);
    };

    register_unary("sqrt", f64::sqrt, "Returns the square root of x.");
    register_unary("sin",  f64::sin,  "Returns the sine of x (radians).");
    register_unary("cos",  f64::cos,  "Returns the cosine of x (radians).");
    register_unary("tan",  f64::tan,  "Returns the tangent of x (radians).");
    register_unary("asin", f64::asin, "Returns the arcsine of x in radians.");
    register_unary("acos", f64::acos, "Returns the arccosine of x in radians.");
    register_unary("atan", f64::atan, "Returns the arctangent of x in radians.");
    register_unary("exp",  f64::exp,  "Returns e raised to the power of x.");
    register_unary("ln",   f64::ln,   "Returns the natural logarithm of x.");
    register_unary("log10",f64::log10,"Returns the base-10 logarithm of x.");
    register_unary("fract",f64::fract,"Returns the fractional part of x.");

    // Rounding functions returning Int
    register_to_int("round", f64::round, "Rounds x to the nearest integer.");
    register_to_int("floor", f64::floor, "Rounds x down to the largest integer ≤ x.");
    register_to_int("ceil",  f64::ceil,  "Rounds x up to the smallest integer ≥ x.");
    register_to_int("trunc", f64::trunc, "Truncates x towards zero.");

    // Binary functions (float, float -> float)
    let register_binary = |name: &'static str, f: fn(f64, f64) -> f64, doc: &'static str| {
        let ty = StaticType::Function(Box::new(Signature {
            params: StaticType::Tuple(vec![StaticType::Float, StaticType::Float]),
            ret: StaticType::Float,
        }));
        env.register_native_fn(name, ty, Purity::Pure, move |args| {
            if let (Some(a), Some(b)) = (to_f64(&args[0]), to_f64(&args[1])) {
                Value::Float(f(a, b))
            } else {
                Value::Void
            }
        }).doc(doc);
    };

    register_binary("atan2", f64::atan2, "Returns the arctangent of y/x, using signs to determine the quadrant.");
    register_binary("pow",   f64::powf,  "Returns x raised to the power y.");
    register_binary("log",   f64::log,   "Returns the logarithm of x in the given base: (log base x).");

    // Specialized abs to handle Int -> Int, Float -> Float
    let abs_ty = StaticType::Function(Box::new(Signature {
        params: StaticType::Tuple(vec![StaticType::Any]),
        ret: StaticType::Any,
    }));
    env.register_native_fn("abs", abs_ty, Purity::Pure, |args| {
        if args.len() != 1 {
            return Value::Void;
        }
        match args[0] {
            Value::Int(i) => Value::Int(i.abs()),
            Value::Float(f) => Value::Float(f.abs()),
            _ => Value::Void,
        }
    }).doc("Returns the absolute value of a number. Works for both int and float.")
      .examples(&["(abs -5)", "(abs -3.14)"]);

    // Variadic min / max — each argument must be numeric
    let variadic_ty = StaticType::Function(Box::new(Signature {
        params: StaticType::Variadic(Box::new(StaticType::Float)),
        ret: StaticType::Float,
    }));

    env.register_native_fn("min", variadic_ty.clone(), Purity::Pure, |args| {
        if args.is_empty() {
            return Value::Void;
        }
        let mut best = args[0].clone();
        for arg in &args[1..] {
            if let Some(ord) = arg.partial_cmp(&best)
                && ord == std::cmp::Ordering::Less
            {
                best = arg.clone();
            }
        }
        best
    }).doc("Returns the minimum of its arguments. Variadic: (min 3 1 2) => 1.");

    env.register_native_fn("max", variadic_ty, Purity::Pure, |args| {
        if args.is_empty() {
            return Value::Void;
        }
        let mut best = args[0].clone();
        for arg in &args[1..] {
            if let Some(ord) = arg.partial_cmp(&best)
                && ord == std::cmp::Ordering::Greater
            {
                best = arg.clone();
            }
        }
        best
    }).doc("Returns the maximum of its arguments. Variadic: (max 3 1 2) => 3.");

    // Random generator factory (Closure-based)
    let generator_sig = Signature {
        params: StaticType::Tuple(vec![]),
        ret: StaticType::Float,
    };

    let make_random_ty = StaticType::FunctionOverloads(vec![
        Signature {
            params: StaticType::Tuple(vec![]),
            ret: StaticType::Function(Box::new(generator_sig.clone())),
        },
        Signature {
            params: StaticType::Tuple(vec![StaticType::Int]),
            ret: StaticType::Function(Box::new(generator_sig)),
        },
    ]);

    env.register_native_fn("make-random", make_random_ty, Purity::SideEffectFree, |args| {
            let rng = if args.is_empty() {
                fastrand::Rng::new()
            } else if let Value::Int(s) = args[0] {
                fastrand::Rng::with_seed(s as u64)
            } else {
                fastrand::Rng::new()
            };

            let prng = Rc::new(RefCell::new(rng));

            Value::Function(Rc::new(NativeFunction {
                func: Rc::new(move |_args| Value::Float(prng.borrow_mut().f64())),
                purity: Purity::Impure,
            }))
    }).doc("Creates a random number generator function returning floats in [0, 1).")
      .description("Optionally seeded: (make-random seed) for reproducible sequences.")
      .examples(&["(def rng (make-random 42))", "(rng)"]);
}