"""
utils.py contains helper function to deal with creating founders,
processing genetic map and processing sample maps.
"""

import numpy as np
import allel
import pandas as pd
import os
import scipy.interpolate

from .person import Person, create_new, create_new_non_rec

def get_chm_info(genetic_map,vcf_data):

    """
    get chromosome length in morgans from genetic map.
    Assumes genetic_map is sorted.
    """

    # getting chm number...
    chm = list(set(vcf_data["variants/CHROM"]))
    if len(chm) == 1:
        chm = chm[0]
    else:
        raise Exception("Multiple chromosomes in this file!!!")
    chm = chm.lstrip("chr") # in some reference files we have this instead of just chr number. 22 or chr22

    # read in genetic map and subset to chm number.
    genetic_df = pd.read_csv(genetic_map,delimiter="\t",header=None,comment="#",dtype=str)
    genetic_df.columns = ["chm","pos","cM"]
    genetic_chm = genetic_df[genetic_df["chm"]==chm]

    if len(genetic_chm) == 0:
        genetic_chm = genetic_df[genetic_df["chm"]=="chr"+chm] # sometimes it is called chr22 instead of 22

    genetic_chm = genetic_chm.astype({"chm":str,"pos":int,"cM":float})

    # get length of chm.
    chm_length_morgans = max(genetic_chm["cM"])/100.0

    # get snp info - snps in the vcf file and their cm values.
    # then compute per position probability of being a breapoint.
    # requires some interpolation and finding closest positions.
    """
    # 1: Minimum in a sorted array approach and implemented inside admix().
        - O(logn) every call to admix. Note that admix() takes O(n) anyway.
    # 2: Find probabilities using span. - One time computation.

    """
    # This adds 0 overhead to code runtime.
    # get interpolated values of all reference snp positions
    genomic_intervals = scipy.interpolate.interp1d(x=genetic_chm["pos"].to_numpy(), y=genetic_chm["cM"].to_numpy(),fill_value="extrapolate")
    genomic_intervals = genomic_intervals(vcf_data["variants/POS"])
    # Method 1: a bit incorrect - to ise midpoints
    # # get their midpoints. goal is to get length spanned by each reference snp in the chromosome
    # midpts = (genomic_intervals[0:-1] + genomic_intervals[1:])/2
    # start_end_pts = np.concatenate((np.array([0]),midpts,np.array([genomic_intervals[-1]])))
    # lengths = start_end_pts[1:] - start_end_pts[0:-1]
    # # normalize the lengths after removing first since first point is not a breaking point by design
    # bp = lengths[1:] / lengths[1:].sum()
    # Method 2: simpler and more accurate
    lengths = genomic_intervals[1:] - genomic_intervals[0:-1]
    bp = lengths / lengths.sum()

    genetic_map_data = {}
    genetic_map_data["chm"] = chm
    genetic_map_data["chm_length_snps"] = vcf_data["calldata/GT"].shape[0]
    genetic_map_data["chm_length_morgans"] = chm_length_morgans
    genetic_map_data["breakpoint_probability"] = bp

    return genetic_map_data

def get_sample_map_data(sample_map, vcf_data, sample_weights = None):

    """
    Inputs:
    sample_map: tab delimited file with sample, population and no header.
    vcf_data: allel.read_vcf output. It is the reference vcf file information.

    Returns:
    sample_map_data: dataframe with sample, population, population_code and index in vcf_data referecnce.

    """

    # reading sample map
    sample_map_data = pd.read_csv(sample_map,delimiter="\t",header=None,comment="#")
    sample_map_data.columns = ["sample","population"]

    # creating ancestry map into integers from strings
    # id is based on order in sample_map file.
    ancestry_map = {}
    curr = 0
    for i in sample_map_data["population"]:
        if i in ancestry_map.keys():
            continue
        else:
            ancestry_map[i] = curr
            curr += 1
    # print("Ancestry map",ancestry_map)
    sample_map_data["population_code"] = np.vectorize(ancestry_map.get)(sample_map_data["population"])

    # getting index of samples in the reference files

    b = vcf_data["samples"]
    a = np.array(list(sample_map_data["sample"]))

    sorter = np.argsort(b)
    indices = sorter[np.searchsorted(b, a, sorter=sorter)]
    sample_map_data["index_in_reference"] = indices

    if sample_weights is not None:
        sample_weights_df = pd.read_csv(sample_weights,delimiter="\t",header=None,comment="#")
        sample_weights_df.columns = ["sample","sample_weight"]
        sample_map_data = pd.merge(sample_map_data, sample_weights_df, on='sample')

    else:
        sample_map_data["sample_weight"] = [1.0/len(sample_map_data)]*len(sample_map_data)

    return sample_map_data

def build_founders(vcf_data, genetic_map_data, sample_map_data):

    """
    Returns founders - a list of Person datatype.
    founders_weight - a list with a weight for each sample in founders

    Inputs
    genetic_map_data - output of get_chm_info

    """

    # information about snps, chm, lengths from reference, genetic map.
    chm = genetic_map_data["chm"]
    chm_length_morgans = genetic_map_data["chm_length_morgans"]
    chm_length_snps = genetic_map_data["chm_length_snps"]

    # building founders
    founders = []

    for i in sample_map_data.iterrows():

        # first get the index of this sample in the vcf_data.
        # if not there, skip and print to log.

        index = i[1]["index_in_reference"]

        name = i[1]["sample"]

        # when creating maternal, paternal make sure it has same keys

        maternal = {}
        paternal = {}

        # let us use the first for maternal in the vcf file...
        maternal["snps"] = vcf_data["calldata/GT"][:,index,0]
        paternal["snps"] = vcf_data["calldata/GT"][:,index,1]

        # single ancestry assumption.
        maternal["anc"] = np.array([i[1]["population_code"]]*chm_length_snps)
        paternal["anc"] = np.array([i[1]["population_code"]]*chm_length_snps)

        # any more info like coordinates, prs can be added here.

        p = Person(chm,chm_length_morgans,chm_length_snps,maternal,paternal,name)

        founders.append(p)

    founders_weight = sample_map_data["sample_weight"]

    return founders,founders_weight

def create_dataset(founders,founders_weight,num_samples_per_gen,gens_to_ret,breakpoint_probability=None,random_seed=42,verbose=True):

    np.random.seed(random_seed)

    if breakpoint_probability is None:
        print("Warning! Genetic map not taken into account when sampling breakpoints.")
        print("Uniformly sampling from given snps.")

    max_gen = max(gens_to_ret)

    overall = {}
    overall[0] = founders

    # could be greater than or equal to 100 to avoid founder bias.
    number_of_admixed_samples = max(400,num_samples_per_gen)

    for gen in range(1,max_gen+1):
        if verbose:
            print("Simulating generation ",gen)
        this_gen_people = []

        for i in range(number_of_admixed_samples):
            # select any 2 parents from prev. gen.
            # if 1st generation, choose from founders
            if gen == 1:
                id1,id2 = np.random.choice(len(founders),size=2,replace=False, p=founders_weight)
            # else choose from num_samples_per_gen
            else:
                id1,id2 = np.random.choice(number_of_admixed_samples,size=2,replace=False)
            p1 = overall[gen-1][id1]
            p2 = overall[gen-1][id2]

            adm = create_new(p1,p2,breakpoint_probability)
            this_gen_people.append(adm)

        overall[gen] = this_gen_people

        if gen-1 not in gens_to_ret and gen-1 !=0:
            del overall[gen-1]

    # for every requested generation, only keep num_samples_per_gen
    for key in overall.keys():
        if key == 0:
            continue
        overall[key] = overall[key][0:num_samples_per_gen]
        # overall[key] = list(np.random.choice(overall[key], \
        #                            size=num_samples_per_gen,replace=False))

    return overall

def create_non_rec_dataset(founders,founders_weight,num_samples_per_gen,gens_to_ret,breakpoint_probability=None,random_seed=42,verbose=True):

    np.random.seed(random_seed)

    if breakpoint_probability is None:
        print("genetic map not taken into account when sampling breakpoints.")
        print("Uniformly sampling from given snps.")

    max_gen = max(gens_to_ret)

    overall = {}
    overall[0] = founders

    # could be greater than or equal to 100 to avoid founder bias.
    number_of_admixed_samples = num_samples_per_gen

    for gen in gens_to_ret:

        # BUG fix: Gen ==0 means do nothing.
        if gen == 0:
            continue

        if verbose:
            print("Simulating generation ",gen)
        this_gen_people = []

        for i in range(number_of_admixed_samples):
            # use the new admixture tool for non-recursive simulation

            adm = create_new_non_rec(founders,founders_weight,gen,breakpoint_probability)
            this_gen_people.append(adm)

        overall[gen] = this_gen_people

    return overall

def write_output(root,dataset,verbose=True):

    """
    creates output numpy files in root directory - under that, we will have gen_{}
    folders and then we have the npy files.
    Make sure root has chm number in it.

    """

    if not os.path.isdir(root):
        os.makedirs(root)

    for gen in dataset.keys():

        if verbose:
            print("Writing generation: {}".format(gen))

        if gen not in dataset.keys():
            if verbose:
                print("Did not simulate gen {}".format(gen))
            continue

        # make directory
        gen_dir = os.path.join(root, "gen_{}".format(gen))
        if not os.path.isdir(gen_dir):
            os.makedirs(gen_dir)

        snps = []
        anc = []
        for person in dataset[gen]:
            snps.append(person.maternal["snps"])
            snps.append(person.paternal["snps"])
            anc.append(person.maternal["anc"])
            anc.append(person.paternal["anc"])

        # create npy files.
        snps = np.stack(snps)
        np.save(gen_dir+"/mat_vcf_2d.npy",snps)

        # create map files.
        anc = np.stack(anc)
        np.save(gen_dir+"/mat_map.npy",anc)


from turtle import *
from random import *

#uma função para mover a tartaruga para uma posição aleatória
def moveToRandomLocation():
    penup()
    setpos( randint(-700,700) , randint(-700,700))
    pendown()

#uma função para desenhar uma estrela de um tamanho específico
def drawStar(starSize, starColour):
    color(starColour)
    pendown()
    begin_fill()
    for side in range(5):
        left(144)
        forward(starSize)
    end_fill()
    penup()

#uma função para desenhar um planeta de um tamanho específico
def drawPlanet(starSize, starColour):
    color(starColour)
    pendown()
    begin_fill()
    for side in range(360):
        left(1)
        forward(starSize)
    end_fill()
    penup()

speed(30)

#isso desenha um fundo preto
bgcolor("black")

#desenha 200 estrelas brancas e 200 estrelas Azuis (tamanho,posições aleatórias)
for star in range(200):
    moveToRandomLocation()
    drawStar( randint(1,5) , "White")

for star in range(200):
    moveToRandomLocation()
    drawStar( randint(1,5) , "SkyBlue")

#desenha Planetas e sol (tamanho,posições especificas)
for Planeta in range(1):
    #SOL
    penup()
    setpos(-700,-100)
    pendown()
    drawPlanet(3, "gold")
    #MERCURIO
    penup()
    setpos(-400,30)
    pendown()
    drawPlanet(0.4, "gray")
    #VENUS
    penup()
    setpos(-300,10)
    pendown()
    drawPlanet(0.7,"DarkGoldenrod")
    #TERRA
    penup()
    setpos(-170,-10)
    pendown()
    drawPlanet(1.1,"LightSkyBlue")
    #MARTE
    setpos(-25,0)
    pendown()
    drawPlanet(0.8,"DarkRed")
    #JUPITE
    penup()
    setpos(130,-30)
    pendown()
    drawPlanet(1.5, "Khaki")
    #SATURNO
    penup()
    setpos(300,0)
    pendown()
    drawPlanet(0.9,"PaleGoldenrod")
    #URANO
    penup()
    setpos(410,20)
    pendown()
    drawPlanet(0.6, "blue")
    #NETUNO
    penup()
    setpos(500,25)
    pendown()
    drawPlanet(0.4,"DeepSkyBlue")

hideturtle()
done()


