from download import download

url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cnnctc_dataset.tar.gz"

download(url, ".", kind="tar.gz", replace=True)

import mindspore as ms
from mindspore import nn, ops
from mindspore.dataset import MindDataset
from mindspore.dataset import transforms, vision
from mindspore.common.initializer import Normal
import math
import numpy as np
import time

def trans_arg_to_list(arg, size = 2):
    if not isinstance(arg, tuple) and not isinstance(arg, list):
        ret = []
        for i in range(size):
            ret.append(arg)
        return ret
    return arg
    
def get_conv2d_output_shape(input_shape, 
                            kernel_size, 
                            stride, 
                            padding_mode="same", 
                            padding = 0,
                            dilation = 1):
    input_shape = trans_arg_to_list(input_shape)        
    kernel_size = trans_arg_to_list(kernel_size)        
    stride = trans_arg_to_list(stride)        
    padding = trans_arg_to_list(padding, 4)        
    dilation = trans_arg_to_list(dilation)        
    h, w = input_shape
    if padding_mode == 'same':
        h_out = math.ceil(h / stride[0])
        w_out = math.ceil(w / stride[1])
    elif padding_mode == 'valid':
        h_out = math.ceil((h - dilation[0] * (kernel_size[0] - 1)) / stride[0])
        w_out = math.ceil((w - dilation[1] * (kernel_size[1] - 1)) / stride[1])
    else:
        h += padding[0] + padding[1]
        w += padding[2] + padding[3]
        h_out = math.floor((h - dilation[0] * (kernel_size[0] - 1) - 1) / stride[0] + 1)
        w_out = math.floor((w - dilation[1] * (kernel_size[1] - 1) - 1) / stride[1] + 1)
    return h_out, w_out

weight_init = Normal(mean=0, sigma=0.02)
gamma_init = Normal(mean=1, sigma=0.02)

class ResidualBlockBase(nn.Cell):
    expansion = 1
    def __init__(self, 
                 in_channel:int, 
                 out_channel:int, 
                 stride = 1,             
                 down_sample:nn.Cell = None,
                 norm:nn.Cell = None,):
        super().__init__()
        
        self.conv1 = nn.Conv2d(in_channel, 
                               out_channel, 
                               kernel_size=3, 
                               stride=stride,
                              weight_init=weight_init)
        self.conv2 = nn.Conv2d(out_channel, 
                               out_channel, 
                               kernel_size=3, 
                              weight_init=weight_init)
        self.norm = norm or nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU()
        self.down_sample = down_sample
        self.kernel_size = 3
        self.stride = stride
        
    def construct(self, x):
        identity = x
        # 第一层
        out = self.conv1(x)
        out = self.norm(out)
        out = self.relu(out)
        # 第二层
        out = self.conv2(out)
        out = self.norm(out)

        # 对identity采样
        if self.down_sample is not None:
            identity = self.down_sample(identity)
        # 做加法
        out = out + identity
        out = self.relu(out)
        return out
        
    def get_output_shape(self, input_shape):
        out_shape = get_conv2d_output_shape(input_shape,
                                           kernel_size=self.conv1.kernel_size,
                                           stride=self.conv1.stride,
                                           padding_mode=self.conv1.pad_mode,
                                           padding=self.conv1.padding,
                                           dilation=self.conv1.dilation)
        out_shape = get_conv2d_output_shape(out_shape,
                                           kernel_size=self.conv2.kernel_size,
                                           stride=self.conv2.stride,
                                           padding_mode=self.conv2.pad_mode,
                                           padding=self.conv2.padding,
                                           dilation=self.conv2.dilation)
        return out_shape


