#!/usr/bin/python

import csv
import urllib
import urllib2
import json
import StringIO
from StringIO import StringIO
import argparse
import sys
import time
import multiprocessing
from os import listdir
from os.path import isfile, join

reload(sys)
sys.setdefaultencoding('utf-8')

if sys.version_info < (2, 7):
    def __writeheader(self):
        header = dict(zip(self.fieldnames, self.fieldnames))
        self.writerow(header)
        csv.DictWriter.writeheader = __writeheader

def getwriter(header, fname_out, delimiter, writer_cache, file_cache):
#    global writer_cache
#    global file_cache
    if not fname_out in writer_cache:
        output_file = open(fname_out,'w+')
        writer = csv.DictWriter(output_file,fieldnames=header,delimiter=delimiter)
        if sys.version_info < (2, 7):
            __writeheader(writer)
        else:
            writer.writeheader()
        writer_cache[fname_out]=writer
        file_cache[fname_out]=output_file
    return writer_cache[fname_out]

# Choose method
def getmethod(line):
    if 'api_method' in line:
        return line['api_method']
    return 'all'

def argparser():
    parser = argparse.ArgumentParser(description='Specify input and output filenames')
    parser.add_argument('-i', '--input', dest='input_folder', default=None, required=True, help='Full path to input folder to read from')
    parser.add_argument('-o', '--output', dest='output_folder', default=None, required=False, help='Full path to output folder to write to')
    parser.add_argument('-r', '--rejected', dest='rejected_filename', default=None, required=False, help='Full path to rejected filename')
    parser.add_argument('-u', '--url', dest='url', default='http://invhdp01:8080/sax/v1/media/', required=False, help='REST API')
    parser.add_argument('-v', '--verbose', dest='verbose', default=False, action='store_true', required=False)
    parser.add_argument('-d', '--delimiter', dest='delimiter', default='\t', required=False)
    parser.add_argument('-b', '--bad', dest='bad', default=100, required=False, help='Number of bad rows allowed')
    parser.add_argument('-s', '--sample', dest='sample', default=1000, required=False, help='Size of sample for pre-validation')
    parser.add_argument('-p','--numProcessors', required=False, type=int,
    					default=multiprocessing.cpu_count(),
    					help='Number of processors to use. ' + \
    					"Default for this machine is %d" % (multiprocessing.cpu_count(),) )
    args = parser.parse_args()
    return args

def validation(fname, fname_out):
    c = 0
    br = 0
    flag = True
    writer_cache = dict()
    file_cache = dict()
#    DEBUG
#    print(fname)
#    print(fname_out)
#    print(params.delimiter)
    with open(fname,'rb') as f:
        reader = csv.DictReader(f,delimiter=params.delimiter)
        for line in reader:
#            DEBUG
#            print(line)
            if (not line["urlip"]): # and (not line["ua"]):
                br = br+1
                writer = getwriter(reader.fieldnames,fname_out,params.delimiter,writer_cache,file_cache)
                writer.writerow(dict(line.items()))
            else:
                pass
            c = c+1
            if (br >= params.bad):
                print("File contains many errors.")
                flag = False
                break
            if c > params.sample:
                break
        print("Number of bad rows: "+str(br))
        print("Number of lines validated: "+str(c))
        print("")
        #Close all open files
        for key, value in file_cache.iteritems():
            value.close()
    return flag

def main(fname, fname_out):
    i = 1
    j = 1
    writer_cache = dict()
    file_cache = dict()
    print("Processing file %s" % fname)
    with open(fname,'rb') as f:
        reader = csv.DictReader(f,delimiter=params.delimiter)
        headers = reader.fieldnames
#        DEBUG
#        print(headers)
        for line in reader:
            request = params.url+getmethod(line)+'?'
            #If we don't have ua column, we need to pass empty ua to REST service
            #otherwise it will return incorrect response
            if not 'ua' in headers:
                request = request + 'ua=&'
            request = request  + urllib.urlencode(line)
            if params.verbose:
                print("Request "+str(i)+": "+request)
                i =i+1
            t = 9 # Number of retry attempts
            while t >= 0:
                success = False
                try:
                    response = urllib2.urlopen(request).read()
                    if params.verbose:
                        print("Response "+str(j)+": "+response)
                        j=j+1
                    io = StringIO(response)
                    parsed = json.load(io)

                    wr_output = 'output'
                    writer = getwriter(reader.fieldnames+sorted(parsed.keys()),fname_out,params.delimiter,writer_cache,file_cache)
#                    DEBUG
#                    print ('Headers: ' + str(reader.fieldnames+sorted(parsed.keys())))
#                    print('Line: ' + str(line))
#                    print('')
#                    print('Line item: ' + str(line.items()))
#                    print('')
#                    print('Parsed item: ' + str(parsed.items()))
                    writer.writerow(dict(line.items()+parsed.items()))
                    success = True
                except Exception as e:
                    print("Eroor message from REST API: " + str(e))
                    print("Request: " + request)
                    print("Response: " + response)
                    time.sleep(5) #Sleep for 5 seconds
                t=t-1
                if success:
                    t = -100 #Condition to exit from while
    print("Output file %s has been created." % fname_out)
#    print("Exiting script.")
#Close all open files
    for key, value in file_cache.iteritems():
        value.close()
    return fname_out

params = argparser()
#Set output_folder parameter if not defined
if params.output_folder == None:
    params.output_folder = params.input_folder

#Get only files from input directory
onlyfiles = [f for f in listdir(params.input_folder) if isfile(join(params.input_folder, f)) and not f.startswith('.')]
if params.verbose:
    print('Number of files in input directory: ' + str(len(onlyfiles)))
    print('')
    print('Files: ')
    print(onlyfiles)

#Validation of files one by one from input directory
i=0
files_validated = []
while i < len(onlyfiles):
    if params.verbose == True:
        print('Validating file: ')
        print(onlyfiles[i])
        print('----------------')
    fname = None
    fname = params.input_folder + onlyfiles[i]
    fname_out = None
    fname_out = params.output_folder + onlyfiles[i] + '.rejected'

    v = validation(fname, fname_out)
    if v == True:
        print("Validation of file %s passed!" % (onlyfiles[i]))
        files_validated.append(onlyfiles[i])
#        DEBUG
#        print(files_validated)
    else:
        print("Validation of file %s failed!" % (onlyfiles[i]))
    i=i+1

print("Only next passed validation files will be processed:")
print(files_validated)

#Multiprocessing files
#Start pool
pool = multiprocessing.Pool(processes=params.numProcessors)

print("Number of processes choosen: " + str(params.numProcessors))

#Build task list
tasks = []
i = 0
while i < len(files_validated):
#DEBUG
#    print('Processing file: ')
#    print(files_validated[i])
#    print('----------------')
    fname = None
    fname = params.input_folder + files_validated[i]
    fname_out = None
    fname_out = params.output_folder + files_validated[i] + '.out'

    tasks.append( (fname, fname_out, ) )
    i=i+1

#DEBUG
#print("Tasks:")
#print(tasks)
#print("======")

#Run tasks
results = [pool.apply_async(main, (t[0],t[1])) for t in tasks]
#DEBUG
#print("Results:")
#print(results)
#print("======")
#Process results
#for result in results:
#    processed_fname = result.get()
#    print("Result: output file %s has been created" % (processed_fname) )

pool.close()
pool.join()

#Processing files one by one from input directory
#i=0
#while i < len(files_validated):
#    DEBUG
#    print('==========')
#    if params.verbose == False:
#        print('Processing file: ')
#        print(files_validated[i])
#        print('----------------')
#    fname = None
#    fname = params.input_folder + files_validated[i]
#    fname_out = None
#    fname_out = params.output_folder + files_validated[i] + '.out'
#    main(fname, fname_out)
#    i=i+1

#Close all open files
#for key, value in file_cache.iteritems():
#    value.close()


import cv2

# Loading Image
image = cv2.imread('lena.jpg')
# Converting Image to Blurry Image(always use odd numbers in second () of Gaussian Blur)
image_blur1 = cv2.GaussianBlur(image, (9, 9), 0)
image_blur2 = cv2.GaussianBlur(image, (1, 23), 0)
image_blur3 = cv2.GaussianBlur(image, (23, 1), 0)
# Displaying Original and Blurred Image
cv2.imshow('Original Coloured Image', image)
cv2.imshow('Blurred Image1', image_blur1)
cv2.imshow('Blurred Image2', image_blur2)
cv2.imshow('Blurred Image3', image_blur3)
cv2.waitKey(0)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 23 16:12:28 2020