# 000 001 000 010 101 0000 0011 0101 0100 1011 0001 0100 1001  01000 01000 01110 01010 01101 01100 01000
#  a   b   a   c   ab   a   d     ab    e   da   b    e   aba    ca     ca  dab     ad   ed   abe    ca

# dict
# a=000
# b=001
# c=010
# d=011
# e=100

# ab=101
# ba=110
# ac=111

# ca=1000
# aba=1001
# ad=1010
# da=1011
# abe=1100
# ed=1101
# dab=1110
# be=1111

# ea=10000
# abac=10001
# cac=10010
# cad=10011
# daba=10100
# ade=10101
# eda=10111
# abec=11000

# abacabadabedabeabacacadabadedabeca

############################################################

# reading file
# (encoding)
def read_file(file_name):
    try:
        file=open(file_name, 'r')
        return file.read()
    except:
        pass

# creating initial dict
# (encoding)
def init_dict(symbols):
    dictionary=[]
    for i in symbols:
        if i not in dictionary:
            dictionary.append(i)
    return dictionary

# get dictionary index of sequence
# convert dictionary index of sequence to bin code
# добавить нули в начало
# (encoding)
def get_index(dictionary, n):
    # количество цифр в индексе последовательности n
    current_index=str(bin(dictionary.index(n)).replace("0b", ""))
    # current_index = str(bin(dictionary.index(n)))

    # добавлять нули в начало, пока длина индекса не равна количеству цифр в самом большом индексе
    while len(current_index)<len(bin(len(dictionary)-1).replace("0b", "")):
    # while len(current_index) < len(bin(len(dictionary) - 1)):
        current_index='0'+current_index
    return current_index

# find max length
# (encoding)
def max_length(dictionary):
    l=0
    for i in dictionary:
        if len(i)>l:
            l=len(i)
    return l

# найти самый длинный код в словаре
# (decoding)
def max_index(dictionary):
    return len(bin(len(dictionary)).replace("0b", ""))

# get sequence from dictionary by index
# (decoding)
def get_sequence(dictionary, bin_ind):
    # convert bin code to decimal index
    # bin_ind=int(bin(int(bin_ind)), 2)
    ind=int('0b'+bin_ind, 2)
    # print('get_sequence: ', ind)

    return dictionary[ind]

### encoding sequence of symbols
def encode(file):
    # создание начального словаря
    dictionary=init_dict(file)

    # зашифрованный файл
    code = ''

    # предыдущая последовательность
    last_seq=''

    # длина самой большой последовательности из словаря
    max_seq=max_length(dictionary)

    # steps
    step=0

    ### пока не закончится файл
    while file!='':
        # steps
        print('step № '+str(step))

        # из файла берется последовательность символов,
        # длина которой равна длине самой большой последовательности в словаре
        current_seq=file[:max_seq]
        print('sequence: '+current_seq)

        # пока последовательности нет в словаре укорачивается на 1
        while current_seq not in dictionary:
            current_seq=current_seq[:-1]

        # добавление новой последовательности (предыдущая+первый символ текущей)
        if (last_seq+current_seq[0]) not in dictionary:
            dictionary.append(last_seq + current_seq[0])
            # длина самой большой записи в словаре
            max_seq = max_length(dictionary)

        # добавляем к коду индекс последовательности
        code+=get_index(dictionary, current_seq)

        # print
        print('current sequence: ' + current_seq)
        print('index: '+get_index(dictionary, current_seq))
        print('code: '+code)
        print('file: '+file)
        print('dictionary: '+str(dictionary))

        # обновляется предыдущая последовательность символов
        last_seq = current_seq
        # из начала файла удаляются зашифрованные символы
        file = file[len(current_seq):]

        # steps
        step+=1
        print('\n')

    return code

### decoding
def decode(code, dictionary):
    # расшифрованный файл
    file=''

    # предыдущая последовательность
    last_seq = ''

    # длина последнего кода в словаре
    max_num=max_index(dictionary)

    # steps
    step=0

    # пока не закончится зашифрованный файл
    while code!='':
        # steps
        print('step № '+str(step))

        # из зашифрованного файла берется часть кода,
        # длина которой равна длине последнего кода в словаре
        current_code=code[:max_num]
        print('index: ' + str(current_code))

        # get sequence from the dictionary by code
        flag=False
        while flag==False:
            try:
                current_seq=get_sequence(dictionary, current_code)
                flag=True
            except:
                current_code = current_code[:-1]

        # add new sequence to the dictionary
        if (last_seq+current_seq[0]) not in dictionary:
            dictionary.append(last_seq + current_seq[0])
            # длина самой большой записи в словаре
            max_num=max_index(dictionary)

        # добавляем к расшифрованному файлу последовательность символов
        file+=current_seq

        # print
        print('current sequence: ' + current_seq)
        print('current index: ' + str(current_code))
        print('code: '+code)
        print('file: '+file)
        print('dictionary: '+str(dictionary))

        # обновляется предыдущая последовательность символов
        last_seq = current_seq                                  ############      !!!!!!!!!!!! last seq
        # из начала файла удаляются зашифрованные символы
        # code = code[max_num + 1:-1]
        code = code[len(str(current_code)):]

        # steps
        step+=1
        print('\n')

    return file

#### MAIN ####

# чтение файла
file=read_file('input_file')

# кодирование файла
code=encode(file)

print('ENCODED FILE: '+code+'\n')

# расшифровка кода
dictionary=init_dict(file)
decoded_file=decode(code, dictionary)

print('DECODED CODE: '+decoded_file+'\n')

# daecbeadcebadbcedabcedeebbabedbebbeabababaedbcdeabcdeabcedacbedabcedebacbabedcbea
# decoded:
# daecbeadcebadbcedabcedeebbabedbebbeabababaedbcdeabcdeabcedacbedabcedebacbabedcbea
#
# dictionary: ['d', 'a', 'e', 'c', 'b', 'de', 'ea', 'ae', 'ed', 'eb', 'be',
#               'ede', 'ebe', 'edb', 'ba', 'ab', 'bac', 'ca', 'abe', 'eae']

# abcdedf
# ['a', 'b', 'c', 'd', 'e', 'f', 'ab', 'bd', 'de', 'ef', 'ff']
#  ['a', 'b', 'c', 'd', 'e', 'f', 'ab', 'bd', 'dc', 'cf', 'ff']


## - YOLO Model Functions

# Define class size range based on model type
def class_range(model_type):
    if model_type == 'species':
        # testing smaller numbers w data augmentation, speed up training time, train more species
        max_per_category = 2100
        min_per_category = 300
    if model_type == 'family':
        max_per_category = 2000
        min_per_category = 200
    if model_type == 'general':
        max_per_category = 30000
        min_per_category = 5000
    if model_type == 'pig_only':
        max_per_category = 1000
        min_per_category = 300
    return max_per_category, min_per_category

# data wrangling
def wrangle_df(df, IMAGE_ROOT):
    # confirm bounding boxes are listed as proportions
    assert df['YMax'].max() <= 1
    assert df['XMax'].max() <= 1

    # fix special characters in file names
    # 07.22.22 AB: uncomment this line on CURC, comment out when local
    # df['filename'] = df['filename'].str.replace("'","").replace(" ","")
    # TODO: ask Ryan to remove special characters from folder names

    # cross check filenames to image files in directory
    # for local jobs, need to be connected to VPN here
    extant = [os.path.join(dp, f).replace(os.sep, '/') for dp, dn, fn in os.walk(IMAGE_ROOT) for f in fn]
    extant = [x.replace(IMAGE_ROOT + '/', '/').strip('/') for x in extant]

    # filter df for existing images
    df = df[df['filename'].isin(extant)]

    # exclude partial images
    df = df[df['partial.image'] == False]

    # exclude species with poor annotations
    # TODO: review these images for updated bboxes
    drops = ['Wood_Stork', 'Sandhill_Crane', 'Great_Egret',
             'Great_Blue_Heron', 'Crested_Caracara', 'Black-crowned_Night_Heron', 'Caribou']
    df = df.loc[~df['common.name'].isin(drops)].reset_index(drop=True)

    # swap y-axis for images where bboxes originate in LL corner (pytorch looks for UL and LR corners)
    df['YMin_org'] = df['YMin']
    df['YMax_org'] = df['YMax']
    df.drop(['YMax', 'YMin'], axis=1, inplace=True)
    df['YMax'] = np.where(df['bbox.origin'] == 'LL', (1 - df['YMin_org']), df['YMax_org'])
    df['YMin'] = np.where(df['bbox.origin'] == 'LL', (1 - df['YMax_org']), df['YMin_org'])

    # check for bboxes that are too small
    df['YDiff'] = df['YMax'] - df['YMin']
    df['XDiff'] = df['XMax'] - df['XMin']
    df_tooSmall = df[(df['YDiff'] < 0.001) | (df['XDiff'] < 0.001)]
    assert df_tooSmall.shape[0] == 0, "Remove image files where bboxes are too small"

    # create new columns for YOLO bboxes
    # Note: this formula applies to *normalized* bbox coordinates!
    df['x_center'] = (df['XMin'] + df['XMax']) / 2
    df['y_center'] = (df['YMin'] + df['YMax']) / 2
    df['w_bbox'] = df['XMax'] - df['XMin']
    df['h_bbox'] = df['YMax'] - df['YMin']

    # update column names for common name
    df['common.name_org'] = df['common.name']
    df['common.name'] = df['common.name.general']

    # combine squirrel species into one group
    squirrels = ["Aberts_Squirrel", 'American_Red_Squirrel', 'Douglas_Squirrel', 'Eastern_Fox_Squirrel',
                 'Eastern_Gray_Squirrel', 'Fox_Squirrel', 'Golden-Mantled_Ground_Squirrel', 'Gray_Squirrel',
                 'Northern_Flying_Squirrel', 'Rock_Squirrel', 'Squirrel', 'Western_Gray_Squirrel']
    df.loc[df['common.name'].isin(squirrels), 'common.name'] = 'squirrel_spp'

    # change the taxonomic classifications for vehicle
    df.loc[df['common.name'] == 'Vehicle', ['genus', 'species', 'family', 'order', 'class']] = 'vehicle'

    # add category for general model
    conditions = [(df['class'] == 'Mammalia') & (df['common.name'] != "Human") & (df['common.name'] != "Vehicle"),
                  (df['class'] == "Aves"),
                  (df['common.name'] == 'Human'),
                  (df['common.name'] == "Vehicle")]
    choices = ['mammal', 'bird', 'human', 'vehicle']
    df['general_category'] = np.select(conditions, choices, default=np.NAN)
    return df

# Create PyTorch dataset
class DetectDataset(torch.utils.data.Dataset):
    def __init__(self, df, w, h, image_dir, transform):
        self.image_dir = image_dir
        self.df = df
        self.image_infos = df.filename.unique()
        self.w = w
        self.h = h
        self.transform = transform

    def __getitem__(self, item):
        # create image id
        image_id = self.image_infos[item]
        # create full path to open each image file
        img_path = os.path.join(self.image_dir, image_id).replace("\\", "/")
        # open image, convert to numpy array
        # resize here so bboxes can also be resized before transformations
        img = Image.open(img_path).convert("RGB").resize((self.w, self.h), resample=Image.Resampling.BILINEAR)
        img = np.array(img, dtype="float32") / 255.
        # extract bbox and label data for image
        data = df[df['filename'] == image_id]
        # extract label names
        labels = data['LabelName'].values.tolist()
        # extract bbox coordinates
        data = data[['XMin', 'YMin', 'XMax', 'YMax']].values
        # convert to absolute values for model input
        data[:, [0, 2]] *= self.w
        data[:, [1, 3]] *= self.h
        # convert coordinates to list
        boxes = data.tolist()
        # convert bboxes and labels to a tensor dictionary
        target = {
            'boxes': torch.tensor(boxes).float(),
            'labels': torch.tensor([label2target[i] for i in labels]).long()
        }
        if self.transform is not None:
            augmented = self.transform(image=img, bboxes=target['boxes'], labels=labels)
            img = augmented['image'].to(device).float()
            target['boxes'] = augmented['bboxes']
        target['boxes'] = torch.tensor(target['boxes']).float()  # for some reason, ToTensorV2() isn't working
        return img, target

    def collate_fn(self, batch):
        return tuple(zip(*batch))

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

