package iwocs.graphs;

import java.io.FileNotFoundException;
import java.util.HashMap;
import java.util.List;
import java.util.Random;
import java.util.Map.Entry;

import org.graphstream.algorithm.Toolkit;
import org.graphstream.algorithm.generator.BarabasiAlbertGenerator;
import org.graphstream.algorithm.generator.Generator;
import org.graphstream.algorithm.generator.RandomGenerator;
import org.graphstream.graph.BreadthFirstIterator;
import org.graphstream.graph.Graph;
import org.graphstream.graph.Node;
import org.graphstream.graph.implementations.SingleGraph;

import iwocs.graphs.propagation.PropagationInfos;
import iwocs.graphs.propagation.PropagationStepInfos;

public class GraphUtils {
  /**
   * Handy method that generates a random graph with a given average degree.
   * 
   * @param id identifier of the resulting graph.
   * @param n the number of nodes in the resulting graph.
   * @param avgDegree the average degree of the resulting graph.
   * @return the resulting random graph.
   */
  public static Graph getRandomGraph(String id, int n, double avgDegree) {
    Generator gen = new RandomGenerator(avgDegree);

    Graph result = new SingleGraph(id);

    gen.addSink(result);

    gen.begin();

    for (int i = 0; i < n; i++)
      gen.nextEvents();

    gen.end();

    return result;
  }

  /**
   * Handy method that generates a preferential
   * graph with a given average degree.
   * 
   * @param id identifier of the resulting graph.
   * @param n the number of nodes in the resulting graph.
   * @param avgDegree the average degree of the resulting graph.
   * @return the resulting preferential graph.
   */
  public static Graph getPreferentialGraph(String id, int n, double avgDegree) {
    int m = (int)Math.round(avgDegree);

    Generator gen = new BarabasiAlbertGenerator(m);

    Graph result = new SingleGraph(id);

    gen.addSink(result);

    gen.begin();

    System.out.printf("Barabasi gen(n=%d, m=%d)\n", n, m);

    for (int i = 0; i < n; i++) {
      System.out.printf("\r(%d/%d)", i, n);
      gen.nextEvents();
    }

    gen.end();

    return result;
  }

  /**
   * Calculate a mean distance from a sample of nodes from a given graph.
   * 
   * @param g          the graph to analyse.
   * @param sampleSize number of nodes to include in the analysis.
   * @return mean distance calculated from the sample.
   */
  public static double calculateMeanDistanceFromSample(Graph g, int sampleSize) {
    List<Node> sampleNodes = Toolkit.randomNodeSet(g, sampleSize);

    double meanSum = 0;

    for (Node n : sampleNodes) {
      System.out.format("\r(%d/%d)", sampleNodes.indexOf(n), sampleSize);
      double localSum = 0;
      int nbVisited;
      BreadthFirstIterator<Node> iter = (BreadthFirstIterator<Node>) n.getBreadthFirstIterator();

      for (nbVisited = 0; iter.hasNext(); nbVisited++)
        localSum += iter.getDepthOf(iter.next());

      meanSum += localSum / nbVisited;
    }

    System.out.println();
    return meanSum / sampleSize;
  }

  /**
   * Creates a distribution of distances from a random set of nodes from a given
   * graph of a given size.
   * 
   * @param g          the graph to analyse.
   * @param sampleSize number of nodes to analyse.
   * @return the distribution calculated.
   */
  public static HashMap<Integer, Double> calculateDistancesDistributionFromSample(Graph g, int sampleSize, boolean normalized) {
    List<Node> sampleNodes = Toolkit.randomNodeSet(g, sampleSize);

    HashMap<Integer, Double> distribution = new HashMap<>();

    int i = 0;

    for (Node n : sampleNodes) {
      System.out.format("\r(%d/%d)", ++i, sampleSize);
      BreadthFirstIterator<Node> iter = (BreadthFirstIterator<Node>) n.getBreadthFirstIterator();

      while (iter.hasNext()) {
        int depth = iter.getDepthOf(iter.next());
        Double depthCount = distribution.get(depth);

        if (depthCount == null)
          distribution.put(depth, 1.0);
        else
          distribution.put(depth, depthCount + 1);
      }
    }

    System.out.println();

    if(normalized) {
      Double totalMeasures = distribution.values().stream().reduce(0.0, Double::sum);

      distribution.forEach((k, v) -> distribution.put(k, v / totalMeasures));
    }

    return distribution;
  }

  /**
   * Calculates variance of a probabilities distribution.
   * 
   * @param distributionFile path to the distribution file.
   * @param isResource file is resource or not.
   * @return the variance calculated from the distribution.
   * @throws FileNotFoundException if the file isn't found.
   */
  public static Double calculateVarianceFromDistribution(String distributionFile, boolean isResource)
  throws FileNotFoundException {
    HashMap<Double, Double> distribution = GraphIO.getDistributionFromFile(distributionFile, isResource);

    Double variance = 0.0;

    for(Entry<Double, Double> e : distribution.entrySet())
      variance += e.getKey() * e.getKey() * e.getValue();

    return variance;
  }

