from cProfile import label
from torch.utils.data import Dataset
import os
import numpy as np
import torch
import csv
import re
import scipy.io as scio
import torch.nn.functional as F
from icecream import  ic

class fMRI_Hidden_Dataset(Dataset):

    def __init__(self, feature_path, exp_path, mode='AD_NC'):
        file_list = []
        feature_list = os.listdir(feature_path)
        with open(exp_path, 'r') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                if row[0]:
                    filename = 'ROISignals_Sub_' + row[0]
                    if filename in feature_list:
                        filepath = os.path.join(feature_path, filename)
                        file_list.append(filepath)
        print("Classification mode: {}, total {} samples.".format(mode, len(file_list)))
        self.file_list = file_list
        self.mode = mode

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        data = torch.load(self.file_list[idx])
        feature = data['feature']
        adj = data['adj']
        adj_discrete = False
        if adj_discrete:
            adj = adj[adj > np.mean(adj)]
        if self.mode == 'MCI_NC':
            label = int(data['label'].item())
        else:
            label = round(int(data['label'].item()) / 2)
        return {'feature': feature.detach().numpy(), 'label': np.array(label), 'adj': adj.detach().numpy()}

class mean_bold_Dataset(Dataset):
    def __init__(self, feature_path, exp_path, mode='AD_NC'):
        file_list = []
        label_list = []
        feature_list = os.listdir(feature_path)
        with open(exp_path, 'r') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                if row[0]:
                    filename = 'ROISignals_Sub_' + row[0] + '.mat'
                    if filename in feature_list:
                        filepath = os.path.join(feature_path, filename)
                        file_list.append(filepath)
                        if row[2] == 'AD':
                            label_list.append(0)
                        elif row[2] == 'NC':
                            label_list.append(1)
        """file_list = []
        label_list = []
        feature_list = os.listdir(feature_path)
        with open(exp_path, 'r') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                if row[0]:
                    filename = 'Sub_' + row[0] + '.pt'
                    if filename in feature_list:
                        filepath = os.path.join(feature_path, filename)
                        file_list.append(filepath)
                        if row[2] == 'AD':
                            label_list.append(0)
                        elif row[2] == 'NC':
                            label_list.append(1)"""

        print("Classification mode: {}, total {} samples.".format(mode, len(file_list)))
        self.file_list = file_list
        self.label_list = label_list
        self.mode = mode

        """with open(exp_path, 'r') as csvfile:
            reader = csv.reader(csvfile)
            for row in reader:
                if row[0]:
                    if ('fMRI' in os.path.basename(feature_path)) \
                            or ('feature_AGNN' in os.path.basename(feature_path)) \
                            or ('AGNN_feature' in os.path.basename(feature_path)):
                        filename = re.sub('Sub', 'Sub_', row[1])
                        filename = re.sub('nii', 'npy', filename)
                    else:
                        filename = 'Sub_' + row[0] + '.npy'
                    if filename in feature_list:
                        filepath = os.path.join(feature_path, filename)
                        file_list.append(filepath)
                        if row[2] == 'AD':
                            label_list.append(2)
                        elif row[2] == 'MCI':
                            label_list.append(1)
                        elif row[2] == 'NC':
                            label_list.append(0)"""

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):

        sample_data = scio.loadmat(self.file_list[idx])
        ROI_signals = sample_data['ROISignals']
        ROI_signals = ROI_signals.T
        mean_sig = np.mean(ROI_signals)
        std_sig = np.std(ROI_signals)
        normal_ROI_signals = (ROI_signals - mean_sig) / std_sig
        ROI_signals = normal_ROI_signals

        adj = np.corrcoef(ROI_signals)
        adj = torch.from_numpy(adj)

        labels = self.label_list[idx]

        return {'feature': torch.Tensor(ROI_signals), 'label': labels, 'adj': torch.Tensor(adj)}

# class fMRI_REHO_ALFF_Dataset(Dataset):

#     def __init__(self, feature_dir,exp_path, cishu, phase = 'ALFF', mode='AD_NC',state='train'):
#         self.mode = mode
#         feature_list = []
#         if (phase == 'ALFF_REHO') or (phase == 'REHO_ALFF'):
#             for i in range(1,91):
#                 feature_name_AD_NC = 'roi-'+str(i)+'AD+NC'+'cishu'+str(cishu)+'.npz'
#                 feature_path_AD_NC_alff = os.path.join(feature_dir,'ALFF',feature_name_AD_NC)
#                 feature_path_AD_NC_reho = os.path.join(feature_dir,'REHO',feature_name_AD_NC)
#                 data_AD_NC_alff = np.load(feature_path_AD_NC_alff, allow_pickle=True)['data'][slice_AD_NC,:]
#                 data_AD_NC_reho = np.load(feature_path_AD_NC_reho, allow_pickle=True)['data'][slice_AD_NC,:]
#                 data_AD_NC = np.hstack((data_AD_NC_reho,data_AD_NC_alff))
#                 feature_name_AD_MCI = 'roi-'+str(i)+'AD+MCI'+'cishu'+str(cishu)+'.npz'
#                 feature_path_AD_MCI_alff = os.path.join(feature_dir,'ALFF',feature_name_AD_MCI)
#                 feature_path_AD_MCI_reho = os.path.join(feature_dir,'REHO',feature_name_AD_MCI)
#                 data_AD_MCI_alff = np.load(feature_path_AD_MCI_alff, allow_pickle=True)['data'][191:286,:]
#                 data_AD_MCI_reho = np.load(feature_path_AD_MCI_reho, allow_pickle=True)['data'][191:286,:]
#                 data_AD_MCI = np.hstack((data_AD_MCI_reho,data_AD_MCI_alff))
#                 data = np.vstack((data_AD_NC,data_AD_MCI))
#                 feature_list.append(data)
#         else:
#             for i in range(1,91):
#                 feature_name_AD_NC = 'roi-'+str(i)+'AD+NC'+'cishu'+str(cishu)+'.npz'
#                 feature_path_AD_NC = os.path.join(feature_dir,phase,feature_name_AD_NC)
#                 data_AD_NC = np.load(feature_path_AD_NC, allow_pickle=True)['data'][slice_AD_NC,:]
#                 feature_name_AD_MCI = 'roi-'+str(i)+'AD+MCI'+'cishu'+str(cishu)+'.npz'
#                 feature_path_AD_MCI = os.path.join(feature_dir,phase,feature_name_AD_MCI)
#                 data_AD_MCI = np.load(feature_path_AD_MCI, allow_pickle=True)['data'][191:286,:]
#                 data = np.vstack((data_AD_NC,data_AD_MCI))
#                 feature_list.append(data)
#         feature_list = np.array(feature_list)
#         label_list = np.hstack((np.ones(149)*2,np.zeros(102),np.ones(95)))
#         # self.label_list = np.hstack((np.ones(191)*2,np.ones(293)))

#         exp_sub_list = []
#         with open(exp_path,'r') as csvfile:
#             reader = csv.reader(csvfile)
#             for row in reader:
#                 if row[0]:
#                     exp_sub_list.append(int(row[0]))
#         if state == 'test':
#             exp_sub_list = exp_sub_list[26:34]
#         self.label_list = label_list[exp_sub_list]
#         self.feature_list = feature_list[:,exp_sub_list,:]
#         # exp_path = os.path.join(feature_dir,'AD+NC'+state+'_mask_list.npy')
#         # sub_list = np.load(exp_path,allow_pickle=True)[cishu]
#         # self.sub_split_list = [idx for idx in sub_list if idx in slice_AD_NC]
#         # self.feature_list = self.feature_list[:,sub_list,:]
#         # self.label_list = self.label_list[sub_list]s
#         print("Classification mode: {}, total {} samples.".format(mode,len(self.label_list)))

#         # print("Classification mode: {}, total {} samples.".format(mode,len(self.label_list)))

#     def __len__(self):
#         return len(self.label_list)

#     def __getitem__(self, idx):
#         feature = self.feature_list[:,idx,:].squeeze()
#         adj = np.corrcoef(feature)
#         adj_discrete=False
#         if adj_discrete:
#             adj = adj[adj>np.mean(adj)]
#         if self.mode == 'MCI_NC' or self.mode == 'AD_MCI_NC':
#             label = round(self.label_list[idx])
#         else:
#             label = round(self.label_list[idx]/2)
#         return {'feature':feature, 'label':np.array(label), 'adj': adj}

class fMRI_REHO_ALFF_Dataset(Dataset):

    def __init__(self, feature_dir, exp_path, cishu, phase='ALFF', mode='AD_NC', state='train'):
        self.mode = mode
        path_mode = re.sub('_', '+', mode)
        feature_name = 'roi-1' + path_mode + 'cishu' + str(cishu) + '.npz'
        sub_num = np.load(os.path.join(feature_dir, phase, feature_name), allow_pickle=True)['data'].shape[0]
        feature_list = np.zeros((sub_num, 90, 64))
        if (phase == 'ALFF_REHO') or (phase == 'REHO_ALFF'):
            for i in range(1, 91):
                feature_name = 'roi-' + str(i) + path_mode + 'cishu' + str(cishu) + '.npz'
                feature_path_alff = os.path.join(feature_dir, 'ALFF', feature_name)
                feature_path_reho = os.path.join(feature_dir, 'REHO', feature_name)
                data_alff = np.load(feature_path_alff, allow_pickle=True)['data']
                data_reho = np.load(feature_path_reho, allow_pickle=True)['data']
                data = np.hstack((data_reho, data_alff))
                feature_list[:, i - 1, :] = data
        else:
            for i in range(1, 91):
                feature_name = 'roi-' + str(i) + path_mode + 'cishu' + str(cishu) + '.npz'
                feature_path = os.path.join(feature_dir, phase, feature_name)
                data = np.load(feature_path, allow_pickle=True)['data']
                feature_list[:, i - 1, :] = data
        if mode == 'AD_NC':
            label_list = np.hstack((np.ones(191) * 2, np.zeros(167)))
            slice_list = np.hstack((np.arange(149), np.arange(191, 293)))
        elif mode == 'AD_MCI':
            label_list = np.hstack((np.ones(191) * 2, np.ones(160)))
            slice_list = np.hstack((np.arange(149), np.arange(191, 286)))
        else:
            label_list = np.hstack((np.zeros(167), np.ones(160)))
            slice_list = np.hstack((np.arange(102), np.arange(167, 262)))

        exp_path = os.path.join(feature_dir, path_mode + state + '_mask_list.npy')
        sub_list = np.load(exp_path, allow_pickle=True)[cishu]
        sub_list = [idx for idx in sub_list if idx in slice_list]
        self.feature_list = feature_list[sub_list, :]
        self.label_list = label_list[sub_list]
        print(feature_list.shape, label_list.shape)
        print("Classification mode: {}, total {} samples.".format(mode, len(self.label_list)))

        # print("Classification mode: {}, total {} samples.".format(mode,len(self.label_list)))

    def __len__(self):
        return len(self.label_list)

    def __getitem__(self, idx):
        feature = self.feature_list[idx, :, :].squeeze()
        adj = np.corrcoef(feature)
        adj_discrete = False
        if adj_discrete:
            adj = adj[adj > np.mean(adj)]
        if (self.mode == 'NC_MCI') or (self.mode == 'AD_MCI_NC') or (self.mode == 'MCI_NC'):
            label = round(self.label_list[idx])
        else:
            label = round(self.label_list[idx] / 2)
        return {'feature': feature, 'label': np.array(label), 'adj': adj}

