tensor.c

#include <math.h>

#include "tensor.h"
#include "dynamic_array.h"
#include "utils.h"

struct Tensor
{
    DynamicArray *data;
    size_t *shape, *strides;
    size_t ndim, size;
};

Tensor *tensorCreate(size_t ndim, const size_t *shape)
{
    if (ndim > 0 && !shape)
    {
        fprintf(stderr, "Error: shape pointer is NULL for non-scalar tensor\n");
        exit(1);
    }

    Tensor *t = safeMalloc(sizeof(Tensor), "Failed to allocate tensor");

    t->ndim = ndim;

    if (ndim == 0)
    {
        t->shape = NULL;
        t->strides = NULL;
        t->size = 1;
        t->data = daCreate(1);
    } else
    {
        t->shape   = safeMalloc(sizeof(size_t) * ndim, "Failed to allocate shape");
        t->strides = safeMalloc(sizeof(size_t) * ndim, "Failed to allocate strides");

        memcpy(t->shape, shape, sizeof(size_t) * ndim);

        tensorComputeStrides(ndim, shape, t->strides);

        t->size = tensorComputeSize(ndim, shape);

        t->data = daCreate(t->size);
    }

    return t;
}

Tensor *tensorCreateFromData(size_t ndim, const size_t *shape, const DynamicArray *src)
{
    daIsNullCheck(src);
    if (ndim > 0 && !shape)
    {
        fprintf(stderr, "Error: shape pointer is NULL for non-scalar tensor\n");
        exit(1);
    }

    /* ---- scalar case ---- */
    if (ndim == 0)
    {
        if (daSize(src) != 1)
        {
            fprintf(stderr, "Error: source size does not match scalar tensor\n");
            exit(1);
        }

        Tensor *t = safeMalloc(sizeof(Tensor), "Failed to allocate tensor");

        t->ndim = 0;
        t->size = 1;
        t->shape = NULL;
        t->strides = NULL;

        t->data = daCreate(1);
        daAppend(t->data, daGet(src, 0));

        return t;
    }

    /* ---- non-scalar case ---- */
    size_t expectedSize = tensorComputeSize(ndim, shape);
    if (daSize(src) != expectedSize)
    {
        fprintf(stderr, "Error: source size does not match tensor shape\n");
        exit(1);
    }

    Tensor *t = safeMalloc(sizeof(Tensor), "Failed to allocate tensor");

    t->ndim = ndim;
    t->size = expectedSize;

    t->shape   = safeMalloc(sizeof(size_t) * ndim, "Failed to allocate shape");
    t->strides = safeMalloc(sizeof(size_t) * ndim, "Failed to allocate strides");

    memcpy(t->shape, shape, sizeof(size_t) * ndim);
    tensorComputeStrides(ndim, shape, t->strides);

    t->data = daCreate(expectedSize);
    for (size_t i = 0; i < expectedSize; i++)
        daAppend(t->data, daGet(src, i));

    return t;
}

Tensor *tensorScalar(double value)
{
    Tensor *t = safeMalloc(sizeof(Tensor), "Failed to allocate scalar");

    t->ndim = 0;
    t->size = 1;
    t->shape = NULL;
    t->strides = NULL;

    t->data = daCreate(1);
    daSet(t->data, 0, value);

    return t;
}

Tensor *tensorVector(size_t n)
{
    size_t shape[1] = { n };
    return tensorCreate(1, shape);
}

Tensor *tensorVectorFromData(size_t n, const double *data)
{
    size_t shape[1] = { n };
    DynamicArray *da = daCreate(n);

    for (size_t i = 0; i < n; i++) daAppend(da, data[i]);

    Tensor *t = tensorCreateFromData(1, shape, da);
    daDestroy(da);
    return t;
}

Tensor *tensorMatrix(size_t rows, size_t cols)
{
    size_t shape[2] = { rows, cols };
    return tensorCreate(2, shape);
}

Tensor *tensorMatrixFromData(size_t rows, size_t cols, const double *data)
{
    size_t shape[2] = { rows, cols };
    DynamicArray *da = daCreate(rows * cols);

    for (size_t i = 0; i < rows * cols; i++) daAppend(da, data[i]);

    Tensor *t = tensorCreateFromData(2, shape, da);
    daDestroy(da);
    return t;
}

void tensorDestroy(Tensor *t)
{
    if (!t) return;

    daDestroy(t->data);
    free(t->shape);
    free(t->strides);
    free(t);
}