  /**
   * Creates a patient zero in a graph,
   * for propagation simulation purposes.
   * 
   * @param g the graph in which we insert a patient zero.
   * @param infectedAttribute the attribute on a node to know if it's infected.
   * @return patient zero node.
   */
  public static Node makePatientZero(Graph g, String infectedAttribute) {
    g.forEach(n -> n.setAttribute(infectedAttribute, false));

    Node zero = Toolkit.randomNode(g);

    zero.setAttribute(infectedAttribute, true);

    return zero;
  }

  /**
   * Simulates a given number of propagation steps in a graph.
   * 
   * @param g the graph in which we run the simulation.
   * @param infectedAttribute the attribute which tells if a node is infected.
   * @param beta the infection probability.
   * @param mu the recovery probability.
   * @param nbSteps the number of steps to simulate.
   * @return the infos about the simulation.
   */
  public static PropagationInfos simulatePropagation(Graph g, String infectedAttribute, double beta, double mu, int nbSteps) {
    PropagationInfos propagationInfos = new PropagationInfos();

    if(!isInfected(g, infectedAttribute))
      makePatientZero(g, infectedAttribute);

    for(int i=0; i < nbSteps; i++)
      propagationInfos.addStep(simulatePropagationStep(g, infectedAttribute, beta, mu));

    return propagationInfos;
  }

  /**
   * Simulates a propagation step in a graph.
   * 
   * @param g the graph in which we run the simulation.
   * @param infectedAttribute the attribute which tells if a node is infected.
   * @param beta the infection probability.
   * @param mu the recovery probability.
   * @return infos about this step.
   */
  public static PropagationStepInfos simulatePropagationStep(Graph g, String infectedAttribute, double beta, double mu) {
    Random rnd = new Random(System.nanoTime());

    PropagationStepInfos stepInfos = new PropagationStepInfos(getNbI(g, infectedAttribute), getNbS(g, infectedAttribute));

    // Infection phase
    g.forEach(n -> {
      double betaRoll = rnd.nextDouble();

      if(betaRoll <= beta && (boolean)n.getAttribute(infectedAttribute))
        n.forEach(e -> {
          Node neighbour = e.getOpposite(n);

          if(!(boolean)neighbour.getAttribute(infectedAttribute))
            stepInfos.incrementNewI();

          neighbour.setAttribute(infectedAttribute, n.getAttribute(infectedAttribute));
        });
    });

    // Recovering phase
    g.forEach(n -> {
      double muRoll = rnd.nextDouble();

      if(muRoll <= mu && (boolean)n.getAttribute(infectedAttribute)) {
        stepInfos.incrementNewS();
        n.setAttribute(infectedAttribute, false);
      }
    });

    return stepInfos;
  }

  /**
   * Tells if a graph is infected.
   * 
   * @param g the graph to test.
   * @param infectedAttribute the attribute that tells whether 
   * a node is infected or not.
   * @return whether the given node is infected or not.
   */
  private static boolean isInfected(Graph g, String infectedAttribute) {
    return g
      .getNodeSet()
      .stream()
      .anyMatch(n -> 
        n.hasAttribute(infectedAttribute) && 
        (boolean)n.getAttribute(infectedAttribute)
      );
  }

  /**
   * Tells if a node is infected.
   * 
   * @param g the node to test.
   * @param infectedAttribute the attribute that tells whether
   * a node is infected or not.
   * @return whether the given node is infected or not.
   */
  private static boolean isInfected(Node n, String infectedAttribute) {
    return n.hasAttribute(infectedAttribute) &&
      (boolean)n.getAttribute(infectedAttribute);
  }

  /**
   * Get the number of infected nodes in the graph.
   * 
   * @param g the graph to analyse.
   * @param infectedAttribute the attribute that
   * tells whether a node is infected or not.
   * @return ths number of infected nodes.
   */
  public static long getNbI(Graph g, String infectedAttribute) {
    return g.getNodeSet()
      .stream()
      .filter(n -> isInfected(n, infectedAttribute))
      .count();
  }

  /**
   * Get the number of susceptible nodes in the graph.
   * 
   * @param g the graph to analyse.
   * @param infectedAttribute the attribute that
   * tells whether a node is infected or not.
   * @return ths number of susceptible nodes.
   */
  public static long getNbS(Graph g, String infectedAttribute) {
    return g.getNodeSet()
      .stream()
      .filter(n -> !isInfected(n, infectedAttribute))
      .count();
  }
}