import gleam/bool
import gleam/dict.{type Dict}
import gleam/float
import gleam/list
import gleam/order.{Eq}
import gleam/pair
import gleam/result
import gleam/set.{type Set}

/// A function that takes the current state and its G score and returns a list
/// of possible [`Action`s](#Action).
///
pub type Successor(action, state) =
  fn(state, Float) -> List(Action(action, state))

/// A transition to a `state` by an `action` with a `weight`.
///
/// The weight could mean a cost, the distance or terrain difficulty, ...
///
pub type Action(action, state) {
  Action(action: action, state: state, weight: Float)
}

type Astarstate(state, action) {
  Astarstate(from_state: state, action: action, g_score: Float, f_score: Float)
}

type Flowfieldstate(action) {
  Flowfieldstate(action: action, weight: Float, g_score: Float)
}

/// Performs an A* search to find the optimal path from an initial `state`.
///
/// This function returns a the sequence of actions to reach the goal when the
/// `heuristic_fun` returns 0.0 (or less), signaling the goal has been
/// reached.
///
/// If the goal can not be reach, it will returns the sequence of actions to
/// the "closest" state reached.
///
/// ## Examples
///
/// ```gleam
/// let actions = [
///   Vec2(-1, 0),
///   Vec2(1, 0),
///   Vec2(0, -1),
///   Vec2(0, 1),
///   Vec2(-1, -1),
///   Vec2(-1, 1),
///   Vec2(1, -1),
///   Vec2(1, 1),
/// ]
///
/// let world0 =
///   vec2i_dict.from_string(
///     ""
///     <> "###########\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "#   ## ####\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "## ##### ##\n"
///     <> "#   # #   #\n"
///     <> "#         #\n"
///     <> "###########\n",
///   )
///
/// let world1 =
///   vec2i_dict.from_string(
///     ""
///     <> "###########\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "#   ##x####\n"
///     <> "#   #     #\n"
///     <> "#   #     #\n"
///     <> "## ##### ##\n"
///     <> "#   # #   #\n"
///     <> "#         #\n"
///     <> "###########\n",
///   )
///
/// let init = Vec2(2, 2)
/// let target = Vec2(7, 2)
///
/// fn successor(
///   state: Vec2i,
///   g_score: Float,
///   world: Dict(Vec2i, String),
/// ) -> List(Action(Vec2i, Vec2i)) {
///   actions
///   |> list.filter_map(fn(action) {
///     let weight = action |> vec2i.length
///     use <- bool.guard(g_score +. weight >. 32.0, Error(Nil))
///
///     let state = state |> vec2i.add(action)
///     use <- bool.guard(world |> dict.has_key(state), Error(Nil))
///     Ok(Action(action:, state:, weight:))
///   })
/// }
///
/// fn heuristic(state: Vec2i) -> Float {
///   vec2i.distance(state, target)
/// }
///
/// astar(
///   init:,
///   successor: fn(state, g_score) {
///     successor(state, g_score, world0)
///   },
///   heuristic:,
/// )
/// // Ok:
/// // ###########
/// // #   #     #
/// // # @ #  *  #
/// // # ↓ # ↗   #
/// // # ↓ ##↑####
/// // # ↓ #  ↖  #
/// // # ↓ #   ↖ #
/// // ##↘#####↑##
/// // #  ↘# #↗  #
/// // #   →→↗   #
/// // ###########
///
/// astar(
///   init:,
///   successor: fn(state, g_score) {
///     successor(state, g_score, world1)
///   },
///   heuristic:,
/// )
/// // Error:
/// // ###########
/// // #   #     #
/// // # @ #  *  #
/// // # ↓ #     #
/// // # ↓ ##x####
/// // # ↓ #     #
/// // # ↓ #   ↖ #
/// // ##↘#####↑##
/// // #  ↘# #↗  #
/// // #   →→↗   #
/// // ###########
/// ```
///
pub fn astar(
  init state: state,
  successor successor_fun: Successor(action, state),
  heuristic heuristic_fun: fn(state) -> Float,
) -> Result(List(action), List(action)) {
  do_astar(
    state,
    successor_fun,
    heuristic_fun,
    set.from_list([state]),
    dict.new(),
  )
}

