package edu.cmu.ml.rtw.pra.features;

import java.rmi.RemoteException;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import junit.framework.TestCase;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;

import edu.cmu.graphchi.walks.distributions.DiscreteDistribution;
import edu.cmu.ml.rtw.pra.data.NodePairInstance;
import edu.cmu.ml.rtw.pra.experiments.Outputter;
import edu.cmu.ml.rtw.pra.graphs.GraphOnDisk;
import com.mattg.util.FakeFileUtil;
import com.mattg.util.TestUtil;
import com.mattg.util.TestUtil.Function;

public class RandomWalkPathFollowerCompanionTest extends TestCase {

  FakeFileUtil fileUtil;
  GraphOnDisk graph;

  @Override
  public void setUp() {
    fileUtil = new FakeFileUtil();
    fileUtil.addFileToBeRead("/graph/node_dict.tsv", "1\tnode1\n");
    fileUtil.addFileToBeRead("/graph/edge_dict.tsv", "1\trel1\n");
    graph = new GraphOnDisk("/graph/", Outputter.justLogger(), fileUtil);
  }

  public void testAcceptableRow() throws RemoteException {
    Set<Integer> allowedTargets = Sets.newHashSet();
    allowedTargets.add(2);
    allowedTargets.add(3);
    allowedTargets.add(10);
    final RandomWalkPathFollowerCompanion companion = new RandomWalkPathFollowerCompanion(
        graph,
        1,
        1024,
        new FakeVertexIdTranslate(),
        new PathType[1],
        MatrixRowPolicy.ALL_TARGETS,
        allowedTargets,
        true);
    final Set<Integer> sourceTargets = Sets.newHashSet();
    sourceTargets.add(2);
    final Set<Integer> allTargets = Sets.newHashSet(sourceTargets);
    allTargets.add(3);

    companion.setAcceptPolicy(MatrixRowPolicy.EVERYTHING);
    assertEquals(true, companion.acceptableRow(1, 20, sourceTargets, allTargets));

    companion.setAcceptPolicy(MatrixRowPolicy.ALL_TARGETS);
    assertEquals(true, companion.acceptableRow(1, 10, sourceTargets, allTargets));
    assertEquals(false, companion.acceptableRow(1, 20, sourceTargets, allTargets));
    companion.setAllowedTargets(null);
    assertEquals(false, companion.acceptableRow(1, 10, sourceTargets, allTargets));

    companion.setAcceptPolicy(MatrixRowPolicy.PAIRED_TARGETS_ONLY);
    assertEquals(true, companion.acceptableRow(1, 2, sourceTargets, allTargets));
    assertEquals(false, companion.acceptableRow(1, 3, sourceTargets, allTargets));

    companion.setAcceptPolicy(null);
    TestUtil.expectError(RuntimeException.class, new Function() {
      @Override
      public void call() {
        companion.acceptableRow(1, 2, sourceTargets, allTargets);
      }
    });
  }

  public void testGetFeatureMatrix() throws RemoteException {
    final RandomWalkPathFollowerCompanion companion = new RandomWalkPathFollowerCompanion(
        graph,
        1,
        1024,
        new FakeVertexIdTranslate(),
        new PathType[1],
        MatrixRowPolicy.EVERYTHING,  // so we don't have to worry about matching
        null,
        true);
    // We'll have one source, one path type, and two targets in this test.
    ConcurrentHashMap<Integer, ConcurrentHashMap<Integer, DiscreteDistribution>> dists =
        new ConcurrentHashMap<Integer, ConcurrentHashMap<Integer, DiscreteDistribution>>();
    int sourceNode = 1;
    int pathType = 1;
    int targetNode1 = 2;
    int targetNode2 = 3;
    DiscreteDistribution targetNodeDist = new DiscreteDistribution(
        new int[]{targetNode1, targetNode1, targetNode1, targetNode2});
    dists.put(sourceNode, new ConcurrentHashMap<Integer, DiscreteDistribution>());
    dists.get(sourceNode).put(pathType, targetNodeDist);
    companion.setDistributions(dists);

    List<NodePairInstance> instances = Lists.newArrayList();
    instances.add(new NodePairInstance(sourceNode, targetNode1, true, graph));
    instances.add(new NodePairInstance(sourceNode, targetNode2, true, graph));
    FeatureMatrix matrix = companion.getFeatureMatrix(instances);
    assertEquals(2, matrix.size());
    // This is a little complicated to test, because the rows could be in any order.  So we check
    // on the target node for each row.  Both of the row tests below have the same values in the
    // asserts, and one of the if/else blocks should be taken for each row.
    MatrixRow row = matrix.getRow(0);
    NodePairInstance instance = (NodePairInstance) row.instance;
    if (instance.target() == targetNode1) {
      assertEquals(1, row.columns);
      assertEquals(.75, row.values[0]);
    } else {
      assertEquals(targetNode2, instance.target());
      assertEquals(1, row.columns);
      assertEquals(.25, row.values[0]);
    }
    row = matrix.getRow(1);
    instance = (NodePairInstance) row.instance;
    if (instance.target() == targetNode1) {
      assertEquals(1, row.columns);
      assertEquals(.75, row.values[0]);
    } else {
      assertEquals(targetNode2, instance.target());
      assertEquals(1, row.columns);
      assertEquals(.25, row.values[0]);
    }
  }
}