@author: cuilab
"""

from tridesclous import DataIO
import quantities as pq
import numpy as np
from tqdm import tqdm
from MongoNeo.user_layer.dict_to_neo import templat_neo
from .BaseSorter2Dict import BaseSorter2Dict

def set_spike_segment(dataio_br, Chn, SpikeDict, RecordingParam, dataio_dir):
    #%% align segments spike data
    memmap = dataio_br.get_spikes(chan_grp=Chn,seg_num = 0)
    spike_time = memmap['index'].copy()
    cluster_label = memmap['cluster_label'].copy()
    #print('cluster:',cluster_num)
    cluster_num = list(dataio_br.load_catalogue(chan_grp=Chn)['cluster_labels'])
    processed_signals = dataio_br.arrays[Chn][0].get('processed_signals')

    #%% set spiketrains list
    for clu in cluster_num:
        unit = {'spk':{}}
        unit['spk'] = templat_neo['spk'].copy()

        spike_index = spike_time[cluster_label==clu]
        if len(spike_index)==0:
            continue
        Select = (spike_index<(len(processed_signals)-32)) * (spike_index>16)
        spike_index = spike_index[Select]

        cluster_spike_times = (spike_index/dataio_br.sample_rate)*pq.s

        #QuatoIndex = int(processed_signals.shape[0]/4)
        peak_amplitude = processed_signals[spike_index]
        mad = np.median(np.abs(processed_signals-np.median(processed_signals)))
        snr = np.mean(np.abs(peak_amplitude))/(mad*1.4826)
        spike_description = {'Clu':str(clu),
                             'snr':snr,
                             'Chn':Chn,
                             'mean_waveform':np.array([processed_signals[(index-16):(index+32)]\
                                               for index in spike_index]).squeeze().mean(0)}

        if isinstance(RecordingParam,dict):
            spike_description.update(RecordingParam[Chn])

        if str(spike_description['Chn']) not in SpikeDict:
            SpikeDict[str(spike_description['Chn'])] = {}

        if str(spike_description['Clu']) not in SpikeDict[str(spike_description['Chn'])]:
            SpikeDict[str(spike_description['Chn'])][str(spike_description['Clu'])] = {}

        unit['spk']['times'] = cluster_spike_times
        unit['spk']['t_stop'] = cluster_spike_times[-1]
        unit['spk']['sampling_rate'] = dataio_br.sample_rate*pq.Hz
        unit['spk']['t_start'] = 0*cluster_spike_times.units
        unit['spk']['description'] = spike_description
        SpikeDict[str(spike_description['Chn'])][str(spike_description['Clu'])] = unit

    dataio_br.arrays[Chn][0].detach_array('spikes')
    dataio_br.arrays[Chn][0].detach_array('processed_signals')

class tdcSorter2Dict(BaseSorter2Dict):
    def _ParseSorterResults(self, sorter_dir):
        self._dataio_br = DataIO(dirname=sorter_dir)
        self._sorter_dir = sorter_dir

    def _ExtractMeanWaveform(self):
        pass

    @classmethod
    def Sorter2Dict(cls, sorter_dir, wavedict = {},
                    RecordingParam = None,sampling_rate = 30000*pq.Hz,**waveform_args):
        self = cls()
        self._ParseSorterResults(sorter_dir)

        SpikeDict = {}
        SpikeDict['name'] = 'tdc_spike_train'
        dataio_br = DataIO(dirname=sorter_dir)
        for Chn in tqdm(self._dataio_br.channel_groups.keys()):
            set_spike_segment(dataio_br, Chn, SpikeDict, wavedict, RecordingParam, sorter_dir)

        return SpikeDict


from django.apps import AppConfig
from watson import search as watson

class MembersConfig(AppConfig):
    name = 'members'

    def ready(self):
        team = self.get_model("Team")
        watson.register(team)
        member = self.get_model("Member")
        watson.register(member)


import json
import re

import valideer as V

from coreapis.utils import (
    LogWrapper, get_platform_admins, AlreadyExistsError, ValidationError,
    json_normalize, userinfo_for_log)
from coreapis.cache import Cache
from coreapis import cassandra_client
from coreapis.crud_base import CrudControllerBase
from coreapis.clientadm.controller import ClientAdmController
from coreapis.ldap.status import ldap_status

VALID_SERVICES = ['auth', 'avtale', 'idporten', 'pilot', 'fsgroups', 'nostatus']
VALID_PREFIXES = ['facebook', 'feide', 'github', 'linkedin', 'nin', 'twitter']
VALID_ROLENAMES = ['admin', 'mercantile', 'technical']

def not_empty(thing):
    return len(thing) > 0

def valid_service(service):
    return service in VALID_SERVICES

def valid_feideid(feideid):
    pattern = r'[a-zA-Z0-9.-_]+@[a-zA-Z0-9-]+(\.[a-zA-Z0-9-]+)'
    return re.match(pattern, feideid)

def valid_identity(identity):
    provider_prefix, separator, user_key = identity.partition(':')
    if separator != ':':
        return False
    if provider_prefix == 'feide':
        return valid_feideid(user_key)
    return provider_prefix in VALID_PREFIXES and user_key != ''

def valid_rolenames(rolenames):
    try:
        return all(name in VALID_ROLENAMES for name in rolenames)
    except TypeError:
        return False

class OrgController(CrudControllerBase):
    schema = {
        # Required
        '+id': 'string',
        '+name': V.AllOf(V.Mapping(key_schema=V.String(min_length=2, max_length=3),
                                   value_schema='string'),
                         not_empty),  # Just len works, but who would understand the message?
        # Other attributes
        'fs_groups': '?boolean',
        'realm': '?string',
        'type': V.Nullable(['string']),
        'organization_number': '?string',
        'uiinfo': V.Nullable({}),
        'services': V.Nullable([valid_service]),
        # Virtual attributes - not stored in database
        'has_ldapgroups': '?boolean',
        'has_peoplesearch': '?boolean',
        'peoplesearch': V.Nullable({})
    }
    geo_schema = {
        '+lat': 'number',
        '+lon': 'number',
    }
    platformadmin_attrs = []
    platformadmin_attrs_update = []
    protected_attrs = ['has_ldapgroups', 'has_peoplesearch', 'peoplesearch']
    protected_attrs_update = ['id']

    def __init__(self, settings):
        contact_points = settings.get('cassandra_contact_points')
        keyspace = settings.get('cassandra_keyspace')
        authz = settings.get('cassandra_authz')
        timer = settings.get('timer')
        maxrows = settings.get('orgadmin_maxrows')
        ldap_config = settings.get('ldap_config_file', 'ldap-config.json')
        super(OrgController, self).__init__(maxrows)
        self.timer = timer
        self.log = LogWrapper('org.OrgController')
        self.session = cassandra_client.Client(contact_points, keyspace, authz=authz)
        self.log.debug('org controller init', keyspace=keyspace)
        self.cadm_controller = ClientAdmController(settings)
        self.ldap_config_file = ldap_config
        self.ldap_certs = settings.get('ldap_ca_certs', None)
        platformadmins_file = settings.get('platformadmins_file')
        self.platformadmins = get_platform_admins(platformadmins_file)
        self.ldap_config_cache = Cache(300, 'org.OrgController.ldap_config_cache')

    def _get_ldap_config(self):
        with open(self.ldap_config_file) as configfile:
            return json.load(configfile)

    @property
    def ldap_config(self):
        return self.ldap_config_cache.get('conf', self._get_ldap_config)

    def format_org(self, org):
        has_ldapgroups = False
        has_peoplesearch = False
        psconf = None
        realm = org['realm']
        if realm in self.ldap_config:
            has_ldapgroups = True
            realmconf = self.ldap_config[realm]
            psconf = realmconf.get('peoplesearch')
            if psconf:
                for val in psconf.values():
                    if val != "none":
                        has_peoplesearch = True
                        break
        org['has_ldapgroups'] = has_ldapgroups
        org['has_peoplesearch'] = has_peoplesearch
        org['peoplesearch'] = psconf
        if 'uiinfo' in org and org['uiinfo']:
            org['uiinfo'] = json.loads(org['uiinfo'])
        return org

    def get(self, orgid):
        self.log.debug('Get org', orgid=orgid)
        org = self.session.get_org(orgid)
        return org

    def show_org(self, orgid):
        org = self.format_org(self.get(orgid))
        return org

    def list_orgs(self, want_peoplesearch=None):
        res = []
        for org in self.session.list_orgs():
            org = self.format_org(org)
            if want_peoplesearch is None or want_peoplesearch == org['has_peoplesearch']:
                res.append(org)
        return res

    def add_org(self, user, org):
        org = self.validate(org)
        if self.exists(org['id']):
            raise AlreadyExistsError('item already exists')
        self.log.info('adding organization',
                      audit=True, orgid=org['id'], user=userinfo_for_log(user))
        self.session.insert_org(org)
        return org

    def update_org(self, user, orgid, attrs):
        org = json_normalize(self.show_org(orgid))
        org.update(attrs)
        org = self.validate(org)
        self.log.info('updating organization',
                      audit=True, orgid=orgid, attrs=attrs,
                      user=userinfo_for_log(user))
        self.session.insert_org(org)
        return org

    def delete_org(self, user, orgid):
        if self.list_mandatory_clients(orgid):
            raise ValidationError('org {} has mandatory clients'.format(orgid))
        if self.list_org_roles(orgid):
            raise ValidationError('org {} has roles'.format(orgid))
        self.log.info('delete organization',
                      audit=True, orgid=orgid, user=userinfo_for_log(user))
        self.session.delete_org(orgid)

    def get_logo(self, orgid):
        logo, updated = self.session.get_org_logo(orgid)
        if logo is None or updated is None:
            return None, None
        return logo, updated

    def _save_logo(self, orgid, data, updated):
        self.session.save_org_logo('organizations', orgid, data, updated)

    def validate_geo(self, geo):
        validator = V.parse(self.geo_schema, additional_properties=False)
        validator.validate(geo)
        if not ((-180. <= geo['lon'] <= 180.) and (-90. <= geo['lat'] <= 90.)):
            raise ValidationError('coordinates outside range: {}'.format(geo))

    def update_geo(self, user, orgid, payload, add):
        try:
            for geo in payload:
                self.validate_geo(geo)
        except V.ValidationError:
            raise ValidationError('payload must be an array of coordinates: {}'.format(payload))

        org = self.session.get_org(orgid)
        self.log.info('updating geo coordinates for organization',
                      audit=True, orgid=orgid, payload=payload, add=add,
                      user=userinfo_for_log(user))
        print("uiinfo:", org.get('uiinfo'))
        uiinfo = json.loads(org.get('uiinfo', "{}"))
        if add:
            geos = uiinfo.get('geo', []) + payload
        else:
            geos = payload
        uiinfo['geo'] = geos
        org['uiinfo'] = uiinfo
        self.session.insert_org(org)

    def list_mandatory_clients(self, orgid):
        org = self.session.get_org(orgid)
        clientids = []
        if org.get('realm'):
            clientids = self.session.get_mandatory_clients(org['realm'])
        cadm = self.cadm_controller
        return cadm.get_public_client_list(self.session.get_clients_by_id(clientids).values())

    def add_mandatory_client(self, user, orgid, clientid):
        org = self.session.get_org(orgid)
        realm = org['realm']
        self.log.info('making client mandatory for organization',
                      audit=True, orgid=orgid, clientid=clientid,
                      user=userinfo_for_log(user))
        self.session.add_mandatory_client(realm, clientid)

    def del_mandatory_client(self, user, orgid, clientid):
        org = self.session.get_org(orgid)
        realm = org['realm']
        self.log.info('making client optional for organization',
                      audit=True, orgid=orgid, clientid=clientid,
                      user=userinfo_for_log(user))
        self.session.del_mandatory_client(realm, clientid)

    def list_services(self, orgid):
        org = self.session.get_org(orgid)
        services = org.get('services', [])
        if services is None:
            services = []
        return list(services)

    def add_service(self, user, orgid, service):
        if not valid_service(service):
            raise ValidationError('payload must be a valid service')
        self.log.info('enabling service for organization',
                      audit=True, orgid=orgid, service=service,
                      user=userinfo_for_log(user))
        services = set()
        services.add(service)
        self.session.add_services(orgid, services)

    def del_service(self, user, orgid, service):
        if not valid_service(service):
            raise ValidationError('not a valid service')
        self.log.info('disabling service for organization',
                      audit=True, orgid=orgid, service=service,
                      user=userinfo_for_log(user))
        services = set()
        services.add(service)
        self.session.del_services(orgid, services)

    def list_org_roles(self, orgid):
        roles = self.session.get_roles(['orgid = ?'], [orgid], self.maxrows)
        return [dict(identity=role['identity'], role=role['role']) for role in roles]

    def add_org_role(self, user, orgid, identity, rolenames):
        if not valid_identity(identity):
            raise ValidationError('{} is not a valid identity'.format(identity))
        if not valid_rolenames(rolenames):
            raise ValidationError('{} is not a list of valid role names'.format(rolenames))
        self.log.info('enabling role for organization',
                      audit=True, orgid=orgid, identity=identity, rolenames=rolenames,
                      user=userinfo_for_log(user))
        role = dict(orgid=orgid, identity=identity, role=rolenames)
        self.session.insert_role(role)

    def del_org_role(self, user, orgid, identity):
        if not valid_identity(identity):
            raise ValidationError('not a valid identity')
        self.log.info('disabling role for organization',
                      audit=True, orgid=orgid, identity=identity,
                      user=userinfo_for_log(user))
        self.session.del_role(orgid, identity)

    def has_permission(self, user, org, needs_platform_admin):
        if user is None or not self.is_admin(user, org['id']):
            print("Fail a")
            print(user)
            return False
        if needs_platform_admin and not self.is_platform_admin(user):
            print("Fail b")
            return False
        return True

    # When no name has been given, one of the local admins
    def get_feideid_for_status(self, orgid, realm):
        for identity in (role['identity'] for role in self.list_org_roles(orgid)):
            prefix, _, user_key = identity.partition(':')
            if prefix == 'feide':
                _, _, irealm = user_key.partition('@')
                if irealm == realm:
                    return user_key
        self.log.debug('get_feideid_for_status - no local admin to look up', realm=realm)
        return None

    def ldap_status(self, orgid, feideid):
        org = self.session.get_org(orgid)
        realm = org.get('realm', None)
        if not feideid:
            feideid = self.get_feideid_for_status(orgid, realm)
        return ldap_status(realm, feideid, self.ldap_config, self.ldap_certs)


import torch
from torch import nn

from .model import BiFPN, Regressor, Classifier, EfficientNet
from .utils import Anchors, BBoxTransform, ClipBoxes
from .process import postprocess
from .loss import FocalLoss

class EfficientDetWithLoss(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.model = EfficientDetBackbone(**kwargs)
        self.criterion = FocalLoss()
        self.threshold = kwargs.get("threshold", 0.2)
        self.iou_threshold = kwargs.get("iou_threshold", 0.2)

    def forward(self, imgs, annotations):
        if self.training:
            features, regression, classification, anchors = self.model(imgs)
            cls_loss, reg_loss = self.criterion(classification, regression, anchors, annotations)
        else:
            with torch.no_grad():
                features, regression, classification, anchors = self.model(imgs)
                cls_loss, reg_loss = self.criterion(classification, regression, anchors, annotations)
        losses = {"cls_loss": cls_loss, "reg_loss": reg_loss}
        return losses

    def infer(self, imgs):
        with torch.no_grad():
            features, regression, classification, anchors = self.model(imgs)
        regressBoxes = BBoxTransform()
        clipBoxes = ClipBoxes()

        out = postprocess(imgs,
                          anchors, regression, classification,
                          regressBoxes, clipBoxes,
                          self.threshold, self.iou_threshold)
        return out

class EfficientDetBackbone(nn.Module):
    def __init__(self, num_classes=80, compound_coef=0, load_weights=False, **kwargs):
        super().__init__()
        self.compound_coef = compound_coef

        self.backbone_compound_coef = [0, 1, 2, 3, 4, 5, 6, 6, 7]
        self.fpn_num_filters = [64, 88, 112, 160, 224, 288, 384, 384, 384]
        self.fpn_cell_repeats = [3, 4, 5, 6, 7, 7, 8, 8, 8]
        self.input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536]
        self.box_class_repeats = [3, 3, 3, 4, 4, 4, 5, 5, 5]
        self.pyramid_levels = [5, 5, 5, 5, 5, 5, 5, 5, 6]
        self.anchor_scale = [4., 4., 4., 4., 4., 4., 4., 5., 4.]
        self.aspect_ratios = kwargs.get('ratios', [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)])
        self.num_scales = len(kwargs.get('scales', [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]))
        conv_channel_coef = {
            # the channels of P3/P4/P5.
            0: [40, 112, 320],
            1: [40, 112, 320],
            2: [48, 120, 352],
            3: [48, 136, 384],
            4: [56, 160, 448],
            5: [64, 176, 512],
            6: [72, 200, 576],
            7: [72, 200, 576],
            8: [80, 224, 640],
        }

        num_anchors = len(self.aspect_ratios) * self.num_scales

        self.bifpn = nn.Sequential(
            *[BiFPN(self.fpn_num_filters[self.compound_coef],
                    conv_channel_coef[compound_coef],
                    True if _ == 0 else False,
                    attention=True if compound_coef < 6 else False,
                    use_p8=compound_coef > 7)
              for _ in range(self.fpn_cell_repeats[compound_coef])])

        self.num_classes = num_classes
        self.regressor = Regressor(in_channels=self.fpn_num_filters[self.compound_coef], num_anchors=num_anchors,
                                   num_layers=self.box_class_repeats[self.compound_coef],
                                   pyramid_levels=self.pyramid_levels[self.compound_coef])
        self.classifier = Classifier(in_channels=self.fpn_num_filters[self.compound_coef], num_anchors=num_anchors,
                                     num_classes=num_classes,
                                     num_layers=self.box_class_repeats[self.compound_coef],
                                     pyramid_levels=self.pyramid_levels[self.compound_coef])

        self.anchors = Anchors(anchor_scale=self.anchor_scale[compound_coef],
                               pyramid_levels=(torch.arange(self.pyramid_levels[self.compound_coef]) + 3).tolist(),
                               **kwargs)

        self.backbone_net = EfficientNet(self.backbone_compound_coef[compound_coef], load_weights)

    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()

    def forward(self, inputs):
        max_size = inputs.shape[-1]

        _, p3, p4, p5 = self.backbone_net(inputs)

        features = (p3, p4, p5)
        features = self.bifpn(features)

        regression = self.regressor(features)
        classification = self.classifier(features)
        anchors = self.anchors(inputs, inputs.dtype)

        return features, regression, classification, anchors

    def init_backbone(self, path):
        state_dict = torch.load(path)
        try:
            ret = self.load_state_dict(state_dict, strict=False)
            print(ret)
        except RuntimeError as e:
            print('Ignoring ' + str(e) + '"')


import numpy as np
from scipy import stats
from math import sqrt
import multiprocessing as mp
from itertools import repeat
from sklearn.mixture import GaussianMixture, BayesianGaussianMixture
from sklearn.cluster import SpectralClustering
from sklearn.metrics import adjusted_rand_score
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt
from low_rank_rnns.modules import SupportLowRankRNN, LowRankRNN, FullRankRNN
from low_rank_rnns.helpers import gram_factorization
from low_rank_rnns.plotting import center_axes
import torch

#%% GAUSSIAN MIXTURES

def gmm_fit(neurons_fs, n_components, algo='bayes', n_init=50, random_state=None, mean_precision_prior=None,
            weight_concentration_prior_type='dirichlet_process', weight_concentration_prior=None):
    """
    Fit a mixture of gaussians to a cloud of n points
    :param neurons_fs: list of numpy arrays of shape n or numpy array of shape n x d
    :param n_components: int
    :param algo: 'em' (expectation-maximization) or 'bayes'
    :param n_init: number of random seeds for the inference algorithm
    :param random_state: random seed for the rng to eliminate randomness
    :return: vector of population labels (of shape n), best fitted model
    """
    if isinstance(neurons_fs, list):
        X = np.vstack(neurons_fs).transpose()
    else:
        X = neurons_fs
    if algo == "em":
        model = GaussianMixture(n_components=n_components, n_init=n_init, random_state=random_state)
    else:
        model = BayesianGaussianMixture(n_components=n_components, n_init=n_init, random_state=random_state,
                                        init_params='random', mean_precision_prior=mean_precision_prior,
                                        weight_concentration_prior_type=weight_concentration_prior_type,
                                        weight_concentration_prior=weight_concentration_prior)
    model.fit(X)
    z = model.predict(X)
    return z, model

def make_vecs(net):
    """
    return a list of vectors (list of numpy arrays of shape n) composing a network
    """
    return [net.m[:, i].detach().cpu().numpy() for i in range(net.rank)] + \
           [net.n[:, i].detach().cpu().numpy() for i in range(net.rank)] + \
           [net.wi[i].detach().cpu().numpy() for i in range(net.input_size)] + \
           [net.wo[:, i].cpu().detach().numpy() for i in range(net.output_size)]

def to_support_net(net, z, new_size=None, take_means=False):
    """
    Generate a SupportLowRankRNN from a LowRankRNN, with populations defined by z
    :param net: LowRankRNN
    :param z: numpy array of size n whose entries are integer labels for populations
    :param new_size: int
    :param take_means: bool, whether to take into account means of components or set them to zero
    :return: the SupportLowRankRNN
    """
    X = np.vstack(make_vecs(net)).transpose()
    _, counts = np.unique(z, return_counts=True)
    n_components = counts.shape[0]
    weights = counts / net.hidden_size
    if take_means:
        means = np.vstack([X[z == i].mean(axis=0) for i in range(n_components)])
    else:
        means = np.zeros((n_components, X.shape[1]))
    covariances = [np.cov(X[z == i].transpose()) for i in range(n_components)]

    rank = net.rank
    basis_dim = 2 * rank + net.input_size + net.output_size
    m_init = torch.zeros(rank, n_components, basis_dim)
    n_init = torch.zeros(rank, n_components, basis_dim)
    wi_init = torch.zeros(net.input_size, n_components, basis_dim)
    wo_init = torch.zeros(net.output_size, n_components, basis_dim)

    if new_size is None:
        new_size = net.hidden_size
    m_means = torch.from_numpy(means[:, :rank]).t() / sqrt(new_size)
    n_means = torch.from_numpy(means[:, rank: 2*rank]).t() / sqrt(new_size)
    wi_means = torch.from_numpy(means[:, 2*rank: 2*rank + net.input_size]).t()

    for i in range(n_components):
        # Compute Gramian matrix of the basis we have to build
        G = covariances[i]
        # compute coefficients on the basis of iid normal vectors
        X_reduced = gram_factorization(G)
        for k in range(rank):
            m_init[k, i] = torch.from_numpy(X_reduced[k]) / sqrt(new_size)
            n_init[k, i] = torch.from_numpy(X_reduced[rank + k]) / sqrt(new_size)
        for k in range(net.input_size):
            wi_init[k, i] = torch.from_numpy(X_reduced[2 * rank + k])
        for k in range(net.output_size):
            wo_init[k, i] = torch.from_numpy(X_reduced[2 * rank + net.input_size + k]) / new_size

    net2 = SupportLowRankRNN(net.input_size, new_size, net.output_size, net.noise_std, net.alpha, rank, n_components,
                             weights, basis_dim, m_init, n_init, wi_init, wo_init, m_means, n_means, wi_means)
    return net2

def inactivate_pop(net, pop):
    size = net.hidden_size - np.sum(pop)
    wi = net.wi_full[:, ~pop]
    wo = net.wo_full[~pop]

    if isinstance(net, LowRankRNN):
        m = net.m[~pop]
        n = net.n[~pop]
        net2 = LowRankRNN(net.input_size, size, net.output_size, net.noise_std, net.alpha, net.rank,
                          wi_init=wi, m_init=m, n_init=n, wo_init=wo)
    elif isinstance(net, FullRankRNN):
        mask = np.ix_(~pop, ~pop)
        wrec = net.wrec[mask]
        net2 = FullRankRNN(net.input_size, size, net.output_size, net.noise_std, net.alpha, net.rho,
                           wi_init=wi, wrec_init=wrec, wo_init=wo)
    else:
        raise NotImplementedError
    return net2

#%%  PLOTTING

def pop_scatter_linreg(vec1, vec2, pops, n_pops=None, colors=('blue', 'red', 'green', 'gray'),
                       linreg=True, figsize=(5, 5), ax=None, dotsize=4, ls=None, lw=None):
    """
    scatter plot of (vec1, vec2) points separated in populations according to int labels in vector pops, with linear
    regressions
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    center_axes(ax)

    # Computing axes limits
    xmax = max(abs(vec1.min()), vec1.max())
    xmin = -xmax
    ax.set_xlim(xmin - .1 * (xmax - xmin), xmax + .1 * (xmax - xmin))
    ymax = max(abs(vec2.min()), vec2.min())
    ymin = -ymax
    ax.set_ylim(ymin - .1 * (ymax - ymin), ymax + .1 * (ymax - ymin))
    xs = np.linspace(xmin - .1 * (xmax - xmin), xmax + .1 * (xmax - xmin), 100)

    if n_pops is None:
        n_pops = np.unique(pops).shape[0]
    for i in range(n_pops):
        ax.scatter(vec1[pops == i], vec2[pops == i], color=colors[i], s=dotsize)
        if linreg:
            slope, intercept, r_value, p_value, std_err = stats.linregress(vec1[pops == i], vec2[pops == i])
            print(f"pop {i}: slope={slope:.2f}, intercept={intercept:.2f}")
            ax.plot(xs, slope * xs + intercept, color=colors[i], zorder=-1, ls=ls, lw=lw)
    ax.set_xticks([])
    ax.set_yticks([])

