import numpy as np

import time

import random

from PIL import Image

import os

import sys



sys.path.append("../")



import mindspore

import mindspore.dataset as dataset

import mindspore.dataset.vision as vision

import mindspore.dataset.transforms as transforms

from mindspore.train import TopKCategoricalAccuracy

from mindspore import Tensor,set_seed,context



sys.path.append('../../')

from methods.utils.models.models import Netwider

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

# mindspore.set_context(mode=mindspore.GRAPH_MODE)

time_now = time.strftime("%Y-%m-%d_%H_%M_%S", time.localtime())



subtask_all_id = {}





class DataLoader(object):

    def __init__(self, dataset, batch_size=1, task_iter=0, epoch_size=None, num_workers=0, shuffle=True):

        self.dataset = dataset

        self.shuffle = shuffle

        self.epoch_size = epoch_size if epoch_size is not None else len(dataset)

        self.batch_size = batch_size

        self.task_iter =task_iter

        self.num_workers = num_workers

        mean_pix = [0.485, 0.456, 0.406]

        std_pix = [0.229, 0.224, 0.225]

        self.transform = transforms.Compose([

            vision.ToTensor(),

            vision.Normalize(mean=mean_pix, std=std_pix,is_hwc=False)

            # vision.HWC2CHW()

        ])



    def __getitem__(self, epoch):

        rand_seed = epoch * self.epoch_size

        random.seed(rand_seed)

        # print(rand_seed)

        dataset.config.set_seed(rand_seed)

        custom_dataset = dataset.GeneratorDataset(self.dataset, ["image", "label"],shuffle=True)

        # custom_dataset =custom_dataset.shuffle(custom_dataset.get_dataset_size())

        custom_dataset = custom_dataset.map(input_columns=["image"], operations=self.transform)

        custom_dataset = custom_dataset.batch(self.batch_size, drop_remainder=False)

        # Create a MindSpore iterator

        iterator = custom_dataset.create_dict_iterator(num_epochs=1, output_numpy=False)

        return iterator



    def __call__(self, epoch):

        return self.__getitem__(epoch)



    def __len__(self):

        return self.epoch_size / self.batch_size





def get_permute(data_name, num_tasks, batch_size, path):

    dataloader_train = {}

    for i in range(num_tasks):

        task_folder_path = os.path.join(path, 'Task_' + str(i))

        # print(task_folder_path)

        if data_name == "cifar100":

            transforms_list = [

                vision.Decode(),

                vision.Resize(32),

                vision.RandomCrop(32),

                vision.RandomHorizontalFlip(),

                lambda x: np.asarray(x),

            ]

        elif data_name == "imagenet":

            transforms_list = [

                vision.Decode(),

                vision.Resize(84),

                vision.RandomCrop(84),

                vision.RandomHorizontalFlip(),

                lambda x: np.asarray(x),

            ]

        image_folder_dataset = dataset.ImageFolderDataset(task_folder_path)

        length=len(image_folder_dataset)

        image_folder_dataset = image_folder_dataset.map(operations=transforms_list, input_columns="image")

        image_folder_dataset = image_folder_dataset.map(input_columns="label",

                                                        operations=transforms.TypeCast(mindspore.int32))

        dataloader_train[i] = DataLoader(image_folder_dataset, batch_size, task_iter=i, epoch_size=None, num_workers=0,

                                         shuffle=True)

    return dataloader_train,length



def accuracy(output, target,k):

    topk = TopKCategoricalAccuracy(k)

    topk.clear()

    topk.update(output, target)

    output = topk.eval()

    return 100.0*output



def currency(data_name, path, batch_size = 16, epochs = 100, lr = 0.01, momentum = 0.9, no_cuda = False,

               seed = 1, num_works = 21, lr_drop = 0.4, epochs_drop = 100, judgement = False):

    num_works = 2

    train_loader,length=get_permute('cifar100', 21,16, r"E:\Yuanmou_Dev-main\learngene\datasets\exp_data\data_cifar100\2024-02-16_18_55_07\continualdataset")



    model = Netwider(13)

    model.equal()

    print(model)

    for p in model.get_parameters(expand=True):

        print(p)



    trainloader=train_loader[0](0)





    for data in trainloader:



        image = data['image']

        image=image.resize(image.shape[0], 4, image.shape[2], image.shape[3])

        label = data['label']

        print(type(image))

        outputs = model(image,-1)

        # print(outputs)

        acc1, acc5 = accuracy(outputs, label, 1), accuracy(outputs, label, 5)

        print(acc1, acc5)