double tensorGet(const Tensor *t, const size_t *indices)
{
    tensorIsNullCheck(t);

    /* ---- scalar case ---- */
    if (t->ndim == 0)
    {
        if (indices)
        {
            fprintf(stderr, "Error: Scalar tensor accessed with indices\n");
            exit(1);
        }
        return daGet(t->data, 0);
    }

    if (!indices)
    {
        fprintf(stderr, "Error: indices pointer is NULL for non-scalar tensor\n");
        exit(1);
    }

    size_t linear = 0;

    for (size_t i = 0; i < t->ndim; i++)
    {
        if (indices[i] >= t->shape[i])
        {
            fprintf(stderr, "Error: index %zu out of bounds for dimension %zu (size %zu)\n", indices[i], i, t->shape[i]);
            exit(1);
        }

        linear += indices[i] * t->strides[i];
    }

    return daGet(t->data, linear);
}

void tensorSet(Tensor *t, const size_t *indices, double value)
{
    tensorIsNullCheck(t);

    /* ---- scalar case ---- */
    if (t->ndim == 0)
    {
        if (indices)
        {
            fprintf(stderr, "Error: Scalar tensor accessed with indices\n");
            exit(1);
        }
        daAppend(t->data, value);
        return;
    }

    if (!indices)
    {
        fprintf(stderr, "Error: indices pointer is NULL for non-scalar tensor\n");
        exit(1);
    }

    size_t linear = 0;

    for (size_t i = 0; i < t->ndim; i++)
    {
        if (indices[i] >= t->shape[i])
        {
            fprintf(stderr, "Error: index %zu out of bounds for dimension %zu (size %zu)\n", indices[i], i, t->shape[i]);
            exit(1);
        }

        linear += indices[i] * t->strides[i];
    }

    daSet(t->data, linear, value);
}

static void tensorPrintRecursive(const Tensor *t, const size_t *indices, size_t dim)
{
    if (dim == t->ndim)
    {
        printf("%g", tensorGet(t, indices));
        return;
    }

    printf("[");
    for (size_t i = 0; i < t->shape[dim]; i++)
    {
        size_t newIndices[ t->ndim ];
        if (indices)
        memcpy(newIndices, indices, sizeof(size_t) * t->ndim);
        newIndices[dim] = i;

        tensorPrintRecursive(t, newIndices, dim + 1);

        if (i + 1 < t->shape[dim]) printf(", ");
    }
    printf("]");
}

void tensorPrint(const Tensor *t)
{
    tensorIsNullCheck(t);

    if (t->ndim == 0)
    {
        printf("%g\n", tensorGet(t, NULL));
        return;
    }

    tensorPrintRecursive(t, NULL, 0);
    printf("\n");
}

static double add_op(double a, double b) { return a + b; }
static double sub_op(double a, double b) { return a - b; }
static double mul_op(double a, double b) { return a * b; }
static double div_op(double a, double b)
{
    if (b == 0.0)
    {
        fprintf(stderr, "Error: Division by zero\n");
        exit(1);
    }
    return a / b;
}

static Tensor *tensorBinaryOp(const Tensor *a, const Tensor *b, BinaryOp op)
{
    tensorIsNullCheck(a);
    tensorIsNullCheck(b);

    if (a->ndim != b->ndim)
    {
        fprintf(stderr, "Error: tensor dimensions do not match (%zu vs %zu)\n", a->ndim, b->ndim);
        exit(1);
    }

    // Scalars
    if (a->ndim == 0)
    return tensorScalar(op(tensorGet(a, NULL), tensorGet(b, NULL)));

    // Validate shapes
    for (size_t i = 0; i < a->ndim; i++)
    {
        if (a->shape[i] != b->shape[i])
        {
            fprintf(stderr, "Error: tensor shapes do not match at dimension %zu (%zu vs %zu)\n",
            i, a->shape[i], b->shape[i]);
            exit(1);
        }
    }

    Tensor *res = tensorCreate(a->ndim, a->shape);

    size_t total = a->size;
    for (size_t i = 0; i < total; i++)
    {
        double va = daGet(a->data, i);
        double vb = daGet(b->data, i);
        daAppend(res->data, op(va, vb));
    }

    return res;
}

Tensor *tensorAdd(const Tensor *a, const Tensor *b) { return tensorBinaryOp(a, b, add_op); }
Tensor *tensorSub(const Tensor *a, const Tensor *b) { return tensorBinaryOp(a, b, sub_op); }
Tensor *tensorHadamardMul(const Tensor *a, const Tensor *b) { return tensorBinaryOp(a, b, mul_op); }
Tensor *tensorHadamardDiv(const Tensor *a, const Tensor *b) { return tensorBinaryOp(a, b, div_op); }

static double add_scalar(double a, double s) { return a + s; }
static double sub_scalar(double a, double s) { return a - s; }
static double mul_scalar(double a, double s) { return a * s; }

