import torch
from .dataset import PandemicDataset

class PandemicProblem:
    def __init__(self, data: PandemicDataset) -> None:
        """Parent class for all pandemic problem classes. Holding the function, that calculates the residuals of the differential system.

        Args:
            data (PandemicDataset): Dataset holding the time values used.
        """
        #store the gradients for each group
        self.gradients = [torch.zeros((len(data.t_raw), data.number_groups)) for _ in range(data.number_groups)]

        for i in range(data.number_groups):
            self.gradients[i][:, i] = 1

    def to_device(self, device):
        for i in range(len(self.gradients)):
            self.gradients[i] = self.gradients[i].to(device)

    def residual():
        """NEEDS TO BE IMPLEMENTED WHEN INHERITING FROM THIS CLASS
        """
        pass