fn do_astar(
  init_state: state,
  successor_fun: Successor(action, state),
  heuristic_fun: fn(state) -> Float,
  queue: Set(state),
  flowfield: Dict(state, Astarstate(state, action)),
) -> Result(List(action), List(action)) {
  case lowest_f_score_state(queue, flowfield) {
    Ok(state) -> {
      let #(g_score, f_score) = case flowfield |> dict.get(state) {
        Ok(Astarstate(g_score:, f_score:, ..)) -> #(g_score, f_score)
        Error(Nil) -> #(0.0, 1.0)
      }

      use <- bool.lazy_guard(f_score -. g_score <=. 0.0, fn() {
        Ok(state |> reconstruct_actions(flowfield |> dict.delete(init_state)))
      })

      let queue = queue |> set.delete(state)

      let actions =
        successor_fun(state, g_score)
        |> astar_actions_folder(state, g_score, heuristic_fun, flowfield)

      let queue = actions |> dict.keys |> set.from_list |> set.union(queue)
      let flowfield = flowfield |> dict.merge(actions)
      do_astar(init_state, successor_fun, heuristic_fun, queue, flowfield)
    }
    Error(Nil) -> {
      closest_state(flowfield)
      |> result.unwrap(init_state)
      |> reconstruct_actions(flowfield |> dict.delete(init_state))
      |> Error
    }
  }
}

fn lowest_f_score_state(
  queue: Set(state),
  flowfield: Dict(state, Astarstate(state, action)),
) -> Result(state, Nil) {
  queue
  |> set.to_list
  |> list.max(fn(a, b) {
    case dict.get(flowfield, a), dict.get(flowfield, b) {
      Ok(a), Ok(b) -> float.compare(b.f_score, a.f_score)
      Ok(a), Error(Nil) -> float.compare(0.0, a.f_score)
      Error(Nil), Ok(b) -> float.compare(b.f_score, 0.0)
      Error(Nil), Error(Nil) -> Eq
    }
  })
}

fn astar_actions_folder(
  actions: List(Action(action, state)),
  from_state: state,
  from_g_score: Float,
  heuristic_fun: fn(state) -> Float,
  flowfield: Dict(state, Astarstate(state, action)),
) -> Dict(state, Astarstate(state, action)) {
  actions
  |> list.fold(dict.new(), fn(acc, action) {
    let g_score = from_g_score +. action.weight

    case flowfield |> dict.merge(acc) |> dict.get(action.state) {
      Ok(old) if old.g_score <=. g_score -> acc
      _ -> {
        let f_score = g_score +. heuristic_fun(action.state)
        Astarstate(from_state:, action: action.action, g_score:, f_score:)
        |> dict.insert(acc, action.state, _)
      }
    }
  })
}

fn closest_state(
  flowfield: Dict(state, Astarstate(state, action)),
) -> Result(state, Nil) {
  flowfield
  |> dict.to_list
  |> list.max(fn(a, b) {
    let #(_, a) = a
    let #(_, b) = b
    float.compare(b.f_score -. b.g_score, a.f_score -. a.g_score)
  })
  |> result.map(pair.first)
}

fn reconstruct_actions(
  state: state,
  flowfield: Dict(state, Astarstate(state, action)),
) -> List(action) {
  do_reconstruct_actions(state, flowfield, [])
}

fn do_reconstruct_actions(
  state: state,
  flowfield: Dict(state, Astarstate(state, action)),
  actions: List(action),
) -> List(action) {
  case flowfield |> dict.get(state) {
    Ok(Astarstate(from_state:, action:, ..)) ->
      do_reconstruct_actions(from_state, flowfield, [action, ..actions])
    Error(Nil) -> actions
  }
}

