// ann.h - Artificial Neural Network Library
//
// Thread safety: Not thread-safe. External synchronization required.
// Memory: Single-allocation design — all neurons, weights, and deltas
//   stored in one contiguous buffer for easy persistence and cloning.
//   Use ann_copy for deep clone; ann_free to release.

#ifndef ANN_H
#define ANN_H

#include "type.h"

#ifdef __cplusplus
extern "C" {
#endif

typedef enum {
  IDENTITY,
  BINARY,
  SIGMOID,
  RELU,
  ELU,
  LRELU,
  TANH,
} ann_activation_t;

typedef struct {
  void *buffer;

  // The full size of the allocated structure
  u64 n;

  // The number of layers in the neural network
  u64 layer_n;

  // The neuron count ( hidden )
  u64 neuron_n;

  // The total number of weights and biases
  u64 weight_n;

  // The number of neurons in each layer
  u64 *layer_neuron_n;

  // The pointer to the neurons
  f64 *neuron;

  // The weights and biases for each neuron
  //   - w_lji denotes the connection between the jth neuron in layer l, and
  //     the ith neuron in the previous layer
  //   - b_lj denotes the bias for the jth neuron in layer l
  // { w_000, w_001, w_002, .... , b_00, ...., w_ijk, b_ij }`
  f64 *weight;

  // The delta between between the actual and the cost function
  f64 *delta;

  // The activation function used in the hidden layer neurons
  f64 (*activation_hidden)(f64);

  // The partial derivative of the activation function used in the hidden
  // layer neurons
  f64 (*activation_hidden_partial)(f64);

  // The activation function used in the ouput neurons
  f64 (*activation_output)(f64);

  // The partial derivative of the activation function used in the output
  // layer neurons
  f64 (*activation_output_partial)(f64);
} ann_t;

int ann_init(ann_t *, u64, u64 *);
void ann_copy(ann_t *, ann_t const *);
void ann_free(ann_t *);
void ann_random(ann_t *);

void ann_propagation_forward(ann_t *, f64 const *const, f64 *);
void ann_propagation_backward(ann_t *, f64 const *, f64 *, f64 const *,
                              f64);
void ann_train_numeric(ann_t *, f64 const *, f64 const *, u64, f64);

f64 ann_error_total(f64 const *, f64 const *, u64);
void ann_set_activation(ann_t *, ann_activation_t, ann_activation_t);

void ann_print_weight(ann_t *);
void ann_print_neuron(ann_t *, f64 const *const, f64 const *const);

#ifdef __cplusplus
}
#endif

#endif // ANN_H

////////////////////////////////////////////////////////////////////////////////
// IMPLEMENTATION SECTION
////////////////////////////////////////////////////////////////////////////////

#ifdef ANN_IMPLEMENTATION

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tgmath.h>

#define PRINT_PRECISION 10

#define SQUARE_ROOT_2 1.4142135623730950488016887242096
#define PI 3.1415926535897932384626433832795
#define SQUARE_ROOT_PI 1.7724538509055160272981674833411
#define SQUARE_ROOT_2_OVER_PI 0.7978845608028653558798921198687

#define ELU_ALPHA 0.2
#define LRELU_ALPHA 0.2

static void ann_zero_weights(ann_t *);
static f64 ann_random_range(f64, f64);
static f64 ann_error(f64, f64);
static f64 ann_error_partial(f64, f64);

static f64 ann_activation_identity(f64);
static f64 ann_activation_identity_partial(f64);
static f64 ann_activation_binary(f64);
static f64 ann_activation_binary_partial(f64);
static f64 ann_activation_sigmoid(f64);
static f64 ann_activation_sigmoid_partial(f64);
static f64 ann_activation_relu(f64);
static f64 ann_activation_relu_partial(f64);
static f64 ann_activation_elu(f64);
static f64 ann_activation_elu_partial(f64);
static f64 ann_activation_lrelu(f64);
static f64 ann_activation_lrelu_partial(f64);
static f64 ann_activation_tanh(f64);
static f64 ann_activation_tanh_partial(f64);

