# start
# find highly selected genes in one species across populations and donors
import os
import glob
import copy
from Bio import SeqIO
from Bio.Seq import Seq
import argparse
############################################ Arguments and declarations ##############################################
parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter)
required = parser.add_argument_group('required arguments')
optional = parser.add_argument_group('optional arguments')
required.add_argument("-i",
                      help="path of folders of WGS of each species",
                      type=str, default='.',
                      metavar='input/')
required.add_argument("-fq",
                      help="file extension of WGS fastq #1 files",
                      type=str, default='_1.fastq',
                      metavar='_1.fastq')
# optional output setup

optional.add_argument("-s",
                      help="a folder to store all scripts",
                      type=str, default='scripts/',
                      metavar='scripts/')
optional.add_argument("-o",
                      help="a folder to store all output",
                      type=str, default='snp_output/',
                      metavar='snp_output/')
# optional search parameters
optional.add_argument('-t',
                      help="Optional: set the thread number assigned for running XXX (default 1)",
                      metavar="1 or more", action='store', default=1, type=int)
optional.add_argument('-job',
                      help="Optional: command to submit jobs",
                      metavar="nohup or customized",
                      action='store', default='jobmit', type=str)
# requirement for software calling
optional.add_argument('-bw', '--bowtie',
                          help="Optional: complete path to bowtie if not in PATH",
                          metavar="/usr/local/bin/bowtie",
                          action='store', default='bowtie', type=str)
optional.add_argument('-sp', '--spades',
                          help="Optional: complete path to spades if not in PATH",
                          metavar="/usr/local/bin/spades",
                          action='store', default='spades', type=str)
optional.add_argument('-pro', '--prodigal',
                      help="Optional: complete path to prodigal if not in PATH, None for no prodigal (default)",
                      metavar="/usr/local/bin/prodigal",
                      action='store', default='None', type=str)
optional.add_argument('-bcf', '--bcftools',
                      help="Optional: complete path to bcftools if not in PATH",
                      metavar="/usr/local/bin/bcftools",
                      action='store', default='bcftools', type=str)
optional.add_argument('-sam', '--samtools',
                      help="Optional: complete path to bwa if not in PATH",
                      metavar="/usr/local/bin/samtools",
                      action='store', default='samtools', type=str)
optional.add_argument('-mini', '--minimap2',
                      help="Optional: complete path to minimap2 if not in PATH",
                      metavar="/usr/local/bin/minimap2",
                      action='store', default='minimap2', type=str)
optional.add_argument('--u','--usearch',
                        help="Optional: cluster genes with SNPs",
                        metavar="usearch",
                        action='store', default='usearch', type=str)
################################################## Definition ########################################################
args = parser.parse_args()

input_script_sub = args.s +'/annotate'
input_script = args.s
genome_root = args.i + '/round*'
output_dir_merge = args.o +'/vcf_round%s/merge_genome/'%(Round)
all_fasta = os.path.join(output_dir_merge + '/summary', 'all.denovo.gene.faa')
all_fasta_HS = os.path.join(output_dir_merge + '/summary', 'all.selected.gene.faa')
input_summary = output_dir_merge + '/summary/all.donor.species.dnds.txt'

# functions
# clustering
def cluster_uc(cluster_input):
    Clusters = dict()
    High_select2 = set()
    High_select2_3SNP = dict()
    High_select2_output = []
    for lines in open(cluster_input, 'r'):
        line_set = lines.split('\n')[0].split('\t')
        cluster = line_set[1]
        record_name = line_set[8]
        record_name2 = line_set[9]
        Clusters.setdefault(record_name, cluster)
        if record_name2!= '*':
            donor = record_name.split('_')[0]
            donor_species = record_name.split('__')[0]
            species = donor_species.replace(donor + '_', '')
            donor2 = record_name2.split('_')[0]
            donor_species2 = record_name2.split('__')[0]
            species2 = donor_species2.replace(donor2 + '_', '')
            if species == species2:
                if species not in SNP_cutoff:
                    High_select2.add(record_name2)
                    High_select2.add(record_name)
                else:
                    High_select2_3SNP.setdefault(record_name,set())
                    High_select2_3SNP[record_name].add(record_name2)
                    High_select2_3SNP.setdefault(record_name2, set())
                    High_select2_3SNP[record_name2].add(record_name)
    if High_select2_3SNP!= dict():
        for record in High_select2_3SNP:
            if len(High_select2_3SNP[record]) >= 3:
                High_select2.append(record)
    return [Clusters,High_select2_output,High_select2]

# correct highly selected genes dNdS
def format_freq(freq):
    freq = freq.split(':')
    return [int(freq[0]),int(freq[1])]

def init_highselect_notHS(line_set):
    allspecies_highselect = line_set
    for i in [3,5,6,7,8,-13,-11,-9,-7,-5,-3]:
        allspecies_highselect[i] = int(allspecies_highselect[i])
    for i in [-12,-10,-8,-6,-4,-2]:
        allspecies_highselect[i] = format_freq(line_set[i])
    return allspecies_highselect

def init_highselect_empty(line_set):
    allspecies_highselect = line_set
    for i in [3]:
        allspecies_highselect[i] = int(allspecies_highselect[i])
    for i in [5,6,7,8,9,10,11,12,-13,-11,-9,-7,-5,-3]:
        allspecies_highselect[i] = 0
    for i in [-12,-10,-8,-6,-4,-2]:
        allspecies_highselect[i] = format_freq('0:0')
    return allspecies_highselect

def add_freq(freq_new,freq_old):
    temp_result = format_freq(freq_new)
    freq_old[0] += temp_result[0]
    freq_old[1] += temp_result[1]

def add_highselect(line_set,allspecies_highselect):
    for i in [3,5,6,7,8,-13,-11,-9,-7,-5,-3]:
        allspecies_highselect[i] += int(line_set[i])
    for i in [-12,-10,-8,-6,-4,-2]:
        add_freq(line_set[i], allspecies_highselect[i])
    return allspecies_highselect

def calculate_NS(allspecies,allspecies_highselect):
    # expect NS
    # expected NS ratio
    tempNS = [0, 0]
    i = -12
    for freq in allspecies:
        tempNS[0] += freq * allspecies_highselect[i][0]
        tempNS[1] += freq * allspecies_highselect[i][1]
        i += 2
    if tempNS[1] > 0:
        allspecies_highselect[11] = tempNS[0] / tempNS[1] # expect
    else:
        allspecies_highselect[11] = 'expect_N_only'
    # NS ratio
    if allspecies_highselect[8] > 0:
        allspecies_highselect[10] = allspecies_highselect[7]/allspecies_highselect[8] # observe
        if allspecies_highselect[11] != 'expect_N_only':
            # dNdS
            allspecies_highselect[12] = allspecies_highselect[10] / allspecies_highselect[11]# dnds
    elif allspecies_highselect[7] > 0:
        allspecies_highselect[10] = 'observe_N_only'
        allspecies_highselect[12] = 'observe_N_only'
    return allspecies_highselect

def output_highlight(allspecies_highselect,allspecies,newoutput):
    allspecies_highselect = calculate_NS(allspecies, allspecies_highselect)
    for i in range(0, len(allspecies_highselect)):
        if i in [14, 16, 18, 20, 22, 24]:
            allspecies_highselect[i] = '%s:%s' % (allspecies_highselect[i][0],
                                                  allspecies_highselect[i][1])
        else:
            allspecies_highselect[i] = str(allspecies_highselect[i])
    newoutput.append('\t'.join(allspecies_highselect) + '\n')
    return newoutput

def genus_species(donor_species):
    genus = donor_species.split('_')[1].replace('BA', 'Bifidobacterium').replace('BL',
                                                                                'Bifidobacterium').replace(
        'PB', 'Parabacteroides')
    if any(item in donor_species for item in ['BA','BL','PB']):
        species = donor_species.split('_')[1].replace('BA', 'Bifidobacterium_adolescentis').replace('BL',
                                                                                'Bifidobacterium_longum').replace(
            'PB', 'Parabacteroides_butyrate')
    else:
        species = (donor_species.split('_')[1] + '_' +
                   donor_species.split('_')[2])
    return [genus,species]

