package ri.scenario;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import org.graphstream.graph.Graph;
import org.graphstream.graph.Node;

public abstract class ScenarioSIS {

    private Graph graph;
    private float beta, mu;
    private Set<Node> infectedNodes;

    public ScenarioSIS(Graph graph, float beta, float mu) {
        this.graph = graph;
        this.beta = beta;
        this.mu = mu;
        this.infectedNodes = new HashSet<>();
    }

    public Graph getGraph() {
        return this.graph;
    }

    public float getBeta() {
        return this.beta;
    }

    public float getMu() {
        return this.mu;
    }
    
    public Collection<Node> getInfectedNodes() {
        return new ArrayList<>(this.infectedNodes);
    }

    public void addInfected(Node node) {
        this.infectedNodes.add(node);
    }

    public void removeInfected(Node node) {
        this.infectedNodes.remove(node);
    }

    public abstract void initialisation();

    public Map<Integer, Collection<Node>> propagation(int days) {
        initialisation();
        Map<Integer, Collection<Node>> result = new HashMap<>(days);
        for (int i = 0; i < days; i++)
            result.put(i, nextDay());
        return result;
    }

    public Collection<Node> nextDay() {
        Collection<Node> infectedNodes = new ArrayList<>(this.infectedNodes);
        infectedNodes.forEach(this::sendMail);
        return infectedNodes;
    }

    public void sendMail(Node node) {
        Iterator<Node> uIte = node.getNeighborNodeIterator();
        while (uIte.hasNext()) {
            Node u = uIte.next();
            if (infection(u)) {
                addInfected(u);
            }
        }
        if (remission(node)) {
            removeInfected(node);
        }
    }

    public boolean infection(Node node) {
        return Math.random() < getBeta();
    }

    public boolean remission(Node node) {
        return Math.random() < getMu();
    }
}