static f64 (*ACTIVATION[][2])(f64) = {
    {ann_activation_identity, ann_activation_identity_partial}, // IDENTITY
    {ann_activation_binary, ann_activation_binary_partial},     // BINARY
    {ann_activation_sigmoid, ann_activation_sigmoid_partial},   // SIGMOID
    {ann_activation_relu, ann_activation_relu_partial},         // RELU
    {ann_activation_elu, ann_activation_elu_partial},           // ELU
    {ann_activation_lrelu, ann_activation_lrelu_partial},       // LRELU
    {ann_activation_tanh, ann_activation_tanh_partial},         // TANH
};

int ann_init(ann_t *ann, u64 layer_n, u64 *layer_neuron_n) {
  if (!ann || !layer_neuron_n || layer_n < 2)
    return 0;

  u64 neuron_n = 0;
  u64 weight_n = 0;

  u64 l = 1;
  for (; l < layer_n - 1; l++) {
    if (neuron_n > UINT64_MAX - layer_neuron_n[l])
      return 0;
    neuron_n += layer_neuron_n[l];
    if (layer_neuron_n[l - 1] == UINT64_MAX)
      return 0;
    {
      u64 rhs = layer_neuron_n[l - 1] + 1;
      if (layer_neuron_n[l] > UINT64_MAX / rhs)
        return 0;
      u64 prod = layer_neuron_n[l] * rhs;
      if (weight_n > UINT64_MAX - prod)
        return 0;
      weight_n += prod;
    }
  }

  if (layer_neuron_n[l - 1] == UINT64_MAX)
    return 0;
  {
    u64 rhs = layer_neuron_n[l - 1] + 1;
    if (layer_neuron_n[l] > UINT64_MAX / rhs)
      return 0;
    u64 prod = layer_neuron_n[l] * rhs;
    if (weight_n > UINT64_MAX - prod)
      return 0;
    weight_n += prod;
  }

  if (layer_n > SIZE_MAX / sizeof(u64))
    return 0;
  size_t layer_neuron_sz = sizeof(u64) * layer_n;
  size_t padding = (sizeof(f64) - (layer_neuron_sz % sizeof(f64))) % sizeof(f64);

  u64 total_fp = neuron_n + weight_n;
  if (total_fp < neuron_n || total_fp < weight_n)
    return 0;
  if (total_fp > UINT64_MAX - neuron_n)
    return 0;
  total_fp += neuron_n;
  if (total_fp > UINT64_MAX - layer_neuron_n[layer_n - 1])
    return 0;
  total_fp += layer_neuron_n[layer_n - 1];
  if (total_fp > SIZE_MAX / sizeof(f64))
    return 0;
  u64 n = (u64)(layer_neuron_sz + padding + sizeof(f64) * total_fp);

  ann->buffer = malloc(n);
  if (!ann->buffer)
    return 0;

  ann->n = n;
  ann->layer_n = layer_n;
  ann->weight_n = weight_n;
  ann->neuron_n = neuron_n;
  ann->layer_neuron_n = (u64 *)ann->buffer;
  memcpy(ann->layer_neuron_n, layer_neuron_n, sizeof(u64) * layer_n);
  ann->neuron = (f64 *)((char *)ann->buffer + layer_neuron_sz + padding);
  ann->weight = ann->neuron + ann->neuron_n;
  ann->delta = ann->weight + ann->weight_n;

  ann_set_activation(ann, SIGMOID, SIGMOID);
  ann_zero_weights(ann);
  return 1;
}

void ann_free(ann_t *ann) {
  if (ann) {
    free(ann->buffer);
    ann->buffer = NULL;
    ann->layer_neuron_n = NULL;
    ann->neuron = NULL;
    ann->weight = NULL;
    ann->delta = NULL;
  }
}