def add_new_selection(input_summary,High_select2):
    Donor_species = dict()
    Donor_species_notHS = dict()
    Genus = dict()
    Genus_notHS = dict()
    Species = dict()
    Species_notHS = dict()
    newoutput = []
    # use selected codon NS ratio (SNP pair) * all genes freq of no single event
    for lines in open(input_summary_nosingleevent, 'r'):
        line_set = lines.split('\n')[0].split('\t')
        Selected = line_set[-1]
        if Selected == 'False':
            donor_species = line_set[0]
            record_id = line_set[1]
            # all species
            if record_id == 'allspecies':
                # all genes freq
                allspecies = [int(line_set[-13]), int(line_set[-11]), int(line_set[-9]),
                              int(line_set[-7]), int(line_set[-5]), int(line_set[-3])]
                allspecies_highselect = init_highselect_empty(line_set)
                allspecies_highselect[1] = 'allspecies_highselect'
                newoutput.append(lines)
            elif record_id in ['allspecies_flexible', 'allspecies_core']:
                newoutput.append(lines)
            elif donor_species == record_id:
                # donor_species not high select set up
                line_set = lines.split('\n')[0].split('\t')
                donor_species_highselect = init_highselect_notHS(line_set)
                Donor_species_notHS.setdefault(donor_species,
                                               donor_species_highselect)
                line_set = lines.split('\n')[0].split('\t')
                donor_species_highselect = init_highselect_empty(line_set)
                donor_species_highselect[1] = donor_species + '_highselect'
                Donor_species.setdefault(donor_species,
                                         donor_species_highselect)
                genus, species = genus_species(donor_species)
                line_set = lines.split('\n')[0].split('\t')
                if genus not in Genus_notHS:
                    genus_highselect = init_highselect_notHS(line_set)
                    genus_highselect[0] = genus
                    genus_highselect[1] = genus
                    Genus_notHS.setdefault(genus, genus_highselect)
                    line_set = lines.split('\n')[0].split('\t')
                    genus_highselect = init_highselect_empty(line_set)
                    genus_highselect[0] = genus
                    genus_highselect[1] = genus + '_highselect'
                    Genus.setdefault(genus, genus_highselect)
                else:
                    genus_highselect = Genus_notHS[genus]
                    Genus_notHS[genus] = add_highselect(line_set, genus_highselect)
                line_set = lines.split('\n')[0].split('\t')
                if species not in Species_notHS:
                    species_highselect = init_highselect_notHS(line_set)
                    species_highselect[0] = species
                    species_highselect[1] = species
                    Species_notHS.setdefault(species, species_highselect)
                    line_set = lines.split('\n')[0].split('\t')
                    species_highselect = init_highselect_empty(line_set)
                    species_highselect[0] = species
                    species_highselect[1] = species + '_highselect'
                    Species.setdefault(species, species_highselect)
                else:
                    species_highselect = Species_notHS[species]
                    Species_notHS[species] = add_highselect(line_set, species_highselect)
    # calculate NS for new HS genes
    for lines in open(input_summary, 'r'):
        line_set = lines.split('\n')[0].split('\t')
        Selected = line_set[-1]
        donor_species = line_set[0]
        record_id = line_set[1]
        # species or genes
        if Selected == 'False':
            if record_id not in ['0', 'gene', 'allspecies', 'allspecies_highselect',
                                 'allspecies_core', 'allspecies_flexible',
                                 'allspecies_highselect_noNS']:
                genus, species = genus_species(donor_species)
                donor_species_set = donor_species.split('_')
                try:
                    donor_species = '%s_%s_%s' % (donor_species_set[0],
                                                  donor_species_set[1][0:min(6, len(donor_species_set[1]))],
                                                  donor_species_set[2][0:min(6, len(donor_species_set[2]))])
                except IndexError:
                    donor_species = '%s_%s_%s' % (donor_species_set[0],
                                                  donor_species_set[1],
                                                  donor_species_set[1])
                if 'cluster' in donor_species_set[-1]:
                    try:
                        donor_species += '_CL' + donor_species_set[-2].split('cluster')[1]
                    except IndexError:
                        donor_species += '_CL' + donor_species_set[-1].split('cluster')[1]
                record_id = '%s__C_%s_G_%s' % (donor_species, record_id.split('_')[1], record_id.split('_')[-1])
                # highly selected genes
                if record_id in High_select2:
                    if (line_set[10] == 'observe_N_only' or float(line_set[10]) > NSratioobserve_cutoff):
                        newoutput.append('\t'.join(line_set[:-1]) + '\tTrue\n')
                        allspecies_highselect = add_highselect(line_set, allspecies_highselect)
                        Donor_species[line_set[0]] = add_highselect(line_set, Donor_species[line_set[0]])
                        Genus[genus] = add_highselect(line_set, Genus[genus])
                        Species[species] = add_highselect(line_set, Species[species])
                    else:
                        newoutput.append('\t'.join(line_set[:-1]) + '\tFalse\n')
                        High_select2.remove(record_id)
                # other genes and species
                elif line_set[0] not in line_set[1]:
                    newoutput.append(lines)
                else:
                    print('not output %s' % (line_set[1]))
        # original highly selected
        elif Selected == 'True':
            newoutput.append(lines)
            allspecies_highselect = add_highselect(line_set, allspecies_highselect)
            Donor_species[line_set[0]] = add_highselect(line_set, Donor_species[line_set[0]])
            Genus[genus] = add_highselect(line_set, Genus[genus])
            Species[species] = add_highselect(line_set, Species[species])
        else:
            newoutput.append(lines)
    # output new highselect
    newoutput = output_highlight(allspecies_highselect, allspecies, newoutput)
    for donor_species in Donor_species:
        line_set = Donor_species_notHS[donor_species]
        line_set_sub = [int(line_set[-13]), int(line_set[-11]), int(line_set[-9]),
                        int(line_set[-7]), int(line_set[-5]), int(line_set[-3])]
        newoutput = output_highlight(Donor_species[donor_species],
                                     line_set_sub, newoutput)
        newoutput = output_highlight(Donor_species_notHS[donor_species],
                                     line_set_sub, newoutput)
    for genus in Genus:
        line_set = Genus_notHS[genus]
        line_set_sub = [int(line_set[-13]), int(line_set[-11]), int(line_set[-9]),
                        int(line_set[-7]), int(line_set[-5]), int(line_set[-3])]
        newoutput = output_highlight(Genus[genus],
                                     line_set_sub, newoutput)
        newoutput = output_highlight(Genus_notHS[genus],
                                     line_set_sub, newoutput)
    for species in Species:
        line_set = Species_notHS[species]
        line_set_sub = [int(line_set[-13]), int(line_set[-11]), int(line_set[-9]),
                        int(line_set[-7]), int(line_set[-5]), int(line_set[-3])]
        newoutput = output_highlight(Species[species],
                                     line_set_sub, newoutput)
        newoutput = output_highlight(Species_notHS[species],
                                     line_set_sub, newoutput)
    # output new summary
    f1 = open(input_summary + '.High_select2.txt', 'w')
    f1.write(''.join(newoutput))
    f1.close()
    return High_select2

def sum_gene(input_summary,High_select2, selected = 1):
    genelist = []
    genelist += High_select2
    for lines in open(input_summary,'r'):
        lines_set = lines.replace('\n','').split('\t')
        if not lines.startswith('#'):
            if selected == 0 or lines_set[-1] == 'True':
                # highly selected genes or all genes
                gene_name = lines_set[1]
                genelist.append(gene_name)
    print(len(genelist))
    genelist = list(set(genelist))
    print(len(genelist))
    Output = []
    output_set = []
    for record in SeqIO.parse(all_fasta_HS, 'fasta'):
        Output.append('>%s\n%s\n' % (str(record.id), str(record.seq)))
        output_set.append(str(record.id))
    for record in SeqIO.parse(all_fasta, 'fasta'):
        if str(record.id) in genelist and str(record.id) not in output_set:
            Output.append('>%s\n%s\n' % (str(record.id), str(record.seq)))
            output_set.append(str(record.id))
    f1 = open(all_fasta_HS + '.High_select2.faa' , 'w')
    f1.write(''.join(Output))
    f1.close()

