# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
from collections import deque
from enum import Enum

import paddle
from paddle.base import log_helper

from .graphs import CUDAGraph

# CUDAGraphedLayer Debug tools
enable_debug_print = bool(
    int(os.getenv('PADDLE_DEBUG_ENABLE_CUDAGRAPH_LAYER_LOGGING', '0'))
)
debug_cudagraphedlayer_fallback_to_default = bool(
    int(os.getenv('PADDLE_DEBUG_CUDAGRAPHEDLAYER_FALLBACK_TO_DEFAULT', '0'))
)

logger = log_helper.get_logger(
    __name__, logging.INFO, fmt='[%(levelname)s] %(message)s'
)


def debug_print(x):
    if not enable_debug_print:
        return
    logger.info(x)


def print_tensor(
    t,
    name="Unnamed",
    print_meta=True,
    print_ptr=False,
    print_hash=True,
    hash=None,
):
    output = []
    if name:
        output.append(name)
    if hash is None:
        hash = lambda t: float((t.astype('float32') * 1000).sum())

    if t is None:
        debug_print(f"{name} is None")
    elif isinstance(t, paddle.Tensor):
        if print_meta:
            output.append(f"shape = {t.shape}")
            output.append(f"place = {t.place}")
        if print_ptr:
            output.append(f"ptr = {hex(t.data_ptr())}")
        if print_hash:
            output.append(f"hash = {hash(t)}")
        debug_print(" | ".join(output))


def printer(x, banner="printer"):
    if not enable_debug_print:
        return
    debug_print(banner.center(100, "-"))
    recursive_apply(print_tensor, x)


# We need this function, for any kind of inputs with iterables
# we recursively apply the function to the leave nodes
def recursive_apply(function, input_var):
    if isinstance(input_var, list):
        return [recursive_apply(function, item) for item in input_var]
    elif isinstance(input_var, tuple):
        return tuple(recursive_apply(function, item) for item in input_var)
    elif isinstance(input_var, dict):
        return {
            key: recursive_apply(function, value)
            for key, value in input_var.items()
        }
    else:
        return function(input_var)


def detach_tensor(tensor):
    # Detach an individual tensor and preserve its 'stop_gradient' property
    if isinstance(tensor, paddle.Tensor):
        detached_tensor = tensor.detach()
        detached_tensor.stop_gradient = tensor.stop_gradient
        return detached_tensor
    return tensor


# We try our best to flatten the input to list of tensors
# example: args = ((t1,t2),(t3,(t4,t5))) -> [t1, t2, t3, t4, t5]
def recursive_flatten(target):
    ret = []

    def append(arg):
        if isinstance(arg, paddle.Tensor):
            # [NOTE] sometimes unnecessary tensors, such as the constant `mask` tensor in the PP layer, is passed into subsequent layers.
            # When a tensor is marked with `stop_gradient=True`, it indicates that it does not contribute to gradient calculations,
            # suggesting it's unrelated to the main computational process.
            # Therefore, I try to eliminate the copying of such tensors in the to optimize performance.
            # if not arg.stop_gradient:
            # [NOTE] However, `stop_gradient=True` propagation rules within the framework appear to be flawed, so directly eliminate stop_gradient may cause bug
            ret.append(arg)

    recursive_apply(append, target)
    return ret


# input any kind of args / kwargs structure, output list of tensor
def recursive_flatten_args_kwargs(args, kwargs):
    return [
        *recursive_flatten(args),
        *recursive_flatten(tuple(kwargs.values())),
    ]


detach = lambda x: recursive_apply(detach_tensor, x)


def get_grad_tensor(x):
    """Returns the gradient of a Paddle Tensor if it's a tensor; otherwise, returns the input."""
    if isinstance(x, paddle.Tensor):
        if x.stop_gradient:
            return None
        else:
            return x.grad
    return None