static double div_scalar(double a, double s)
{
    if (s == 0.0)
    {
        fprintf(stderr, "Error: Division by zero\n");
        exit(1);
    }
    return a / s;
}

static Tensor *tensorScalarOp(const Tensor *t, double scalar, ScalarOp op)
{
    tensorIsNullCheck(t);

    /* ---- scalar tensor ---- */
    if (t->ndim == 0) return tensorScalar(op(tensorGet(t, NULL), scalar));

    Tensor *res = tensorCreate(t->ndim, t->shape);

    for (size_t i = 0; i < t->size; i++)
    {
        double v = daGet(t->data, i);
        daAppend(res->data, op(v, scalar));
    }

    return res;
}

Tensor *tensorAddScalar(const Tensor *t, double scalar) { return tensorScalarOp(t, scalar, add_scalar); }
Tensor *tensorSubScalar(const Tensor *t, double scalar) { return tensorScalarOp(t, scalar, sub_scalar); }
Tensor *tensorMulScalar(const Tensor *t, double scalar) { return tensorScalarOp(t, scalar, mul_scalar); }
Tensor *tensorDivScalar(const Tensor *t, double scalar) { return tensorScalarOp(t, scalar, div_scalar); }

Tensor *tensorApply(const Tensor *t, UnaryOp op)
{
    tensorIsNullCheck(t);

    if (!op)
    {
        fprintf(stderr, "Error: function pointer is NULL\n");
        exit(1);
    }

    /* ---- scalar tensor ---- */
    if (t->ndim == 0) return tensorScalar(op(tensorGet(t, NULL)));

    Tensor *res = tensorCreate(t->ndim, t->shape);

    for (size_t i = 0; i < t->size; i++)
    {
        double v = daGet(t->data, i);
        daAppend(res->data, op(v));
    }

    return res;
}

double tensorReduce(const Tensor *t, ReduceOp op, double init)
{
    tensorIsNullCheck(t);

    if (!op)
    {
        fprintf(stderr, "Error: reduce operator is NULL\n");
        exit(1);
    }

    /* ---- scalar ---- */
    if (t->ndim == 0)
        return op(init, tensorGet(t, NULL));

    double acc = init;

    for (size_t i = 0; i < t->size; i++)
        acc = op(acc, daGet(t->data, i));

    return acc;
}

static double sum_op(double a, double b)  { return a + b; }
static double prod_op(double a, double b) { return a * b; }
static double square(double x) { return x * x; }
static double min_op(double a, double b) { return a < b ? a : b; }
static double max_op(double a, double b) { return a > b ? a : b; }
static double and_op(double a, double b) { return (a != 0.0) && (b != 0.0); }
static double or_op(double a, double b) { return (a != 0.0) || (b != 0.0); }
static double neg(double x) { return -x; }
static double abs_val(double x) { return fabs(x); }
static double exp_func(double x) { return exp(x); }
static double log_func(double x) { return log(x); }

double tensorSum(const Tensor *t) { return tensorReduce(t, sum_op, 0.0); }
double tensorProd(const Tensor *t) { return tensorReduce(t, prod_op, 1.0); }
double tensorMin(const Tensor *t) { return tensorReduce(t, min_op, daGet(t->data, 0)); }
double tensorMax(const Tensor *t) { return tensorReduce(t, max_op, daGet(t->data, 0)); }
double tensorMean(const Tensor *t) { return tensorSum(t) / t->size; }
double tensorAll(const Tensor *t) { return tensorReduce(t, and_op, 1.0); }
double tensorAny(const Tensor *t) { return tensorReduce(t, or_op, 0.0); }
Tensor *tensorNeg(const Tensor *t) { return tensorApply(t, neg); }
Tensor *tensorAbs(const Tensor *t) { return tensorApply(t, abs_val); }
Tensor *tensorExp(const Tensor *t) { return tensorApply(t, exp_func); }
Tensor *tensorLog(const Tensor *t) { return tensorApply(t, log_func); }

size_t tensorArgMin(const Tensor *t)
{
    tensorIsNullCheck(t);

    size_t idx = 0;
    double best = daGet(t->data, 0);

    for (size_t i = 1; i < t->size; i++)
    {
        double v = daGet(t->data, i);
        if (v < best)
        {
            best = v;
            idx = i;
        }
    }

    return idx;
}

size_t tensorArgMax(const Tensor *t)
{
    tensorIsNullCheck(t);

    size_t idx = 0;
    double best = daGet(t->data, 0);

    for (size_t i = 1; i < t->size; i++)
    {
        double v = daGet(t->data, i);
        if (v > best)
        {
            best = v;
            idx = i;
        }
    }

    return idx;
}

