//! Reverse-mode automatic differentiation over scalar values.  Every
//! `Value` tracks its own gradient and its position in the computation graph,
//! enabling backpropagation through arbitrary arithmetic and activation
//! functions.

use std::cell::RefCell;
use std::collections::HashSet;
use std::ops::{Add, Div, Mul, Neg, Sub};
use std::rc::Rc;

use crate::activations::{Activation, derivative};
use crate::error::{TuphiumError, TuphiumResult};

#[derive(Clone)]
pub struct Value(Rc<RefCell<ValueInner>>);

#[derive(Clone)]
struct ValueInner {
    data: f32,
    grad: f32,
    children: Vec<Value>,
    op: Op,
}

#[derive(Clone)]
enum Op {
    Leaf,
    Add,
    Mul,
    Sub,
    Div,
    Pow(f32),
    Neg,
    Exp,
    Log,
    Activate(Activation),
}

impl Value {
    /// Creates a new leaf `Value` holding `data` with a zero gradient.
    ///
    /// Leaves are the inputs of the computation graph (parameters, constants);
    /// every other `Value` is produced by an operation on existing nodes.
    ///
    /// # Examples
    ///
    /// ```
    /// use tuphium::autograd::Value;
    ///
    /// let x = Value::new(3.0).unwrap();
    /// assert_eq!(x.data(), 3.0);
    /// assert_eq!(x.grad(), 0.0);
    /// ```
    pub fn new(data: f32) -> TuphiumResult<Self> {
        if !data.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::new",
                value: data,
            });
        }
        Ok(Self::new_unchecked(data))
    }

    fn new_unchecked(data: f32) -> Self {
        Value(Rc::new(RefCell::new(ValueInner {
            data,
            grad: 0.0,
            children: Vec::new(),
            op: Op::Leaf,
        })))
    }

    /// Returns the scalar value currently stored in this node.
    pub fn data(&self) -> f32 {
        self.0.borrow().data
    }

    /// Returns the accumulated gradient of this node.
    ///
    /// The gradient is `0.0` until [`Value::backward`] is run on an output that
    /// depends on this node.
    pub fn grad(&self) -> f32 {
        self.0.borrow().grad
    }

    /// Resets the gradient of this node and of every node it depends on to `0.0`.
    ///
    /// Call this before each new backward pass so gradients from previous passes
    /// do not accumulate. Shared subgraphs are visited only once.
    pub fn zero_grad(&self) {
        let mut visited = HashSet::new();
        zero_grad_recursive(self, &mut visited);
    }

    /// Returns a new node equal to `self + other`, recording both as children.
    ///
    /// The four reference/owned combinations of the `+` operator are also
    /// implemented, so `&a + &b`, `a + b`, etc. all forward to this method.
    pub fn add(&self, other: &Value) -> Value {
        let left = self.0.borrow().data;
        let right = other.0.borrow().data;
        Value(Rc::new(RefCell::new(ValueInner {
            data: left + right,
            grad: 0.0,
            children: vec![self.clone(), other.clone()],
            op: Op::Add,
        })))
    }

    /// Returns a new node equal to `self * other`, recording both as children.
    ///
    /// Also reachable through the `*` operator in all reference/owned combinations.
    pub fn mul(&self, other: &Value) -> Value {
        let left = self.0.borrow().data;
        let right = other.0.borrow().data;
        Value(Rc::new(RefCell::new(ValueInner {
            data: left * right,
            grad: 0.0,
            children: vec![self.clone(), other.clone()],
            op: Op::Mul,
        })))
    }

    /// Returns a new node equal to `self - other`, recording both as children.
    ///
    /// Also reachable through the `-` operator in all reference/owned combinations.
    pub fn sub(&self, other: &Value) -> Value {
        let left = self.0.borrow().data;
        let right = other.0.borrow().data;
        Value(Rc::new(RefCell::new(ValueInner {
            data: left - right,
            grad: 0.0,
            children: vec![self.clone(), other.clone()],
            op: Op::Sub,
        })))
    }

    /// Returns a new node equal to `self / other`, recording both as children.
    ///
    /// Also reachable through the `/` operator in all reference/owned
    /// combinations. The caller is responsible for keeping `other` non-zero.
    pub fn div(&self, other: &Value) -> TuphiumResult<Value> {
        let numerator = self.0.borrow().data;
        let denominator = other.0.borrow().data;
        if !numerator.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::div",
                value: numerator,
            });
        }
        if !denominator.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::div",
                value: denominator,
            });
        }
        if denominator == 0.0 {
            return Err(TuphiumError::DivisionByZero {
                operation: "Value::div",
                numerator,
            });
        }
        Ok(Value(Rc::new(RefCell::new(ValueInner {
            data: numerator / denominator,
            grad: 0.0,
            children: vec![self.clone(), other.clone()],
            op: Op::Div,
        }))))
    }

    fn div_unchecked(&self, other: &Value) -> Value {
        let numerator = self.0.borrow().data;
        let denominator = other.0.borrow().data;
        Value(Rc::new(RefCell::new(ValueInner {
            data: numerator / denominator,
            grad: 0.0,
            children: vec![self.clone(), other.clone()],
            op: Op::Div,
        })))
    }

    /// Returns a new node equal to `self` raised to the constant power `exp`.
    ///
    /// The exponent is a plain `f32` (not a tracked `Value`), so only `self`
    /// receives a gradient during the backward pass.
    pub fn pow(&self, exp: f32) -> TuphiumResult<Value> {
        let base = self.0.borrow().data;
        if !base.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::pow",
                value: base,
            });
        }
        if !exp.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::pow",
                value: exp,
            });
        }
        let data = base.powf(exp);
        if !data.is_finite() {
            return Err(TuphiumError::NumericDomain {
                operation: "Value::pow",
                input: base,
                reason: "power result must be finite in the real domain",
            });
        }
        Ok(Value(Rc::new(RefCell::new(ValueInner {
            data,
            grad: 0.0,
            children: vec![self.clone()],
            op: Op::Pow(exp),
        }))))
    }

    /// Returns a new node equal to `-self`.
    ///
    /// Also reachable through the unary `-` operator on `Value` and `&Value`.
    pub fn neg(&self) -> Value {
        let input = self.0.borrow().data;
        Value(Rc::new(RefCell::new(ValueInner {
            data: -input,
            grad: 0.0,
            children: vec![self.clone()],
            op: Op::Neg,
        })))
    }

    /// Returns a new node equal to `e^self`.
    pub fn exp(&self) -> TuphiumResult<Value> {
        let input = self.0.borrow().data;
        if !input.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::exp",
                value: input,
            });
        }
        let data = input.exp();
        if !data.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::exp",
                value: data,
            });
        }
        Ok(Value(Rc::new(RefCell::new(ValueInner {
            data,
            grad: 0.0,
            children: vec![self.clone()],
            op: Op::Exp,
        }))))
    }

    /// Returns a new node equal to the natural logarithm of `self`.
    ///
    /// `self` must be strictly positive; otherwise the stored data becomes
    /// `NaN`/`-inf`, as with [`f32::ln`].
    pub fn log(&self) -> TuphiumResult<Value> {
        let input = self.0.borrow().data;
        if !input.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "Value::log",
                value: input,
            });
        }
        if input <= 0.0 {
            return Err(TuphiumError::NumericDomain {
                operation: "Value::log",
                input,
                reason: "logarithm requires a positive input",
            });
        }
        Ok(Value(Rc::new(RefCell::new(ValueInner {
            data: input.ln(),
            grad: 0.0,
            children: vec![self.clone()],
            op: Op::Log,
        }))))
    }

    /// Returns a new node with the logistic sigmoid of `self` applied.
    ///
    /// Convenience wrapper around [`Value::activate`] with [`Activation::Sigmoid`].
    pub fn sigmoid(&self) -> Value {
        self.activate(&Activation::Sigmoid)
    }

    /// Returns a new node with `ReLU` applied to `self`.
    ///
    /// Convenience wrapper around [`Value::activate`] with [`Activation::ReLU`].
    pub fn relu(&self) -> Value {
        self.activate(&Activation::ReLU)
    }

    /// Returns a new node with the given activation applied to `self`.
    ///
    /// Records the activation in the graph so the chain rule is applied
    /// automatically during [`Value::backward`].
    pub fn activate(&self, act: &Activation) -> Value {
        let input = self.0.borrow().data;
        let data = match act {
            Activation::Sigmoid => 1.0 / (1.0 + (-input).exp()),
            Activation::ReLU => {
                if input > 0.0 {
                    input
                } else {
                    0.0
                }
            }
        };
        Value(Rc::new(RefCell::new(ValueInner {
            data,
            grad: 0.0,
            children: vec![self.clone()],
            op: Op::Activate(act.clone()),
        })))
    }

    /// Runs reverse-mode backpropagation starting from this node.
    ///
    /// Builds a topological order of the graph reachable from `self`, seeds
    /// `self`'s gradient with `1.0`, and propagates gradients to every child by
    /// accumulating the local gradient of each operation. After it returns, each
    /// node's [`Value::grad`] holds `d(self) / d(node)`.
    ///
    /// Gradients are accumulated, not overwritten; call [`Value::zero_grad`]
    /// between independent backward passes.
    ///
    /// # Examples
    ///
    /// ```
    /// use tuphium::autograd::Value;
    ///
    /// let x = Value::new(3.0).unwrap();
    /// let y = &x * &x; // y = x^2  ->  dy/dx = 2x = 6
    /// y.backward();
    /// assert_eq!(x.grad(), 6.0);
    /// ```
    pub fn backward(&self) {
        let mut topological_order = Vec::new();
        let mut visited = HashSet::new();
        build_topo(self, &mut topological_order, &mut visited);

        self.0.borrow_mut().grad = 1.0;

        for node in topological_order.iter().rev() {
            let (op, grad, children, output_data) = {
                let inner = node.0.borrow();
                (
                    inner.op.clone(),
                    inner.grad,
                    inner.children.clone(),
                    inner.data,
                )
            };

            // How much v.data changes respect to each child? This is the local gradient
            let local_grads = match &op {
                Op::Leaf => continue,
                Op::Add => vec![1.0, 1.0],
                Op::Mul => {
                    vec![children[1].0.borrow().data, children[0].0.borrow().data]
                }
                Op::Sub => vec![1.0, -1.0],
                Op::Div => {
                    let numerator = children[0].0.borrow().data;
                    let denominator = children[1].0.borrow().data;
                    vec![1.0 / denominator, -numerator / (denominator * denominator)]
                }
                Op::Pow(exp) => {
                    let base = children[0].0.borrow().data;
                    vec![exp * base.powf(exp - 1.0)]
                }
                Op::Neg => vec![-1.0],
                Op::Exp => vec![output_data],
                Op::Log => {
                    let input = children[0].0.borrow().data;
                    vec![1.0 / input]
                }
                Op::Activate(act) => {
                    vec![derivative(output_data, act)]
                }
            };

            for (child, local_grad) in children.iter().zip(local_grads.iter()) {
                child.0.borrow_mut().grad += local_grad * grad;
            }
        }
    }
}

