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

import java.io.IOException;
import java.rmi.RemoteException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;

import edu.cmu.graphchi.ChiLogger;
import edu.cmu.graphchi.ChiVertex;
import edu.cmu.graphchi.EdgeDirection;
import edu.cmu.graphchi.EmptyType;
import edu.cmu.graphchi.datablocks.IntConverter;
import edu.cmu.graphchi.preprocessing.VertexIdTranslate;
import edu.cmu.graphchi.walks.BucketsToSend;
import edu.cmu.graphchi.walks.DrunkardContext;
import edu.cmu.graphchi.walks.DrunkardDriver;
import edu.cmu.graphchi.walks.DrunkardJob;
import edu.cmu.graphchi.walks.DrunkardMobEngine;
import edu.cmu.graphchi.walks.LongDrunkardContext;
import edu.cmu.graphchi.walks.LongDrunkardDriver;
import edu.cmu.graphchi.walks.LongDumperThread;
import edu.cmu.graphchi.walks.LongDrunkardFactory;
import edu.cmu.graphchi.walks.LongWalkArray;
import edu.cmu.graphchi.walks.LongWalkManager;
import edu.cmu.graphchi.walks.WalkArray;
import edu.cmu.graphchi.walks.WalkManager;
import edu.cmu.graphchi.walks.WalkUpdateFunction;
import edu.cmu.ml.rtw.pra.data.NodePairInstance;
import edu.cmu.ml.rtw.pra.graphs.GraphOnDisk;
import com.mattg.util.FileUtil;
import com.mattg.util.MutableConcurrentIndex;
import com.mattg.util.Pair;

/**
 * Finds common path types between pairs of nodes in a given set, using DrunkardMobEngine.
 * @author Matt Gardner
 */
public class RandomWalkPathFinder implements WalkUpdateFunction<EmptyType, Integer> {

  public static final int MAX_HOPS = 10;
  private static final Logger logger = ChiLogger.getLogger("path-finder");
  private final DrunkardMobEngine<EmptyType, Integer>  drunkardMobEngine;
  private final RandomWalkPathFinderCompanion companion;
  private final PathTypeFactory pathTypeFactory;
  private final int numWalksPerSource;
  private final Path[] walkPaths;
  private final int[][][] encodedWalkPaths;
  private final MutableConcurrentIndex<PathType> pathDict;
  private final int[] sourceIds;
  private final List<Integer> origSources;
  private final List<Integer> origTargets;
  private final List<NodePairInstance> instances;
  private final EdgeExcluder edgeExcluder;
  private final VertexIdTranslate vertexIdTranslate;

  private double resetProbability = 0.35;

  public RandomWalkPathFinder(GraphOnDisk graph,
                              List<NodePairInstance> instances,
                              EdgeExcluder edgeExcluder,
                              int walksPerSource,
                              PathTypePolicy policy,
                              PathTypeFactory pathTypeFactory) {
    try {
      this.drunkardMobEngine =
          new DrunkardMobEngine<EmptyType, Integer>(graph.graphFile(), graph.numShards(), new Factory());
      this.drunkardMobEngine.setEdataConverter(new IntConverter());
    } catch (IOException e) {
      e.printStackTrace();
      throw new RuntimeException("IOException when creating DrunkardMobEngine", e);
    }
    vertexIdTranslate = this.drunkardMobEngine.getVertexIdTranslate();
    this.edgeExcluder = edgeExcluder;
    edgeExcluder.prepUnallowedWalks(vertexIdTranslate);
    this.pathTypeFactory = pathTypeFactory;
    this.numWalksPerSource = walksPerSource;
    this.instances = instances;
    origSources = Lists.newArrayList();
    origTargets = Lists.newArrayList();
    for (NodePairInstance instance : instances) {
      origSources.add(instance.source());
      origTargets.add(instance.target());
    }
    this.pathDict = new MutableConcurrentIndex<PathType>(pathTypeFactory, new FileUtil());
    // We add these to a set first, so we don't start twice as many walks from a node that shows up
    // twice in the training data.  You could argue that those nodes should have more influence on
    // the resultant paths, and that's fair, but it also slows down the computation by quite a bit,
    // because the walks aren't evenly distributed, so there is often one thread that takes a long
    // time to finish each iteration.
    Set<Integer> allSourceNodes = new HashSet<Integer>();
    for (int i=0; i<origSources.size(); i++) {
      if (origSources.get(i) > drunkardMobEngine.getEngine().numVertices()) {
        logger.warning("Found a node with a higher index than is in the graph: " + origSources.get(i));
        continue;
      }
      int translatedSource = vertexIdTranslate.forward(origSources.get(i));
      allSourceNodes.add(translatedSource);
      if (origTargets != null) {
        if (origTargets.get(i) > drunkardMobEngine.getEngine().numVertices()) {
          logger.warning("Found a node with a higher index than is in the graph: " + origTargets.get(i));
          continue;
        }
        int translatedTarget = vertexIdTranslate.forward(origTargets.get(i));
        allSourceNodes.add(translatedTarget);
      }
    }

    List<Integer> sources = new ArrayList<Integer>(allSourceNodes);
    Collections.sort(sources);
    try {
      companion = new RandomWalkPathFinderCompanion(graph,
                                                    4,  // numThreads
                                                    Runtime.getRuntime().maxMemory() / 3,
                                                    vertexIdTranslate,
                                                    pathDict,
                                                    pathTypeFactory,
                                                    policy);
    } catch(RemoteException e) {
      throw new RuntimeException(e);
    }
    DrunkardJob drunkardJob =
        this.drunkardMobEngine.addJob("pathFinder", EdgeDirection.IN_AND_OUT_EDGES, this, companion);

    /* Configure walk sources. Note, GraphChi's internal ids are used. */
    drunkardJob.configureWalkSources(sources, numWalksPerSource);
    sourceIds = drunkardJob.getWalkManager().getSources();
    companion.setSources(sourceIds);
    int numWalks = sources.size() * numWalksPerSource;
    walkPaths = new Path[numWalks];
    encodedWalkPaths = new int[numWalks][MAX_HOPS][];
  }

