#include "tensor.h"
#include <stdlib.h>

Function* create_add_backward() {
    Function* fn = malloc(sizeof(Function));
    fn->backward = add_backward;
    fn->num_inputs = 2;
    fn->saved_tensors = NULL;
    fn->num_saved = 0;
    return fn;
}

Tensor** add_backward(Function* self, Tensor* grad_output){
    Tensor** grads = malloc(sizeof(Tensor*) * 2);
    grads[0] = grad_output;
    grads[1] = grad_output;
    return grads;
}

Function* create_sub_backward() {
    Function* fn = malloc(sizeof(Function));
    fn->backward = sub_backward;
    fn->num_inputs = 2;
    fn->saved_tensors = NULL;
    fn->num_saved = 0;
    return fn;
}

Tensor** sub_backward(Function* self, Tensor* grad_output){
    Tensor** grads = malloc(sizeof(Tensor*) * 2);
    grads[0] = grad_output;
    // For subtraction, second gradient is negated
    grads[1] = tensor_scalar_op(grad_output, -1.0, TENSOR_MUL);
    return grads;
}

Function* create_mul_backward(Tensor* a, Tensor* b) {
    Function* fn = malloc(sizeof(Function));
    fn->backward = mul_backward;
    fn->num_inputs = 2;
    fn->saved_tensors = malloc(sizeof(Tensor*) * 2);
    fn->saved_tensors[0] = a;
    fn->saved_tensors[1] = b;
    fn->num_saved = 2;
    return fn;
}

Tensor** mul_backward(Function* self, Tensor* grad_output) {
    Tensor* a = self->saved_tensors[0];
    Tensor* b = self->saved_tensors[1];

    // grad_a = grad_output * b
    Tensor* grad_a = tensor_op(grad_output, b, TENSOR_MUL);
    
    // grad_b = grad_output * a
    Tensor* grad_b = tensor_op(grad_output, a, TENSOR_MUL);

    Tensor** grads = malloc(2 * sizeof(Tensor*));
    grads[0] = grad_a;
    grads[1] = grad_b;
    return grads;
}

Function* create_matmul_backward(Tensor* a, Tensor* b) {
    Function* fn = malloc(sizeof(Function));
    fn->backward = matmul_backward;
    fn->num_inputs = 2;
    fn->saved_tensors = malloc(sizeof(Tensor*) * 2);
    fn->saved_tensors[0] = a;
    fn->saved_tensors[1] = b;
    fn->num_saved = 2;
    return fn;
}

Tensor** matmul_backward(Function* self, Tensor* grad_output) {
    Tensor* a = self->saved_tensors[0];
    Tensor* b = self->saved_tensors[1];
    
    // grad_a = grad_output @ b.T
    Tensor* b_T = tensor_transpose(b, -2, -1);
    Tensor* grad_a = tensor_matmul(grad_output, b_T);
    
    // grad_b = a.T @ grad_output
    Tensor* a_T = tensor_transpose(a, -2, -1);
    Tensor* grad_b = tensor_matmul(a_T, grad_output);
    
    Tensor** grads = malloc(2 * sizeof(Tensor*));
    grads[0] = grad_a;
    grads[1] = grad_b;
    return grads;
}

int is_visited(Tensor* t, Tensor** visited, int visited_size) {
    for (int i = 0; i < visited_size; i++) {
        if (visited[i] == t) return 1;
    }
    return 0;
}

void build_topo(Tensor* t, Tensor** topo, int* topo_size, Tensor** visited, int* visited_size) {
    // If already visited, skip
    if (is_visited(t, visited, *visited_size)) {
        return;
    }
    
    // Mark as visited
    visited[*visited_size] = t;
    (*visited_size)++;
    
    // Recursively visit all inputs
    for (int i = 0; i < t->num_inputs; i++) {
        build_topo(t->inputs[i], topo, topo_size, visited, visited_size);
    }
    
    // Add to topo list (post-order)
    topo[*topo_size] = t;
    (*topo_size)++;
}

void tensor_backward(Tensor* t) {
    // Allocate arrays (assume max 1000 tensors for now)
    Tensor** topo = malloc(1000 * sizeof(Tensor*));
    Tensor** visited = malloc(1000 * sizeof(Tensor*));
    int topo_size = 0;
    int visited_size = 0;
    
    // Build topological order
    build_topo(t, topo, &topo_size, visited, &visited_size);
    
    // Initialize gradient (dL/dL = 1)
    t->grad = tensor_ones(t->shape, t->ndim, t->dtype, t->device);
    
    // Backward pass
    for (int i = topo_size - 1; i >= 0; i--) {
        Tensor* node = topo[i];
        if (node->grad_fn) {
            Tensor** input_grads = node->grad_fn->backward(node->grad_fn, node->grad);
            
            for (int j = 0; j < node->num_inputs; j++) {
                // Accumulate gradient
                if (node->inputs[j]->grad == NULL) {
                    node->inputs[j]->grad = input_grads[j];
                } else {
                    // Add to existing gradient
                    Tensor* new_grad = tensor_op(node->inputs[j]->grad, input_grads[j], TENSOR_ADD);
                    tensor_free(node->inputs[j]->grad);
                    node->inputs[j]->grad = new_grad;
                }
            }
            free(input_grads);
        }
    }
    
    free(visited);
    free(topo);
}

void function_free(Function* fn) {
    if (!fn) return;
    
    if (fn->saved_tensors) {
        free(fn->saved_tensors);
    }
    
    if (fn->metadata) {
        free(fn->metadata);
    }
    
    free(fn);
}