if __name__ == "__main__":
    feature_path = '/home/ding/exp_4/data/AGNN_feature_set/2022-07-12_03-58-46/ch50_chll50_nu66/Sub_Dataset'
    # feature_path = '/home/ding/exp_4/data/AGNN_MoCo'
    # feature_path = '/home/ding/exp_4/data/fMRI_Sub_avg_norm'
    # feature_path = '/home/ding/exp_4/data/fMRI_feature/feature_AGNN'
    # feature_path = '/home/ding/exp_4/data/MoCo_KMeans/MoCo_Sub_dataset'
    # feature_path = '/home/ding/exp_4/data/fMRI_norm_DeepDPM/fMRI_DeepDPM_aal_output/AGNN_feature'
    # feature_path = '/home/ding/exp_4/data/脑区图(AGNN)_DTI+fMRI/hiddenfMRI'

    # feature_path = '/home/ding/exp_4/data/MoCo_KMeans/cluster2_Sub_dataset'
    # feature_path='/home/ding/exp_4/data/MoCo_KMeans/GMM_cluster10_Sub_dataset'
    feature_path = '/home/ding/exp_4/data/MoCo_KMeans/Birch_cluster10_Sub_dataset'
    # exp_path = '/home/ding/exp_4/data/fold_split_MCI_NC/exp_0.csv'
    exp_train_path = '/home/ding/exp_4/data/fold_split_AD_NC/trainset_0.csv'
    exp_test_path = '/home/ding/exp_4/data/fold_split_AD_NC/testset_0.csv'
    # exp_train_path='/home/ding/exp_4/data/fold_split_AD_MCI/trainset_1.csv'
    # exp_test_path='/home/ding/exp_4/data/fold_split_AD_MCI/testset_1.csv'

    # testset = fMRI_Avg_Dataset(feature_path, exp_test_path, mode='AD_NC')
    # trainset = fMRI_Avg_Dataset(feature_path, exp_train_path, mode='AD_NC')
    # testset = fMRI_Hidden_Dataset(feature_path,exp_test_path,mode='MCI_NC')
    # trainset = fMRI_Hidden_Dataset(feature_path,exp_train_path,mode='MCI_NC')
    cishu = 0
    mode = 'NC_MCI'
    phase = 'REHO'
    state = 'train'
    # trainset = fMRI_REHO_ALFF_Dataset(feature_path,exp_train_path,cishu,phase,mode,'train')
    # testset = fMRI_REHO_ALFF_Dataset(feature_path,exp_test_path,cishu,phase,mode,'test')

    # print(len(trainset))
    # print(len(testset))
    # for i in range(40):
    #     print(testset[i]['feature'].shape,testset[i]['adj'].shape,testset[i]['label'])
    # for i in range(43):
    #     print(trainset[i]['feature'].shape,trainset[i]['adj'].shape,trainset[i]['label'])

    # print(testset[0]['feature'].shape, testset[0]['adj'].shape, testset[0]['label'])
    # print(testset[20]['feature'].shape, testset[20]['adj'].shape, testset[20]['label'])
    # print(testset[0]['adj'])
    # print(testset[47]['adj'])
    # print(testset[48]['adj'])

    # import numpy as np
    # from sklearn import svm
    # from sklearn.model_selection import train_test_split as ts
    # slice_AD_NC = np.hstack((np.arange(149),np.arange(191,293)))
    # feature_path = '/home/ding/exp_4/data/脑区图(AGNN)_DTI+fMRI/hiddenfMRI'
    # feature_dir = feature_path
    # feature_list = np.zeros((251,90,64))
    # cishu = 20
    # for i in range(1,91):
    #     feature_name_AD_NC = 'roi-'+str(i)+'AD+NC'+'cishu'+str(cishu)+'.npz'
    #     feature_path_AD_NC = os.path.join(feature_dir,phase,feature_name_AD_NC)
    #     data_AD_NC = np.load(feature_path_AD_NC, allow_pickle=True)['data'][slice_AD_NC,:]
    #     # feature_name_AD_MCI = 'roi-'+str(i)+'AD+MCI'+'cishu'+str(cishu)+'.npz'
    #     # feature_path_AD_MCI = os.path.join(feature_dir,phase,feature_name_AD_MCI)
    #     # data_AD_MCI = np.load(feature_path_AD_MCI, allow_pickle=True)['data'][191:286,:]
    #     # data = np.vstack((data_AD_NC,data_AD_MCI))
    #     # feature_list.append(data)
    #     feature_list[:,i-1,:] = data_AD_NC
    # feature_list = feature_list.reshape((251,90*64))
    # label_list = np.hstack((np.ones(149)*2,np.zeros(102)))
    # print('Total feature shape {} label shape {}'
    #       .format(feature_list.shape,label_list.shape))

    # exp_path_train = os.path.join(feature_dir,'AD+NCtrain_mask_list.npy')
    # sub_list_train = np.load(exp_path_train,allow_pickle=True)[cishu]
    # sub_split_list_train = [idx for idx in sub_list_train if idx in slice_AD_NC]
    # feature_list = feature_list[sub_split_list_train,:]
    # label_list = label_list[sub_split_list_train]
    # print('feature shape {} label shape {}'.format(feature_list.shape,label_list.shape))

    # ratio = 0.9
    # X_train, X_test, Y_train, Y_test = ts(feature_list,label_list,test_size=ratio)
    # print('Train set shape of {} train label shape of {} \nTest set shape of {} Test label shape of {}'
    #       .format(X_train.shape,Y_train.shape,X_test.shape,Y_test.shape))

    # classifier = svm.SVC()
    # classifier.fit(X_train,Y_train)
    # for i in range(17):
    #     predict = classifier.predict(X_test[i,np.newaxis])
    #     print(i,X_test[i],predict,Y_test[i])
    # print('cishu {}\ttest size of {}\tclassification avg score {}'.format(cishu,ratio,classifier.score(X_test,Y_test)))

    # AD_list = []
    # MCI_list = []
    # for i in range(len(trainset)):
    #     if trainset[i]['label'] == 1:
    #         AD_list.append(trainset[i]['adj'])
    #     else:
    #         MCI_list.append(trainset[i]['adj'])
    # for j in range(len(testset)):
    #     if testset[j]['label'] == 1:
    #         AD_list.append(testset[j]['adj'])
    #     else:
    #         MCI_list.append(testset[j]['adj'])
    # AD_list = np.array(AD_list)
    # MCI_list = np.array(MCI_list)
    # print(AD_list.shape,MCI_list.shape)
    # AD_avg = np.mean(AD_list,axis=0)
    # MCI_avg = np.mean(MCI_list,axis=0)
    # print(AD_avg.shape,MCI_avg.shape)
    # # np.save('/home/ding/exp_4/test/debug/AD_adj_mean.npy',AD_avg)
    # np.save('/home/ding/exp_4/test/debug/NC_adj_mean.npy',MCI_avg)


# -*- coding: utf-8 -*-

import unittest

import pytest

from docker.constants import DEFAULT_DOCKER_API_VERSION
from docker.errors import InvalidArgument, InvalidVersion
from docker.types import (
    EndpointConfig, HostConfig, IPAMConfig, IPAMPool, LogConfig, Mount,
    ServiceMode, Ulimit,
)

try:
    from unittest import mock
except:
    import mock

def create_host_config(*args, **kwargs):
    return HostConfig(*args, **kwargs)

