from datetime import datetime
from datetime import timezone
import numpy


class TimeSeriesAccumulator:
    def __init__(self, start_time, end_time, n_bins, alignment=None):
        self.start_time = start_time
        self.start_timestamp = start_time.replace(
            tzinfo=timezone.utc).timestamp()
        self.end_time = end_time
        self.end_timestamp = end_time.replace(tzinfo=timezone.utc).timestamp()
        self.n_bins = n_bins

        if alignment:
            self.start_timestamp = (
                self.start_timestamp//alignment) * alignment
            self.end_timestamp = (self.end_timestamp//alignment) * alignment

        self.bin_time = (self.end_timestamp -
                         self.start_timestamp)/self.n_bins
        self.data = [[] for _ in range(n_bins)]

    def add(self, time, value):
        timestamp = time.replace(tzinfo=timezone.utc).timestamp()
        if timestamp < self.start_timestamp or timestamp > self.end_timestamp:
            return
        idx = int((timestamp - self.start_timestamp)/self.bin_time)
        self.data[idx].append(value)

    def get_f(self, f, initial=0):
        prev = initial
        result = []
        for i in self.data:
            new = f(i)
            if numpy.isnan(new):
                result.append(prev)
            else:
                result.append(new)
                prev = new
        return result

    def get_ts(self):
        return [datetime.utcfromtimestamp(int(self.start_timestamp+(i+1)*self.bin_time)) for i in range(self.n_bins)]


if __name__ == "__main__":
    tsa = TimeSeriesAccumulator(start_time=datetime(
        2020, 7, 27, 23, 0, 33), end_time=datetime(2020, 7, 28, 0, 0, 43), n_bins=2, alignment=5)
    tsa.add(datetime(2020, 7, 27, 23, 15, 0), 1)
    tsa.add(datetime(2020, 7, 27, 23, 45, 0), 2)
    tsa.add(datetime(2020, 7, 27, 22, 45, 0), 3)
    print(tsa.get_ts())
    print(tsa.get_f(lambda x: sum(x)))