package iwocs.graphs;

import java.io.FileNotFoundException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map.Entry;
import java.util.Random;
import java.util.Set;
import java.util.stream.Collectors;

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.Edge;
import org.graphstream.graph.Graph;
import org.graphstream.graph.Node;
import org.graphstream.graph.implementations.SingleGraph;

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

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.
   */
  private static PropagationStepInfos simulatePropagationStep(Graph g, String infectedAttribute, double beta,
      double mu) {
    Random rnd = new Random(System.nanoTime());

    Set<Node> infectedNodes = getInfectedNodes(g, infectedAttribute);
    PropagationStepInfos stepInfos = new PropagationStepInfos(infectedNodes.size(),
        g.getNodeCount() - infectedNodes.size());

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

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

          if (!(boolean) neighbour.getAttribute(infectedAttribute)) {
            stepInfos.incrementNewI();
            neighbour.setAttribute(infectedAttribute, (boolean) n.getAttribute(infectedAttribute));
          }
        });
    });

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

      if (muRoll <= mu) {
        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();
  }

  /**
   * Get the set of infected nodes in a graph.
   * 
   * @param g                 the graph.
   * @param infectedAttribute the attribute that tells whether a node is infected
   *                          or not.
   * @return the set of infected nodes.
   */
  public static Set<Node> getInfectedNodes(Graph g, String infectedAttribute) {
    return g.getNodeSet().stream().filter(n -> (boolean) n.getAttribute(infectedAttribute)).collect(Collectors.toSet());
  }

  /**
   * Immunize random nodes in the graph. This is done by removing them.
   * 
   * @param g      the graph to immunize.
   * @param p      the proportion of the graph to immunize.
   * @param method the immunization method.
   * @return the randomly immunized graph.
   */
  public static Graph randomlyImmunize(Graph g, double p, ImmunizationMethod method) {
    for (Node node : Toolkit.randomNodeSet(g, p))
      switch (method) {
      case NODE:
        g.removeNode(Toolkit.randomNode(g));
        break;
      case EDGE:
        Edge randomEdge = Toolkit.randomEdge(node);

        if (randomEdge != null) {
          Node neighbour = randomEdge.getOpposite(node);
          g.removeNode(neighbour);
        }
        break;
      }

    return g;
  }

  public static SelectiveImmunization getSelectiveImmunizationGroups(Graph g, double p) {
    Set<Node> group0 = new HashSet<>(Toolkit.randomNodeSet(g, p));

    Set<Node> group1 = group0.stream().map(n -> {
      Edge randomEdge = Toolkit.randomEdge(n);
      Node neighbour = null;

      if (randomEdge != null)
        neighbour = randomEdge.getOpposite(n);

      return neighbour;
    }).filter(n -> n != null).collect(Collectors.toSet());

    SelectiveImmunization result = new SelectiveImmunization();

    result.setGroup0(group0);
    result.setGroup1(group1);

    return result;
  }

  public static double meanDegreeOfNodeSet(Set<Node> set) {
    return set.stream().map(n -> (double) n.getDegree()).reduce(0.0, Double::sum) / set.size();
  }
}