def annotation(all_filter_gene_fasta_file,pre_cluster = ''):
    # run cluster
    cutoff = 0.7
    cmd_cluster = ('%s -sort length -cluster_fast %s -id %s -centroids %s.cluster.aa -uc %s.uc -threads %s\n'
                   % ('usearch', all_filter_gene_fasta_file, cutoff, all_filter_gene_fasta_file,
                      all_filter_gene_fasta_file, 40))
    os.system(cmd_cluster)
    all_filter_gene_fasta_file = all_filter_gene_fasta_file + '.cluster.aa'
    if pre_cluster!= '':
        os.system('#%s -makeudb_usearch %s -output %s.udb' %
                  ('usearch', pre_cluster, pre_cluster))
        os.system('%s -ublast %s -db %s.udb  -evalue 1e-2 -accel 0.5 -blast6out %s -threads 2'%
                  ('usearch', all_filter_gene_fasta_file,pre_cluster, all_filter_gene_fasta_file + '.ref.out.txt'))
    # run prokka
    cmdsprokka = 'py37\nprokka --kingdom Bacteria --outdir %s/prokka_%s  --protein %s --locustag Bacter %s/%s\n' % \
                 (output_dir_merge + '/summary', os.path.split(all_filter_gene_fasta_file)[-1],
                  all_filter_gene_fasta_file,
                  output_dir_merge + '/summary',
                  os.path.split(all_filter_gene_fasta_file)[-1].replace('.faa', '.fna'))
    f1 = open(os.path.join(input_script_sub, 'prokka.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s' % (cmdsprokka))
    f1.close()
    # run metacyc
    cutoff = 50
    cutoff2 = 80
    database = '/scratch/users/mit_alm/database/metacyc/protseq.fsa'
    cmds = ("diamond blastp --query %s --db %s.dmnd --out %s.metacyc.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    f1 = open(os.path.join(input_script_sub, 'metacyc.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s'%(cmds))
    f1.close()
    # run eggnog
    cutoff = 0.01
    database = '/scratch/users/mit_alm/database/eggnog/xaa.hmm'
    cmds = ('hmmsearch --tblout %s.eggnog.1.txt --cpu 40 -E %s %s %s\n') %(all_filter_gene_fasta_file,cutoff,database,all_filter_gene_fasta_file)
    f1 = open(os.path.join(input_script_sub, 'eggnog.1.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s'%(cmds))
    f1.close()
    database = '/scratch/users/mit_alm/database/eggnog/xab.hmm'
    cmds = ('hmmsearch --tblout %s.eggnog.2.txt --cpu 40 -E %s %s %s\n') % (
    all_filter_gene_fasta_file, cutoff, database, all_filter_gene_fasta_file)
    f1 = open(os.path.join(input_script_sub, 'eggnog.2.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s' % (cmds))
    f1.close()
    database = '/scratch/users/mit_alm/database/eggnog/xac.hmm'
    cmds = ('hmmsearch --tblout %s.eggnog.3.txt --cpu 40 -E %s %s %s\n') % (
    all_filter_gene_fasta_file, cutoff, database, all_filter_gene_fasta_file)
    f1 = open(os.path.join(input_script_sub, 'eggnog.3.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s' % (cmds))
    f1.close()
    # run kegg
    cutoff = 0.01
    database = '/scratch/users/mit_alm/database/kegg/kofam/profiles/prokaryote/prokaryote.hmm'
    cmds = ('hmmsearch --tblout %s.kegg.txt --cpu 40 -E %s %s %s\n') %(all_filter_gene_fasta_file,cutoff,database,all_filter_gene_fasta_file)
    f1 = open(os.path.join(input_script_sub, 'kegg.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s'%(cmds))
    f1.close()
    # run customed database
    cutoff = 80
    cutoff2 = 80
    cmds = ''
    database = '/scratch/users/anniz44/scripts/database/SARG.db.fasta'
    cmds += ("diamond blastp --query %s --db %s.dmnd --out %s.SARG.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    cutoff = 50
    cutoff2 = 50
    database = '/scratch/users/anniz44/scripts/database/AHR.aa.db'
    cmds += ("diamond blastp --query %s --db %s.dmnd --out %s.AHR.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    cutoff = 60
    cutoff2 = 80
    database = '/scratch/users/anniz44/scripts/database/Butyrate.pro.aa'
    cmds += ("diamond blastp --query %s --db %s.dmnd --out %s.buty.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    cutoff = 50
    cutoff2 = 80
    database = '/scratch/users/anniz44/scripts/database/IntI1_database.fasta'
    cmds += ("diamond blastp --query %s --db %s.dmnd --out %s.int.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    cutoff = 50
    cutoff2 = 80
    database = '/scratch/users/anniz44/scripts/database/SRB.AA'
    cmds += ("diamond blastp --query %s --db %s.dmnd --out %s.SRB.txt --id %s --query-cover %s --outfmt 6 --max-target-seqs 2 --evalue 1e-1 --threads 40\n"
            %(all_filter_gene_fasta_file,database,all_filter_gene_fasta_file,cutoff,cutoff2))
    cutoff = 0.01
    database = '/scratch/users/anniz44/scripts/database/NR.hmm'
    cmds += ('hmmsearch --tblout %s.NR.txt --cpu 40 -E %s %s %s\n') %(all_filter_gene_fasta_file,cutoff,database,all_filter_gene_fasta_file)
    f1 = open(os.path.join(input_script_sub, 'customed.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n%s'%(cmds))
    f1.close()
    # all scripts
    f1 = open(os.path.join(input_script, 'allannotate.sh'), 'w')
    f1.write('#!/bin/bash\nsource ~/.bashrc\n')
    for sub_scripts in glob.glob(os.path.join(input_script_sub, '*.sh')):
        f1.write('jobmit %s %s\n' % (sub_scripts, os.path.split(sub_scripts)[-1]))
    f1.close()

# clustering
Clusters_gene, High_select2_output,High_select2 = cluster_uc(all_fasta + '.uc')
f1 = open(all_fasta + '.High_select2.txt', 'w')
f1.write(''.join(High_select2_output))
f1.close()
# correcting
High_select2 = add_new_selection(input_summary,High_select2)

# run clustering
sum_gene(input_summary,High_select2,1)
annotation(all_fasta_HS + '.High_select2.faa',pre_cluster)

################################################### END ########################################################


import logging
from typing import Any, Dict, Hashable, TypeVar, Union

logger = logging.getLogger(__name__)

D = TypeVar("D")

def get_by_path(d: Dict, *path: Hashable, default: D = None) -> Union[Any, D]:
    """
    Given a nested dict of dicts, traverse a given path and return the result or the default if it is not found.
    This is used as a replacement for the pattern
    >>> d.get("a", {}).get("b", {}).get("c", {}).get("d", 1)
    with
    >>> get_by_path(d, "a", "b", "c", "d", default=1)

    Does not traverse lists. Should be dict all the way down.
    """
    if len(path) == 0:
        raise ValueError("No path given")
    head, *tail = path
    logger.debug("head: %s, tail: %s", head, tail)
    if not tail:
        return d.get(head, default)
    try:
        return get_by_path(d[head], *tail, default=default)
    except KeyError:
        return default


"""Dataset Handler."""
import os
import torch
from datasets import cifar_handler
from datasets import tinyimagenet_handler
from datasets import imagenet2012_handler
from datasets import stl10_handler
import numpy as np

class DataHandler:
  """Handler for datasets."""

  def __init__(
      self,
      dataset_name,
      data_root,
      test_dataset_root,
      experiment_root,
      grayscale,
      gauss_noise,
      gauss_noise_std,
      blur,
      blur_std,
      val_split=0.1,
      test_split=0.1,
      split_idxs_root="split_idxs",
      noise_type=None,
      load_previous_splits=True,
      verbose=True,
      **kwargs
    ):
    """Initialize dataset handler."""
    self.dataset_name = dataset_name
    self.data_root = data_root
    self.test_dataset_root = test_dataset_root
    self.experiment_root = experiment_root
    self.test_split = test_split
    self.val_split = val_split
    self.noise_type = noise_type
    self._verbose = verbose
    self.load_previous_splits = load_previous_splits
    self.grayscale = grayscale #pg_grayscale
    self.gauss_noise = gauss_noise
    self.gauss_noise_std = gauss_noise_std
    self.blur = blur
    self.blur_std = blur_std
    self._kwargs = kwargs

    self._set_num_classes(dataset_name)

    # Set idx with dataset_name
    split_idxs_root = os.path.join(split_idxs_root, dataset_name)
    if not os.path.exists(split_idxs_root):
      os.makedirs(split_idxs_root)

    if split_idxs_root and val_split:
      self.split_idxs_root = self._build_split_idx_root(
        split_idxs_root,
        dataset_name
      )
    else:
      self.split_idxs_root = None

    # Create datasets
    self.datasets = self._build_datasets()

  def _set_num_classes(self, dataset_name):
    """Set number of classes in dataset."""
    if dataset_name == "CIFAR10":
      self.num_classes = 10
    elif dataset_name == "CIFAR100":
      self.num_classes = 100
    elif dataset_name == "TinyImageNet":
      self.num_classes = 200
    elif dataset_name == "STL10":
      self.num_classes = 10

  def get_transform(self, dataset_key=None):
    """Build dataset transform."""
    if dataset_key is None:
      dataset_key = list(self.datasets.keys())[0]

    normalize_transform = None
    # Grab transforms - location varies depending on base dataset.
    try:
      transforms = self.datasets[dataset_key].transform.transforms
      found = True
    except AttributeError:
      found = False

    if not found:
      try:
        transforms = self.datasets[dataset_key].dataset.transform.transforms
        found = True
      except AttributeError:
        found = False

    if not found:
      print("Transform list not found!")
    else:
      found = False
      for xform in transforms:
        if "normalize" in str(xform).lower():
          normalize_transform = xform
          found = True
          break

    if not found:
      print("Normalization transform not found!")
    return normalize_transform

  def _build_split_idx_root(self, split_idxs_root, dataset_name):
    """Build directory for split idxs."""
    if ".json" in split_idxs_root and not os.path.exists(split_idxs_root):
      split_idxs_root = os.path.join(split_idxs_root, dataset_name)
    print(f"Setting split idxs root to {split_idxs_root}")
    if not os.path.exists(split_idxs_root):
      print(f"{split_idxs_root} does not exist!")
      os.makedirs(split_idxs_root)
      print("Complete.")
    return split_idxs_root

  def _build_datasets(self):
    """Build dataset."""
    print("_build_datasets")
    print("self.grayscale", self.grayscale)
    print("self.gauss_noise", self.gauss_noise)
    print("self.gauss_noise_std", self.gauss_noise_std)
    print("self.blur", self.blur)
    print("self.blur_std", self.blur_std)
    if "cifar" in self.dataset_name.lower():
      dataset_dict = cifar_handler.create_datasets(
        self.data_root,
        dataset_name=self.dataset_name,
        val_split=self.val_split,
        grayscale=self.grayscale,
        gauss_noise=self.gauss_noise,
        gauss_noise_std=self.gauss_noise_std,
        blur=self.blur,
        blur_std=self.blur_std,
        split_idxs_root=self.split_idxs_root,
        noise_type=self.noise_type,
        load_previous_splits=self.load_previous_splits,
        verbose=self._verbose
      )

    elif self.dataset_name.lower() == "tinyimagenet":
      dataset_dict = tinyimagenet_handler.create_datasets(
        self.data_root,
        self.val_split,
        self.split_idxs_root
      )

    elif "stl" in self.dataset_name.lower():
        dataset_dict = stl10_handler.create_datasets(
          self.data_root,
          dataset_name=self.dataset_name,
          val_split=self.val_split,
          grayscale=self.grayscale,
          gauss_noise=self.gauss_noise,
          gauss_noise_std=self.gauss_noise_std,
          blur=self.blur,
          blur_std=self.blur_std,
          split_idxs_root=self.split_idxs_root,
          noise_type=self.noise_type,
          load_previous_splits=self.load_previous_splits,
          verbose=self._verbose
        )

    elif str.find(self.dataset_name.lower(), "imagenet2012")>-1:
      # Build path dataframe
      path_df = imagenet2012_handler.build_path_df(self.data_root, self.experiment_root)
      test_path_df = imagenet2012_handler.build_test_path_df(self.test_dataset_root, self.grayscale, self.gauss_noise, self.blur, self.gauss_noise_std, self.blur_std)

      assert len(path_df), "Failed to load path df"
      # Subset data
      if "imagenet_params" in self._kwargs:
        path_df = imagenet2012_handler.subset_path_df(
          path_df,
          self._kwargs["imagenet_params"]
        )
        test_path_df = imagenet2012_handler.subset_path_df(
          test_path_df,
          self._kwargs["imagenet_params"]
        )

      # Set number of classes!
      self.num_classes = path_df.class_lbl.unique().shape[0]
      # Build dataset dict
      dataset_dict = imagenet2012_handler.create_datasets(
        path_df,
        self.val_split,
        self.test_split,
        self.split_idxs_root,
        self.experiment_root,
        self.grayscale,
        self.gauss_noise,
        self.gauss_noise_std,
        self.blur,
        self.blur_std,
        test_path_df
      )

    return dataset_dict

  def build_loader(
      self,
      dataset_key,
      flags,
      dont_shuffle_train=False
    ):
    """Build dataset loader."""
    # Get dataset source
    dataset_src = self.datasets[dataset_key]
    # Specify shuffling
    if dont_shuffle_train:
      shuffle = False
    else:
      shuffle = dataset_key == "train"

    #create weighted sampler for 16 class dataset
    if self.dataset_name == 'ImageNet2012_16classes_rebalanced' and dataset_key == 'train':
      class_sample_count = np.array([len(dataset_src.path_df[dataset_src.path_df.y == t]) for t in dataset_src.path_df.y.unique()])
      weights = 1./class_sample_count
      samples_weight = torch.from_numpy(np.array([weights[t] for t in dataset_src.path_df.y])).double()
      weighted_sampler = torch.utils.data.WeightedRandomSampler(samples_weight, len(dataset_src.path_df), replacement=True)

      # Creates dataloaders, which load data in batches
      loader = torch.utils.data.DataLoader(
        dataset=dataset_src,
        batch_size=flags.batch_size,
        sampler = weighted_sampler,
        num_workers=flags.num_workers,
        drop_last=flags.drop_last,
        pin_memory=True)
    else:

      # Creates dataloaders, which load data in batches
      loader = torch.utils.data.DataLoader(
          dataset=dataset_src,
          batch_size=flags.batch_size,
          shuffle=shuffle,
          num_workers=flags.num_workers,
          drop_last=flags.drop_last,
          pin_memory=True)

    return loader


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
from django.conf import settings

class Migration(migrations.Migration):

    dependencies = [
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
        ('promotions_app', '0002_promotionmanager'),
    ]

    operations = [
        migrations.DeleteModel(
            name='PromotionManager',
        ),
        migrations.AddField(
            model_name='promotion',
            name='account',
            field=models.ForeignKey(default='', to=settings.AUTH_USER_MODEL),
            preserve_default=False,
        ),
        migrations.AlterField(
            model_name='promotion',
            name='desc',
            field=models.TextField(),
            preserve_default=True,
        ),
        migrations.AlterField(
            model_name='promotion',
            name='image',
            field=models.ImageField(upload_to=b'promotions'),
            preserve_default=True,
        ),
    ]


from typing import Optional, List
from fastapi import Body, Depends, HTTPException, status, Query
from pydantic import EmailStr
from .app_helper import (
    create_student,
    get_student,
    update_student,
    delete_student,
)
from .models import (
    Student,
    StudentDelete,
    StudentUpdate,
    StudentReturn,
    Message,
    TokenData,
)
from fastapi import APIRouter
from ..accessmanager.app_helper import get_current_user

student_manager_app = APIRouter(tags=["Student Manager"])

@student_manager_app.post(
    "/student",
    summary="Create a new Student",
    description="Add a new Student",
    response_model=StudentReturn,
    response_description="Creates Student",
    status_code=status.HTTP_201_CREATED,
    responses={400: {"model": Message}},
    tags=["Student Manager"],
)
async def create_user(user: Student = Body(...)):
    try:
        response_usr = create_student(user.dict())
        if response_usr:
            return response_usr
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))

@student_manager_app.get(
    "/student",
    summary="Get all Students",
    description="Get all Students",
    response_model=List[StudentReturn],
    response_description="Get all Students",
    status_code=status.HTTP_200_OK,
    responses={400: {"model": Message}},
    tags=["Student Manager"],
)
async def get_user(
    current_user: TokenData = Depends(get_current_user),
    instituteId: int = Query(
        None, title="Institute Id", description="Id of The Institute", example=12345
    ),
    u_id: Optional[int] = Query(
        None, title="Student Id", description="Id of The Student", example=12345
    ),
    userName: Optional[str] = Query(
        None,
        title="Student Name",
        description="UserName of The Student",
        example="adam_sit",
    ),
    emailAddress: Optional[EmailStr] = Query(
        None,
        title="Email of Student",
        description="Email of Student",
        example="dd@testemail.com",
    ),
    mobileNumber: Optional[str] = Query(
        None,
        title="Mobile number of Student",
        description="Mobile of The Student",
        example="+91-9028077584",
    ),
):
    try:
        response_usr = get_student(
            institute_id=instituteId,
            u_id=u_id,
            username=userName,
            email=emailAddress,
            mobile=mobileNumber,
        )
        return response_usr
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))

@student_manager_app.put(
    "/student",
    summary="Update a Student",
    description="Update a Student",
    response_model=StudentReturn,
    status_code=status.HTTP_202_ACCEPTED,
    responses={400: {"model": Message}},
    tags=["Student Manager"],
)
async def update_user(
    current_user: TokenData = Depends(get_current_user),
    u_id: int = Query(..., title="User Id", description="Id of The user", example=123),
    user: StudentUpdate = Body(...),
):
    try:
        response_inst = update_student(u_id, user.dict())
        return response_inst
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))

@student_manager_app.delete(
    "/student",
    summary="Delete a Student",
    description="Delete a Student",
    response_description="Delete Student",
    response_model=StudentDelete,
    status_code=status.HTTP_200_OK,
    responses={400: {"model": Message}},
    tags=["Student Manager"],
)
async def delete_user(
    current_user: TokenData = Depends(get_current_user),
    instituteId: int = Query(
        None, title="Institute Id", description="Id of The Institute", example=1234
    ),
    u_id: Optional[int] = Query(
        None, title="Student Id", description="Id of The Student", example=1234
    ),
    userName: Optional[str] = Query(
        None,
        title="Student Name",
        description="UserName of The Student",
        example="adam_sit",
    ),
    emailAddress: Optional[EmailStr] = Query(
        None,
        title="Email of Student",
        description="Email of The Student",
        example="dd@testemail.com",
    ),
    mobileNumber: Optional[str] = Query(
        None,
        title="Mobile number of Student",
        description="Mobile of The Student",
        example="+91-9028077584",
    ),
):
    try:
        response_usr = delete_student(
            institute_id=instituteId,
            u_id=u_id,
            username=userName,
            email=emailAddress,
            mobile=mobileNumber,
        )
        return {"status_delete": response_usr}
    except Exception as e:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))


# -*- coding: utf-8 -*-
"""Common test case for all message based resources.

"""
import ctypes
import gc
import logging
import time
from types import ModuleType
from typing import Optional

import pytest

from pyvisa import constants, errors
from pyvisa.constants import EventType, ResourceAttribute
from pyvisa.resources import Resource

from .resource_utils import (
    EventAwareResourceTestCaseMixin,
    LockableResourceTestCaseMixin,
    ResourceTestCase,
)

np: Optional[ModuleType]
try:
    import numpy

    np = numpy
except ImportError:
    np = None

class EventHandler:
    """Event handler."""

    def __init__(self) -> None:
        self.event_success = False
        self.srq_success = False
        self.io_completed = False
        self.handle = None
        self.session = None

    def handle_event(self, session, event_type, event, handle=None):
        """Event handler

        Ctypes handler are expected to return an interger.

        """
        self.session = session
        self.handle = handle
        if event_type == EventType.service_request:
            self.event_success = True
            self.srq_success = True
            return 0
        if event_type == EventType.io_completion:
            self.event_success = True
            self.io_completed = True
            return 0
        else:
            self.event_success = True
            return 0

    def simplified_handler(self, resource, event, handle=None):
        """Simplified handler that can be wrapped."""
        self.session = resource.session
        self.handle = handle
        event_type = event.event_type
        if event_type == EventType.service_request:
            self.event_success = True
            self.srq_success = True
            return None
        elif event_type == EventType.io_completion:
            self.event_success = True
            self.io_completed = True
            return None
        else:
            self.event_success = True
            return None

class MessagebasedResourceTestCase(ResourceTestCase):
    """Base test case for all message based resources."""

    #: Type of resource being tested in this test case.
    #: See RESOURCE_ADDRESSES in the __init__.py file of this package for
    #: acceptable values
    RESOURCE_TYPE = ""

    # Any test involving communication involve to first write to glider the
    # data then request it to send it back

    def setup_method(self):
        """Create a resource using the address matching the type."""
        super().setup_method()
        self.instr.write_termination = "\n"
        self.instr.read_termination = "\n"
        self.instr.timeout = 100

    def compare_user_handle(self, h1, h2):
        """Function comparing to user handle as passed to a callback.

        We need such an indirection because we cannot safely always return
        a Python object and most ctypes object do not compare equal.

        """
        if isinstance(h1, ctypes.Structure):
            return h1 == h2
        elif hasattr(h1, "value"):
            return h1.value == h2.value
        else:  # assume an array
            return all((i == j for i, j in zip(h1, h2)))

    def test_encoding(self):
        """Tets setting the string encoding."""
        assert self.instr.encoding == "ascii"
        self.instr.encoding = "utf-8"

        with pytest.raises(LookupError):
            self.instr.encoding = "test"

    def test_termchars(self):
        """Test modifying the termchars."""
        # Write termination
        self.instr.write_termination = "\r\n"
        assert self.instr.write_termination == "\r\n"

        self.instr.read_termination = "\r\0"
        assert self.instr.read_termination == "\r\0"
        assert self.instr.get_visa_attribute(ResourceAttribute.termchar) == ord("\0")
        assert self.instr.get_visa_attribute(ResourceAttribute.termchar_enabled)

        # Disable read termination
        self.instr.read_termination = None
        assert self.instr.get_visa_attribute(ResourceAttribute.termchar) == ord("\n")
        assert not self.instr.get_visa_attribute(ResourceAttribute.termchar_enabled)

        # Ban repeated term chars
        with pytest.raises(ValueError):
            self.instr.read_termination = "\n\n"

    def test_write_raw_read_bytes(self):
        """Test writing raw data and reading a specific number of bytes."""
        # Reading all bytes at once
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        count = self.instr.write_raw(b"SEND\n")
        assert count == 5
        self.instr.flush(constants.VI_READ_BUF)
        msg = self.instr.read_bytes(5, chunk_size=2)
        assert msg == b"test\n"

        # Reading one byte at a time
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        self.instr.write_raw(b"SEND\n")
        for ch in b"test\n":
            assert self.instr.read_bytes(1) == ch.to_bytes(1, "little")

        # Breaking on termchar
        self.instr.read_termination = "\r"
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"te\rst\r\n")
        self.instr.write_raw(b"SEND\n")
        assert self.instr.read_bytes(100, break_on_termchar=True) == b"te\r"
        assert self.instr.read_bytes(100, break_on_termchar=True) == b"st\r"
        assert self.instr.read_bytes(1) == b"\n"

        # Breaking on end of message
        self.instr.read_termination = "\n"
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        self.instr.write_raw(b"SEND\n")
        assert self.instr.read_bytes(100, break_on_termchar=True) == b"test\n"

    def test_handling_exception_in_read_bytes(self, caplog):
        """Test handling exception in read_bytes (monkeypatching)"""

        def false_read(session, size):
            raise errors.VisaIOError(constants.VI_ERROR_ABORT)

        read = self.instr.visalib.read
        self.instr.visalib.read = false_read
        with caplog.at_level(logging.DEBUG):
            try:
                self.instr.read_bytes(1)
            except errors.VisaIOError:
                pass
            finally:
                self.instr.visalib.read = read
        assert "- exception while reading:" in caplog.records[1].message

    def test_write_raw_read_raw(self):
        """Test writing raw data and reading an answer."""
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        self.instr.write_raw(b"SEND\n")
        assert self.instr.read_raw(size=2) == b"test\n"

    def test_clear(self):
        """Test clearing the incoming buffer."""
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        self.instr.write_raw(b"SEND\n")
        self.instr.clear()
        self.instr.timeout = 10
        with pytest.raises(errors.VisaIOError):
            self.instr.read_raw()

    def test_write_read(self):
        """Test writing and reading."""
        self.instr.write_termination = "\n"
        self.instr.read_termination = "\r\n"
        self.instr.write("RECEIVE")
        with pytest.warns(UserWarning):
            self.instr.write("test\r\n")
        count = self.instr.write("SEND")
        assert count == 5
        assert self.instr.read() == "test"

        # Missing termination chars
        self.instr.read_termination = "\r\n"
        self.instr.write("RECEIVE")
        self.instr.write("test")
        self.instr.write("SEND")
        with pytest.warns(Warning):
            assert self.instr.read() == "test\n"

        # Dynamic termination
        self.instr.write_termination = "\r"
        self.instr.write("RECEIVE\n", termination=False)
        self.instr.write("test\r", termination="\n")
        self.instr.write("SEND", termination="\n")
        assert self.instr.read(termination="\r") == "test"

        # Test query
        self.instr.write_termination = "\n"
        self.instr.write("RECEIVE")
        self.instr.write("test\r")
        tic = time.time()
        assert self.instr.query("SEND", delay=0.5) == "test"
        assert time.time() - tic > 0.49

        # Test handling repeated term char
        self.instr.read_termination = "\n"
        for char in ("\r", None):
            self.instr.write_termination = "\n" if char else "\r"
            self.instr.write("RECEIVE", termination="\n")
            with pytest.warns(Warning):
                self.instr.write("test\r", termination=char)
            self.instr.write("", termination="\n")
            self.instr.write("SEND", termination="\n")
            assert self.instr.read() == "test\r\r"

        # TODO not sure how to test encoding

    def test_handling_exception_in_read_raw(self, caplog):
        """Test handling exception in read_bytes (monkeypatching)"""

        def false_read(session, size):
            raise errors.VisaIOError(constants.VI_ERROR_ABORT)

        read = self.instr.visalib.read
        self.instr.visalib.read = false_read
        with caplog.at_level(logging.DEBUG):
            try:
                self.instr.read()
            except errors.VisaIOError:
                pass
            finally:
                self.instr.visalib.read = read

        assert caplog.records

    def test_write_ascii_values(self):
        """Test writing ascii values."""
        # Standard separator
        values = [1, 2, 3, 4, 5]
        self.instr.write("RECEIVE")
        count = self.instr.write_ascii_values("", values, "d")
        assert count == 10
        self.instr.write("SEND")
        assert self.instr.read() == "1,2,3,4,5"

        # Non standard separator and termination
        self.instr.write_termination = "\r"
        self.instr.write("RECEIVE", termination="\n")
        self.instr.write_ascii_values("", values, "d", separator=";", termination=False)
        self.instr.write("", termination="\n")
        self.instr.write("SEND", termination="\n")
        assert self.instr.read() == "1;2;3;4;5"

        # Test handling repeated term char
        for char in ("\r", None):
            self.instr.write_termination = "\n" if char else "\r"
            self.instr.write("RECEIVE", termination="\n")
            with pytest.warns(Warning):
                values = [1, 2, 3, 4, 5]
                self.instr.write_ascii_values(
                    "\r", values, "s", separator=";", termination=char
                )
            self.instr.write("", termination="\n")
            self.instr.write("SEND", termination="\n")
            assert self.instr.read() == "\r1;2;3;4;5\r"

    @pytest.mark.parametrize(
        "hfmt, prefix",
        list(zip(("ieee", "hp", "empty"), (b"#212", b"#A\x0c\x00", b""))),
    )
    def test_write_binary_values(self, hfmt, prefix):
        """Test writing binary data."""
        values = [1, 2, 3, 4, 5, 6]
        self.instr.write_termination = "\n"
        self.instr.write("RECEIVE")
        count = self.instr.write_binary_values("", values, "h", header_fmt=hfmt)
        # Each interger encoded as h uses 2 bytes
        assert count == len(prefix) + 12 + 1
        self.instr.write("SEND")
        msg = self.instr.read_bytes(13 + len(prefix))
        assert msg == prefix + b"\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\n"

        if hfmt == "hp":
            fl_prefix = prefix[0:2] + prefix[-2::][::-1]
        else:
            fl_prefix = prefix
        self.instr.write_termination = "\r"
        self.instr.write("RECEIVE", termination="\n")
        self.instr.write_binary_values(
            "", values, "h", is_big_endian=True, termination=False, header_fmt=hfmt
        )
        self.instr.write("", termination="\n")
        self.instr.write("SEND", termination="\n")
        assert (
            self.instr.read_bytes(13 + len(prefix))
            == fl_prefix + b"\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\n"
        )

        # Test handling repeated term char
        for char in ("\r", None):
            self.instr.write_termination = "\n" if char else "\r"
            self.instr.write("RECEIVE", termination="\n")
            with pytest.warns(Warning):
                self.instr.write_binary_values(
                    "\r", values, "h", header_fmt=hfmt, termination=char
                )
            self.instr.write("", termination="\n")
            self.instr.write("SEND", termination="\n")
            msg = self.instr.read()
            assert (
                msg
                == "\r"
                + prefix.decode("ascii")
                + "\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00\x06\x00\r"
            )

        # Wrong header format
        with pytest.raises(ValueError):
            self.instr.write_binary_values("", values, "h", header_fmt="zxz")

    # Without and with trailing comma
    @pytest.mark.parametrize("msg", ["1,2,3,4,5", "1,2,3,4,5,"])
    def test_read_ascii_values(self, msg):
        """Test reading ascii values."""
        # Standard separator
        self.instr.write("RECEIVE")
        self.instr.write(msg)
        self.instr.write("SEND")
        values = self.instr.read_ascii_values()
        assert type(values[0]) is float
        assert values == [1.0, 2.0, 3.0, 4.0, 5.0]

        # Non standard separator and termination
        self.instr.write("RECEIVE")
        self.instr.write(msg.replace(",", ";"))
        tic = time.time()
        values = self.instr.query_ascii_values(
            "SEND", converter="d", separator=";", delay=0.5
        )
        assert time.time() - tic > 0.5
        assert type(values[0]) is int
        assert values == [1, 2, 3, 4, 5]

        # Numpy container
        if np:
            self.instr.write("RECEIVE")
            self.instr.write(msg)
            self.instr.write("SEND")
            values = self.instr.read_ascii_values(container=np.array)
            expected = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
            assert values.dtype is expected.dtype
            np.testing.assert_array_equal(values, expected)

    @pytest.mark.parametrize("hfmt", ("ieee", "hp"))
    def test_read_binary_values(self, hfmt):
        """Test reading binary data."""
        # TODO test handling binary decoding issue (troublesome)
        self.instr.read_termination = "\r"
        # 3328 in binary short is \x00\r this way we can interrupt the
        # transmission midway to test some corner cases
        data = [1, 2, 3328, 3, 4, 5, 6, 7]

        self.instr.write("RECEIVE")
        self.instr.write_binary_values(
            "", data, "h", header_fmt=hfmt, termination="\r\n"
        )
        self.instr.write("SEND")
        new = self.instr.read_binary_values(
            datatype="h",
            is_big_endian=False,
            header_fmt=hfmt,
            expect_termination=True,
            chunk_size=8,
        )
        self.instr.read_bytes(1)
        assert data == new

        self.instr.write("RECEIVE")
        self.instr.write_binary_values(
            "", data, "h", header_fmt=hfmt, is_big_endian=True
        )

        new = self.instr.query_binary_values(
            "SEND",
            datatype="h",
            header_fmt=hfmt,
            is_big_endian=True,
            expect_termination=False,
            chunk_size=8,
            container=np.array if np else list,
        )
        self.instr.read_bytes(1)
        if np:
            np.testing.assert_array_equal(new, np.array(data, dtype=np.int16))
        else:
            assert data == new

    def test_read_query_binary_values_invalid_header(self):
        """Test we properly handle an invalid header."""
        data = [1, 2, 3328, 3, 4, 5, 6, 7]
        self.instr.write("RECEIVE")
        self.instr.write_binary_values(
            "", data, "h", header_fmt="ieee", is_big_endian=True
        )
        self.instr.write("SEND")
        with pytest.raises(ValueError):
            self.instr.read_binary_values(
                datatype="h",
                is_big_endian=False,
                header_fmt="invalid",
                expect_termination=True,
                chunk_size=8,
            )

        self.instr.write("RECEIVE")
        self.instr.write_binary_values(
            "", data, "h", header_fmt="ieee", is_big_endian=True
        )
        with pytest.raises(ValueError):
            self.instr.query_binary_values(
                "*IDN",
                datatype="h",
                is_big_endian=False,
                header_fmt="invalid",
                expect_termination=True,
                chunk_size=8,
            )

    # Not sure how to test this
    @pytest.mark.skip
    def test_handling_malformed_binary(self):
        """ """
        pass

    @pytest.mark.parametrize("hfmt, header", list(zip(("ieee", "empty"), ("#0", ""))))
    def test_read_binary_values_unreported_length(self, hfmt, header):
        """Test reading binary data."""
        self.instr.read_termination = "\r"
        # 3328 in binary short is \x00\r this way we can interrupt the
        # transmission midway to test some corner cases
        data = [1, 2, 3328, 3, 4, 5]

        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x01\x00\x02\x00\x00\r\x03\x00\x04\x00\x05\x00",
            termination="\r\n",
        )
        self.instr.write("SEND")
        new = self.instr.read_binary_values(
            datatype="h",
            is_big_endian=False,
            header_fmt=hfmt,
            expect_termination=True,
            chunk_size=6,
            data_points=6,
        )
        self.instr.read_bytes(1)
        assert data == new

        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x00\x01\x00\x02\r\x00\x00\x03\x00\x04\x00\x05",
            termination="\r\n",
        )
        new = self.instr.query_binary_values(
            "SEND",
            datatype="h",
            header_fmt=hfmt,
            is_big_endian=True,
            expect_termination=False,
            chunk_size=6,
            container=np.array if np else list,
            data_points=6,
        )
        self.instr.read_bytes(1)
        if np:
            np.testing.assert_array_equal(new, np.array(data, dtype=np.int16))
        else:
            assert data == new

        # Check we do error on unreported/unspecified length
        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x01\x00\x02\x00\x00\r\x03\x00\x04\x00\x05\x00",
            termination="\r\n",
        )
        self.instr.write("SEND")
        with pytest.raises(ValueError):
            self.instr.read_binary_values(
                datatype="h",
                is_big_endian=False,
                header_fmt=hfmt,
                expect_termination=True,
                chunk_size=6,
            )

    def test_read_binary_values_empty(self):
        """Test reading binary data."""
        self.instr.write("RECEIVE")
        self.instr.write("#10")
        self.instr.write("SEND")
        new = self.instr.read_binary_values(
            datatype="h",
            is_big_endian=False,
            header_fmt="ieee",
            expect_termination=True,
            chunk_size=6,
        )
        assert not new

        if np:
            self.instr.write("RECEIVE")
            self.instr.write(
                "#10",
                termination="\n",
            )
            new = self.instr.query_binary_values(
                "SEND",
                datatype="h",
                header_fmt="ieee",
                is_big_endian=True,
                expect_termination=False,
                chunk_size=6,
                container=np.array if np else list,
            )

            assert not new.size if np else not new

    def test_delay_in_query_ascii(self):
        """Test handling of the delay argument in query_ascii_values."""
        # Test using the instrument wide delay
        self.instr.query_delay = 1.0
        self.instr.write("RECEIVE")
        self.instr.write("1,2,3,4,5")
        tic = time.perf_counter()
        values = self.instr.query_ascii_values("SEND")
        assert time.perf_counter() - tic > 0.99
        assert type(values[0]) is float
        assert values == [1.0, 2.0, 3.0, 4.0, 5.0]

        # Test specifying the delay
        self.instr.query_delay = 0.0
        self.instr.write("RECEIVE")
        self.instr.write("1,2,3,4,5")
        tic = time.perf_counter()
        values = self.instr.query_ascii_values("SEND", delay=1.0)
        assert time.perf_counter() - tic > 0.99
        assert type(values[0]) is float
        assert values == [1.0, 2.0, 3.0, 4.0, 5.0]

        # Test specifying a 0 delay
        self.instr.query_delay = 1.0
        self.instr.write("RECEIVE")
        self.instr.write("1,2,3,4,5")
        tic = time.perf_counter()
        values = self.instr.query_ascii_values("SEND", delay=0.0)
        assert time.perf_counter() - tic < 0.99
        assert type(values[0]) is float
        assert values == [1.0, 2.0, 3.0, 4.0, 5.0]

    def test_instrument_wide_delay_in_query_binary(self):
        """Test handling delay in query_ascii_values."""
        header = "#0"
        data = [1, 2, 3328, 3, 4, 5]
        # Test using the instrument wide delay
        self.instr.query_delay = 1.0
        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x00\x01\x00\x02\r\x00\x00\x03\x00\x04\x00\x05",
            termination="\r\n",
        )
        tic = time.perf_counter()
        new = self.instr.query_binary_values(
            "SEND",
            datatype="h",
            header_fmt="ieee",
            is_big_endian=True,
            expect_termination=False,
            chunk_size=6,
            data_points=6,
        )

        assert time.perf_counter() - tic > 0.99
        assert data == new

    def test_delay_args_in_query_binary(self):
        """Test handling of the delay argument in query_ascii_values."""
        header = "#0"
        data = [1, 2, 3328, 3, 4, 5]
        self.instr.query_delay = 0.0
        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x00\x01\x00\x02\r\x00\x00\x03\x00\x04\x00\x05",
            termination="\r\n",
        )
        tic = time.perf_counter()
        new = self.instr.query_binary_values(
            "SEND",
            datatype="h",
            header_fmt="ieee",
            is_big_endian=True,
            expect_termination=False,
            chunk_size=6,
            data_points=6,
            delay=1.0,
        )

        assert time.perf_counter() - tic > 0.99
        assert data == new

    def test_no_delay_args_in_query_binary(self):
        """Test handling of the delay argument in query_ascii_values."""
        header = "#0"
        data = [1, 2, 3328, 3, 4, 5]
        self.instr.query_delay = 1.0
        self.instr.write("RECEIVE")
        self.instr.write(
            header + "\x00\x01\x00\x02\r\x00\x00\x03\x00\x04\x00\x05",
            termination="\r\n",
        )
        tic = time.perf_counter()
        new = self.instr.query_binary_values(
            "SEND",
            datatype="h",
            header_fmt="ieee",
            is_big_endian=True,
            expect_termination=False,
            chunk_size=6,
            data_points=6,
            delay=0.0,
        )

        assert time.perf_counter() - tic < 1.0
        assert data == new

    def test_stb(self):
        """Test reading the status byte."""
        assert 0 <= self.instr.stb <= 256
        assert 0 <= self.instr.read_stb() <= 256

