#include "tensor.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>
#include <time.h>

/* ─────────────────────────────────────────
   LIFECYCLE
   ───────────────────────────────────────── */

Tensor tensor_create(int rows, int cols) {
    Tensor t;
    t.rows = rows;
    t.cols = cols;
    t.data = (float*)malloc(rows * cols * sizeof(float));
    if (!t.data) {
        fprintf(stderr, "[TENSOR] malloc failed for %dx%d\n", rows, cols);
        exit(1);
    }
    return t;
}

Tensor tensor_zeros(int rows, int cols) {
    Tensor t = tensor_create(rows, cols);
    memset(t.data, 0, rows * cols * sizeof(float));
    return t;
}

Tensor tensor_copy(const Tensor* t) {
    Tensor c = tensor_create(t->rows, t->cols);
    memcpy(c.data, t->data, t->rows * t->cols * sizeof(float));
    return c;
}

void tensor_free(Tensor* t) {
    if (t->data) {
        free(t->data);
        t->data = NULL;
    }
    t->rows = 0;
    t->cols = 0;
}

/* ─────────────────────────────────────────
   RANDOM — Box-Muller gaussian
   ───────────────────────────────────────── */

static float randn_single(void) {
    static int has_spare = 0;
    static float spare;
    if (has_spare) { has_spare = 0; return spare; }
    float u, v, s;
    do {
        u = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;
        v = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;
        s = u * u + v * v;
    } while (s >= 1.0f || s == 0.0f);
    float mul = sqrtf(-2.0f * logf(s) / s);
    spare = v * mul;
    has_spare = 1;
    return u * mul;
}

void tensor_fill(Tensor* t, float val) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) t->data[i] = val;
}

void tensor_randn(Tensor* t, float std) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) t->data[i] = randn_single() * std;
}

void tensor_xavier(Tensor* t) {
    float std = sqrtf(2.0f / (float)(t->rows + t->cols));
    tensor_randn(t, std);
}

/* ─────────────────────────────────────────
   ELEMENTWISE OPS
   ───────────────────────────────────────── */

void tensor_add(Tensor* dst, const Tensor* a, const Tensor* b) {
    int n = a->rows * a->cols;
    for (int i = 0; i < n; i++) dst->data[i] = a->data[i] + b->data[i];
}

void tensor_add_inplace(Tensor* a, const Tensor* b) {
    int n = a->rows * a->cols;
    for (int i = 0; i < n; i++) a->data[i] += b->data[i];
}

void tensor_sub(Tensor* dst, const Tensor* a, const Tensor* b) {
    int n = a->rows * a->cols;
    for (int i = 0; i < n; i++) dst->data[i] = a->data[i] - b->data[i];
}

void tensor_mul_scalar(Tensor* t, float s) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) t->data[i] *= s;
}

void tensor_add_scalar(Tensor* t, float s) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) t->data[i] += s;
}

void tensor_elementwise_mul(Tensor* dst, const Tensor* a, const Tensor* b) {
    int n = a->rows * a->cols;
    for (int i = 0; i < n; i++) dst->data[i] = a->data[i] * b->data[i];
}

/* ─────────────────────────────────────────
   MATRIX OPS
   ───────────────────────────────────────── */

/* dst = a @ b   (a: MxK, b: KxN, dst: MxN) */
void tensor_matmul(Tensor* dst, const Tensor* a, const Tensor* b) {
    int M = a->rows, K = a->cols, N = b->cols;
    memset(dst->data, 0, M * N * sizeof(float));
    for (int i = 0; i < M; i++) {
        for (int k = 0; k < K; k++) {
            float aik = a->data[i * K + k];
            for (int j = 0; j < N; j++) {
                dst->data[i * N + j] += aik * b->data[k * N + j];
            }
        }
    }
}

/* dst = a @ b.T   (a: MxK, b: NxK, dst: MxN) */
void tensor_matmul_transB(Tensor* dst, const Tensor* a, const Tensor* b) {
    int M = a->rows, K = a->cols, N = b->rows;
    memset(dst->data, 0, M * N * sizeof(float));
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += a->data[i * K + k] * b->data[j * K + k];
            }
            dst->data[i * N + j] = sum;
        }
    }
}

void tensor_transpose(Tensor* dst, const Tensor* a) {
    for (int i = 0; i < a->rows; i++)
        for (int j = 0; j < a->cols; j++)
            dst->data[j * a->rows + i] = a->data[i * a->cols + j];
}

/* ─────────────────────────────────────────
   ACTIVATIONS
   ───────────────────────────────────────── */

