1from .recorder import Recorder
 2from aegis_sim.dataclasses.population import Population
 3
 4
 5class PopsizeRecorder(Recorder):
 6    def __init__(self, odir):
 7        self.odir = odir
 8        self._buffers = {}  # filename -> list of lines
 9
10    def write(self, popsize: int, filename: str):
11        if filename not in self._buffers:
12            self._buffers[filename] = []
13        self._buffers[filename].append(str(popsize))
14        # Flush every 100 entries to avoid unbounded memory growth
15        if len(self._buffers[filename]) >= 100:
16            self._flush(filename)
17
18    def _flush(self, filename):
19        if filename in self._buffers and self._buffers[filename]:
20            path = self.odir / filename
21            with open(path, "a") as file_:
22                file_.write("\n".join(self._buffers[filename]) + "\n")
23            self._buffers[filename] = []
24
25    def flush_all(self):
26        """Flush all buffered data to disk."""
27        for filename in list(self._buffers.keys()):
28            self._flush(filename)
29
30    def write_before_reproduction(self, population):
31        """
32        # OUTPUT SPECIFICATION
33        path: /popsize_before_reproduction.csv
34        filetype: csv
35        description: Number of live individuals before reproduction.
36        category: demography
37        time granularity: every step
38        frequency parameter: N/A
39        structure: A vector of integers.
40        header: None
41        """
42        self.write(len(population), "popsize_before_reproduction.csv")
43
44    def write_after_reproduction(self, population):
45        """
46        # OUTPUT SPECIFICATION
47        path: /popsize_after_reproduction.csv
48        filetype: csv
49        description: Number of live individuals after reproduction.
50        category: demography
51        time granularity: every step
52        frequency parameter: N/A
53        structure: A vector of integers.
54        header: None
55        """
56        self.write(len(population), "popsize_after_reproduction.csv")
57
58    def write_egg_num_after_reproduction(self, eggs):
59        """
60        # OUTPUT SPECIFICATION
61        path: /eggnum_after_reproduction.csv
62        filetype: csv
63        description: Number of eggs which have not yet hatched.
64        category: demography
65        time granularity: every step
66        frequency parameter: N/A
67        structure: A vector of integers.
68        header: None
69        """
70        eggnum = len(eggs) if eggs is not None else 0
71        self.write(eggnum, "eggnum_after_reproduction.csv")