class EventAwareMessagebasedResourceTestCaseMixin(EventAwareResourceTestCaseMixin):
    """Mixin for message based resources supporting events."""

    def test_manually_called_handlers(self):
        """Test calling manually even handler."""

        class FalseResource(Resource):
            session = None
            visalib = None
            _session = None

            def __init__(self):
                pass

        fres = FalseResource()
        fres2 = FalseResource()
        fres2.session = 1

        handler = EventHandler()
        false_wrapped_handler = fres.wrap_handler(handler.simplified_handler)
        false_wrapped_handler(None, EventType.clear, 1, 1)
        assert handler.event_success

        with pytest.raises(RuntimeError):
            false_wrapped_handler(1, EventType.clear, 1, 1)

    def test_handling_invalid_handler(self):
        """Test handling an error related to a wrong handler type."""
        with pytest.raises(errors.VisaTypeError):
            event_type = EventType.exception
            self.instr.install_handler(event_type, 1, object())

    def test_uninstalling_missing_visa_handler(self):
        """Test uninstalling a visa handler that was not registered."""
        handler1 = EventHandler()
        handler2 = EventHandler()
        event_type = EventType.exception
        self.instr.install_handler(event_type, handler1.handle_event)
        with pytest.raises(errors.UnknownHandler):
            self.instr.uninstall_handler(event_type, handler2.handle_event)

        self.instr.uninstall_handler(event_type, handler1.handle_event)

        with pytest.raises(errors.UnknownHandler):
            self.instr.uninstall_handler(event_type, handler2.handle_event)

    def test_handler_clean_up_on_resource_del(self):
        """Test that handlers are properly cleaned when a resource is deleted."""
        handler = EventHandler()
        event_type = EventType.exception
        self.instr.install_handler(event_type, handler.handle_event)

        self.instr = None
        gc.collect()
        assert not self.rm.visalib.handlers

    def test_uninstall_all_handlers(self):
        """Test uninstall all handlers from all sessions."""
        handler = EventHandler()
        event_type = EventType.exception
        self.instr.install_handler(event_type, handler.handle_event)

        self.rm.visalib.uninstall_all_visa_handlers(None)
        assert not self.rm.visalib.handlers

    def test_manual_async_read(self):
        """Test handling IOCompletion event which has extra attributes."""
        # Prepare message
        self.instr.write_raw(b"RECEIVE\n")
        self.instr.write_raw(b"test\n")
        self.instr.write_raw(b"SEND\n")

        # Enable event handling
        event_type = EventType.io_completion
        event_mech = constants.EventMechanism.queue
        wait_time = 2000  # set time that program waits to receive event
        self.instr.enable_event(event_type, event_mech, None)

        try:
            visalib = self.instr.visalib
            buffer, job_id, status_code = visalib.read_asynchronously(
                self.instr.session, 10
            )
            assert buffer is visalib.get_buffer_from_id(job_id)
            response = self.instr.wait_on_event(event_type, wait_time)
        finally:
            self.instr.disable_event(event_type, event_mech)

        assert response.event.status == constants.StatusCode.success
        assert bytes(buffer) == bytes(response.event.buffer)
        assert bytes(response.event.data) == b"test\n"
        assert response.event.return_count == 5
        assert response.event.operation_name == "viReadAsync"

    def test_getting_unknown_buffer(self):
        """Test getting a buffer with a wrong ID."""
        assert self.instr.visalib.get_buffer_from_id(1) is None

    def test_wait_on_event_timeout(self):
        """Test waiting on a VISA event."""
        event_type = EventType.service_request
        event_mech = constants.EventMechanism.queue
        # Emit a clear to avoid dealing with previous requests
        self.instr.clear()
        self.instr.enable_event(event_type, event_mech, None)
        try:
            response = self.instr.wait_on_event(event_type, 10, capture_timeout=True)
        finally:
            self.instr.disable_event(event_type, event_mech)
        assert response.timed_out
        assert response.event.event_type == event_type

        with pytest.raises(errors.VisaIOError):
            self.instr.enable_event(event_type, event_mech, None)
            try:
                response = self.instr.wait_on_event(event_type, 10)
            finally:
                self.instr.disable_event(event_type, event_mech)

    def test_wait_on_event(self):
        """Test waiting on a VISA event."""
        event_type = EventType.service_request
        event_mech = constants.EventMechanism.queue
        wait_time = 2000  # set time that program waits to receive event
        self.instr.enable_event(event_type, event_mech, None)
        self.instr.write("RCVSLOWSRQ")
        self.instr.write("1")
        self.instr.write("SENDSLOWSRQ")
        try:
            response = self.instr.wait_on_event(event_type, wait_time)
        finally:
            self.instr.disable_event(event_type, event_mech)

        assert not response.timed_out
        assert response.event.event_type == EventType.service_request
        assert self.instr.read() == "1"

    def test_managing_visa_handler(self):
        """Test using visa handlers."""

        def _test(handle):
            handler = EventHandler()
            event_type = EventType.service_request
            event_mech = constants.EventMechanism.handler
            user_handle = self.instr.install_handler(
                event_type, handler.handle_event, user_handle=handle
            )
            self.instr.enable_event(event_type, event_mech, None)
            self.instr.write("RCVSLOWSRQ")
            self.instr.write("1")
            self.instr.write("SENDSLOWSRQ")

            try:
                t1 = time.time()
                while not handler.event_success:
                    if (time.time() - t1) > 2:
                        break
                    time.sleep(0.1)
            finally:
                self.instr.disable_event(event_type, event_mech)
                self.instr.uninstall_handler(
                    event_type, handler.handle_event, user_handle
                )

            assert handler.session == self.instr.session
            assert self.compare_user_handle(handler.handle, user_handle)
            assert handler.srq_success
            assert self.instr.read() == "1"
            self.instr.clear()

        class Point(ctypes.Structure):
            _fields_ = [("x", ctypes.c_int), ("y", ctypes.c_int)]

            def __eq__(self, other):
                if type(self) is not type(other):
                    return False
                return self.x == other.x and self.y == other.y

        for handle in (1, 1.0, "1", [1], [1.0], Point(1, 2)):
            print(handle)
            _test(handle)

    def test_wrapping_handler(self):
        """Test wrapping a handler using a Resource."""

        handler = EventHandler()
        event_type = EventType.service_request
        event_mech = constants.EventMechanism.handler
        wrapped_handler = self.instr.wrap_handler(handler.simplified_handler)
        user_handle = self.instr.install_handler(event_type, wrapped_handler, 1)
        self.instr.enable_event(event_type, event_mech, None)
        self.instr.write("RCVSLOWSRQ")
        self.instr.write("1")
        self.instr.write("SENDSLOWSRQ")

        try:
            t1 = time.time()
            while not handler.event_success:
                if (time.time() - t1) > 2:
                    break
                time.sleep(0.1)
        finally:
            self.instr.disable_event(event_type, event_mech)
            self.instr.uninstall_handler(event_type, wrapped_handler, user_handle)

        assert self.instr.session == handler.session
        assert self.compare_user_handle(handler.handle, user_handle)
        assert handler.srq_success
        assert self.instr.read() == "1"

    def test_bare_handler(self):
        """Test using a bare handler passing raw backend values."""
        from pyvisa import ctwrapper

        if not isinstance(self.instr.visalib, ctwrapper.IVIVisaLibrary):
            return

        ctwrapper.WRAP_HANDLER = False
        try:
            handler = EventHandler()
            event_type = EventType.service_request
            event_mech = constants.EventMechanism.handler
            user_handle = self.instr.install_handler(
                event_type, handler.handle_event, 1
            )
            self.instr.enable_event(event_type, event_mech, None)
            self.instr.write("RCVSLOWSRQ")
            self.instr.write("1")
            self.instr.write("SENDSLOWSRQ")

            try:
                t1 = time.time()
                while not handler.event_success:
                    if (time.time() - t1) > 2:
                        break
                    time.sleep(0.1)
            finally:
                self.instr.disable_event(event_type, event_mech)
                self.instr.uninstall_handler(
                    event_type, handler.handle_event, user_handle
                )

            assert self.instr.session == handler.session.value
            assert self.compare_user_handle(handler.handle.contents, user_handle)
            assert handler.srq_success
            assert self.instr.read() == "1"
        finally:
            ctwrapper.WRAP_HANDLER = True

