//  ================================================================
//
//    A* graph search algorithm
//
//  ================================================================

//  TODO
//  - Nearest: save the nearest node when the search is failed or not
//    yet finished.
//  - Sight: when two nodes are in direct sight of each other, skip
//    nodes between them (Theta*).

#ifndef KIT_ASTAR_INL_H
#define KIT_ASTAR_INL_H

#include "types.h"
#include "status.h"

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

#ifdef __cplusplus
extern "C" {
#endif

enum {
  ASTAR_PROGRESS = 0,
  ASTAR_SUCCESS,
  ASTAR_FAIL,
  ASTAR_OUT_OF_MEMORY,
  ASTAR_INVALID_ARGUMENT,
};

typedef struct {
  b8  stop;
  b8  skip;
  i64 neighbor;
  i64 distance;
} astar_link_t;

typedef struct {
  i64 id;
  i64 previous;
  i64 exact_source_to_node;
  i64 estimated_source_to_destination;
} astar_node_t;

typedef struct {
  i64           capacity;
  i64           size;
  astar_node_t *values;
} astar_set_t;

typedef struct {
  astar_set_t open;
  astar_set_t closed;
  void       *user_data;
  i64         source;
  i64         destination;
} astar_state_t;

#ifdef __GNUC__
#  pragma GCC diagnostic push
#  pragma GCC diagnostic ignored "-Wunused-function"
#  pragma GCC diagnostic ignored "-Wunknown-pragmas"
#  pragma GCC            push_options
#  pragma GCC            optimize("O3")
#endif

static s32 astar_init(astar_state_t *state, astar_set_t open,
                      astar_set_t closed, void *user_data, i64 source,
                      i64 destination) {
  assert(state != NULL && open.capacity > 0 && open.values != NULL &&
         closed.capacity > 0 && closed.values != NULL);

  if (state == NULL)
    return KIT_ERROR_INVALID_ARGUMENT;

  memset(state, 0, sizeof *state);

  state->source      = source;
  state->destination = destination;
  state->open        = open;
  state->closed      = closed;
  state->user_data   = user_data;

  if (state->open.capacity <= 0)
    return KIT_ERROR_INVALID_ARGUMENT;

  state->open.values[0] = (astar_node_t) {
    .id                              = source,
    .previous                        = -1,
    .exact_source_to_node            = 0,
    .estimated_source_to_destination = -1,
  };

  state->open.size   = 1;
  state->closed.size = 0;

  return KIT_OK;
}

static s32 astar_path(astar_state_t *state, i64 *size, i64 *path) {
  assert(state != NULL && size != NULL);

  if (state == NULL || size == NULL)
    return KIT_ERROR_INVALID_ARGUMENT;

  i64 path_size = *size;

  *size = 1;

  i64 id = state->destination;

  while (id != state->source) {
    if (*size > state->closed.size)
      return KIT_ERROR_INTERNAL;

    i64 index = 0;
    for (; index < state->closed.size; index++)
      if (state->closed.values[index].id == id)
        break;
    if (index == state->closed.size)
      return KIT_ERROR_INTERNAL;

    if (path != NULL && *size <= path_size)
      path[*size - 1] = id;

    id = state->closed.values[index].previous;
    ++*size;
  }

  if (path != NULL && *size <= path_size)
    path[*size - 1] = id;

  if (path != NULL) {
    if (*size < path_size)
      path_size = *size;

    for (i64 i = 0; i < path_size / 2; ++i) {
      i64 z                   = path[i];
      path[i]                 = path[path_size - 1 - i];
      path[path_size - 1 - i] = z;
    }
  }

  return KIT_OK;
}

#ifdef __GNUC__
#  pragma GCC            pop_options
#  pragma GCC diagnostic pop
#endif

#ifdef __cplusplus
}
#endif

#endif

//  ================================================================
//
//    Algorithm template
//
//  ================================================================

#ifdef ASTAR_TEMPLATE