# define data augmentation pipelines
train_transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.Affine(rotate=(-20, 20), fit_output=True, p=0.3),
    A.Affine(shear=(-20, 20), fit_output=True, p=0.3),
    A.RandomBrightnessContrast(brightness_by_max=True, p=0.3),
    A.RandomSizedBBoxSafeCrop(height=307, width=408, erosion_rate=0.2, p=0.5),
    ToTensorV2(),
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels'], min_visibility=0),
)
val_transform = A.Compose([
    ToTensorV2(),
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']),
)

# Plot images
COLORS = np.random.randint(0, 255, size=(80, 3), dtype="uint8")

# manually adjust score threshold?
def show_img_bbox(img, targets, score_threshold=0.7):
    # convert image format to PIL
    if torch.is_tensor(img):
        img = to_pil_image(img)
    elif isinstance(img, np.ndarray):
        img = Image.fromarray(img, 'RGB')
    # convert bboxes and labels to np array
    if torch.is_tensor(targets):
        boxes = targets.numpy()[:, 1:]
    else:
        # if using the input of `output`.
        if 'scores' in targets:
            boxes, scores, labels = decode_output(targets, labels_as_numbers=True)
        else:
            boxes = targets['boxes']
            labels = targets['labels']
            scores = [1] * len(boxes)
    # plot image
    draw = ImageDraw.Draw(img)
    # plot bboxes
    for i, tg in enumerate(boxes):
        if scores[i] > score_threshold:
            id_ = int(labels[i])
            bbox = boxes[i]
            xmin, ymin, xmax, ymax = bbox
            color = [int(c) for c in COLORS[id_]]
            name = target2label[id_]
            draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=tuple(color), width=3)
            draw.text((xmin, ymin), name, fill=(255, 255, 255, 0))
    plt.imshow(np.array(img))

# define model

# Define PyTorch data loaders
def get_dataloaders(train_df, train_ds, val_ds, model_type, num_classes, batch_size):
    # for pig model, oversample Suidae to give it weight equal to other classes combined
    if model_type == 'pig_only':
        # set up class weights
        s = dict(Counter(train_df['LabelName']))
        # weight all other classes equally by num of classes
        # note : prob a way to do this by samples as well, but ok for now
        s = {x: (1 / (2 * (num_classes - 2))) for x in s}
        # weight pig class by sum of all other samples
        s['Suidae'] = 0.5
        # sanity check weighting was done correctly
        assert sum(s.values()) == 1
        # assign sample weight to each image in the dataset
        sample_weights = train_df.LabelName.map(s).tolist()
        # create weighted random sampler
        sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True)
        # define data loaders
        train_loader = DataLoader(train_ds, batch_size=batch_size, collate_fn=train_ds.collate_fn, drop_last=True,
                                  sampler=sampler)
        # do not oversample for validation, just for training
        val_loader = DataLoader(val_ds, batch_size=batch_size, collate_fn=train_ds.collate_fn, drop_last=True)
    # No oversampling for other models at this time
    else:
        train_loader = DataLoader(train_ds, batch_size=batch_size, collate_fn=train_ds.collate_fn, shuffle=True,
                                  drop_last=True)
        val_loader = DataLoader(val_ds, batch_size=batch_size, collate_fn=train_ds.collate_fn, drop_last=True)
    return train_loader, val_loader

# obtain learning rate
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

# define checkpoint functions
def create_checkpoint(model, optimizer, epoch, lr_scheduler, loss_history, best_loss, model_type, num_classes,
                      label2target):
    checkpoint = {'state_dict': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  'epoch': epoch + 1,
                  'lr': current_lr,
                  'scheduler': lr_scheduler.state_dict(),
                  'loss_history': loss_history,
                  'best_loss': best_loss,
                  'model_type': model_type,
                  'num_classes': num_classes,
                  'label2target': label2target}
    return checkpoint

def save_checkpoint(checkpoint, checkpoint_file):
    print(" Saving model state")
    torch.save(checkpoint, checkpoint_file)

def load_checkpoint(checkpoint_path):
    print(" Loading saved model state")
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    lr_scheduler.load_state_dict(checkpoint['scheduler'])
    epoch = checkpoint['epoch']
    loss_history = checkpoint['loss_history']
    best_loss = checkpoint['best_loss']
    model_type = checkpoint['model_type']
    label2target = checkpoint['label2target']
    return model, optimizer, lr_scheduler, epoch, loss_history, best_loss, model_type, label2target

# make predictions
def decode_output(output, labels_as_numbers=False):
    bbs = output['boxes'].cpu().detach().numpy().astype(np.uint16)
    if labels_as_numbers:
        labels = np.array(output['labels'].cpu().detach().numpy())
    else:
        labels = np.array([target2label[i] for i in output['labels'].cpu().detach().numpy()])
    confs = output['scores'].cpu().detach().numpy()
    ixs = nms(torch.tensor(bbs.astype(np.float32)), torch.tensor(confs), 0.5)
    bbs, confs, labels = [tensor[ixs] for tensor in [bbs, confs, labels]]
    if len(ixs) == 1:
        bbs, confs, labels = [np.array([tensor]) for tensor in [bbs, confs, labels]]
    return bbs.tolist(), confs.tolist(), labels.tolist()

# deploy function
def deploy(df, sample=False, w=408, h=307):
    model.eval()
    model.to(device)
    image_infos = df.filename.unique()
    if sample:
        image_infos = random.sample(list(image_infos), 10)
    gt_df = []
    pred_df = []
    for i in tqdm(range(len(image_infos))):
        dfi = df[df['filename'] == image_infos[i]]
        dsi = DetectDataset(df=dfi, image_dir=IMAGE_ROOT, w=w, h=h, transform=val_transform)
        dli = DataLoader(dsi, batch_size=1, collate_fn=dsi.collate_fn, drop_last=True)
        input, target = next(iter(dli))
        image = list(image.to(device) for image in input)
        output = model(image)[0]
        bbs, confs, labels = decode_output(output, labels_as_numbers=True)
        boxes = bbs
        if len(bbs) == 0:
            pred_df_i = pd.DataFrame({
                'filename': image_infos[i],
                'file_id': image_infos[i][:-4],
                'class_name': 'empty',
                'confidence': 1,
                'bbox': [0, 0, w, h]
            })
        else:
            pred_df_i = pd.DataFrame({
                'filename': image_infos[i],
                'file_id': image_infos[i][:-4],
                'class_name': [target2label[a] for a in labels],
                'confidence': confs,
                'bbox': bbs
            })
        gt_df_i = pd.DataFrame({
            'filename': image_infos[i],
            'file_id': image_infos[i][:-4],
            'class_name': [target2label[a] for a in dsi[0][1]['labels'].tolist()],
            'bbox': dsi[0][1]['boxes'].tolist(),
            'used': False
        })
        gt_df_i = gt_df_i.join(pd.DataFrame(dsi[0][1]['boxes'].tolist()))
        gt_df.append(gt_df_i)
        pred_df.append(pred_df_i)
    gt_df = pd.concat(gt_df)
    pred_df = pd.concat(pred_df)
    gt_df.to_csv(eval_path + "gt_df.csv")
    pred_df.to_csv(eval_path + "pred_df.csv")
    return gt_df, pred_df

# define intersection over union function
def IoU(bbs, tbs):
    """
    Calculates intersection over union
    :param bbs: set of prediction bounding boxes for an image
    :param tbs: set of true target bounding boxes for an image
    :return: intersection over union
    """
    # shape is (N,4) where N is number of bboxes per image
    # Convert coordinates from (x,y,w,h) to (x1, y1, x2, y2)
    pred_xmin = bbs[..., 0:1] - bbs[..., 2:3] / 2
    pred_ymin = bbs[..., 1:2] - bbs[..., 3:4] / 2
    pred_xmax = bbs[..., 2:3] + bbs[..., 2:3] / 2
    pred_ymax = bbs[..., 3:4] + bbs[..., 3:4] / 2
    target_xmin = tbs[..., 0:1] - tbs[..., 2:3] / 2
    target_ymin = tbs[..., 1:2] - tbs[..., 3:4] / 2
    target_xmax = tbs[..., 2:3] + tbs[..., 2:3] / 2
    target_ymax = tbs[..., 3:4] + tbs[..., 3:4] / 2

    # find area of each box
    pred_area = abs((pred_xmax - pred_xmin) * (pred_ymax - pred_ymin))
    target_area = abs((target_xmax - target_xmin) * (target_ymax - target_ymin))

    # find intersection area
    xmin = torch.max(pred_xmin, target_xmin)
    ymin = torch.max(pred_ymin, target_ymin)
    xmax = torch.min(pred_xmax, target_xmax)
    ymax = torch.min(pred_ymax, target_ymax)
    intersect = (xmax - xmin).clamp(0) * (ymax - ymin).clamp(0)

    # find area of union
    union = pred_area + target_area - intersect + 1e-6  # add numeric stabilizer in case union = 0

    return intersect / union


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import tensorflow as tf

import resnet_preprocessing

