package org.sosy_lab.java_smt.example;

import com.google.common.collect.Lists;
import java.util.List;
import java.util.logging.Level;
import org.sosy_lab.common.ShutdownNotifier;
import org.sosy_lab.common.configuration.Configuration;
import org.sosy_lab.common.configuration.InvalidConfigurationException;
import org.sosy_lab.common.log.BasicLogManager;
import org.sosy_lab.common.log.LogManager;
import org.sosy_lab.common.rationals.Rational;
import org.sosy_lab.java_smt.SolverContextFactory;
import org.sosy_lab.java_smt.SolverContextFactory.Solvers;
import org.sosy_lab.java_smt.api.BooleanFormulaManager;
import org.sosy_lab.java_smt.api.IntegerFormulaManager;
import org.sosy_lab.java_smt.api.NumeralFormula.IntegerFormula;
import org.sosy_lab.java_smt.api.OptimizationProverEnvironment;
import org.sosy_lab.java_smt.api.OptimizationProverEnvironment.OptStatus;
import org.sosy_lab.java_smt.api.SolverContext;
import org.sosy_lab.java_smt.api.SolverException;

/**
 * Example for optimizing the weight of some constraints. For a given set of formulas, the weight of
 * the satisfied formulas should be maximal and the weight of unsatisfied formulas should be
 * minimal.
 */
public class OptimizationFormulaWeights {

  public static void main(String... args)
      throws InvalidConfigurationException, SolverException, InterruptedException {
    Configuration config = Configuration.defaultConfiguration();
    LogManager logger = BasicLogManager.create(config);
    ShutdownNotifier notifier = ShutdownNotifier.createDummy();
    Solvers solver = Solvers.Z3; // Z3 works for optimization

    try (SolverContext context =
            SolverContextFactory.createSolverContext(config, logger, notifier, solver);
        OptimizationProverEnvironment prover = context.newOptimizationProverEnvironment(); ) {

      BooleanFormulaManager bmgr = context.getFormulaManager().getBooleanFormulaManager();
      IntegerFormulaManager imgr = context.getFormulaManager().getIntegerFormulaManager();

      optimizeWithWeights(prover, bmgr, imgr, logger);
    }
  }

  /** add some constraints and then solve them with optimization according to weights. */
  private static void optimizeWithWeights(
      OptimizationProverEnvironment prover,
      BooleanFormulaManager bmgr,
      IntegerFormulaManager imgr,
      LogManager logger)
      throws InterruptedException, SolverException {

    // create some symbols and formulas
    IntegerFormula x = imgr.makeVariable("x");
    IntegerFormula y = imgr.makeVariable("y");
    IntegerFormula z = imgr.makeVariable("z");

    IntegerFormula zero = imgr.makeNumber(0);
    IntegerFormula four = imgr.makeNumber(4);
    IntegerFormula ten = imgr.makeNumber(10);

    IntegerFormula scoreBasic = imgr.makeNumber(0);
    IntegerFormula scoreLow = imgr.makeNumber(2);
    IntegerFormula scoreMedium = imgr.makeNumber(4);
    IntegerFormula scoreHigh = imgr.makeNumber(10);

    // add some very important constraints: x<10, y<10, z<10, 10=x+y+z
    prover.addConstraint(
        bmgr.and(
            imgr.lessOrEquals(x, ten), // very important -> direct constraint
            imgr.lessOrEquals(y, ten), // very important -> direct constraint
            imgr.lessOrEquals(z, ten), // very important -> direct constraint
            imgr.equal(ten, imgr.add(x, imgr.add(y, z)))));

    // generate weighted formulas: if a formula should be satisfied,
    // use higher weight for the positive instance than for its negated instance.
    List<IntegerFormula> weights =
        Lists.newArrayList(
            bmgr.ifThenElse(imgr.lessOrEquals(x, zero), scoreHigh, scoreBasic), // important
            bmgr.ifThenElse(imgr.lessOrEquals(x, four), scoreHigh, scoreBasic), // important
            bmgr.ifThenElse(imgr.lessOrEquals(y, zero), scoreMedium, scoreBasic), // less important
            bmgr.ifThenElse(imgr.lessOrEquals(y, four), scoreMedium, scoreBasic), // less important
            bmgr.ifThenElse(imgr.lessOrEquals(x, zero), scoreLow, scoreBasic), // not important
            bmgr.ifThenElse(imgr.lessOrEquals(z, four), scoreHigh, scoreBasic) // important
            );

    // Maximize sum of weights
    int handle = prover.maximize(imgr.sum(weights));

    OptStatus response = prover.check();
    assert response == OptStatus.OPT;

    // for integer theory we get the optimal solution directly as model.
    // ideal solution: sum=30 with e.g. x=0,y=6,z=4
    logger.log(
        Level.INFO,
        "maximal sum ",
        prover.upper(handle, Rational.ZERO).get(),
        "with model",
        prover.getModel());
  }
}

add example for optimization of weights of formulas.