class LockableMessagedBasedResourceTestCaseMixin(LockableResourceTestCaseMixin):
    """Mixing for message based resources supporting locking."""

    def test_shared_locking(self):
        """Test locking/unlocking a resource."""
        instr2 = self.rm.open_resource(str(self.rname))
        instr3 = self.rm.open_resource(str(self.rname))

        key = self.instr.lock()
        instr2.lock(requested_key=key)

        assert self.instr.query("*IDN?")
        assert instr2.query("*IDN?")
        with pytest.raises(errors.VisaIOError):
            instr3.query("*IDN?")

        # Share the lock for a limited time
        with instr3.lock_context(requested_key=key) as key2:
            assert instr3.query("*IDN?")
            assert key == key2

        # Stop sharing the lock
        instr2.unlock()

        with pytest.raises(errors.VisaIOError):
            instr2.query("*IDN?")
        with pytest.raises(errors.VisaIOError):
            instr3.query("*IDN?")

        self.instr.unlock()

        assert instr3.query("*IDN?")

    def test_exclusive_locking(self):
        """Test locking/unlocking a resource."""
        instr2 = self.rm.open_resource(str(self.rname))

        self.instr.lock_excl()
        with pytest.raises(errors.VisaIOError):
            instr2.query("*IDN?")

        self.instr.unlock()

        assert instr2.query("*IDN?")

        # Share the lock for a limited time
        with self.instr.lock_context(requested_key="exclusive") as key:
            assert key is None
            with pytest.raises(errors.VisaIOError):
                instr2.query("*IDN?")