class PBRSceneNetDepthInput(object):
    """Generates PBRNet and Scenenet data input_pbrscenent for training or evaluation.

    The training data is assumed to be in TFRecord format with keys as specified
    in the dataset_parser below, sharded across 366 files, named sequentially:
        train-0
        train-1
        ...
        train-365

    The validation data is in the same format but sharded in 50 files.

    Args:
      is_training: `bool` for whether the input is for training
      data_dir: `str` for the directory of the training and validation data
      num_cores: `int` for the number of TPU cores
    """

    def __init__(self, is_training, pbr_dir, scene_dir, num_cores=8):
        # self.image_preprocessing_fn = resnet_preprocessing.full_imagenet_augment
        self.image_preprocessing_fn = resnet_preprocessing.depth_image_preprocess_image
        self.depth_preprocessing_fn = resnet_preprocessing.depth_preprocess_image
        self.is_training = is_training
        self.pbr_dir = pbr_dir
        self.scene_dir = scene_dir
        self.num_cores = num_cores
        self.training_data_num = 366 + 843

    def dataset_parser(self, value):
        """Parse an ImageNet record from a serialized string Tensor."""
        keys_to_features = {
            'mlt': tf.FixedLenFeature((), tf.string, ''),
            'depth': tf.FixedLenFeature((), tf.string, '')
        }

        parsed = tf.parse_single_example(value, keys_to_features)
        print("parsed example", parsed)

        original_image = tf.reshape(parsed['mlt'], shape=[])
        original_image = tf.image.decode_png(original_image, channels=3)
        original_image = tf.image.convert_image_dtype(original_image, dtype=tf.float32)
        original_image.set_shape((240, 320, 3))

        depth_image = tf.reshape(parsed['depth'], shape=[])
        depth_image = tf.image.decode_png(depth_image, dtype=tf.uint16)
        depth_image.set_shape((240, 320, 1))
        depth_image = tf.cast(depth_image, tf.int32)

        original_image, depth_image = self.image_preprocessing_fn(
            original_image=original_image,
            depth_image=depth_image,
            is_training=self.is_training,
            image_height=240,
            image_width=320,
        )
        depth_image = self.depth_preprocessing_fn(
            image=depth_image,
        )
        return original_image, depth_image

    def input_fn(self, params):
        """Input function which provides a single batch for train or eval.

        Args:
          params: `dict` of parameters passed from the `TPUEstimator`.
              `params['batch_size']` is always provided and should be used as the
              effective batch size.

        Returns:
          A (images, labels) tuple of `Tensor`s for a batch of samples.
        """
        # Retrieves the batch size for the current shard. The # of shards is
        # computed according to the input pipeline deployment. See
        # tf.contrib.tpu.RunConfig for details.

        batch_size = params['batch_size']

        # Shuffle the filenames to ensure better randomization.
        pbr_file_pattern = os.path.join(
            self.pbr_dir, 'train-*' if self.is_training else 'validation-*')
        dataset = tf.data.Dataset.list_files(pbr_file_pattern)

        print("dataset", dataset)

        scene_file_pattern = os.path.join(
            self.scene_dir, 'train-*' if self.is_training else 'validation-*')
        dataset_tmp = tf.data.Dataset.list_files(scene_file_pattern)

        print("dataset_tmp", dataset_tmp)

        dataset.concatenate(dataset_tmp)

        print("dataset", dataset)

        if self.is_training:
            dataset = dataset.shuffle(buffer_size=self.training_data_num)  # 1024 files in dataset
            dataset = dataset.repeat()

        def fetch_dataset(filename):
            buffer_size = 1024 * 1024 * 1024  # 1024 MiB per file
            dataset = tf.data.TFRecordDataset(filename, buffer_size=buffer_size)
            return dataset

        dataset = dataset.apply(
            tf.contrib.data.parallel_interleave(
                fetch_dataset, cycle_length=self.num_cores, sloppy=True))

        dataset = dataset.shuffle(batch_size * 20)

        print("before parsing", dataset)
        dataset = dataset.map(
            self.dataset_parser,
            num_parallel_calls=64)
        print("after parsing", dataset)
        dataset = dataset.prefetch(batch_size)

        # For training, batch as usual. When evaluating, prevent accidentally
        # evaluating the same image twice by dropping the final batch if it is less
        # than a full batch size. As long as this validation is done with
        # consistent batch size, exactly the same images will be used.
        dataset = dataset.apply(
            tf.contrib.data.batch_and_drop_remainder(batch_size))

        dataset = dataset.prefetch(2)  # Prefetch overlaps in-feed with training

        original_images, depth_images = dataset.make_one_shot_iterator().get_next()
        print("original image and depth type:")
        print(original_images.shape)
        print(depth_images.shape)

        return original_images, depth_images

    def postproc_flag(self, images,
                      NOW_SIZE1=256,
                      NOW_SIZE2=256,
                      seed_random=0,
                      curr_batch_size=None,
                      with_noise=0,
                      noise_level=10,
                      with_flip=0,
                      is_normal=0,
                      eliprep=0,
                      thprep=0,
                      sub_mean=0,
                      mean_path=None,
                      color_norm=0,
                      size_vary_prep=0,
                      with_color_noise=0,
                      shape_undefined=0,
                      size_minval=0.08,
                      sm_full_size=0,  # sm_add
                      ):

        if curr_batch_size == None:
            curr_batch_size = self.batch_size

        orig_dtype = images.dtype
        norm = tf.cast(images, tf.float32)

        if eliprep == 1:
            def prep_each(norm_):
                _RESIZE_SIDE_MIN = 256
                _RESIZE_SIDE_MAX = 512

                if self.group == 'train':
                    im = preprocess_image(norm_, self.crop_size, self.crop_size, is_training=True,
                                          resize_side_min=_RESIZE_SIDE_MIN,
                                          resize_side_max=_RESIZE_SIDE_MAX)
                else:
                    im = preprocess_image(norm_, self.crop_size, self.crop_size, is_training=False,
                                          resize_side_min=_RESIZE_SIDE_MIN,
                                          resize_side_max=_RESIZE_SIDE_MAX)

                return im

            crop_images = tf.map_fn(prep_each, norm)
        elif thprep == 1:
            def prep_each(norm_):
                im = preprocessing_th(norm_, self.crop_size, self.crop_size,
                                      is_training=self.group == 'train', seed_random=seed_random)
                return im

            crop_images = prep_each(images)
            crop_images = tf.expand_dims(crop_images, axis=0)
        elif sm_full_size == 1:
            if with_color_noise == 1 and self.group == 'train':
                order_temp = tf.constant([0, 1, 2], dtype=tf.int32)
                order_rand = tf.random_shuffle(order_temp, seed=seed_random)

                fn_pred_fn_pairs = lambda x, image: [
                    (tf.equal(x, order_temp[0]),
                     lambda: tf.image.random_saturation(image, 0.6, 1.4, seed=seed_random)),
                    (tf.equal(x, order_temp[1]), lambda: tf.image.random_brightness(image, 0.4, seed=seed_random)),
                ]
                default_fn = lambda image: tf.image.random_contrast(image, 0.6, 1.4, seed=seed_random)

                def _color_jitter_one(_norm):
                    orig_shape = _norm.get_shape().as_list()
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[0], _norm), default=lambda: default_fn(_norm))
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[1], _norm), default=lambda: default_fn(_norm))
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[2], _norm), default=lambda: default_fn(_norm))
                    _norm.set_shape(orig_shape)

                    return _norm

                norm = tf.map_fn(_color_jitter_one, norm)

            if sub_mean == 1:
                IMAGENET_MEAN = tf.constant(np.load(mean_path).swapaxes(0, 1).swapaxes(1, 2)[:, :, ::-1],
                                            dtype=tf.float32)
                orig_dtype = tf.float32
                norm = norm - IMAGENET_MEAN

            if self.group == 'train':
                if self.withflip == 1 or with_flip == 1:
                    def _postprocess_flip(im):
                        # Original way of flipping, changing to random_uniform way to be more controllable
                        # im = tf.image.random_flip_left_right(im, seed = seed_random)
                        # return im
                        do_flip = tf.random_uniform(shape=[1], minval=0, maxval=1, dtype=tf.float32,
                                                    seed=seed_random)

                        def __left_right_flip(im):
                            flipped = tf.image.flip_left_right(im)
                            if is_normal == 1:
                                # flipped = 256 - flipped
                                flipped_x, flipped_y, flipped_z = tf.unstack(flipped, axis=2)
                                flipped = tf.stack([256 - flipped_x, flipped_y, flipped_z], axis=2)
                            return flipped

                        return tf.cond(tf.less(do_flip[0], 0.5), fn1=lambda: __left_right_flip(im), fn2=lambda: im)

                    norm = tf.map_fn(_postprocess_flip, norm, dtype=norm.dtype)

                if with_noise == 1:
                    def _postprocess_noise(im):
                        do_noise = tf.random_uniform(shape=[1], minval=0, maxval=1, dtype=tf.float32, seed=None)

                        def __add_noise(im):
                            curr_level = tf.random_uniform(shape=[1], minval=0, maxval=noise_level,
                                                           dtype=tf.float32,
                                                           seed=None)
                            curr_noise = tf.random_normal(shape=tf.shape(im), mean=0.0, stddev=curr_level,
                                                          dtype=tf.float32)

                            return tf.add(im, curr_noise)

                        # return tf.cond(tf.less(do_noise[0], 0.5), true_fn = lambda: __add_noise(im), false_fn = lambda: im)
                        return tf.cond(tf.less(do_noise[0], 0.5), fn1=lambda: __add_noise(im), fn2=lambda: im)

                    norm = tf.map_fn(_postprocess_noise, norm, dtype=norm.dtype)
            crop_images = tf.cast(norm, orig_dtype)
        else:
            if with_color_noise == 1 and self.group == 'train':
                order_temp = tf.constant([0, 1, 2], dtype=tf.int32)
                order_rand = tf.random_shuffle(order_temp, seed=seed_random)

                fn_pred_fn_pairs = lambda x, image: [
                    (tf.equal(x, order_temp[0]),
                     lambda: tf.image.random_saturation(image, 0.6, 1.4, seed=seed_random)),
                    (tf.equal(x, order_temp[1]), lambda: tf.image.random_brightness(image, 0.4, seed=seed_random)),
                ]
                default_fn = lambda image: tf.image.random_contrast(image, 0.6, 1.4, seed=seed_random)

                def _color_jitter_one(_norm):
                    orig_shape = _norm.get_shape().as_list()
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[0], _norm), default=lambda: default_fn(_norm))
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[1], _norm), default=lambda: default_fn(_norm))
                    _norm = tf.case(fn_pred_fn_pairs(order_rand[2], _norm), default=lambda: default_fn(_norm))
                    _norm.set_shape(orig_shape)

                    return _norm

                norm = tf.map_fn(_color_jitter_one, norm)

            if sub_mean == 1:
                IMAGENET_MEAN = tf.constant(np.load(mean_path).swapaxes(0, 1).swapaxes(1, 2)[:, :, ::-1],
                                            dtype=tf.float32)
                orig_dtype = tf.float32
                norm = norm - IMAGENET_MEAN
            if self.group == 'train':

                if self.size_vary_prep == 0 and size_vary_prep == 0:
                    shape_tensor = norm.get_shape().as_list()
                    if self.crop_each == 0:
                        crop_images = tf.random_crop(norm, [curr_batch_size, self.crop_size, self.crop_size,
                                                            shape_tensor[3]], seed=seed_random)
                    else:
                        # original implementation is not useful, deleted, see the end of this file
                        crop_images = tf.random_crop(norm, [curr_batch_size, self.crop_size, self.crop_size,
                                                            shape_tensor[3]], seed=seed_random)
                else:  # self.size_vary_prep==1
                    if shape_undefined == 0:
                        channel_num = norm.get_shape().as_list()[-1]
                    else:
                        channel_num = 3
                    RandomSizedCrop_with_para = lambda image: RandomSizedCrop(
                        image=image,
                        out_height=self.crop_size,
                        out_width=self.crop_size,
                        seed_random=seed_random,
                        channel_num=channel_num,
                        fix_asp_ratio=self.fix_asp_ratio,
                        size_minval=size_minval,
                    )
                    if shape_undefined == 0:
                        crop_images = tf.map_fn(RandomSizedCrop_with_para, norm)
                        curr_shape = crop_images.get_shape().as_list()
                        crop_images.set_shape([curr_batch_size] + curr_shape[1:])
                    else:
                        crop_images = RandomSizedCrop_with_para(norm)
                        crop_images = tf.expand_dims(crop_images, axis=0)

                if self.withflip == 1 or with_flip == 1:
                    def _postprocess_flip(im):
                        # Original way of flipping, changing to random_uniform way to be more controllable
                        # im = tf.image.random_flip_left_right(im, seed = seed_random)
                        # return im
                        do_flip = tf.random_uniform(shape=[1], minval=0, maxval=1, dtype=tf.float32,
                                                    seed=seed_random)

                        def __left_right_flip(im):
                            flipped = tf.image.flip_left_right(im)
                            if is_normal == 1:
                                # flipped = 256 - flipped
                                flipped_x, flipped_y, flipped_z = tf.unstack(flipped, axis=2)
                                flipped = tf.stack([256 - flipped_x, flipped_y, flipped_z], axis=2)
                            return flipped

                        return tf.cond(tf.less(do_flip[0], 0.5), fn1=lambda: __left_right_flip(im), fn2=lambda: im)

                    crop_images = tf.map_fn(_postprocess_flip, crop_images, dtype=crop_images.dtype)

                if with_noise == 1:
                    def _postprocess_noise(im):
                        do_noise = tf.random_uniform(shape=[1], minval=0, maxval=1, dtype=tf.float32, seed=None)

                        def __add_noise(im):
                            curr_level = tf.random_uniform(shape=[1], minval=0, maxval=noise_level,
                                                           dtype=tf.float32, seed=None)
                            curr_noise = tf.random_normal(shape=tf.shape(im), mean=0.0, stddev=curr_level,
                                                          dtype=tf.float32)

                            return tf.add(im, curr_noise)

                        # return tf.cond(tf.less(do_noise[0], 0.5), true_fn = lambda: __add_noise(im), false_fn = lambda: im)
                        return tf.cond(tf.less(do_noise[0], 0.5), fn1=lambda: __add_noise(im), fn2=lambda: im)

                    crop_images = tf.map_fn(_postprocess_noise, crop_images, dtype=crop_images.dtype)

            else:  # not self.group=='train'

                if shape_undefined == 0:
                    off = np.zeros(shape=[curr_batch_size, 4])
                    off[:, 0] = int((NOW_SIZE1 - self.crop_size) / 2)
                    off[:, 1] = int((NOW_SIZE2 - self.crop_size) / 2)
                    off[:, 2:4] = off[:, :2] + self.crop_size
                    off[:, 0] = off[:, 0] * 1.0 / (NOW_SIZE1 - 1)
                    off[:, 2] = off[:, 2] * 1.0 / (NOW_SIZE1 - 1)

                    off[:, 1] = off[:, 1] * 1.0 / (NOW_SIZE2 - 1)
                    off[:, 3] = off[:, 3] * 1.0 / (NOW_SIZE2 - 1)

                    box_ind = tf.constant(range(curr_batch_size))

                    crop_images = tf.image.crop_and_resize(norm, off, box_ind,
                                                           tf.constant([self.crop_size, self.crop_size]))
                else:
                    image = _aspect_preserving_resize(norm, 256)
                    image = _central_crop([image], self.crop_size, self.crop_size)[0]
                    image.set_shape([self.crop_size, self.crop_size, 3])
                    crop_images = image
                    crop_images = tf.expand_dims(crop_images, axis=0)

            crop_images = tf.cast(crop_images, orig_dtype)
        if curr_batch_size == 1:
            crop_images = tf.squeeze(crop_images, axis=[0])

        return crop_images


