defmodule CPSolver.Constraint.Factory do
  alias CPSolver.Constraint.{
    Sum,
    ElementVar,
    Element2D,
    Modulo,
    Absolute,
    LessOrEqual,
    Equal,
    Reified
  }

  alias CPSolver.Propagator.Modulo, as: ModuloPropagator
  alias CPSolver.IntVariable, as: Variable
  alias CPSolver.BooleanVariable
  alias CPSolver.Variable.Interface
  alias CPSolver.DefaultDomain, as: Domain
  import CPSolver.Variable.View.Factory

  def element(array, x, y) do
    ElementVar.new(array, x, y)
  end

  def element(array, x) do
    y_domain =
      Enum.reduce(array, MapSet.new(), fn el, acc ->
        Interface.domain(el) |> Domain.to_list() |> MapSet.union(acc)
      end)
      |> MapSet.to_list()

    y = Variable.new(y_domain)
    result(y, element(array, x, y))
  end

  def element2d(array2d, x, y) do
    domain = array2d |> List.flatten()
    z = Variable.new(domain)
    result(z, element2d(array2d, x, y, z))
  end

  def element2d(array2d, x, y, z) do
    Element2D.new([array2d, x, y, z])
  end

  def element2d_var(array2d, x, y, z) do
    num_rows = length(array2d)
    num_cols = length(hd(array2d))
    Interface.removeBelow(x, 0)
    Interface.removeAbove(x, num_rows - 1)
    Interface.removeBelow(y, 0)
    Interface.removeAbove(y, num_cols - 1)

    {flat_idx_var, sum_constraint} = add(mul(x, num_cols), y)
    Interface.removeBelow(flat_idx_var, 0)
    Interface.removeAbove(flat_idx_var, num_rows * num_cols - 1)
    element_constraint = element(List.flatten(array2d), flat_idx_var, z)
    [sum_constraint, element_constraint]
  end

  def element2d_var(array2d, x, y) do
    domain =
      Enum.reduce(array2d |> List.flatten(), MapSet.new(), fn el, acc ->
        Interface.domain(el)
        |> Domain.to_list()
        |> MapSet.union(acc)
      end)
      |> MapSet.to_list()

    z = Variable.new(domain)
    result(z, element2d_var(array2d, x, y, z))
  end

  def sum(vars, sum_var) do
    Sum.new(sum_var, vars)
  end

  def sum(vars) do
    {domain_min, domain_max} =
      Enum.reduce(vars, {0, 0}, fn var, {min_acc, max_acc} ->
        domain = Interface.domain(var) |> Domain.to_list()
        {min_acc + Enum.min(domain), max_acc + Enum.max(domain)}
      end)

    domain = domain_min..domain_max

    sum_var = Variable.new(domain)
    result(sum_var, Sum.new(sum_var, vars))
  end

  def count(array, y, c) do
    {b_vars, reif_propagators} =
      for a <- array, reduce: {[], []} do
        {vars_acc, propagators_acc} ->
          b = BooleanVariable.new()
          equal_p = Reified.new([Equal.new(a, y), b])
          {[b | vars_acc], [equal_p | propagators_acc]}
      end

    Interface.removeBelow(c, 0)
    Interface.removeAbove(c, length(array))
    [Sum.new(c, b_vars) | reif_propagators]
  end

  def add(var1, var2) do
    sum([var1, var2])
  end

  def subtract(var1, var2) do
    add(var1, linear(var2, -1, 0))
  end

  def mod(x, y) do
          {lb, ub} = ModuloPropagator.mod_bounds(x, y)
          domain =
            lb..ub

    mod_var = Variable.new(domain)
    result(mod_var, Modulo.new(mod_var, x, y))
  end

  def mod(mod_var, x, y) do
    Modulo.new(mod_var, x, y)
  end

  def absolute(x) do
    abs_min = abs(Interface.min(x))
    abs_max = abs(Interface.max(x))
    domain = 0..max(abs_min, abs_max)

    abs_var = Variable.new(domain)
    result(abs_var, Absolute.new(x, abs_var))
  end

  def absolute(x, abs_var) do
    Absolute.new(x, abs_var)
  end

  defp compose(constraint1, constraint2, relation) do
    b1 = BooleanVariable.new()
    b2 = BooleanVariable.new()
    reif_c1 = Reified.new([constraint1, b1])
    reif_c2 = Reified.new([constraint2, b2])
    %{constraints: [reif_c1, reif_c2, relation.new([b1, b2])], derived_variables: [b1, b2]}
  end

  ## Implication, equivalence, inverse implication.
  ## These function produce the list of constraints:
  ## - 2 reified constraints for constraint1 and constraint2
  ## - relational constraint (LessOrEqual for implications, Equal for equivalence)
  ## over control variables induced by reified constraints.
  ##
  def impl(constraint1, constraint2) do
    compose(constraint1, constraint2, LessOrEqual)
  end

  def equiv(constraint1, constraint2) do
    compose(constraint1, constraint2, Equal)
  end

  def inverse_impl(constraint1, constraint2) do
    impl(constraint2, constraint1)
  end

  defp result(derived_variable, constraint) do
    {derived_variable, constraint}
  end
end