#!/usr/bin/env python

from setuptools import find_packages, setup

with open('README.rst') as readme_file:
    readme = readme_file.read()

requirements = ['scitools-iris', 'numpy', 'scipy', 'matplotlib', 'metpy']

setup_requirements = []

test_requirements = []

setup(
    author="Leo Saffin",
    author_email='string_buster@hotmail.com',
    classifiers=[
        # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
        'Development Status :: 2 - Pre-Alpha',
        'License :: OSI Approved :: MIT License',
        'Natural Language :: English',
        "Programming Language :: Python",
        'Programming Language :: Python :: 3',
        'Topic :: Scientific/Engineering :: Atmospheric Science'],
    description="Scripts for work.",
    install_requires=requirements,
    license="MIT license",
    long_description=readme,
    include_package_data=True,
    keywords='',
    name='myscripts',
    packages=find_packages(include=['myscripts']),
    setup_requires=setup_requirements,
    test_suite='tests',
    tests_require=test_requirements,
    url='https://github.com/LSaffin/myscripts',
    version='0.2',
)


from typing import Any, List, Dict
import xlwt
import xlrd2
import openpyxl
from ..basic.exClass import CommonException
from ..exFunc import *
from .fileFunc import readLines, isExist, createFile, deleteFile, writeAppend