  public void setResetProbability(double resetProbability) {
    this.resetProbability = resetProbability;
  }

  public void execute(int numIters) {
    try {
      drunkardMobEngine.run(numIters);
    } catch(IOException e) {
      throw new RuntimeException(e);
    }
  }

  public Map<PathType, Integer> getPathCounts() {
    return companion.getPathCounts(origSources, origTargets);
  }

  public Map<NodePairInstance, Map<PathType, Integer>> getPathCountMap() {
    return companion.getPathCountMap(instances);
  }

  public Map<NodePairInstance, Map<PathType, Set<Pair<Integer, Integer>>>> getLocalSubgraphs() {
    return companion.getLocalSubgraphs(instances);
  }

  public void shutDown() {
    logger.info("Path Finder shutting down");
    companion.close();
    pathDict.clear();
  }

  public int[] encodePath(Path path) {
    PathType[] encodedPathTypes = pathTypeFactory.encode(path);
    int[] encodedPaths = new int[encodedPathTypes.length];
    for (int i = 0; i < encodedPaths.length; i++) {
      encodedPaths[i] = pathDict.getIndex(encodedPathTypes[i]);
    }
    return encodedPaths;
  }

  @VisibleForTesting
  protected MutableConcurrentIndex<PathType> getPathDictionary() {
    return pathDict;
  }

  @VisibleForTesting
  protected void setEncodedWalkPaths(int[] encoded, int walkId, int hopNum) {
    encodedWalkPaths[walkId][hopNum] = encoded;
  }

  @VisibleForTesting
  protected void setWalkPath(Path path, int walkId) {
    walkPaths[walkId] = path;
  }

  @VisibleForTesting
  protected Path getWalkPath(int walkId) {
    return walkPaths[walkId];
  }

  /**
   * WalkUpdateFunction interface implementations
   */
  @Override
  public void processWalksAtVertex(WalkArray walkArray,
                                   ChiVertex<EmptyType, Integer> chiVertex,
                                   DrunkardContext drunkardContext_,
                                   Random random) {
    long[] walks = ((LongWalkArray)walkArray).getArray();
    LongDrunkardContext drunkardContext = (LongDrunkardContext) drunkardContext_;
    Vertex vertex = new Vertex(chiVertex, false);

    // If there are no edges connected to this node, we just give up.  No need to restart any
    // walks, as they'll just come back here and waste our time again.  If this happens, if means
    // you gave a source (or target) node that was completely disconnected from the rest of the
    // graph.
    if (vertex.getNumEdges() == 0) {
      return;
    }

    // Advance each walk to a random edge.
    for (long walk : walks) {
      processSingleWalkAtVertex(walk, vertex, drunkardContext, random);
    }
  }