#%% SPECTRAL CLUSTERING AND STABILITY ANALYSIS

def generate_subsamples(neurons_fs, fraction=.8):
    indexes = np.random.choice(neurons_fs.shape[0], int(fraction * neurons_fs.shape[0]), replace=False)
    indexes = np.sort(indexes)
    return neurons_fs[indexes], indexes

def spectral_clustering(neurons_fs, n_clusters, metric='euclidean', n_neighbors=10):
    if metric == 'euclidean':
        model = SpectralClustering(n_clusters, affinity='nearest_neighbors')
        model.fit(neurons_fs)
    elif metric == 'cosine':
        model = SpectralClustering(n_clusters, affinity='precomputed')
        nn = NearestNeighbors(n_neighbors=n_neighbors, algorithm='brute', metric='cosine')
        nn.fit(neurons_fs)
        knn_graph = nn.kneighbors_graph()
        knn_graph = 0.5 * (knn_graph + knn_graph.transpose())
        model.fit(knn_graph)
    return model

def clustering_stability_task(neurons_fs, algo, n_clusters, metric, n_neighbors, mean_precision_prior=1e5):
    sample, indexes = generate_subsamples(neurons_fs)
    if algo == 'spectral':
        model = spectral_clustering(sample, n_clusters, metric, n_neighbors)
        labels = model.labels_
    else:
        labels, _ = gmm_fit(sample, n_clusters, mean_precision_prior=mean_precision_prior,
                            weight_concentration_prior_type='dirichlet_distribution')
    return labels, indexes

def clustering_stability(neurons_fs, n_clusters, n_bootstrap, algo='gmm', metric='cosine', n_neighbors=10,
                         mean_precision_prior=1e5, normalize=None):
    """
    Compute clustering stability distribution for a particular number of clusters and algorithm
    :param neurons_fs: numpy array of shape Nxd (neurons embedded in some feature space)
    :param n_clusters: int
    :param n_bootstrap: int
    :param algo: 'spectral' or 'gmm'
    :param metric: 'euclidean' or 'cosine' for spectral clustering
    :param n_neighbors: int, for spectral clustering
    :param mean_precision_prior:
    :param normalize: None, 'normal' or 'uniform'
    :return: list of n_bootstrap x (n_bootstrap - 1) / 2 ARI values for bootstrapped clusterings (possibly normalized)
    """

    with mp.Pool(mp.cpu_count()) as pool:
        args = repeat((neurons_fs, algo, n_clusters, metric, n_neighbors, mean_precision_prior), n_bootstrap)
        res = pool.starmap(clustering_stability_task, args)
        labels_list, indexings = zip(*res)

    aris = []
    # Align bootstrap samples and compute pairwise Rand indexes
    for i in range(n_bootstrap):
        for j in range(i+1, n_bootstrap):
            # build aligned labellings
            indexes_i = indexings[i]
            indexes_j = indexings[j]
            labels_i = []
            labels_j = []
            l = 0
            for k in range(len(indexes_i)):
                if l > len(indexes_j):
                    break
                while l < len(indexes_j) and indexes_j[l] < indexes_i[k]:
                    l += 1
                if l < len(indexes_j) and indexes_j[l] == indexes_i[k]:
                    labels_i.append(labels_list[i][k])
                    labels_j.append(labels_list[j][l])
                    l += 1
            aris.append(adjusted_rand_score(labels_i, labels_j))

    if normalize is not None:
        if normalize == 'normal':
            X_base = np.random.randn(neurons_fs.shape[0], neurons_fs.shape[1])
        elif normalize == 'uniform':
            X_base = (np.random.rand(neurons_fs.shape[0], neurons_fs.shape[1]) - 0.5) * 2
        base_aris = clustering_stability(X_base, n_clusters, n_bootstrap, metric, normalize=None)
        base_mean, base_std = np.mean(base_aris), np.std(base_aris)
        aris = [(ari - base_mean) / base_std for ari in aris]
    return aris

