import argparse
import numpy as np
import os

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import mindspore.dataset.transforms.vision.c_transforms as transforms
import mindspore.dataset as ds
import mindspore.dataset.transforms.c_transforms as C
import mindspore.dataset.transforms.vision.c_transforms as CV
import mindspore.nn as nn

from mindspore import context
from mindspore.common import dtype as mstype
from mindspore.common.initializer import TruncatedNormal
from mindspore.model_zoo.resnet import resnet50
from mindspore.nn.metrics import Accuracy
from mindspore.train import Model
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits


def create_train_dataset(args, num_parallel_workers = 1):
    dataset = ds.ImageFolderDatasetV2(args.paths_train, decode=True, shuffle=True)

    dataset = dataset.map(input_columns="label", operations=C.TypeCast(mstype.int32), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.RandomResizedCrop((args.resize_height, args.resize_width)), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.RandomRotation(30), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.RandomHorizontalFlip(), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.Rescale(1.0/255.0, 0.0), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.RandomColorAdjust(), num_parallel_workers=num_parallel_workers)
    dataset = dataset.map(input_columns="image", operations=CV.HWC2CHW(), num_parallel_workers=num_parallel_workers)
    dataset = dataset.batch(args.batch, drop_remainder=True)
    dataset = dataset.repeat(1)

    return dataset

def weight_variable():
    """Weight initial."""
    return TruncatedNormal(0.02)

def conv(in_channels, out_channels, kernel_size, stride=1):
    """Conv layer weight initial."""
    weight = weight_variable()
    return nn.Conv2d(in_channels, out_channels,
                     kernel_size=kernel_size, stride=stride,
                     weight_init=weight, has_bias=False, pad_mode="same")

def fc_with_initialize(input_channels, out_channels):
    """Fc layer weight initial."""
    weight = weight_variable()
    bias = weight_variable()
    return nn.Dense(input_channels, out_channels, weight, bias)

class LeNet5(nn.Cell):
    """Lenet network structure."""
    # define the operator required
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = conv(3, 6, 5)
        self.conv1BatchNorm = nn.BatchNorm2d(6)
        self.conv2 = conv(6, 16, 5)
        self.fc1 = fc_with_initialize(16 * 32 * 32, 120)
        self.fc2 = fc_with_initialize(120, 84)
        self.fc3 = fc_with_initialize(84, 2)
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()

    # use the preceding operators to construct networks
    def construct(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv1BatchNorm(x) #works fine without batchnorm
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Fire detection example')
    parser.add_argument('--device_target', type=str, default="CPU", choices=['Ascend', 'GPU', 'CPU'],
                        help='device where the code will be implemented (default: CPU)')
    parser.add_argument('--paths_train', default="./data/train", type=str, help="paths of the train folder")
    parser.add_argument('--paths_test', default="./data/test", type=str, help="paths of the test folder")
    parser.add_argument('--resize_height', type=int, default=128, help="height of the input image of model")
    parser.add_argument('--resize_width', type=int, default=128, help="width of the input image of model")
    parser.add_argument('--epoch', type=int, default=1, help="Epoch of training")
    parser.add_argument('--batch', type=int, default=32, help="Batch size")

    args = parser.parse_args()
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
    dataset_train = create_train_dataset(args)

    net = LeNet5()#fire_classify_net.fire_classify_net()
    loss = SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean")
    optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
    #optim = nn.Adam(params=net.trainable_params())
    model = Model(net, loss_fn=loss, optimizer=optim, metrics={"Accuracy": Accuracy()})
    config_ck = CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=10)
    # save the network model and parameters for subsequence fine-tuning
    ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet", config=config_ck)
    model.train(args.epoch, dataset_train, callbacks=[ckpoint_cb, LossMonitor()], dataset_sink_mode=False)