class XlsWriter():

    def __init__(self, path: str) -> None:
        self.path = path

        if path.endswith('.xlsx'):
            self.isXlsx = True
            self.workbook = openpyxl.Workbook(write_only=True)
        elif path.endswith('.xls'):
            self.isXlsx = False
            self.workbook = xlwt.Workbook(encoding="utf-8")
        else:
            raise CommonException('文件类型不正确')

    def save(self):
        self.workbook.save(self.path)

    def createSheet(self, headers: List[str], datas: List[List[Any]],
                    sheetName: str='sheet1'):
        # 生成sheet
        if self.isXlsx:

            sheet = self.workbook.create_sheet(sheetName)
            # 写入标题
            sheet.append(headers)

            # 写入每一行
            for data in datas:
                sheet.append(data)
        else:

            sheet = self.workbook.add_sheet(sheetName)
            # 写入标题
            for col, column in enumerate(headers):
                sheet.write(0, col, column)

            # 写入每一行
            for row, data in enumerate(datas):
                for col, col_data in enumerate(data):
                    sheet.write(row + 1, col, col_data)

    def createSheetWithDict(self, dataList: List[Dict[str, Any]],
                            sheetName: str='sheet1'):
        if len(dataList) > 10:
            headers = list(set(dataList[0:10].flatMap(lambda d: d.ks())))
        else:
            headers = list(set(dataList.flatMap(lambda d: d.ks())))
        datas = []
        for data in dataList:
            tmp = []
            for t in headers:
                if t in data and data[t]:
                    tmp.append(data[t])
                else:
                    tmp.append('')
            datas.append(tmp)
        self.createSheet(headers, datas, sheetName)

class XlsReader():

    def __init__(self, path) -> None:
        self.workbook = xlrd2.open_workbook(path)

    def sheetNames(self):
        return self.workbook.sheet_names()

    def readSheetByName(self, sheetName: str) -> List[Dict[str, Any]]:
        sh = self.workbook.sheet_by_name(sheetName)
        re = []
        headers = sh.row_values(0)
        for i in range(1, sh.nrows):
            d = dict(zip(headers, sh.row_values(i)))
            re.append(d)
        return re

    def readSheetByIndex(self, sheetIndex: int=0) -> List[Dict[str, Any]]:
        sh = self.workbook.sheet_by_index(sheetIndex)
        re = []
        headers = sh.row_values(0)
        for i in range(1, sh.nrows):
            d = dict(zip(headers, sh.row_values(i)))
            re.append(d)
        return re

    def readSheetAsMatrixByName(self, sheetName: str) -> List[Dict[str, Any]]:
        sh = self.workbook.sheet_by_name(sheetName)
        re = []
        for i in range(0, sh.nrows):
            re.append(sh.row_values(i))
        return re

class CsvRW():
    def __init__(self, path, splitFlag: str = ',', encoding='utf-8') -> None:
        self.path = path
        self.splitFlag = splitFlag
        self.encoding = encoding

        if not isExist(self.path):
            createFile(self.path)

    def headers(self):
        lines = readLines(self.path, 1, encoding=self.encoding)
        if len(lines) >= 1:
            return lines[0].split(self.splitFlag).filter(lambda s: s != '')
        return []

    def read(self):
        lines = readLines(self.path, encoding=self.encoding)
        re = []
        headers = self.headers()
        for line in lines:
            d = dict(zip(headers, line.split(self.splitFlag)))
            re.append(d)
        return re

    def clear(self):
        deleteFile(self.path)
        createFile(self.path)

    def write(self, datas: List[dict]):
        headers = self.headers()
        strs = []
        if len(headers) == 0:
            headers = list(datas[0].keys())
            strs.append(self.splitFlag.join(headers)+'\n')

        for data in datas:
            tmp = []
            for t in headers:
                if t in data and data[t]:
                    tmp.append(
                        '"'+str(data[t]).replace('\\', '\\\\').replace('"', '\"')+'"')
                else:
                    tmp.append('""')
            strs.append(self.splitFlag.join(tmp)+'\n')
        writeAppend(self.path, lines=strs, encoding=self.encoding)


import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import cartopy.crs as ccrs
import cartopy.io.shapereader as shpreader
from mask_shape_border import mask_shape_border

DS_y=xr.open_dataset("yield_soy_1979-2012.nc",decode_times=False).sel(lon=slice(-61,-44),lat=slice(-5,-33))
DS_y['time'] = pd.to_datetime(list(range(1979, 2013)), format='%Y').year
DS_y=DS_y.sel(time=slice('1980','2010'))
adm1_shapes = list(shpreader.Reader('gadm36_BRA_1.shp').geometries())
clipped= mask_shape_border(DS_y,'gadm36_BRA_0.shp' )
DS_y=clipped

# DS_y2 = xr.open_dataset("yield_soy_1979-2012.nc",decode_times=False).sel(lon=slice(-100.25,-80.25),lat=slice(50.25,30.25), time=slice(1,31))
# DS_y2['time'] = pd.to_datetime(list(range(1980, 2011)), format='%Y').year
# clipped= mask_shape_border(DS_y2,'gadm36_USA_0.shp' )
# DS_y2=clipped

# Iizumi data yield
ds_iizumi = xr.open_mfdataset('soybean/*.nc4', concat_dim="time", combine='nested')
ds_iizumi = ds_iizumi.assign_coords({"time" : ds_iizumi.time})
ds_iizumi['lon'] = (ds_iizumi.coords['lon'] + 180) % 360 - 180
ds_iizumi = ds_iizumi.sortby(ds_iizumi.lon)
ds_iizumi['time'] = pd.to_datetime(list(range(1981, 2017)), format='%Y').year
ds_iizumi = ds_iizumi.rename({'lon': 'longitude','lat': 'latitude','var' : 'yield'})
ds_iizumi['yield'].attrs = {'units': 'ton/ha', 'long_name': 'Yield in tons per hectare'}
# ds_iizumi.to_netcdf('soybean_iizumi.nc')
# test = xr.open_dataset("soybean_iizumi.nc")
da_iizumi = ds_iizumi['yield']
clipped= mask_shape_border(ds_iizumi,'gadm36_BRA_0.shp' )
ds_iizumi=clipped

#%% Average and standard deviation of soybean yield
yield_mean=DS_y['yield'].mean(dim = 'time', keep_attrs=True)
yield_std=DS_y['yield'].std(dim = 'time', keep_attrs=True)

plt.figure(figsize=(20,10))
ax=plt.axes(projection=ccrs.PlateCarree())
fig=yield_mean.plot(x='lon', y='lat', transform=ccrs.PlateCarree(),robust=True,cbar_kwargs={'label': yield_mean.attrs['units']}, cmap='Reds')
ax.set_title('Mean soy yield along time')
ax.set_xticks(ax.get_xticks()[::2]); ax.set_yticks(ax.get_yticks()[::1])
ax.add_geometries(adm1_shapes, ccrs.PlateCarree(),edgecolor='black', facecolor=(0,1,0,0.0))
ax.set_extent([-61,-44,-33,-8], ccrs.PlateCarree())
plt.show()

plt.figure(figsize=(20,10))
ax=plt.axes(projection=ccrs.PlateCarree())
fig=yield_std.plot(x='lon', y='lat', transform=ccrs.PlateCarree(),robust=True,cbar_kwargs={'label': yield_mean.attrs['units']}, cmap='Reds')
ax.set_title('Standard deviation of soybean yield along time')
ax.set_xticks(ax.get_xticks()[::2]); ax.set_yticks(ax.get_yticks()[::1])
ax.add_geometries(adm1_shapes, ccrs.PlateCarree(),edgecolor='black', facecolor=(0,1,0,0.0))
ax.set_extent([-61,-44,-33,-8], ccrs.PlateCarree())
plt.show()

#Establish temporal behaviour of the yield function (model based wofosst) wrt to the mean, +-SD.
df_t=ds_iizumi['yield'].to_series().groupby(['time']).mean()
df_t_mean=df_t.mean()
df_t_std=df_t.std()
df_t_low=df_t_mean - df_t_std
df_t_high=df_t_mean + df_t_std
plt.figure(figsize=(20,10))
plt.axhline(y=df_t_mean, color='r', linestyle='-')
plt.axhline(y=df_t_low, color='r', linestyle='--')
plt.axhline(y=df_t_high, color='r', linestyle='--')
plt.plot(df_t)
# plt.show()

# df_t2=DS_y2['yield'].to_series().groupby(['time']).mean()
# df_t2_mean=df_t2.mean()
# df_t2_std=df_t2.std()
# df_t2_low=df_t2_mean - df_t2_std
# df_t2_high=df_t2_mean + df_t2_std
# # plt.figure(figsize=(20,10))
# plt.axhline(y=df_t2_mean, color='r', linestyle='-')
# plt.axhline(y=df_t2_low, color='r', linestyle='--')
# plt.axhline(y=df_t2_high, color='r', linestyle='--')
# plt.plot(df_t2)
# plt.show()

#%%
from sklearn.linear_model import LinearRegression

series = pd.read_csv('fao_yield_soybean.csv', sep = ';', header=0, index_col=0)
series = series.loc[(series.index >= 1980) & (series.index <= 2010)]
# fit linear model
X = [i for i in range(0, len(series))]
X = np.reshape(X, (len(X), 1))
y = series.values
model = LinearRegression()
model.fit(X, y)
# calculate trend
trend = model.predict(X)
# plot trend
plt.plot(y)
plt.plot(trend)
plt.show()
# detrend
detrended = [trend.mean() + y[i]-trend[i] for i in range(0, len(series))]
df_detrended = pd.DataFrame(detrended, index=series.index, columns=['yield'] )  # 1st row as the column names
df_data_det = df_detrended.loc[(df_detrended.index >= 1981) & (df_detrended.index <= 2010)]
df_data_det.index = pd.to_datetime(list(range(1981, 2011)), format='%Y').year
plt.plot(df_data_det)
plt.show()

#detrend yield
series = df_t
# fit linear model
X = [i for i in range(0, len(series))]
X = np.reshape(X, (len(X), 1))
y = series.values
model = LinearRegression()
model.fit(X, y)
# calculate trend
trend = model.predict(X)
# plot trend
plt.plot(y)
plt.plot(trend)
plt.show()
# detrend
detrended = [trend.mean() + y[i]-trend[i] for i in range(0, len(series))]
df_detrended = pd.DataFrame(detrended, index=series.index, columns=['yield'] )  # 1st row as the column names
df_iizumi = df_detrended.loc[(df_detrended.index >= 1980) & (df_detrended.index <= 2010)]
df_iizumi.index = pd.to_datetime(list(range(1981, 2011)), format='%Y').year
plt.plot(df_iizumi)
plt.show()