def boxplot_clustering_stability(neurons_fs, clusters_nums, aris=None, algo='gmm', n_bootstrap=20, metric='cosine',
                                 n_neighbors=10, ax=None):
    """
    :param neurons_fs: numpy array of shape Nxd (neurons embedded in some feature space)
    :param clusters_nums: list of ints
    :param aris: if precomputed aris, list of lists with ari distributions
    :param algo: 'spectral' or 'gmm'
    :param n_bootstrap: int
    :param metric: 'euclidean' or 'cosine' for spectral clustering
    :param n_neighbors: int, for spectral clustering
    :param ax: matplotlib axes if already setup
    :return: axes
    """
    if aris is None:
        aris = [clustering_stability(neurons_fs, k, n_bootstrap, algo, metric, n_neighbors) for k in clusters_nums]
    aris = np.array(aris)
    if ax is None:
        fig, ax = plt.subplots()
    col_lines = 'indianred'
    bp = ax.boxplot(aris.T, patch_artist=True)
    for box in bp['boxes']:
        box.set(color='steelblue', facecolor='steelblue')
    for med in bp['medians']:
        med.set(color=col_lines)
    for l in bp['whiskers']:
        l.set(color=col_lines)
    for l in bp['caps']:
        l.set(color=col_lines)
    ax.set_xticks(list(range(1, aris.shape[0] + 1)))
    ax.set_xticklabels(clusters_nums)
    ax.set(xlabel='number of clusters', ylabel='stability', ylim=(-.1, 1.1))
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    return ax


from odoo import api, models

class ProductChangeQuantity(models.TransientModel):

    _inherit = "stock.change.product.qty"

    @api.model
    def default_get(self, fields):
        res = super(ProductChangeQuantity, self).default_get(fields)

        if self.env.user.company_id.stock_update_qty_location_default:
            res['location_id'] = self.env.user.company_id \
                .stock_update_qty_location_default.id

        return res


import numpy as np
import matplotlib.pyplot as plt

Fd = ['1.350', '1.440', '1.465', '1.770']
p = []
for f in Fd:
    p.append(np.loadtxt("data_{}000.dat".format(f)))
p_delta = []
for f in Fd:
    p_delta.append(np.loadtxt("data_{}000_delta.dat".format(f)))

#%%

fig1 = plt.figure(figsize=(15,15))
for i in range(4):
    ax = fig1.add_subplot(2,2,i+1)
    ax.plot(p[i][:,0],p[i][:,1], label="Fd={}".format(Fd[i]))
    ax.set_xlabel("$time$")
    ax.set_ylabel(u"$\theta$")
    ax.legend()
plt.savefig("theta_vs_time.png")

#%%

fig2 = plt.figure(figsize=(15,15))
for i in range(4):
    ax = fig2.add_subplot(2,2,i+1)
    ax.scatter(p[i][:,1],p[i][:,2], s=5, label="Fd={}".format(Fd[i]))
    ax.set_xlabel(u"$\theta$")
    ax.set_ylabel("$\omega$")
    ax.legend()
plt.savefig("omega_vs_theta.png")

#%%

fig3 = plt.figure(figsize=(15,15))
for i in range(4):
    ax = fig3.add_subplot(2,2,i+1)
    ax.semilogy(p[i][:,0],np.abs(p_delta[i][:,1] - p[i][:,1]),marker="+", label="Fd={}".format(Fd[i]))
    ax.set_xlabel(u"$time$")
    ax.set_ylabel("$\Delta\Theta$")
    ax.legend()
plt.savefig("lyapunov.png")
#%%

bifurcacion = np.loadtxt("bifurcacion.dat")

#%%

fig4 = plt.figure(figsize=(15,15))
ax = fig4.add_subplot(111)

ax.scatter(bifurcacion[:,0], bifurcacion[:,1], marker=".", s=1)
ax.set_title("Diagrama de Bifurcación")
ax.set_ylim(1,3)

plt.savefig("bifurcacion.png")


import open3d as o3d
import open3d.visualization.gui as gui
import numpy as np
import time
import sys
import json
import pandas as pd
sys.path.append("../")
from testingStageViz import jsonToDf
from positionAnalysis import jointsNames
from realTimeHandRotationCompute import jointsNames as handJointsNames

# 全身骨架的結構
fullBodyBoneStrcuture = [
    [jointsNames.Hip, jointsNames.LeftUpperLeg], [jointsNames.Hip, jointsNames.RightUpperLeg],
    [jointsNames.LeftUpperLeg, jointsNames.LeftLowerLeg], [jointsNames.RightUpperLeg, jointsNames.RightLowerLeg],
    [jointsNames.LeftLowerLeg, jointsNames.LeftFoot], [jointsNames.RightLowerLeg, jointsNames.RightFoot],
    [jointsNames.Hip, jointsNames.Spine], [jointsNames.Spine, jointsNames.Chest], [jointsNames.Chest, jointsNames.UpperChest],
    [jointsNames.UpperChest, jointsNames.LeftUpperArm], [jointsNames.UpperChest, jointsNames.RightUpperArm], [jointsNames.UpperChest, jointsNames.Head],
    [jointsNames.LeftUpperArm, jointsNames.LeftLowerArm], [jointsNames.RightUpperArm, jointsNames.RightLowerArm],
    [jointsNames.LeftLowerArm, jointsNames.LeftHand], [jointsNames.RightLowerArm, jointsNames.RightHand],
]

# 手部骨架結構
handBoneStructure = [
    [handJointsNames.wrist, handJointsNames.thunmbCMC], [handJointsNames.wrist, handJointsNames.indexMCP],
    [handJointsNames.wrist, handJointsNames.middleMCP], [handJointsNames.wrist, handJointsNames.ringMCP],
    [handJointsNames.wrist, handJointsNames.pinkyMCP],
    [handJointsNames.thunmbCMC, handJointsNames.thunmbMCP], [handJointsNames.thunmbMCP, handJointsNames.thunmbIP],
    [handJointsNames.thunmbIP, handJointsNames.thunmbTIP],
    [handJointsNames.indexMCP, handJointsNames.indexPIP], [handJointsNames.indexPIP, handJointsNames.indexDIP],
    [handJointsNames.indexDIP, handJointsNames.indexTIP],
    [handJointsNames.middleMCP, handJointsNames.middlePIP], [handJointsNames.middlePIP, handJointsNames.middleDIP],
    [handJointsNames.middleDIP, handJointsNames.middleTIP],
    [handJointsNames.ringMCP, handJointsNames.ringPIP], [handJointsNames.ringPIP, handJointsNames.ringDIP],
    [handJointsNames.ringDIP, handJointsNames.ringTIP],
    [handJointsNames.pinkyMCP, handJointsNames.pinkyPIP], [handJointsNames.pinkyPIP, handJointsNames.pinkyDIP],
    [handJointsNames.pinkyDIP, handJointsNames.pinkyTIP],
]

# 目前顯示的frame index
curTimeInd = 0

def readAppliedRotPos(rotAppliedFilePath='../positionData/leftFrontKick_quat_directMapping.json', TPoseFilePath = ''):
    # 可以參考testingStageViz.py
    with open(rotAppliedFilePath, 'r') as WFile:
        lowerBodyPosition=json.load(WFile)
    timeCount = len(lowerBodyPosition)
    jointsIndices = list(lowerBodyPosition[0]['data'].keys())
    lowerBodyPosition = [{'data': {int(i): lowerBodyPosition[t]['data'][i] for i in jointsIndices}} for t in range(timeCount)]
    lowerBodyDf = jsonToDf(lowerBodyPosition)
    concatDfs = pd.concat(lowerBodyDf.values(), axis=1)
    concatArr = concatDfs.values
    concatArr = concatArr.reshape((concatArr.shape[0], -1, 3))
    print('lower body joint indices:', jointsIndices)
    print('timeCount: ', timeCount)
    print('lower body array shape', concatArr.shape)
    # Hip校正為原點
    concatArr = concatArr - concatArr[:, -1:, :]

    # read T pose
    # Ref: realTimeRotToAvatarPos.py
    saveDirPath = '../positionData/fromDB/genericAvatar/'
    TPoseJson = None
    with open(saveDirPath+'TPose.json', 'r') as fileIn:
        TPoseJson = json.load(fileIn)['results']
    jointCount = len(TPoseJson[0]['data'])
    TPoseDfs = jsonToDf(TPoseJson)
    TPoseConcatDfs = pd.concat(TPoseDfs.values(), axis=1)
    TPoseArr = TPoseConcatDfs.values
    TPoseArr = TPoseArr.reshape((TPoseArr.shape[0], -1, 3))
    TPoseArr = TPoseArr[3:4, :, :]  # Just dont use the first frame, since it might be wrong
    TPoseArr = np.repeat(TPoseArr, concatArr.shape[0], axis=0)
    print('T Pose joint count: ', jointCount)
    print('T Pose time count: ', TPoseArr.shape[0])
    print('T Pose array shape: ', TPoseArr.shape)
    # Hip校正回原點
    TPoseArr = TPoseArr - TPoseArr[:, 6:7, :]

    # combine T pose to lower body motions
    retArr = np.concatenate((concatArr, TPoseArr[:, 7:, :]), axis=1)
    print('return array shape: ', retArr.shape)

    return retArr

def readSynthesisPos(synthesisFilePath = '../positionData/afterSynthesis/NoVelAccOverlap/leftFrontKick_quat_direct_075_EWMA.json'):
    # TODO
    # read synthesis result
    # TODO: 好像會少一個點, 應該就是少hip, 應為預設他就是0, 需要把hip是0補回來
    synthesisPos = None
    with open(synthesisFilePath, 'r') as WFile:
        synthesisPos=json.load(WFile)
    timeCount = len(synthesisPos)
    jointCount = len(synthesisPos[0]['data'])
    synthesisDfs = jsonToDf(synthesisPos)
    synthesisConcatDfs = pd.concat(synthesisDfs.values(), axis=1)
    synthesisArr = synthesisConcatDfs.values
    synthesisArr = synthesisArr.reshape((synthesisArr.shape[0], -1, 3))
    print('time count: ', timeCount)
    print('joint count: ', jointCount)

    # 補Hip是0回array當中
    _hip = np.zeros((synthesisArr.shape[0], 3))
    synthesisArr = np.insert(synthesisArr, 6, _hip, axis=1)
    print('output array shape: ', synthesisArr.shape)
    return synthesisArr

def readHandPos(
    handLandmarkFilePath = '../complexModel/newRecord/twoLegJump_rgb.json',
    scale = [1, 1, 1], negate = [False, False, False]
):
    handPos = None
    with open(handLandmarkFilePath, 'r') as RFile:
        handPos = json.load(RFile)
    timeCount = len(handPos)
    jointCount = len(handPos[0]['data'])
    handDfs = jsonToDf(handPos)
    handConcatDfs = pd.concat(handDfs.values(), axis=1)
    handArr = handConcatDfs.values
    handArr = handArr.reshape((handArr.shape[0], -1, 3))
    print('time count: ', timeCount)
    print('joint count: ', jointCount)
    print('output array shape: ', handArr.shape)

    # 轉換成wrist在原點
    handArr = handArr - handArr[:, handJointsNames.wrist:handJointsNames.wrist+1, :]
    # 放大數值範圍
    for _axis, _scale in enumerate(scale):
        handArr[:, :, _axis] = handArr[:, :, _axis] * _scale
    # 反轉坐標軸
    for _axis, _neg in enumerate(negate):
        if _neg:
            handArr[:, :, _axis] = handArr[:, :, _axis] * -1
    return handArr