#ifdef __cplusplus
extern "C" {
#endif

#ifndef ASTAR_SUFFIX
#  define ASTAR_SUFFIX
#endif

#ifndef ASTAR_NEIGHBOR
#  define ASTAR_NEIGHBOR(user_data_, node_, index_) \
    (astar_link_t) {                                \
      .stop = 1,                                    \
    }
#endif

#ifndef ASTAR_HEURISTIC
#  define ASTAR_HEURISTIC(user_data_, node_0_, node_2_) (-1)
#endif

#ifdef __GNUC__
#  pragma GCC diagnostic push
#  pragma GCC diagnostic ignored "-Wunused-function"
#  pragma GCC diagnostic ignored "-Wunknown-pragmas"
#  pragma GCC            push_options
#  pragma GCC            optimize("O3")
#endif

#define CAT1_(x, y) x##y
#define CAT2_(x, y) CAT1_(x, y)
#define NAME_(x) CAT2_(x, ASTAR_SUFFIX)

static s32 NAME_(astar_iteration)(astar_state_t *state) {
  assert(state != NULL);

  if (state == NULL)
    return ASTAR_INVALID_ARGUMENT;

  if (state->open.size == 0)
    return ASTAR_FAIL;

  if (state->open.values == NULL || state->closed.values == NULL)
    return ASTAR_OUT_OF_MEMORY;

  //  Find a node in the open set that is closest to the destination
  //
  astar_node_t nearest_node;
  {
    i64 index_in_open = 0;
    for (i64 i = 1; i < state->open.size; i++)
      if (state->open.values[i].estimated_source_to_destination <
          state->open.values[index_in_open]
              .estimated_source_to_destination)
        index_in_open = i;
    nearest_node = state->open.values[index_in_open];
    if (index_in_open != state->open.size - 1)
      state->open.values[index_in_open] =
          state->open.values[state->open.size - 1];
    --state->open.size;
  }

  //  Check if we reached the destination
  //
  if (nearest_node.id == state->destination) {
    if (state->closed.size + 1 > state->closed.capacity) {
      //  Restore the state
      state->open.values[state->open.size - 1] = nearest_node;
      ++state->open.size;

      return ASTAR_OUT_OF_MEMORY;
    }

    //  Add the nearest node to the closed set
    //
    state->closed.values[state->closed.size] = nearest_node;
    ++state->closed.size;

    //  Finish the search
    return ASTAR_SUCCESS;
  }

  //  Enumerate all neighbors
  //
  for (i64 k = 0;; ++k) {
    //  Get a link to the next neighbor node
    //
    (void) state->user_data;
    (void) nearest_node.id;
    (void) k;
    astar_link_t link = ASTAR_NEIGHBOR(state->user_data,
                                       nearest_node.id, k);

    //  If there is no more neighbors, break the loop
    if (link.stop)
      break;
    //  If there is no link, proceed to the next link
    if (link.skip)
      continue;

    astar_node_t neighbor_node = {
      .id                   = link.neighbor,
      .previous             = nearest_node.id,
      .exact_source_to_node = nearest_node.exact_source_to_node +
                              link.distance,
      .estimated_source_to_destination = -1,
    };

    //  Calculate distance estimations
    //
    (void) state->user_data;
    (void) neighbor_node.id;
    (void) state->destination;
    i64 estimated_node_to_destination = ASTAR_HEURISTIC(
        state->user_data, neighbor_node.id, state->destination);

    neighbor_node.estimated_source_to_destination =
        neighbor_node.exact_source_to_node +
        estimated_node_to_destination;

    //  Check if we reached the destination
    //
    if (neighbor_node.id == state->destination) {
      if (state->closed.size + 2 > state->closed.capacity) {
        //  Restore the state
        state->open.values[state->open.size - 1] = nearest_node;
        ++state->open.size;

        return ASTAR_OUT_OF_MEMORY;
      }

      //  Add the nearest node to the closed set
      //
      state->closed.values[state->closed.size] = nearest_node;
      ++state->closed.size;

      //  Add the neighbor node to the closed set
      //
      state->closed.values[state->closed.size] = neighbor_node;
      ++state->closed.size;

      //  Finish the search
      return ASTAR_SUCCESS;
    }

    //  Check if this node is already in the closed set
    //
    {
      i64 index_in_closed = -1;
      for (i64 i = 0; i < state->closed.size; ++i)
        if (state->closed.values[i].id == neighbor_node.id) {
          index_in_closed = i;
          break;
        }
      if (index_in_closed != -1) {
        if (state->closed.values[index_in_closed]
                .exact_source_to_node <
            neighbor_node.exact_source_to_node)
          //  Skip this node
          continue;
        //  Replace the node
        state->closed.values[index_in_closed] = neighbor_node;
      }
    }

    //  Check if this node is already in the open set
    //
    {
      i64 index_in_open = -1;
      for (i64 i = 0; i < state->closed.size; ++i)
        if (state->open.values[i].id == neighbor_node.id) {
          index_in_open = i;
          break;
        }
      if (index_in_open != -1) {
        //  Check if this node has a better estimate
        //
        if (neighbor_node.estimated_source_to_destination <
            state->open.values[index_in_open]
                .estimated_source_to_destination)
          //  Replace the node
          state->open.values[index_in_open] = neighbor_node;
        continue;
      }
    }

    if (state->open.size + 1 > state->open.capacity) {
      //  Restore the state
      state->open.values[state->open.size - 1] = nearest_node;
      ++state->open.size;

      return ASTAR_OUT_OF_MEMORY;
    }

    //  Add this node to the open set
    //
    state->open.values[state->open.size] = neighbor_node;
    ++state->open.size;
  }

  //  Check if we can add a node to the closed set
  //
  if (state->closed.size + 1 > state->closed.capacity) {
    //  Restore the state
    state->open.values[state->open.size - 1] = nearest_node;
    ++state->open.size;

    return ASTAR_OUT_OF_MEMORY;
  }

  //  Add the nearest node to the closed set
  //
  state->closed.values[state->closed.size] = nearest_node;
  ++state->closed.size;

  //  Continue the search
  return ASTAR_PROGRESS;
}

#undef NAME_
#undef CAT1_
#undef CAT2_

#ifdef __GNUC__
#  pragma GCC            pop_options
#  pragma GCC diagnostic pop
#endif

#ifdef __cplusplus
}
#endif

#undef ASTAR_TEMPLATE
#endif