/*
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 */

#include <assert.h>
#include <stdint.h>
#include <string.h>

#include "av1/common/blockd.h"
#include "av1/encoder/palette.h"
#include "av1/encoder/random.h"

#ifndef AV1_K_MEANS_DIM
#error "This template requires AV1_K_MEANS_DIM to be defined"
#endif

#define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y)
#define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM)

static int RENAME(calc_dist)(const int *p1, const int *p2) {
  int dist = 0;
  for (int i = 0; i < AV1_K_MEANS_DIM; ++i) {
    const int diff = p1[i] - p2[i];
    dist += diff * diff;
  }
  return dist;
}

void RENAME(av1_calc_indices)(const int *data, const int *centroids,
                              uint8_t *indices, int64_t *dist, int n, int k) {
  if (dist) {
    *dist = 0;
  }
  for (int i = 0; i < n; ++i) {
    int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
    indices[i] = 0;
    for (int j = 1; j < k; ++j) {
      const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
                                              centroids + j * AV1_K_MEANS_DIM);
      if (this_dist < min_dist) {
        min_dist = this_dist;
        indices[i] = j;
      }
    }
    if (dist) {
      *dist += min_dist;
    }
  }
}

static void RENAME(calc_centroids)(const int *data, int *centroids,
                                   const uint8_t *indices, int n, int k) {
  int i, j;
  int count[PALETTE_MAX_SIZE] = { 0 };
  unsigned int rand_state = (unsigned int)data[0];
  assert(n <= 32768);
  memset(centroids, 0, sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);

  for (i = 0; i < n; ++i) {
    const int index = indices[i];
    assert(index < k);
    ++count[index];
    for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
      centroids[index * AV1_K_MEANS_DIM + j] += data[i * AV1_K_MEANS_DIM + j];
    }
  }

  for (i = 0; i < k; ++i) {
    if (count[i] == 0) {
      memcpy(centroids + i * AV1_K_MEANS_DIM,
             data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM,
             sizeof(centroids[0]) * AV1_K_MEANS_DIM);
    } else {
      for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
        centroids[i * AV1_K_MEANS_DIM + j] =
            DIVIDE_AND_ROUND(centroids[i * AV1_K_MEANS_DIM + j], count[i]);
      }
    }
  }
}

void RENAME(av1_k_means)(const int *data, int *centroids, uint8_t *indices,
                         int n, int k, int max_itr) {
  int centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
  uint8_t indices_tmp[MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT];
  int *meta_centroids[2] = { centroids, centroids_tmp };
  uint8_t *meta_indices[2] = { indices, indices_tmp };
  int i, l = 0, prev_l, best_l = 0;
  int64_t this_dist;

  assert(n <= MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT);

#if AV1_K_MEANS_DIM - 2
  av1_calc_indices_dim1(data, centroids, indices, &this_dist, n, k);
#else
  av1_calc_indices_dim2(data, centroids, indices, &this_dist, n, k);
#endif

  for (i = 0; i < max_itr; ++i) {
    const int64_t prev_dist = this_dist;
    prev_l = l;
    l = (l == 1) ? 0 : 1;

    RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k);
#if AV1_K_MEANS_DIM - 2
    av1_calc_indices_dim1(data, meta_centroids[l], meta_indices[l], &this_dist,
                          n, k);
#else
    av1_calc_indices_dim2(data, meta_centroids[l], meta_indices[l], &this_dist,
                          n, k);
#endif

    if (this_dist > prev_dist) {
      best_l = prev_l;
      break;
    }
    if (!memcmp(meta_centroids[l], meta_centroids[prev_l],
                sizeof(centroids[0]) * k * AV1_K_MEANS_DIM))
      break;
  }
  if (i == max_itr) best_l = l;
  if (best_l != 0) {
    memcpy(centroids, meta_centroids[1],
           sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
    memcpy(indices, meta_indices[1], sizeof(indices[0]) * n);
  }
}
#undef RENAME_
#undef RENAME