import typing
from typing import NamedTuple

import torch
from torch import Tensor

import torch_geometric.typing
from torch_geometric import is_compiling
from torch_geometric.utils import is_sparse

from {{module}} import *


class CollectArgs(NamedTuple):
{%- for param in collect_param_dict.values() %}
    {{param.name}}: {{param.type_repr}}
{%- endfor %}


def _collect(
    self,
    edge_index: Union[Tensor, SparseTensor],
{%- for param in propagate_signature.param_dict.values() %}
    {{param.name}}: {{param.type_repr}},
{%- endfor %}
    size: List[Optional[int]],
) -> CollectArgs:

    i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1)
{% for name in collect_param_dict %}
{%- if name.endswith('_i') or name.endswith('_j') %}
    # Collect `{{name}}`:
    if isinstance({{name[:-2]}}, (tuple, list)):
        assert len({{name[:-2]}}) == 2
        _{{name[:-2]}}_0, _{{name[:-2]}}_1 = {{name[:-2]}}[0], {{name[:-2]}}[1]
        if isinstance(_{{name[:-2]}}_0, Tensor):
            self._set_size(size, 0, _{{name[:-2]}}_0)
{%- if name.endswith('_j') %}
            {{name}} = self._lift(_{{name[:-2]}}_0, edge_index, {{name[-1]}})
        else:
            {{name}} = None
{%- endif %}
        if isinstance(_{{name[:-2]}}_1, Tensor):
            self._set_size(size, 1, _{{name[:-2]}}_1)
{%- if name.endswith('_i') %}
            {{name}} = self._lift(_{{name[:-2]}}_1, edge_index, {{name[-1]}})
        else:
            {{name}} = None
{%- endif %}
    elif isinstance({{name[:-2]}}, Tensor):
        self._set_size(size, 0, {{name[:-2]}})
        self._set_size(size, 1, {{name[:-2]}})
        {{name}} = self._lift({{name[:-2]}}, edge_index, {{name[-1]}})
    else:
        {{name}} = None
{%- endif %}
{%- endfor %}

    assert isinstance(edge_index, Tensor)
    adj_t = None
    edge_index_i = edge_index[i]
    edge_index_j = edge_index[j]
    ptr = None
    index = edge_index_i

    size_i = size[i]
    size_j = size[j]
    dim_size = size_i

    return CollectArgs(
{%- for name in collect_param_dict %}
        {{name}},
{%- endfor %}
    )