def readExampleAnimPos(filePath = ''):
    print('read example animation')
    # read example animation positions
    synthesisPos = None
    with open(filePath, 'r') as WFile:
        synthesisPos=json.load(WFile)['results']
    timeCount = len(synthesisPos)
    jointCount = len(synthesisPos[0]['data'])
    synthesisDfs = jsonToDf(synthesisPos)
    synthesisConcatDfs = pd.concat(synthesisDfs.values(), axis=1)
    synthesisArr = synthesisConcatDfs.values
    synthesisArr = synthesisArr.reshape((synthesisArr.shape[0], -1, 3))

    # Hip校正回原點
    synthesisArr = synthesisArr - synthesisArr[:, 6:7, :]
    print('time count: ', timeCount)
    print('joint count: ', jointCount)
    print('output array shape: ', synthesisArr.shape)
    return synthesisArr

def vizMotions(jointData, jointHeirarchy, hipPos, axisJointInd, frameRate=0.05, return_imgs=False, cameraParam=None):
    '''
    Objective
        輸入joint data, 使用open 3d顯示骨架運動資訊
    Input
        - jointData: list of 3 dimention array. 多組關節點的3d數值. time * number of joints * 3
        - jointHeirarchy: list of lists. 關節點連結資訊. each element is a list with two numbers in it.
            First number indicates parent joint's index, the other indicates child joint's index.
        - hipPos: list of 1 dimention array. 每一組關節點的hip位置. XYZ values in the array.
        - axisJointInd: list of list. 哪一些點需要在它上面繪製座標軸
        - frameRate: 更新率
    '''
    # 輸入joint data, 使用open 3d顯示骨架運動資訊

    # create visualizer and window.
    # vis = o3d.visualization.Visualizer()
    vis = o3d.visualization.VisualizerWithKeyCallback()
    # vis.create_window(height=480, width=640)
    vis.create_window(height=768, width=1024)
    opt = vis.get_render_option()
    opt.point_size = 10
    opt.line_width = 5

    # 註冊callback function, 讓按鍵觸發可以被偵測
    ## 印出當前顯示的frame index -> 按下空格鍵
    vis.register_key_callback(ord(' '), printCurFrameIdxCallback)

    # create axis at origin
    # origin_axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2)
    # vis.add_geometry(origin_axis)

    # 根據hipPos改變每一組關節點的3d positions資訊
    for i, _hipPos in enumerate(hipPos):
        jointData[i] = jointData[i] + _hipPos[np.newaxis, np.newaxis, :]

    # 在部分節點繪製坐標軸標示 (e.g. 下半身的upper leg and lower leg)
    axisFrameList = []
    axisFrameOriginPosList = []
    for i, _indList in enumerate(axisJointInd):
        _axis_frame_list = []
        _axis_frame_pos_list = []
        for _ind in _indList:
            _axis_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(
                size=0.1, origin=jointData[i][0, _ind, :]
            )
            vis.add_geometry(_axis_frame)
            _axis_frame_list.append(_axis_frame)
            _axis_frame_pos_list.append(np.array(_axis_frame.vertices))
        axisFrameList.append(_axis_frame_list)
        axisFrameOriginPosList.append(_axis_frame_pos_list)

    # initial point clouds list
    # initial line set list
    pcdList = []
    lineSetList = []
    for _jointData, _jointHeirarchy in zip(jointData, jointHeirarchy):
        _pcd = o3d.geometry.PointCloud()
        _pcd.points = o3d.utility.Vector3dVector(_jointData[0, :, :])
        pcdList.append(_pcd)

        colors = [[1, 0, 0] for i in range(len(_jointHeirarchy))]
        _line_set = o3d.geometry.LineSet()
        _line_set.points = o3d.utility.Vector3dVector(_jointData[0, :, :])
        _line_set.lines = o3d.utility.Vector2iVector(_jointHeirarchy)
        _line_set.colors = o3d.utility.Vector3dVector(colors)
        lineSetList.append(_line_set)
        # print(_jointData[0, :, :])
        # print(np.asarray(_line_set.points).shape)
        # exit()

        vis.add_geometry(_pcd)
        vis.add_geometry(_line_set)

    # 將右腿的顏色改成與其他點不同, 並且把線段/骨頭顏色改為黑色
    # 右腿的三個點的index = {0, 1, 2}, 三個骨頭index = {0, 2, 4}
    # 假設前面兩個輸入的point cloud是full body point cloud
    # bodyPointCloudInd = [0, 1]
    bodyPointCloudInd = [0]
    bodyBoneInd = [0, 2, 4]
    bodyJointInd = [0, 1, 2]
    for _pcdInd in bodyPointCloudInd:
        _colors = np.array([[0, 0, 0] for i in range(17)])
        _boneColors = np.array([[0, 0, 0] for i in range(16)])
        for i in bodyBoneInd:
            _boneColors[i, :] = [1, 0, 0]
        lineSetList[_pcdInd].colors = o3d.utility.Vector3dVector(_boneColors)
        for i in bodyJointInd:
            _colors[i, :] = [1, 0, 0]
        pcdList[_pcdInd].colors = o3d.utility.Vector3dVector(_colors)

    # 手的食指的點與線段也要修改顏色
    # 食指四個點index = {5, 6, 7, 8}, 三個骨頭的index = {8, 9, 10}
    # handPointCloudInd = [2]
    # handPointCloudInd = [0]
    # handBoneInd = [8, 9, 10]
    # handJointInd = [5, 6, 7, 8]
    # _colors = np.array([[0, 0, 0] for i in range(21)])
    # _boneColors = np.array([[0, 0, 0] for i in range(20)])
    # for i in handBoneInd:
    #     _boneColors[i, :] = [1, 0, 0]
    # for i in handJointInd:
    #     _colors[i, :] = [1, 0, 0]
    # for i in handPointCloudInd:
    #     pcdList[i].colors = o3d.utility.Vector3dVector(_colors)
    #     lineSetList[i].colors = o3d.utility.Vector3dVector(_boneColors)

    # Trajectory改變顏色, 改成藍色
    trajectoryInd = [1]
    # trajectoryInd = [0]
    for _pcdInd in trajectoryInd:
        _colors = np.array([[0, 0, 1] for i in range(jointData[_pcdInd].shape[1])])
        pcdList[_pcdInd].colors = o3d.utility.Vector3dVector(_colors)

    # 修改camera參數, 包含內參與外參
    if cameraParam is not None:
        res=vis.get_view_control().convert_from_pinhole_camera_parameters(cameraParam, allow_arbitrary=True)
        print('Is camera parameter adjust? ', res)

    # 紀錄visualize的影像畫面
    imgs = []

    # to add new points each dt secs.
    dt = frameRate
    previous_t = time.time()
    global curTimeInd
    curTimeInd = 0

    # run non-blocking visualization.
    # To exit, press 'q' or click the 'x' of the window.
    keep_running = True
    while keep_running:

        if time.time() - previous_t > dt:

            # update joints
            # update bones
            # TODO update coordinate frame
            for i, _jointData in enumerate(jointData):
                pcdList[i].points = o3d.utility.Vector3dVector(_jointData[curTimeInd, :, :])
                vis.update_geometry(pcdList[i])

                lineSetList[i].points = o3d.utility.Vector3dVector(_jointData[curTimeInd, :, :])
                vis.update_geometry(lineSetList[i])

                for _j, _axisJointInd in enumerate(axisJointInd[i]):
                    axisFrameList[i][_j].vertices = \
                        o3d.utility.Vector3dVector(
                            axisFrameOriginPosList[i][_j] - _jointData[0, _axisJointInd:_axisJointInd+1, :] + _jointData[curTimeInd, _axisJointInd:_axisJointInd+1, :]
                        )
                    vis.update_geometry(axisFrameList[i][_j])

            previous_t = time.time()
            curTimeInd += 1
            if curTimeInd >= jointData[0].shape[0]:
                break
            if return_imgs:
                imgs.append(vis.capture_screen_float_buffer())

        keep_running = vis.poll_events()
        vis.update_renderer()

    vis.destroy_window()

    # return 擷取的圖片
    if return_imgs:
        imgs = [np.asarray(i) for i in imgs]
        return imgs
    pass

def printCurFrameIdxCallback(vis):
    global curTimeInd
    print('cur frame index: ', curTimeInd)
    # print(o3d.camera.PinholeCameraParameters.extrinsic)
    print('Intrinsic matrix: ', vis.get_view_control().convert_to_pinhole_camera_parameters().intrinsic.intrinsic_matrix)
    print('Extrinsic matrix: ', vis.get_view_control().convert_to_pinhole_camera_parameters().extrinsic)
    return True

def main():
    # create visualizer and window.
    vis = o3d.visualization.Visualizer()
    vis.create_window(height=480, width=640)

    # create axis at origin
    origin_axis = o3d.geometry.TriangleMesh.create_coordinate_frame()
    vis.add_geometry(origin_axis)

    # create plane by box mesh
    # Ref: https://github.com/isl-org/Open3D/issues/3618
    plane = o3d.geometry.TriangleMesh.create_box(width=1.0, height=1.0, depth=0.01)
    vis.add_geometry(plane)

    # initialize pointcloud instance.
    pcd = o3d.geometry.PointCloud()
    # *optionally* add initial points
    points = np.random.rand(10, 3)
    pcd.points = o3d.utility.Vector3dVector(points)

    # initialize line set
    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [1, 0, 1],
              [0, 1, 1], [1, 1, 1]]).astype(float)
    lines = [[0, 1], [0, 2], [1, 3], [2, 3], [4, 5], [4, 6], [5, 7], [6, 7],
             [0, 4], [1, 5], [2, 6], [3, 7]]
    colors = [[1, 0, 0] for i in range(len(lines))]
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(points)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector(colors)
    vis.add_geometry(line_set)

    # include it in the visualizer before non-blocking visualization.
    vis.add_geometry(pcd)

    # to add new points each dt secs.
    dt = 0.05
    # number of points that will be added
    n_new = 10

    previous_t = time.time()

    # run non-blocking visualization.
    # To exit, press 'q' or click the 'x' of the window.
    keep_running = True
    while keep_running:

        if time.time() - previous_t > dt:
            s = time.time()

            # Options (uncomment each to try them out):
            # 1) extend with ndarrays.
            pcd.points.extend(np.random.rand(n_new, 3))

            # 取代舊的points
            # print(np.asarray(pcd.points).shape[0])
            if np.asarray(pcd.points).shape[0] > 10:
                pcd.points = o3d.utility.Vector3dVector(np.random.rand(n_new, 3))

            # 線段當中某個點的x, y, z軸不斷增加
            points[0,0]+=0.001
            points[0,1]+=0.001
            points[0,2]+=0.001
            line_set.points = o3d.utility.Vector3dVector(points)
            # print(np.asarray(line_set.points).shape)

            # 2) extend with Vector3dVector instances.
            # pcd.points.extend(
            #     o3d.utility.Vector3dVector(np.random.rand(n_new, 3)))

            # 3) other iterables, e.g
            # pcd.points.extend(np.random.rand(n_new, 3).tolist())

            vis.update_geometry(pcd)
            vis.update_geometry(line_set)
            previous_t = time.time()

        keep_running = vis.poll_events()
        vis.update_renderer()

    vis.destroy_window()

    pass

# 目標是展示全身動作, 並且也展示腳的trajectory
# 正在考慮要不要實作
if __name__=='__main01__':
    # example animation joints position隨時間變動資料, 透過重複相同資料延長播放時間
    exampleAnimMotion = readExampleAnimPos('../positionData/fromDB/genericAvatar/leftFrontKickPositionFullJointsWithHead_withoutHip_075_quat_direct_normalized.json')
    exampleAnimMotion = np.tile(exampleAnimMotion, (3, 1, 1))
    # 只有腳的部分隨時間變動資料, 重複複製原始frame數量的次數
    frameCount = exampleAnimMotion.shape[0]
    leftFootTrajectory = exampleAnimMotion[:, 2, :]
    leftFootTrajectory = np.tile(leftFootTrajectory, (frameCount, 1, 1))

    # 修改成顯示單一frame
    # spFrameInd = 80
    # # spFrameInd = 105
    # repeatShowingTime = 10000
    # exampleAnimMotion = np.repeat(
    #     exampleAnimMotion[spFrameInd:spFrameInd+1, :, :],
    #     repeatShowingTime,
    #     axis=0
    # )
    # leftFootTrajectory = np.repeat(
    #     leftFootTrajectory[spFrameInd:spFrameInd+1, :, :],
    #     repeatShowingTime,
    #     axis=0
    # )

    # 讀取camera參數 (目的: 預先設定錄製mp4的拍攝角度)
    cameraParamFilePath = 'ScreenCamera_2023-05-02-13-22-03.json'
    cameraParam = o3d.io.read_pinhole_camera_parameters(cameraParamFilePath)

    capture_imgs = vizMotions(
        [exampleAnimMotion, leftFootTrajectory],
        [fullBodyBoneStrcuture, [[5, 10]]],
        [np.array([0, 0, 0.5]), np.array([0, 0, 0.5])],
        [[], [], []],
        0.05,
        True,
        cameraParam
    )

    # 將animation記錄成mp4
    # import imageio
    # outputVideoFilePath = 'trajectoryAndPose.mp4'
    # imageio.mimwrite(outputVideoFilePath, capture_imgs, fps=30, quality=8)
    # pass

