package models;

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

/**
 * FlowModel implements a MIP model for the problem based on flows between
 * nodes.
 *
 * @author Paul A. Rubin (rubin@msu.edu)
 */
public final class FlowModel implements AutoCloseable {
  private static final double HALF = 0.5;  // for rounding tests
  private final int nVertices;  // number of vertices in the graph
  private final int nEdges;     // number of edges

  // CPLEX objects.
  private final IloCplex mip;        // the MIP model
  private final IloNumVar[] select;  // indicators for selecting vertices
  private final IloNumVar[] flowF;   // forward edge flows
  private final IloNumVar[] flowR;   // reverse edge flows

  /**
   * Constructor.
   * @param graph the graph to solve
   * @param nSelect the number of nodes to select
   * @throws IloException if the model cannot be constructed
   */
  public FlowModel(final Graph graph, final int nSelect) throws IloException {
    nVertices = graph.getnVertices();
    nEdges = graph.getnEdges();
    // Initialize the model.
    mip = new IloCplex();
    // Create the variables.
    select = new IloNumVar[nVertices];
    flowF = new IloNumVar[nEdges];
    flowR = new IloNumVar[nEdges];
    for (int i = 0; i < nVertices; i++) {
      select[i] = mip.boolVar("select_" + i);
    }
    for (int i = 0; i < nEdges; i++) {
      flowF[i] = mip.numVar(0, nVertices, "forward_" + i);
      flowR[i] = mip.numVar(0, nVertices, "backward_" + i);
    }
    // The objective is to minimize the sum of all flows.
    mip.addMinimize(mip.sum(mip.sum(flowF), mip.sum(flowR)));
    // The number of nodes selected must match the requirement.
    mip.addEq(mip.sum(select), nSelect, "selection_count");
    // At every node, there must be exactly one unit more of outflow than of
    // inflow unless the node is selected.
    int m = nVertices - nSelect;  // "big M"
    for (int i = 0; i < nVertices; i++) {
      // Form an expression of the form flow out of i minus flow into i.
      IloLinearNumExpr expr = mip.linearNumExpr();
      // Check each edge incident at i.
      for (int j : graph.incidentAt(i)) {
        if (graph.isHead(i, j)) {
          // The forward direction enters i, the reverse direction leaves i.
          expr.addTerm(1.0, flowR[j]);
          expr.addTerm(-1.0, flowF[j]);
        } else {
          // The forward direction leaves i, the reverse direction enters i.
          expr.addTerm(-1.0, flowR[j]);
          expr.addTerm(1.0, flowF[j]);
        }
      }
      // Flow out - flow in must be >= 1 - m * select[i].
      mip.addGe(expr, mip.diff(1.0, mip.prod(m, select[i])), "balance_" + i);
    }
  }

  /**
   * 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 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();
  }

}