import typing
from typing import *

import torch
from torch import Tensor

import torch_geometric.typing
from torch_geometric.typing import *


class Sequential(torch.nn.Module):
    def reset_parameters(self) -> None:
{%- for child in children %}
        if hasattr(self.{{child.name}}, 'reset_parameters'):
            self.{{child.name}}.reset_parameters()
{%- endfor %}

    def forward(self, {{ input_types|join(', ') }}) -> {{return_type}}:
{%- for child in children %}
        {{child.return_names|join(', ')}} = self.{{child.name}}({{child.param_names|join(', ')}})
{%- endfor %}
        return {{children[-1].return_names|join(', ')}}

    def __getitem__(self, idx: int) -> torch.nn.Module:
        return getattr(self, self._module_names[idx])

    def __len__(self) -> int:
        return {{children|length}}

    def __repr__(self) -> str:
        module_reprs = [
            f'  ({i}) - {self[i]}: {self._module_descs[i]}'
            for i in range(len(self))
        ]
        return 'Sequential(\n{}\n)'.format('\n'.join(module_reprs))