void ann_copy(ann_t *dest, ann_t const *source) {
  if (!dest || !source || !source->buffer)
    return;
  void *new_buffer = malloc(source->n);
  if (!new_buffer)
    return;

  memcpy(new_buffer, source->buffer, source->n);

  if (dest->buffer)
    free(dest->buffer);

  size_t layer_neuron_sz = sizeof(u64) * source->layer_n;
  size_t padding = (sizeof(f64) - (layer_neuron_sz % sizeof(f64))) % sizeof(f64);

  dest->buffer = new_buffer;
  dest->n = source->n;
  dest->layer_n = source->layer_n;
  dest->weight_n = source->weight_n;
  dest->neuron_n = source->neuron_n;
  dest->layer_neuron_n = (u64 *)dest->buffer;
  dest->neuron = (f64 *)((char *)dest->buffer + layer_neuron_sz + padding);
  dest->weight = dest->neuron + source->neuron_n;
  dest->delta = dest->weight + source->weight_n;
  dest->activation_hidden = source->activation_hidden;
  dest->activation_hidden_partial = source->activation_hidden_partial;
  dest->activation_output = source->activation_output;
  dest->activation_output_partial = source->activation_output_partial;
}

void ann_propagation_forward(ann_t *ann, f64 const *const input,
                              f64 *output) {
  if (!ann || !input || !output)
    return;
  f64 *w_ij = ann->weight;
  f64 const *x = input;
  f64 *y = ann->neuron;

  u64 l = 1;

  for (; l < ann->layer_n - 1; l++) {
    for (u64 j = 0; j < ann->layer_neuron_n[l]; j++) {
      y[j] = 0;
      for (u64 i = 0; i < ann->layer_neuron_n[l - 1]; i++) {
        y[j] += x[i] * *w_ij++;
      }

      y[j] += *w_ij++;
      y[j] = ann->activation_hidden(y[j]);
    }

    x = y;
    y += ann->layer_neuron_n[l];
  }

  for (u64 j = 0; j < ann->layer_neuron_n[l]; j++) {
    output[j] = 0;

    for (u64 i = 0; i < ann->layer_neuron_n[l - 1]; i++) {
      output[j] += x[i] * *w_ij++;
    }

    output[j] += *w_ij++;
    output[j] = ann->activation_output(output[j]);
  }
}

void ann_propagation_backward(ann_t *ann, f64 const *input, f64 *output,
                               f64 const *target, f64 rate) {
  if (!ann || !input || !output || !target)
    return;
  i64 l = (i64)ann->layer_n - 1;
  u64 i = 0, j = 0, q = 0;

  f64 *d_j = ann->delta + ann->neuron_n;

  for (j = 0; j < ann->layer_neuron_n[l]; j++) {
    d_j[j] = ann_error_partial(output[j], target[j]) * ann->activation_output_partial(output[j]);
  }

 {
    u64 last_layer_n = ann->layer_neuron_n[l];
    u64 prev_layer_plus_bias = ann->layer_neuron_n[l - 1] + 1;
    if (last_layer_n > UINT64_MAX / prev_layer_plus_bias)
      return;
    f64 *w_jq = ann->weight + ann->weight_n - last_layer_n * prev_layer_plus_bias;

    f64 *o_j = ann->neuron + ann->neuron_n;
    f64 *d_q;

    l--;

    for (; l > 0; --l) {
      d_q = d_j;
      d_j -= ann->layer_neuron_n[l];
      o_j -= ann->layer_neuron_n[l];

      for (j = 0; j < ann->layer_neuron_n[l]; j++) {
        d_j[j] = 0;
        for (q = 0; q < ann->layer_neuron_n[l + 1]; q++) {
          d_j[j] += w_jq[q * (ann->layer_neuron_n[l] + 1) + j] * d_q[q];
        }
        d_j[j] *= ann->activation_hidden_partial(o_j[j]);
      }

      if (l > 1) {
        u64 shift = ann->layer_neuron_n[l];
        u64 prev_plus_bias = ann->layer_neuron_n[l - 1] + 1;
        if (shift <= UINT64_MAX / prev_plus_bias) {
          w_jq -= shift * prev_plus_bias;
        } else {
          return;
        }
      }
    }
  }

  f64 *w_ij = ann->weight;
  d_j = ann->delta;

  l = 1;

  for (j = 0; j < ann->layer_neuron_n[l]; j++) {
    for (i = 0; i < ann->layer_neuron_n[l - 1]; i++) {
      *w_ij -= rate * input[i] * d_j[j];
      w_ij++;
    }
    *w_ij -= rate * d_j[j];
    w_ij++;
  }

  l++;
  f64 *i_i = ann->neuron;

  for (; l < (i64)ann->layer_n; l++) {
    d_j += ann->layer_neuron_n[l - 1];
    for (j = 0; j < ann->layer_neuron_n[l]; j++) {
      for (i = 0; i < ann->layer_neuron_n[l - 1]; i++) {
        *w_ij -= rate * i_i[i] * d_j[j];
        w_ij++;
      }
      *w_ij -= rate * d_j[j];
      w_ij++;
    }
    i_i += ann->layer_neuron_n[l - 1];
  }
}

