package solvers;

import ilog.concert.IloException;
import ilog.concert.IloIntVar;
import ilog.concert.IloLinearNumExpr;
import ilog.concert.IloNumVar;
import ilog.cplex.IloCplex;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import xmtrclusters.Problem;

/**
 * MIPModel2 implements a mixed-integer programming model for the clustering
 * problem.
 *
 * It uses the same linearization of the objective that MIPModel uses, but
 * handles cluster creation as an assignment model with some symmetry-breaking
 * constraints.
 *
 * The model winds up being considerably larger than MIPModel and takes several
 * times as long to solve.
 *
 * @author Paul A. Rubin (rubin@msu.edu)
 */
public final class MIPModel2 {
  private static final double HALF = 0.5;  // for rounding binary values

  private final Problem problem;  // the parent problem

  // Model components.
  private final IloCplex mip;     // the MIP model
  private final IloIntVar[][] x;  // x[i][j] = 1 if trans i & j are in the same
                                  // cluster; x[i][i] = 1 if i anchors a cluster
  private final IloIntVar[][] y;  // y[i][j] = 1 if trans i & user j are
                                  // in the same cluster
  private final IloNumVar[][] z;  // z[i][j] is the objective contribution of
                                  // the trans i / user j pairing
  private final IloNumVar[] q;    // q[j] is the solution quality for user j
  private final IloNumVar[][][] v;
    // v[t][u][c] is an indicator for transmitter t and user u both belonging
    // to cluster c

  /**
   * Constructor.
   * @param prob the problem to solve
   * @throws IloException if CPLEX throws an error building the model
   */
  public MIPModel2(final Problem prob) throws IloException {
    problem = prob;
    int nU = problem.getNUsers();
    int nT = problem.getNTrans();
    int nC = problem.getMaxClusters();
    int maxC = problem.getMaxSize();
    double[] qMax = problem.getMaxQuality();
    // Initialize the model.
    mip = new IloCplex();
    // Add the variables.
    x = new IloIntVar[nT][nC];
    y = new IloIntVar[nT][nU];
    z = new IloNumVar[nT][nU];
    q = new IloNumVar[nU];
    v = new IloNumVar[nT][nU][nC];
    for (int t = 0; t < nT; t++) {
      for (int c = 0; c < nC; c++) {
        x[t][c] = mip.boolVar("x_" + t + "_" + c);
      }
      for (int u = 0; u < nU; u++) {
        y[t][u] = mip.boolVar("y_" + t + "_" + u);
        z[t][u] = mip.numVar(0, qMax[u], "z_" + t + "_" + u);
        for (int c = 0; c < nC; c++) {
          v[t][u][c] = mip.numVar(0, 1, "v_" + t + "_" + u + "_" + c);
        }
      }
    }
    for (int u = 0; u < nU; u++) {
      q[u] = mip.numVar(0, qMax[u], "q_" + u);
    }
    // The objective is to maximize the total userQuality.
    mip.addMaximize(mip.sum(q));
    // Every transmitter is assigned to exactly one cluster.
    for (int t = 0; t < nT; t++) {
      mip.addEq(mip.sum(x[t]), 1.0, "assign_" + t + "_once");
    }
    // Observe cluster capacity limits.
    for (int c = 0; c < nC; c++) {
      IloLinearNumExpr expr = mip.linearNumExpr();
      for (int t = 0; t < nT; t++) {
        expr.addTerm(1.0, x[t][c]);
      }
      mip.addLe(expr, maxC, "capacity_" + c);
    }
    // Antisymmetry: Automatically assign transmitter 0 to cluster 0.
    x[0][0].setLB(1.0);
    // Antisymmetry: To assign transmitter t to cluster c, cluster c-1 must
    // contain at least one transmitter with index less than t.
    for (int t = 1; t < nT; t++) {
      for (int c = 1; c < nC; c++) {
        IloLinearNumExpr expr = mip.linearNumExpr();
        for (int t0 = 0; t0 < t; t0++) {
          expr.addTerm(1.0, x[t0][c - 1]);
        }
        mip.addLe(x[t][c], expr, "asym_" + t + "_" + c);
      }
    }
    // Every user is assigned to the cluster containing its best transmitter.
    int[] tau = problem.getBest();
    for (int u = 0; u < nU; u++) {
      mip.addEq(y[tau[u]][u], 1.0, "favorite_transmitter_" + u);
    }
    // Use the auxiliary (v) variables to define y[t][u] when t is not tau[u].
    for (int t = 0; t < nT; t++) {
      for (int u = 0; u < nU; u++) {
        IloLinearNumExpr expr = mip.linearNumExpr();
        for (int c = 0; c < nC; c++) {
          // v[t][u][c] = 1 if t is in c and if u is in c (because tau[u] is
          // in c).
          mip.addLe(v[t][u][c], x[tau[u]][c], "v1_" + t + "_" + u + "_" + c);
          mip.addLe(v[t][u][c], x[t][c], "v2_" + t + "_" + u + "_" + c);
          mip.addGe(v[t][u][c], mip.diff(mip.sum(x[tau[u]][c], x[t][c]), 1.0),
                    "v3_" + t + "_" + u + "_" + c);
          expr.addTerm(1.0, v[t][u][c]);
        }
        // y[t][u] = 1 if v[t][u][c] = 1 for some c.
        mip.addEq(y[t][u], expr, "def_y_" + t + "_" + u);
      }
    }
    // Define the z variables.
    for (int t = 0; t < nT; t++) {
      for (int u = 0; u < nU; u++) {
        mip.addLe(z[t][u], mip.prod(qMax[u], y[t][u]), "zdef1_" + t + "_" + u);
        IloLinearNumExpr expr = mip.linearNumExpr();
        expr.addTerm(1.0, z[t][u]);
        expr.addTerm(-1.0, q[u]);
        expr.addTerm(qMax[u], y[t][u]);
        mip.addLe(expr, qMax[u], "zdef2_" + t + "_" + u);
        expr = mip.linearNumExpr();
        expr.addTerm(1.0, z[t][u]);
        expr.addTerm(-1.0, q[u]);
        expr.addTerm(-qMax[u], y[t][u]);
        mip.addGe(expr, -qMax[u], "zdef3_" + t + "_" + u);
      }
    }
    // Linearize q.
    for (int u = 0; u < nU; u++) {
      IloLinearNumExpr expr = mip.linearNumExpr();
      double[] w = problem.getWeights(u);
      for (int i = 0; i < nT; i++) {
        expr.addTerm(w[i], q[u]);
        expr.addTerm(-w[i], z[i][u]);
        expr.addTerm(-w[i], y[i][u]);
      }
      mip.addEq(expr, 0.0, "qdef_" + u);
    }
  }