# 最初始的展示. 展示手的作與全身的動作
if __name__=='__main__':
    # main()
    appliedRotMotion = readAppliedRotPos('../positionData/fromAfterMappingHand/leftSideKickStreamLinearMapping_FTTFFF.json')
    # appliedRotMotion = readAppliedRotPos('../positionData/leftSideKick_quat_directMapping.json')
    # 因為synthesis motion會少前面10個frame, 所以applied rotation版本需要捨去前面10個frame
    appliedRotMotion = appliedRotMotion[10:, :, :]

    # synthesisMotion = readSynthesisPos('../positionData/afterSynthesis/leftFrontKickStreamLinearMapping_TFFTTT_075_EWMA.json')
    exampleAnimMotion = readExampleAnimPos('../positionData/fromDB/genericAvatar/leftSideKickPositionFullJointsWithHead_withoutHip.json')
    synthesisMotion = exampleAnimMotion

    fingerMotion = readHandPos('../complexModel/leftSideKick.json', scale=[3.5, 1.5, 7], negate=[True, True, True])
    fingerMotion = fingerMotion[2600:3500, :, :]    # 側踢只有選部分區間的資料
    fingerMotion = fingerMotion[10:, :, :]
    # capture_imgs = vizMotions(
    #     [appliedRotMotion, synthesisMotion, fingerMotion],
    #     [fullBodyBoneStrcuture, fullBodyBoneStrcuture, handBoneStructure],
    #     [np.array([0, 0, 0.5]), np.array([-1, 0, 0.5]), np.array([1, 0, 0.5])],
    #     # [[], [jointsNames.LeftUpperLeg, jointsNames.LeftLowerLeg], [handJointsNames.indexMCP, handJointsNames.indexPIP]],
    #     [[], [], []],
    #     0.03333,
    #     True
    # )

    # 只展示單一個骨架的動作
    ## 讀取特定的camera設定檔案
    cameraParamFilePath = 'ScreenCamera_2023-03-30-14-17-05.json'   # For body motion
    # cameraParamFilePath = 'ScreenCamera_2023-03-30-15-44-09.json'   # For finger motion
    # cameraParamFilePath = 'ScreenCamera_2023-04-26-14-46-59.json'   # For finger motion (side kick)
    cameraParam = o3d.io.read_pinhole_camera_parameters(cameraParamFilePath)
    capture_imgs = vizMotions(
        [synthesisMotion],
        [fullBodyBoneStrcuture],
        # [fingerMotion],
        # [handBoneStructure],
        [np.array([0, 0, 0.5])],
        [[], [], []],
        0.03333,
        True,
        cameraParam
    )

    print(len(capture_imgs))
    print(capture_imgs[0].shape)
    # 輸出骨架運動過程的影片
    # import imageio
    # outputVideoFilePath = 'sideKickExampleAnimWithoutHip.mp4'
    # imageio.mimwrite(outputVideoFilePath, capture_imgs, fps=30, quality=8)

    # 顯示單一個frame的資訊, 方便拍攝論文的展示圖片
    # specificFrameInd = 57
    # # specificFrameInd = 67
    # repeatShowingTime = 10000
    # appliedRotMotion = np.repeat(
    #     appliedRotMotion[specificFrameInd:specificFrameInd+1, :, :], repeatShowingTime, axis=0
    # )
    # synthesisMotion = np.repeat(
    #     synthesisMotion[specificFrameInd:specificFrameInd+1, :, :], repeatShowingTime, axis=0
    # )
    # fingerMotion = np.repeat(
    #     fingerMotion[specificFrameInd:specificFrameInd+1, :, :], repeatShowingTime, axis=0
    # )
    # vizMotions(
    #     [appliedRotMotion, synthesisMotion, fingerMotion],
    #     [fullBodyBoneStrcuture, fullBodyBoneStrcuture, handBoneStructure],
    #     [np.array([-5, 0, 0.5]), np.array([-8, 0, 0.5]), np.array([5, 0, 0.5])],
    #     # [[], [jointsNames.LeftUpperLeg, jointsNames.LeftLowerLeg], [handJointsNames.indexMCP, handJointsNames.indexPIP]],
    #     [[], [], []],
    #     0.05
    # )

    # 錄製只有全身與手的立正狀態對比, 使用front kick放鬆姿態作為例子.
    # 需要繪製座標軸在特定的joint上方, 兩者需要距離近一些
    # vizMotions(
    #     [synthesisMotion, fingerMotion],
    #     [fullBodyBoneStrcuture, handBoneStructure],
    #     [np.array([-0.5, 0, 0.5]), np.array([0.5, 0, 0.5])],
    #     [[jointsNames.LeftUpperLeg, jointsNames.LeftLowerLeg], [handJointsNames.indexMCP, handJointsNames.indexPIP]],
    #     0.05
    # )


import os
import fnmatch
import shutil
from sklearn.model_selection import train_test_split
import random

imgPath = "/Users/aaronyu/Desktop/Project14_Image/surface-based-samples/other/canvas"
templatePath = "/Users/aaronyu/Desktop/Project14_Image/surface-based-samples/other/fragment"
allPath = "/Users/aaronyu/Desktop/Project14_Image/WorkSpace/All"
trainPath = "/Users/aaronyu/Desktop/Project14_Image/WorkSpace/Train"
testPath = "/Users/aaronyu/Desktop/Project14_Image/WorkSpace/Test"

def createAll(imgPath, templatePath, allPath):
    """

    :param imgPath: Path containing all images (canvas to be searched on)
    :param templatePath: Path containing all template (to be searched for in canvas)
    :param allPath: Path for image, yml, template
    :return: 3 folders under allPath, image, yml, template respectively
    """

    allImgPath = os.path.join(allPath, "img")
    allTemplatePath = os.path.join(allPath, "template")
    allYmlPath = os.path.join(allPath, "yml")
    if not os.path.exists(allImgPath):
        os.makedirs(allImgPath)
    if not os.path.exists(allTemplatePath):
        os.makedirs(allTemplatePath)
    if not os.path.exists(allYmlPath):
        os.makedirs(allYmlPath)

    imgFiles = fnmatch.filter(os.listdir(imgPath), "*.png")
    for img in imgFiles:
        imgFile = os.path.join(imgPath, img)
        ymlFile = os.path.join(imgPath, img.replace("png", "yml"))
        templateFile = os.path.join(templatePath, img)
        if not os.path.isfile(ymlFile):
            continue
        if not os.path.isfile(templateFile):
            continue
        shutil.copy(imgFile, allImgPath)
        shutil.copy(ymlFile, allYmlPath)
        shutil.copy(templateFile, allTemplatePath)

def createTrainTest(allPath, trainPath, testPath, testPerc = 0.2, seed = 2019):
    """

    :param allPath:
    :param trainPath:
    :param testPath:
    :return:
    """
    ## Create train/test path for each type - img, template, yml
    folders = ["img", "template", "yml"]
    for folder in folders:
        tmpTrainPath = os.path.join(trainPath, folder)
        if not os.path.exists(tmpTrainPath):
            os.makedirs(tmpTrainPath)
        tmpTestpath = os.path.join(testPath, folder)
        if not os.path.exists(tmpTestpath):
            os.makedirs(tmpTestpath)

    path1 = os.path.join(allPath, folders[0])
    imgFiles = fnmatch.filter(os.listdir(path1), "*.png")
    # path2 = os.path.join(all)
    totalSize = len(imgFiles)
    testSize = int(totalSize * testPerc)
    randomShuffler = random.Random(seed)
    randomShuffler.shuffle(imgFiles)

    ## Create Test
    for img in imgFiles[:testSize]:
        imgFile = os.path.join(allPath,folders[0], img)
        ymlFile = os.path.join(allPath,folders[2],  img.replace("png", "yml"))
        templateFile = os.path.join(allPath,folders[1], img)
        shutil.copy(imgFile, os.path.join(testPath, folders[0], img))
        shutil.copy(ymlFile, os.path.join(testPath, folders[2], img.replace("png", "yml")))
        shutil.copy(templateFile, os.path.join(testPath, folders[1], img))

    ## Create Train
    for img in imgFiles[testSize:]:
        imgFile = os.path.join(allPath,folders[0], img)
        ymlFile = os.path.join(allPath,folders[2],  img.replace("png", "yml"))
        templateFile = os.path.join(allPath,folders[1], img)
        shutil.copy(imgFile, os.path.join(trainPath, folders[0], img))
        shutil.copy(ymlFile, os.path.join(trainPath, folders[2], img.replace("png", "yml")))
        shutil.copy(templateFile, os.path.join(trainPath, folders[1], img))

if __name__ == '__main__':
    createAll(imgPath, templatePath, allPath)
    createTrainTest(allPath, trainPath, testPath)


def solution(N, stages):
    answer = []

    #stages.sort()

    #idx = 0
    total_player = len(stages)
    rate = {}

    '''
    for i in range(1, N + 1):
        cnt = 0

        while idx < len(stages) and i == stages[idx]:
            cnt += 1
            idx += 1

        if total_player <= 0:
            rate[i] = 0
        else:
            rate[i] = cnt / total_player
            total_player -= cnt
    '''

    for stage in range(1, N + 1):
        if total_player <= 0:
            rate[stage] = 0
        else:
            cnt = stages.count(stage)
            rate[stage] = cnt / total_player
            total_player -= cnt

    '''
    rate = sorted(rate.items(), key=lambda x: x[1], reverse=True)

    for i in rate:
        answer.append(i[0])
    '''

    # dict를 list에 넣을경우 key값이 들어가게된다.
    answer = sorted(rate, key=lambda x: rate[x], reverse=True)

    return answer

# [3,4,2,1,5]
print(solution(5, [2, 1, 2, 6, 2, 4, 3, 3]))
# [4,1,2,3]
print(solution(4, [4, 4, 4, 4, 4]))


import os
import shutil
import argparse
import record
import tensorflow as tf

def convert2tfrecord(work_dir, filepath=''):
    # work_dir: '/assets/datasets'
    pascal_root = os.path.join(work_dir, 'data')
    output_dir = os.path.join(work_dir, 'tfrecord')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    image_folder = os.path.join(pascal_root, 'JPEGImages')
    list_folder = os.path.join(pascal_root, 'ImageSets/Segmentation')
    if filepath != '':
        fn = os.path.basename(filepath)
        local_fn = os.path.join(image_folder, fn)
        local_fn_list = os.path.join(list_folder, 'val.txt')
        if not os.path.exists(local_fn):
            shutil.copyfile(filepath, local_fn)
        with open(local_fn_list, 'w') as of:
            print(fn[:fn.rfind('.')])
            of.write(fn[:fn.rfind('.')])
    # image_folder: 'data/JPEGImages'
    # list_folder:  'data/ImageSets/Segmentation'
    # output_dir:   'tfrecord'
    command = 'python seg/deeplab/datasets/build_voc2012_data.py \
                --image_folder={} \
                --list_folder={} \
                --image_format="jpg"\
                --output_dir={}'.format(image_folder, list_folder, output_dir)
    os.system(command)

def segment(checkpoint_dir, work_dir):
    # checkpoint_dir: 'exp/train_on_train_set/train'
    # dataset_dir:    'tfrecord'
    dataset_dir = os.path.join(work_dir, 'tfrecord')
    command = 'python seg/deeplab/vis.py \
                --logtostderr \
                --vis_split="val" \
                --model_variant="xception_65" \
                --atrous_rates=6 \
                --atrous_rates=12 \
                --atrous_rates=18 \
                --output_stride=16 \
                --decoder_output_stride=4 \
                --vis_crop_size=1505 \
                --vis_crop_size=2049 \
                --dataset="homeplus" \
                --colormap_type=homeplus \
                --vis_batch_size=1 \
                --checkpoint_dir={} \
                --vis_logdir={} \
                --dataset_dir={}'.format(checkpoint_dir, work_dir, dataset_dir)
    os.system(command)