currency('cifar100', r"E:\Yuanmou_Dev-main\learngene\datasets\exp_data\data_cifar100\2024-02-16_18_55_07\continualdataset") ```

Traceback (most recent call last):

File "C:\Users\LENOVO\Desktop\heur_leargene\methocis\train_learngene.py", line 340, in <module>

heru_train('cifar100', r"E : \Yuanmou_Dev-main\learngene\datasetslexp_data\data_cifar100\2024-02 -16_18_55_07\continualdataset")

File "C:\Users\LENOVO\Desktoplheur_leargenelmethods\train_learngene.py", line 267, in heru_train

snapshot=snapshot, len=length,name=train_name)

File "C:\Users\LENOVO\Desktoplheur leargenelmethodslutil:sltrain.py", line 117, in train

loss = train_step(image , label)

File "C:\Users\LENOVO\Desktop)heur leargene\methodslutils\train.py", line 106, in train_step

(loss, -), grads = grad_fn(data, label)

File "F:\ananconda3\envs\heur liblsite-packages\mindsporelopslcomposite\base.py", line 625, in after_grad

return grad_(fn-, weights)(*args, **kwargs)

F1le "F:\ananconda3\envs\heurlliblsite-peckagesmindsporelcommonlapi.py", line 121, in wrapper

results = fn(*arg, **kwargs)

File "F:\ananconda3\envslheurlliblsite-packageslmindsporelopslcompositelbase.py", line600, in after_grad

res = self. pynative forward run(fn, grad_, weights, args, kwargs)

File "F:\ananconda3\envslheurllib\site-packages\mindsporelopslcoompositelbase.py", line 650, in _pynative_forward_run

outputs = fn(*args, **new_kwargs)

File "F:\ananconda3\envs\heur|lib\site-packages\mindsporelopslcompositelbase.py", line 561, in aux_fn

outputs = fn(*args)

File "C:\Users\LENOVO\Desktop\heur leargenelmethodslutilsltrain.py", line 95, in forward_fn

outputs = model(data,task)

File "F: lananconda3\envs\heurllib\site-packages\mindspore\nnlcell.py", line 705, in __call_-

raise err

File "F:\ananconda3\envslheurllib\site-packages\mindspore\nnlcell.py", line 701, in -_call_-

output = self._run_construct(args , kwargs)

File "F:\ananconda3\envs\heurllib\site-packages\mindspore\nnlcell.py", line 482, in _run_construct

output = self.construct(*cast_inputs, **kwargs)

File "C:\Users\LENOVO\Desktop\hetur leargenelmethodslutils\modelslmodels.py", line 96 , in construct

x = layer(x)

File "F: lananconda3lenvs\heurllib\site-packages\mindspore\nnlcell.py", line 705, in --call_-

raise err

F1le "F:\ananconda3\envs\heurl1iblsite-packageslmindsporelnnlcell.py", Line 701, in -_call--

output = self._run_construct(args, kwargs)

File "F:\ananconda3\envs\heur|lib\site-packages\mindsporelnnlcell.py", line 482, in _run_construct

output = self.construct(*cast_inputs, **kwargs)

File "F:lananconda3\envslheurlliblsite-packageslmindsporelnnllayerlbasic.py", line 623,in construct

x_shape = self.shape_op(x)

File "F:\ananconda3lenvslheurlliblssite-packages\mindsporelopslopeerationslarray ops.py", line 701, in __call__

return x.shape

File "F: \ananconda3lenvs\heurlliblsite-packageslmindsporelcommonl stub tensor.py" line85, in shape

self.stub_shape = self.stub.get_shape()

RuntimeError: could L not create descriptor for dilated convolution forward propagation primitive