/*
 * Licensed to Julian Hyde under one or more contributor license
 * agreements.  See the NOTICE file distributed with this work
 * for additional information regarding copyright ownership.
 * Julian Hyde licenses this file to you under the Apache
 * License, Version 2.0 (the "License"); you may not use this
 * file except in compliance with the License.  You may obtain a
 * copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied.  See the License for the specific
 * language governing permissions and limitations under the
 * License.
 */
package net.hydromatic.morel.util;

import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;
import static net.hydromatic.morel.util.Pair.forEachIndexed;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSortedMap;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;
import org.checkerframework.checker.nullness.qual.NonNull;

/**
 * Given pairs of terms, finds a substitution to minimize those pairs of terms.
 */
public abstract class Unifier {
  private int varId;
  private final Map<String, Variable> variableMap = new HashMap<>();
  private final Map<String, Sequence> atomMap = new HashMap<>();
  private final Map<String, Sequence> sequenceMap = new HashMap<>();

  /** Whether this unifier checks for cycles in substitutions. */
  public boolean occurs() {
    return false;
  }

  /** Creates a sequence, or returns an existing one with the same terms. */
  public Sequence apply(String operator, Term... args) {
    return apply(operator, ImmutableList.copyOf(args));
  }

  /** Creates a sequence, or returns an existing one with the same terms. */
  public Sequence apply(String operator, Iterable<? extends Term> args) {
    final Sequence sequence =
        new Sequence(operator, ImmutableList.copyOf(args));
    return sequenceMap.computeIfAbsent(sequence.toString(), n -> sequence);
  }

  /** Creates a variable, or returns an existing one with the same name. */
  public Variable variable(String name) {
    return variableMap.computeIfAbsent(name, Variable::new);
  }

  /** Creates a variable, or returns an existing one with the same ordinal. */
  public Variable variable(int ordinal) {
    String name = "T" + ordinal;
    return variableMap.computeIfAbsent(name, name2 -> new Variable(ordinal));
  }

  /** Creates a new variable, with a new name. */
  public Variable variable() {
    for (; ; ) {
      final int ordinal = varId++;
      final String name = "T" + ordinal;
      if (!variableMap.containsKey(name)) {
        final Variable variable = new Variable(ordinal);
        assert variable.name.equals(name);
        variableMap.put(name, variable);
        return variable;
      }
    }
  }

  /** Creates an atom, or returns an existing one with the same name. */
  public Term atom(String name) {
    return atomMap.computeIfAbsent(name, Sequence::new);
  }

  /**
   * Creates a substitution.
   *
   * <p>The arguments are alternating variable / term pairs. For example, {@code
   * substitution(a, x, b, y)} becomes [a/X, b/Y].
   */
  public Substitution substitution(Term... varTerms) {
    final ImmutableMap.Builder<Variable, Term> mapBuilder =
        ImmutableMap.builder();
    if (varTerms.length % 2 != 0) {
      throw new AssertionError();
    }
    for (int i = 0; i < varTerms.length; i += 2) {
      mapBuilder.put((Variable) varTerms[i + 1], varTerms[i]);
    }
    return new Substitution(mapBuilder.build());
  }

  static Sequence sequenceApply(
      String operator,
      Map<Variable, Term> substitutions,
      Iterable<Term> terms) {
    final ImmutableList.Builder<Term> newTerms = ImmutableList.builder();
    for (Term term : terms) {
      newTerms.add(term.apply(substitutions));
    }
    return new Sequence(operator, newTerms.build());
  }

  public abstract @NonNull Result unify(
      List<TermTerm> termPairs,
      Map<Variable, Action> termActions,
      Tracer tracer);

  private static void checkCycles(
      Map<Variable, Term> map, Map<Variable, Variable> active)
      throws CycleException {
    for (Term term : map.values()) {
      term.checkCycle(map, active);
    }
  }

  protected Failure failure(String reason) {
    return new Failure() {
      @Override
      public String toString() {
        return reason;
      }
    };
  }

  /** Called by the unifier when a Term's type becomes known. */
  @FunctionalInterface
  public interface Action {
    void accept(
        Variable variable,
        Term term,
        Substitution substitution,
        BiConsumer<Term, Term> termPairs);
  }

  /**
   * Result of attempting unification. A success is {@link Substitution}, but
   * there are other failures.
   */
  public interface Result {}

  /** Result indicating that unification was not possible. */
  public static class Failure implements Result {}