def segment2skels(work_dir):
    segment_folder = os.path.join(work_dir, 'segmentation_results')
    command = 'python seg2skel/seg2skels.py \
                --segment-folder={}'.format(segment_folder)
    os.system(command)
    input('before load skel')
    # command = 'python seg2skel/load_skels.py \
    #             --segment-folder={}'.format(segment_folder)
    # os.system(command)

def main(args):
    # convert2tfrecord(args.work_dir, args.filepath)
    segment(args.checkpoint_dir, args.work_dir)
    input('before seg2skel')
    segment2skels(args.work_dir)

if __name__ == '__main__':
    program_dir = '/home/chr/sgy/web/program/'
    parser = argparse.ArgumentParser(description='Segmentation and Skeleton')
    parser.add_argument('--work-dir', default=program_dir+'assets/datasets/', type=str)
    parser.add_argument('--checkpoint-dir', default=program_dir+'assets/datasets/exp/train_on_train_set/train/', type=str)
    parser.add_argument('--filepath', default='', type=str)
    args = parser.parse_args()

    # print(os.environ['PYTHONPATH'])

    main(args)
    # with tf.Session() as sess:
    #     record.write_record(None, sess)
    #     input('write finish')

[1505, 2049]
1
None
None
None
val
xception_65
[6, 12, 18]
16

[1505, 2049]
1
None
None
None
val
xception_65
[6, 12, 18]
16


import torch
import torch.nn as nn
import torch.nn.functional as F

class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(ResBlock, self).__init__()

        self.model = nn.Sequential(nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1),
                                   nn.BatchNorm2d(out_ch),
                                   nn.ReLU(),
                                   nn.Conv2d(out_ch, out_ch, kernel_size=1, stride=1),
                                   nn.BatchNorm2d(out_ch))

    def forward(self, x):
        return x + self.model(x)

class IntModel(nn.Module):
    def __init__(self):
        super(IntModel, self).__init__()
        self.conv_A = self._build_conv()
        self.conv_B = self._build_conv()

        self.lstm_A = nn.LSTM(512, 256, 2, batch_first=True)
        self.lstm_B = nn.LSTM(512, 256, 2, batch_first=True)

        self.classifier = self._build_cls()

    def forward(self, x1, x2):
        bs = x1.shape[0]
        depth = x1.shape[1]
        chw = x1.shape[-3:]

        x1 = torch.reshape(x1, [-1] + list(chw))
        x2 = torch.reshape(x2, [-1] + list(chw))

        encoded_A = self.conv_A(x1)
        encoded_B = self.conv_B(x2)

        encoded_A = torch.reshape(encoded_A, (bs, depth, -1))
        encoded_B = torch.reshape(encoded_B, (bs, depth, -1))

        out_A, (h_A, c_A) = self.lstm_A(encoded_A)
        out_B, (h_B, c_B) = self.lstm_B(encoded_B)

        h_A = torch.reshape(torch.stack(torch.unbind(h_A, dim=1), dim=0), (bs, -1))
        h_B = torch.reshape(torch.stack(torch.unbind(h_B, dim=1), dim=0), (bs, -1))

        feat_AB = torch.cat((h_A, h_B), dim=1)

        prob = self.classifier(feat_AB)

        return prob

    def _build_cls(self, in_size=1024, num_cls=6):
        model = [nn.Linear(in_size, num_cls)]
        return nn.Sequential(*model)

    def _build_conv(self, in_ch=32, npf=64):
        n_blocks = 3
        model = [nn.Conv2d(in_ch, 64, kernel_size=3, stride=2, padding=1),
                 nn.ReLU(True)]

        for cnt in range(n_blocks):
            previous = npf * 2 ** cnt
            model += [nn.Conv2d(previous, previous * 2, kernel_size=3, stride=2, padding=1),
                      nn.BatchNorm2d(previous * 2),
                      nn.ReLU(True),
                      ResBlock(previous * 2, previous * 2)]

        model += [nn.AdaptiveMaxPool2d(1)]
        return nn.Sequential(*model)

if __name__ == '__main__':
    m = ResBlock(64,64)

    # x1 = torch.randn(4,5,32,64,64)
    # x2 = torch.randn(4,5,32,64,64)
    data = torch.randn(4,64,128,128)
    out = m(data)
    import ipdb
    ipdb.set_trace()


from random import randint
import time
from functools import wraps
import sys

sys.setrecursionlimit(100000000)

def timeit(func):
    @wraps(func)
    def timeit_wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        total_time = end_time - start_time
        print(f'Function {func.__name__} Took {total_time:.4f} seconds')
        return result
    return timeit_wrapper

def produce_random_dataset(dataset_size):

    dataset = []

    for i in range(dataset_size):

        dataset.append(randint(0, dataset_size))

    return dataset

# Heap Sort
def heapify(arr, n, i):
    largest = i  # Initialize largest as root
    l = 2 * i + 1  # left = 2*i + 1
    r = 2 * i + 2  # right = 2*i + 2

 # See if left child of root exists and is
 # greater than root

    if l < n and arr[i] < arr[l]:
        largest = l

 # See if right child of root exists and is
 # greater than root

    if r < n and arr[largest] < arr[r]:
        largest = r

 # Change root, if needed

    if largest != i:
        (arr[i], arr[largest]) = (arr[largest], arr[i])  # swap

  # Heapify the root.

        heapify(arr, n, largest)

# The main function to sort an array of given size

@timeit
def heapSort(arr):
    n = len(arr)

 # Build a maxheap.
 # Since last parent will be at ((n//2)-1) we can start at that location.

    for i in range(n // 2 - 1, -1, -1):
        heapify(arr, n, i)

 # One by one extract elements

    for i in range(n - 1, 0, -1):
        (arr[i], arr[0]) = (arr[0], arr[i])  # swap
        heapify(arr, i, 0)
#-----------------------------------

def merge(arr, l, m, r):
    n1 = m - l + 1
    n2 = r - m

    # create temp arrays
    L = [0] * (n1)
    R = [0] * (n2)

    # Copy data to temp arrays L[] and R[]
    for i in range(0, n1):
        L[i] = arr[l + i]

    for j in range(0, n2):
        R[j] = arr[m + 1 + j]

    # Merge the temp arrays back into arr[l..r]
    i = 0     # Initial index of first subarray
    j = 0     # Initial index of second subarray
    k = l     # Initial index of merged subarray

    while i < n1 and j < n2:
        if L[i] <= R[j]:
            arr[k] = L[i]
            i += 1
        else:
            arr[k] = R[j]
            j += 1
        k += 1

    # Copy the remaining elements of L[], if there
    # are any
    while i < n1:
        arr[k] = L[i]
        i += 1
        k += 1

    # Copy the remaining elements of R[], if there
    # are any
    while j < n2:
        arr[k] = R[j]
        j += 1
        k += 1

# l is for left index and r is right index of the
# sub-array of arr to be sorted

def mergeSort(arr, l, r):
    if l < r:

        # Same as (l+r)//2, but avoids overflow for
        # large l and h
        m = l+(r-l)//2

        # Sort first and second halves
        mergeSort(arr, l, m)
        mergeSort(arr, m+1, r)
        merge(arr, l, m, r)

@timeit
def do_merge_sort(dataset):

    mergeSort(dataset, 0, len(dataset)-1)

#-----------------------------------

#quicksort

def partition(array, low, high):

    # choose the rightmost element as pivot
    pivot = array[high]

    # pointer for greater element
    i = low - 1

    # traverse through all elements
    # compare each element with pivot
    for j in range(low, high):
        if array[j] <= pivot:

            # If element smaller than pivot is found
            # swap it with the greater element pointed by i
            i = i + 1

            # Swapping element at i with element at j
            (array[i], array[j]) = (array[j], array[i])

    # Swap the pivot element with the greater element specified by i
    (array[i + 1], array[high]) = (array[high], array[i + 1])

    # Return the position from where partition is done
    return i + 1

# function to perform quicksort

def quickSort(array, low, high):
    if low < high:

        # Find pivot element such that
        # element smaller than pivot are on the left
        # element greater than pivot are on the right
        pi = partition(array, low, high)

        # Recursive call on the left of pivot
        quickSort(array, low, pi - 1)

        # Recursive call on the right of pivot
        quickSort(array, pi + 1, high)

@timeit
def do_quick_sort(dataset):
    quickSort(dataset, 0, len(dataset) - 1)
    return dataset

#-----------------------------------

#Eray sort ---------------
def compare(dataset):

    data_length = len(dataset)

    if data_length <= 1:

        if data_length == 0:

            return

        return dataset

    last_index = 1
    compare_index = 0

    while last_index < data_length:

        if dataset[last_index] > dataset[compare_index]:

            compare_index = last_index

        else:

            if dataset[last_index] < dataset[compare_index]:

                pass

            else:

                compare_index = last_index

        last_index += 1

    return compare_index

def find_best(best, dataset, new_dataset):
    if best == 0:
        return new_dataset
    else :
        new_dataset.append(dataset.pop(compare(dataset)))
        best -= 1
        find_best(best, dataset, new_dataset)

@timeit
def eray_sort(dataset):
    new_dataset = []
    find_best(len(dataset)-1, dataset, new_dataset)
    return new_dataset
#-----------------------------------

def main():
    dataset_lenght = 2000

    print("--------------------")

    for i in range(100, dataset_lenght + 100, 100):
        dataset = produce_random_dataset(i)
        print(f"Dataset size: {i}")
        do_quick_sort(dataset)
        do_merge_sort(dataset)
        heapSort(dataset)
        eray_sort(dataset)
        print("--------------------")

if __name__ == '__main__':
    main()


T = int(input())

def BubbleSort(a):
    for i in range(len(a)):
        for j in range(0,len(a)-i-1):
            if a[j]>a[j+1]:
                a[j],a[j+1]=a[j+1],a[j]
    return a

for test_case in range(1, T + 1):
    N=int(input())
    numlist=list(map(int,(input().rstrip(" ").split(" "))))
    s_numlist=BubbleSort(numlist)

    print(f'#{test_case} {s_numlist[-1]-s_numlist[0]}')


# -*- coding: utf-8 -*-
"""
Created on Tue May  9 16:43:04 2017

@author: vurga
"""

import keras
import numpy as np
import sys
from sklearn.utils import shuffle
import sklearn.metrics

def read_data(file):

    with open(file, 'r') as source:
        lines = source.readlines()

    cvxs = []
    for line in lines:
        cvxs.append(np.fromstring(line.strip(), sep=" "))

    return np.concatenate([cvxs], axis=1)

if __name__ == '__main__':

    ## get data here
    data_X_a = read_data(sys.argv[1])
    data_X_b = read_data(sys.argv[2])
    data_y = np.concatenate((np.tile([1, 0], (int(len(data_X_a)/4), 1)), np.tile([0, 1], (int(len(data_X_a)/4), 1)), np.tile([1, 0], (int(len(data_X_b)/4), 1)), np.tile([0, 1], (int(len(data_X_b)/4), 1))))

    ## strategy for feature combination
    ## simple concatenation
    data_X_conc= np.concatenate((data_X_a, data_X_b), axis=1)
    ## element-wise multiplication
    data_X_mult = np.multiply(data_X_a, data_X_b)
    ## subtraction
    data_X_sub = data_X_a - data_X_b

    data_X = np.concatenate((data_X_conc, data_X_mult, data_X_sub), axis=1)
    #data_X = data_X_mult

    ## paste & shuffle
    full_data = np.concatenate((data_X, data_y), axis=1)
    full_data = shuffle(full_data, random_state=3)

    training_portion = int(len(full_data) * 0.875)
    test_portion = int(len(full_data) * 0.025)

    ## layer sizes
    n_input = data_X.shape[1] ## number of input neurons
    n_hidden = 4096 ## hidden layer size
    n_output = data_y.shape[1] ## number of labels

    train_X = full_data[:training_portion,:n_input]
    train_y = full_data[:training_portion,n_input:]

    test_X = full_data[training_portion:training_portion+test_portion,:n_input]
    test_y = full_data[training_portion:training_portion+test_portion,n_input:]

    learning_rate = 0.0001

    # define the architecture of the network
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(n_hidden, input_dim=n_input, init="uniform", activation="sigmoid"))
    model.add(keras.layers.Dense(int(n_hidden/2), init="uniform", activation="sigmoid"))
    #model.add(keras.layers.Dense(int(n_hidden/4), init="uniform", activation="sigmoid"))
    model.add(keras.layers.Dense(n_output))
    model.add(keras.layers.Activation("softmax"))

    # train the model using SGD
    print("[INFO] compiling model...")
    opt = keras.optimizers.Adagrad(lr=learning_rate)
    model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=['accuracy'])
    model.fit(train_X, train_y, nb_epoch=15, batch_size=32)

    # show the accuracy on the testing set
    test_preds = np.argmax(model.predict(test_X), axis=1)
    test_y = np.argmax(test_y, axis=1)
    print(sklearn.metrics.precision_score(test_preds, test_y))
    print(sklearn.metrics.recall_score(test_preds, test_y))
    print(sklearn.metrics.f1_score(test_preds, test_y))

    model.save('ctx_allcomb_4096-2048.h5')