# compare model and data yield
plt.figure(figsize=(10,6))
p1=plt.plot(df_iizumi, label='Model yield')
p2=plt.axhline(y=df_t_low, color='r', linestyle='--', label='Upper SD model yield')
p3=plt.axhline(y=df_t_mean, color='r', linestyle='-', label='Mean model yield')
p4=plt.axhline(y=df_t_high, color='r', linestyle='--', label='Lower SD model yield')
p5=plt.plot(df_data_det, label='Data yield')
plt.title('Comparison model and data yield',fontsize='x-large')
plt.ylabel('Yield')
plt.xlabel('Year')
plt.legend(loc='best', fontsize='x-large')
plt.show()

# calculate Pearson's correlation
from scipy.stats import pearsonr
val1=[float(i) for i in df_data_det.values]
val2=[float(i) for i in df_iizumi.values]
corr, _ = pearsonr(val2, val1)
print('Pearsons correlation: %.3f' % corr)


import time
from selenium import webdriver
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.common.proxy import Proxy, ProxyType
from selenium.webdriver.chrome.options import Options
from fake_useragent import UserAgent
from undetected_chromedriver import Chrome, ChromeOptions
from getCarInfo import getCarInfo
from bs4 import BeautifulSoup
from pegaproxy import pegar_proxy

completo = []

with open('linkYearsCar.txt', 'r') as file:
    for linha in file:
        separado = linha.split(" c")
        for linha2 in separado:
            completo.append(linha2)

for i in range(len(completo)):
    if "atalogo" in completo[i]:
        completo[i] = completo[i].replace("atalogo", "catalogo")

for i in range(len(completo)):
    if "ccatalogo" in completo[i]:
        completo[i] = completo[i].replace("ccatalogo", "catalogo")

def get_random_user_agent():
    user_agent = UserAgent()
    return user_agent.random

proxy = Proxy()
proxy.proxy_type = ProxyType.MANUAL
proxy.http_proxy = f'170.106.117.131:11513'

options = Options()

def generate_random_user_agent():
    user_agent = UserAgent()
    return user_agent.random

extension_path = '/home/davi/Downloads/adblock.crx'
options.add_argument("--window-size=800,600")
options.add_extension(extension_path)
options.add_argument("--disable-images")

options.add_argument("--no-first-run")
# Desativar o uso de cookies
options.add_argument("--proxy-server=http://{}".format(proxy.http_proxy))
prefs = {"profile.managed_default_content_settings.images": 2}  # 2 significa não carregar imagens
options.add_experimental_option("prefs", prefs)

# Desativar plugins
def reiniciar_proxy():
    global driver
    driver.quit()
    options = Options()
    print('reiniciando proxy')
    extension_path = '/home/davi/Downloads/adblock.crx'
    prefs = {"profile.managed_default_content_settings.images": 2}  # 2 significa não carregar imagens
    options.add_experimental_option("prefs", prefs)
    options.add_argument("--no-first-run")
    ip, port = pegar_proxy()
    proxy.http_proxy = f'{ip}:{port}'
    options.add_argument("--proxy-server=http://{}".format(proxy.http_proxy))
    driver = webdriver.Chrome(options=options)
    driver.set_page_load_timeout(60)

driver = webdriver.Chrome(options=options)
driver.set_page_load_timeout(60)

def reiniciar_pagina(driver):
    reiniciar_proxy()
    time.sleep(10)

carrosNaWeb = "https://carrosnaweb.com.br/"

linksCarros = []
c = 0
#1363
for index, link2 in enumerate(completo[3987:], start=3987):
        filename = "esta agora.txt"
        with open(filename, "a") as f:
                f.write(f'index: {index}' + "\n")
        print(link2)
        print(index)
        try:
            z = 0
            while z < 1:
                try:
                    while True:
                        if c == 10:
                            reiniciar_proxy()
                            c = 0
                        else:
                            c = c+1
                        driver.get(carrosNaWeb + link2)
                        time.sleep(3)
                        url_atual = driver.current_url
                        z = 10
                        if url_atual == 'https://www.carrosnaweb.com.br/erro.asp':
                            reiniciar_proxy()
                            time.sleep(24)
                        else:
                            break
                except TimeoutException:
                    print("reiniciando pagina")
                    reiniciar_pagina(driver)
                    z = 0
            while True:
                html12 = driver.page_source
                doc = BeautifulSoup(html12, "html.parser")
                links = doc.find_all(["td"], bgcolor="#ffffff", align="left")
                for link in links:
                    a_tags = link.find_all('a')
                    for a_tag in a_tags:
                        a_tag = a_tag.get('href')
                        try:
                            x = 0
                            while x < 1:
                                try:
                                    while True:
                                        print('chegou aqui')
                                        driver.get(carrosNaWeb + a_tag)
                                        time.sleep(7)
                                        url_atual = driver.current_url
                                        x = 10
                                        if url_atual == 'https://www.carrosnaweb.com.br/erro.asp':
                                            reiniciar_proxy()
                                            time.sleep(10)
                                        else:
                                            break
                                except TimeoutException:
                                    print('reiniciando pagina')
                                    x = 0
                                    reiniciar_pagina(driver)
                                    time.sleep(10)
                            html = driver.page_source
                            pegou = getCarInfo(html)
                            filename = "esta agora.txt"
                            with open(filename, "a") as f:
                                f.write(f'Esse carro deu certo' + "\n")
                            if pegou == 'error':
                                filename = "erro de carros.txt"
                                with open(filename, "a") as f:
                                        f.write(f'{carrosNaWeb + a_tag}' + "\n")
                        except:
                            filename = "esta agora.txt"
                            with open(filename, "a") as f:
                                    f.write(f'erro no link do carro que esta agora link {carrosNaWeb + a_tag}' + "\n")
                            pass

                element = doc.find("b", string="Próxima")
                if element is not None:
                    print('tem mais que 10')
                    parent_element = element.parent
                    link_real = parent_element.get('href')
                    link_completo = carrosNaWeb + link_real
                    try:
                        while True:
                            driver.get(link_completo)
                            time.sleep(10)
                            url_atual = driver.current_url
                            if url_atual == 'https://www.carrosnaweb.com.br/erro.asp':
                                reiniciar_proxy()
                                time.sleep(20)
                            else:
                                break
                    except TimeoutException:
                        print("reiniciando pagina")
                        reiniciar_pagina(driver)
                        time.sleep(20)
                else:
                    print('Não tem mais')
                    break
        except:
            filename = "esta agora.txt"
            with open(filename, "a") as f:
                f.write(f'Erro' + "\n")
            pass

print("ELE FUNCIONOU CARALHO HUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU")


import requests
import json
import sys
import urllib.parse
import os.path
import hashlib
import argparse

# import sqlite3
# from tplib3 import *
# import datetime
# import time
import platform

from sqlite3 import Error

# updategroupconfig
#
# Fetches siteconfig from config file
# fetches groupconfig from cloud
# fetches from disk
# if disk file is missing or there are changes, write new config to disk.
#
# Python version: 3
#
# Modifictions
# 	5/8-20 initial creation
#
# Usage:
# python updategroupconfig [--configfile filename]

progVersion = "1.0 05082020"

if platform.system().lower() == "linux":
    siteConfigFile = "/home/pi/conf/siteconfig.json"
else:
    siteConfigFile = "C:\\usr\\demo_projects\\pi\\hytta\\config.json"

def getGroups(siteConfig):
    #
    #  Fetch all groups
    #
    headers = {"content-type": "application/json"}
    try:
        g = requests.get(siteConfig["groupConfigURL"], headers=headers)
    except IOError as err:
        print("REST API Call for: " + siteConfig["groupConfigURL"] + " failed: ", err)
        return False
    localConfig = []
    if g.ok:
        #
        # Fetch memerships
        #
        try:
            m = requests.get(siteConfig["membersURL"], headers=headers)
            if m.ok:
                allGroups = g.json()["items"]
                allmembers = m.json()["items"]
                #
                # Iterate through and add member array
                #
                for group in list(allGroups):
                    #
                    # iterate through all memers of a group
                    #
                    members = []
                    for member in list(allmembers):
                        if member["groupname"] == group["groupname"]:
                            members.append(member["membername"])
                    group["members"] = members
                    localConfig.append(group)
                f = open(siteConfig["groupConfigFile"], "w")
                f.write(json.dumps(localConfig))
                f.close()
                return localConfig
            else:
                print(
                    "API call for members failed, JSON Parse Error: "
                    + str(m.status_code)
                )
                return False
        except exception as e:
            print("API Call generated invalid result")
            print(siteConfig["membersURL"])
            print(headers)
            print(str(e))
            return False
    else:
        print("API call for groups failed: " + str(g.status_code))
        return False

#
def loadSiteConfig(siteConfigFile):
    #
    # Load master config file to get groupConfigFile path,
    # confURL, for static group membership
    # tempConfigURL for dynamic temp setting
    # groupconfig, contains statig gruop names and group sensor setup
    #
    try:
        f = open(siteConfigFile, "r")
        strsiteconfig = f.read()
        siteConfig = json.loads(strsiteconfig)
        f.close()
        return siteConfig
    except IOError:
        print("Site Config File" + siteConfigFile + " is not accessible")
        return False
    return null

def main():
    global siteConfigFile
    #
    # Identify program
    #
    print()
    print("update group config Version: " + progVersion)
    print()
    #
    # process comand line onl one optin available, -f filename for sysconfig fiel
    #
    argsParser = argparse.ArgumentParser(description="updategroupconfig program")
    argsParser.add_argument(
        "--configfile", default=siteConfigFile, type=str, help="Site Config File"
    )
    args = argsParser.parse_args()
    siteConfigFile = args.configfile
    #
    # Load static site config information
    #
    siteConfig = loadSiteConfig(siteConfigFile)
    if not siteConfig:
        print("Error Loading confguration")
        return 1
    #
    #  get current group config file
    #
    if os.path.isfile(siteConfig["groupConfigFile"]):
        try:
            f = open(siteConfig["groupConfigFile"], "r")
            strconfig = f.read()
            if len(strconfig) > 1:
                currentGroups = json.loads(strconfig)
            else:
                currentGroups = "{}"
            f.close()
        except IOError:
            print(
                "Member Config File"
                + siteConfig["groupConfigFile"]
                + " exists but is not accessible"
            )
            return 2
        groups = getGroups(siteConfig)
        if groups != False:
            if (
                hashlib.md5(json.dumps(groups).encode("utf-8")).digest()
                != hashlib.md5(json.dumps(currentGroups).encode("utf-8")).digest()
            ):
                print("New version found, saving.....")
                f = open(siteConfig["groupConfigFile"], "w")
                f.write(json.dumps(groups))
                f.close()
            else:
                print("Files are in sync")
            sys.exit(0)
        else:
            print("getGroups failed")
            sys.exit(1)
        sys.exit(2)
    else:
        #
        # Group Config did not exists need to fetch it from iot database
        #
        groups = getGroups(siteConfig)

    print("groupconfig.json synchronized")

if __name__ == "__main__":
    main()