Epoch

class EpochDataMap(dict):
    """
    A class for managing epoch data mapping, including timestamps and associated images.

    Args:
        image_dir (Union[str, Path]): The directory containing image data.
        master_camera (str): The name of the master camera (optional).
        time_tolerance_sec (int): The maximum time difference allowed to consider two images taken by different cameras as simultaneous (this allows for considering a non-perfect time synchronization between different cameras). Default is 180 seconds (3 minutes).
        min_images (int): The minimum number of images required for an epoch to be included (optional).

    Attributes:
        _image_dir (Path): The path to the image directory.
        _master_camera (str): The name of the master camera.
        _timetolerance (timedelta): The time tolerance for timestamp matching.
        _cams (List[str]): The list of camera names.
        _map (dict): The mapping of epoch data.

    Methods:
        __init__(self, image_dir: Union[str, Path], master_camera=None, time_tolerance_sec: timedelta = 180, min_images: int = 2):
            Initialize the EpochDataMap with image directory, master camera, and time tolerance.
        __getitem__(self, key):
            Retrieve epoch data by key.
        __repr__(self) -> str:
            Return a string representation of the EpochDataMap.
        __len__(self) -> int:
            Return the number of epochs in the EpochDataMap.
        __iter__(self):
            Initialize an iterator for the EpochDataMap.
        __next__(self):
            Get the next element in the EpochDataMap.
        __contains__(self, timestamp: Union[str, dt]) -> bool:
            Check if a timestamp is present in the EpochDataMap.
        cams(self) -> List[str]:
            Return the list of camera names.
        get_timestamp(self, epoch_id: int) -> dt:
            Get the timestamp of a specific epoch.
        get_timestamp_str(self, epoch_id: int) -> str:
            Get the timestamp of a specific epoch as a formatted string.
        get_images(self, epoch_id: int) -> List[Path]:
            Get the images associated with a specific epoch.
        get_images_by_timestamp(self, timestamp: Union[str, dt]) -> List[Path]:
            Get the images associated with an epoch by timestamp.
        _get_timestamps(self, folder: Union[str, Path]) -> Tuple[List[dt], List[Path]]:
            Get timestamps and image paths from a specified folder.
        _build_map(self, min_images: int = None):
            Build the mapping of epoch data.
        _write_map(self, filename: str, sep: str = ",", header: bool = True) -> None:
            Write the mapping data to a CSV file.
    """

    def __init__(
        self,
        image_dir: Union[str, Path],
        master_camera=None,
        time_tolerance_sec: timedelta = 180,
        min_images: int = 2,
    ):
        """
        Initialize the EpochDataMap with image directory, master camera, and time tolerance.
        """
        self._image_dir = Path(image_dir)
        assert self._image_dir.exists(), f"{self._image_dir} does not exist"

        # if a master camera is specified, check if it exists in the image dir
        if master_camera:
            assert self._master_camera in os.scandir(
                self._image_dir
            ), f"Master camera not found in image directory {self._image_dir}"
        # if master camera is not specified, use the first camera found
        else:
            master_camera = sorted(
                [f.name for f in os.scandir(self._image_dir) if f.is_dir()]
            )[0]
        self._master_camera = master_camera

        self._timetolerance = timedelta(seconds=time_tolerance_sec)

        # Get camera names
        self._cams = sorted([f.name for f in os.scandir(self._image_dir) if f.is_dir()])

        # Build dict
        self._map = {}
        self._build_map(min_images)

        # Write dict to file
        self._write_map(self._image_dir / "epoch_map.csv")

    def __getitem__(self, key):
        return self._map[key]

    def __repr__(self) -> str:
        return f"{self.__class__.__name__} with {len(self._map)} epochs"

    def __len__(self):
        return len(self._map)

    def __iter__(self):
        self._elem = 0
        return self

    def __next__(self):
        while self._elem < len(self._map):
            file = self._map[self._elem]
            self._elem += 1
            return file
        else:
            self._elem = 0
            raise StopIteration

    def __contains__(self, timestamp: Union[str, dt]) -> bool:
        timestamp = parse_str_to_datetime(timestamp)
        timestamps = [x["timestamp"] for x in self._map.values()]
        return timestamp in timestamps

    @property
    def cams(self):
        return self._cams

    def get_timestamp(self, epoch_id: int) -> dt:
        return self._map[epoch_id]["timestamp"]

    def get_timestamp_str(self, epoch_id: int) -> str:
        return self._map[epoch_id]["timestamp"].strftime(DATETIME_FMT)

    def get_images(self, epoch_id: int) -> List[Path]:
        return self._map[epoch_id]["images"]

    def get_images_by_timestamp(self, timestamp: Union[str, dt]) -> List[Path]:
        timestamp = parse_str_to_datetime(timestamp)
        timestamps = [x["timestamp"] for x in self._map.values()]
        idx = timestamps.index(timestamp)
        return self._map[idx]["images"]

    def _get_timestamps(self, folder: Union[str, Path]) -> Tuple[List[dt], List[Path]]:
        imageDS = ImageDS(folder)
        paths = list(imageDS.files)
        timestamps = list(imageDS.timestamps.values())
        return timestamps, paths

    def _build_map(self, min_images: int = None):
        # Re-initialize map
        self._map = {}

        # Build imageDS for master camera and get timestamps
        timestamps, paths = self._get_timestamps(self._image_dir / self._master_camera)

        # build mapping dict for master camera
        for i, (ts, path) in enumerate(zip(timestamps, paths)):
            img = Image(path)
            self._map[i] = DotDict(
                {
                    "timestamp": ts,
                    "images": {self._master_camera: img},
                    "dt": {self._master_camera: timedelta(seconds=0)},
                }
            )

        # Find closest timestamp for each camera
        slave_cameras = set(self._cams)
        slave_cameras.discard(self._master_camera)

        for cam in slave_cameras:
            timestamps1, paths1 = self._get_timestamps(self._image_dir / cam)
            for key, value in self._map.items():
                ref_ts = value["timestamp"]
                _, closest_idx, dt = find_closest_timestamp(
                    ref_ts, timestamps1, self._timetolerance
                )
                if closest_idx is not None:
                    self._map[key]["images"][cam] = Image(paths1[closest_idx])
                    self._map[key]["dt"][cam] = dt
        logger.info(f"Building EpochDataMap: found {len(self._map)} epochs")

        # Get max dt for each epoch
        self._max_dt_sec = {
            ep: max(list(x["dt"].values())).seconds for ep, x in self._map.items()
        }
        dtmax = np.array([x for x in self._max_dt_sec.values()])
        logger.info(
            f"Mean max dt: {np.mean(dtmax):.2f} seconds (max: {np.max(dtmax):.2f} seconds))"
        )

        # Remove epochs with less than min_images images
        removed = 0
        if min_images is not None:
            for key, value in self._map.copy().items():
                if len(value["images"]) < min_images:
                    del self._map[key]
                    removed += 1
            logger.info(f"Removed {removed} epochs with less than {min_images} images")

    def _write_map(self, filename: str, sep: str = ",", header: bool = True) -> None:
        file = open(filename, "w")
        if header:
            columns = ["epoch", "date", "time"]
            for cam in self._cams:
                columns.append(cam)
                columns.append(f"{cam}_timestamp")
            file.write(f"{sep}".join(columns) + "\n")
        for key, value in self._map.items():
            value = self._map[key]
            date = value["timestamp"].strftime("%Y-%m-%d")
            time = value["timestamp"].strftime("%H:%M:%S")
            str_2_add = []
            for cam in self._cams:
                str_2_add.append(value["images"][cam].name)
                str_2_add.append(
                    f"{value['images'][cam].date}_{value['images'][cam].time}"
                )
            line = [str(key), date, time] + str_2_add
            file.write(f"{sep}".join(line) + "\n")
        file.close()

