PlagueField.java

import java.util.Random;

/**
 * PlagueField - a stochastic SIR/SIRS epidemic cellular automaton rendered
 * live to the terminal, or run headless for reproducible batch stats.
 *
 * Run with: java PlagueField.java [flags]
 * See README.md for the full flag list.
 */
public class PlagueField {

    enum State { S, I, R }

    /** The grid itself: per-cell epidemic state plus per-cell timers. */
    static final class Grid {
        final int width, height;
        final boolean wrap = true;
        State[][] state;
        int[][] infectedFor;   // ticks spent infected, once state == I
        int[][] recoveredFor;  // ticks spent recovered, once state == R (for waning immunity)
        boolean[][] everInfected; // sticky flag; survives waning so cumulative stats stay correct
        long everInfectedCount;
        final Random rng;
        final double beta;          // per-neighbor infection probability
        final int infectionDuration; // ticks until I -> R
        final int waning;           // ticks until R -> S, 0 = permanent immunity

        Grid(int width, int height, int initialInfected, long seed,
             double beta, int infectionDuration, int waning) {
            this.width = width;
            this.height = height;
            this.beta = beta;
            this.infectionDuration = infectionDuration;
            this.waning = waning;
            this.rng = new Random(seed);
            this.state = new State[height][width];
            this.infectedFor = new int[height][width];
            this.recoveredFor = new int[height][width];
            this.everInfected = new boolean[height][width];
            for (State[] row : state) java.util.Arrays.fill(row, State.S);

            int placed = 0;
            int total = width * height;
            initialInfected = Math.min(initialInfected, total);
            while (placed < initialInfected) {
                int x = rng.nextInt(width);
                int y = rng.nextInt(height);
                if (state[y][x] == State.S) {
                    state[y][x] = State.I;
                    everInfected[y][x] = true;
                    everInfectedCount++;
                    placed++;
                }
            }
        }

        /** Count infected Moore neighbors of (x, y), with optional wraparound. */
        private int infectedNeighbors(int x, int y) {
            int count = 0;
            for (int dy = -1; dy <= 1; dy++) {
                for (int dx = -1; dx <= 1; dx++) {
                    if (dx == 0 && dy == 0) continue;
                    int nx = x + dx, ny = y + dy;
                    if (wrap) {
                        nx = ((nx % width) + width) % width;
                        ny = ((ny % height) + height) % height;
                    } else if (nx < 0 || ny < 0 || nx >= width || ny >= height) {
                        continue;
                    }
                    if (state[ny][nx] == State.I) count++;
                }
            }
            return count;
        }