# -*- coding: utf-8 -*-
"""
Created on Wed Mar 10 11:51:30 2021

@author: perger
"""

import numpy as np
import pandas as pd
from pyomo.environ import *

def run_LP(load, PV, prosumer_data, grid_data, weight,
            distances, battery, solver_name, sharing=False):

    time_steps = load.index.tolist()
    index_time = list(range(len(time_steps)))
    prosumer = load.columns.tolist()
    weight=weight['weight'].tolist()

    # Define some parameters and variables
    SoC_max = 'Maximum Storage|Electricity|Energy Storage System'
    SoC_min = 'Minimum Storage|Electricity|Energy Storage System'
    q_bat_max = 'Maximum Charge|Electricity|Energy Storage System'
    q_bat_min = 'Maximum Discharge|Electricity|Energy Storage System'
    w = 'Price|Carbon'

    # deactivate BESS
    if battery == False:
        for i in prosumer:
            prosumer_data.loc[SoC_max,i] = 0
            prosumer_data.loc[q_bat_max,i] = 0

    eta_battery = 0.9

    #index_time = list(range(len(time_steps)))

    wtp={}
    for i in prosumer:
        wtp[i] = {}
        for j in prosumer:
            wtp[i][j] = (
                grid_data['Residential']
                + prosumer_data.loc[w,j]
                * (1 - distances.loc[i,j])
                * grid_data['Emissions']
                / 1000000)

    # Define model as concrete model
    model = ConcreteModel()

    #Define optimization variables
    model.q_G_in = Var(time_steps,
                       prosumer,
                       within = NonNegativeReals)
    model.q_G_out = Var(time_steps,
                        prosumer,
                        within = NonNegativeReals)
    model.q_share = Var(time_steps,
                        prosumer,
                        prosumer,
                        within = NonNegativeReals)
    model.q_B_in = Var(time_steps,
                       prosumer,
                       within = NonNegativeReals)
    model.q_B_out = Var(time_steps,
                        prosumer,
                        within = NonNegativeReals)
    model.SoC = Var(time_steps,
                    prosumer,
                    within = NonNegativeReals)
    model.community_welfare = Var(time_steps,
                                  prosumer,
                                  within = Reals)

    # Define constraints

    def no_sharing_rule(model, i, j, t):
         if i is not j:
             return(model.q_share[t,i,j] == 0)
         else:
             return (model.q_share[t,i,j] >= 0)

    if sharing == False:
         model.no_sharing_con = Constraint(prosumer,
                                           prosumer,
                                           time_steps,
                                           rule = no_sharing_rule)

    def load_constraint_rule(model, i, t):
         return (model.q_G_in[t,i]
                 + model.q_B_out[t,i]
                 + sum(model.q_share[t,j,i] for j in prosumer)
                 - load.loc[t,i] == 0)
    model.load_con = Constraint(prosumer,
                                 time_steps,
                                 rule = load_constraint_rule)

    def PV_constraint_rule(model, i, t):
         return (model.q_G_out[t,i]
                 + model.q_B_in[t,i]
                 + sum(model.q_share[t,i,j] for j in prosumer)
                 - PV.loc[t,i] == 0)
    model.PV_con = Constraint(prosumer,
                               time_steps,
                               rule = PV_constraint_rule)

    def SoC_min_constraint_rule(model, i, t):
         return (model.SoC[t,i] >= prosumer_data.loc[SoC_min][i])
    model.SoC_min_con = Constraint(prosumer,
                                    time_steps,
                                    rule = SoC_min_constraint_rule)

    def SoC_max_constraint_rule(model, i, t):
         return (model.SoC[t,i] <= prosumer_data.loc[SoC_max][i])
    model.SoC_max_con = Constraint(prosumer,
                                    time_steps,
                                    rule = SoC_max_constraint_rule)

    def q_B_in_constraint_rule(model, i, t):
         return (model.q_B_in[t,i] <= prosumer_data.loc[q_bat_max][i])
    model.q_B_in_con = Constraint(prosumer,
                                  time_steps,
                                  rule = q_B_in_constraint_rule)

    def q_B_out_constraint_rule(model, i, t):
         return (model.q_B_out[t,i] <= prosumer_data.loc[q_bat_min][i])
    model.q_B_out_con = Constraint(prosumer,
                                   time_steps,
                                   rule = q_B_out_constraint_rule)

    def SoC_constraint_rule(model, i, t):
         if t == 0:
             return (model.SoC[time_steps[-1],i]
                     + model.q_B_in[time_steps[t],i] * eta_battery
                     - model.q_B_out[time_steps[t],i] / eta_battery
                     - model.SoC[time_steps[t],i] == 0)
         elif t > 0:
             return (model.SoC[time_steps[t-1],i]
                     + model.q_B_in[time_steps[t],i] * eta_battery
                     - model.q_B_out[time_steps[t],i] / eta_battery
                     - model.SoC[time_steps[t],i] == 0)
    model.SoC_con = Constraint(prosumer,
                               index_time,
                               rule = SoC_constraint_rule)

    def CW_constraint_rule(model, i, t):
         return model.community_welfare[t,i] == (
             - grid_data.loc[t,'Residential'] * model.q_G_in[t,i] * weight[t]
             + grid_data.loc[t, 'DA'] * model.q_G_out[t,i] * weight[t]
             + sum(wtp[j][i][t]
                   * model.q_share[t,j,i]
                   * weight[t]
                   for j in prosumer))
    model.CW_con = Constraint(prosumer,
                              time_steps,
                              rule = CW_constraint_rule)

    # Objective function
    community_welfare = {new_list: [] for new_list in prosumer}
    prosumer_welfare = {new_list: [] for new_list in prosumer}
    prosumer_welfare2 = {new_list: [] for new_list in prosumer}

    for i in prosumer:
        community_welfare[i] = sum(- grid_data.loc[t,'Residential'] * model.q_G_in[t,i]*weight[t]
                                    + grid_data.loc[t, 'DA'] * model.q_G_out[t,i]*weight[t]
                                    for t in time_steps)
        prosumer_welfare[i] = sum(wtp[i][j][t]
                                  * model.q_share[t,i,j]
                                  * weight[t]
                                  for j in prosumer
                                  for t in time_steps)
        prosumer_welfare2[i] = sum(wtp[j][i][t]
                                   * model.q_share[t,j,i]
                                   * weight[t]
                                   for j in prosumer
                                   for t in time_steps)

    #     # 1. prosumer i sells to prosumer j
    #     # 2. prosumer i buys from prosumer j

    # model.obj = Objective(
    #     expr = sum(model.community_welfare[t,i]
    #                 for t in time_steps
    #                 for i in prosumer),
    #     sense = maximize)

    model.obj = Objective(
        expr = sum(community_welfare[i] + prosumer_welfare2[i]
                   for i in prosumer),
        sense = maximize)

    opt = SolverFactory(solver_name)
    opt_success = opt.solve(model)

    # Evaluate the results
    community_welfare = value(model.obj)

    q_share_total = pd.DataFrame(index=prosumer)
    for j in prosumer:
        a = []
        for i in prosumer:
            a.append(value(sum(model.q_share[t,i,j]*weight[t]
                               for t in time_steps)))
        q_share_total[j] = a

    results= pd.DataFrame(index=prosumer)
    for i in prosumer:
        results.loc[i,'buying grid'] = value(sum(model.q_G_in[t,i]
                                                 * weight[t]
                                                 for t in time_steps))
        results.loc[i,'selling grid'] = value(sum(model.q_G_out[t,i]
                                                  * weight[t]
                                                  for t in time_steps))
        results.loc[i,'battery charging'] = value(sum(model.q_B_in[t,i]
                                                      * weight[t]
                                                      for t in time_steps))
        results.loc[i,'battery discharging'] = value(sum(model.q_B_out[t,i]
                                                         * weight[t]
                                                         for t in time_steps))
        results.loc[i,'self-consumption'] = q_share_total.loc[i,i]
        results.loc[i,'buying community'] = (sum(q_share_total.loc[j,i]
                                                 for j in prosumer)
                                             - q_share_total.loc[i,i])
        results.loc[i,'selling community'] = (sum(q_share_total.loc[i,j]
                                                  for j in prosumer)
                                              - q_share_total.loc[i,i])
        results.loc[i,'emissions'] = value(sum(model.q_G_in[t,i]
                                               * weight[t]
                                               * grid_data.loc[t,'Emissions']
                                               / 1000000
                                                for t in time_steps))
        results.loc[i,'costs'] = value(sum(
            grid_data.loc[t,'Residential'] * model.q_G_in[t,i] * weight[t]
            + sum(wtp[j][i][t] * model.q_share[t,j,i] * weight[t] for j in prosumer)
            - grid_data.loc[t,'DA'] * model.q_G_out[t,i] * weight[t]
            - sum(wtp[i][j][t] * model.q_share[t,i,j] * weight[t] for j in prosumer)
            for t in time_steps))

    return results, q_share_total, community_welfare


import logging

import paddle

import paddle.vision.transforms as transforms
import paddle.vision.datasets as datasets

from paddle.io import DataLoader,RandomSampler, SequenceSampler,DistributedBatchSampler,Dataset

from paddle.distributed import fleet

logger = logging.getLogger(__name__)

class MyDataset(Dataset):
    def __init__(self, trainset):
        self.num_samples = len(trainset)
        self.trainset=trainset
    def __getitem__(self,idx):

        if(isinstance(idx,int)):
            print(1)
            image,label=self.trainset[idx]
        else:
            print(2)
            image,label=self.trainset[idx[0]]

        return image,label
    def __len__(self):
        return self.num_samples

def get_loader(args):
    if args.local_rank not in [-1, 0]:
        paddle.distributed.barrier()

    transform_train = transforms.Compose([
        transforms.RandomResizedCrop((args.img_size, args.img_size), scale=(0.05, 1.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
    transform_test = transforms.Compose([
        transforms.Resize((args.img_size, args.img_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])

    if args.dataset == "cifar10":
        trainset = datasets.Cifar10(mode='train',
                                    download=True,
                                    transform=transform_train)
        testset = datasets.Cifar10(mode='test',
                                   download=True,
                                   transform=transform_test) if args.local_rank in [-1, 0] else None

    elif args.dataset == "cifar100":
        trainset = datasets.Cifar100(mode='train',
                                     download=True,
                                     transform=transform_train)
        train_dataset = MyDataset(trainset)

        testset = datasets.Cifar100(mode='test',
                                    download=True,
                                    transform=transform_test) if args.local_rank in [-1, 0] else None
        test_dataset = MyDataset(testset)

    if args.local_rank == 0:
        fleet.barrier_worker()

    tra_sampler = RandomSampler(train_dataset)if args.local_rank == -1 else DistributedBatchSampler(train_dataset,
                                                                                batch_size=1,
                                                                                shuffle=True)

    train_sampler=paddle.io.BatchSampler(sampler=tra_sampler,
                                        batch_size=args.train_batch_size,
                                        drop_last=True)
    tes_sampler =  SequenceSampler(test_dataset)
    test_sampler=paddle.io.BatchSampler(sampler=tes_sampler,
                                        batch_size=args.eval_batch_size,
                                        drop_last=True)

    train_loader = DataLoader(train_dataset,
                              batch_sampler=train_sampler,
                              num_workers=0)
    test_loader = DataLoader(test_dataset,
                             batch_sampler=test_sampler,
                             num_workers=0) if testset is not None else None

    return train_loader, test_loader