class HostConfigTest(unittest.TestCase):
    def test_create_host_config_no_options(self):
        config = create_host_config(version='1.19')
        self.assertFalse('NetworkMode' in config)

    def test_create_host_config_no_options_newer_api_version(self):
        config = create_host_config(version='1.20')
        self.assertEqual(config['NetworkMode'], 'default')

    def test_create_host_config_invalid_cpu_cfs_types(self):
        with pytest.raises(TypeError):
            create_host_config(version='1.20', cpu_quota='0')

        with pytest.raises(TypeError):
            create_host_config(version='1.20', cpu_period='0')

        with pytest.raises(TypeError):
            create_host_config(version='1.20', cpu_quota=23.11)

        with pytest.raises(TypeError):
            create_host_config(version='1.20', cpu_period=1999.0)

    def test_create_host_config_with_cpu_quota(self):
        config = create_host_config(version='1.20', cpu_quota=1999)
        self.assertEqual(config.get('CpuQuota'), 1999)

    def test_create_host_config_with_cpu_period(self):
        config = create_host_config(version='1.20', cpu_period=1999)
        self.assertEqual(config.get('CpuPeriod'), 1999)

    def test_create_host_config_with_blkio_constraints(self):
        blkio_rate = [{"Path": "/dev/sda", "Rate": 1000}]
        config = create_host_config(version='1.22',
                                    blkio_weight=1999,
                                    blkio_weight_device=blkio_rate,
                                    device_read_bps=blkio_rate,
                                    device_write_bps=blkio_rate,
                                    device_read_iops=blkio_rate,
                                    device_write_iops=blkio_rate)

        self.assertEqual(config.get('BlkioWeight'), 1999)
        self.assertTrue(config.get('BlkioWeightDevice') is blkio_rate)
        self.assertTrue(config.get('BlkioDeviceReadBps') is blkio_rate)
        self.assertTrue(config.get('BlkioDeviceWriteBps') is blkio_rate)
        self.assertTrue(config.get('BlkioDeviceReadIOps') is blkio_rate)
        self.assertTrue(config.get('BlkioDeviceWriteIOps') is blkio_rate)
        self.assertEqual(blkio_rate[0]['Path'], "/dev/sda")
        self.assertEqual(blkio_rate[0]['Rate'], 1000)

    def test_create_host_config_with_shm_size(self):
        config = create_host_config(version='1.22', shm_size=67108864)
        self.assertEqual(config.get('ShmSize'), 67108864)

    def test_create_host_config_with_shm_size_in_mb(self):
        config = create_host_config(version='1.22', shm_size='64M')
        self.assertEqual(config.get('ShmSize'), 67108864)

    def test_create_host_config_with_oom_kill_disable(self):
        config = create_host_config(version='1.20', oom_kill_disable=True)
        self.assertEqual(config.get('OomKillDisable'), True)
        self.assertRaises(
            InvalidVersion, lambda: create_host_config(version='1.18.3',
                                                       oom_kill_disable=True))

    def test_create_host_config_with_userns_mode(self):
        config = create_host_config(version='1.23', userns_mode='host')
        self.assertEqual(config.get('UsernsMode'), 'host')
        self.assertRaises(
            InvalidVersion, lambda: create_host_config(version='1.22',
                                                       userns_mode='host'))
        self.assertRaises(
            ValueError, lambda: create_host_config(version='1.23',
                                                   userns_mode='host12'))

    def test_create_host_config_with_oom_score_adj(self):
        config = create_host_config(version='1.22', oom_score_adj=100)
        self.assertEqual(config.get('OomScoreAdj'), 100)
        self.assertRaises(
            InvalidVersion, lambda: create_host_config(version='1.21',
                                                       oom_score_adj=100))
        self.assertRaises(
            TypeError, lambda: create_host_config(version='1.22',
                                                  oom_score_adj='100'))

    def test_create_host_config_with_dns_opt(self):

        tested_opts = ['use-vc', 'no-tld-query']
        config = create_host_config(version='1.21', dns_opt=tested_opts)
        dns_opts = config.get('DnsOptions')

        self.assertTrue('use-vc' in dns_opts)
        self.assertTrue('no-tld-query' in dns_opts)

        self.assertRaises(
            InvalidVersion, lambda: create_host_config(version='1.20',
                                                       dns_opt=tested_opts))

    def test_create_host_config_with_mem_reservation(self):
        config = create_host_config(version='1.21', mem_reservation=67108864)
        self.assertEqual(config.get('MemoryReservation'), 67108864)
        self.assertRaises(
            InvalidVersion, lambda: create_host_config(
                version='1.20', mem_reservation=67108864))

    def test_create_host_config_with_kernel_memory(self):
        config = create_host_config(version='1.21', kernel_memory=67108864)
        self.assertEqual(config.get('KernelMemory'), 67108864)
        self.assertRaises(
            InvalidVersion, lambda: create_host_config(
                version='1.20', kernel_memory=67108864))

    def test_create_host_config_with_pids_limit(self):
        config = create_host_config(version='1.23', pids_limit=1024)
        self.assertEqual(config.get('PidsLimit'), 1024)

        with pytest.raises(InvalidVersion):
            create_host_config(version='1.22', pids_limit=1024)
        with pytest.raises(TypeError):
            create_host_config(version='1.23', pids_limit='1024')

    def test_create_host_config_with_isolation(self):
        config = create_host_config(version='1.24', isolation='hyperv')
        self.assertEqual(config.get('Isolation'), 'hyperv')

        with pytest.raises(InvalidVersion):
            create_host_config(version='1.23', isolation='hyperv')
        with pytest.raises(TypeError):
            create_host_config(
                version='1.24', isolation={'isolation': 'hyperv'}
            )

    def test_create_host_config_pid_mode(self):
        with pytest.raises(ValueError):
            create_host_config(version='1.23', pid_mode='baccab125')

        config = create_host_config(version='1.23', pid_mode='host')
        assert config.get('PidMode') == 'host'
        config = create_host_config(version='1.24', pid_mode='baccab125')
        assert config.get('PidMode') == 'baccab125'

    def test_create_host_config_invalid_mem_swappiness(self):
        with pytest.raises(TypeError):
            create_host_config(version='1.24', mem_swappiness='40')

class UlimitTest(unittest.TestCase):
    def test_create_host_config_dict_ulimit(self):
        ulimit_dct = {'name': 'nofile', 'soft': 8096}
        config = create_host_config(
            ulimits=[ulimit_dct], version=DEFAULT_DOCKER_API_VERSION
        )
        self.assertIn('Ulimits', config)
        self.assertEqual(len(config['Ulimits']), 1)
        ulimit_obj = config['Ulimits'][0]
        self.assertTrue(isinstance(ulimit_obj, Ulimit))
        self.assertEqual(ulimit_obj.name, ulimit_dct['name'])
        self.assertEqual(ulimit_obj.soft, ulimit_dct['soft'])
        self.assertEqual(ulimit_obj['Soft'], ulimit_obj.soft)

    def test_create_host_config_dict_ulimit_capitals(self):
        ulimit_dct = {'Name': 'nofile', 'Soft': 8096, 'Hard': 8096 * 4}
        config = create_host_config(
            ulimits=[ulimit_dct], version=DEFAULT_DOCKER_API_VERSION
        )
        self.assertIn('Ulimits', config)
        self.assertEqual(len(config['Ulimits']), 1)
        ulimit_obj = config['Ulimits'][0]
        self.assertTrue(isinstance(ulimit_obj, Ulimit))
        self.assertEqual(ulimit_obj.name, ulimit_dct['Name'])
        self.assertEqual(ulimit_obj.soft, ulimit_dct['Soft'])
        self.assertEqual(ulimit_obj.hard, ulimit_dct['Hard'])
        self.assertEqual(ulimit_obj['Soft'], ulimit_obj.soft)

    def test_create_host_config_obj_ulimit(self):
        ulimit_dct = Ulimit(name='nofile', soft=8096)
        config = create_host_config(
            ulimits=[ulimit_dct], version=DEFAULT_DOCKER_API_VERSION
        )
        self.assertIn('Ulimits', config)
        self.assertEqual(len(config['Ulimits']), 1)
        ulimit_obj = config['Ulimits'][0]
        self.assertTrue(isinstance(ulimit_obj, Ulimit))
        self.assertEqual(ulimit_obj, ulimit_dct)

    def test_ulimit_invalid_type(self):
        self.assertRaises(ValueError, lambda: Ulimit(name=None))
        self.assertRaises(ValueError, lambda: Ulimit(name='hello', soft='123'))
        self.assertRaises(ValueError, lambda: Ulimit(name='hello', hard='456'))

class LogConfigTest(unittest.TestCase):
    def test_create_host_config_dict_logconfig(self):
        dct = {'type': LogConfig.types.SYSLOG, 'config': {'key1': 'val1'}}
        config = create_host_config(
            version=DEFAULT_DOCKER_API_VERSION, log_config=dct
        )
        self.assertIn('LogConfig', config)
        self.assertTrue(isinstance(config['LogConfig'], LogConfig))
        self.assertEqual(dct['type'], config['LogConfig'].type)

    def test_create_host_config_obj_logconfig(self):
        obj = LogConfig(type=LogConfig.types.SYSLOG, config={'key1': 'val1'})
        config = create_host_config(
            version=DEFAULT_DOCKER_API_VERSION, log_config=obj
        )
        self.assertIn('LogConfig', config)
        self.assertTrue(isinstance(config['LogConfig'], LogConfig))
        self.assertEqual(obj, config['LogConfig'])

    def test_logconfig_invalid_config_type(self):
        with pytest.raises(ValueError):
            LogConfig(type=LogConfig.types.JSON, config='helloworld')

class EndpointConfigTest(unittest.TestCase):
    def test_create_endpoint_config_with_aliases(self):
        config = EndpointConfig(version='1.22', aliases=['foo', 'bar'])
        assert config == {'Aliases': ['foo', 'bar']}

        with pytest.raises(InvalidVersion):
            EndpointConfig(version='1.21', aliases=['foo', 'bar'])

class IPAMConfigTest(unittest.TestCase):
    def test_create_ipam_config(self):
        ipam_pool = IPAMPool(subnet='192.168.52.0/24',
                             gateway='192.168.52.254')

        ipam_config = IPAMConfig(pool_configs=[ipam_pool])
        self.assertEqual(ipam_config, {
            'Driver': 'default',
            'Config': [{
                'Subnet': '192.168.52.0/24',
                'Gateway': '192.168.52.254',
                'AuxiliaryAddresses': None,
                'IPRange': None,
            }]
        })

class ServiceModeTest(unittest.TestCase):
    def test_replicated_simple(self):
        mode = ServiceMode('replicated')
        assert mode == {'replicated': {}}
        assert mode.mode == 'replicated'
        assert mode.replicas is None

    def test_global_simple(self):
        mode = ServiceMode('global')
        assert mode == {'global': {}}
        assert mode.mode == 'global'
        assert mode.replicas is None

    def test_global_replicas_error(self):
        with pytest.raises(InvalidArgument):
            ServiceMode('global', 21)

    def test_replicated_replicas(self):
        mode = ServiceMode('replicated', 21)
        assert mode == {'replicated': {'Replicas': 21}}
        assert mode.mode == 'replicated'
        assert mode.replicas == 21

    def test_invalid_mode(self):
        with pytest.raises(InvalidArgument):
            ServiceMode('foobar')

class MountTest(unittest.TestCase):
    def test_parse_mount_string_ro(self):
        mount = Mount.parse_mount_string("/foo/bar:/baz:ro")
        assert mount['Source'] == "/foo/bar"
        assert mount['Target'] == "/baz"
        assert mount['ReadOnly'] is True

    def test_parse_mount_string_rw(self):
        mount = Mount.parse_mount_string("/foo/bar:/baz:rw")
        assert mount['Source'] == "/foo/bar"
        assert mount['Target'] == "/baz"
        assert not mount['ReadOnly']

    def test_parse_mount_string_short_form(self):
        mount = Mount.parse_mount_string("/foo/bar:/baz")
        assert mount['Source'] == "/foo/bar"
        assert mount['Target'] == "/baz"
        assert not mount['ReadOnly']

    def test_parse_mount_string_no_source(self):
        mount = Mount.parse_mount_string("foo/bar")
        assert mount['Source'] is None
        assert mount['Target'] == "foo/bar"
        assert not mount['ReadOnly']

    def test_parse_mount_string_invalid(self):
        with pytest.raises(InvalidArgument):
            Mount.parse_mount_string("foo:bar:baz:rw")

    def test_parse_mount_named_volume(self):
        mount = Mount.parse_mount_string("foobar:/baz")
        assert mount['Source'] == 'foobar'
        assert mount['Target'] == '/baz'
        assert mount['Type'] == 'volume'

    def test_parse_mount_bind(self):
        mount = Mount.parse_mount_string('/foo/bar:/baz')
        assert mount['Source'] == "/foo/bar"
        assert mount['Target'] == "/baz"
        assert mount['Type'] == 'bind'

    @pytest.mark.xfail
    def test_parse_mount_bind_windows(self):
        with mock.patch('docker.types.services.IS_WINDOWS_PLATFORM', True):
            mount = Mount.parse_mount_string('C:/foo/bar:/baz')
        assert mount['Source'] == "C:/foo/bar"
        assert mount['Target'] == "/baz"
        assert mount['Type'] == 'bind'