double tensorL2Norm(const Tensor *t)
{
    Tensor *tmp = tensorApply(t, square);
    double sum = tensorSum(tmp);
    tensorDestroy(tmp);
    return sqrt(sum);
}

Tensor *tensorMul(const Tensor *a, const Tensor *b)
{
    tensorIsNullCheck(a);
    tensorIsNullCheck(b);

    /* scalar cases */
    if (a->ndim == 0)
        return tensorScalarOp(b, daGet(a->data,0), mul_scalar);

    if (b->ndim == 0)
        return tensorScalarOp(a, daGet(b->data,0), mul_scalar);

    /* vector dot */
    if (a->ndim == 1 && b->ndim == 1)
        return tensorDot(a, b);

    /* matrix */
    if (a->ndim == 2 && b->ndim == 2)
        return tensorMatMul(a, b);

    /* batched */
    return tensorBatchedMatMul(a, b);
}

Tensor *tensorMatMul(const Tensor *a, const Tensor *b)
{
    tensorIsNullCheck(a);
    tensorIsNullCheck(b);
    
    if (a->ndim != 2 || b->ndim != 2)
    {
        fprintf(stderr, "Error: matmul requires 2D tensors\n");
        exit(1);
    }
    
    size_t m = a->shape[0];
    size_t k1 = a->shape[1];
    size_t k2 = b->shape[0];
    size_t n = b->shape[1];
    
    if (k1 != k2)
    {
        fprintf(stderr, "Error: incompatible shapes for matmul\n");
        exit(1);
    }
    
    size_t result_shape[2] = {m, n};
    Tensor *res = tensorCreate(2, result_shape);
    
    for (size_t i = 0; i < m; i++)
    {
        for (size_t j = 0; j < n; j++)
        {
            double sum = 0.0;
            for (size_t k = 0; k < k1; k++)
            {
                size_t ai = i * k1 + k;
                size_t bi = k * n + j;
                sum += daGet(a->data, ai) * daGet(b->data, bi);
            }
            daAppend(res->data, sum);
        }
    }
    
    return res;
}

Tensor *tensorDot(const Tensor *a, const Tensor *b)
{
    tensorIsNullCheck(a);
    tensorIsNullCheck(b);

    if (a->shape[0] != b->shape[0])
    {
        fprintf(stderr, "Error: dot product size mismatch\n");
        exit(1);
    }

    double sum = 0.0;

    for (size_t i = 0; i < a->shape[0]; i++)
        sum += daGet(a->data,i) * daGet(b->data,i);

    return tensorScalar(sum);
}

static size_t tensorBatchCount(const Tensor *t)
{
    size_t count = 1;
    for (size_t i = 0; i < t->ndim-2; i++)
        count *= t->shape[i];
    return count;
}

Tensor *tensorBatchedMatMul(const Tensor *a, const Tensor *b)
{
    tensorIsNullCheck(a);
    tensorIsNullCheck(b);

    if (a->ndim != b->ndim || a->ndim < 3)
    {
        fprintf(stderr, "Error: invalid ranks for batched matmul\n");
        exit(1);
    }

    size_t ndim = a->ndim;

    for (size_t i = 0; i < ndim-2; i++)
        if (a->shape[i] != b->shape[i])
        {
            fprintf(stderr, "Error: batch dimension mismatch\n");
            exit(1);
        }

    size_t m = a->shape[ndim-2];
    size_t k1 = a->shape[ndim-1];
    size_t k2 = b->shape[ndim-2];
    size_t n = b->shape[ndim-1];

    if (k1 != k2)
    {
        fprintf(stderr, "Error: inner dimension mismatch\n");
        exit(1);
    }

    /* build output shape */
    size_t *out_shape = safeMalloc(sizeof(size_t)*ndim, "shape alloc");

    for (size_t i = 0; i < ndim-2; i++)
        out_shape[i] = a->shape[i];

    out_shape[ndim-2] = m;
    out_shape[ndim-1] = n;

    Tensor *res = tensorCreate(ndim, out_shape);
    free(out_shape);

    size_t batch = tensorBatchCount(a);

    size_t Ablock = m * k1;
    size_t Bblock = k1 * n;

    for (size_t bidx = 0; bidx < batch; bidx++)
    {
        size_t aBase = bidx * Ablock;
        size_t bBase = bidx * Bblock;

        for (size_t i = 0; i < m; i++)
            for (size_t j = 0; j < n; j++)
            {
                double sum = 0;

                for (size_t p = 0; p < k1; p++)
                {
                    size_t ai = aBase + i*k1 + p;
                    size_t bi = bBase + p*n + j;

                    sum += daGet(a->data, ai) *
                           daGet(b->data, bi);
                }

                daAppend(res->data, sum);
            }
    }

    return res;
}

