#include "erl_nif.h"
#include "mmath.h"

#include <stdio.h>
#include <math.h>

typedef ffloat (*aggr_func) (ffloat, ffloat);
typedef ffloat (*emit_func) (ffloat, double);
typedef ffloat (*threshold_func) (ffloat*, uint32_t, double);

void print(ffloat* vs, uint32_t count) {
  for (uint32_t i = 0; i < count; i++) {
    printf("%.2f(%.2f) ", vs[i].value, vs[i].confidence);
  };
  printf("\n\r");
}

static int
load(ErlNifEnv* env, void** priv, ERL_NIF_TERM load_info)
{
  return 0;
}

static int
upgrade(ErlNifEnv* env, void** priv, void** old_priv, ERL_NIF_TERM load_info)
{
  return 0;
}

static ERL_NIF_TERM
aggr2(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[], aggr_func f, emit_func g)
{
  ErlNifBinary bin;
  ErlNifSInt64 chunk;         // size to be compressed

  ERL_NIF_TERM r;
  ffloat* vs;
  ffloat* target;
  ffloat aggr;          // Aggregator
  double confidence;

  uint32_t target_i = 0;      // target position
  uint32_t count;
  uint32_t pos = 0;
  uint32_t target_size;

  if (argc != 2)
    return enif_make_badarg(env);

  GET_CHUNK(chunk);
  GET_BIN(0, bin, count, vs);

  target_size = ceil((double) count / chunk) * sizeof(ffloat);
  if (! (target = (ffloat*) enif_make_new_binary(env, target_size, &r)))
    return enif_make_badarg(env); // TODO return propper error
  if (count > 0) {
    aggr = vs[0];
    confidence = aggr.confidence;
    pos = 1;
    //We will be overwriting the confidence generated by dec_add because
    //it would give a false impression based on the later values having
    //a higher influence.
    for (uint32_t i = 1; i < count; i++, pos++) {
      if (pos == chunk) {
        aggr.confidence = confidence / chunk;
        target[target_i] = g(aggr, chunk);
        target_i++;
        aggr = vs[i];
        confidence = aggr.confidence;
        pos = 0;
      } else {
        confidence += vs[i].confidence;
        aggr = f(aggr, vs[i]);
      }
    }

    if (count % chunk) {
      for (uint32_t i = 0; i < (chunk - (count % chunk)); i++) {
          aggr = f(aggr, vs[count-1]);
      }
    }

    aggr.confidence = confidence / chunk;
    target[target_i] = g(aggr, chunk);
  }

  return r;
}

static ERL_NIF_TERM
aggr(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[], aggr_func f)
{
    return aggr2(env, argc, argv, f, float_const);
}

static ERL_NIF_TERM
min(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return aggr(env, argc, argv, float_min);
}

static ERL_NIF_TERM
max(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return aggr(env, argc, argv, float_max);
}

static ERL_NIF_TERM
sum(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return aggr(env, argc, argv, float_add);
}

static ERL_NIF_TERM
avg(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return aggr2(env, argc, argv, float_add, float_divc);
}

int comp (const void * elem1, const void * elem2)
{
  double f = ((ffloat*)elem1) ->value;
  double s = ((ffloat*)elem2) ->value;
  if (f > s) return  1;
  if (f < s) return -1;
  return 0;
}

void sort(ffloat* vs, uint32_t count) {
  qsort(vs, count, sizeof(ffloat), comp);
}