#!/usr/local/bin/python3

import sys;
import time;
import asyncio;
import os;
import pickle;

sys.path.append('/usr/local/lib/python3.9/site-packages');

from kasa import Discover, SmartPlug;

CACHE_PATH = '/tmp/automation';
CACHE_FILE = CACHE_PATH + '/smart_plug_devices';

#DEVICES_NOT = [ 'Pluggy' ];
#DEVICES_NOT = [ 'Smuggy' ];
DEVICES_NOT = [];

DEVICE_CACHE = [];
PLUG_CACHE = {};

async def load_into_cache():
   global DEVICE_CACHE;
   if not os.path.exists(CACHE_PATH):
      os.mkdir(CACHE_PATH);

   with open(CACHE_FILE, 'w+') as f:
      devices = await Discover.discover();
      for ip_addr in devices:
         DEVICE_CACHE.append(ip_addr);
         f.write(f'{ip_addr}\n');

   return None;

async def load_from_cache():
   global DEVICE_CACHE;
   with open(CACHE_FILE, 'r') as f:
      DEVICE_CACHE = f.read().splitlines();

   return None;

async def main(cmd):
   if not DEVICE_CACHE:
      if os.path.exists(CACHE_FILE):
         await load_from_cache()
      else:
         await load_into_cache();

   #sys.exit();
   cmd = cmd.upper();

   if cmd not in ['ON', 'OFF']:
      print('Error: Unknown command!');
      return False;

   tasks = [];
   for ip_addr in DEVICE_CACHE:
      # create concurrent tasks
      task = asyncio.create_task(operate(ip_addr, cmd));
      tasks.append(task);

   # run all the tasks concurrently now
   for task in tasks:
      await task;

async def operate(ip_addr, cmd):
   #print(f'{time.strftime("%X")}');

   if ip_addr not in PLUG_CACHE:
      plug = SmartPlug(ip_addr);
      await plug.update();
      PLUG_CACHE[ip_addr] = plug;

   plug = PLUG_CACHE[ip_addr];

   alias = plug.alias;
   if alias in DEVICES_NOT:
      #print(f'skipping {alias} ...');
      return None;

   sys_cmd = None;

   if cmd == 'ON':
      sys_cmd = plug.turn_on;
   elif cmd == 'OFF':
      sys_cmd = plug.turn_off;

   #print(f'{alias} going { cmd } ...', end='');
   await sys_cmd();
   #print('done');
   #await asyncio.sleep(5);

if __name__ == '__main__':
   command = None;
   if len(sys.argv) == 2:
      cmd = sys.argv[1];
      asyncio.run(main(cmd));
   else:
      print(f'\nUsage: {sys.argv[0]} <on|off>\n');

   exit(0);


class ParkingGarage:
    def __init__(self):
        self.availableTicket = [1,1,1,1,1,1,1,1,1,1]
        self.parkingSpaces = [1,1,1,1,1,1,1,1,1,1]
        self.currentTicket = {
            "paid" : False
        }

    def kiosk(self):
        self.park = input("\nWelcome to our parking garage! \nOur fee is $20 for the day... \nWhat would you like to do? Park/Pay/Leave \n")
        self.park = self.park.strip().lower()
        if self.park == "park":
            garage.takeTicket()
        elif self.park == "pay":
            garage.payForParking()
        elif self.park == "leave":
            garage.leaveGarage()
        else:
            print("\nDid not receive a valid response.")
            garage.kiosk()

    # Decrease available tickets and Parking spaces by one
    def takeTicket(self):
        print("\nNew ticket has been taken!")
        self.availableTicket.pop()
        self.parkingSpaces.pop()

        self.availableTicketLength = len(self.availableTicket)
        self.parkingSpacesLength = len(self.parkingSpaces)
        print("\nAvailable tickets: " + str(self.availableTicketLength) + " \nAvailable spaces: " + str(self.parkingSpacesLength))
        garage.kiosk()

    def payForParking(self):
        self.payment = input("\nReady to pay $20? Insert/Back ")

        # Check if payment variable is not empty
        if self.payment == "insert":
            print("\nYour ticket has been paid and you have 15 minutes to leave")
            self.payment = True
            self.currentTicket["paid"] = True
            print(self.currentTicket)
        elif self.payment == "back":
            garage.kiosk()
        else:
            print("\nDid not receive a valid response.")
            garage.payForParking()

    def leaveGarage(self):
        if bool(self.currentTicket["paid"]) == True:
            print("\nThank you, have a nice day!")

            self.parkingSpaces.append(1)
            self.availableTicket.append(1)

            self.availableTicketLength = len(self.availableTicket)
            self.parkingSpacesLength = len(self.parkingSpaces)

            print("\nAvailable tickets: " + str(self.availableTicketLength) + " \nAvailable spaces: " + str(self.parkingSpacesLength))
        else:
            print("\nYou must make a payment first...")
            garage.payForParking()
garage = ParkingGarage()
garage.kiosk()


import sys

import transformers

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

from transformers import BertTokenizer

from model.bert_dataloader import BertDataLoader
from model.bert_model import BertHotModel
from model.bert_process import BertProcess
from bean.arg_bean import BERTArgBean
from util.arg_parse import CustomArgParser
from util.log_util import CustomLogger

from accelerate import Accelerator

class BERTBase(object):
    def __init__(self, args):
        self.args = args
        self.model = BertHotModel(self.args)
        self.bert_tokenizer = BertTokenizer.from_pretrained(self.args.pretrain_model_path)
        self.accelerator = Accelerator()
        self.logger = CustomLogger.logger

        self.bert_dataloader = BertDataLoader(
            self.args,
            self.bert_tokenizer,
            self.logger
        )

        self.bert_process = BertProcess(
            self.args,
            self.model,
            self.bert_tokenizer,
            self.accelerator,
            self.logger
        )

    def train(self):
        # 加载数据
        train_dataloader, dev_dataloader = self.bert_dataloader.load_data(
            self.args.train_data_path, self.args.per_device_train_batch_size, is_train=True
        )
        print(train_dataloader)
        print(dev_dataloader)
        self.logger.info("Finished loading data !!!")

        # 初始化模型前固定种子，保证每次运行结果一致
        transformers.set_seed(self.args.seed)

        # 训练模型
        self.logger.info("Training model ...")
        self.bert_process.train(train_dataloader, dev_dataloader)
        self.logger.info("Finished Training model !!!")

    def test(self):
        test_dataloader = self.bert_dataloader.load_data(
            self.args.test_data_path, self.args.per_device_train_batch_size, is_train=False
        )
        # 测试模型
        self.logger.info("Testing model ...")
        self.bert_process.test(test_dataloader)
        self.logger.info("识别出热点事件数目：9386个")
        self.logger.info("Finished Testing model !!!")

if __name__ == '__main__':
    # 解析命令行参数
    args = CustomArgParser(BERTArgBean).parse_args_into_dataclass()
    bertbase = BERTBase(args)
    if args.do_train:
        bertbase.train()
    else:
        bertbase.test()


"""
Sherman Yan
Assignment 5 - P5.4
Input:  String
Output: Middle character(s)
"""

##
# Main function
# ask for user to input a word and output the middle characters of that input
#
def main():
    userInput = input("Please enter a word: ")

    print("The middle character(s) of \"" + userInput + " is: " + middle(userInput))

##
# This function calculates for the middle two characters of the inputted word
# @param word
# @return middle character(s) as string
#
def middle(word):
    length = len(word)

    middle_index = length // 2

    if length % 2 == 0:
        return str(word[middle_index - 1] + word[middle_index])
    else:
        return str(word[middle_index])

if __name__ == '__main__':
    main()

"""
Sample Run
Please enter a word: middle
The middle character(s) of "middle is: dd
"""


import typing as t
from dataclasses import dataclass

from docopt import docopt  # type: ignore

from glue.chain import Chain

@dataclass
class DocOptChain(Chain):
    name: str
    command_doc: str

class Generator:
    def __call__(self, chains: t.List[DocOptChain]) -> t.Dict[str, t.Callable]:
        commands = {}
        for chain in chains:
            command = create_command(chain)
            command_name = chain.name
            commands[command_name] = command
        return commands

def create_command(link: DocOptChain) -> t.Callable:
    def command(args):
        args = docopt(link.command_doc, argv=args)
        link.compose()(args)

    return command


import streamlit as st
import pandas as pd

DATA_URL = 'uber-raw-data-sep14.csv'
DATE_COLUMN ='Date/Time'

st.title("Activate  Caso UBER A01702048")
st.subheader("Desarrollar un código para una aplicación web que permita visualizar en un mapa de la ciudad de Nueva York los viajes realizados por la empresa Uber. ")
#filter

@st.cache
def load_data(nrows):
    data = pd.read_csv(DATA_URL, nrows=nrows)
    lowercase = lambda x: str(x).lower()
    data.rename(lowercase, axis='columns', inplace=True)
    data['date/time'] = pd.to_datetime(data['date/time'])
    data['time_hour'] = pd.to_datetime(data['date/time']).dt.hour

    return data

data = load_data(1000)

Horas = data['time_hour'].unique().tolist()
hora = st.selectbox('¿Que hora de viajes quieres ver?', Horas,0)
data = data[data['time_hour']==hora]

st.dataframe(data)
st.map(data)


#READ TXT FILE
# Path: p1\input.txt

with open('2022-01/input.txt', 'r') as file:
    lines = file.readlines()

#sacar los \n de cada linea
lines1 = [line.strip() for line in lines]

# definimos la lista de elfos
elfos = []
acumulado = 0
for i in lines1:

    if i:
        acumulado = acumulado + int(i)
    else:
        elfos.append(acumulado)
        acumulado = 0

# extract max of list and index
maximo = max(elfos)
index = elfos.index(maximo)
print(f'El elfo {index+1} és el que tiene mas calorias con: {maximo}')
elfos.sort(reverse=True)
print(f'Los 3 elfos con mas calorias tienen: {sum(elfos[:3])}')


#!/usr/bin/env python3
# -*- coding utf-8 -*-

from PIL import Image
from transformers import CLIPProcessor, CLIPModel

"""

"""

if __name__ == '__main__':
    # 加载clip-vit-base-patch32模型
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    # 加载了一张猫的图片
    image_file = "./data/cat.jpg"
    image =  Image.open(image_file)

    # 用文本标识分成四类
    categories = ["cat", "dog", "truck", "couch"]
    categories_text = list(map(lambda x: f"a photo of a {x}", categories))

    # 用softmax算法来计算图片分类到某一类的概率
    inputs = processor(text=categories_text, images=image, return_tensors="pt", padding=True)
    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)

    # 输出概率
    for i in range(len(categories)):
        print(f"{categories[i]}\t{probs[0][i].item():.2%}")

"""
输出：
cat  74.51%         # 图片中有两只猫
dog  0.39%
truck  0.04%
couch  25.07%       # 图片中有沙发
"""


import numpy as np
from astropy.io import fits

