import dataclasses from datetime import datetime from decimal import Decimal from typing import Generator import matplotlib.pyplot @dataclasses.dataclass(kw_only=True, frozen=True) class Flow: """Time-discrete flow of money paid on the first day of a month""" amount: Decimal since: None | datetime until: None | datetime def integrate(self, start: datetime, end: datetime) -> Decimal: """Integrate the flow between two dates to an amount of money""" payments: int = 0 if self.since is not None: if start < self.since: start = self.since if self.until is not None: if end > self.until: end = self.until for candidate in monthly_candidates(start): if start <= candidate <= end: payments += 1 if candidate > end: break return self.amount * Decimal(payments) def monthly_candidates(start: datetime) -> Generator[datetime, None, None]: current = datetime(start.year, start.month, 1) while True: yield current if current.month == 12: current = datetime(current.year + 1, 1, 1) else: current = datetime(current.year, current.month + 1, 1) def simulate( start: datetime, end: datetime, flows: tuple[Flow, ...] ) -> tuple[list[datetime], list[Decimal]]: dates: list[datetime] = [] values: list[Decimal] = [] for candidate in monthly_candidates(start): if start <= candidate <= end: dates.append(candidate) if candidate > end: break for date in dates: value = Decimal(0.0) for flow in flows: value += flow.integrate(start, date) values.append(value) return (dates, values) def display(data: tuple[list[datetime], list[Decimal]]) -> None: matplotlib.pyplot.plot( list(data[0]), # type: ignore [float(i) for i in data[1]], label="Money", ) matplotlib.pyplot.xlabel("Time") matplotlib.pyplot.legend() matplotlib.pyplot.xticks(rotation=45) matplotlib.pyplot.show()