# CUDA Graph with Static Input and Output
class CUDAGraphWithStaticInputOutput:
    def __init__(self, num_warmup_steps):
        self.num_warmup_steps = num_warmup_steps
        self.graph = CUDAGraph()

        self.has_recorded = False

        self.has_preserved_inputs = False
        self.args_static = None
        self.kwargs_static = None

        # inputs is the recursively flattened args and kwargs
        self.inputs_static = None
        self.outputs_static = None

    def preserve_or_copy(self, args, kwargs):
        """
        For the CUDA Graph, it is crucial that the buffer remains address-stable,
        meaning that the buffer addresses for any inputs to the CUDA Graph should not change.
        One solution to achieve this is to preserve all input tensors.

        This function attempts to recursively flatten the input arguments and keyword arguments
        to identify all tensors passed to the layer (though it may still miss some due to other implicit
        ways inputs can be passed to a layer). It then preserves references to these input tensors
        as `self.inputs_static` so that the buffer pointers can be reused later.

        When this method is called subsequently, it copies the values back to the preserved input tensors
        to ensure the buffers are reused.
        """
        if not self.has_preserved_inputs:
            self.has_preserved_inputs = True
            self.args_static = args
            self.kwargs_static = kwargs
            self.inputs_static = recursive_flatten_args_kwargs(
                self.args_static, self.kwargs_static
            )
        else:
            inputs = recursive_flatten_args_kwargs(args, kwargs)
            for x_staic, x in zip(self.inputs_static, inputs):
                x_staic.copy_(x, True)

    def record(self, f, *args, **kwargs):
        self.preserve_or_copy(args, kwargs)

        self.graph.capture_begin()
        self.outputs_static = f(*self.args_static, **self.kwargs_static)
        self.graph.capture_end()
        debug_print(
            "[CUDAGraph] Record-Replay Start (Graph is replayed for the first time)"
        )
        self.graph.replay()

        self.has_recorded = True

        return self.outputs_static

    def set_output_static(self, outputs_static):
        self.outputs_static = outputs_static

    def replay(self, *args, **kwargs):
        if not self.has_recorded:
            raise RuntimeError("Graph should be recorded first")

        self.preserve_or_copy(args, kwargs)

        debug_print("[CUDAGraph] Replay Start")
        self.graph.replay()
        return self.outputs_static

    def save(self, name):
        logging.info(f"save graph to {name}")
        self.graph.print_to_dot_files(name)


# CUDA Graph Layer Status Enumeration
class CUDAGraphLayerStatus(Enum):
    """Enum to represent the status of a CUDA Graph Layer."""

    WARMUP = 1
    RECORD = 2
    CUDAGRAPH = 3


class CUDAGraphForwardBackward:
    def __init__(self, num_warmup_steps):
        self.forward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
        self.backward_graph = CUDAGraphWithStaticInputOutput(num_warmup_steps)
        self.status = CUDAGraphLayerStatus.RECORD

    def capture_end(self):
        self.status = CUDAGraphLayerStatus.CUDAGRAPH

    def is_record_step(self):
        return self.status == CUDAGraphLayerStatus.RECORD

    def is_cuda_graph_step(self):
        return self.status == CUDAGraphLayerStatus.CUDAGRAPH


class CUDAGraphContext:
    """
    Manages the context for CUDA graph execution in layers. This includes handling
    the state of CUDA graph layers, managing forward and backward graphs, and
    tracking the execution steps.
    """

    def __init__(self, layer, num_warmup_steps):
        """
        Initializes the CUDA graph context.
        :param layer: The layer to be used in the CUDA graph.
        :param num_warmup_steps: Number of warmup steps before recording starts.
        """
        self.layer = layer
        self.num_warmup_steps = num_warmup_steps

        # The state of context is in either WARMUP or CUDAGRAPH
        self._step = 0
        self.status = CUDAGraphLayerStatus.WARMUP

        # Queue to support 1f1b/interleaved scheduler, assuming FIFO order
        # data queue
        self.data_queue = deque()

        # graph queue
        self.graph_queue = deque()

    # Graph Operations
    def get_graph(self):
        if len(self.graph_queue) == 0:
            return CUDAGraphForwardBackward(self.num_warmup_steps)
        else:
            return self.graph_queue.popleft()

    def reuse_graph(self, g):
        self.graph_queue.append(g)

    # Tensor Queue Operations
    def push_data(self, args):
        self.data_queue.append(args)

    def pop_data(self):
        return self.data_queue.popleft()

    # Finite State Machine of Layer State
    def warmup_step(self):
        self._step += 1
        if self._step == self.num_warmup_steps:
            self.status = CUDAGraphLayerStatus.CUDAGRAPH

    def is_warmup_step(self):
        return self.status == CUDAGraphLayerStatus.WARMUP

    def is_cuda_graph_step(self):
        return self.status == CUDAGraphLayerStatus.CUDAGRAPH