from PIL import Image
import numpy as np
import json
import torch
from torch import nn, optim
from torchvision import datasets, transforms, models
import matplotlib.pyplot as plt
import argparse

def load_checkpoint(filename):

    parser=argparse.ArgumentParser()

    parser.add_argument('-i','--image_dir', help='Path to image', default='flowers/train/1/image_06734.jpg')
    parser.add_argument('-c','--check_dir', help='Directory to checkpoint', default="")
    parser.add_argument('-t','--top_k', help='Choose number of top K', type=int, default='5')
    parser.add_argument('-f','--category_name', help='The mapping full file name for converting categories to real      names', default='cat_to_name.json')
    parser.add_argument('-g','--gpu', help='Use GPU for training gpu / cpu', default='gpu')

    args=parser.parse_args()
    print(args)

    if args.gpu == 'gpu':
        device = 'cuda'
    else:
        device = 'cpu'
    device = torch.device(device)
    checkpoint=torch.load(args.check_dir+filename)

    model=checkpoint['model']
    model.to(device)
    model.load_state_dict (checkpoint['state_dict'])
    optimizer= checkpoint['optimizer']
    epochs=checkpoint['epochs']
    model.to('cpu')

    return model, epochs, optimizer, args

model, epochs, optimizer, args=load_checkpoint('checkpoints.pth')

with open(args.category_name, 'r') as f:
    cat_to_name = json.load(f)

def process_image(image):

    image = Image.open(image)
    image= image.resize((224, 224))
    #box=(224, 224, 224, 224)
    #image= image.crop(box)

    image = np.array(image)

    image = np.divide(image,255)
    mean= np.array([0.485, 0.456, 0.406])
    std= np.array([0.229, 0.224, 0.225])

    image= (image - mean)/ std

    image = image.transpose((2 ,0 ,1))
    image= torch.from_numpy(image)
    model.to('cpu')

    return image

def predict(image_path, model, topk):
    model.eval()
    #model.float()
    # TODO: Implement the code to predict the class from an image file
    image = process_image(image_path)

    image = image.type(torch.FloatTensor)
    image=torch.unsqueeze(image,dim=0)

    #print(image.shape)
    logps = model.forward(image)
    ps= torch.exp(logps)
    probs, classes = ps.topk(topk, dim=1)

    return probs, classes

file_path = args.image_dir

probs, classes = predict(file_path, model, args.top_k)

print(probs)
print(classes)

probs_np=probs.detach().numpy()
max_index = np.argmax(probs_np)
max_probability = probs[max_index]
label = classes[0][max_index]
label_np=label.numpy()
classes_np=classes.numpy()

fig = plt.figure(figsize=(6,6))
ax1 = plt.subplot2grid((15,9), (0,0), colspan=9, rowspan=9)
ax2 = plt.subplot2grid((15,9), (9,2), colspan=5, rowspan=5)
#print (z)
image = Image.open(file_path)
ax1.axis('off')
ax1.set_title(cat_to_name[str(label_np+1)])
ax1.imshow(image)

labels = []
for i in classes_np[0]:
    labels.append(cat_to_name[str(i+1)])

y_pos = np.arange(args.top_k)
ax2.set_yticks(y_pos)
ax2.set_yticklabels(labels)
ax2.set_xlabel('Probability')
ax2.invert_yaxis()
ax2.barh(y_pos, probs_np[0], xerr=0, align='center', color='blue')

plt.show()


import torch
import torch.nn as nn
import torch.nn.functional as F
from .basic_modules import Bottleneck, SpatialGRU, ConvBlock, GRUCell

import pdb
import copy

class DistributionModule(nn.Module):
    """
    A convolutional net that parametrises a diagonal Gaussian distribution.
    """

    def __init__(
            self, in_channels, latent_dim, min_log_sigma, max_log_sigma):
        super().__init__()
        self.compress_dim = in_channels // 2
        self.latent_dim = latent_dim
        self.min_log_sigma = min_log_sigma
        self.max_log_sigma = max_log_sigma

        self.encoder = DistributionEncoder(
            in_channels,
            self.compress_dim,
        )
        self.last_conv = nn.Sequential(
            nn.AdaptiveAvgPool2d(1), nn.Conv2d(
                self.compress_dim, out_channels=2 * self.latent_dim, kernel_size=1)
        )

        self.fp16_enabled = False

    def forward(self, s_t):
        b, s = s_t.shape[:2]
        assert s == 1
        encoding = self.encoder(s_t[:, 0])

        mu_log_sigma = self.last_conv(encoding).view(b, 1, 2 * self.latent_dim)
        mu = mu_log_sigma[:, :, :self.latent_dim]
        log_sigma = mu_log_sigma[:, :, self.latent_dim:]

        # clip the log_sigma value for numerical stability
        log_sigma = torch.clamp(
            log_sigma, self.min_log_sigma, self.max_log_sigma)
        return mu, log_sigma

class SpatialDistributionModule(nn.Module):
    """
    A convolutional net that parametrises a diagonal Gaussian distribution.
    """

    def __init__(
            self, in_channels, latent_dim, min_log_sigma, max_log_sigma):
        super().__init__()
        self.compress_dim = in_channels // 2
        self.latent_dim = latent_dim
        self.min_log_sigma = min_log_sigma
        self.max_log_sigma = max_log_sigma

        self.encoder = DistributionEncoder(
            in_channels,
            self.compress_dim,
        )

        # difference: remove global average pooling
        self.last_conv = nn.Sequential(
            nn.Conv2d(self.compress_dim, out_channels=2 *
                      self.latent_dim, kernel_size=1)
        )
        self.fp16_enabled = False

    def forward(self, s_t):
        b, s = s_t.shape[:2]
        assert s == 1
        encoding = self.encoder(s_t[:, 0])

        # [batch, latent_dim, h, w]
        mu_log_sigma = self.last_conv(encoding)
        mu = mu_log_sigma[:, :self.latent_dim]
        log_sigma = mu_log_sigma[:, self.latent_dim:]

        # clip the log_sigma value for numerical stability
        log_sigma = torch.clamp(
            log_sigma, self.min_log_sigma, self.max_log_sigma)

        return mu, log_sigma

class DistributionEncoder(nn.Module):
    """Encodes s_t or (s_t, y_{t+1}, ..., y_{t+H}).
    """

    def __init__(self, in_channels, out_channels, num_layer=2):
        super().__init__()

        layers = []
        for _ in range(num_layer):
            layers.append(Bottleneck(in_channels=in_channels,
                          out_channels=out_channels, downsample=True))
            in_channels = out_channels

        self.model = nn.Sequential(*layers)
        self.fp16_enabled = False

    def forward(self, s_t):
        return self.model(s_t)

class FuturePrediction(torch.nn.Module):
    def __init__(self, in_channels, latent_dim, n_gru_blocks=3, n_res_layers=3):
        super().__init__()
        self.n_gru_blocks = n_gru_blocks

        # Convolutional recurrent model with z_t as an initial hidden state and inputs the sample
        # from the probabilistic model. The architecture of the model is:
        # [Spatial GRU - [Bottleneck] x n_res_layers] x n_gru_blocks
        self.spatial_grus = []
        self.res_blocks = []

        # 第一个 gru_block 将低维的 random_variable 转为高维的 特征
        for i in range(self.n_gru_blocks):
            gru_in_channels = latent_dim if i == 0 else in_channels
            self.spatial_grus.append(SpatialGRU(gru_in_channels, in_channels))
            self.res_blocks.append(torch.nn.Sequential(*[Bottleneck(in_channels)
                                                         for _ in range(n_res_layers)]))

        self.spatial_grus = torch.nn.ModuleList(self.spatial_grus)
        self.res_blocks = torch.nn.ModuleList(self.res_blocks)

    def forward(self, x, hidden_state):
        # x 为隐分布, hidden_state 为当前时刻的 bev_features
        # 在每个 gru_block 中，x 为 future_states, hidden_state 为当前的特征
        # grucell 实现了 future_states 沿着时序方向的 updates
        # res_blocks 则是在空域上对特征进行进一步的提取

        # x has shape (b, n_future, c_latent_dim, h, w), hidden_state (b, c, h, w)
        for i in range(self.n_gru_blocks):
            x = self.spatial_grus[i](x, hidden_state, flow=None)
            b, n_future, c, h, w = x.shape

            x = self.res_blocks[i](x.view(b * n_future, c, h, w))
            x = x.view(b, n_future, c, h, w)

        return x

class ResFuturePrediction(torch.nn.Module):
    def __init__(self,
                 in_channels,
                 latent_dim,
                 n_future,
                 detach_state=True,
                 n_gru_blocks=3,
                 flow_warp=True,
                 prob_each_future=False,
                 ):
        super().__init__()
        self.n_future = n_future
        self.detach_state = detach_state
        self.flow_warp = flow_warp
        self.prob_each_future = prob_each_future

        # 每个时刻，都以 sample_distribution 和 当前帧的 bev features 作为输入
        # 1. offset prediction: 预测 feature flow ==> warp features
        # 2. gru_cell: reset & update
        # 3. spatial convolution: 输出当前帧特征

        # offset prediction
        if self.flow_warp:
            self.offset_conv = ConvBlock(in_channels=latent_dim + in_channels)
            self.offset_pred = nn.Conv2d(
                latent_dim + in_channels, 2, kernel_size=1, padding=0)

        # gru_cell
        self.gru_cells = nn.ModuleList()
        gru_in_channels = in_channels + latent_dim
        for _ in range(n_gru_blocks):
            self.gru_cells.append(
                GRUCell(input_size=gru_in_channels, hidden_size=in_channels))
            gru_in_channels = in_channels

        # spatial conv
        self.spatial_conv = ConvBlock(in_channels=in_channels)

        def _get_clones(module, N):
            return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

        # init weights
        self.init_weights()

    def init_weights(self):
        if self.flow_warp:
            self.offset_pred.weight.data.normal_(0.0, 0.02)
            self.offset_pred.bias.data.fill_(0)

    def forward(self, sample_distribution, hidden_state):
        # x has shape (b, c_latent_dim, h, w), hidden_state (b, c, h, w)
        res = []
        current_state = hidden_state
        for i in range(self.n_future):

            if self.flow_warp:
                combine = torch.cat(
                    (sample_distribution, current_state), dim=1)
                flow = self.offset_pred(self.offset_conv(combine))
                warp_state = warp_with_flow(current_state, flow=flow)
                warp_state = torch.cat(
                    (warp_state, sample_distribution), dim=1)
            else:
                warp_state = torch.cat(
                    (sample_distribution, current_state), dim=1)

            for gru_cell in self.gru_cells:
                warp_state = gru_cell(warp_state, state=current_state)

            warp_state = self.spatial_conv(warp_state)
            res.append(warp_state)

            # updating current states
            if self.detach_state:
                current_state = warp_state.detach()
            else:
                current_state = warp_state.clone()

        return torch.stack(res, dim=1)