def propagate(
    self,
    edge_index: Union[Tensor, SparseTensor],
{%- for param in propagate_signature.param_dict.values() %}
    {{param.name}}: {{param.type_repr}},
{%- endfor %}
    size: Size = None,
) -> {{propagate_signature.return_type_repr}}:

    assert self.decomposed_layers == 1  # TODO

    # Begin Propagate Forward Pre Hook #########################################
    if not torch.jit.is_scripting() and not is_compiling():
        for hook in self._propagate_forward_pre_hooks.values():
            hook_kwargs = dict(
{%- for name in propagate_signature.param_dict %}
                {{name}}={{name}},
{%- endfor %}
            )
            res = hook(self, (edge_index, size, hook_kwargs))
            if res is not None:
                edge_index, size, hook_kwargs = res
{%- for name in propagate_signature.param_dict %}
                {{name}} = hook_kwargs['{{name}}']
{%- endfor %}
    # End Propagate Forward Pre Hook ###########################################

    mutable_size = self._check_input(edge_index, size)
    fuse = is_sparse(edge_index) and self.fuse and self.explain is not True

    if fuse:
        assert isinstance(edge_index, SparseTensor)  # TODO

        # Begin Message and Aggregate Forward Pre Hook #########################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._message_and_aggregate_forward_pre_hooks.values():
                hook_kwargs = dict(
{%- for name in message_and_aggregate_args %}
                    {{name}}={{name}},
{%- endfor %}
                )
                res = hook(self, (edge_index, hook_kwargs))
                if res is not None:
                    edge_index, hook_kwargs = res
{%- for name in message_and_aggregate_args %}
                    {{name}} = hook_kwargs['{{name}}']
{%- endfor %}
        # End Message And Aggregate Forward Pre Hook ##########################

        out = self.message_and_aggregate(
            edge_index,
{%- for name in message_and_aggregate_args %}
            {{name}},
{%- endfor %}
        )

        # Begin Message And Aggregate Forward Hook #############################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._message_and_aggregate_forward_hooks.values():
                hook_kwargs = dict(
{%- for name in message_and_aggregate_args %}
                    {{name}}={{name}},
{%- endfor %}
                )
                res = hook(self, (edge_index, hook_kwargs, ), out)
                out = res if res is not None else out
        # End Message And Aggregate Forward Hook ###############################

        out = self.update(
            out,
{%- for name in update_args %}
            {{name}}={{name}},
{%- endfor %}
        )

    else:
        kwargs = self._collect(
            edge_index,
{%- for name in propagate_signature.param_dict %}
            {{name}},
{%- endfor %}
            mutable_size,
        )

        # Begin Message Forward Pre Hook #######################################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._message_forward_pre_hooks.values():
                hook_kwargs = dict(
{%- for name in message_args %}
                    {{name}}=kwargs.{{name}},
{%- endfor %}
                )
                res = hook(self, (hook_kwargs, ))
                hook_kwargs = res[0] if isinstance(res, tuple) else res
                if res is not None:
{%- for name in message_args %}
                    kwargs.{{name}} = hook_kwargs['{{name}}']
{%- endfor %}
        # End Message Forward Pre Hook #########################################

        out = self.message(
{%- for name in message_args %}
            {{name}}=kwargs.{{name}},
{%- endfor %}
        )

        # Begin Message Forward Hook ###########################################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._message_forward_hooks.values():
                hook_kwargs = dict(
{%- for name in message_args %}
                    {{name}}=kwargs.{{name}},
{%- endfor %}
                )
                res = hook(self, (hook_kwargs, ), out)
                out = res if res is not None else out
        # End Message Forward Hook #############################################

        if self.explain is True:  # TODO
            raise NotImplementedError

        # Begin Aggregate Forward Pre Hook #####################################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._aggregate_forward_pre_hooks.values():
                hook_kwargs = dict(
{%- for name in aggregate_args %}
                    {{name}}=kwargs.{{name}},
{%- endfor %}
                )
                res = hook(self, (hook_kwargs, ))
                hook_kwargs = res[0] if isinstance(res, tuple) else res
                if res is not None:
{%- for name in aggregate_args %}
                    kwargs.{{name}} = hook_kwargs['{{name}}']
{%- endfor %}
        # End Aggregate Forward Pre Hook #######################################

        out = self.aggregate(
            out,
{%- for name in aggregate_args %}
            {{name}}=kwargs.{{name}},
{%- endfor %}
        )

        # Begin Aggregate Forward Hook #########################################
        if not torch.jit.is_scripting() and not is_compiling():
            for hook in self._aggregate_forward_hooks.values():
                hook_kwargs = dict(
{%- for name in aggregate_args %}
                    {{name}}=kwargs.{{name}},
{%- endfor %}
                )
                res = hook(self, (hook_kwargs, ), out)
                out = res if res is not None else out
        # End Aggregate Forward Hook ###########################################

        out = self.update(
            out,
{%- for name in update_args %}
            {{name}}=kwargs.{{name}},
{%- endfor %}
        )

    # Begin Propagate Forward Hook ############################################
    if not torch.jit.is_scripting() and not is_compiling():
        for hook in self._propagate_forward_hooks.values():
            hook_kwargs = dict(
{%- for name in message_args %}
                {{name}}=kwargs.{{name}},
{%- endfor %}
            )
            res = hook(self, (edge_index, mutable_size, hook_kwargs), out)
            out = res if res is not None else out
    # End Propagate Forward Hook ##############################################

    return out