Tensor *tensorTranspose(const Tensor *t, size_t axis1, size_t axis2)
{
    tensorIsNullCheck(t);

    if (axis1 >= t->ndim || axis2 >= t->ndim)
    {
        fprintf(stderr, "Error: invalid transpose axes\n");
        exit(1);
    }

    if (axis1 == axis2)
        return tensorCreateFromData(t->ndim, t->shape, t->data);

    /* build new shape */
    size_t *newShape = safeMalloc(sizeof(size_t)*t->ndim, "shape alloc");

    for (size_t i = 0; i < t->ndim; i++) newShape[i] = t->shape[i];

    size_t tmp = newShape[axis1];
    newShape[axis1] = newShape[axis2];
    newShape[axis2] = tmp;

    Tensor *res = tensorCreate(t->ndim, newShape);
    free(newShape);

    size_t *idx = safeMalloc(sizeof(size_t)*t->ndim, "index buffer");
    size_t *swapped = safeMalloc(sizeof(size_t)*t->ndim, "index buffer");

    double *temp = safeMalloc(sizeof(double) * t->size, "temp array");

    for (size_t i = 0; i < t->size; i++)
    {
        linearToIndices(i, t->shape, t->strides, t->ndim, idx);
        
        for (size_t k = 0; k < t->ndim; k++) swapped[k] = idx[k];
        tmp = swapped[axis1];
        swapped[axis1] = swapped[axis2];
        swapped[axis2] = tmp;
        
        size_t outLinear = indicesToLinear(swapped, res->strides, res->ndim);
        temp[outLinear] = daGet(t->data, i);
    }

    for (size_t i = 0; i < t->size; i++) {
        daAppend(res->data, temp[i]);
    }

    free(temp);

    free(idx);
    free(swapped);

    return res;
}

Tensor *tensorTranspose2D(const Tensor *t)
{
    if (t->ndim != 2)
    {
        fprintf(stderr, "Error: not a matrix\n");
        exit(1);
    }

    return tensorTranspose(t, 0, 1);
}

Tensor *tensorReshape(const Tensor *t, size_t new_ndim, const size_t *new_shape)
{
    tensorIsNullCheck(t);

    if (new_ndim > 0 && !new_shape)
    {
        fprintf(stderr, "Error: new shape is NULL\n");
        exit(1);
    }

    size_t new_size = tensorComputeSize(new_ndim, new_shape);

    if (new_size != t->size)
    {
        fprintf(stderr, "Error: reshape changes total size (%zu -> %zu)\n", t->size, new_size);
        exit(1);
    }

    Tensor *res = tensorCreate(new_ndim, new_shape);

    for (size_t i = 0; i < t->size; i++) daAppend(res->data, daGet(t->data, i));

    return res;
}

Tensor *tensorReshapeInfer(const Tensor *t, size_t new_ndim, const size_t *shape)
{
    tensorIsNullCheck(t);

    if (new_ndim > 0 && !shape)
    {
        fprintf(stderr, "Error: new shape is NULL\n");
        exit(1);
    }

    size_t infer_index = SIZE_MAX;
    size_t known_product = 1;

    for (size_t i = 0; i < new_ndim; i++)
    {
        if (shape[i] == SIZE_MAX)
        {
            if (infer_index != SIZE_MAX)
            {
                fprintf(stderr, "Error: multiple inferred dimensions\n");
                exit(1);
            }
            infer_index = i;
        }
        else
        {
            known_product *= shape[i];
        }
    }

    size_t *final_shape = safeMalloc(sizeof(size_t)*new_ndim, "shape alloc");

    memcpy(final_shape, shape, sizeof(size_t)*new_ndim);

    if (infer_index != SIZE_MAX)
    {
        if (known_product == 0 || t->size % known_product != 0)
        {
            fprintf(stderr, "Error: cannot infer reshape dimension\n");
            exit(1);
        }

        final_shape[infer_index] = t->size / known_product;
    }

    size_t new_size = tensorComputeSize(new_ndim, final_shape);

    if (new_size != t->size)
    {
        fprintf(stderr, "Error: reshape changes total size (%zu -> %zu)\n", t->size, new_size);
        exit(1);
    }

    Tensor *res = tensorCreate(new_ndim, final_shape);

    for (size_t i = 0; i < t->size; i++) daAppend(res->data, daGet(t->data, i));

    free(final_shape);
    return res;
}