# Copyright 2022 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""STGCN layer"""
import mindspore as ms
import mindspore.nn as nn
from mindspore import ops
from mindspore_gl.nn.conv.chebconv import ChebConv
from mindspore_gl.nn import GNNCell
from mindspore_gl import Graph

class TemporalConv(ms.nn.Cell):
    """temporal convolution layer
    from the paper `A deep learning framework for traffic forecasting
    arXiv preprint arXiv:1709.04875, 2017. <https://arxiv.org/pdf/1709.04875.pdf>`_ .

    Args:
        in_channels (int): Input node feature size.
        out_channels (int): Output node feature size.
        kernel_size (int): Convolutional kernel size.

    Inputs:
        - **x** (Tensor) - The input node features. The shape is :math:`(B, T, N, (D_{in}))`
          where :math:`B` is the size of batch, :math:`T` is the number of input time steps,
          :math:`N` is the number of nodes,
          :math:`(D_{in})` should be equal to `in_channels` in `Args`.

    Outputs:
        - Tensor, output node features with shape of :math:`(B, D_{out}, N, T)`,
        where :math:`B` is the size of batch, :math:`(D_{out})` should be the same as
        `out_channels` in `Args`, :math:`N` is the number of nodes,
        :math:`T` is the number of input time steps.

    Raises:
        TypeError: If `in_channels` or `out_channels` or `kernel_size` is not an int.

    Examples:
        >>> import mindspore as ms
        >>> from mindspore_gl.nn.temporal import TemporalConv
        >>> import numpy as np
        >>> batch_size = 4
        >>> input_time_steps = 6
        >>> num_nodes = 2
        >>> in_channels = 2
        >>> out_channels = 1
        >>> temprol_conv = TemporalConv(in_channels, out_channels)
        >>> input = ms.Tensor(np.ones((batch_size, input_time_steps, num_nodes, in_channels)), ms.float32)
        >>> out = temprol_conv(input)
        >>> print(out.shape)
        (4, 4, 2, 1)
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):
        super(TemporalConv, self).__init__()
        if in_channels <= 0 or not isinstance(in_channels, int):
            raise ValueError("in_channels must be positive int")
        if out_channels <= 0 or not isinstance(out_channels, int):
            raise ValueError("out_channels must be positive int")
        if kernel_size <= 0 or not isinstance(kernel_size, int):
            raise ValueError("kernel_size must be positive int")

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size

        self.conv_1 = ms.nn.Conv2d(in_channels, out_channels, (1, kernel_size), pad_mode='valid', has_bias=True)
        self.conv_2 = ms.nn.Conv2d(in_channels, out_channels, (1, kernel_size), pad_mode='valid', has_bias=True)
        self.conv_3 = ms.nn.Conv2d(in_channels, out_channels, (1, kernel_size), pad_mode='valid', has_bias=True)

        self.sigmoid = ms.nn.Sigmoid()
        self.relu = ms.nn.ReLU()

    def construct(self, x):
        """
        Construct function for temporal convolution layer.
        """
        x = ops.Transpose()(x, (0, 3, 2, 1))
        p = self.conv_1(x)
        q = self.sigmoid(self.conv_2(x))
        pq = p * q
        h = self.relu(pq + self.conv_3(x))
        h = ops.Transpose()(h, (0, 3, 2, 1))
        return h

class STConv(GNNCell):
    r"""
    Spatial-Temporal Graph Convolutional layer.
    From the paper `A deep learning framework for traffic forecasting
    arXiv preprint arXiv:1709.04875, 2017. <https://arxiv.org/pdf/1709.04875.pdf>`_ .
    The STGCN layer contains 2 temporal convolution layer and 1
    graph convolution layer (ChebyNet).

    Args:
        num_nodes (int): number of nodes.
        in_channels (int): Input node feature size.
        hidden_channels (int): hidden feature size.
        out_channels (int): Output node feature size.
        kernel_size (int, optional): Convolutional kernel size. Default: ``3``.
        k (int, optional): Chebyshev filter size. Default: ``3``.
        bias (bool, optional): Whether use bias. Default: ``True``.

    Inputs:
        - **x** (Tensor) - The input node features. The shape is :math:`(B, T, N, (D_{in}))`
          where :math:`B` is the size of batch, :math:`T` is the number of input time steps,
          :math:`N` is the number of nodes,
          :math:`(D_{in})` should be equal to `in_channels` in `Args`.
        - **edge_weight** (Tensor) - Edge weights. The shape is :math:`(N\_e,)`
          where :math:`N\_e` is the number of edges.
        - **g** (Graph) - The input graph.

    Outputs:
        - Tensor, output node features with shape of :math:`(B, D_{out}, N, T)`,
          where :math:`B` is the size of batch, :math:`(D_{out})` should be the same as
          `out_channels` in `Args`, :math:`N` is the number of nodes,
          :math:`T` is the number of input time steps.

    Raises:
        TypeError: If `num_nodes` or `in_channels` or `out_channels` or `hidden_channels`
            or `kernel_size` or is `k` not an int.
        TypeError: If `bias` is not a bool.

    Supported Platforms:
        ``Ascend`` ``GPU``

    Examples:
        >>> import numpy as np
        >>> import mindspore as ms
        >>> from mindspore_gl.nn.temporal import STConv
        >>> from mindspore_gl import GraphField
        >>> from mindspore_gl.graph import norm
        >>> n_nodes = 4
        >>> n_edges = 6
        >>> feat_size = 2
        >>> edge_attr = ms.Tensor([1, 1, 1, 1, 1, 1], ms.float32)
        >>> edge_index = ms.Tensor([[1, 1, 2, 2, 3, 3],
        >>>                         [0, 2, 1, 3, 0, 1]], ms.int32)
        >>> edge_index, edge_weight = norm(edge_index, n_nodes, edge_attr, 'sym')
        >>> edge_weight = ms.ops.Reshape()(edge_weight, ms.ops.Shape()(edge_weight) + (1,))
        >>> batch_size = 2
        >>> input_time_steps = 5
        >>> feat = ms.Tensor(np.ones((batch_size, input_time_steps, n_nodes, feat_size)), ms.float32)
        >>> graph_field = GraphField(edge_index[0], edge_index[1], n_nodes, n_edges)
        >>> stconv = STConv(num_nodes=n_nodes, in_channels=feat_size,
        >>>                 hidden_channels=3, out_channels=2,
        >>>                 kernel_size=2, k=2)
        >>> out = stconv(feat, edge_weight, *graph_field.get_graph())
        >>> print(out.shape)
        (2, 3, 4, 2)
    """
    def __init__(self,
                 num_nodes: int,
                 in_channels: int,
                 hidden_channels: int,
                 out_channels: int,
                 kernel_size: int = 3,
                 k: int = 3,
                 bias: bool = True):
        super().__init__()
        if num_nodes <= 0 or not isinstance(num_nodes, int):
            raise ValueError("num_nodes must be positive int")
        if in_channels <= 0 or not isinstance(in_channels, int):
            raise ValueError("in_channels must be positive int")
        if hidden_channels <= 0 or not isinstance(hidden_channels, int):
            raise ValueError("hidden_channels must be positive int")
        if out_channels <= 0 or not isinstance(out_channels, int):
            raise ValueError("out_channels must be positive int")
        if kernel_size <= 0 or not isinstance(kernel_size, int):
            raise ValueError("kernel_size must be positive int")
        if k <= 0 or not isinstance(k, int):
            raise ValueError("k must be positive int")
        if not isinstance(bias, bool):
            raise ValueError("bias must be bool")


        self.num_nodes = num_nodes
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.k = k

        self.num_nodes = num_nodes
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.k = k
        self.bias = bias

        self.temporala_conv1 = TemporalConv(
            in_channels=self.in_channels,
            out_channels=self.hidden_channels,
            kernel_size=self.kernel_size
        )
        self.cheb_conv = ChebConv(
            in_channels=self.hidden_channels,
            out_channels=self.hidden_channels,
            k=self.k,
            bias=self.bias,
        )
        self.temporala_conv2 = TemporalConv(
            in_channels=self.hidden_channels,
            out_channels=self.out_channels,
            kernel_size=self.kernel_size
        )
        self.batch_norm = nn.BatchNorm2d(num_nodes)
        self.relu = nn.ReLU()

    def construct(self, x, edge_weight, g: Graph):
        """
        Construct function for STConv.
        """
        t0 = self.temporala_conv1(x)
        t = ops.ZerosLike()(t0)
        for b in range(t0.shape[0]):
            for s in range(t0.shape[1]):
                t[b][s] = self.cheb_conv(t0[b][s], edge_weight, g)
        t = self.relu(t0)
        t = self.temporala_conv2(t)
        t = ops.Transpose()(t, (0, 2, 1, 3))
        t = self.batch_norm(t)
        t = ops.Transpose()(t, (0, 2, 1, 3))
        return t