// TODO: investigate quickselect https://en.wikipedia.org/wiki/Quickselect
static ERL_NIF_TERM
percentile(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  ErlNifBinary bin;
  ErlNifSInt64 chunk;         // size to be compressed
  double p;

  ERL_NIF_TERM r;
  ffloat* vs;
  ffloat* target;

  uint32_t count;
  uint32_t pos = 0;
  uint32_t target_size;

  if (argc != 3)
    return enif_make_badarg(env);


  if (!enif_get_double(env, argv[1], &p) || p < 0 || p > 1)
    return enif_make_badarg(env);

  if (!enif_get_int64(env, argv[2], &chunk) || chunk < 1)
    return enif_make_badarg(env);

  GET_BIN(0, bin, count, vs);

  target_size = ceil((double) count / chunk) * sizeof(ffloat);
  if (! (target = (ffloat*) enif_make_new_binary(env, target_size, &r)))
    return enif_make_badarg(env); // TODO return propper error
  uint32_t offset;
  for (offset = 0; (offset + chunk) <= count; offset += chunk) {
    uint32_t n = (uint32_t)((chunk - 1) * p);
    sort(&vs[offset], chunk);
    target[pos] = vs[offset + n];
    target[pos].confidence = 0;
    for (uint32_t i = 0; i < chunk; i ++) {
      target[pos].confidence += vs[offset + i].confidence;
    }
    target[pos].confidence /= chunk;
    pos++;
  }
  // if the last chunk is incomplete
  uint32_t leftover = count % chunk;
  if (leftover) {
    uint32_t n = (uint32_t)((leftover - 1) * p);
    sort(&vs[offset], leftover);
    target[pos].value = vs[offset + n].value;
    target[pos].confidence = 0;
    for (uint32_t i = 0; i < leftover; i ++) {
      double c = vs[offset + i].confidence;
      target[pos].confidence += c;
    }
    target[pos].confidence /= leftover;
  }
  return r;
}

static ERL_NIF_TERM
threshold(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[], threshold_func f)
{
  ErlNifBinary bin;
  ErlNifSInt64 chunk;         // size to be compressed
  double level;

  ERL_NIF_TERM r;
  ffloat* vs;
  ffloat* target;

  uint32_t count;
  uint32_t pos = 0;
  uint32_t target_size;

  if (argc != 3)
    return enif_make_badarg(env);

  if (!enif_get_double(env, argv[1], &level))
    return enif_make_badarg(env);

  if (!enif_get_int64(env, argv[2], &chunk) || chunk < 1)
    return enif_make_badarg(env);

  GET_BIN(0, bin, count, vs);

  target_size = ceil((double) count / chunk) * sizeof(ffloat);
  if (! (target = (ffloat*) enif_make_new_binary(env, target_size, &r)))
    return enif_make_badarg(env); // TODO return propper error
  uint32_t offset;
  for (offset = 0; (offset + chunk) <= count; offset += chunk) {
    target[pos] = f(&vs[offset], chunk, level);
    pos++;
  }

  // if the last chunk is incomplete
  uint32_t leftover = count % chunk;
  if (leftover) {
    target[pos] = f(&vs[offset], leftover, level);
  }
  return r;
}

ffloat first_below(ffloat* vs, uint32_t count, double level)
{
  uint32_t i = 0;
  double confidence = 0;
  while(i < count) {
    if (vs[i].value < level) {
      return (ffloat) {
        .value = vs[i].value,
        .confidence = confidence / (i + 1)
      };
    }
    confidence += vs[i].confidence;
    i++;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_first_below(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, first_below);
}

ffloat first_above(ffloat* vs, uint32_t count, double level)
{
  uint32_t i = 0;
  double confidence = 0;
  while(i < count) {
    if (vs[i].value > level) {
      return (ffloat) {
        .value = vs[i].value,
        .confidence = confidence / (i + 1)
      };
    }
    confidence += vs[i].confidence;
    i++;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_first_above(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, first_above);
}

ffloat last_below(ffloat* vs, uint32_t count, double level)
{
  int32_t i = count - 1;
  double confidence = 0;
  while(i >= 0) {
    if (vs[i].value < level) {
      return (ffloat) {
        .value = vs[i].value,
        .confidence = confidence / (i + 1)
      };
    }
    confidence += vs[i].confidence;
    i--;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_last_below(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, last_below);
}

ffloat last_above(ffloat* vs, uint32_t count, double level)
{
  int32_t i = count - 1;
  double confidence = 0;
  while(i >= 0) {
    if (vs[i].value > level) {
      return (ffloat) {
        .value = vs[i].value,
        .confidence = confidence / (i + 1)
      };
    }
    confidence += vs[i].confidence;
    i--;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_last_above(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, last_above);
}

ffloat count_above(ffloat* vs, uint32_t count, double level)
{
  uint32_t counter = 0;
  double confidence = 0;

  for (uint32_t i = 0; i < count; i ++) {
    if (vs[i].value > level) {
      counter++;
    }
    confidence += vs[i].confidence;
  }

  if (counter) {
    return (ffloat) {
      .confidence = confidence / count,
      .value = counter
    };
  } else {
    return (ffloat) {.confidence = 0, .value = 0};
  }
}

static ERL_NIF_TERM
nif_count_above(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, count_above);
}

ffloat count_below(ffloat* vs, uint32_t count, double level)
{
  uint32_t counter = 0;
  double confidence = 0;

  for (uint32_t i = 0; i < count; i ++) {
    if (vs[i].value < level) {
      counter++;
    }
    confidence += vs[i].confidence;
  }

  if (counter) {
    return (ffloat) {
      .confidence = confidence / count,
      .value = counter
    };
  } else {
    return (ffloat) {.confidence = 0, .value = 0};
  }
}

static ERL_NIF_TERM
nif_count_below(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, count_below);
}

ffloat first_above_conf(ffloat* vs, uint32_t count, double level)
{
  uint32_t i = 0;
  while(i < count) {
    if (vs[i].confidence > level) {
      return vs[i];
    }
    i++;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_first_above_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, first_above_conf);
}

ffloat first_below_conf(ffloat* vs, uint32_t count, double level)
{
  uint32_t i = 0;
  while(i < count) {
    if (vs[i].confidence < level) {
      return vs[i];
    }
    i++;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_first_below_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, first_below_conf);
}

ffloat last_above_conf(ffloat* vs, uint32_t count, double level)
{
  int32_t i = count - 1;
  while(i >= 0) {
    if (vs[i].confidence > level) {
      return vs[i];
    }
    i--;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_last_above_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, last_above_conf);
}

ffloat last_below_conf(ffloat* vs, uint32_t count, double level)
{
  int32_t i = count - 1;
  double confidence = 0;
  while(i >= 0) {
    if (vs[i].confidence < level) {
      return vs[i];
    }
    i--;
  }
  return (ffloat) {.confidence = 0, .value = 0};
}

static ERL_NIF_TERM
nif_last_below_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, last_below_conf);
}