  /**
   * The results of a successful unification. Gives access to the raw variable
   * mapping that resulted from the algorithm, but can also resolve a variable
   * to the fullest extent possible with the {@link #resolve} method.
   */
  public static final class SubstitutionResult extends Substitution
      implements Result {
    private SubstitutionResult(Map<Variable, Term> resultMap) {
      super(ImmutableSortedMap.copyOf(resultMap));
    }

    /** Empty substitution result. */
    public static final SubstitutionResult EMPTY =
        create(ImmutableSortedMap.of());

    /** Creates a substitution result from a map. */
    public static SubstitutionResult create(Map<Variable, Term> resultMap) {
      return new SubstitutionResult(ImmutableSortedMap.copyOf(resultMap));
    }

    /** Creates a substitution result with one (variable, term) entry. */
    public static SubstitutionResult create(Variable v, Term t) {
      return new SubstitutionResult(ImmutableSortedMap.of(v, t));
    }
  }

  /**
   * Map from variables to terms.
   *
   * <p>Quicker to create than its sub-class {@link SubstitutionResult} because
   * the map is mutable and not sorted.
   */
  public static class Substitution {
    /**
     * The result of the unification algorithm proper. This does not have
     * everything completely resolved: some variable substitutions are required
     * before getting the most atom-y representation.
     */
    public final Map<Variable, Term> resultMap;

    Substitution(Map<Variable, Term> resultMap) {
      this.resultMap = resultMap;
    }

    @Override
    public int hashCode() {
      return resultMap.hashCode();
    }

    @Override
    public boolean equals(Object obj) {
      return this == obj
          || obj instanceof Substitution
              && resultMap.equals(((Substitution) obj).resultMap);
    }

    @Override
    public String toString() {
      return accept(new StringBuilder()).toString();
    }

    public StringBuilder accept(StringBuilder buf) {
      buf.append("[");
      forEachIndexed(
          resultMap,
          (i, variable, term) ->
              buf.append(i > 0 ? ", " : "")
                  .append(term)
                  .append("/")
                  .append(variable));
      return buf.append("]");
    }

    public Term resolve(Term term) {
      Term previous;
      Term current = term;
      do {
        previous = current;
        current = current.apply(resultMap);
      } while (!current.equals(previous));
      return current;
    }

    private boolean hasCycles(Map<Variable, Term> map) {
      try {
        checkCycles(map, new IdentityHashMap<>());
        return false;
      } catch (CycleException e) {
        return true;
      }
    }

    public Substitution resolve() {
      if (hasCycles(resultMap)) {
        return this;
      }
      final ImmutableMap.Builder<Variable, Term> builder =
          ImmutableMap.builder();
      resultMap.forEach((key, value) -> builder.put(key, resolve(value)));
      return new Substitution(builder.build());
    }
  }

  /** Term (variable, symbol or node). */
  public interface Term {
    /** Applies a substitution to this term. */
    Term apply(Map<Variable, Term> substitutions);

    /** Applies a single-variable substitution to this term. */
    Term apply1(Variable variable, Term term);

    /** Returns whether this term references a given variable. */
    boolean contains(Variable variable);

    /** Throws CycleException if expanding this term leads to a cycle. */
    void checkCycle(Map<Variable, Term> map, Map<Variable, Variable> active)
        throws CycleException;

    /** Accepts a visitor. */
    <R> R accept(TermVisitor<R> visitor);
  }

  /**
   * Visitor for terms.
   *
   * @param <R> return type
   * @see Term#accept(TermVisitor)
   */
  public interface TermVisitor<R> {
    R visit(Sequence sequence);

    R visit(Variable variable);
  }

  /**
   * Control flow exception, thrown by {@link Term#checkCycle(Map, Map)} if it
   * finds a cycle in a substitution map.
   */
  public static class CycleException extends Exception {}

  /**
   * A variable that represents a symbol or a sequence; unification's task is to
   * find the substitutions for such variables.
   */
  public static final class Variable implements Term, Comparable<Variable> {
    final String name;
    final int ordinal;

    Variable(String name, int ordinal) {
      this.name = requireNonNull(name);
      this.ordinal = ordinal;
      checkArgument(
          name.equals(name.toUpperCase(Locale.ROOT)),
          "must be upper case: %s",
          name);
    }

    /**
     * Creates a variable with a name. The name must not be like "T0" or "T123",
     * because those are the names created for variables with ordinals.
     */
    Variable(String name) {
      this(name, -1);
      checkArgument(!name.matches("T[0-9]+"), name);
    }

    /**
     * Creates a variable with an ordinal. If the ordinal is "34" the name will
     * be "T34". The ordinal must be non-negative.
     */
    Variable(int ordinal) {
      this("T" + ordinal, ordinal);
      checkArgument(ordinal >= 0, ordinal);
    }

    @Override
    public String toString() {
      return name;
    }