void tensor_relu(Tensor* dst, const Tensor* src) {
    int n = src->rows * src->cols;
    for (int i = 0; i < n; i++)
        dst->data[i] = src->data[i] > 0.0f ? src->data[i] : 0.0f;
}

void tensor_relu_inplace(Tensor* t) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++)
        if (t->data[i] < 0.0f) t->data[i] = 0.0f;
}

/* GELU approximation: x * 0.5 * (1 + tanh(sqrt(2/pi)*(x + 0.044715*x^3))) */
void tensor_gelu_inplace(Tensor* t) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) {
        float x = t->data[i];
        float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
        t->data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
}

void tensor_gelu(Tensor* dst, const Tensor* src) {
    int n = src->rows * src->cols;
    for (int i = 0; i < n; i++) {
        float x = src->data[i];
        float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
        dst->data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
}

void tensor_softmax_rows(Tensor* t) {
    for (int i = 0; i < t->rows; i++) {
        float* row = t->data + i * t->cols;
        float max_val = row[0];
        for (int j = 1; j < t->cols; j++)
            if (row[j] > max_val) max_val = row[j];
        float sum = 0.0f;
        for (int j = 0; j < t->cols; j++) {
            row[j] = expf(row[j] - max_val);
            sum += row[j];
        }
        for (int j = 0; j < t->cols; j++) row[j] /= sum;
    }
}

void tensor_softmax_inplace(Tensor* t, int len) {
    float max_val = t->data[0];
    for (int i = 1; i < len; i++)
        if (t->data[i] > max_val) max_val = t->data[i];
    float sum = 0.0f;
    for (int i = 0; i < len; i++) {
        t->data[i] = expf(t->data[i] - max_val);
        sum += t->data[i];
    }
    for (int i = 0; i < len; i++) t->data[i] /= sum;
}

/* ─────────────────────────────────────────
   LAYER NORM
   ───────────────────────────────────────── */

void tensor_layernorm(Tensor* dst, const Tensor* src,
                      const Tensor* w, const Tensor* b,
                      int row, float eps) {
    int D = src->cols;
    float* x   = src->data + row * D;
    float* out = dst->data + row * D;

    /* mean */
    float mean = 0.0f;
    for (int i = 0; i < D; i++) mean += x[i];
    mean /= D;

    /* variance */
    float var = 0.0f;
    for (int i = 0; i < D; i++) {
        float diff = x[i] - mean;
        var += diff * diff;
    }
    var /= D;

    float inv_std = 1.0f / sqrtf(var + eps);
    for (int i = 0; i < D; i++)
        out[i] = (x[i] - mean) * inv_std * w->data[i] + b->data[i];
}

/* ─────────────────────────────────────────
   HELPERS
   ───────────────────────────────────────── */

float tensor_get(const Tensor* t, int row, int col) {
    return t->data[row * t->cols + col];
}

void tensor_set(Tensor* t, int row, int col, float val) {
    t->data[row * t->cols + col] = val;
}

void tensor_add_to(Tensor* t, int row, int col, float val) {
    t->data[row * t->cols + col] += val;
}

void tensor_embed(Tensor* dst, const Tensor* embed,
                  const int* token_ids, int n_tokens) {
    int D = embed->cols;
    for (int i = 0; i < n_tokens; i++) {
        int id = token_ids[i];
        memcpy(dst->data + i * D,
               embed->data + id * D,
               D * sizeof(float));
    }
}

void tensor_add_row(Tensor* a, int a_row, const Tensor* b, int b_row) {
    int D = a->cols;
    float* ap = a->data + a_row * D;
    float* bp = b->data + b_row * D;
    for (int i = 0; i < D; i++) ap[i] += bp[i];
}

void tensor_clip(Tensor* t, float min_val, float max_val) {
    int n = t->rows * t->cols;
    for (int i = 0; i < n; i++) {
        if (t->data[i] < min_val) t->data[i] = min_val;
        if (t->data[i] > max_val) t->data[i] = max_val;
    }
}

void tensor_print(const Tensor* t, const char* name, int max_rows) {
    int rows = t->rows < max_rows ? t->rows : max_rows;
    printf("[%s] shape=(%d,%d)\n", name, t->rows, t->cols);
    for (int i = 0; i < rows; i++) {
        printf("  row%d: ", i);
        int cols = t->cols < 8 ? t->cols : 8;
        for (int j = 0; j < cols; j++)
            printf("%.4f ", t->data[i * t->cols + j]);
        if (t->cols > 8) printf("...");
        printf("\n");
    }
}