fn build_topo(
    node: &Value,
    topological_order: &mut Vec<Value>,
    visited: &mut HashSet<*const RefCell<ValueInner>>,
) {
    if !visited.insert(Rc::as_ptr(&node.0)) {
        return;
    }
    let children = node.0.borrow().children.clone();
    for child in &children {
        build_topo(child, topological_order, visited);
    }
    topological_order.push(node.clone());
}

fn zero_grad_recursive(node: &Value, visited: &mut HashSet<*const RefCell<ValueInner>>) {
    if !visited.insert(Rc::as_ptr(&node.0)) {
        return;
    }
    node.0.borrow_mut().grad = 0.0;
    let children = node.0.borrow().children.clone();
    for child in &children {
        zero_grad_recursive(&child, visited);
    }
}

impl Neg for Value {
    type Output = Value;

    fn neg(self) -> Value {
        Value::neg(&self)
    }
}

impl<'a> Neg for &'a Value {
    type Output = Value;

    fn neg(self) -> Value {
        Value::neg(self)
    }
}

// Generate possible binary operator implementations for Value
macro_rules! binary_op_impl {
    ($trait:ident, $method:ident) => {
        impl $trait for Value {
            type Output = Value;

            fn $method(self, other: Value) -> Value {
                Value::$method(&self, &other)
            }
        }

        impl<'a> $trait<&'a Value> for Value {
            type Output = Value;

            fn $method(self, other: &'a Value) -> Value {
                Value::$method(&self, other)
            }
        }

        impl<'a> $trait<Value> for &'a Value {
            type Output = Value;

            fn $method(self, other: Value) -> Value {
                Value::$method(self, &other)
            }
        }

        impl<'a, 'b> $trait<&'b Value> for &'a Value {
            type Output = Value;

            fn $method(self, other: &'b Value) -> Value {
                Value::$method(self, other)
            }
        }
    };
}

binary_op_impl!(Add, add);
binary_op_impl!(Mul, mul);
binary_op_impl!(Sub, sub);

impl Div for Value {
    type Output = Value;

    fn div(self, other: Value) -> Value {
        self.div_unchecked(&other)
    }
}

impl<'a> Div<&'a Value> for Value {
    type Output = Value;

    fn div(self, other: &'a Value) -> Value {
        self.div_unchecked(other)
    }
}

impl<'a> Div<Value> for &'a Value {
    type Output = Value;

    fn div(self, other: Value) -> Value {
        self.div_unchecked(&other)
    }
}

impl<'a, 'b> Div<&'b Value> for &'a Value {
    type Output = Value;

    fn div(self, other: &'b Value) -> Value {
        self.div_unchecked(other)
    }
}

impl std::fmt::Debug for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let inner = self.0.borrow();
        f.debug_struct("Value")
            .field("data", &inner.data)
            .field("grad", &inner.grad)
            .finish()
    }
}

impl std::fmt::Display for Value {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let inner = self.0.borrow();
        write!(f, "Value(data={}, grad={})", inner.data, inner.grad)
    }
}