        /** Advance the whole grid by one tick. Returns true if any cell is still infected. */
        boolean step() {
            State[][] next = new State[height][width];
            int[][] nextInfectedFor = new int[height][width];
            int[][] nextRecoveredFor = new int[height][width];
            boolean anyInfected = false;

            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    State s = state[y][x];
                    switch (s) {
                        case S -> {
                            int k = infectedNeighbors(x, y);
                            double pInfect = 1.0 - Math.pow(1.0 - beta, k);
                            if (k > 0 && rng.nextDouble() < pInfect) {
                                next[y][x] = State.I;
                                if (!everInfected[y][x]) {
                                    everInfected[y][x] = true;
                                    everInfectedCount++;
                                }
                            } else {
                                next[y][x] = State.S;
                            }
                        }
                        case I -> {
                            int dur = infectedFor[y][x] + 1;
                            if (dur >= infectionDuration) {
                                next[y][x] = State.R;
                            } else {
                                next[y][x] = State.I;
                                nextInfectedFor[y][x] = dur;
                            }
                        }
                        case R -> {
                            if (waning > 0) {
                                int dur = recoveredFor[y][x] + 1;
                                if (dur >= waning) {
                                    next[y][x] = State.S;
                                } else {
                                    next[y][x] = State.R;
                                    nextRecoveredFor[y][x] = dur;
                                }
                            } else {
                                next[y][x] = State.R;
                            }
                        }
                    }
                    if (next[y][x] == State.I) anyInfected = true;
                }
            }
            state = next;
            infectedFor = nextInfectedFor;
            recoveredFor = nextRecoveredFor;
            return anyInfected;
        }

        /** Returns {susceptible, infected, recovered} counts. */
        int[] counts() {
            int s = 0, i = 0, r = 0;
            for (State[] row : state) {
                for (State c : row) {
                    switch (c) {
                        case S -> s++;
                        case I -> i++;
                        case R -> r++;
                    }
                }
            }
            return new int[] { s, i, r };
        }

        String render(boolean color) {
            StringBuilder sb = new StringBuilder();
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    char ch = switch (state[y][x]) {
                        case S -> '.';
                        case I -> '#';
                        case R -> 'o';
                    };
                    if (color) {
                        String code = switch (state[y][x]) {
                            case S -> ""; // grey
                            case I -> ""; // red
                            case R -> ""; // blue
                        };
                        sb.append(code).append(ch).append("");
                    } else {
                        sb.append(ch);
                    }
                }
                sb.append('\n');
            }
            return sb.toString();
        }
    }

    static final class Args {
        int width = 60, height = 24;
        int initialInfected = 3;
        long seed = System.nanoTime();
        double beta = 0.22;
        int infectionDuration = 6;
        int waning = 0;
        int ticks = 0;      // 0 = live mode: run until extinction or safety cap
        int delayMs = 90;
        boolean headless = false;
        boolean noColor = false;
        int maxTicks = 2000; // safety cap for live/extinction mode

        static Args parse(String[] argv) {
            Args a = new Args();
            for (int i = 0; i < argv.length; i++) {
                String flag = argv[i];
                String val = (i + 1 < argv.length) ? argv[i + 1] : null;
                switch (flag) {
                    case "--width" -> { a.width = Integer.parseInt(val); i++; }
                    case "--height" -> { a.height = Integer.parseInt(val); i++; }
                    case "--initial-infected" -> { a.initialInfected = Integer.parseInt(val); i++; }
                    case "--seed" -> { a.seed = Long.parseLong(val); i++; }
                    case "--beta" -> { a.beta = Double.parseDouble(val); i++; }
                    case "--duration" -> { a.infectionDuration = Integer.parseInt(val); i++; }
                    case "--waning" -> { a.waning = Integer.parseInt(val); i++; }
                    case "--ticks" -> { a.ticks = Integer.parseInt(val); i++; }
                    case "--delay" -> { a.delayMs = Integer.parseInt(val); i++; }
                    case "--headless" -> a.headless = true;
                    case "--no-color" -> a.noColor = true;
                    case "--help" -> { printHelp(); System.exit(0); }
                    default -> {
                        System.err.println("Unknown flag: " + flag);
                        printHelp();
                        System.exit(1);
                    }
                }
            }
            return a;
        }

        static void printHelp() {
            System.out.println("""
                PlagueField - stochastic SIR/SIRS epidemic cellular automaton

                Usage: java PlagueField.java [flags]

                  --width N              grid width            (default 60)
                  --height N             grid height           (default 24)
                  --initial-infected N   seed infections       (default 3)
                  --seed N               RNG seed              (default random)
                  --beta F               per-neighbor infect chance, 0..1 (default 0.22)
                  --duration N           ticks infected before recovery (default 6)
                  --waning N             ticks until immunity fades, 0=permanent (default 0)
                  --ticks N              run exactly N ticks headless, print summary
                  --delay MS             ms between frames in live mode (default 90)
                  --headless             no animation, print only the final summary
                  --no-color             disable ANSI color in rendered frames
                  --help                 show this message
                """);
        }
    }

    public static void main(String[] args) throws InterruptedException {
        Args a = Args.parse(args);
        Grid grid = new Grid(a.width, a.height, a.initialInfected, a.seed,
                a.beta, a.infectionDuration, a.waning);

        int peakInfected = grid.counts()[1];
        int peakTick = 0;
        int totalCells = a.width * a.height;

        boolean liveMode = a.ticks == 0 && !a.headless;
        int tick = 0;
        boolean infectedRemain = true;

        int limit = a.ticks > 0 ? a.ticks : a.maxTicks;

        while (tick < limit && (a.ticks > 0 || infectedRemain)) {
            if (liveMode) {
                System.out.print("");
                int[] c = grid.counts();
                System.out.printf("PlagueField  tick=%d  S=%d I=%d R=%d%n", tick, c[0], c[1], c[2]);
                System.out.print(grid.render(!a.noColor));
                System.out.flush();
                Thread.sleep(a.delayMs);
            }
            infectedRemain = grid.step();
            tick++;
            int infectedNow = grid.counts()[1];
            if (infectedNow > peakInfected) {
                peakInfected = infectedNow;
                peakTick = tick;
            }
            if (a.ticks == 0 && a.waning == 0 && !infectedRemain) break; // classic SIR: stop once extinct
        }

        int[] finalCounts = grid.counts();
        if (liveMode) {
            System.out.print("");
            System.out.printf("PlagueField  tick=%d  S=%d I=%d R=%d%n", tick, finalCounts[0], finalCounts[1], finalCounts[2]);
            System.out.print(grid.render(!a.noColor));
        } else if (!a.headless) {
            System.out.print(grid.render(!a.noColor));
        }

        long everInfected = grid.everInfectedCount;
        System.out.println("---- SIMULATION COMPLETE ----");
        System.out.printf("seed=%d ticks=%d grid=%dx%d beta=%.3f duration=%d waning=%d%n",
                a.seed, tick, a.width, a.height, a.beta, a.infectionDuration, a.waning);
        System.out.printf("final: S=%d I=%d R=%d (total=%d, conserved=%b)%n",
                finalCounts[0], finalCounts[1], finalCounts[2],
                finalCounts[0] + finalCounts[1] + finalCounts[2],
                finalCounts[0] + finalCounts[1] + finalCounts[2] == totalCells);
        System.out.printf("peak infected: %d at tick %d%n", peakInfected, peakTick);
        System.out.printf("ever infected: %d (%.1f%% of population)%n",
                everInfected, 100.0 * everInfected / totalCells);
        if (a.waning > 0) {
            System.out.printf("mode: SIRS (waning=%d) -- endemic/oscillating, not expected to reach S=0%n", a.waning);
        }
    }
}