from abc import ABC, abstractmethod
from typing import Dict, List, Tuple, Union

import torch
from torch.fx.node import Node

from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
    OperationData,
    OperationDataType,
    ShardingSpec,
    ShardingStrategy,
    StrategiesVector,
    TrainCycleItem,
)
from colossalai.auto_parallel.tensor_shard.utils import check_sharding_spec_validity
from colossalai.device.device_mesh import DeviceMesh
from colossalai.tensor.shape_consistency import ShapeConsistencyManager

from .strategy import StrategyGenerator


class NodeHandler(ABC):
    '''
    The NodeHandler is an abstract class used to generate every possible strategies for an operator node.

    Args:
        node (Node): the input node in node argument list.
        device_mesh (DeviceMesh): A logical view of a physical mesh.
        strategies_vector (StrategiesVector): all the strategies generated in this handler will be recorded into the strategies_vector.
    '''

    def __init__(
        self,
        node: Node,
        device_mesh: DeviceMesh,
        strategies_vector: StrategiesVector,
    ) -> None:
        self.node = node
        self.predecessor_node = list(node._input_nodes.keys())
        self.successor_node = list(node.users.keys())
        self.device_mesh = device_mesh
        self.strategies_vector = strategies_vector

    def update_resharding_cost(self, strategy: ShardingStrategy) -> None:
        """
        Compute the resharding costs and save the costs in the ShardingStrategy object.
        """
        # TODO: test this function when other handlers are ready
        resharding_costs = {}
        shape_consistency_manager = ShapeConsistencyManager()

        for node in self.predecessor_node:
            node_name = str(node)
            # get the current sharding spec generated by this node handler

            # we will not compute the resharding costs for the node not counted in the strategy.
            # And the node with tuple or list output need to be handled below.
            node_in_strategy = [op_data.name for op_data in strategy.sharding_specs.keys()]
            if str(node) not in node_in_strategy:
                continue

            op_data = strategy.get_op_data_by_name(node_name)
            current_sharding_spec = strategy.sharding_specs[op_data]
            # get the sharding specs for this node generated
            # in its own node handler
            assert hasattr(node, 'strategies_vector'), \
                f'The predecessor node {node_name} has no strategy vector to compute the resharding cost.'
            prev_strategy_vector = node.strategies_vector
            prev_sharding_specs = [
                prev_strategy.get_sharding_spec_by_name(node_name) for prev_strategy in prev_strategy_vector
            ]

            # create data structrure to store costs
            if node not in resharding_costs:
                resharding_costs[node] = []

            def _compute_resharding_cost(
                    prev_sharding_spec: Union[ShardingSpec,
                                              List[ShardingSpec]], current_sharding_spec: Union[ShardingSpec,
                                                                                                List[ShardingSpec]],
                    data: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]) -> TrainCycleItem:
                """
                This is a helper function to compute the resharding cost for a specific strategy of a node.
                """
                if prev_sharding_spec is None:
                    return TrainCycleItem(fwd=0, bwd=0, total=0)
                elif isinstance(prev_sharding_spec, ShardingSpec):
                    if isinstance(data, torch.nn.parameter.Parameter):
                        # we won't compute the resharding cost for the parameters,
                        # since the parameters will be sharded before runtime and
                        # not converted during runtime.
                        return TrainCycleItem(fwd=0, bwd=0, total=0)
                    elif isinstance(data, torch.Tensor):
                        dtype = data.dtype
                        size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()
                        _, _, consistency_cost = shape_consistency_manager.shape_consistency(
                            prev_sharding_spec, current_sharding_spec)

                        resharding_cost = TrainCycleItem(fwd=consistency_cost["forward"] * size_per_elem_bytes,
                                                         bwd=consistency_cost["backward"] * size_per_elem_bytes,
                                                         total=consistency_cost["total"] * size_per_elem_bytes)
                        return resharding_cost
                    else:
                        # This raise is used to check if we have missed any type of data.
                        # It could be merged into Parameter branch, which means we won't handle
                        # non-tensor arguments.
                        raise ValueError(f'Unsupported data type {type(data)}')
                else:
                    assert isinstance(prev_sharding_spec, (tuple, list)), \
                        f'prev_sharding_spec should be in type of ShardingSpec, List[ShardingSpec], \
                            or Tuple[ShardingSpec], but got {type(prev_sharding_spec)}'

                    fwd_cost = 0
                    bwd_cost = 0
                    total_cost = 0
                    for index, (prev_sharding_spec_item,
                                current_sharding_spec_item) in enumerate(zip(prev_sharding_spec,
                                                                             current_sharding_spec)):
                        item_cost = _compute_resharding_cost(prev_sharding_spec_item, current_sharding_spec_item,
                                                             data[index])
                        fwd_cost += item_cost.fwd
                        bwd_cost += item_cost.bwd
                        total_cost += item_cost.total
                    resharding_cost = TrainCycleItem(fwd=fwd_cost, bwd=bwd_cost, total=total_cost)
                    return resharding_cost

            # for each sharding spec generated by the predecessor's node handler
            # compute the resharding cost to switch to the sharding spec generated
            # by the current node handler
            for prev_sharding_spec in prev_sharding_specs:
                resharding_cost = _compute_resharding_cost(prev_sharding_spec, current_sharding_spec, op_data.data)
                resharding_costs[node].append(resharding_cost)
        strategy.resharding_costs = resharding_costs
        return strategy

    def register_strategy(self, compute_resharding_cost: bool = True) -> StrategiesVector:
        """
        Register different sharding strategies for the current node.
        """
        strategy_generators = self.get_strategy_generator()
        for generator in strategy_generators:
            strategies = generator.generate()

            # postprocess a strategy
            # postprocess can produce one strategy or multiple strategies
            post_processed_strategies_map = map(self.post_process, strategies)
            post_processed_strategies = []

            for strategy in post_processed_strategies_map:
                if isinstance(strategy, (list, tuple)):
                    post_processed_strategies.extend(strategy)
                else:
                    post_processed_strategies.append(strategy)

            # compute the resharding costs based on the previous node
            # strategies if specified
            if compute_resharding_cost:
                updated_strategies = map(self.update_resharding_cost, post_processed_strategies)
                post_processed_strategies = list(updated_strategies)

            self.strategies_vector.extend(post_processed_strategies)

        # validating the correctness of the sharding strategy
        for strategy in self.strategies_vector:
            for op_data, sharding_spec in strategy.sharding_specs.items():
                if op_data.data is not None and isinstance(op_data.data, torch.Tensor):
                    check_sharding_spec_validity(sharding_spec, op_data.data)

        return self.strategies_vector

    def post_process(self, strategy: ShardingStrategy) -> Union[ShardingStrategy, List[ShardingStrategy]]:
        # tranform the strategy generated
        # e.g. to process the sharding strategy for the transposed weights
        return strategy

    @abstractmethod
    def get_strategy_generator(self) -> List[StrategyGenerator]:
        """
        Define which generators should be used by this NodeHandler object.
        """
        pass

    @abstractmethod
    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        """
        Returns the mapping between the logical operation data to its physical data.
        A logical operation data is a data associated with an operation, which can be input and output. It is
        defined by the strategy generator, for example, a matrix multiplication operation has two operands "input"
        and "other" and one result "output". For a nn.Linear module, the physical operand for "input" is
        the module input, the physical operand for "other" is the module weight, and the physical result for "output"
        is the module output.
        Note that the operand name is specified by the StrategyGenerator object.

        For example:

            # for a linear layer
            mapping = {
                "input": Operand(name=str(self.node.args[0]), type=OperationDataType.ARG, data=self.node.args[0]._meta_data),
                "other": Operand(name="weight", type=OperationDataType.PARAM, data=self.named_parameters['weight']),
                "bias": Operand(name="bias", type=OperationDataType.PARAM, data=self.named_parameters['bias']),
                "output": Operand(name=str(self.node), type=OperationDataType.OUTPUT, data=self.node._meta_data),
            }
        """
        pass


class ModuleHandler(NodeHandler):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        # set attributes to access module parameters for convenience
        assert self.node.graph.owning_module is not None, \
            f'The graph is not associated with a module, please make sure it can be used to instantiate a GraphModule object.'
        module = self.node.graph.owning_module.get_submodule(self.node.target)
        named_parameters = list(module.named_parameters(recurse=False))
        named_buffers = list(module.named_buffers(recurse=False))
        # convert named parameters from list to dict
        named_parameters = {k: v for k, v in named_parameters}
        named_buffers = {k: v for k, v in named_buffers}
        self.module = module
        self.named_parameters = named_parameters
        self.named_buffers = named_buffers