try:
  import matplotlib.pyplot as plt
  from matplotlib import colors
  mpl=True
except:
  print("Warning: matplotlib not found, plotting will not work...")
  mpl=False

try:
  from pyds9 import DS9
  pyds9=True
except:
  pyds9=False

fitter=None
fitters=[]
#Had issues getting iminuit to work, curve_fit works fine for now
try:
  import iminuit
  import probfit
  fitter="iminuit"
  fitters.append("iminuit")
  print("Using iminuit as fitter")
except:
  pass
try:
  from scipy.optimize import curve_fit
  if fitter is None:
    print("Iminuit not found, using scipy as fitter...")
    fitter="scipy"
    fitters.append("scipy")
except:
    print("Warning, no fitter found, cannot fit to noise...")

def mad(data):
  return np.median(np.abs(data-np.median(data)))

def plot_2d(data,cmap="rainbow",xlim=None, ylim=None, title=None, units="ADU", vmin=-50, vmax=None):
  if not mpl:
    return None, None
  data=np.array(data)
  fig=plt.figure()
  # mesh=plt.pcolormesh(data-np.min(data),cmap=plt.get_cmap(cmap), norm=colors.LogNorm())
  mesh=plt.pcolormesh(data,cmap=plt.get_cmap(cmap), norm=colors.SymLogNorm(10, vmin=vmin, vmax=vmax))

  ax=mesh.axes
  ax.set_aspect("equal")
  if xlim is not None:
    ax.set_xlim(right=xlim)
  if ylim is not None:
    ax.set_ylim(top=ylim)
  cbar=plt.colorbar()
  cbar.set_label(units)
  plt.xlabel("X"); plt.ylabel("Y");
  plt.title(title)
  plt.show(False)
  return fig, mesh

#For fitting
def gauss(x, norm,mean, sigma):
  return norm*np.exp(-(x-mean)**2/(2.*sigma**2))

#For fitting
def gauss_norm(x, mean, sigma):
  return 1./(sigma*(2*np.pi)**.5)*np.exp(-(x-mean)**2/(2.*sigma**2))

def fit_gauss(data, nbins=None, plot_fit=False, fitter=fitter, fit_range=None, rms_range_factor=8, *args, **kwargs):
  '''
  Wrapper function to perform binned fit of data to func. If given, will fit in the range "fit_range" (should be a tuple/array if [min,max]). Otherwise fits in the range median-rms_range_factor*MAD.
  '''
  if len(fitters) < 1:
    print("Error, no valid fitters loaded, exiting...")
    return None, None
  #Guess for initial parameters and to get a reasonable fitting range
  mean_guess=np.median(data)
  rms_guess=mad(data)
  if fit_range is None:
    fit_range=[mean_guess-rms_guess*rms_range_factor, mean_guess+rms_guess*rms_range_factor]
  fit_data=data[(data > fit_range[0]) & (data < fit_range[1])]

  #Get a reasonable number of bins
  if nbins is None:
    nbins=int((fit_range[1]-fit_range[0])/2)
    print("Bins for fit: " + str(nbins))
    if nbins < 40:
      print("Warning, number of bins is fairly small, fit may not be good...")

  if fitter not in fitters:
    print("Error, selected fitter is not loaded, loading another fitter")
    fitter=fitters[0]

  if fitter=="scipy":
    hist, bin_edges=np.histogram(fit_data, density=True, bins=nbins)
    bin_centres = (bin_edges[:-1] + bin_edges[1:])/2
    coeff, var_matrix = curve_fit(gauss, bin_centres, hist)

    if plot_fit and mpl:
      hist_fit = gauss(bin_centres, *coeff)
      fig=plt.figure()
      plt.plot(bin_centres, hist, label="Data")
      plt.plot(bin_centres, hist_fit, label="Fit")
      plt.show(False)
    #To keep compatibility between scipy fitter (which requires norm parameter A) and iminuit (which does not),
    #don't return the normalization coeff.
    coeff=coeff[1:]
    var_matrix=var_matrix[1:,1:]
    if plot_fit:
      return coeff, var_matrix, fig
    else:
      return coeff, var_matrix

  elif fitter=="iminuit":
    #ext_gauss=probfit.Extended(probfit.gaussian)
    #likelihood=probfit.BinnedLH(ext_gauss, fit_data,bins=nbins, extended=True)#, extended_bound=(-1000,1000))
    likelihood=probfit.BinnedLH(probfit.gaussian, fit_data,bins=nbins)#, extended_bound=(-1000,1000))
    #likelihood=probfit.UnbinnedLH(probfit.gaussian, fit_data, extended=False)
    minuit = iminuit.Minuit(likelihood,mean=mean_guess,sigma=rms_guess, print_level=1)
    minuit.migrad()
    #      minuit.minos()
    if plot_fit:
      fig=plt.figure()
      likelihood.draw(minuit)
    #Probably a better way to handle this
    coeff=[]
    coeff.append(minuit.values[0])
    coeff.append(minuit.values[1])
    var_matrix=[]
    var_matrix.append(minuit.matrix()[0])
    var_matrix.append(minuit.matrix()[1])
#    var_matrix.append(minuit.matrix()[2])
    if plot_fit:
      return coeff, var_matrix, fig
    else:
      return coeff, var_matrix
  else:
    print("Error, fitter not valid, exiting without doing fit...")
    return None, None

class SkipperImage:
  def _reset_derived(self):
    self.image_means=None
    self.image_rmses=None
    self.baseline=None
    self.charge_mask=None

    self.baseline_subtracted=False
    self.noise=None
    self.noise_coeffs=None
    self.noise_errs=None
    self.gen_charge=None
    self.image_gen_charge=None
    return True

  def __init__(self, fname, extension=0):
    self.hdu=fits.open(fname)[extension]
    self.header=self.hdu.header.copy()
    self.ndcms=self.header.get("NDCMS")
    self.nrows=self.header.get("NAXIS2")
    self.ncols=self.header.get("NAXIS1")
    self.ncols_phys=int(self.ncols/self.ndcms)
    self.npix=self.nrows*self.ncols_phys
    self.data=self.hdu.data
    #Total exposure time in seconds
    self.exptot=self.header.get("MEXP")/1000+self.header.get("MREAD")/1000
    self.nskips=self.ndcms
    #Default parameters
    self.pre_skips=0
    self.post_skips=0
    self.invert=False

    #Derived values
    self._reset_derived()

    self.figures=[]
    self.ds9=None

  def set_params(self,nskips=None, pre_skips=None, post_skips=None, invert=None, *args, **kwars):
    '''
    Simple handler function to change parameters for handling the skipper image. Returns number of changed parameters.
    '''
    changed=0
    if nskips is not None and nskips > 0 and nskips != self.nskips:
      self.nskips=nskips
      changed+=1

    if pre_skips is not None and pre_skips >= 0 and pre_skips != self.pre_skips:
      self.pre_skips=pre_skips
      changed+=1

    if post_skips is not None and post_skips >= 0 and post_skips != self.post_skips:
      self.post_skips=post_skips
      changed+=1

    if invert is not None and type(invert) is bool and invert != self.invert:
      self.invert=invert
      changed+=1

    total_skips=self.pre_skips+self.post_skips
    if total_skips >= self.ndcms:
      print("Warning, trying to skip more skips than exist, reducing skipped skips...")
      #Reduced both proportional to the original setting
      self.pre_skips=int((self.ndcms-1)*self.pre_skips/(total_skips))
      self.post_skips=int((self.ndcms-1)*self.post_skips/(total_skips))
      print("New pre_skips is: " + str(self.pre_skips) + ". New post_skips is: " +str(self.post_skips))

    if changed>0:
      #All our old computed values become useless, so clear them out
      self._reset_derived()

    return changed

  def combine_skips(self,recompute=False,*args, **kwargs):
    '''
    Creates an image from the raw skipper data, skipping the first "pre_skips" and last "post_skip" skips.
    '''
    new_params=self.set_params(*args, **kwargs)

    if new_params==0 and self.image_means is not None and not recompute:
      return False

    means=np.zeros((self.nrows,self.ncols_phys))
    rmses=np.zeros((self.nrows,self.ncols_phys))

    #Actually make the image
    for i in range(self.ncols_phys):
      xslice=np.s_[i*self.ndcms+self.pre_skips:(i+1)*self.ndcms-self.post_skips]
      means[:,i]=np.mean(self.data[:,xslice],axis=1)
      rmses[:,i]=np.std(self.data[:,xslice],axis=1)

    if self.invert:
      means=-1*means
    #Store the result
    self.image_means=means
    self.image_rmses=rmses
    self.baseline_subtracted=False

    return True

  def write_combined_fits(self,fname, *args, **kwargs):
    '''
    Write the combined skips to a .fits file.
    '''

    header=self.header.copy()
    header["NAXIS1"]=self.ncols/self.ndcms

    self.combine_skips(*args,**kwargs)

    mean_hdu=fits.PrimaryHDU(self.image_means)
    mean_hdu.header=header
    rms_hdu=fits.ImageHDU(self.image_rmses, header, "RMSES")

    hdul=fits.HDUList([mean_hdu, rms_hdu])
    hdul.writeto(fname,clobber=False)

  def compute_baseline(self, *args, **kwargs):
    '''
    Computes the baseline for the CCD image, as the median value along each row and column. Also forms the baseline per pixel as the sum of the baselines in x and y.
    '''
    new_image=self.combine_skips(*args, **kwargs)

    if self.baseline is not None and not new_image:
      return False
    self.baseline=np.zeros(self.image_means.shape)

    self.baseline_x=np.median(self.image_means-self.baseline, axis=0)
    self.baseline[:,:]+=self.baseline_x[np.newaxis,:]
    self.baseline_y=np.median(self.image_means-self.baseline, axis=1)
    self.baseline[:,:]+=self.baseline_y[:,np.newaxis]
    return True

  def subtract_baseline(self, *args, **kwargs):
    '''
    Subtracts baseline from the image. Computes the baseline and skipper means, if not already done.
    '''
    new_image=self.compute_baseline(*args, **kwargs)
    if self.baseline_subtracted and not new_image:
      return False

    self.image_means-=self.baseline
    self.baseline_subtracted=True
    return True

  def draw_image(self, cmap="spectral", *args,**kwargs):
    '''
    Function to draw the image after combining skips.
    '''

    if not mpl:
      print("Error, matplotlib not found, cannot draw image...")
      return False

    self.combine_skips(*args, **kwargs)
    if pyds9:
      if self.ds9 is None:
        self.ds9=DS9()
      print("Drawing to DS9")
      self.ds9.set_np2arr(self.image_means)
      self.ds9.set("scale histequ")
      return True
    else:
      print("Drawing using matplotlib")
      fig, mesh=plot_2d(self.image_means)
      self.figures.append(fig)
      return True

  def compute_charge_mask(self, *args, **kwargs):
    '''
    Not implemented
    '''
    self.compute_baseline(*args, **kwargs)

    if self.baseline_subtracted:
      pass

  def compute_noise(self,nbins=None, plot_fit=False,*args, **kwargs):
    '''
    Performs a binned fit to the data to estimate the pixel noise of the CCD image
    '''

    #self.combine_skips(*args, **kwargs)
    #Need to subtract the baseline for the fit to work properly
    #TODO: fix that so this doesn't need to be done (maybe split left/right sides of image? Estimate mu?)
    self.subtract_baseline(*args, **kwargs)
    plot_fit=plot_fit and mpl

    fit_results=fit_gauss(self.image_means.flatten(), nbins=nbins, plot_fit=plot_fit, *args, **kwargs)
    coeff=fit_results[0]
    var_matrix=fit_results[1]
    if plot_fit:
      self.figures.append(fit_results[2])
    self.noise=coeff[-1]
    self.noise_coeffs=coeff
    self.noise_errs=var_matrix

    return coeff, var_matrix

  def compute_gen_charge(self, *args, **kwargs):
    '''
    Estimates the amount of spurious charge generated in each pixel from the skipping procedure
    '''
    self.set_params(*args, **kwargs)
    self.image_gen_charge=np.zeros((self.nrows, self.ncols_phys))
    #TODO: reimplement this using slices to eliminate the loop
    for i in range(self.ncols_phys):
      x1=i*self.ndcms+self.pre_skips
      x2=(i+1)*self.ndcms-self.post_skips-1
      self.image_gen_charge[:,i]=self.data[:,x2]-self.data[:,x1]
    self.gen_charge=np.median(self.image_gen_charge)
    return self.gen_charge

  def compute_charge_loss(*args, **kwargs):
    self.set_params(*args, **kwargs)

  def compute_statistics(self, *args, **kwargs):
    self.combine_skips(*args, **kwargs)
    self.compute_gen_charge(*args, **kwargs)
    self.compute_charge_loss(*args, **kwargs)