/// Generates a flowfield (also known as vector field or Dijkstra map).
///
/// ## Examples
///
/// ```gleam
/// let actions = [
///   Vec2(-1, 0),
///   Vec2(1, 0),
///   Vec2(0, -1),
///   Vec2(0, 1),
///   Vec2(-1, -1),
///   Vec2(-1, 1),
///   Vec2(1, -1),
///   Vec2(1, 1),
/// ]
///
/// let world =
///   vec2i_dict.from_string(
///     ""
///     <> "     #     \n"
///     <> " ###   ### \n"
///     <> "           \n"
///     <> "### ##### #\n"
///     <> "# # #   # #\n"
///     <> "# # # # # #\n"
///     <> "# # ### # #\n"
///     <> "# #     # #\n"
///     <> "# ####### #\n"
///     <> "#         #\n"
///     <> "###########\n",
///   )
///
/// let target = Vec2(5, 5)
///
/// fn successor(state: Vec2i, g_score: Float) -> List(Action(Vec2i, Vec2i)) {
///   actions
///   |> list.filter_map(fn(action) {
///     let weight = action |> vec2i.length
///     use <- bool.guard(g_score +. weight >. 32.0, Error(Nil))
///
///     let state = state |> vec2i.subtract(action)
///     use <- bool.guard(vec2i.distance(state, target) >. 16.0, Error(Nil))
///     use <- bool.guard(world |> dict.has_key(state), Error(Nil))
///     Ok(Action(action:, state:, weight:))
///   })
/// }
///
/// astar(init:, successor:)
/// // ↓↙→↘↓#↙←←←←
/// // ↘###↓↙←###↙
/// // →→↘↓↙←←←←←←
/// // ###↓#####↖#
/// // # #↓#↓↙←#↑#
/// // # #↓#*#↖#↑#
/// // # #↘###↑#↑#
/// // # #→→→↗↑#↑#
/// // # #######↑#
/// // #→→→→→→→↗↑#
/// // ###########
/// ```
///
pub fn flowfield(
  init state: state,
  successor successor_fun: Successor(action, state),
) -> Dict(state, action) {
  do_flowfield(state, successor_fun, [state], dict.new())
}

fn do_flowfield(
  start_state: state,
  successor_fun: Successor(action, state),
  queue: List(state),
  flowfield: Dict(state, Flowfieldstate(action)),
) -> Dict(state, action) {
  case queue {
    [state, ..queue] -> {
      let g_score = case flowfield |> dict.get(state) {
        Ok(Flowfieldstate(g_score:, ..)) -> g_score
        Error(Nil) -> 0.0
      }

      let actions =
        successor_fun(state, g_score)
        |> flowfield_actions_folder(g_score, flowfield)

      let queue =
        actions
        |> dict.drop(flowfield |> dict.keys)
        |> dict.keys
        |> list.append(queue, _)

      let flowfield = flowfield |> dict.merge(actions)
      do_flowfield(start_state, successor_fun, queue, flowfield)
    }
    _ ->
      flowfield
      |> dict.delete(start_state)
      |> dict.map_values(fn(_key, value) { value.action })
  }
}

fn flowfield_actions_folder(
  actions: List(Action(action, state)),
  from_g_score: Float,
  flowfield: Dict(state, Flowfieldstate(action)),
) -> Dict(state, Flowfieldstate(action)) {
  actions
  |> list.fold(dict.new(), fn(acc, action) {
    let g_score = from_g_score +. action.weight

    case flowfield |> dict.merge(acc) |> dict.get(action.state) {
      Ok(old)
        if old.g_score <. g_score
        || { old.g_score <=. g_score && old.weight <=. action.weight }
      -> acc
      _ ->
        Flowfieldstate(action: action.action, weight: action.weight, g_score:)
        |> dict.insert(acc, action.state, _)
    }
  })
}