static f64 ann_activation_identity(f64 x) { return x; }
static f64 ann_activation_identity_partial(f64 x) { (void)x; return 1.0; }

static f64 ann_activation_binary(f64 x) {
  return (x >= 0.0) ? 1.0 : 0.0;
}
static f64 ann_activation_binary_partial(f64 x) { (void)x; return 0.0; }

static f64 ann_activation_sigmoid(f64 x) {
  return (x >= 0) ? (1.0 / (1.0 + exp(-x))) : (exp(x) / (1.0 + exp(x)));
}
static f64 ann_activation_sigmoid_partial(f64 x) {
  // x is the already-activated sigmoid output σ(z), so σ'(z) = σ(z)(1-σ(z)) = x(1-x)
  return x * (1.0 - x);
}

static f64 ann_activation_relu(f64 x) { return (x > 0.0) ? x : 0.0; }
static f64 ann_activation_relu_partial(f64 x) {
  return (x > 0.0) ? 1.0 : 0.0;
}

static f64 ann_activation_elu(f64 x) {
  return (x > 0.0) ? x : ELU_ALPHA * (expm1(x));
}
static f64 ann_activation_elu_partial(f64 x) {
  return (x > 0.0) ? 1.0 : x + ELU_ALPHA;
}

static f64 ann_activation_lrelu(f64 x) {
  return (x > 0.0) ? x : LRELU_ALPHA * x;
}
static f64 ann_activation_lrelu_partial(f64 x) {
  return (x > 0.0) ? 1.0 : LRELU_ALPHA;
}

static f64 ann_activation_tanh(f64 x) { return tanh(x); }
static f64 ann_activation_tanh_partial(f64 x) { return 1.0 - (x * x); }

static f64 ann_error(f64 output, f64 target) {
  f64 error = output - target;
  return 0.5 * error * error;
}

static f64 ann_error_partial(f64 output, f64 target) { return (output - target); }

f64 ann_error_total(f64 const *output, f64 const *target, u64 n) {
  if (!output || !target)
    return 0;
  f64 error = 0;
  for (u64 i = 0; i < n; i++) {
    error += ann_error(output[i], target[i]);
  }
  return error;
}

void ann_train_numeric(ann_t *ann, f64 const *input, f64 const *target,
                            u64 n, f64 rate) {
  if (!ann || !input || !target || rate == 0)
    return;
  if (n > SIZE_MAX / sizeof(f64))
    return;
  f64 *output = malloc(n * sizeof(f64));
  if (!output)
    return;

  for (u64 w = 0; w < ann->weight_n; w++) {
    f64 original = ann->weight[w];
    ann->weight[w] = original + rate;
    ann_propagation_forward(ann, input, output);
    f64 error_p = ann_error_total(output, target, n);
    ann->weight[w] = original - rate;
    ann_propagation_forward(ann, input, output);
    f64 error_m = ann_error_total(output, target, n);
    f64 gradient = (error_p - error_m) / (2.0 * rate);
    ann->weight[w] = original - rate * gradient;
  }

  free(output);
}