# Define a procedure, find_element,
# that takes as its inputs a list
# and a value of any type, and
# returns the index of the first
# element in the input list that
# matches the value.

# If there is no matching element,
# return -1.

#print find_element([1,2,3],3)
#>>> 2

#print find_element(['alpha','beta'],'gamma')
#>>> -1

def find_element(search_string, target_string):
    i = 0
    for element in search_string:
        if element == target_string:
            return i
        else:
            i = i + 1
    return -1


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

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://doc.scrapy.org/en/latest/topics/item-pipeline.html

import pymongo
from scrapy.utils.project import get_project_settings
from xiecheng_scrapy.items import XiechengScrapyItem,XiechengComments,XiechengTicket

class XiechengScrapyPipeline(object):
    def __init__(self):
        settings = get_project_settings()
        # 链接数据库
        self.client = pymongo.MongoClient(host=settings['MONGO_HOST'], port=settings['MONGO_PORT'])

        self.db = self.client[settings['MONGO_DB']]
        self.scenic_comments_coll = self.db[settings['SCENIC_COMMENTS_COLL']]
        self.scenic_info_coll = self.db[settings['SCENIC_INFO_COLL']]
        self.scenic_tickes_coll = self.db[settings['SCENIC_TICKES_COLL']]

    def process_item(self,item ,spider):
        # self.scenic_data = XiechengScrapyItem
        if isinstance(item, XiechengScrapyItem):
            self.scenic_info_coll.insert(dict(item))
        if isinstance(item, XiechengTicket):
            self.scenic_tickes_coll.insert(dict(item))
        if isinstance(item, XiechengComments):
            self.scenic_comments_coll.insert(dict(item))
        return item


import pytest
from middle_of_the_linked_list import ListNode, middle_node

def list_from_array(nums):
    if not nums:
        return None
    head = ListNode(nums[0])
    current = head
    for num in nums[1:]:
        current.next = ListNode(num)
        current = current.next
    return head

def array_from_list(node):
    result = []
    while node:
        result.append(node.val)
        node = node.next
    return result

@pytest.mark.parametrize("head, expected", [
    ([0], [0]),
    ([6], [6]),
    ([0, 1], [1]),
    ([1, 2, 3], [2, 3]),
    ([0, 1], [1]),
    ([1, 0], [0]),
    ([1, 0, 2, 0, 3, 0, 4], [0, 3, 0, 4]),
    ([1, 0, -2, 6, 3, 0, 5, 4, 0], [3, 0, 5, 4, 0]),
    ([1, 0, -2, 6, 3, 7, 5, 4, 0, 8], [7, 5, 4, 0, 8]),
    ([1, 2, 3, 4, 5], [3, 4, 5]),
    ([1, 2, 3, 4, 5, 6], [4, 5, 6])
])
def test_middle_of_the_linked_list(head, expected):
    linked_list = list_from_array(head)
    middle_nodes = middle_node(linked_list)
    assert array_from_list(middle_nodes) == expected


#
# @lc app=leetcode.cn id=567 lang=python3
#
# [567] 字符串的排列
#

# @lc code=start
"""
固定单位滑动窗口，感觉效率不咋滴，算是暴力枚举了，但是能通过。
Accepted
106/106 cases passed (5452 ms)
Your runtime beats 7.31 % of python3 submissions
Your memory usage beats 13.36 % of python3 submissions (15.2 MB)
"""

"""
class Solution:
    def checkInclusion(self, s1: str, s2: str) -> bool:
        length = len(s1)
        left, right = 0, length
        need = {}
        for char in s1:
            count = need.get(char, 0)
            need[char] = count + 1

        while right <= len(s2):
            sub_str = s2[left:right]
            p = 0
            nd = need.copy()
            while p < length:
                char = sub_str[p]
                if char in need:
                    count = nd.get(char)
                    if count < 0:
                        break
                    nd[char] = nd.get(char) - 1
                else:
                    break
                p += 1
            if all(i == 0 for i in list(nd.values())):
                return True

            right += 1
            left += 1

        return False
"""

"""
滑动窗口
"""

"""
滑动窗口算法框架
window: 窗口中的字符
need：需要凑齐的字符
valid标识窗口中满足need条件的字符的个数，如果valid和len(need)相同，说明窗口已经满足条件，完全覆盖了字串。
- 一个字符进入窗口时，增加window计数器
- 一个字符移出窗口的时候，应该减少window计数器
- 当valid满足need时，应该收缩窗口
- 在收缩窗口的时候更新最终结构
"""

"""
class Solution:
    def slidingWindow(self, s1: str, s2: str) -> bool:
        need = {}
        window = {}
        for char in s2:
            count = need.get(char, 0)
            need[char] = count + 1

        left = right = valid = 0

        while right < len(s1):
            c = s1[right]
            right += 1
            # 进行窗口内数据的一系列更新
            # ...

            # debug 输出的位置
            print('widow: [{}, {}]\n'.format(left, right))

            # 判断左窗口是否要收缩（循环窗口是否搜索的判断条件）
            while 'window needs shrink':
                d = s1[left]
                left += 1
                # 进行窗口内数据的一系列更新
                # ...
"""

"""
套模板重新写这道题
Accepted
106/106 cases passed (64 ms)
Your runtime beats 65.4 % of python3 submissions
Your memory usage beats 24.64 % of python3 submissions (15.1 MB)
"""

class Solution:
    def checkInclusion(self, s1: str, s2: str) -> bool:
        need = {}
        window = {}
        for char in s1:
            count = need.get(char, 0)
            need[char] = count + 1

        left = right = valid = 0

        while right < len(s2):
            c = s2[right]
            right += 1
            if c in need:
                window[c] = window.get(c, 0) + 1
                if window[c] == need[c]:
                    valid += 1

            while right - left >= len(s1):
                if valid == len(need):
                    return True
                d = s2[left]
                left += 1
                if d in need:
                    if window.get(d) == need.get(d):
                        valid -= 1
                    window[d] = window.get(d) - 1

        return False

# @lc code=end


import numpy as np
from astropy import units as u
from astropy.coordinates import Angle
from astropy.time import Time

from ctapipe.instrument.guess import guess_telescope
from ctapipe.instrument import TelescopeDescription, CameraGeometry, \
    OpticsDescription, SubarrayDescription
from ctapipe.io.eventsource import EventSource

from target_io import WaveformArrayReader
from target_calib import CameraConfiguration

from .containers import TargetIODataContainer
from .calibration import TargetIOR1Calibrator

__all__ = [
    'TargetIOEventSource',
    'TargetIOR1Calibrator',
]

class TargetIOEventSource(EventSource):
    """
    EventSource for the targetio unofficial data format, the data
    format used by cameras containing TARGET modules, such as CHEC for
    the GCT SST.

    Extract waveform information from `target_io` to store them
    into `ctapipe.io.containers`.

    This Extractor can fill either the R0 or R1 event container, depending on
    the file being read (The header of the file is checked for a flag
    indicating that it has had R1 calibration applied to it).

    This reader requires the TARGET libraries. The instructions to install
    these libraries can be found here:
    https://forge.in2p3.fr/projects/gct/wiki/Installing_CHEC_Software

    Attributes
    ----------
    _tio_reader : target_io.TargetIOEventReader()
        C++ event reader for TargetIO files. Handles the event building into
        an array of (n_pixels, n_samples) in C++, avoiding loops in Python
    _n_events : int
        number of events in the fits file
    _n_samples : int
        number of samples in the waveform
    _r0_samples : ndarray
        three dimensional array to store the R0 level waveform for each pixel
        (n_channels, n_pixels, n_samples)
    _r1_samples : ndarray
        three dimensional array to store the R1 level waveform for each pixel
        (n_channels, n_pixels, n_samples)
    _samples : ndarray
        pointer to the first index of either r0_samples or r1_samples
        (depending if the file has been R1 calibrated) for passing to
        TargetIO to be filled
    """

    def __init__(self, config=None, parent=None, **kwargs):
        super().__init__(config=config, parent=parent, **kwargs)

        self._data = None
        self._event_index = None
        self._event_id = 0
        self._time_tack = None
        self._time_sec = None
        self._time_ns = None

        self._reader = WaveformArrayReader(self.input_url, 2, 1)

        self._n_events = self._reader.fNEvents
        self._first_event_id = self._reader.fFirstEventID
        self._last_event_id = self._reader.fLastEventID
        self._obs_id = self._reader.fRunID
        n_modules = self._reader.fNModules
        n_pix = self._reader.fNPixels
        n_samples = self._reader.fNSamples
        self.camera_config = CameraConfiguration(self._reader.fCameraVersion)
        self._n_cells = self.camera_config.GetNCells()
        m = self.camera_config.GetMapping(n_modules == 1)

        self._optical_foclen = u.Quantity(2.15, u.m)
        self._mirror_area = u.Quantity(14.126, u.m ** 2)
        self._n_pixels = m.GetNPixels()
        self._xpix = np.array(m.GetXPixVector()) * u.m
        self._ypix = np.array(m.GetYPixVector()) * u.m
        self._refshape = np.zeros(10)  # TODO: Get correct values for CHEC-S
        self._refstep = 0  # TODO: Get correct values for CHEC-S
        self._time_slice = 0  # TODO: Get correct values for CHEC-S
        self._chec_tel = 0

        # Init fields
        self._r0_samples = None
        self._r1_samples = None
        self._first_cell_ids = np.zeros(n_pix, dtype=np.uint16)

        # Check if file is already r1 (Information obtained from a flag
        # in the file's header)
        is_r1 = self._reader.fR1
        if is_r1:
            self._r1_samples = np.zeros(
                (1, n_pix, n_samples),
                dtype=np.float32
            )
            self._get_tio_event = self._reader.GetR1Event
            self._samples = self._r1_samples[0]
        else:
            self._r0_samples = np.zeros(
                (1, n_pix, n_samples),
                dtype=np.uint16
            )
            self._get_tio_event = self._reader.GetR0Event
            self._samples = self._r0_samples[0]

        self._init_container()

    @staticmethod
    def is_compatible(file_path):
        return file_path.endswith('.tio')

    def _init_container(self):
        """
        Prepare the ctapipe event container, and fill it with the information
        that does not change with event, including the instrument information.
        """
        chec_tel = 0

        data = TargetIODataContainer()
        data.meta['origin'] = "targetio"

        data.meta['input'] = self.input_url
        data.meta['max_events'] = self.max_events

        # Instrument information

        camera = CameraGeometry(
            "CHEC",
            pix_id=np.arange(self._n_pixels),
            pix_x=self._xpix,
            pix_y=self._ypix,
            pix_area=None,
            pix_type='rectangular',
        )

        optics = OpticsDescription(
            name="ASTRI",
            num_mirrors=2,
            equivalent_focal_length=self._optical_foclen,
            mirror_area=self._mirror_area,
            num_mirror_tiles=2,
        )

        tel_descriptions = {
            chec_tel: TelescopeDescription(
                name="ASTRI",
                type="SST",
                camera=camera,
                optics=optics,
            )
        }
        tel_positions = {
            chec_tel: u.Quantity(0, u.m)
        }

        data.inst.subarray =SubarrayDescription(
            "CHECMonoArray",
            tel_positions=tel_positions,
            tel_descriptions=tel_descriptions,
        )

        self._data = data

    def _update_container(self):
        """
        Update the ctapipe event containers with the information from the
        current event being pointed to in TargetIO.
        """
        data = self._data
        chec_tel = 0

        obs_id = self._obs_id
        event_id = self._event_id
        tels = {self._chec_tel}

        data.r0.obs_id = obs_id
        data.r0.event_id = event_id
        data.r0.tels_with_data = tels
        data.r1.obs_id = obs_id
        data.r1.event_id = event_id
        data.r1.tels_with_data = tels
        data.dl0.obs_id = obs_id
        data.dl0.event_id = event_id
        data.dl0.tels_with_data = tels

        data.trig.tels_with_trigger = [chec_tel]

        data.meta['tack'] = self._time_tack
        data.meta['sec'] = self._time_sec
        data.meta['ns'] = self._time_ns
        data.trig.gps_time = Time(self._time_sec * u.s, self._time_ns * u.ns,
                                  format='unix', scale='utc', precision=9)

        data.count = self._event_index

        data.r0.tel.clear()
        data.r1.tel.clear()
        data.dl0.tel.clear()
        data.dl1.tel.clear()
        data.mc.tel.clear()
        data.targetio.tel.clear()

        # load the data per telescope/chan
        data.r0.tel[chec_tel].waveform = self._r0_samples
        data.r1.tel[chec_tel].waveform = self._r1_samples

        # Load the TargetIO specific data per telescope/chan
        data.targetio.tel[chec_tel].first_cell_ids = self._first_cell_ids
        data.r0.tel[chec_tel].num_samples = self._samples.shape[-1]

        # Some information that currently exists in the mc container, but is
        # useful for real data (essentially the reference pulse shape,
        # which may be used in charge extraction methods)
        data.mc.tel[chec_tel].reference_pulse_shape = self._refshape
        data.mc.tel[chec_tel].meta['refstep'] = self._refstep
        data.mc.tel[chec_tel].time_slice = self._time_slice

    @property
    def _current_event_index(self):
        return self._event_index

    @_current_event_index.setter
    def _current_event_index(self, val):
        """
        Setting the event index will cause the event to be saught from
        TargetIO, and the Containers to point to
        the correct event. The ctapipe event containers are then updated
        with this new event's information.
        """
        self._event_index = val
        self._get_tio_event(val, self._samples, self._first_cell_ids)
        self._event_id = self._reader.fCurrentEventID
        self._time_tack = self._reader.fCurrentTimeTack
        self._time_sec = self._reader.fCurrentTimeSec
        self._time_ns = self._reader.fCurrentTimeNs
        self._update_container()

    def _generator(self):
        for self._current_event_index in range(self._n_events):
            yield self._data
        return

    def __len__(self):
        num = self._n_events
        if self.max_events and self.max_events < num:
            num = self.max_events
        return num

    def _get_event_by_index(self, index):
        self._current_event_index = index
        return self._data

    def _get_event_by_id(self, event_id):
        if ((event_id < self._first_event_id) |
                (event_id > self._last_event_id)):
            raise IndexError(f"Event id {event_id} not found in file")
        index = self._reader.GetEventIndex(event_id)
        return self._get_event_by_index(index)