class ResFuturePredictionV2(torch.nn.Module):
    def __init__(self,
                 in_channels,
                 latent_dim,
                 n_future,
                 detach_state=True,
                 n_gru_blocks=3,
                 flow_warp=True,
                 prob_each_future=False,
                 with_state_refine=True,
                 ):
        super().__init__()
        self.n_future = n_future
        self.detach_state = detach_state
        self.flow_warp = flow_warp
        self.prob_each_future = prob_each_future
        self.with_state_refine = with_state_refine

        # 每个时刻，都以 sample_distribution 和 当前帧的 bev features 作为输入
        # 1. offset prediction: 预测 feature flow ==> warp features
        # 2. gru_cell: reset & update
        # 3. spatial convolution: 输出当前帧特征

        # offset prediction
        if self.flow_warp:
            flow_pred = nn.Sequential(
                ConvBlock(in_channels=latent_dim + in_channels),
                nn.Conv2d(latent_dim + in_channels, 2,
                          kernel_size=1, padding=0),
            )

        # gru_cell
        self.gru_cells = nn.ModuleList()
        gru_in_channels = in_channels + latent_dim
        for _ in range(n_gru_blocks):
            self.gru_cells.append(
                GRUCell(input_size=gru_in_channels, hidden_size=in_channels))
            gru_in_channels = in_channels

        # spatial conv
        spatial_conv = ConvBlock(in_channels=in_channels)

        def _get_clones(module, N):
            return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

        if self.with_state_refine:
            if self.flow_warp:
                self.flow_preds = _get_clones(flow_pred, self.n_future)

            self.spatial_convs = _get_clones(spatial_conv, self.n_future)
        else:
            if self.flow_warp:
                self.flow_preds = nn.ModuleList(
                    [flow_pred for _ in range(self.n_future)])

            self.spatial_convs = nn.ModuleList(
                [spatial_conv for _ in range(self.n_future)])

        self.init_weights()

    def init_weights(self):
        if self.flow_warp:
            for flow_pred in self.flow_preds:
                flow_pred[1].weight.data.normal_(0.0, 0.02)
                flow_pred[1].bias.data.fill_(0)

    def forward(self, sample_distribution, hidden_state):
        # x has shape (b, c_latent_dim, h, w), hidden_state (b, c, h, w)
        res = []
        current_state = hidden_state
        for i in range(self.n_future):
            if self.flow_warp:
                combine = torch.cat(
                    (sample_distribution, current_state), dim=1)
                flow = self.flow_preds[i](combine)

                warp_state = warp_with_flow(current_state, flow=flow)
                warp_state = torch.cat(
                    (warp_state, sample_distribution), dim=1)
            else:
                warp_state = torch.cat(
                    (current_state, sample_distribution), dim=1)

            for gru_cell in self.gru_cells:
                warp_state = gru_cell(warp_state, state=current_state)

            warp_state = self.spatial_convs[i](warp_state)
            res.append(warp_state)

            # updating current states
            if self.detach_state:
                current_state = warp_state.detach()
            else:
                current_state = warp_state.clone()

        return torch.stack(res, dim=1)

class ResFuturePredictionV1(torch.nn.Module):
    def __init__(
        self,
        in_channels,
        latent_dim,
        n_future,
        n_gru_blocks=3,
        flow_warp=True,
        detach_state=False,
        prob_each_future=False,
    ):
        super().__init__()
        self.n_future = n_future
        self.detach_state = detach_state
        self.flow_warp = flow_warp
        self.latent_dim = latent_dim
        self.hidden_size = in_channels
        self.n_gru_blocks = n_gru_blocks
        self.prob_each_future = prob_each_future

        # offset prediction
        if self.flow_warp:
            # 输入为上一个时刻的特征 + 采样的分布
            self.flow_pred = nn.Sequential(
                ConvBlock(in_channels=latent_dim + in_channels),
                nn.Conv2d(latent_dim + in_channels, 2,
                          kernel_size=1, padding=0),
            )

        # 为每个时刻都生成一个分布 ?
        self.spatial_grus = nn.ModuleList()
        self.spatial_convs = nn.ModuleList()

        gru_in_channels = in_channels + latent_dim
        for i in range(n_gru_blocks):
            if i == 0:
                self.spatial_grus.append(
                    GRUCell(input_size=gru_in_channels, hidden_size=in_channels))
            else:
                self.spatial_grus.append(
                    SpatialGRU(input_size=gru_in_channels, hidden_size=in_channels))

            gru_in_channels = in_channels
            # spatial conv
            self.spatial_convs.append(
                nn.Sequential(
                    ConvBlock(in_channels=in_channels),
                    nn.Conv2d(in_channels, in_channels,
                              kernel_size=1, padding=0),
                ),
            )

        # init weights
        self.init_weights()

    def init_weights(self):
        self.flow_pred[1].weight.data.normal_(0.0, 0.02)
        self.flow_pred[1].bias.data.fill_(0)

    def forward(self, sample_distribution, hidden_state):
        current_state = hidden_state
        future_states = [current_state]

        if self.prob_each_future:
            future_distributions = torch.split(
                sample_distribution, self.latent_dim, dim=1)
        else:
            future_distributions = [sample_distribution] * self.n_future

        # initialize hidden state
        b, _, h, w = hidden_state.shape
        rnn_state = torch.zeros(b, self.hidden_size, h,
                                w).type_as(hidden_state)
        for i in range(self.n_future):
            if self.flow_warp:
                combine = torch.cat(
                    (future_distributions[i], current_state), dim=1)
                flow = self.flow_pred(combine)
                current_state = warp_with_flow(current_state, flow=flow)

            combine = torch.cat(
                (future_distributions[i], current_state), dim=1)
            rnn_state = self.spatial_grus[0](combine, rnn_state)

            # update
            if self.detach_state:
                current_state = rnn_state.detach()
            else:
                current_state = rnn_state
            future_states.append(rnn_state)

        # [b, t, c, h, w]
        future_states = torch.stack(future_states, dim=1)
        b, t, c, h, w = future_states.shape
        future_states = self.spatial_convs[0](future_states.flatten(0, 1))
        future_states = future_states.view(b, t, c, h, w)

        # further updating
        for k in range(1, self.n_gru_blocks):
            future_states = self.spatial_grus[k](future_states)
            future_states = self.spatial_convs[k](
                future_states.flatten(0, 1)).view(b, t, c, h, w)

        return future_states

def warp_with_flow(x, flow):
    B, C, H, W = x.size()
    # mesh grid
    xx = torch.arange(0, W).view(1, -1).repeat(H, 1)
    yy = torch.arange(0, H).view(-1, 1).repeat(1, W)
    xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
    yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
    grid = torch.cat((xx, yy), dim=1).float()
    flow += grid.type_as(flow)
    flow = flow.permute(0, 2, 3, 1)

    flow[..., 0] = flow[..., 0] / (W - 1) * 2 - 1.0
    flow[..., 1] = flow[..., 1] / (H - 1) * 2 - 1.0
    x = F.grid_sample(x, flow, mode='bilinear', align_corners=True)

    return x


"""Load a layout in Blender."""

import json
from pathlib import Path
from pprint import pformat
from typing import Dict, Optional

import bpy

from openpype.pipeline import (
    discover_loader_plugins,
    remove_container,
    load_container,
    get_representation_path,
    loaders_from_representation,
    AVALON_CONTAINER_ID,
)
from openpype.hosts.blender.api.pipeline import (
    AVALON_INSTANCES,
    AVALON_CONTAINERS,
    AVALON_PROPERTY,
)
from openpype.hosts.blender.api import plugin