def select_y_with_grad(ys, dys):
    # [TODO] when there is multiple output tensor, we support only one y that allows backward
    y, dy = None, None
    if isinstance(ys, paddle.Tensor):
        y, dy = ys, dys[0]
    elif isinstance(ys, (list, tuple)):
        for v, dv in zip(ys, dys):
            if isinstance(v, paddle.Tensor) and (not v.stop_gradient):
                y, dy = v, dv
                break
    assert isinstance(y, paddle.Tensor) and isinstance(dy, paddle.Tensor)
    return y, dy


# we get the output of the backward from the detached inputs after the backward is calculated
# we save it to the graph itself
def get_args_grad(inputs):
    grad_inputs, detached_grad_inputs = inputs
    args_grad = []
    for x, detached_x in zip(grad_inputs, detached_grad_inputs):
        # if required grad
        if not x.stop_gradient:
            if detached_x.grad is None:
                # if input requires grad but we don't have grad, we just allocate some zeros
                # x.stop_gradient = True
                args_grad.append(paddle.zeros(detached_x.shape))
                # args_grad.append(None)
            else:
                args_grad.append(detached_x.grad)
        else:
            args_grad.append(None)
    return tuple(args_grad)


class _CUDAGraphedLayer(paddle.autograd.PyLayer):
    """
    A custom layer that integrates CUDA Graph recording and execution into PaddlePaddle's autograd system.
    It handles forward and backward operations differently based on the CUDA graph layer status.
    """

    @staticmethod
    def forward(ctx, context, arg_tuple, *grad_inputs):
        """
        Handles the forward pass of the layer. It operates differently based on the
        context's status: warmup, recording, or CUDA graph step.
        """
        args, kwargs = arg_tuple
        # Detach all inputs from the computational graph
        args = detach(args)
        kwargs = detach(kwargs)

        detached_grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
        inputs = (grad_inputs, detached_grad_inputs)

        printer(detached_grad_inputs, "Forward input")
        if (
            context.is_warmup_step()
            or debug_cudagraphedlayer_fallback_to_default
        ):
            debug_print("[CUDAGraph] Forward Step (Default)")

            with paddle.enable_grad():
                y = context.layer(*args, **kwargs)

            context.push_data((CUDAGraphLayerStatus.WARMUP, None, inputs, y))
        else:
            graph = context.get_graph()
            if graph.is_record_step():
                # In record step, record the forward pass in CUDA graph
                debug_print(f"[CUDAGraph] Forward Step (Record) id {id(graph)}")

                def forward(*args, **kwargs):
                    with paddle.enable_grad():
                        return context.layer(*args, **kwargs)

                y = graph.forward_graph.record(forward, *args, **kwargs)

                context.push_data(
                    (CUDAGraphLayerStatus.RECORD, graph, inputs, y)
                )
            else:
                debug_print(f"[CUDAGraph] Forward Step (Graph) id {id(graph)}")
                y = graph.forward_graph.replay(*args, **kwargs)

                context.push_data(
                    (CUDAGraphLayerStatus.CUDAGRAPH, graph, None, y)
                )

        debug_print("[CUDAGraph] Forward Step End")

        ctx.save_for_backward(context)
        printer(y, "Forward output")
        return detach(y)

    @staticmethod
    def backward(ctx, *dys):
        """
        Handles the backward pass of the layer. Similar to forward, it handles
        backward based on the context's status: warmup, record, or CUDAGraph.
        """
        (context,) = ctx.saved_tensor()

        (status, graph, inputs, ys) = context.pop_data()
        y, dy = select_y_with_grad(ys, dys)

        printer((y, dy), "Backward input")

        if status == CUDAGraphLayerStatus.WARMUP:
            debug_print("[CUDAGraph] Backward Step (Default)")

            # In warmup step, perform standard backward operation
            y.backward(dy)
            args_grad = get_args_grad(inputs)

            context.warmup_step()
        elif status == CUDAGraphLayerStatus.RECORD:
            debug_print(f"[CUDAGraph] Backward Step (Record) id {id(graph)}")

            # In record step, record the backward pass in CUDA graph
            def backward(y, dy):
                y.backward(dy)

            graph.backward_graph.record(backward, y, dy)

            # [NOTE] the get_args_grad should not put inside backward
            # the args_grad should be calculated after graph is replayed
            args_grad = get_args_grad(inputs)
            graph.backward_graph.set_output_static(args_grad)
            graph.capture_end()

            context.reuse_graph(graph)
        elif status == CUDAGraphLayerStatus.CUDAGRAPH:
            debug_print(f"[CUDAGraph] Backward Step (Graph) id {id(graph)}")

            # In CUDA graph step, replay the recorded graph for backward pass
            args_grad = graph.backward_graph.replay(y, dy)
            context.reuse_graph(graph)
        else:
            raise RuntimeError("Unknown cuda graph status")

        debug_print("[CUDAGraph] Backward Step End")

        printer(args_grad, "Backward output")
        return args_grad