def make_layer(block:nn.Cell,
                last_out_channel:int, 
                channel:int,  
                nums, 
                stride:int = 1, 
                in_shape=None):
    down_sample = None
    if stride != 1 or last_out_channel != channel * block.expansion:
        down_sample = nn.SequentialCell([
            nn.Conv2d(last_out_channel, 
                      channel * block.expansion, 
                      kernel_size=1, stride=stride),
            nn.BatchNorm2d(channel * block.expansion, 
                           gamma_init=gamma_init)
        ])
    layers = []
    layers.append(block(last_out_channel, channel, stride, down_sample))
    in_channel = channel * block.expansion
    for i in range(1, nums):
        layers.append(block(in_channel, channel))

    if in_shape is not None:
        for layer in layers:
            in_shape = layer.get_output_shape(in_shape)
        return nn.SequentialCell(layers), in_shape

    return nn.SequentialCell(layers)

class CNNCTC(nn.Cell):
    def __init__(self, num_classes, final_feature_width):
        super().__init__()
        self.num_classes = num_classes
        self.final_feature_width = final_feature_width
        self.conv1 = nn.Conv2d(3, 512, kernel_size=1)
        outshape = (32, 100)
        self.conv2_x, outshape = make_layer(ResidualBlockBase, 512, 512, 1, 1, outshape)
        print(512, outshape)
        self.conv3_x, outshape = make_layer(ResidualBlockBase, 512, 512, 1, 2, outshape)
        print(512, outshape)
        self.conv4_x, outshape = make_layer(ResidualBlockBase, 512, final_feature_width, 1, 4, outshape)
        print(final_feature_width, outshape)
        self.flatten = nn.Flatten(2)
        self.fc1 = nn.Dense(outshape[0] * outshape[1], num_classes)
        self.fc1_activation = nn.LogSoftmax()
        print(final_feature_width, self.fc1.out_channels)
        
    def construct(self, x:ms.Tensor):
        x = self.conv1(x)
        x = self.conv2_x(x)
        x = self.conv3_x(x)
        x = self.conv4_x(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc1_activation(x)
        return x
final_feature_width = 26
num_class = 37
epochs = 10        
net = CNNCTC(num_class, final_feature_width)
input = ops.ones((1, 3, 32, 100))
output = net(input)
print(output.shape)

batch_size = 64
dataset_train = MindDataset('./cnnctc_dataset/STMJ_train_dataset.mindrecord')
dataset_train = dataset_train.batch(batch_size, drop_remainder=True)
dataset_train = dataset_train.map(lambda x: x.squeeze(), input_columns='img')    
dataset_train = dataset_train.map(lambda x: x.squeeze(), input_columns='target_lengths')    
dataset_train = dataset_train.map(lambda x: x.squeeze(), input_columns='text')    
dataset_size = dataset_train.get_dataset_size()
print(dataset_size)
d1 = next(dataset_train.create_tuple_iterator())
print(d1[0].shape, d1[1].shape, d1[2].shape)

loss_fn = nn.CTCLoss()
sequence_length = np.array([final_feature_width] * batch_size, dtype=np.int32)
sequence_length = ms.Tensor.from_numpy(sequence_length)
def forward_fn(img, target_lengths, text):
    logits = net(img)
    logits = ops.transpose(logits, (1, 0, 2))
    loss = loss_fn(logits, text, sequence_length, target_lengths)
    return loss, logits
opt = nn.Adam(net.trainable_params(), 1e-6)
grad_fn = ms.value_and_grad(forward_fn, None, net.trainable_params(), has_aux=True)

def train_step(img, target_lengths, text):
    (loss, logits), grads = grad_fn(img, target_lengths, text)
    opt(grads)
    return loss, logits
    
def train_loop(epochs, dataset):
    net.set_train(True)
    for epoch in range(epochs):
        for step, (img, target_lengths, text) in enumerate(dataset.create_tuple_iterator()):
            loss, logits = train_step(img, target_lengths, text)
            if step % 100 == 0:
                print('epoch', epoch, 'step', step, 'loss', loss)

train_loop(1, dataset_train)

dataset_eval = MindDataset('cnnctc_dataset/IIIT_eval_dataset.mindrecord')
net.set_train(False)
count = 0
for data in dataset_eval.create_tuple_iterator():
    logits = net(data[0])
    print(logits.argmax(2))
    print(data[2])
    count += 1
    if count>=10:
        break