def printmap(knots, visited):
    ax = min(x for x, y in visited) - 5
    bx = max(x for x, y in visited) + 5
    ay = min(y for x, y in visited) - 5
    by = max(y for x, y in visited) + 5
    for y in range(ay, by)[::-1]:
        for x in range(ax, bx):
            xy = Position(x, y)
            if knots[0] == xy:
                print("H", end='')
            else:
                isk = False
                k = 0
                for k, K in enumerate(knots[1:], start=1):
                    if K == xy:
                        isk = True
                        break
                if isk:
                    print(str(k), end='')
                elif (xy.x, xy.y) in visited:
                    print("x", end='')
                elif y == 0:
                    print("-", end='')
                elif x == 0:
                    print("|", end='')
                else:
                    print(".", end='')
        print(f" # {y}")

class Position:
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

    def move(self, other: "Position"):
        self.x += other.x
        self.y += other.y

    def __eq__(self, other: "Position"):
        return self.x == other.x and self.y == other.y

    def touching(self, other: "Position"):
        return (
            abs(self.x - other.x) <= 1
            and abs(self.y - other.y) <= 1
        )

    def follow(self, other: "Position"):
        if not self.touching(other):
            self.x += max(-1, min(1, other.x - self.x))
            self.y += max(-1, min(1, other.y - self.y))
            return True
        return False

DIRECTIONS = {
    "L": Position(-1, 0),
    "R": Position(1, 0),
    "U": Position(0, 1),
    "D": Position(0, -1)
}

def readfile():
    with open('d9_input.txt') as f:
        for line in f:
            if not line:
                continue
            parts = line.strip().split()
            direction = DIRECTIONS[parts[0]]
            steps = int(parts[1])
            yield direction, steps

def solve(num_knots=1):
    knots = [Position() for _ in range(num_knots + 1)]
    visited = {(0, 0)}
    for i, (direction, steps) in enumerate(readfile()):
        while steps > 0:
            knots[0].move(direction)
            for k in range(len(knots)-1):
                k1 = knots[k]
                k2 = knots[k+1]
                if not k2.follow(k1):
                    break
            visited.add((knots[-1].x, knots[-1].y))
            steps -= 1
    printmap(knots, visited)
    print("Total positions:", len(visited))

solve(num_knots=1)
solve(num_knots=9)


import os

from acts.libs.proto.proto_utils import parse_proto_to_ascii
from acts.libs.proto.proto_utils import to_descriptor_proto
from acts.utils import dump_string_to_file

class ProtoMetric(object):
    """A wrapper around a protobuf containing metrics data.

    This is the primary data structure used as the output of MetricLoggers. It
    is generally intended to be used as-is as a simple wrapper structure, but
    can be extended to provide self-populating APIs.

    Attributes:
        name: The name of the metric.
        data: The data of the metric.
    """

    def __init__(self, name=None, data=None):
        """Initializes a metric with given name and data.

        Args:
            name: The name of the metric. Used in identifiers such as filename.
            data: The data of the metric. Should be a protobuf object.
        """
        self.name = name
        if not data:
            raise ValueError("Parameter 'data' cannot be None.")
        self.data = data

    def get_binary(self):
        """Gets the binary representation of the protobuf data."""
        return self.data.SerializeToString()

    def get_ascii(self):
        """Gets the ascii representation of the protobuf data."""
        return parse_proto_to_ascii(self.data)

    def get_descriptor_binary(self):
        """Gets the binary representation of the descriptor protobuf."""
        return to_descriptor_proto(self.data).SerializeToString()

    def get_descriptor_ascii(self):
        """Gets the ascii representation of the descriptor protobuf."""
        return parse_proto_to_ascii(to_descriptor_proto(self.data))

class MetricPublisher(object):
    """A helper object for publishing metric data.

    This is a base class intended to be implemented to accommodate specific
    metric types and output formats.

    Attributes:
        context: The context in which the metrics are being published.
    """

    def __init__(self, context):
        """Initializes a publisher for the give context.

        Args:
            context: The context in which the metrics are being published.
                     Typically matches that of a containing MetricLogger.
        """
        if not context:
            raise ValueError("Parameter 'context' cannot be None.")
        self.context = context

    def publish(self, metrics):
        """Publishes a list of metrics.

        Args:
            metrics: A list of metrics to publish. The requirements on the
            object type of these metrics is up to the implementing class.
        """
        raise NotImplementedError()

class ProtoMetricPublisher(MetricPublisher):
    """A MetricPublisher that will publish ProtoMetrics to files.

    Attributes:
        publishes_binary: Whether to publish the binary proto.
        publishes_ascii: Whether to publish the ascii proto.
        publishes_descriptor_binary: Whether to publish the binary descriptor.
        publishes_descriptor_ascii: Whether to publish the ascii descriptor.
    """

    ASCII_EXTENSION = 'proto.data'
    BINARY_EXTENSION = 'proto.bin'
    ASCII_DESCRIPTOR_EXTENSION = 'proto.desc'
    BINARY_DESCRIPTOR_EXTENSION = 'proto.desc.bin'

    METRICS_DIR = 'metrics'

    def __init__(self,
                 context,
                 publishes_binary=True,
                 publishes_ascii=True,
                 publishes_descriptor_binary=True,
                 publishes_descriptor_ascii=True):
        """Initializes a ProtoMetricPublisher.

        Args:
            context: The context in which the metrics are being published.
            publishes_binary: Whether to publish the binary proto.
            publishes_ascii: Whether to publish the ascii proto.
            publishes_descriptor_binary: Whether to publish the binary
                                         descriptor.
            publishes_descriptor_ascii: Whether to publish the ascii
                                        descriptor.
        """
        super().__init__(context)
        self.publishes_binary = publishes_binary
        self.publishes_ascii = publishes_ascii
        self.publishes_descriptor_binary = publishes_descriptor_binary
        self.publishes_descriptor_ascii = publishes_descriptor_ascii

    def get_output_path(self):
        """Gets the output directory path of the metrics."""
        return os.path.join(self.context.get_full_output_path(),
                            self.METRICS_DIR)

    def publish(self, metrics):
        """Publishes the given list of metrics.

        Based on the publish_* attributes of this class, this will publish
        the varying data formats provided by the metric object. Data is written
        to files on disk named according to the name of the metric.

        Args:
            metrics: The list metric to publish. Assumed to be a list of
                     ProtoMetric objects.
        """
        if isinstance(metrics, list):
            for metric in metrics:
                self._publish_single(metric)
        else:
            self._publish_single(metrics)

    def _publish_single(self, metric):
        """Publishes a single metric.

        Based on the publish_* attributes of this class, this will publish
        the varying data formats provided by the metric object. Data is written
        to files on disk named according to the name of the metric.

        Args:
            metric: The metric to publish. Assumed to be a ProtoMetric object.
        """
        output_path = self.get_output_path()

        os.makedirs(output_path, exist_ok=True)

        if self.publishes_binary:
            self.write_binary(metric, output_path)
        if self.publishes_ascii:
            self.write_ascii(metric, output_path)
        if self.publishes_descriptor_binary:
            self.write_descriptor_binary(metric, output_path)
        if self.publishes_descriptor_ascii:
            self.write_descriptor_ascii(metric, output_path)

    def write_binary(self, metric, output_path):
        """Writes the binary format of the protobuf to file.

        Args:
            metric: The metric to write.
            output_path: The output directory path to write the file to.
        """
        filename = self._get_output_file(
            output_path, metric.name, self.BINARY_EXTENSION)
        dump_string_to_file(metric.get_binary(), filename, mode='wb')

    def write_ascii(self, metric, output_path):
        """Writes the ascii format of the protobuf to file.

        Args:
            metric: The metric to write.
            output_path: The output directory path to write the file to.
        """
        filename = self._get_output_file(
            output_path, metric.name, self.ASCII_EXTENSION)
        dump_string_to_file(metric.get_ascii(), filename)

    def write_descriptor_binary(self, metric, output_path):
        """Writes the binary format of the protobuf descriptor to file.

        Args:
            metric: The metric to write.
            output_path: The output directory path to write the file to.
        """
        filename = self._get_output_file(
            output_path, metric.name, self.BINARY_DESCRIPTOR_EXTENSION)
        dump_string_to_file(metric.get_descriptor_binary(), filename, mode='wb')

    def write_descriptor_ascii(self, metric, output_path):
        """Writes the ascii format of the protobuf descriptor to file.

        Args:
            metric: The metric to write.
            output_path: The output directory path to write the file to.
        """
        filename = self._get_output_file(
            output_path, metric.name, self.ASCII_DESCRIPTOR_EXTENSION)
        dump_string_to_file(metric.get_descriptor_ascii(), filename)

    def _get_output_file(self, output_path, filename, extension):
        """Gets the full output file path.

        Args:
            output_path: The output directory path.
            filename: The base filename of the file.
            extension: The extension of the file, without the leading '.'

        Returns:
            The full file path.
        """
        return os.path.join(output_path, "%s.%s" % (filename, extension))