package tpri;

import org.graphstream.algorithm.Toolkit;
import org.graphstream.algorithm.generator.BarabasiAlbertGenerator;
import org.graphstream.algorithm.generator.RandomGenerator;
import org.graphstream.graph.Graph;
import org.graphstream.graph.Node;
import org.graphstream.graph.implementations.SingleGraph;
import org.graphstream.stream.file.FileSourceEdge;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.*;

public class SimuReseauxTheoriques {

    // Paramètres SIS identiques
    private static final double BETA = 0.14;
    private static final double MU   = 0.07;
    private static final int    TMAX = 90;
    private static final int    REPEATS = 10; 

    private enum NetType { DBLP, ER, BA }
    private enum Scenario { NONE, RANDOM_VACCINATION, SELECTIVE_VACCINATION }

    public static void main(String[] args) throws IOException {
        // 1. Chargement et Génération
        Graph dblp = chargerDBLP("com-dblp.ungraph.txt");
        int n = dblp.getNodeCount();
        double k = Toolkit.averageDegree(dblp);

        System.out.println("Génération des réseaux théoriques (N=" + n + ")...");
        Graph er = genererER(n, k);
        Graph ba = genererBA(n, k);

        Map<NetType, Graph> reseaux = new LinkedHashMap<>();
        reseaux.put(NetType.DBLP, dblp);
        reseaux.put(NetType.ER, er);
        reseaux.put(NetType.BA, ba);

        // Stockage des résultats [Réseau][Scénario][Jour]
        double[][][] donnees = new double[NetType.values().length][Scenario.values().length][TMAX];

        // 2. Boucle de simulation
        for (int r = 0; r < REPEATS; r++) {
            System.out.println("Cycle de répétition : " + (r + 1) + "/" + REPEATS);
            
            for (NetType type : NetType.values()) {
                Graph g = reseaux.get(type);
                List<Node> listeNoeuds = g.nodes().toList(); // Cache pour la vitesse

                for (Scenario sce : Scenario.values()) {
                    double[] resultat = lancerSimu(listeNoeuds, sce);
                    for (int t = 0; t < TMAX; t++) {
                        donnees[type.ordinal()][sce.ordinal()][t] += resultat[t];
                    }
                }
            }
        }

        // 3. Écriture du fichier final
        exporterDonnees(donnees);
    }

    private static double[] lancerSimu(List<Node> noeuds, Scenario sce) {
        int n = noeuds.size();
        byte[] etats = new byte[n]; // 0:S, 1:I, 2:V
        Random rng = new Random();

        // Vaccination
        if (sce != Scenario.NONE) {
            int cible = n / 2;
            if (sce == Scenario.RANDOM_VACCINATION) {
                List<Integer> ids = new ArrayList<>(n);
                for(int i=0; i<n; i++) ids.add(i);
                Collections.shuffle(ids);
                for(int i=0; i<cible; i++) etats[ids.get(i)] = 2;
            } else {
                Set<Integer> vax = new HashSet<>();
                while(vax.size() < cible) {
                    Node u = noeuds.get(rng.nextInt(n));
                    List<Node> amis = u.neighborNodes().toList();
                    if(!amis.isEmpty()) vax.add(amis.get(rng.nextInt(amis.size())).getIndex());
                }
                for(int id : vax) etats[id] = 2;
            }
        }

        // Patient zéro
        int p0; do { p0 = rng.nextInt(n); } while (etats[p0] != 0);
        etats[p0] = 1;

        double[] jours = new double[TMAX];
        for (int t = 0; t < TMAX; t++) {
            List<Integer> infections = new ArrayList<>();
            List<Integer> guerisons = new ArrayList<>();
            int inf = 0; int sainsMalades = 0;

            for (int i = 0; i < n; i++) {
                if (etats[i] == 2) continue;
                sainsMalades++;
                if (etats[i] == 1) {
                    inf++;
                    if (rng.nextDouble() < MU) guerisons.add(i);
                    noeuds.get(i).neighborNodes().forEach(v -> {
                        int j = v.getIndex();
                        if (etats[j] == 0 && rng.nextDouble() < BETA) infections.add(j);
                    });
                }
            }
            for (int id : guerisons) etats[id] = 0;
            for (int id : infections) etats[id] = 1;
            jours[t] = (double) inf / sainsMalades;
        }
        return jours;
    }

    // --- Utilitaires de génération et export ---

    private static Graph chargerDBLP(String path) throws IOException {
        Graph g = new SingleGraph("DBLP");
        FileSourceEdge fs = new FileSourceEdge();
        fs.addSink(g);
        fs.readAll(path);
        return g;
    }

    private static Graph genererER(int n, double k) {
        Graph g = new SingleGraph("ER");
        RandomGenerator rg = new RandomGenerator(k);
        rg.addSink(g);
        rg.begin();
        for (int i = 0; i < n; i++) rg.nextEvents();
        rg.end();
        return g;
    }

    private static Graph genererBA(int n, double k) {
        Graph g = new SingleGraph("BA");
        int m = (int) Math.round(k / 2.0);
        BarabasiAlbertGenerator ba = new BarabasiAlbertGenerator(m);
        ba.addSink(g);
        ba.begin();
        for (int i = 0; i < n; i++) ba.nextEvents();
        ba.end();
        return g;
    }

    private static void exporterDonnees(double[][][] d) throws IOException {
        try (PrintWriter pw = new PrintWriter(new FileWriter("virus_compare.dat"))) {
            pw.println("# t DBLP_none DBLP_rand DBLP_sel ER_none ER_rand ER_sel BA_none BA_rand BA_sel");
            for (int t = 0; t < TMAX; t++) {
                pw.print(t);
                for (int net = 0; net < 3; net++) {
                    for (int sce = 0; sce < 3; sce++) {
                        pw.printf(Locale.US, " %.6f", d[net][sce][t] / REPEATS);
                    }
                }
                pw.println();
            }
        }
    }
}