class Epoch:
    """
    Class for storing, saving, and reading ICEpy4D solution at one epoch

    Attributes:
        cameras (classes.Camera): The dictionary of camera parameters
        images (classes.ImagesDict): The dictionary of images and their metadata
        features (classes.Features): The dictionary of feature points
        points (classes.Points): The dictionary of 3D points
    """

    def __init__(
        self,
        timestamp: Union[str, dt],
        epoch_dir: Union[str, Path] = None,
        images: Dict[str, Image] = None,
        cameras: Dict[str, Camera] = None,
        features: Dict[str, Features] = None,
        points: Points = None,
        targets: Targets = None,
        point_cloud: PointCloud = None,
        datetime_format: str = DATETIME_FMT,
    ) -> None:
        """
        Initializes a Epcoh object with the provided data

        Args:
            cameras (classes.Camera): The dictionary of camera parameters
            images (classes.ImagesDict): The dictionary of images and their metadata
            features (classes.Features): The dictionary of feature points
            points (classes.Points): The dictionary of 3D points

        TODO: Simplify the initialization of the Epoch object if no parameters are provided.
        """

        self._timestamp = parse_str_to_datetime(timestamp, datetime_format)

        if images is not None:
            assert (
                isinstance(images, dict) and isinstance(x, Image)
                for x in images.values()
            ), "Input images must be a dictionary mapping each camera name to the image (of type Image)"
            self.images = images
        self.cameras = cameras
        self.features = features
        self.targets = targets
        self.point_cloud = point_cloud

        if points is not None:
            assert isinstance(points, Points), "Input points must be of type Points"
            self.points = points
        else:
            self.points = Points()

        if epoch_dir is not None:
            self.epoch_dir = Path(epoch_dir)
        else:
            logger.info("Epoch directory not provided. Using epoch timestamp.")
            self.epoch_dir = Path(str(self._timestamp).replace(" ", "_"))
        self.epoch_dir.mkdir(parents=True, exist_ok=True)

    @property
    def timestamp(self):
        return self._timestamp

    @property
    def date_str(self) -> str:
        """
        Returns the date and time of the epoch in a string.

        Returns:
            str: The date and time of the epoch in the format "YYYY:MM:DD HH:MM:SS".
        """
        return self._timestamp.strftime("%Y:%m:%d")

    @property
    def time_str(self) -> str:
        """
        Returns the time of the epoch as a string.

        Returns:
            str: The time of the epoch in the format "HH:MM:SS".

        """
        return self._timestamp.strftime("%H:%M:%S")

    def __str__(self) -> str:
        """
        Returns a string representation of the Epoch object

        Returns:
            str: The string representation of the Epoch object
        """
        return f"{self._timestamp.strftime(DATETIME_FMT).replace(' ', '_')}"

    def __repr__(self):
        """
        Returns a string representation of the Epoch object

        Returns:
            str: The string representation of the Epoch object
        """
        return f"Epoch {self._timestamp}"

    def __iter__(self):
        """
        Returns an iterator over the four dictionaries of Epoch object

        Yields:
            dict: The dictionary of camera parameters
            dict: The dictionary of images and their metadata
            dict: The dictionary of feature points
            dict: The dictionary of 3D points
        """
        yield self.cameras
        yield self.images
        yield self.features
        yield self.points

    def __hash__(self):
        """
        Computes the hash value of the Epoch object

        Returns:
            int: The hash value of the Epoch object
        """
        return hash((self.cameras, self.images, self.features, self.points))

    def save_pickle(self, path: Union[str, Path]) -> bool:
        """
        Saves the Epoch object to a binary file

        Args:
            path (Union[str, Path]): The path to the binary file

        Returns:
            bool: True if the object was successfully saved to file, False otherwise
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        try:
            with open(path, "wb") as f:
                pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
            return True
        except:
            logger.error("Unable to save the Solution as Pickle object")
            return False

    @staticmethod
    def read_pickle(path: Union[str, Path]):
        """
        Load a Epoch object from a binary file

        Args:
            path (Union[str, Path]): The path to the binary file

        Returns:
            Epoch: An Epoh object
        """
        path = Path(path)
        assert path.exists(), f"Input path {path} does not exists"

        try:
            logger.info(f"Loading epoch from {path}")
            with open(path, "rb") as inp:
                epoch = pickle.load(inp)
            if epoch is not None:
                logger.info(f"Epoch loaded from {path}")
                return epoch
            else:
                logger.error(f"Unable to load epoch from {path}")
                return None
        except Exception as e:
            raise e(f"Unable to read Epoch from file {path}")

class Epoches:
    """Class for storing all the epochs in ICEpy4D processing"""

    def __init__(self, starting_epoch: int = 0) -> None:
        self._starting_epoch = starting_epoch
        self._last_epoch: int = -1
        self._epochs: Dict[int, Epoch] = {}
        self._epoches_map: Dict[int, dt] = {}
        self._elem = 0

    def __repr__(self):
        """
        Returns a string representation of the Epoches object

        Returns:
            str: The string representation of the Epoches object
        """

        return f"Epoches with {len(self._epochs)} epochs"

    def __len__(self) -> int:
        """Get number of epoches in the Epoches object"""
        return len(self._epochs)

    def __iter__(self):
        self._elem = self._starting_epoch
        return self

    def __next__(self):
        while self._elem <= self._last_epoch:
            file = self._epochs[self._elem]
            self._elem += 1
            return file
        else:
            self._elem
            raise StopIteration

    def __getitem__(self, epoch_id):
        """
        Returns the epoch object with the provided epoch_id

        Args:
            epoch_id (int): The numeric key of the epoch

        Returns:
            Epoch: The epoch object
        """
        return self._epochs[epoch_id]

    def __contains__(
        self, epoch_date: Union[str, dt], datetime_format: str = DATETIME_FMT
    ) -> bool:
        """Check if an epoch is in the Epoch objet"""
        timestamp = parse_str_to_datetime(epoch_date, datetime_format)
        timestamps = [x for x in self._epoches_map.values()]
        return timestamp in timestamps

    def add_epoch(self, epoch: Epoch):
        """
        Adds an epoch to the Epoches object

        Args:
            epoch_id (int): The numeric key of the epoch
            epoch_date (str): The corresponding date of the epoch
            epoch (Epoch): The epoch object to be added
        """
        assert isinstance(epoch, Epoch), "Input epoch must be of type Epoch"
        assert hasattr(epoch, "timestamp"), "Epoch must have a timestamp attribute"
        assert isinstance(
            epoch.timestamp, dt
        ), "Epoch timestamp must be of type datetime"
        assert (
            epoch.timestamp not in self._epoches_map.values()
        ), "Epoch with the same timestamp already exists"
        if self._last_epoch == -1:
            epoch_id = self._starting_epoch
        else:
            epoch_id = self._last_epoch + 1
        self._epoches_map[epoch_id] = epoch.timestamp
        self._epochs[epoch_id] = epoch
        self._last_epoch = epoch_id

    def get_epoch_date(self, epoch_id: int) -> str:
        """
        Retrieves the date corresponding to an epoch from the Epoches object

        Args:
            epoch_id (int): The numeric key of the epoch

        Returns:
            str: The date corresponding to the epoch_id
        """
        return self._epochs.get(epoch_id).timestamp

    def get_epoch_id(
        self, epoch_date: Union[str, dt], datetime_format: str = DATETIME_FMT
    ) -> int:
        timestamp = parse_str_to_datetime(epoch_date, datetime_format)
        for i, ep in self._epochs.items():
            if ep.timestamp == timestamp:
                return i

    def get_epoch_by_date(
        self, timestamp: Union[str, dt], datetime_format: str = DATETIME_FMT
    ) -> Epoch:
        timestamp = parse_str_to_datetime(timestamp, datetime_format)
        for ep in self._epochs.values():
            if ep.timestamp == timestamp:
                return ep
        logger.warning(f"Epoch with timestamp {timestamp} not found")
        return None