import mindspore.numpy as np
from mindspore import dtype as mstype
from mindspore import ops, Tensor, context, nn

class AdaptiveAvgPool2d(nn.Cell):
    def init(self, output_size):
    """Initialize AdaptiveAvgPool2d."""
        super(AdaptiveAvgPool2d, self).init()
        self.output_size = output_size

def adaptive_avgpool2d(self, inputs):
    """ NCHW """
    H = self.output_size[0]
    W = self.output_size[1]

    H_start = ops.Cast()(np.arange(start=0, stop=H, dtype=mstype.float32) * (inputs.shape[-2] / H), mstype.int64)
    H_end = ops.Cast()(np.ceil(((np.arange(start=0, stop=H, dtype=mstype.float32)+1) * (inputs.shape[-2] / H))), mstype.int64)

    W_start = ops.Cast()(np.arange(start=0, stop=W, dtype=mstype.float32) * (inputs.shape[-1] / W), mstype.int64)
    W_end = ops.Cast()(np.ceil(((np.arange(start=0, stop=W, dtype=mstype.float32)+1) * (inputs.shape[-1] / W))), mstype.int64)

    pooled2 = []
    for idx_H in range(H):
        pooled1 = []
        for idx_W in range(W):
            h_s = int(H_start[idx_H].asnumpy())
            h_e = int(H_end[idx_H].asnumpy())
            w_s = int(W_start[idx_W].asnumpy())
            w_e = int(W_end[idx_W].asnumpy())
            res = inputs[:, :, h_s:h_e, w_s:w_e]
            pooled1.append(ops.ReduceMean(keep_dims=True)(res, (-2,-1)))
        pooled1 = ops.Concat(-1)(pooled1)
        pooled2.append(pooled1)
    pooled2 = ops.Concat(-2)(pooled2)

    return pooled2

def construct(self, x):
    x = self.adaptive_avgpool2d(x)
    return x