#ifndef TENSOR_H
#define TENSOR_H

#include <stddef.h>
#include <stdint.h>

/* ─────────────────────────────────────────
   TENSOR  —  the core math object
   Everything in JOY is built on this.
   ───────────────────────────────────────── */

typedef struct {
    float*  data;
    int     rows;
    int     cols;
} Tensor;

/* lifecycle */
Tensor  tensor_create(int rows, int cols);
Tensor  tensor_zeros(int rows, int cols);
Tensor  tensor_copy(const Tensor* t);
void    tensor_free(Tensor* t);

/* init */
void    tensor_fill(Tensor* t, float val);
void    tensor_randn(Tensor* t, float std);   /* gaussian random */
void    tensor_xavier(Tensor* t);             /* xavier init for weights */

/* elementwise */
void    tensor_add(Tensor* dst, const Tensor* a, const Tensor* b);
void    tensor_add_inplace(Tensor* a, const Tensor* b);
void    tensor_sub(Tensor* dst, const Tensor* a, const Tensor* b);
void    tensor_mul_scalar(Tensor* t, float s);
void    tensor_add_scalar(Tensor* t, float s);
void    tensor_elementwise_mul(Tensor* dst, const Tensor* a, const Tensor* b);

/* matrix ops */
void    tensor_matmul(Tensor* dst, const Tensor* a, const Tensor* b);
void    tensor_matmul_transB(Tensor* dst, const Tensor* a, const Tensor* b);
void    tensor_transpose(Tensor* dst, const Tensor* a);

/* activations */
void    tensor_relu(Tensor* dst, const Tensor* src);
void    tensor_relu_inplace(Tensor* t);
void    tensor_gelu(Tensor* dst, const Tensor* src);
void    tensor_gelu_inplace(Tensor* t);
void    tensor_softmax_rows(Tensor* t);       /* softmax each row */
void    tensor_softmax_inplace(Tensor* t, int len); /* softmax over len elements */

/* normalization */
void    tensor_layernorm(Tensor* dst, const Tensor* src,
                         const Tensor* w, const Tensor* b,
                         int row, float eps);

/* reduction */
float   tensor_get(const Tensor* t, int row, int col);
void    tensor_set(Tensor* t, int row, int col, float val);
void    tensor_add_to(Tensor* t, int row, int col, float val);

/* embedding lookup: dst[i] = embed[token_ids[i]] */
void    tensor_embed(Tensor* dst, const Tensor* embed,
                     const int* token_ids, int n_tokens);

/* add row of b into row of a */
void    tensor_add_row(Tensor* a, int a_row, const Tensor* b, int b_row);

/* clip gradients */
void    tensor_clip(Tensor* t, float min_val, float max_val);

/* debug */
void    tensor_print(const Tensor* t, const char* name, int max_rows);

#endif /* TENSOR_H */