import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops

import numpy as np

from typing import *

# ms.set_context(mode=ms.GRAPH_MODE, device_target="CPU")
ms.set_context(mode=ms.PYNATIVE_MODE, device_target="CPU")

class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 1, 3, 1)

    def construct(self, x):
        return self.conv1(x)



class Callback1(ms.train.callback.Callback):
    def __init__(self):
        super().__init__()

    def on_train_epoch_begin(self, run_context):
        print("on_train_epoch_begin")

    def on_train_step_end(self, run_context):
        print("🍓", end='')

    def on_train_epoch_end(self, run_context, *args):
        print()
        print("on_train_epoch_end")

    def on_eval_epoch_begin(self, run_context):
        print("on_eval_epoch_begin")

    def on_eval_step_end(self, run_context):
        print("🧃", end='')

    def on_eval_epoch_end(self, run_context, *args):
        print()
        print("on_eval_epoch_end")



if __name__ == "__main__":
    dataset = ms.dataset.NumpySlicesDataset(
        (np.random.randn(100, 1, 16, 16).astype(np.float32), np.random.randn(100, 1, 16, 16).astype(np.float32)),
        column_names=["image", "label"], shuffle=False)
    train_ds, val_ds = dataset.split([0.5, 0.5], randomize=False)

    train_ds = train_ds.batch(4)
    val_ds = val_ds.batch(4)

    net = Net()
    loss = nn.MSELoss()
    opt = nn.optim.Adam(net.trainable_params(), 0.001)
    callback1 = Callback1()

    model = ms.Model(
        network=net,
        loss_fn=loss,
        optimizer=opt,
        metrics={"MSE": nn.metrics.MSE()},
    )

    model.fit(
        epoch=4,
        train_dataset=train_ds,
        valid_dataset=val_ds,
        callbacks=[callback1],
        dataset_sink_mode=False,
        valid_dataset_sink_mode=False,
        sink_size=-1,
        initial_epoch=0
    )