class JsonLayoutLoader(plugin.AssetLoader):
    """Load layout published from Unreal."""

    families = ["layout"]
    representations = ["json"]

    label = "Load Layout"
    icon = "code-fork"
    color = "orange"

    animation_creator_name = "CreateAnimation"

    def _remove(self, asset_group):
        objects = list(asset_group.children)

        for obj in objects:
            remove_container(obj.get(AVALON_PROPERTY))

    def _remove_animation_instances(self, asset_group):
        instances = bpy.data.collections.get(AVALON_INSTANCES)
        if instances:
            for obj in list(asset_group.children):
                anim_collection = instances.children.get(
                    obj.name + "_animation")
                if anim_collection:
                    bpy.data.collections.remove(anim_collection)

    def _get_loader(self, loaders, family):
        name = ""
        if family == 'rig':
            name = "BlendRigLoader"
        elif family == 'model':
            name = "BlendModelLoader"

        if name == "":
            return None

        for loader in loaders:
            if loader.__name__ == name:
                return loader

        return None

    def _process(self, libpath, asset, asset_group, actions):
        plugin.deselect_all()

        with open(libpath, "r") as fp:
            data = json.load(fp)

        all_loaders = discover_loader_plugins()

        for element in data:
            reference = element.get('reference')
            family = element.get('family')

            loaders = loaders_from_representation(all_loaders, reference)
            loader = self._get_loader(loaders, family)

            if not loader:
                continue

            instance_name = element.get('instance_name')

            action = None

            if actions:
                action = actions.get(instance_name, None)

            options = {
                'parent': asset_group,
                'transform': element.get('transform'),
                'action': action,
                'create_animation': True if family == 'rig' else False,
                'animation_asset': asset
            }

            if element.get('animation'):
                options['animation_file'] = str(Path(libpath).with_suffix(
                    '')) + "." + element.get('animation')

            # This should return the loaded asset, but the load call will be
            # added to the queue to run in the Blender main thread, so
            # at this time it will not return anything. The assets will be
            # loaded in the next Blender cycle, so we use the options to
            # set the transform, parent and assign the action, if there is one.
            load_container(
                loader,
                reference,
                namespace=instance_name,
                options=options
            )

        # Camera creation when loading a layout is not necessary for now,
        # but the code is worth keeping in case we need it in the future.
        # # Create the camera asset and the camera instance
        # creator_plugin = get_legacy_creator_by_name("CreateCamera")
        # if not creator_plugin:
        #     raise ValueError("Creator plugin \"CreateCamera\" was "
        #                      "not found.")

        # legacy_create(
        #     creator_plugin,
        #     name="camera",
        #     # name=f"{unique_number}_{subset}_animation",
        #     asset=asset,
        #     options={"useSelection": False}
        #     # data={"dependencies": str(context["representation"]["_id"])}
        # )

    def process_asset(self,
                      context: dict,
                      name: str,
                      namespace: Optional[str] = None,
                      options: Optional[Dict] = None):
        """
        Arguments:
            name: Use pre-defined name
            namespace: Use pre-defined namespace
            context: Full parenthood of representation to load
            options: Additional settings dictionary
        """
        libpath = self.filepath_from_context(context)
        asset = context["asset"]["name"]
        subset = context["subset"]["name"]

        asset_name = plugin.asset_name(asset, subset)
        unique_number = plugin.get_unique_number(asset, subset)
        group_name = plugin.asset_name(asset, subset, unique_number)
        namespace = namespace or f"{asset}_{unique_number}"

        avalon_container = bpy.data.collections.get(AVALON_CONTAINERS)
        if not avalon_container:
            avalon_container = bpy.data.collections.new(name=AVALON_CONTAINERS)
            bpy.context.scene.collection.children.link(avalon_container)

        asset_group = bpy.data.objects.new(group_name, object_data=None)
        asset_group.empty_display_type = 'SINGLE_ARROW'
        avalon_container.objects.link(asset_group)

        self._process(libpath, asset, asset_group, None)

        bpy.context.scene.collection.objects.link(asset_group)

        asset_group[AVALON_PROPERTY] = {
            "schema": "openpype:container-2.0",
            "id": AVALON_CONTAINER_ID,
            "name": name,
            "namespace": namespace or '',
            "loader": str(self.__class__.__name__),
            "representation": str(context["representation"]["_id"]),
            "libpath": libpath,
            "asset_name": asset_name,
            "parent": str(context["representation"]["parent"]),
            "family": context["representation"]["context"]["family"],
            "objectName": group_name
        }

        self[:] = asset_group.children
        return asset_group.children

    def exec_update(self, container: Dict, representation: Dict):
        """Update the loaded asset.

        This will remove all objects of the current collection, load the new
        ones and add them to the collection.
        If the objects of the collection are used in another collection they
        will not be removed, only unlinked. Normally this should not be the
        case though.
        """
        object_name = container["objectName"]
        asset_group = bpy.data.objects.get(object_name)
        libpath = Path(get_representation_path(representation))
        extension = libpath.suffix.lower()

        self.log.info(
            "Container: %s\nRepresentation: %s",
            pformat(container, indent=2),
            pformat(representation, indent=2),
        )

        assert asset_group, (
            f"The asset is not loaded: {container['objectName']}"
        )
        assert libpath, (
            "No existing library file found for {container['objectName']}"
        )
        assert libpath.is_file(), (
            f"The file doesn't exist: {libpath}"
        )
        assert extension in plugin.VALID_EXTENSIONS, (
            f"Unsupported file: {libpath}"
        )

        metadata = asset_group.get(AVALON_PROPERTY)
        group_libpath = metadata["libpath"]

        normalized_group_libpath = (
            str(Path(bpy.path.abspath(group_libpath)).resolve())
        )
        normalized_libpath = (
            str(Path(bpy.path.abspath(str(libpath))).resolve())
        )
        self.log.debug(
            "normalized_group_libpath:\n  %s\nnormalized_libpath:\n  %s",
            normalized_group_libpath,
            normalized_libpath,
        )
        if normalized_group_libpath == normalized_libpath:
            self.log.info("Library already loaded, not updating...")
            return

        actions = {}

        for obj in asset_group.children:
            obj_meta = obj.get(AVALON_PROPERTY)
            if obj_meta.get('family') == 'rig':
                rig = None
                for child in obj.children:
                    if child.type == 'ARMATURE':
                        rig = child
                        break
                if not rig:
                    raise Exception("No armature in the rig asset group.")
                if rig.animation_data and rig.animation_data.action:
                    namespace = obj_meta.get('namespace')
                    actions[namespace] = rig.animation_data.action

        mat = asset_group.matrix_basis.copy()

        self._remove_animation_instances(asset_group)

        self._remove(asset_group)

        self._process(str(libpath), asset_group, actions)

        asset_group.matrix_basis = mat

        metadata["libpath"] = str(libpath)
        metadata["representation"] = str(representation["_id"])

    def exec_remove(self, container: Dict) -> bool:
        """Remove an existing container from a Blender scene.

        Arguments:
            container (openpype:container-1.0): Container to remove,
                from `host.ls()`.

        Returns:
            bool: Whether the container was deleted.
        """
        object_name = container["objectName"]
        asset_group = bpy.data.objects.get(object_name)

        if not asset_group:
            return False

        self._remove_animation_instances(asset_group)

        self._remove(asset_group)

        bpy.data.objects.remove(asset_group)

        return True


import gensim
from gensim.models.doc2vec import LabeledSentence
from gensim.models import Doc2Vec
import re
from nltk.tokenize import RegexpTokenizer
from stop_words import get_stop_words
from nltk.stem.porter import PorterStemmer
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
import seaborn
import time

start_time = time.time()
count = 0
lablist = []
for i in range(len(abstract_tot)): #abstract_tot is the list of all the scraped abstracts
    lablist.append(count)
    count+=1
tokenizer  = RegexpTokenizer(r'\w+')
en_stop = get_stop_words('en')
p_stemmer = PorterStemmer()

texts = []
taggeddoc = []
for index, i in enumerate(abstract_tot):
    raw = i.lower()
    tokens = tokenizer.tokenize(raw)
    stopped_tokens = [j for j in tokens if not j in en_stop]
    number_tokens = [re.sub(r'[\d]', ' ', k) for k in stopped_tokens]
    number_tokens = ' '.join(number_tokens).split()
    stemmed_tokens = [p_stemmer.stem(p) for p in number_tokens]
    texts.append(number_tokens)
    #td = TaggedDocument(gensim.utils.to_unicode(str.encode(' '.join(stemmed_tokens))).split(),str(index))
    #taggeddoc.append(td)

LabeledSentence = gensim.models.doc2vec.LabeledSentence
class LabeledLineSentence(object):
    def __init__(self, doc_list, labels_list):
       self.labels_list = labels_list
       self.doc_list = doc_list
    def __iter__(self):
        for idx, doc in enumerate(self.doc_list):
            #yield LabeledSentence(words=doc.split(),tags=[self.labels_list[idx]])
            yield LabeledSentence(words=doc,tags=[self.labels_list[idx]])

#it = LabeledLineSentence(claimcalc_tot, lablist)
it = LabeledLineSentence(texts, lablist)

#model = gensim.models.Doc2Vec(size=300, window=10, min_count=5, workers=11,alpha=0.025, min_alpha=0.025) # use fixed learning rate
model = gensim.models.Doc2Vec(dm = 1, size=100, min_count=10, workers=11,alpha=0.025, min_alpha=0.025) # use fixed learning rate
model.build_vocab(it)
for epoch in range(10):
    print("epoch is: " + str(epoch))
    model.train(it)
    model.alpha -= 0.002 # decrease the learning rate
    model.min_alpha = model.alpha # fix the learning rate
    model.train(it)

end_time = time.time()
print("time taken is" + str(end_time-start_time))

model.save('trained.model')
model.save_word2vec_format('trained.word2vec')

#plot word2vec example
w2v = gensim.models.Doc2Vec.load_word2vec_format('trained.word2vec')
words_np = []
words_label = []
for word in w2v.vocab.keys():
    words_np.append(w2v[word])
    words_label.append(word)
print('Added %s words. Shape %s'%(len(words_np),np.shape(words_np)))

pca = PCA(n_components=2)
pca.fit(words_np)
reduced= pca.transform(words_np)
plt.figure(figsize=(10, 10))

for index,vec in enumerate(reduced):
        # print ('%s %s'%(words_label[index],vec))
        if index <100:
            x,y=vec[0],vec[1]
            plt.scatter(x,y)
            plt.annotate(words_label[index],xy=(x,y), fontsize=14)
plt.axis([0,.2,-.5,.5])
plt.show()


import logging
import sys

# citest modules.
from citest.service_testing import HttpContractBuilder
from citest.service_testing import NoOpOperation
from citest.base import JournalLogger
import citest.base
import citest.gcp_testing as gcp
import citest.json_predicate as jp
import citest.service_testing as st
from citest.json_contract import ObservationPredicateFactory

ov_factory = ObservationPredicateFactory()

# Spinnaker modules.
import spinnaker_testing as sk
import spinnaker_testing.kato as kato

GCP_STANDARD_IMAGES = {
    "centos-cloud": ["centos-7", "centos-6"],
    "coreos-cloud": ["coreos-stable", "coreos-beta", "coreos-alpha"],
    "debian-cloud": ["debian-8"],
    "opensuse-cloud": [""],
    "rhel-cloud": ["rhel-7", "rhel-6"],
    "suse-cloud": [""],
    "ubuntu-os-cloud": [
        "ubuntu-1604-lts",
        "ubuntu-1404-lts",
        "ubuntu-1204-lts",
        "ubuntu-1510",
    ],
    "windows-cloud": ["windows-2012-r2", "windows-2008-r2"],
}

