fixpoint

defmodule CPSolver.Examples.TSP do
  @doc """
  The Traveling Salesman problem.
  Given:
  - a set of n locations;
  - for each pair of locations, a distance between them.

  Find the shortest possible route that visits each location exactly once and returns to the origin location.

  <a href="https://en.wikipedia.org/wiki/Travelling_salesman_problem">Wikipedia</a>.

  """
  alias CPSolver.IntVariable, as: Variable
  alias CPSolver.Model
  alias CPSolver.Constraint.{Circuit, Less}
  alias CPSolver.Objective
  import CPSolver.Constraint.Factory
  import CPSolver.Utils
  alias CPSolver.Utils.TupleArray

  alias CPSolver.Variable.UnfixedTracker, as: Tracker

  require Logger

  @checkmark_symbol "\u2713"
  @failure_symbol "\u1D350"

  def run(instance, opts \\ []) do
    Logger.configure(level: :notice)

    model = model(instance, opts)

    opts =
      Keyword.merge(
        [
          search: search(model),
          solution_handler: solution_handler(model),
          timeout: :timer.minutes(5)
        ],
        opts
      )

    Logger.warning("Bounds: #{model.extra.lb}, #{model.extra.ub}")
    {:ok, _res} = CPSolver.solve(model, opts)
  end

  ## Read and compile data from instance file
  def model(data, opts \\ [])

  def model(data, opts) when is_binary(data) do
    {_n, distances} = parse_instance(data)
    model(distances, opts)
  end

  def model(distances, opts) do
    n = length(distances)

    symmetry_breaking = Keyword.get(opts, :symmetry_breaking, true)
    {lb, ub} = get_bounds(distances)
    ## successor[i] = j <=> location j follows location i
    successors =
      Enum.map(0..(n - 1), fn i ->
        for j <- 0..(n - 1), j != i do
          j
        end
        |> Variable.new(name: "succ_#{i}")
      end)

    ## Element constraints
    ## For each i, distance between i and it's successor must be in i-row of distance matrix
    {dist_succ, element_constraints} =
      Enum.map(0..(n - 1), fn i ->
        element(Enum.at(distances, i), Enum.at(successors, i))
      end)
      |> Enum.unzip()

    {total_distance, sum_constraint} = sum(dist_succ)
    ## Apply bounds
    Variable.removeBelow(total_distance, lb)
    Variable.removeAbove(total_distance, ub)

    Model.new(
      successors,
      [
        Circuit.new(successors),
        sum_constraint
      ] ++
        element_constraints ++
        ((symmetry_breaking && symmetry_constraints(successors, n)) || []),
      objective: Objective.minimize(total_distance),
      extra: %{n: n, distances: distances, lb: lb, ub: ub}
    )
  end

  defp symmetry_constraints(successors, n) do
    zero_succ = hd(successors)
    zero_pred = Variable.new(0..(n - 2))
    pred_index_constraint = element(successors, zero_pred, zero_succ)
    ## For the start of the cycle, the index of predessor is less than the index of successor
    ordering_constraint = Less.new(zero_pred, zero_succ)

    [pred_index_constraint, ordering_constraint]
  end

  defp get_bounds(distances) do
    l = length(distances)

    graph =
      Enum.reduce(1..(l - 1), BitGraph.new(), fn v1, acc ->
        Enum.reduce((v1 + 1)..l, acc, fn v2, acc2 ->
          BitGraph.add_edge(acc2, v1, v2)
        end)
      end)

    dist_fun = fn from, to -> Enum.at(distances, to - 1) |> Enum.at(from - 1) end
    {mst_edges, lb} = BitGraph.mst(graph, dist_fun: dist_fun)
    path_upper_bound = path_upper_bound(mst_edges, graph, dist_fun)
    {lb, min(path_upper_bound, 2 * lb)}
  end

  defp path_upper_bound(mst_edges, graph, dist_fun) do
    BitGraph.Algorithm.dfs(graph,
      process_edge_fun: fn %{acc: acc} = _state, from, to ->
        if {from, to} in mst_edges do
          if acc do
            acc = if from in acc do
              acc
            else
              [from | acc]
            end

            if to in acc do
              acc
            else
              [to | acc]
            end
          else
            [to, from]
          end
        else
          acc
        end
      end
    )
    |> Map.get(:acc)
    |> then(fn path ->
      circuit = [List.last(path) | path]
      Enum.reduce(0..length(circuit) - 2, 0, fn idx, acc ->
        acc + dist_fun.(Enum.at(circuit, idx), Enum.at(circuit, idx + 1))
      end)
    end)
  end

  def check_solution(solution, %{extra: %{distances: distances}} = _model) do
    n = length(distances)
    successors = Enum.take(solution, n)
    ## In the solution, total cost follows assignments
    total_distance = Enum.at(solution, n)

    sum_distances =
      successors
      |> Enum.with_index()
      |> Enum.reduce(0, fn {succ, idx}, acc ->
        acc + (Enum.at(distances, idx) |> Enum.at(succ))
      end)

    hamiltonian?(successors) &&
      total_distance == sum_distances && n == MapSet.new(successors) |> MapSet.size()
  end

  def search(%{extra: %{distances: distances, n: n}} = _model) do
    tuple_matrix = TupleArray.new(distances)

    choose_value_fun = fn %{index: idx} = var ->
      d_values = domain_values(var)

      (idx in 1..n &&
         Enum.min_by(d_values, fn dom_idx -> TupleArray.at(tuple_matrix, [idx - 1, dom_idx]) end)) ||
        Enum.random(d_values)
    end

    choose_variable_fun = fn %{
                               unfixed_variables_tracker: tracker,
                               variables: variables
                             } = _space_data ->
      circuit_vars =
        Tracker.iterate(
          tracker,
          variables,
          [],
          fn v, acc ->
            if v.index <= n do
              [v | acc]
            else
              acc
            end
          end,
          false
        )

      if !Enum.empty?(circuit_vars) do
        difference_between_closest_distances(circuit_vars, tuple_matrix)
      end
    end

    {choose_variable_fun, choose_value_fun}
  end

  def solution_handler(model) do
    fn solution, space_state ->
      solution
      |> Enum.at(model.extra.n)
      |> tap(fn {_ref, objective} ->
        if check_solution(
             Enum.map(solution, fn {_, val} -> val end),
             model
           ) do
          Logger.notice("#{@checkmark_symbol} #{objective}")
          Logger.notice(inspect(CPSolver.statistics(space_state.shared)))
        else
          Logger.error("#{@failure_symbol} #{objective}" <> ": wrong -((")
        end
      end)
    end
  end

  ## Choose the variable with the maximum difference between closest and second closest distance to its successors
  ##
  defp difference_between_closest_distances(circuit_vars, distances) do
    Enum.max_by(circuit_vars, fn %{index: idx} = var ->
      dom = domain_values(var)

      (MapSet.size(dom) < 2 && 0) ||
        dom
        |> Enum.map(fn value ->
          TupleArray.at(distances, [idx - 1, value])
        end)
        |> Enum.sort(:desc)
        |> then(fn dists -> abs(Enum.at(dists, 1) - hd(dists)) end)
    end)
  end

  ## solution -> sequence of visits
  def to_route(solution, %{extra: %{n: n}} = _model) do
    circuit = Enum.take(solution, n)

    Enum.reduce(0..(n - 1), [0], fn _idx, [next | _rest] = acc ->
      [Enum.at(circuit, next) | acc]
    end)
    |> Enum.reverse()
  end

  def hamiltonian?(sequence) do
    {cycle_length, _current} =
      Enum.reduce_while(sequence, {1, 0}, fn _succ, {length_acc, succ_acc} = acc ->
        next = Enum.at(sequence, succ_acc)

        if next == 0 do
          {:halt, acc}
        else
          {:cont, {length_acc + 1, next}}
        end
      end)

    cycle_length == length(sequence)
  end

  def parse_instance(filename) do
    filename
    |> File.read!()
    |> String.split("\n", trim: true)
    |> then(fn [n_str | lines] ->
      n = String.to_integer(String.trim(n_str))

      distances =
        lines
        |> Enum.take(n)
        |> parse_matrix()

      {n, distances}
    end)
  end

  defp parse_matrix(lines) do
    Enum.map(lines, fn line ->
      line
      |> String.replace("\t", " ")
      |> String.split(" ", trim: true)
      |> Enum.map(fn num_str -> String.to_integer(String.trim(num_str)) end)
    end)
  end

  def to_dot(distances) do
    n = length(distances)

    graph =
      for i <- 0..(n - 1), j <- 0..(n - 1), reduce: Graph.new() do
        acc ->
          weight = Enum.at(distances, i) |> Enum.at(j)
          Graph.add_edge(acc, i, j, weight: weight, label: weight)
      end

    Graph.to_dot(graph)
  end
end