  /**
   * Take a single walk, and advance it.
   *
   * We take a bunch of parameters here that we probably don't need to, but it's in the hope of
   * reducing the cost of computation in this method, because it gets called a _lot_.
   */
  @VisibleForTesting
  protected void processSingleWalkAtVertex(long walk,
                                           Vertex vertex,
                                           LongDrunkardContext drunkardContext,
                                           Random random) {
    int walkId = Manager.walkId(walk);
    int sourceVertex = sourceIds[staticSourceIdx(walk)];

    // Reset?
    if (random.nextDouble() < resetProbability) {
      resetWalk(walk, walkId, drunkardContext);
      return;
    }

    Path path = walkPaths[walkId];
    if (path != null && path.getHops() >= MAX_HOPS - 1) {
      resetWalk(walk, walkId, drunkardContext);
      return;
    }
    if (path == null) {
      path = new Path(sourceIds[staticSourceIdx(walk)], MAX_HOPS);
      walkPaths[walkId] = path;
    }
    boolean reverse;
    int edgeNum = random.nextInt(vertex.getNumEdges());
    if (edgeNum < vertex.getNumInEdges()) {
      reverse = true;
    } else {
      reverse = false;
    }
    int nextVertex  = vertex.getEdgeNode(edgeNum);
    // Prohibit cyclical walks
    if (path.alreadyVisited(nextVertex)) {
      resetWalk(walk, walkId, drunkardContext);
      return;
    }
    int edgeType = vertex.getEdgeType(edgeNum);
    if (edgeExcluder.shouldExcludeEdge(sourceVertex, nextVertex, vertex, edgeType)) {
      resetWalk(walk, walkId, drunkardContext);
      return;
    }
    path.addHop(nextVertex, edgeType, reverse);
    walk = Manager.incrementHopNum(walk);
    int hopNum = Manager.hopNum(walk);
    encodedWalkPaths[walkId][hopNum] = encodePath(path);
    drunkardContext.forwardWalkTo(walk, nextVertex, true);
  }

  private void resetWalk(long walk, int walkId, LongDrunkardContext drunkardContext) {
    walkPaths[walkId] = null;
    walk = Manager.resetHopNum(walk);
    drunkardContext.resetWalk(walk, true);
  }

  /**
   * We don't want to ignore any vertices for this walk.
   */
  @Override
  public int[] getNotTrackedVertices(ChiVertex<EmptyType, Integer> vertex) {
    int[] notCounted = new int[0];
    return notCounted;
  }

  /////////////////////////////////////////////////////////////////////////
  // Some boilerplate to modify the behavior of DrunkardMobEngine slightly
  /////////////////////////////////////////////////////////////////////////

  @VisibleForTesting
  protected long[] encodeWalkForCompanion(long walk) {
    int walkId = Manager.walkId(walk);
    int hopNum = Manager.hopNum(walk);
    int[] walkPathTypes = encodedWalkPaths[walkId][hopNum];
    if (walkPathTypes == null) {
      return new long[0];
    }
    int numPathTypes = walkPathTypes.length;
    long[] encodedWalks = new long[numPathTypes];
    for (int i = 0; i < numPathTypes; i++) {
      // Replace the walk id with the path type - the companion needs to be able to get
      // (source, target, path) triples.  Our regular encoding already has source and target, so we
      // just need to put path type in.
      int pathType = walkPathTypes[i];
      encodedWalks[i] = Manager.encodeForCompanion(pathType,  // path
                                                   staticSourceIdx(walk),  // source
                                                   staticTrackBit(walk),
                                                   staticOff(walk));  // target
    }
    return encodedWalks;
  }

  class Factory extends LongDrunkardFactory<EmptyType, Integer> {
    public WalkManager createWalkManager(int numVertices, int numSources) {
      return new Manager(numVertices, numSources);
    }
    public DrunkardDriver<EmptyType, Integer> createDrunkardDriver(DrunkardJob job,
                                                                   WalkUpdateFunction<EmptyType,
                                                                   Integer> callback) {
      return new Driver(job, callback);
    }
  }

  class Driver extends LongDrunkardDriver<EmptyType, Integer> {
    Driver(final DrunkardJob job, WalkUpdateFunction<EmptyType, Integer> callback) {
      super(job, callback);
    }

    @Override
    protected LongDumperThread createDumperThread() {
      return new DumperThread(bucketQueue, pendingWalksToSubmit, finished, job);
    }

    class DumperThread extends LongDumperThread {

      DumperThread(LinkedBlockingQueue<BucketsToSend> bucketQueue,
                          AtomicLong pendingWalksToSubmit,
                          AtomicBoolean finished,
                          DrunkardJob job) {
        super(bucketQueue, pendingWalksToSubmit, finished, job);
      }

      @Override
      protected void processWalks(BucketsToSend bucket, int i) {
        LongWalkArray bucketWalks = (LongWalkArray) bucket.walks;
        long w = bucketWalks.getArray()[i];
        LongWalkManager manager = (LongWalkManager) job.getWalkManager();
        int v = manager.off(w) + bucket.firstVertex;


        // Skip walks with the track-bit not set
        boolean trackBit = manager.trackBit(w);

        if (!trackBit) {
          return;
        }

        // The only difference with LongDrunkardDriver.LongDumperThread
        for (long encodedWalk : encodeWalkForCompanion(w)) {
          walks[idx] = encodedWalk;
          vertices[idx] = v;
          idx++;

          if (idx >= walks.length) {
            try {
              job.getCompanion().processWalks(new LongWalkArray(walks), vertices);
            } catch (Exception err) {
              err.printStackTrace();
            }
            idx = 0;
          }
        }
      }
    }
  }