  /**
   * Exports the model to a file.
   * @param target the target file path/name
   * @throws IloException if CPLEX encounters an error during export
   */
  public void export(final String target) throws IloException {
    mip.exportModel(target);
  }

  /**
   * Solves the model.
   * @param sec time limit in seconds
   * @return the final objective value
   * @throws IloException if CPLEX encounters an error
   */
  public double solve(final double sec) throws IloException {
    mip.setParam(IloCplex.DoubleParam.TimeLimit, sec);
    mip.setParam(IloCplex.Param.Emphasis.MIP, 2);
    mip.solve();
    return mip.getObjValue();
  }
  /**
   * Gets the transmitter clusters in the final solution.
   * @return the collection of clusters
   * @throws IloException if CPLEX cannot extract the solution.
   */
  public Collection<Collection<Integer>> getClusters() throws IloException {
    // Make sure a solution exists.
    IloCplex.Status status = mip.getStatus();
    if (status != IloCplex.Status.Optimal
        && status != IloCplex.Status.Feasible) {
      throw new IllegalArgumentException("Cannot extract a"
                                         + " nonexistent solution!");
    }
    int nT = problem.getNTrans();
    int nC = problem.getMaxClusters();
    // Extract the values of x[][] in the solution.
    boolean[][] xx = new boolean[nT][nC];
    for (int t = 0; t < nT; t++) {
      for (int c = 0; c < nC; c++) {
        xx[t][c] = mip.getValue(x[t][c]) > HALF;
      }
    }
    // Identify the clusters.
    ArrayList<Collection<Integer>> clusters = new ArrayList<>();
    for (int c = 0; c < nC; c++) {
      HashSet<Integer> cluster = new HashSet<>();
      for (int t = 0; t < nT; t++) {
        if (xx[t][c]) {
          cluster.add(t);
        }
      }
      clusters.add(cluster);
    }
    return clusters;
  }
}