package ri.scenario;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

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

public class ScenarioSIR extends ScenarioSIS {

    private Set<Node> removedNodes;

    public ScenarioSIR(Graph graph, float beta, float mu) {
        super(graph, beta, mu);
        this.removedNodes = new HashSet<>();
    }

    @Override
    public void initialisation() {
        Graph g = getGraph();
        int count = g.getNodeCount();
        int amount = (count / 2) + 1;
        if (amount > count)
            throw new IllegalStateException("Graph must have multiple elements.");
        List<Node> nodes = Toolkit.randomNodeSet(g, amount);
        addInfected(nodes.remove(0));
        this.removedNodes.addAll(nodes);
    }

    public boolean infection(Node node) {
        return super.infection(node) && !this.removedNodes.contains(node);
    }
}