void ann_set_activation(ann_t *ann, ann_activation_t activation_hidden,
                          ann_activation_t activation_output) {
  if (!ann)
    return;
  if ((int)activation_hidden < 0 || (int)activation_output < 0)
    return;
  if (activation_hidden > TANH || activation_output > TANH)
    return;
  ann->activation_hidden = ACTIVATION[activation_hidden][0];
  ann->activation_hidden_partial = ACTIVATION[activation_hidden][1];
  ann->activation_output = ACTIVATION[activation_output][0];
  ann->activation_output_partial = ACTIVATION[activation_output][1];
}

static void ann_zero_weights(ann_t *ann) {
  if (!ann)
    return;
  memset(ann->weight, 0, ann->weight_n * sizeof(f64));
}

void ann_random(ann_t *ann) {
  if (!ann)
    return;
  u64 l, j, i;
  f64 *wb = ann->weight;
  for (l = 1; l < ann->layer_n; l++) {
    if (ann->layer_neuron_n[l - 1] == 0)
      continue;
    if (ann->layer_neuron_n[l] == 0)
      continue;
    f64 fan_in = (f64)ann->layer_neuron_n[l - 1];
    f64 fan_out = (f64)ann->layer_neuron_n[l];
    f64 range = sqrt(6.0 / (fan_in + fan_out));
    for (j = 0; j < ann->layer_neuron_n[l]; j++) {
      for (i = 0; i < ann->layer_neuron_n[l - 1]; i++) {
        *wb++ = ann_random_range(-range, range);
      }
      *wb++ = 0;
    }
  }
}

static f64 ann_random_range(f64 low, f64 high) {
  return (low + ((f64)rand()) * (high - low) / (f64)RAND_MAX);
}

void ann_print_neuron(ann_t *ann, f64 const *const input,
                       f64 const *const output) {
  if (!ann || !input || !output)
    return;
  fputs("-> ", stderr);
  for (u64 i = 0; i < ann->layer_neuron_n[0]; i++) {
    fprintf(stderr, "  %+.*f", PRINT_PRECISION, input[i]);
  }
  fputs("\n", stderr);

  f64 *n = ann->neuron;
  for (u64 l = 1; l < ann->layer_n - 1; l++) {
    fprintf(stderr, "   ");
    for (u64 i = 0; i < ann->layer_neuron_n[l]; i++) {
      fprintf(stderr, "  %+.*f", PRINT_PRECISION, *n);
      n++;
    }
    fputs("\n", stderr);
  }

  fputs("=  ", stderr);
  for (u64 i = 0; i < ann->layer_neuron_n[ann->layer_n - 1]; i++) {
    fprintf(stderr, "  %+.*f", PRINT_PRECISION, output[i]);
  }
  fputs("\n\n", stderr);
}

void ann_print_weight(ann_t *ann) {
  if (!ann)
    return;
  f64 *weight = ann->weight;
  for (u64 i = 1; i < ann->layer_n; i++) {
    for (u64 j = 0; j < ann->layer_neuron_n[i - 1] + 1; j++) {
      for (u64 k = 0; k < ann->layer_neuron_n[i]; k++) {
        fprintf(stderr, "  %+.*f", PRINT_PRECISION,
                *(weight + k * (ann->layer_neuron_n[i - 1] + 1) + j));
      }
      fputs("\n", stderr);
    }
    weight += ann->layer_neuron_n[i] * (ann->layer_neuron_n[i - 1] + 1);
    fputs("\n", stderr);
  }
}

#endif // ANN_IMPLEMENTATION