# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

import logging
import os
from threading import Event
from typing import Dict, List, Optional, Union

import torch

from .async_schedule import AsyncEventLoop, ModuleWrapper
from .messages import MakeTransport
from .microbatch import Batch
from .skip.layout import SkipLayout
from .skip.tracker import SkipTrackerThroughPotals


class AsyncPipeline:
    """The async pipeline parallelism for Pipe."""

    def __init__(
        self,
        partitions: List[ModuleWrapper],
        skip_layout: SkipLayout,
        checkpoint_stop: int,
        group: torch.distributed.ProcessGroup,
        *,
        worker_map: Optional[Dict[int, str]] = None,
        input_device: Union[None, int, str, torch.device] = None,
        final_stage: bool = False,
    ) -> None:
        self.partitions = partitions
        self.skip_layout = skip_layout
        self.__checkpoint_stop = checkpoint_stop
        self.group = group
        self.training: bool
        self.transport = MakeTransport(
            use_rpc=("OMPI_COMM_WORLD_RANK" not in os.environ) or ("FORCE_RPC" in os.environ),
            worker_map=worker_map,
            input_device=input_device,
        )
        self.input_device = input_device
        self.final_stage = final_stage

    @property
    def checkpoint_stop(self) -> int:
        # Disable checkpointing if in eval mode.
        training = self.partitions[0].module.training
        if not training:
            return 0
        return self.__checkpoint_stop

    def run(self, training: bool, batches: List[Batch], event: Optional[Event]) -> None:

        """Runs pipeline parallelism.

        It modifies the given batches in place.

        """
        self.training = training

        skip_trackers = [SkipTrackerThroughPotals(self.skip_layout, i) for i in range(len(batches))]

        rank = self.group.rank()
        event_loop = AsyncEventLoop(
            self.partitions,
            self.group,
            self.transport,
            self.training,
            self.checkpoint_stop,
        )
        if rank == 0 and not self.final_stage:
            logging.debug(f"{torch.distributed.get_rank()}: entered event head")
            event_loop.event_loop_head(batches, skip_trackers, event)
            logging.debug(f"{torch.distributed.get_rank()}: exited event head")
        elif self.final_stage:
            logging.debug(f"{torch.distributed.get_rank()}: entered event tail")
            event_loop.event_loop_tail(batches, skip_trackers)
            logging.debug(f"{torch.distributed.get_rank()}: exited event tail")
        else:
            logging.debug(f"{torch.distributed.get_rank()}: entered event loop")
            event_loop.event_loop(len(batches), skip_trackers)
            logging.debug(f"{torch.distributed.get_rank()}: exited event loop")

    def back_helper(self, output: List[Batch]) -> None:
        pass