  static class Manager extends LongWalkManager {
    private int nextWalkId;
    public static final int MAX_ENCODABLE_WALKS = (int) Math.pow(2, 27);
    public static final int MAX_ENCODABLE_HOPS = (int) Math.pow(2, 5);
    public static final int MAX_SOURCES = (int) Math.pow(2, 24);

    Manager(int numVertices, int numSources) {
      super(numVertices, numSources);
    }

    /**
     * We take the standard (sourceId, hop, off) encoding from IntWalkManager and put some walk
     * stuff in front of it - 27 bits for the walk id, and 5 bits for the hop number.
     *
     * Because we're using the same lower 32 bits, also, we don't have to override the sourceIdx,
     * trackBit, and off methods of the base LongWalkManager, which doesn't use the top 32 bits.
     */
    static long encode(int walkId, int hopNum, int sourceId, boolean track, int off) {
      if (off >= 128) throw new AssertionError("off too high: " + off);
      if (walkId >= MAX_ENCODABLE_WALKS) throw new AssertionError("walkId too high: " + walkId);
      if (hopNum >= MAX_ENCODABLE_HOPS) throw new AssertionError("hopNum too high: " + hopNum);
      if (sourceId >= MAX_SOURCES) throw new AssertionError("sourceId too high: " + sourceId);
      int trackBit = (track ? 1 : 0);
      return ((long) walkId & 0x7ffffff) << 37 |
          ((long)(hopNum & 0x1f)) << 32 |
          (long)((sourceId & 0xffffffL) << 8) |
          (long)((off & 0x7f) << 1) | (long)trackBit;
    }

    static long encodeForCompanion(int pathType, int sourceId, boolean track, int off) {
      assert(off < 128);
      int trackBit = (track ? 1 : 0);
      return ((long) pathType) << 32 |
          (long)((sourceId & 0xffffffL) << 8) |
          (long)((off & 0x7f) << 1) | (long)trackBit;
    }

    public static int walkId(long walk) {
      return (int) (walk >> 37) & 0x7ffffff;
    }

    public static int hopNum(long walk) {
      return (int) (walk >> 32) & 0x1f;
    }

    public static long incrementHopNum(long walk) {
      int hopNum = hopNum(walk);
      return encode(walkId(walk),
                    hopNum+1,
                    staticSourceIdx(walk),
                    staticTrackBit(walk),
                    staticOff(walk));
    }

    public static long resetHopNum(long walk) {
      return encode(walkId(walk), 0, staticSourceIdx(walk), staticTrackBit(walk), staticOff(walk));
    }

    public static int pathType(long walkEncodedForCompanion) {
      return (int) (walkEncodedForCompanion >> 32);
    }

    @Override
    protected long encode(int sourceId, boolean hop, int off) {
      throw new RuntimeException("This implementation needs a walk id!");
    }

    @Override
    protected long encodeV(int sourceId, boolean hop, int vertexId) {
      throw new RuntimeException("Use encodeNewWalk instead of this method, so we can do " +
                                 "the walkId correctly");
    }

    @Override
    protected long encodeNewWalk(int sourceId, int sourceVertex, boolean hop) {
      return encode(nextWalkId++, 0, sourceId, hop, sourceVertex % bucketSize);
    }

    @Override
    protected long reencodeWalk(long walk, int toVertex, boolean trackBit) {
      int walkId = walkId(walk);
      int hopNum = hopNum(walk);
      return encode(walkId, hopNum, sourceIdx(walk), trackBit, toVertex % bucketSize);
    }
  }

  // These next three are for convenience, so that Path.encode doesn't need an instance of
  // Manager.  Kind of ugly, but oh well...
  public static int staticSourceIdx(long walk) {
    return (int) ((walk & 0xffffff00) >> 8) & 0xffffff;
  }

  public static boolean staticTrackBit(long walk) {
    return ((walk & 1) != 0);
  }

  public static int staticOff(long walk) {
    return (int) (walk >> 1) & 0x7f;
  }

  public static long setTrackBit(long walk, boolean trackBit) {
    if (walk % 2 == 0) {
      // Track bit not set
      if (trackBit) {
        return walk + 1;
      } else {
        return walk;
      }
    } else {
      // Track bit is currently set
      if (trackBit) {
        return walk;
      } else {
        return walk - 1;
      }
    }
  }
}
