package models;

import ilog.concert.IloException;
import ilog.concert.IloLinearNumExpr;
import ilog.concert.IloNumVar;
import ilog.cplex.IloCplex;
import java.util.HashSet;
import java.util.Set;
import socialnet.Graph;

/**
 * AssignmentModel tries a formulation proposed in an answer to the OR SE post,
 * using binary variables to assign nodes to selected nodes.
 *
 * @author Paul A. Rubin (rubin@msu.edu)
 */
public final class AssignmentModel implements AutoCloseable {
  private static final double HALF = 0.5;  // for rounding tests
  private final int nVertices;  // number of vertices in the graph

  // CPLEX objects.
  private final IloCplex mip;          // the MIP model
  private final IloNumVar[] select;    // indicators for selecting vertices
  private final IloNumVar[][] assign;  // assign[i][j] = 1 if node i is
                                       // assigned to selected node j

  /**
   * Constructor.
   * @param graph the graph to solve
   * @param nSelect the number of nodes to select
   * @param prog print progress reports every this many variables/constraints
   * @throws IloException if the model cannot be constructed
   */
  public AssignmentModel(final Graph graph, final int nSelect, final int prog)
         throws IloException {
    System.out.println("Building the model ...");
    nVertices = graph.getnVertices();
    // Initialize the model.
    mip = new IloCplex();
    // Create the variables and simultaneously construct the objective function.
    select = new IloNumVar[nVertices];
    assign = new IloNumVar[nVertices][nVertices];
    IloLinearNumExpr obj = mip.linearNumExpr();
    long ctr = 0;
    for (int i = 0; i < nVertices; i++) {
      select[i] = mip.boolVar("select_" + i);
      for (int j = 0; j < nVertices; j++) {
        assign[j][i] = mip.boolVar("assign_" + j + "_to_" + i);
        obj.addTerm(graph.getDistance(i, j), assign[j][i]);
        ctr += 1;
        if (ctr % prog == 0) {
          System.out.println("... " + ctr + " assignment variables added");
        }
      }
    }
    // The objective is to minimize the sum of the assigned distances.
    mip.addMinimize(obj);
    // The number of nodes selected must match the requirement.
    mip.addEq(mip.sum(select), nSelect, "selection_count");
    // Every node must be assigned to exactly one selected node (possibly
    // itself).
    for (int i = 0; i < nVertices; i++) {
      mip.addEq(mip.sum(assign[i]), 1.0, "assign_" + i + "_once");
    }
    ctr = 0;
    // Assignments can only be made to selected nodes.
    for (int i = 0; i < nVertices; i++) {
      for (int j = 0; j < nVertices; j++) {
        mip.addLe(assign[i][j], select[j], "select_" + j + "_to_assign_" + i);
        ctr += 1;
        if (ctr % prog == 0) {
          System.out.println("... " + ctr + " assignment constraints added");
        }
      }
    }
  }

  /**
   * Solves the model.
   * @param timeLimit the time limit for the solver (in seconds)
   * @return the final solver status
   * @throws IloException if the solution fails
   */
  public IloCplex.Status solve(final double timeLimit) throws IloException {
    mip.setParam(IloCplex.DoubleParam.TimeLimit, timeLimit);
    mip.solve();
    return mip.getStatus();
  }

  /**
   * Gets the final objective value.
   * @return the objective value
   * @throws IloException if there is no objective value
   */
  public double getObjValue() throws IloException {
    return mip.getObjValue();
  }

  /**
   * Gets the set of nodes selected by the solution.
   * @return the set of selected nodes
   * @throws IloException if no solution exists
   */
  public Set<Integer> getSolution() throws IloException {
    HashSet<Integer> sol = new HashSet<>();
    double[] x = mip.getValues(select);
    for (int i = 0; i < nVertices; i++) {
      if (x[i] > HALF) {
        sol.add(i);
      }
    }
    return sol;
  }

  /**
   * Closes the model.
   */
  @Override
  public void close() {
    mip.close();
  }
}