#ifndef AUTOGRAD_H
#define AUTOGRAD_H

#include <stddef.h>

// Forward declarations
typedef struct Tensor Tensor;
typedef struct Function Function;

struct Function {
    // Function pointer for backward pass
    // Takes: output_grad, returns: array of input_grads
    Tensor** (*backward)(struct Function* self, Tensor* output_grad);
    
    // Saved tensors needed for backward computation
    Tensor** saved_tensors;
    int num_saved;

    // Number of inputs this operation had
    int num_inputs;
    
    // Optional: metadata for the operation (e.g., transpose dims)
    void* metadata;
};

// Function to trigger backward pass
void tensor_backward(Tensor* t);

// Helper to free a Function
void function_free(Function* fn);

Function* create_add_backward();
Function* create_sub_backward();
Function* create_mul_backward(Tensor* a, Tensor* b);
Function* create_matmul_backward(Tensor* a, Tensor* b);

Tensor** add_backward(Function* self, Tensor* grad_output);
Tensor** sub_backward(Function* self, Tensor* grad_output);
Tensor** mul_backward(Function* self, Tensor* grad_output);
Tensor** matmul_backward(Function* self, Tensor* grad_output);


#endif