class GoogleKatoTestScenario(sk.SpinnakerTestScenario):
    # _instance_names and _instance_zones will be set in create_instances_.
    # We're breaking them out so that they can be shared by other methods,
    # especially terminate.
    use_instance_names = []
    use_instance_zones = []
    __use_lb_name = ""  # The load balancer name.
    __use_lb_tp_name = ""  # The load balancer's target pool name.
    __use_lb_hc_name = ""  # The load balancer's health check name.
    __use_lb_target = ""  # The load balancer's 'target' resource.
    __use_http_lb_name = ""  # The HTTP load balancer name.
    __use_http_lb_proxy_name = ""  # The HTTP load balancer target proxy name.
    __use_http_lb_hc_name = ""  # The HTTP load balancer health check name.
    __use_http_lb_bs_name = ""  # The HTTP load balancer backend service name.
    __use_http_lb_fr_name = ""  # The HTTP load balancer forwarding rule.
    __use_http_lb_map_name = ""  # The HTTP load balancer url map name.
    __use_http_lb_http_proxy_name = ""  # The HTTP load balancer target http proxy.

    @classmethod
    def new_agent(cls, bindings):
        """Implements the base class interface to create a new agent.

        This method is called by the base classes during setup/initialization.

        Args:
          bindings: The bindings dictionary with configuration information
            that this factory can draw from to initialize. If the factory would
            like additional custom bindings it could add them to initArgumentParser.

        Returns:
          A citest.service_testing.BaseAgent that can interact with Kato.
          This is the agent that test operations will be posted to.
        """
        return kato.new_agent(bindings)

    def create_instances(self):
        """Creates test adding instances to GCE.

         Create three instances.
           * The first two are of different types and zones, which
             we'll check. Future tests will also be using these
             from different zones (but same region).

           * The third is a duplicate in the same zone as another
             so we can check duplicate deletes (which limit one zone per call).

         We'll set the class properties use_instance_names and use_instance_zones
         so that they can be communicated to future tests to reference.

        Returns:
          st.OperationContract
        """
        # We're going to make specific instances so we can refer to them later
        # in tests involving instances. The instances are decorated to trace back
        # to this particular run so as not to conflict with other tests that may
        # be running.
        self.use_instance_names = [
            "katotest%sa" % self.test_id,
            "katotest%sb" % self.test_id,
            "katotest%sc" % self.test_id,
        ]

        # Put the instance in zones. Force one zone to be different
        # to ensure we're testing zone placement. We arent bothering
        # with different regions at this time.
        test_zone = self.bindings["TEST_GCE_ZONE"]
        other_zone = self.bindings["TEST_GCE_REGION"] + (
            "-b" if test_zone[-1] != "b" else "-c"
        )

        self.use_instance_zones = [test_zone, other_zone, test_zone]

        # Give the instances images and machine types. Again we're forcing
        # one to be different to ensure that we're using the values.
        image_name = [
            self.bindings["TEST_GCE_IMAGE_NAME"],
            "debian-9-stretch-v20190116",
            self.bindings["TEST_GCE_IMAGE_NAME"],
        ]
        if image_name[0] == image_name[1]:
            image_name[1] = "ubuntu-minimal-1604-xenial-v20190122"
        machine_type = ["f1-micro", "g1-small", "f1-micro"]

        # The instance_spec will turn into the payload of instances we request.
        instance_spec = []
        builder = gcp.GcpContractBuilder(self.gcp_observer)
        for i in range(3):
            # pylint: disable=bad-continuation
            instance_spec.append(
                {
                    "createGoogleInstanceDescription": {
                        "instanceName": self.use_instance_names[i],
                        "image": image_name[i],
                        "instanceType": machine_type[i],
                        "zone": self.use_instance_zones[i],
                        "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
                    }
                }
            )

            # Verify we created an instance, whether or not it boots.
            (
                builder.new_clause_builder(
                    "Instance %d Created" % i, retryable_for_secs=90
                )
                .aggregated_list_resource("instances")
                .EXPECT(
                    ov_factory.value_list_path_contains(
                        "name", jp.STR_SUBSTR(self.use_instance_names[i])
                    )
                )
            )
            if i < 2:
                # Verify the details are what we asked for.
                # Since we've finished the created clause, this already exists.
                # Note we're only checking the first two since they are different
                # from one another. Anything after that isnt necessary for the test.
                # The clause above already checked that they were all created so we
                # can assume from this test that the details are ok as well.
                (
                    builder.new_clause_builder("Instance %d Details" % i)
                    .inspect_resource(
                        "instances",
                        self.use_instance_names[i],
                        zone=self.use_instance_zones[i],
                    )
                    .contains_path_value("machineType", machine_type[i])
                )
                # Verify the instance eventually boots up.
                # We can combine this with above, but we'll probably need
                # to retry this, but not the above, so this way if the
                # above is broken (wrong), we wont retry thinking it isnt there yet.
                (
                    builder.new_clause_builder(
                        "Instance %d Is Running" % i, retryable_for_secs=360
                    )
                    .inspect_resource(
                        "instances",
                        self.use_instance_names[i],
                        zone=self.use_instance_zones[i],
                    )
                    .EXPECT(
                        ov_factory.value_list_path_contains(
                            "status", jp.STR_EQ("RUNNING")
                        )
                    )
                )

        payload = self.agent.make_json_payload_from_object(instance_spec)

        return st.OperationContract(
            self.new_post_operation(title="create_instances", data=payload, path="ops"),
            contract=builder.build(),
        )

    def terminate_instances(self, names, zone):
        """Creates test for removing specific instances.

        Args:
          names: A list of instance names to be removed.
          zone: The zone containing the instances.

        Returns:
          st.OperationContract
        """
        builder = gcp.GcpContractBuilder(self.gcp_observer)
        clause = builder.new_clause_builder(
            "Instances Deleted", retryable_for_secs=30, strict=True
        ).aggregated_list_resource("instances")
        for name in names:
            # If one of our instances still exists, it should be STOPPING.
            name_matches_pred = jp.PathContainsPredicate("name", name)
            is_stopping_pred = jp.PathEqPredicate("status", "STOPPING")

            # We want the condition to apply to all the observed objects so we'll
            # map the constraint over the observation. Otherwise, if dont map it,
            # then we'd expect the constraint to hold somewhere among the observed
            # objects, but not necessarily all of them.
            clause.AND(
                ov_factory.value_list_matches(
                    [jp.IF(name_matches_pred, is_stopping_pred)], strict=True
                )
            )

        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "terminateInstances",
            {
                "instanceIds": names,
                "zone": zone,
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
            },
        )

        return st.OperationContract(
            self.new_post_operation(
                title="terminate_instances", data=payload, path="gce/ops"
            ),
            contract=builder.build(),
        )

    def upsert_google_server_group_tags(self):
        # pylint: disable=bad-continuation
        server_group_name = "katotest-server-group"
        payload = self.agent.type_to_payload(
            "upsertGoogleServerGroupTagsDescription",
            {
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
                "zone": self.bindings["TEST_GCE_ZONE"],
                "serverGroupName": "katotest-server-group",
                "tags": ["test-tag-1", "test-tag-2"],
            },
        )

        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder("Server Group Tags Added")
            .inspect_resource("instanceGroupManagers", server_group_name)
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(server_group_name),
                            jp.build_path("tags", "items"): jp.LIST_MATCHES(
                                ["test-tag-1", "test-tag-2"]
                            ),
                        }
                    )
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="upsert_server_group_tags", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def create_http_load_balancer(self):
        logical_http_lb_name = "katotest-httplb-" + self.test_id
        self.__use_http_lb_name = logical_http_lb_name

        self.__use_http_lb_hc_name = logical_http_lb_name + "-health-check"
        self.__use_http_lb_bs_name = logical_http_lb_name + "-backend-service"
        self.__use_http_lb_fr_name = logical_http_lb_name
        self.__use_http_lb_map_name = logical_http_lb_name + "-url-map"
        self.__use_http_lb_proxy_name = logical_http_lb_name + "-target-http-proxy"

        interval = 231
        healthy = 8
        unhealthy = 9
        timeout = 65
        path = "/hello/world"
        port_range = "123-456"

        health_check = {
            "checkIntervalSec": interval,
            "healthyThreshold": healthy,
            "unhealthyThreshold": unhealthy,
            "timeoutSec": timeout,
            "requestPath": path,
        }

        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "createGoogleHttpLoadBalancerDescription",
            {
                "healthCheck": health_check,
                "portRange": port_range,
                "loadBalancerName": logical_http_lb_name,
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
            },
        )

        hc_dict = dict(health_check)
        del hc_dict["requestPath"]
        hc_match = {name: jp.NUM_EQ(value) for name, value in health_check.items()}
        hc_match["requestPath"] = jp.STR_EQ(path)
        hc_match["name"] = (jp.STR_SUBSTR(self.__use_http_lb_hc_name),)
        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder("Http Health Check Added")
            .list_resource("httpHealthChecks")
            .EXPECT(ov_factory.value_list_contains(jp.DICT_MATCHES(hc_match)))
        )
        (
            builder.new_clause_builder(
                "Global Forwarding Rule Added", retryable_for_secs=15
            )
            .list_resource("globalForwardingRules")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_http_lb_fr_name),
                            "portRange": jp.STR_EQ(port_range),
                        }
                    )
                )
            )
        )
        (
            builder.new_clause_builder("Backend Service Added")
            .list_resource("backendServices")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_http_lb_bs_name),
                            "healthChecks": jp.STR_SUBSTR(self.__use_http_lb_hc_name),
                        }
                    )
                )
            )
        )
        (
            builder.new_clause_builder("Url Map Added")
            .list_resource("urlMaps")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_http_lb_map_name),
                            "defaultService": jp.STR_SUBSTR(self.__use_http_lb_bs_name),
                        }
                    )
                )
            )
        )
        (
            builder.new_clause_builder("Target Http Proxy Added")
            .list_resource("targetHttpProxies")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_http_lb_proxy_name),
                            "urlMap": jp.STR_SUBSTR(self.__use_http_lb_map_name),
                        }
                    )
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="create_http_load_balancer", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def delete_http_load_balancer(self):
        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "deleteGoogleHttpLoadBalancerDescription",
            {
                "loadBalancerName": self.__use_http_lb_name,
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
            },
        )

        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder("Health Check Removed")
            .list_resource("httpHealthChecks")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_http_lb_hc_name)
                )
            )
        )
        (
            builder.new_clause_builder("Global Forwarding Rules Removed")
            .list_resource("globalForwardingRules")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_http_lb_fr_name)
                )
            )
        )
        (
            builder.new_clause_builder("Backend Service Removed")
            .list_resource("backendServices")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_http_lb_bs_name)
                )
            )
        )
        (
            builder.new_clause_builder("Url Map Removed")
            .list_resource("urlMaps")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_http_lb_map_name)
                )
            )
        )
        (
            builder.new_clause_builder("Target Http Proxy Removed")
            .list_resource("targetHttpProxies")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_http_lb_proxy_name)
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="delete_http_load_balancer", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def upsert_load_balancer(self):
        self.__use_lb_name = "katotest-lb-" + self.test_id
        self.__use_lb_hc_name = "%s-hc" % self.__use_lb_name
        self.__use_lb_tp_name = "%s-tp" % self.__use_lb_name
        self.__use_lb_target = "{}/targetPools/{}".format(
            self.bindings["TEST_GCE_REGION"], self.__use_lb_tp_name
        )

        interval = 123
        healthy = 4
        unhealthy = 5
        timeout = 78
        path = "/" + self.__use_lb_target

        health_check = {
            "checkIntervalSec": interval,
            "healthyThreshold": healthy,
            "unhealthyThreshold": unhealthy,
            "timeoutSec": timeout,
            "requestPath": path,
        }

        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "upsertGoogleLoadBalancerDescription",
            {
                "healthCheck": health_check,
                "region": self.bindings["TEST_GCE_REGION"],
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
                "loadBalancerName": self.__use_lb_name,
            },
        )

        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder("Forwarding Rules Added", retryable_for_secs=30)
            .list_resource("forwardingRules")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_lb_name),
                            "target": jp.STR_SUBSTR(self.__use_lb_target),
                        }
                    )
                )
            )
        )
        (
            builder.new_clause_builder("Target Pool Added", retryable_for_secs=15)
            .list_resource("targetPools")
            .EXPECT(
                ov_factory.value_list_path_contains(
                    "name", jp.STR_SUBSTR(self.__use_lb_tp_name)
                )
            )
        )

        # We list the resources here because the name isnt exact
        # and the list also returns the details we need.
        hc_dict = dict(health_check)
        del hc_dict["requestPath"]

        hc_match = {name: jp.NUM_EQ(value) for name, value in hc_dict.items()}
        hc_match["requestPath"] = jp.STR_EQ(path)
        hc_match["name"] = jp.STR_SUBSTR(self.__use_http_lb_hc_name)
        (
            builder.new_clause_builder("Health Check Added", retryable_for_secs=15)
            .list_resource("httpHealthChecks")
            .EXPECT(ov_factory.value_list_contains(jp.DICT_MATCHES(hc_match)))
        )

        return st.OperationContract(
            self.new_post_operation(
                title="upsert_load_balancer", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def delete_load_balancer(self):
        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "deleteGoogleLoadBalancerDescription",
            {
                "region": self.bindings["TEST_GCE_REGION"],
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
                "loadBalancerName": self.__use_lb_name,
            },
        )

        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder("Health Check Removed")
            .list_resource("httpHealthChecks")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_lb_hc_name)
                )
            )
        )
        (
            builder.new_clause_builder("Target Pool Removed")
            .list_resource("targetPools")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_lb_tp_name)
                )
            )
        )
        (
            builder.new_clause_builder("Forwarding Rule Removed")
            .list_resource("forwardingRules")
            .EXPECT(
                ov_factory.value_list_path_excludes(
                    "name", jp.STR_SUBSTR(self.__use_lb_name)
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="delete_load_balancer", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def register_load_balancer_instances(self):
        """Creates test registering the first two instances with a load balancer.

           Assumes that create_instances test has been run to add
           the instances. Note by design these were in two different zones
           but same region as required by the API this is testing.

           Assumes that upsert_load_balancer has been run to
           create the load balancer itself.
        Returns:
          st.OperationContract
        """
        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "registerInstancesWithGoogleLoadBalancerDescription",
            {
                "loadBalancerNames": [self.__use_lb_name],
                "instanceIds": self.use_instance_names[:2],
                "region": self.bindings["TEST_GCE_REGION"],
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
            },
        )

        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder(
                "Instances in Target Pool", retryable_for_secs=15
            )
            .list_resource("targetPools")
            .EXPECT(
                ov_factory.value_list_contains(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_lb_tp_name),
                            "instances": jp.LIST_MATCHES(
                                [
                                    jp.STR_SUBSTR(self.use_instance_names[0]),
                                    jp.STR_SUBSTR(self.use_instance_names[1]),
                                ]
                            ),
                        }
                    )
                )
            )
            .AND(
                ov_factory.value_list_excludes(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_lb_tp_name),
                            "instances": jp.LIST_MATCHES(
                                [jp.STR_SUBSTR(self.use_instance_names[2])]
                            ),
                        }
                    )
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="register_load_balancer_instances", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def deregister_load_balancer_instances(self):
        """Creates a test unregistering instances from load balancer.

        Returns:
          st.OperationContract
        """
        # pylint: disable=bad-continuation
        payload = self.agent.type_to_payload(
            "deregisterInstancesFromGoogleLoadBalancerDescription",
            {
                "loadBalancerNames": [self.__use_lb_name],
                "instanceIds": self.use_instance_names[:2],
                "region": self.bindings["TEST_GCE_REGION"],
                "credentials": self.bindings["SPINNAKER_GOOGLE_ACCOUNT"],
            },
        )

        # NOTE(ewiseblatt): 20150530
        # This displays an error that 'instances' field doesnt exist.
        # That's because it was removed because all the instances are gone.
        # I dont have a way to express that the field itself is optional,
        # just the record. Leaving it as is because displaying this type of
        # error is usually helpful for development.
        builder = gcp.GcpContractBuilder(self.gcp_observer)
        (
            builder.new_clause_builder(
                "Instances not in Target Pool", retryable_for_secs=30
            )
            .list_resource("targetPools", region=self.bindings["TEST_GCE_REGION"])
            .EXPECT(
                ov_factory.value_list_excludes(
                    jp.DICT_MATCHES(
                        {
                            "name": jp.STR_SUBSTR(self.__use_lb_tp_name),
                            "instances": jp.LIST_MATCHES(
                                [
                                    jp.STR_SUBSTR(self.use_instance_names[0]),
                                    jp.STR_SUBSTR(self.use_instance_names[1]),
                                ]
                            ),
                        }
                    )
                )
            )
        )

        return st.OperationContract(
            self.new_post_operation(
                title="deregister_load_balancer_instances", data=payload, path="ops"
            ),
            contract=builder.build(),
        )

    def list_available_images(self):
        """Creates a test that confirms expected available images.

        Returns:
          st.OperationContract
        """
        logger = logging.getLogger(__name__)

        # Get the list of images available (to the service account we are using).
        context = citest.base.ExecutionContext()
        gcp_agent = self.gcp_observer
        JournalLogger.begin_context("Collecting expected available images")
        relation_context = "ERROR"
        try:
            logger.debug("Looking up available images.")

            json_doc = gcp_agent.list_resource(context, "images")
            for project in GCP_STANDARD_IMAGES.keys():
                logger.info("Looking for images from project=%s", project)
                found = gcp_agent.list_resource(context, "images", project=project)
                for image in found:
                    if not image.get("deprecated", None):
                        json_doc.append(image)

            # Produce the list of images that we expect to receive from spinnaker
            # (visible to the primary service account).
            spinnaker_account = self.bindings["SPINNAKER_GOOGLE_ACCOUNT"]

            logger.debug('Configured with Spinnaker account "%s"', spinnaker_account)
            expect_images = [
                {"account": spinnaker_account, "imageName": image["name"]}
                for image in json_doc
            ]
            expect_images = sorted(expect_images, key=lambda k: k["imageName"])
            relation_context = "VALID"
        finally:
            JournalLogger.end_context(relation=relation_context)

        # pylint: disable=bad-continuation
        builder = HttpContractBuilder(self.agent)
        (
            builder.new_clause_builder("Has Expected Images")
            .get_url_path("/gce/images/find")
            .EXPECT(
                ov_factory.value_list_matches(
                    [jp.DICT_SUBSET(image_entry) for image_entry in expect_images],
                    strict=True,
                    unique=True,
                )
            )
        )

        return st.OperationContract(
            NoOpOperation("List Available Images"), contract=builder.build()
        )