    @Override
    public int compareTo(Variable o) {
      int c = Integer.compare(ordinal, o.ordinal);
      if (c == 0) {
        c = name.compareTo(o.name);
      }
      return c;
    }

    @Override
    public Term apply(Map<Variable, Term> substitutions) {
      return substitutions.getOrDefault(this, this);
    }

    @Override
    public Term apply1(Variable variable, Term term) {
      return variable == this ? term : this;
    }

    @Override
    public boolean contains(Variable variable) {
      return variable == this;
    }

    @Override
    public void checkCycle(
        Map<Variable, Term> map, Map<Variable, Variable> active)
        throws CycleException {
      final Term term = map.get(this);
      if (term != null) {
        if (active.put(this, this) != null) {
          throw new CycleException();
        }
        term.checkCycle(map, active);
        active.remove(this);
      }
    }

    @Override
    public <R> R accept(TermVisitor<R> visitor) {
      return visitor.visit(this);
    }
  }

  /** A pair of terms. */
  public static final class TermTerm {
    final Term left;
    final Term right;

    public TermTerm(Term left, Term right) {
      this.left = requireNonNull(left);
      this.right = requireNonNull(right);
    }

    @Override
    public String toString() {
      return left + " = " + right;
    }
  }

  /**
   * A sequence of terms.
   *
   * <p>A sequence [a b c] is often printed "a(b, c)", as if "a" is the type of
   * node and "b" and "c" are its children.
   */
  public static final class Sequence implements Term {
    public final String operator;
    public final List<Term> terms;

    Sequence(String operator, List<Term> terms) {
      this.operator = requireNonNull(operator);
      this.terms = ImmutableList.copyOf(terms);
    }

    Sequence(String operator) {
      this(operator, ImmutableList.of());
    }

    @Override
    public int hashCode() {
      return Objects.hash(operator, terms);
    }

    @Override
    public boolean equals(Object obj) {
      return this == obj
          || obj instanceof Sequence
              && operator.equals(((Sequence) obj).operator)
              && terms.equals(((Sequence) obj).terms);
    }

    @Override
    public String toString() {
      if (terms.isEmpty()) {
        return operator;
      }
      final StringBuilder builder = new StringBuilder(operator).append('(');
      for (int i = 0; i < terms.size(); i++) {
        Term term = terms.get(i);
        if (i > 0) {
          builder.append(", ");
        }
        builder.append(term);
      }
      return builder.append(')').toString();
    }

    public Term apply(Map<Variable, Term> substitutions) {
      if (terms.isEmpty()) {
        return this;
      }
      final List<Term> newTerms = new ArrayList<>(terms.size());
      for (Term term : terms) {
        newTerms.add(term.apply(substitutions));
      }
      if (listEqual(terms, newTerms)) {
        return this;
      }
      return new Sequence(operator, newTerms);
    }

    @Override
    public Term apply1(Variable variable, Term term) {
      if (terms.isEmpty()) {
        return this;
      }
      if (!contains(variable)) {
        return this;
      }
      final ImmutableList.Builder<Term> newTerms = ImmutableList.builder();
      for (Term term1 : terms) {
        newTerms.add(term1.apply1(variable, term));
      }
      return new Sequence(operator, newTerms.build());
    }

    public void checkCycle(
        Map<Variable, Term> map, Map<Variable, Variable> active)
        throws CycleException {
      for (Term term : terms) {
        term.checkCycle(map, active);
      }
    }

    /**
     * Compares whether two sequences have the same terms. Compares addresses,
     * not contents, to avoid hitting cycles if any of the terms are cyclic
     * (e.g. "X = f(X)").
     */
    private boolean equalsShallow(Sequence sequence) {
      return this == sequence || listEqual(terms, sequence.terms);
    }

    private static <E> boolean listEqual(List<E> list0, List<E> list1) {
      if (list0.size() != list1.size()) {
        return false;
      }
      for (int i = 0; i < list0.size(); i++) {
        if (list0.get(i) != list1.get(i)) {
          return false;
        }
      }
      return true;
    }

    public boolean contains(Variable variable) {
      for (Term term : terms) {
        if (term.contains(variable)) {
          return true;
        }
      }
      return false;
    }

    public <R> R accept(TermVisitor<R> visitor) {
      return visitor.visit(this);
    }
  }

  /** Called on various events during unification. */
  public interface Tracer {
    void onDelete(Term left, Term right);

    void onConflict(Sequence left, Sequence right);

    void onSequence(Sequence left, Sequence right);

    void onSwap(Term left, Term right);

    void onCycle(Variable variable, Term term);

    void onVariable(Variable variable, Term term);

    void onSubstitute(Term left, Term right, Term left2, Term right2);
  }
}

// End Unifier.java
