//  ================================================================
//
//    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_H
#define KIT_ASTAR_H

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

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

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;

typedef struct {
  i64 id;
  i64 previous;
  i64 exact_distance;
  i64 estimate;
} Astar_Node;

typedef struct {
  i64         capacity;
  i64         size;
  Astar_Node *values;
} Astar_Set;

typedef struct {
  Astar_Set open;
  Astar_Set closed;
  void *    user_data;
  i64       source;
  i64       destination;
} Astar_State;

#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

#ifdef __cplusplus
extern "C" {
#endif

static s32 astar_init(
  Astar_State *state,
  Astar_Set    open,
  Astar_Set    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) {
    .id             = source,
    .previous       = -1,
    .exact_distance = 0,
    .estimate       = 0x7fffffffffffffffll,
  };

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

  return KIT_OK;
}

static s32 astar_path(Astar_State *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 __cplusplus
}
#endif

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

#endif

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

#ifdef ASTAR_TEMPLATE

#ifndef ASTAR_SUFFIX
#  define ASTAR_SUFFIX
#endif

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

#ifndef ASTAR_HEURISTIC
#  define ASTAR_HEURISTIC(user_data_, node_0_, node_2_) \
    0x7fffffffffffffffll
#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)

#ifdef __cplusplus
extern "C" {
#endif

static s32 NAME_(astar_iteration)(Astar_State *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 nearest_node;
  {
    i64 index_in_open = 0;
    for (i64 i = 1; i < state->open.size; i++)
      if (state->open.values[i].estimate < state->open.values[index_in_open].estimate)
        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] = 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) k;
    Astar_Link 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;

    //  Calculate distance estimations
    //
    (void) state->user_data;
    (void) state->destination;
    
    Astar_Node neighbor_node = {
      .id             = link.neighbor,
      .previous       = nearest_node.id,
      .exact_distance = nearest_node.exact_distance + link.distance,
      .estimate       = ASTAR_HEURISTIC(state->user_data, link.neighbor, state->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] = 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) {
        //  Check if this node has a better distance
        //
        if (neighbor_node.exact_distance < state->closed.values[index_in_closed].exact_distance)
          //  Replace the node
          state->closed.values[index_in_closed] = neighbor_node;

        //  Skip this node
        continue;
      }
    }

    //  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 distance
        //
        if (neighbor_node.exact_distance < state->open.values[index_in_open].exact_distance)
          //  Replace the node
          state->open.values[index_in_open] = neighbor_node;
        
        //  Skip this node
        continue;
      }
    }

    if (state->open.size + 1 > state->open.capacity) {
      //  Restore the state
      state->open.values[state->open.size] = 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;
}

#ifdef __cplusplus
}
#endif

#undef NAME_
#undef CAT1_
#undef CAT2_

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

#undef ASTAR_TEMPLATE
#endif