class GoogleKatoIntegrationTest(st.AgentTestCase):
    @property
    def scenario(self):
        return citest.base.TestRunner.global_runner().get_shared_data(
            GoogleKatoTestScenario
        )

    def Xtest_a_upsert_server_group_tags(self):
        self.run_test_case(self.scenario.upsert_google_server_group_tags())

    def test_a_upsert_load_balancer(self):
        self.run_test_case(self.scenario.upsert_load_balancer())

    def test_b_create_instances(self):
        self.run_test_case(self.scenario.create_instances())

    def test_c_register_load_balancer_instances(self):
        self.run_test_case(self.scenario.register_load_balancer_instances())

    def Xtest_d_create_http_load_balancer(self):
        self.run_test_case(self.scenario.create_http_load_balancer())

    def Xtest_v_delete_http_load_balancer(self):
        self.run_test_case(
            self.scenario.delete_http_load_balancer(),
            timeout_ok=True,
            retry_interval_secs=10,
            max_retries=9,
        )

    def test_w_deregister_load_balancer_instances(self):
        self.run_test_case(self.scenario.deregister_load_balancer_instances())

    def test_x_terminate_instances(self):
        # delete 1 which was in a different zone than the other two.
        # Then delete [0,2] together, which were in the same zone.
        try:
            self.run_test_case(
                self.scenario.terminate_instances(
                    [self.scenario.use_instance_names[1]],
                    self.scenario.use_instance_zones[1],
                )
            )
        finally:
            # Always give this a try, even if the first test fails.
            # that increases our chances of cleaning everything up.
            self.run_test_case(
                self.scenario.terminate_instances(
                    [
                        self.scenario.use_instance_names[0],
                        self.scenario.use_instance_names[2],
                    ],
                    self.scenario.use_instance_zones[0],
                )
            )

    def test_z_delete_load_balancer(self):
        # TODO(ewiseblatt): 20151220
        # The retry here is really due to the 400 "not ready" race condition
        # within GCP. Would be better to couple this to the agent and not the
        # test so that it is easier to maintain. Need to add a generalization
        # so the agent can see this is a delete test, got a 400, and only
        # in that condition override the default retry parameters, then stick
        # with the defaults here.
        self.run_test_case(self.scenario.delete_load_balancer(), max_retries=5)

    def Xtest_available_images(self):
        self.run_test_case(self.scenario.list_available_images())

def main():
    return citest.base.TestRunner.main(
        parser_inits=[GoogleKatoTestScenario.initArgumentParser],
        test_case_list=[GoogleKatoIntegrationTest],
    )

if __name__ == "__main__":
    sys.exit(main())