import json
import os
import random
import shutil
import time

from tests.utils.lineage_writer._summary_record import LineageSummary


def generate_graph(dataset_name='MnistDataset', batch_size=16, buffer_size=10, rescale=0.003921, num_samples=100, normalize_weight=0.48):
    dataset_graph = {
        'num_parallel_workers': None,
        'op_type': 'BatchDataset',
        'op_module': 'minddata.dataengine.datasets',
        'drop_remainder': True,
        'batch_size': batch_size,
        'children': [{
            'op_type': 'ShuffleDataset',
            'op_module': 'minddata.dataengine.datasets',
            'num_parallel_workers': None,
            'buffer_size': buffer_size,
            'children': [{
                'op_type': 'MapDataset',
                'op_module': 'minddata.dataengine.datasets',
                'num_parallel_workers': None,
                'output_columns': [None],
                'input_columns': ['label'],
                'operations': [{
                    'tensor_op_module': 'minddata.transforms.c_transforms',
                    'tensor_op_name': 'OneHot',
                    'num_classes': 10
                }],
                'children': [{
                    'op_type': 'MapDataset',
                    'op_module': 'minddata.dataengine.datasets',
                    'num_parallel_workers': None,
                    'input_columns': ['image'],
                    'operations': [
                        {
                            'tensor_op_module': 'minddata.transforms.c_transforms',
                            'tensor_op_name': 'RandomCrop',
                            'weight': [32, 32, 4, 4, 4, 4],
                            'padding_mode': "constant",
                            'pad_if_needed': False,
                            'fill_value': 0
                        },
                        {
                            'tensor_op_module': 'minddata.transforms.c_transforms',
                            'tensor_op_name': 'Rescale',
                            'rescale': rescale,
                            'shift': 0,
                            'num_classes': 10
                        },
                        {
                            'tensor_op_module': 'minddata.transforms.c_transforms',
                            'tensor_op_name': 'Normalize',
                            'weights': [normalize_weight]
                        }
                    ],
                    'children': [{
                        'dataset_dir': '/home/anthony/MindData/tests/dataset/data/testMnistData',
                        'op_module': 'minddata.dataengine.datasets',
                        'num_shards': None,
                        'num_parallel_workers': None,
                        'shuffle': None,
                        'op_type': dataset_name,
                        'shard_id': None,
                        'num_samples': num_samples,
                        'sampler': {
                            'sampler_module': 'minddata.dataengine.samplers',
                            'sampler_name': 'RandomSampler',
                            'replacement': True,
                            'num_samples': num_samples
                        },
                        'children': []
                    }]
                }]
            }]
        }]
    }
    return dataset_graph


class LineageWriter:
    def __init__(self):
        self.train_args = dict()
        self.train_args["train_network"] = "LeNet5"
        self.train_args["loss"] = 0.01
        self.train_args["learning_rate"] = 0.01
        self.train_args["optimizer"] = "Momentum"
        self.train_args["loss_function"] = "SoftmaxCrossEntropyWithLogits"
        self.train_args["epoch"] = 500
        self.train_args["parallel_mode"] = ""
        self.train_args["device_num"] = 1
        self.train_args["batch_size"] = 32
        self.train_args["map_model"] = 1
        self.train_args["train_dataset_path"] = "/home/data/train"
        self.train_args["train_dataset_size"] = 301234
        self.train_args["model_path"] = "/home/demo/demo_model.pkl"
        self.train_args["model_size"] = 100 * 1024 * 1024

        self.eval_args = dict()
        self.eval_args["metrics"] = json.dumps({"acc": 0.88})
        self.eval_args["valid_dataset_path"] = "/home/data/test"
        self.eval_args["valid_dataset_size"] = 5000

        # self.base_dir = "./demo_lineage_r0.3_0530"
        self.base_dir = "/data/luopengting/summaries/demo_param_importances/"
        self._run_num = 1000

        # self._file_lineage = f'events.out.summary.{int(time.time())}.local_MS_lineage'

    def init_dirs(self, clean_base_dir=False):
        if clean_base_dir and os.path.exists(self.base_dir):
            shutil.rmtree(self.base_dir)
        if not os.path.exists(self.base_dir):
            os.makedirs(self.base_dir)

        for i in range(1, self._run_num):
            sub_dir = os.path.join(self.base_dir, f'run{i}')
            os.makedirs(sub_dir, exist_ok=True)

    def create_eval_event(self):
        eval_args = self.eval_args
        for i in range(1, self._run_num):
            summary_dir_path = os.path.join(self.base_dir, f'run{i}')
            lineage_summary = LineageSummary(summary_dir_path)
            eval_args["valid_dataset_size"] = random.choice([13, 24, 28]) * 100
            eval_args["metrics"] = json.dumps({'Accuracy': random.uniform(0.85, 0.96)})
            lineage_summary.record_evaluation_lineage(eval_args)
            # lineage_summary.record_user_defined_info({'Version': f'v{i%3}'})
            user_defined_info = {}
            for i in range(105):
                user_defined_info.update({f'n{i}': random.random()})
            lineage_summary.record_user_defined_info(user_defined_info)
            # time.sleep(1)

    def create_train_event(self):
        network = ['ResNet', 'LeNet5', 'AlexNet']
        optimizer = ['SGD', 'Adam', 'Momentum']
        train_args = self.train_args

        for i in range(1, self._run_num):
            summary_dir_path = os.path.join(self.base_dir, f'run{i}')
            lineage_summary = LineageSummary(summary_dir_path)
            train_args['learning_rate'] = random.uniform(0.001, 0.005)
            train_args["loss"] = random.uniform(0.001, 0.005)
            train_args["epoch"] = random.choice([100, 200, 300])
            train_args["batch_size"] = random.choice([16, 32, 64])
            train_args["model_size"] = random.randint(350, 450) * 1024 * 1024
            train_args["train_network"] = random.choice(network)
            train_args["optimizer"] = random.choice(optimizer)
            train_args["device_num"] = random.choice([1, 2, 4, 6, 8])
            train_args["loss_function"] = random.choice(["SoftmaxCrossEntropyWithLogits", "CrossEntropyLoss"])
            train_args["train_dataset_size"] = random.choice([56, 67, 78]) * 10000
            lineage_summary.record_train_lineage(train_args)
            # lineage_summary.record_user_defined_info({'train_time': time.strftime('%m-%d %H:%M:%S', time.localtime(time.time()))})

            dataset_graph = generate_graph(
                dataset_name=random.choice(['MindDataset']),  # 'MnistDataset', 'Cifar10Datset'
                batch_size=random.choice([8, 16, 32, 64]),
                buffer_size=random.choice([10, 20, 30]),
                rescale=random.choice([0.003921, 0.005632, 0.0078, 0.005678]),
                num_samples=random.choice([100, 200, 300]),
                normalize_weight=random.choice([0.20, 0.50])  # random.uniform(0.2, 0.5)
            )
            lineage_summary.record_dataset_graph(dataset_graph)
            # time.sleep(1)

    def create_summary(self):
        self.init_dirs()
        self.create_train_event()
        self.create_eval_event()


lw = LineageWriter()
lw.create_summary()