ffloat count_above_conf(ffloat* vs, uint32_t count, double level)
{
  uint32_t counter = 0;
  double confidence = 0;

  for (uint32_t i = 0; i < count; i ++) {
    if (vs[i].confidence > level) {
      counter++;
    }
    confidence += vs[i].confidence;
  }

  if (counter) {
    return (ffloat) {
      .confidence = confidence / count,
      .value = counter
    };
  } else {
    return (ffloat) {.confidence = 0, .value = 0};
  }
}

static ERL_NIF_TERM
nif_count_above_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, count_above_conf);
}

ffloat count_below_conf(ffloat* vs, uint32_t count, double level)
{
  uint32_t counter = 0;
  double confidence = 0;

  for (uint32_t i = 0; i < count; i ++) {
    if (vs[i].confidence < level) {
      counter++;
    }
    confidence += vs[i].confidence;
  }

  if (counter) {
    return (ffloat) {
      .confidence = confidence / count,
      .value = counter
    };
  } else {
    return (ffloat) {.confidence = 0, .value = 0};
  }
}

static ERL_NIF_TERM
nif_count_below_conf(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{
  return threshold(env, argc, argv, count_below_conf);
}

static ErlNifFunc nif_funcs[] = {
  {"min",              2, min},
  {"max",              2, max},
  {"sum",              2, sum},
  {"avg",              2, avg},
  {"percentile",       3, percentile},
  {"first_below",      3, nif_first_below},
  {"first_above",      3, nif_first_above},
  {"last_below",       3, nif_last_below},
  {"last_above",       3, nif_last_above},
  {"count_below",      3, nif_count_below},
  {"count_above",      3, nif_count_above},
  {"first_below_conf", 3, nif_first_below_conf},
  {"first_above_conf", 3, nif_first_above_conf},
  {"last_below_conf",  3, nif_last_below_conf},
  {"last_above_conf",  3, nif_last_above_conf},
  {"count_below_conf", 3, nif_count_below_conf},
  {"count_above_conf", 3, nif_count_above_conf}
};

// Initialize this NIF library.
//
// Args: (MODULE, ErlNifFunc funcs[], load, reload, upgrade, unload)
// Docs: http://erlang.org/doc/man/erl_nif.html#ERL_NIF_INIT

ERL_NIF_INIT(mmath_aggr, nif_funcs, &load, NULL, &upgrade, NULL);