class CUDAGraphedLayer(paddle.nn.Layer):
    """
    CUDAGraphedLayer: A PaddlePaddle Layer to convert an eager mode model to utilize CUDA Graphs.

    CUDA Graphs provide a way to capture kernel-level operations of a model and play
    them back efficiently, allowing for potential speedups in repetitive computations,
    such as those during training iterations. This layer is a wrapper that enables
    the usage of CUDA Graphs with PaddlePaddle models.

    Overview:
    - The layer encapsulates another layer (the model to be converted).
    - During the first few (num_warmup_steps) iterations, the layer operates in
      eager mode without any CUDA Graphs.
    - After the warmup steps, the layer captures the forward and backward computations
      and replays them using CUDA Graphs in subsequent iterations.

    Usage:
        model = Model()
        graphed_model = CUDAGraphedLayer(model)

    Parameters:
    - layer (paddle.nn.Layer): The PaddlePaddle model/layer to be converted.
    - num_warmup_steps (int): The number of iterations before the CUDA Graph
      capture begins. Default is 3.

    Notes:
    - Restrictions:
        * CPU-GPU Synchronization: Operations that synchronize the CPU with the GPU, like device to host transfers, are not allowed.
        * CPU Work: Any operations on the CPU within the captured graph are not recorded.
        * Memory Address (Pointer) Consistency: Replays consistently read from and write to identical virtual memory addresses.
        * Dynamic Operations:
            - Control Flow: Dynamic control flows, especially those based on CPU data like if/else statements, are prohibited.
            - Tensor Shapes: Dynamic tensor shapes are not supported.

    - Allowed Operations:
        * CUDA RNG Operations: CUDA-based Random Number Generation operations are allowed.
    """

    def __init__(self, layer: paddle.nn.Layer, num_warmup_steps=3):
        super().__init__()
        self.context = CUDAGraphContext(layer, num_warmup_steps)
        self.add_sublayer(f"Graphed {type(layer).__name__}", layer)

    def forward(self, *args, **kwargs):
        # We collect them into a list of tensor that required grad
        grad_inputs = recursive_flatten_args_kwargs(args, kwargs)
        return _CUDAGraphedLayer.apply(
            self.context, (args, kwargs), *grad_inputs
        )

    def is_warmup_step(self):
        return self.context.is_warmup_step()

    def is_cuda_graph_step(self):
        return self.context.is_cuda_graph_step()
