#!/usr/bin/env python3
# Christopher Vollmers
# Roger Volden

import os
import sys
import numpy as np
import mappy as mp
import argparse

parser = argparse.ArgumentParser()

parser.add_argument('--infile', '-i', type=str, action='store')
parser.add_argument('--path', '-p', type=str, action='store')
parser.add_argument('--downstream_buffer', '-d', type=int, action='store')
parser.add_argument('--upstream_buffer', '-u', type=int, action='store')
parser.add_argument('--subreads', '-b', type=str, action='store')
parser.add_argument('--fasta_files', '-f', type=str, action='store')
parser.add_argument('--minimum_feature_count', '-n', type=int, action='store')
parser.add_argument('--minimum_read_count', '-R', type=int, action='store')
parser.add_argument('--consensusMode', '-C', type=str, default='P', action='store')

if len(sys.argv) == 1:
    parser.print_help()
    sys.exit(0)
args = parser.parse_args()
infile = args.infile
path = args.path + '/'  # path where you want your output files to go
downstream_buffer = args.downstream_buffer
upstream_buffer = args.upstream_buffer
subreads = args.subreads
fasta_files = args.fasta_files
minimum_feature_count = args.minimum_feature_count
minimum_read_count = args.minimum_read_count
consensusMode = args.consensusMode
extend=True
outbed = open(path + 'start_end.bed', 'w')

if '.fofn' in fasta_files:
    fastaList=[]
    for line in open(fasta_files):
        fasta=line.strip()
        fastaList.append(fasta)
else:
    fastaList=fasta_files.split(',')

if '.fofn' in subreads:
    subreadList=[]
    for line in open(subreads):
        sub=line.strip()
        subreadList.append(sub)
else:
    subreadList=subreads.split(',')

def find_peaks(starts, ends, identity, count_dict):
    start_peaks, end_peaks = {}, {}
    if '+' in identity:
        count_dict['+'].add(identity)
    else:
        count_dict['-'].add(identity)

    start_count = {}
    end_count = {}
    for position in starts:
        if position not in start_count:
            start_count[position] = 0
        start_count[position] += 1

    for position in ends:
        if position not in end_count:
            end_count[position] = 0
        end_count[position] += 1

    show=False
    if '5l5292' in identity:
        show=True

    for position in sorted(starts):
        if position - upstream_buffer not in start_peaks:
            window_count = 0
            for i in np.arange(0,10):
                window_position = position + i
                if window_position in start_count:
                    window_count += start_count[window_position]
            if window_count >= minimum_feature_count:
                original_bin=[]
                for shift in np.arange(-upstream_buffer, downstream_buffer):
                    start_peaks[position + shift] = position
                    original_bin.append(position + shift)
                if extend:
                    bins=[]
                    for i in np.arange(min(original_bin),max(original_bin)):
                        bin=0
                        for shift in np.arange(0,10):
                            if i+shift in start_count:
                                 bin+=start_count[i+shift]
                        bins.append(bin)

                    best_bin=max(bins)
                    extended=True
                    adjacent=position - upstream_buffer
                    while extended:
                        window_count=0
                        adjacent_list=[]
                        for i in np.arange(1,11):
                            adjacent_pos = adjacent - i
                            adjacent_list.append(adjacent_pos)
                            if adjacent_pos in start_count:
                                window_count += start_count[adjacent_pos]
                        if best_bin > window_count >= minimum_feature_count:
                            for element in adjacent_list:
                                if element not in start_peaks:
                                    start_peaks[element] = position
                                    original_bin.append(element)
                                else:
                                    extended=False
                        else:
#                            if window_count > best_bin:
#                                print(window_count, best_bin, identity)
                            extended=False
                        adjacent = adjacent_pos
                        if extended:
                            count_dict['start_left'].add(identity)

                    extended = True
                    adjacent = position + downstream_buffer - 1
                    while extended:
                        window_count=0
                        adjacent_list=[]
                        for i in np.arange(1,11):
                            adjacent_pos = adjacent + i
                            adjacent_list.append(adjacent_pos)
                            if adjacent_pos in start_count:
                                window_count += start_count[adjacent_pos]
                        if best_bin > window_count >= minimum_feature_count:
                            for element in adjacent_list:
                                if element not in start_peaks:
                                    start_peaks[element] = position
                                    original_bin.append(element)
                                else:
                                     extended=False
                        else:
#                            if window_count > best_bin:
#                                print(identity)
                            extended=False
                        adjacent = adjacent_pos
                        if extended:
                            count_dict['start_right'].add(identity)
                outbed.write(identity.split('_')[0]+'\t'+str(max(0,min(original_bin)))+'\t'+str(max(original_bin))+'\t'+identity+'\n')

    for position in sorted(ends, reverse=True):
        if position + upstream_buffer -1 not in end_peaks:
            window_count = 0
            for i in np.arange(0,10):
                window_position = position - i
                if window_position in end_count:
                    window_count += end_count[window_position]
            if window_count >= minimum_feature_count:
                original_bin=[]
                for shift in np.arange(-downstream_buffer, upstream_buffer):
                    end_peaks[position + shift] = position
                    original_bin.append(position + shift)
                if extend:
                    bins=[]
                    for i in np.arange(min(original_bin),max(original_bin)):
                        bin=0
                        for shift in np.arange(0,10):
                            if i+shift in end_count:
                                 bin+=end_count[i+shift]
                        bins.append(bin)

                    best_bin=max(bins)
                    extended = True
                    adjacent = position - downstream_buffer
                    while extended:
                        window_count=0
                        adjacent_list=[]
                        for i in np.arange(1,11):
                            adjacent_pos = adjacent - i
                            adjacent_list.append(adjacent_pos)
                            if adjacent_pos in end_count:
                                window_count += end_count[adjacent_pos]
                        if best_bin > window_count >= minimum_feature_count:
                            for element in adjacent_list:
                                if element not in end_peaks:
                                    end_peaks[element] = position
                                    original_bin.append(element)
                                else:
                                    extended=False
                        else:
#                            if window_count > best_bin:
#                                print(identity)
                            extended=False
                        adjacent=adjacent_pos
                        if extended:
                            count_dict['end_left'].add(identity)

                    extended=True
                    adjacent=position + upstream_buffer - 1
                    while extended:
                        window_count=0
                        adjacent_list=[]
                        for i in np.arange(1,11):
                            adjacent_pos = adjacent + i
                            adjacent_list.append(adjacent_pos)
                            if adjacent_pos in end_count:
                                window_count += end_count[adjacent_pos]
                        if best_bin > window_count >= minimum_feature_count:
                            for element in adjacent_list:
                                if element not in end_peaks:
                                    original_bin.append(element)
                                    end_peaks[element] = position
                                else:
                                    extended=False
                        else:
#                            if window_count > best_bin:
#                                print(identity)
                            extended=False
                        adjacent = adjacent_pos
                        if extended:
                            count_dict['end_right'].add(identity)
                outbed.write(identity.split('_')[0]+'\t'+str(max(0,min(original_bin)))+'\t'+str(max(original_bin))+'\t'+identity+'\n')

    return start_peaks, end_peaks, count_dict

def collect_splice_events(path):
    splice_dict = {}
    for line in open(path + '/SS.bed'):
        a = line.strip().split('\t')
        chromosome = a[0]
        if not splice_dict.get(chromosome):
            splice_dict[chromosome] = {}
        splice_left = int(a[1])
        splice_right = int(a[2])
        for base in np.arange(splice_left, splice_right + 1):
            splice_dict[chromosome][base] = a[3].split('_')[0]
    return splice_dict

def sort_reads_into_splice_junctions(splice_dict, fastaList, infile):
    start_end_dict, start_end_dict_mono, readDict = {}, {}, {}
    for fasta_file in fastaList:
        tempSeqs, headers, sequences = [], [], []
        for name,seq,qual in mp.fastx_read(fasta_file):
            readDict[('-').join(name.split('_')[0].split('-')[:5])] = [name, seq]

    for line in open(infile):
        a = line.strip().split('\t')
        read_chromosome, read_direction = a[13], a[8]
        name = a[9].split('_')[0]
        name = ('-').join(name.split('-')[:5])
        show = False
        read_direction = '+'  # ignores read direction
        start, end = int(a[15]), int(a[16])
        if read_direction == '+':
            left_extra, right_extra = int(a[11]), int(a[10]) - int(a[12])
        if read_direction == '-':
            right_extra, left_extra = int(a[11]), int(a[10]) - int(a[12])

        failed = False
        identity = read_chromosome + '_'

        blocksizes = a[18].split(',')[:-1]
        blockstarts = a[20].split(',')[:-1]

        for x in range(0, len(blocksizes) - 1):
            blockstart = int(blockstarts[x])
            blocksize = int(blocksizes[x])
            left_splice = blockstart + blocksize
            right_splice = int(blockstarts[x + 1])
            if right_splice - left_splice > 50:
                if read_chromosome not in splice_dict:
                    failed = True
                    break
                else:
                    if not splice_dict[read_chromosome].get(left_splice) or not splice_dict[read_chromosome].get(right_splice):
                        failed = True
                        break
                left_splice_site = splice_dict[read_chromosome][left_splice]
                right_splice_site = splice_dict[read_chromosome][right_splice]
                identity += str(left_splice_site) + '-' + str(right_splice_site) + '~'
        if show:
            print(failed)
        if not failed:
            if identity.split('_')[1] != '':
                if identity not in start_end_dict:
                    start_end_dict[identity] = []
                start_end_dict[identity].append((start, end,
                                                 '>' + readDict[name][0] + '\n'
                                                 + readDict[name][1] + '\n',
                                                 left_extra,
                                                 right_extra,
                                                 read_direction))
            else:
                if identity not in start_end_dict_mono:
                    start_end_dict_mono[identity] = []

                start_end_dict_mono[identity].append((start, end,
                                                      '>' + readDict[name][0] + '\n'
                                                      + readDict[name][1] + '\n',
                                                      left_extra,
                                                      right_extra,
                                                      read_direction))

    return start_end_dict, start_end_dict_mono

def group_mono_exon_transcripts(start_end_dict, start_end_dict_mono):
    for identity in start_end_dict_mono:
        previous_end = 0
        iso_counter = 0
        new_identity = identity + 'M' + str(iso_counter)
        positions = start_end_dict_mono[identity]
        for position in sorted(positions):
            start = position[0]
            end = position[1]
            if start > previous_end:
                iso_counter += 1
                new_identity = identity + 'M' + str(iso_counter)
                if new_identity not in start_end_dict:
                    start_end_dict[new_identity] = []
                start_end_dict[new_identity].append(position)
                previous_end = max(end, previous_end)
            else:
                if new_identity not in start_end_dict:
                    start_end_dict[new_identity] = []
                start_end_dict[new_identity].append(position)
                previous_end = end

    return start_end_dict

def define_start_end_sites(start_end_dict, start_end_dict_mono, individual_path):
    left_extras, right_extras = {}, {}
    file_set = {}
    isoform_counter, isoform_dict, subread_pointer = 0, {}, {}

    start_end_dict = group_mono_exon_transcripts(start_end_dict, start_end_dict_mono)

    number_of_isoforms = len(start_end_dict)
    counter = 0
    count_dict={}
    count_dict['start_left']=set()
    count_dict['start_right']=set()
    count_dict['end_left']=set()
    count_dict['end_right']=set()
    count_dict['+']=set()
    count_dict['-']=set()

    for identity in sorted(start_end_dict):
        counter += 1
        print('\tprocessing reads assigned to spliceform ', counter,
              ' of ', number_of_isoforms, ' to generate isoform models', end='\r')
        starts, ends = [], []
        ends = []
        positions = start_end_dict[identity]
        length_of_positions = len(positions)
        indexes = np.random.choice(range(0, length_of_positions, 1),
                                   size=min(10000, length_of_positions),
                                   replace=False)
        subsampled_positions = []
        for index in indexes:
            subsampled_positions.append(positions[index])
        for position in subsampled_positions:
            starts.append(int(position[0]))
            ends.append(int(position[1]))

        start_dict, end_dict, count_dict = find_peaks(starts, ends, identity, count_dict)
        matched_positions = []
        left_extras[identity], right_extras[identity] = {}, {}
        for start, end, read, left_extra, right_extra, read_direction in positions:
            if int(start) in start_dict and int(end) in end_dict:
                left = start_dict[int(start)]
                right = end_dict[int(end)]
                if not left_extras[identity].get((left, right)):
                    left_extras[identity][(left, right)] = []
                    right_extras[identity][(left, right)] = []

                left_extras[identity][(left, right)].append(int(left_extra))
                right_extras[identity][(left, right)].append(int(right_extra))
                matched_positions.append((left, right, read, read_direction))

        median_left = {}
        median_right = {}
        for combination, values in left_extras[identity].items():
            median = np.median(values)
            median_left[combination] = median
        for combination, values in right_extras[identity].items():
            median = np.median(values)
            median_right[combination] = median

        for left, right, read, read_direction in matched_positions:
            medianLeft = median_left[(left, right)]
            medianRight = median_right[(left, right)]
            new_identity = identity + '_' + read_direction + '_' + str(left) \
                + '_' + str(right) + '_' \
                + str(round(medianLeft, 2)) \
                + '_' + str(round(medianRight, 2))
            if not isoform_dict.get(new_identity):
                isoform_counter += 1
                isoform_dict[new_identity] = isoform_counter

            subfolder = str(int(isoform_dict[new_identity] / 4000))
            if subfolder not in os.listdir(individual_path + '/parsed_reads/'):
                os.makedirs(individual_path + '/parsed_reads/' + subfolder)
            filename = subfolder + '/Isoform' + str(isoform_dict[new_identity])
            out_reads_fasta = open(individual_path + '/parsed_reads/' + filename + '.fasta', 'a')
            out_reads_subreads = open(individual_path + '/parsed_reads/' + filename + '_subreads.fastq', 'w')
            out_reads_fasta.write(read)
            out_reads_fasta.close()
            out_reads_subreads.close()

            file_entry = individual_path + '/parsed_reads/' + filename\
                         + '.fasta' + '\t' + individual_path\
                         + '/parsed_reads/' + filename + '_subreads.fastq'\
                         + '\t' + new_identity + '\n'

            if file_entry not in file_set:
                file_set[file_entry]=0

            file_set[file_entry]+=1
            read = read.split()[0].split('_')[0][1:]

            subread_pointer[read] = individual_path + '/parsed_reads/' + filename + '_subreads.fastq'
            # for subread, sequence, qual in subread_list:
            #     out_reads_subreads.write(subread + '\n' + sequence
            #                              + '\n+\n' + qual + '\n')
    print('\n\tfinished generating isoforms models')
    out = open(individual_path + 'isoform_list', 'a')

#    print('start_left',len(count_dict['start_left']))
#    print('start_right',len(count_dict['start_right']))
#    print('end_left',len(count_dict['end_left']))
#    print('end_right',len(count_dict['end_right']))

    oute = open(individual_path + 'extended_isoforms', 'w')
    for spliceform in count_dict['start_right']:
        oute.write(spliceform+'\n')
    oute.close()

    for item,repeats in file_set.items():
        if repeats >= minimum_read_count:
            out.write(item)
    out.close()
    return subread_pointer

def read_subreads(seq_file, subread_pointer):
    for name, seq, qual in mp.fastx_read(seq_file):
        root_name = ('-').join(name.split('_')[0].split('-')[:5])
        filepath = subread_pointer.get(root_name)
        if filepath:
            outfile = open(filepath, 'a')
            outfile.write('@%s\n%s\n+\n%s\n' % (name, seq, qual))
            outfile.close()

def main():
    individual_path = path
    print('\terasing files from previous run')
    os.system('rm -r ' + individual_path + '/parsed_reads/')
    os.system('rm -r ' + individual_path + '/mp/')
    os.system('mkdir ' + individual_path + '/parsed_reads')
    listDir = os.listdir(individual_path + '/parsed_reads')
    for file in listDir:
        print(file)
        os.system('rm -r ' + individual_path + '/parsed_reads/' + file)
    out = open(individual_path + 'isoform_list', 'w')
    out.close()

    print('\treading splice sites')
    splice_dict = splice_dict = collect_splice_events(path)
    print('\tdefining spliceforms')
    start_end_dict, start_end_dict_mono = sort_reads_into_splice_junctions(splice_dict, fastaList, infile)
    print('\tdefining isoforms')
    subread_pointer = define_start_end_sites(start_end_dict, start_end_dict_mono, individual_path)

    if 'C' in consensusMode:
        print('\treading subreads from files:')
        for subread_file in subreadList:
            print('\t' + subread_file)
            read_subreads(subread_file, subread_pointer)

if __name__ == '__main__':
    main()


import torch
import torch.nn as nn
from torch.nn.modules.utils import _single, _pair, _triple
from torch.autograd import Variable
import math
import torch.nn.functional as F
import quantize
import numpy as np
import actquant
from uniq import UNIQNet

__all__ = ['mobilenet']

def nearby_int(n):
    return int(round(n))

class Conv2d(nn.Conv2d):
    """ Create registered buffer for layer base for quantization"""

    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, dilation=1, groups=1, bias=True):

        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size,
                                      stride, padding, dilation, groups, bias)
        self.register_buffer('layer_b', torch.ones(1))  # Attempt to enable multi-GPU
        self.register_buffer('layer_basis', torch.ones(1))  # Attempt to enable multi-GPU
        self.register_buffer('initial_clamp_value', torch.ones(1))  # Attempt to enable multi-GPU

    def forward(self, input):

        output = F.conv2d(input, self.weight, self.bias, self.stride,
                          self.padding, self.dilation, self.groups)

        return output

class Linear(nn.Linear):
    """ Create registered buffer for layer base for quantization"""

    def __init__(self, in_features, out_features, bias=True):

        super(Linear, self).__init__(in_features, out_features, bias=True)
        self.register_buffer('layer_b', torch.ones(1))  # Attempt to enable multi-GPU
        self.register_buffer('layer_basis', torch.ones(1))  # Attempt to enable multi-GPU
        self.register_buffer('initial_clamp_value', torch.ones(1))  # Attempt to enable multi-GPU

    def forward(self, input):

        output = F.linear(input, self.weight, self.bias)

        return output

class DepthwiseSeparableFusedConv2d(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, act_bitwidth=32, act_quant=False,
                 act_noise=False, uniq=True):
        super(DepthwiseSeparableFusedConv2d, self).__init__()
        if act_bitwidth == 32:
            self.components = nn.Sequential(
                Conv2d(in_channels, in_channels, kernel_size,
                          stride=stride, padding=padding, groups=in_channels),
                nn.BatchNorm2d(in_channels),
                # nn.ReLU(inplace=True),
                actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=uniq),

                Conv2d(in_channels, out_channels, 1, bias=False),
                nn.BatchNorm2d(out_channels),
                # nn.ReLU(inplace=True)
                actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=uniq)
            )
        else:
            self.components = nn.Sequential(
                Conv2d(in_channels, in_channels, kernel_size,
                          stride=stride, padding=padding, groups=in_channels),
                nn.BatchNorm2d(in_channels),
                actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=uniq),

                Conv2d(in_channels, out_channels, 1, bias=False),
                nn.BatchNorm2d(out_channels),
                actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=uniq)
            )

    def forward(self, x):
        return self.components(x)

class MobileNet(UNIQNet):

    def __init__(self, quant_epoch_step,quant_start_stage, alpha=1.0, shallow=False, num_classes=1000, quant=False, noise=False, bitwidth=32, step=2,
                 quant_edges=True, act_bitwidth=32, act_noise=True, step_setup=[15, 9], act_quant=False, uniq=True,
                 normalize=False):
        super(MobileNet, self).__init__(quant_epoch_step=quant_epoch_step,quant_start_stage=quant_start_stage, quant=quant, noise=noise, bitwidth=bitwidth, step=step, quant_edges=quant_edges,
                                     act_noise=act_noise, step_setup=step_setup, act_bitwidth=act_bitwidth,
                                     act_quant=act_quant, uniq=uniq)

        # if act_bitwidth == 32:
        #     act = nn.ReLU(inplace=True)
        # else:
        self.preprocess =  lambda x: x
        act = actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=self.uniq)

        layers = [
            Conv2d(3, nearby_int(alpha * 32),
                      kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(nearby_int(alpha * 32)),
            act,

            DepthwiseSeparableFusedConv2d(
                nearby_int(alpha * 32), nearby_int(alpha * 64), kernel_size=3, padding=1, act_bitwidth=act_bitwidth,
                act_quant=act_quant, act_noise=act_noise, uniq=uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 64), nearby_int(alpha * 128), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 128), nearby_int(alpha * 128), kernel_size=3, padding=1,
                                          act_bitwidth=act_bitwidth, act_quant=act_quant, act_noise=act_noise, uniq=uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 128), nearby_int(alpha * 256), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise,uniq=uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 256), nearby_int(alpha * 256), kernel_size=3, padding=1,
                                          act_bitwidth=act_bitwidth, act_quant=act_quant, act_noise=act_noise, uniq=uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 256), nearby_int(alpha * 512), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise,uniq=uniq)
        ]
        if not shallow:
            # 5x 512->512 DW-separable convolutions
            layers += [
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=uniq),
            ]
        layers += [
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 1024), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=uniq),
            # Paper specifies stride-2, but unchanged size.
            # Assume its a typo and use stride-1 convolution
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 1024), nearby_int(alpha * 1024), kernel_size=3, stride=1,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=uniq)
        ]
        self.features = nn.Sequential(*layers)
        self.avg_pool = nn.AvgPool2d(7)
        self.classifier = Linear(nearby_int(alpha * 1024), num_classes)

        self.regime = [
            {'epoch': 0, 'optimizer': 'SGD', 'lr': 1e-3,
             # 'weight_decay': 5e-4, 'momentum': 0.9},
             'weight_decay': 1e-4, 'momentum': 0.9},
            {'epoch': 30, 'lr': 5e-3},
            {'epoch': 60, 'lr': 5e-4, 'weight_decay': 0},
            {'epoch': 90, 'lr': 5e-5}
        ]

    def forward(self, x):
        x = self.preprocess(x)
        x = self.features(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

"""
class MobileNet(nn.Module):

    def __init__(self, alpha=1.0, shallow=False, num_classes=1000, quant=False, noise=False, bitwidth=32, step=2,
                 quant_edges=True, act_bitwidth=32, act_noise=True, step_setup=[15, 9], act_quant=False, uniq=True):
        super(MobileNet, self).__init__()

        self.quant = quant
        self.noise = noise
        self.bitwidth = bitwidth
        self.training_stage = 0
        self.step = step
        self.act_bitwidth = act_bitwidth
        self.act_noise = act_noise
        self.quant_edges = quant_edges
        self.uniq = uniq

        if act_bitwidth == 32:
            act = nn.ReLU(inplace=True)
        else:
            act = actquant.ActQuantBuffers(quant=act_quant, noise=act_noise, bitwidth=act_bitwidth, uniq=self.uniq)

        layers = [
            nn.Conv2d(3, nearby_int(alpha * 32),
                      kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(nearby_int(alpha * 32)),
            act,

            DepthwiseSeparableFusedConv2d(
                nearby_int(alpha * 32), nearby_int(alpha * 64), kernel_size=3, padding=1, act_bitwidth=act_bitwidth,
                act_quant=act_quant, act_noise=act_noise, uniq=self.uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 64), nearby_int(alpha * 128), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=self.uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 128), nearby_int(alpha * 128), kernel_size=3, padding=1,
                                          act_bitwidth=act_bitwidth, act_quant=act_quant, act_noise=act_noise, uniq=self.uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 128), nearby_int(alpha * 256), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise,uniq=self.uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 256), nearby_int(alpha * 256), kernel_size=3, padding=1,
                                          act_bitwidth=act_bitwidth, act_quant=act_quant, act_noise=act_noise, uniq=self.uniq),
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 256), nearby_int(alpha * 512), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise,uniq=self.uniq)
        ]
        if not shallow:
            # 5x 512->512 DW-separable convolutions
            layers += [
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=self.uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=self.uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=self.uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise, uniq=self.uniq),
                DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 512), kernel_size=3,
                                              padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                              act_noise=act_noise),
            ]
        layers += [
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 512), nearby_int(alpha * 1024), kernel_size=3, stride=2,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=self.uniq),
            # Paper specifies stride-2, but unchanged size.
            # Assume its a typo and use stride-1 convolution
            DepthwiseSeparableFusedConv2d(nearby_int(alpha * 1024), nearby_int(alpha * 1024), kernel_size=3, stride=1,
                                          padding=1, act_bitwidth=act_bitwidth, act_quant=act_quant,
                                          act_noise=act_noise, uniq=self.uniq)
        ]
        self.features = nn.Sequential(*layers)
        self.avg_pool = nn.AvgPool2d(7)
        self.classifier = nn.Linear(nearby_int(alpha * 1024), num_classes)

        print(step_setup)
        self.stages = list(range(step_setup[0], 1000, step_setup[1]))
        self.regime = [
            {'epoch': 0, 'optimizer': 'SGD', 'lr': 1e-3,
             # 'weight_decay': 5e-4, 'momentum': 0.9},
             'weight_decay': 1e-4, 'momentum': 0.9},
            {'epoch': 30, 'lr': 5e-3},
            {'epoch': 60, 'lr': 5e-4, 'weight_decay': 0},
            {'epoch': 90, 'lr': 5e-5}
        ]
        # self.regime = {
        #    0: {'optimizer': 'SGD', 'lr': 1e-1,
        #        'weight_decay': 1e-4, 'momentum': 0.9},
        #    30: {'lr': 1e-2},
        #    60: {'lr': 1e-3},
        #    90: {'lr': 1e-4}
        # }

        modules_list = list(self.modules())
        self.layers_list = [x for x in modules_list if
                            isinstance(x, nn.Conv2d) or isinstance(x, nn.Linear) or isinstance(x, actquant.ActQuantBuffers)]
        if not self.quant_edges:
            self.layers_list = self.layers_list[1:-1]

        self.layers_steps = np.array_split(self.layers_list, self.step)

        if self.act_noise:
            for layer in self.layers_steps[0]:  # Turn on noise for first stage
                if isinstance(layer, actquant.ActQuantBuffers):
                    layer.noise_during_training = True

    def switch_stage(self):
        if self.training_stage + 1 >= self.step:
            return
        print("Switching stage")
        self.training_stage += 1
        for step in self.layers_steps[:self.training_stage]:
            for layer in step:
                if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                    for param in layer.parameters():
                        param.requires_grad = False
                elif isinstance(layer, actquant.ActQuantBuffers):
                    layer.quatize_during_training = True
                    layer.noise_during_training = False

        if self.act_noise:
            for layer in self.layers_steps[self.training_stage]:  # Turn on noise only for current stage
                if isinstance(layer, actquant.ActQuantBuffers):
                    layer.noise_during_training = True

    def forward(self, x):

        temp_saved = {}

        if self.quant and not self.training:
            temp_saved = quantize.backup_weights(self.layers_list, {})
            quantize.quantize(self.layers_list, bitwidth=self.bitwidth)

        elif self.noise and self.training:

            temp_saved = quantize.backup_weights(self.layers_steps[self.training_stage], {})
            quantize.add_noise(self.layers_steps[self.training_stage], bitwidth=self.bitwidth, training=self.training)

            for i in range(self.training_stage):
                temp_saved = quantize.backup_weights(self.layers_steps[i], temp_saved)
                quantize.quantize(self.layers_steps[i], bitwidth=self.bitwidth)

        x = self.features(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)

        if self.quant and not self.training:
            quantize.restore_weights(self.layers_list, temp_saved)

        elif self.noise and self.training:
            quantize.restore_weights(self.layers_steps[self.training_stage], temp_saved)  # Restore the noised layers
            for i in range(self.training_stage):
                quantize.restore_weights(self.layers_steps[i], temp_saved)  # Restore the quantized layers
        return x
"""
def mobilenet(**kwargs):
    num_classes, depth, dataset, batch_norm, shortcut, quantize, noise, bitwidth, step, \
    act_bitwidth, act_noise, act_quant, quant_edges, step_setup, uniq , quant_epoch_step,quant_start_stage, \
    normalize= map(kwargs.get,
                                                                      ['num_classes', 'depth', 'dataset', 'batch_norm',
                                                                       'shortcut', 'quantize', 'noise', 'bitwidth',
                                                                       'step', 'act_bitwidth', 'act_noise', 'act_quant',
                                                                       'quant_edges', 'step_setup', 'uniq',
                                                                       'quant_epoch_step','quant_start_stage',
                                                                       'normalize'])

    return MobileNet(quant_epoch_step=quant_epoch_step, quant_start_stage=quant_start_stage, quant=quantize, noise=noise, bitwidth=bitwidth, step=step, act_bitwidth=act_bitwidth,
                     act_quant=act_quant, quant_edges=quant_edges, act_noise=act_noise, step_setup=step_setup,
                     uniq=uniq, normalize=normalize)

# if __name__ == '__main__':
#    from time import time
#    model = model().cuda()
#    x = Variable(torch.rand(16, 3, 224, 224).cuda())
#    t = time()
#   y = model(x)
#   print(time() - t)


# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Compute minibatch blobs for training a Fast R-CNN network."""

import numpy as np
import numpy.random as npr
import cv2
from fast_rcnn.config import cfg
from utils.blob import prep_im_for_blob, im_list_to_blob
from utils.cython_bbox import bbox_overlaps
import math

def get_minibatch(roidb, num_classes, mask_h_w):
    """Given a roidb, construct a minibatch sampled from it."""
    num_images = len(roidb)
    # Sample random scales to use for each image in this batch
    random_scale_inds = npr.randint(0, high=len(cfg.TRAIN.SCALES),
                                    size=num_images)
    assert(cfg.TRAIN.BATCH_SIZE % num_images == 0), \
        'num_images ({}) must divide BATCH_SIZE ({})'. \
        format(num_images, cfg.TRAIN.BATCH_SIZE)
    rois_per_image = cfg.TRAIN.BATCH_SIZE / num_images
    fg_rois_per_image = np.round(cfg.TRAIN.FG_FRACTION * rois_per_image)

    # Get the input image blob, formatted for caffe
    im_blob, im_scales = _get_image_blob(roidb, random_scale_inds)

    # Get the input seg image blob, formatted for caffe
    # seg_blob, im_scales = _get_seg_blob(roidb, random_scale_inds)

    # Get the input ins image blob, formatted for caffe

    # add seg_blob
    # blobs = {'data': im_blob, 'seg': seg_blob, 'ins': ins_blob}
    blobs = {'data': im_blob}

    if cfg.TRAIN.HAS_RPN:
        assert len(im_scales) == 1, "Single batch only"
        assert len(roidb) == 1, "Single batch only"
        # gt boxes: (x1, y1, x2, y2, cls)
        gt_inds = np.where(roidb[0]['gt_classes'] != 0)[0]
        gt_boxes = np.empty((len(gt_inds), 5), dtype=np.float32)
        gt_boxes[:, 0:4] = roidb[0]['boxes'][gt_inds, :] * im_scales[0]
        gt_boxes[:, 4] = roidb[0]['gt_classes'][gt_inds]
        blobs['gt_boxes'] = gt_boxes
        blobs['im_info'] = np.array(
            [[im_blob.shape[2], im_blob.shape[3], im_scales[0]]],
            dtype=np.float32)
        mask_rois_blob = np.zeros((0, 5), dtype=np.float32)
        masks_blob = np.zeros((0, 14, 14), dtype=np.float32)
        for im_i in xrange(num_images):
            labels, overlaps, im_rois, bbox_targets, bbox_inside_weights, mask_rois, masks \
                = _sample_rois(roidb[im_i], fg_rois_per_image, rois_per_image,
                               num_classes, mask_h_w)
            batch_ind_mask =  im_i * np.ones((mask_rois.shape[0], 1))
            mask_rois_blob_this_image = np.hstack((batch_ind_mask, mask_rois))
            mask_rois_blob = np.vstack((mask_rois_blob, mask_rois_blob_this_image))
            masks_blob = np.vstack((masks_blob, masks))

        blobs['mask_rois'] = mask_rois_blob
        blobs['masks'] = masks_blob
    else: # not using RPN
        # Now, build the region of interest and label blobs
        rois_blob = np.zeros((0, 5), dtype=np.float32)
        mask_rois_blob = np.zeros((0, 5), dtype=np.float32)
        labels_blob = np.zeros((0), dtype=np.float32)
        bbox_targets_blob = np.zeros((0, 4 * num_classes), dtype=np.float32)
        bbox_inside_blob = np.zeros(bbox_targets_blob.shape, dtype=np.float32)
        masks_blob = np.zeros((0, 14, 14), dtype=np.float32)
        # all_overlaps = []
        for im_i in xrange(num_images):
            labels, overlaps, im_rois, bbox_targets, bbox_inside_weights, mask_rois, masks \
                = _sample_rois(roidb[im_i], fg_rois_per_image, rois_per_image,
                               num_classes, mask_h_w)

            # Add to RoIs blob
            rois = im_rois
            # _project_im_rois(im_rois, im_scales[im_i])
            batch_ind = im_i * np.ones((rois.shape[0], 1))
            batch_ind_mask =  im_i * np.ones((mask_rois.shape[0], 1))
            rois_blob_this_image = np.hstack((batch_ind, rois))
            mask_rois_blob_this_image = np.hstack((batch_ind_mask, mask_rois))

            rois_blob = np.vstack((rois_blob, rois_blob_this_image))
            mask_rois_blob = np.vstack((mask_rois_blob, mask_rois_blob_this_image))

            # Add to labels, bbox targets, and bbox loss blobs
            labels_blob = np.hstack((labels_blob, labels))
            bbox_targets_blob = np.vstack((bbox_targets_blob, bbox_targets))
            bbox_inside_blob = np.vstack((bbox_inside_blob, bbox_inside_weights))
            masks_blob = np.vstack((masks_blob, masks))
            # all_overlaps = np.hstack((all_overlaps, overlaps))

        # For debug visualizations
        # _vis_minibatch(im_blob, rois_blob, labels_blob, all_overlaps)

        blobs['rois'] = rois_blob
        blobs['labels'] = labels_blob

        if cfg.TRAIN.BBOX_REG:
            blobs['bbox_targets'] = bbox_targets_blob
            blobs['bbox_inside_weights'] = bbox_inside_blob
            blobs['bbox_outside_weights'] = \
                np.array(bbox_inside_blob > 0).astype(np.float32)

        blobs['mask_rois'] = mask_rois_blob
        blobs['masks'] = masks_blob

    return blobs

def _sample_rois(roidb, fg_rois_per_image, rois_per_image, num_classes, mask_h_w):
    """Generate a random sample of RoIs comprising foreground and background
    examples.
    """
    # label = class RoI has max overlap with
    labels = roidb['max_classes']
    overlaps = roidb['max_overlaps']
    rois = roidb['boxes']

    # Select foreground RoIs as those with >= FG_THRESH overlap
    fg_inds = np.where(overlaps >= cfg.TRAIN.FG_THRESH)[0]
    # Guard against the case when an image has fewer than fg_rois_per_image
    # foreground RoIs
    fg_rois_per_this_image = int(np.minimum(fg_rois_per_image, fg_inds.size))
    # Sample foreground regions without replacement
    if fg_inds.size > 0:
        fg_inds = npr.choice(
                fg_inds, size=fg_rois_per_this_image, replace=False)

    # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI)
    bg_inds = np.where((overlaps < cfg.TRAIN.BG_THRESH_HI) &
                       (overlaps >= cfg.TRAIN.BG_THRESH_LO))[0]
    # Compute number of background RoIs to take from this image (guarding
    # against there being fewer than desired)
    bg_rois_per_this_image = rois_per_image - fg_rois_per_this_image
    bg_rois_per_this_image = np.minimum(bg_rois_per_this_image,
                                        bg_inds.size)
    # Sample foreground regions without replacement
    if bg_inds.size > 0:
        bg_inds = npr.choice(
                bg_inds, size=int(bg_rois_per_this_image), replace=False)

    # The indices that we're selecting (both fg and bg)
    keep_inds = np.append(fg_inds, bg_inds)
    # Select sampled values from various arrays:
    labels = labels[keep_inds]
    # Clamp labels for the background RoIs to 0
    labels[fg_rois_per_this_image:] = 0
    overlaps = overlaps[keep_inds]
    sampled_boxes = rois[keep_inds]
    rois = sampled_boxes

    bbox_targets, bbox_inside_weights = _get_bbox_regression_labels(
            roidb['bbox_targets'][keep_inds, :], num_classes)

    mask_rois, roi_has_mask, masks = _get_mask_rcnn_blobs(sampled_boxes, roidb, labels, mask_h_w)
    #mask_rois = mask_rois[np.newaxis, :]
    return labels, overlaps, rois, bbox_targets, bbox_inside_weights, mask_rois, masks

def _get_mask_rcnn_blobs(sampled_boxes, roidb, labels, mask_h_w):
    M = mask_h_w

    mask_file = cv2.imread(roidb["ins"], cv2.IMREAD_GRAYSCALE)
    if roidb['flipped']:
        mask_file = mask_file[:, ::-1]

    polys_gt_inds = np.where(
        (roidb['gt_classes'] > 0)
    )[0]
    # print(roidb.keys())

    boxes= [roidb['boxes'][i] for i in polys_gt_inds]
    boxes_from_masks = np.zeros((len(boxes), 4), dtype=np.float32)
    for i in range(len(boxes)):
        boxes_from_masks[i, :] = boxes[i]

    fg_inds = np.where(labels> 0)[0]
    roi_has_mask = labels.copy()
    roi_has_mask[roi_has_mask > 0] = 1
    # print("################")
    # print(mask_file.shape)
    # mask_file = (mask_file[:,:,0] != 0 | mask_file[:,:,1] != 0 | mask_file[:,:,2] != 0)
    # mask_file = (mask_file != [0,0,0])[:,:,0]

    if fg_inds.shape[0] > 0:

        # Class labels for the foreground rois
        masks = np.zeros((fg_inds.shape[0], M, M), dtype=np.int32)

        # Find overlap between all foreground rois and the bounding boxes
        # enclosing each segmentation
        rois_fg = sampled_boxes[fg_inds]
        overlaps_bbfg_bbpolys = bbox_overlaps(
            np.ascontiguousarray(rois_fg, dtype=np.float),
            np.ascontiguousarray(boxes_from_masks, dtype=np.float)
        )
        # Map from each fg rois to the index of the mask with highest overlap
        # (measured by bbox overlap)
        fg_bbox_inds = np.argmax(overlaps_bbfg_bbpolys, axis=1)

        # add fg targets
        for i in range(rois_fg.shape[0]):
            fg_bbox_ind = fg_bbox_inds[i]
            boxes_from_masks_now=boxes_from_masks[fg_bbox_ind]
            roi_fg = rois_fg[i]
            # im = cv2.imread(roidb['image'])
            # cv2.rectangle(im, (boxes_from_masks_now[0], boxes_from_masks_now[1]), (boxes_from_masks_now[2], boxes_from_masks_now[3]), (0, 255, 0), 2)
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # cv2.putText(im, 'boxes_from_masks_now', (boxes_from_masks_now[0], boxes_from_masks_now[1]), font, 2, (255, 255, 255), 2)
            # cv2.imwrite("boxes_from_masks_now.jpg", im)
            # im = cv2.imread(roidb['image'])
            # cv2.rectangle(im, (roi_fg[0], roi_fg[1]), (roi_fg[2], roi_fg[3]), (0, 122, 122), 2)
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # cv2.putText(im, 'roi_fg_now', (roi_fg[0], roi_fg[1]), font, 2, (122, 122, 122), 2)
            # cv2.imwrite("roi_fg_now.jpg", im)
            # im = mask_file.copy()
            # cv2.rectangle(im, (boxes_from_masks_now[0], boxes_from_masks_now[1]), (boxes_from_masks_now[2], boxes_from_masks_now[3]), (0, 255, 0), 2)
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # cv2.putText(im, 'boxes_from_masks_now', (boxes_from_masks_now[0], boxes_from_masks_now[1]), font, 2, (255, 255, 255), 2)
            # cv2.imwrite("boxes_from_masks_now_mask.jpg", im*999)
            # im = mask_file.copy()
            # cv2.rectangle(im, (roi_fg[0], roi_fg[1]), (roi_fg[2], roi_fg[3]), (0, 122, 122), 2)
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # cv2.putText(im, 'roi_fg_now', (roi_fg[0], roi_fg[1]), font, 2, (122, 122, 122), 2)
            # cv2.imwrite("roi_fg_now_mask.jpg", im*999)
            # # Rasterize the portion of the polygon mask within the given fg roi
            # to an M x M binary image
            mask = get_mask(mask_file, roi_fg, boxes_from_masks_now, M)
            mask = np.array(mask > 0, dtype=np.int32)  # Ensure it's binary
            # cv2.imwrite("mask.png", mask*999)
            masks[i, :] = mask
            # input()
            # masks[i, :] = np.reshape(mask, (M, M))
    else:  # If there are no fg masks (it does happen)
        # The network cannot handle empty blobs, so we must provide a mask
        # We simply take the first bg roi, given it an all -1's mask (ignore
        # label), and label it with class zero (bg).
        bg_inds = np.where(labels == 0)[0]
        # rois_fg is actually one background roi, but that's ok because ...
        rois_fg = sampled_boxes[bg_inds[0]].reshape((1, -1))
        # We give it an -1's blob (ignore label)
        masks = -np.ones((rois_fg.shape[0], M, M), dtype=np.int32)
        # We label it with class = 0 (background)
        mask_class_labels = np.zeros((1, ))
        # Mark that the first roi has a mask
        roi_has_mask[0] = 1
    return rois_fg, roi_has_mask, masks

def get_mask(mask_in, roi, gt_rois, size):
    x_start = int(math.floor(roi[1]))
    x_end = int(math.ceil(roi[3]))
    y_start = int(math.floor(roi[0]))
    y_end = int(math.ceil(roi[2]))

    width = mask_in.shape[0]
    height = mask_in.shape[1]
    x_start = min(max(0,x_start), width)
    x_end = min(max(0,x_end), width)
    y_start = min(max(0,y_start), height)
    y_end = min(max(0,y_end), height)

    if x_start == x_end:
        x_end += 1
    if y_start == y_end:
        y_end += 1
    if x_start == mask_in.shape[0]:
        x_start -= 1
    if y_start == mask_in.shape[1]:
        y_start -= 1

    patch_cropped = mask_in[x_start:x_end, y_start:y_end].copy()

    x_start_gt = int(math.floor(gt_rois[1]))
    x_end_gt = int(math.ceil(gt_rois[3]))
    y_start_gt = int(math.floor(gt_rois[0]))
    y_end_gt = int(math.ceil(gt_rois[2]))

    gt_patch_cropped = mask_in[x_start_gt:x_end_gt, y_start_gt:y_end_gt].copy()
    # find the main obj by count the number of pixel
    ids = np.unique(gt_patch_cropped)
    size_ = -1
    id_pick = -1
    for id in ids:
        if id == 0:
            continue
        mask = (gt_patch_cropped == id)
        arr_new = gt_patch_cropped[mask]
        if arr_new.size > size_:
            size_ = arr_new.size
            id_pick = id

    patch_cropped_temp = np.array(patch_cropped == id_pick, dtype=np.int32)
    # cv2.imwrite("patch_cropped_temp.jpg", patch_cropped_temp*99)

    mask = cv2.resize(patch_cropped_temp, (size,size), interpolation=cv2.INTER_NEAREST)
    # mask = np.array(patch_resized == id_pick, dtype=np.int32)
    return mask

def _get_image_blob(roidb, scale_inds):
    """Builds an input blob from the images in the roidb at the specified
    scales.
    """
    num_images = len(roidb)
    processed_ims = []
    im_scales = []
    for i in xrange(num_images):
        im = cv2.imread(roidb[i]['image'])
        if roidb[i]['flipped']:
            im = im[:, ::-1, :]
        target_size = cfg.TRAIN.SCALES[scale_inds[i]]
        im, im_scale = prep_im_for_blob(im, cfg.PIXEL_MEANS, target_size,
                                        cfg.TRAIN.MAX_SIZE)
        im_scales.append(im_scale)
        processed_ims.append(im)

    # Create a blob to hold the input images
    blob = im_list_to_blob(processed_ims)

    return blob, im_scales

# added for reading seg data
def _get_seg_blob(roidb, scale_inds):
    """Builds an input blob from the images in the roidb at the specified
    scales.
    """
    num_images = len(roidb)
    processed_ims = []
    im_scales = []
    for i in xrange(num_images):
        seg = cv2.imread(roidb[i]['seg'])
        seg = seg[:, :, :1]
        if roidb[i]['flipped']:
            seg = seg[:, ::-1, :]
        target_size = cfg.TRAIN.SCALES[scale_inds[i]]
        seg, im_scale = prep_seg_for_blob(seg, cfg.PIXEL_MEANS, target_size,
                                        cfg.TRAIN.MAX_SIZE)
        im_scales.append(im_scale)
        processed_ims.append(seg)

    # Create a blob to hold the input images
    blob = seg_list_to_blob(processed_ims)

    return blob, im_scales

def _project_im_rois(im_rois, im_scale_factor):
    """Project image RoIs into the rescaled training image."""
    rois = im_rois * im_scale_factor
    return rois

def _get_bbox_regression_labels(bbox_target_data, num_classes):
    """Bounding-box regression targets are stored in a compact form in the
    roidb.

    This function expands those targets into the 4-of-4*K representation used
    by the network (i.e. only one class has non-zero targets). The loss weights
    are similarly expanded.

    Returns:
        bbox_target_data (ndarray): N x 4K blob of regression targets
        bbox_inside_weights (ndarray): N x 4K blob of loss weights
    """
    clss = bbox_target_data[:, 0]
    bbox_targets = np.zeros((clss.size, 4 * num_classes), dtype=np.float32)
    bbox_inside_weights = np.zeros(bbox_targets.shape, dtype=np.float32)
    inds = np.where(clss > 0)[0]
    for ind in inds:
        cls = clss[ind]
        start = int(4 * cls)
        end = start + 4
        bbox_targets[ind, start:end] = bbox_target_data[ind, 1:]
        bbox_inside_weights[ind, start:end] = cfg.TRAIN.BBOX_INSIDE_WEIGHTS
    return bbox_targets, bbox_inside_weights

def _vis_minibatch(im_blob, rois_blob, labels_blob, overlaps):
    """Visualize a mini-batch for debugging."""
    import matplotlib.pyplot as plt
    for i in xrange(rois_blob.shape[0]):
        rois = rois_blob[i, :]
        im_ind = rois[0]
        roi = rois[1:]
        im = im_blob[im_ind, :, :, :].transpose((1, 2, 0)).copy()
        im += cfg.PIXEL_MEANS
        im = im[:, :, (2, 1, 0)]
        im = im.astype(np.uint8)
        cls = labels_blob[i]
        plt.imshow(im)
        print 'class: ', cls, ' overlap: ', overlaps[i]
        plt.gca().add_patch(
            plt.Rectangle((roi[0], roi[1]), roi[2] - roi[0],
                          roi[3] - roi[1], fill=False,
                          edgecolor='r', linewidth=3)
            )
        plt.show()


import os
import numpy as np
import time
import tensorflow as tf
import sklearn.model_selection

from khan.data.dataset import RawDataset
from khan.training.trainer_multi_tower import TrainerMultiTower, flatten_results, initialize_module
from data_utils import HARTREE_TO_KCAL_PER_MOL
from data_loaders import DataLoader
from concurrent.futures import ThreadPoolExecutor

import multiprocessing
import argparse

def main():

    parser = argparse.ArgumentParser(description="Run ANI1 neural net training.", formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--ani-lib', required=True, help="Location of the shared object for GPU featurization")
    parser.add_argument('--fitted', default=False, action='store_true', help="Whether or use fitted or self-ixn")
    parser.add_argument('--add_ffdata', default=False, action='store_true', help="Whether or not to add the forcefield data")
    parser.add_argument('--gpus', default=1, help="Number of gpus we use")
    parser.add_argument('--train_forces', default=True, help="If we train to the forces")

    parser.add_argument('--save-dir', default='~/work', help="location where save data is dumped")
    parser.add_argument('--train-dir', default='~/ANI-1_release', help="location where work data is dumped")

    args = parser.parse_args()

    print("Arguments", args)

    lib_path = os.path.abspath(args.ani_lib)
    print("Loading custom kernel from", lib_path)
    initialize_module(lib_path)

    ANI_TRAIN_DIR = args.train_dir
    ANI_SAVE_DIR = args.save_dir

    save_dir = os.path.join(ANI_SAVE_DIR, "save")

    use_fitted = args.fitted
    add_ffdata = args.add_ffdata

    data_loader = DataLoader(False)

    all_Xs, all_Ys = data_loader.load_gdb8(ANI_TRAIN_DIR)

    # todo: ensure disjunction in train_test_valid
    X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(all_Xs, all_Ys, test_size=0.25) # stratify by UTT would be good to try here
    rd_train, rd_test = RawDataset(X_train, y_train), RawDataset(X_test,  y_test)

    X_gdb11, y_gdb11 = data_loader.load_gdb11(ANI_TRAIN_DIR)
    rd_gdb11 = RawDataset(X_gdb11, y_gdb11)

    batch_size = 1024

    config = tf.ConfigProto(allow_soft_placement=True)

    all_Xs_f, all_Ys_f, all_Fs_f = data_loader.load_gdb8_forces(ANI_TRAIN_DIR) # todo: figure out how to split this consistently later

    rd_train_forces = RawDataset(all_Xs_f, all_Ys_f, all_Fs_f)

    with tf.Session(config=config) as sess:

        # This training code implements cross-validation based training, whereby we determine convergence on a given
        # epoch depending on the cross-validation error for a given validation set. When a better cross-validation
        # score is detected, we save the model's parameters as the putative best found parameters. If after more than
        # max_local_epoch_count number of epochs have been run and no progress has been made, we decrease the learning
        # rate and restore the best found parameters.

        n_gpus = int(args.gpus)
        if n_gpus > 0:
            towers = ["/gpu:"+str(i) for i in range(n_gpus)]
        else:
            towers = ["/cpu:"+str(i) for i in range(multiprocessing.cpu_count())]

        print("towers:", towers)

        trainer = TrainerMultiTower(
            sess,
            towers=towers,
            precision=tf.float32,
            layer_sizes=(128, 128, 64, 1),
            # fit_charges=True,
        )

        # if os.path.exists(save_dir):
            # print("Restoring existing model from", save_dir)
            # trainer.load(save_dir)
        # else:
        trainer.initialize() # initialize to random variables

        max_local_epoch_count = 10

        train_ops = [
            trainer.global_epoch_count,
            trainer.learning_rate,
            trainer.local_epoch_count,
            trainer.unordered_l2s,
            trainer.train_op,
        ]

        print("------------Starting Training--------------")

        start_time = time.time()

        train_forces = bool(int(args.train_forces)) # python is retarded

        # training with forces
        while sess.run(trainer.learning_rate) > 5e-10: # this is to deal with a numerical error, we technically train to 1e-9

            while sess.run(trainer.local_epoch_count) < max_local_epoch_count:

                start_time = time.time()
                # train to forces
                if train_forces:
                    train_results_forces = list(trainer.feed_dataset(
                        rd_train_forces,
                        shuffle=True,
                        target_ops=[trainer.train_op_forces, trainer.tower_force_rmses],
                        batch_size=batch_size,
                        before_hooks=trainer.max_norm_ops))
                    print(train_results_forces, end=" | ")

                #train to energies
                train_results_energies = list(trainer.feed_dataset(
                    rd_train,
                    shuffle=True,
                    target_ops=train_ops,
                    batch_size=batch_size,
                    before_hooks=trainer.max_norm_ops))

                train_abs_rmse = np.sqrt(np.mean(flatten_results(train_results_energies, pos=3))) * HARTREE_TO_KCAL_PER_MOL
                test_abs_rmse = trainer.eval_abs_rmse(rd_test)
                gdb11_abs_rmse = trainer.eval_abs_rmse(rd_gdb11)

                print(time.time()-start_time, train_abs_rmse, test_abs_rmse, gdb11_abs_rmse)

            print("==========Decreasing learning rate==========")
            sess.run(trainer.decr_learning_rate)
            sess.run(trainer.reset_local_epoch_count)
            trainer.load_best_params()

    return

if __name__ == "__main__":
    main()


import numpy as np
import os, pickle
from tqdm import tqdm
def get_word_emb(word2coef_dict, word, default_value):
    return word2coef_dict.get(word, default_value)

def get_phrase_emb(word2coef_dict, phrase, default_value):
    words = phrase.split(' ')
    embs = [ get_word_emb(word2coef_dict, word, default_value) for word in words ]
    return np.mean(embs, axis=0)

def init_glove_data(glove_fname, glove_outname):
    print(f'Constructing glove dictionary from {glove_fname} ...... ')
    word2coef_dict = {}
    running_sum = np.zeros((300,))
    with open(os.path.join(glove_fname), 'r') as f:
        for idx, line in enumerate(tqdm( list(f) )):
            values = line.split()
            word = ''.join(values[0:-300])
            coefs = np.asarray(values[-300:], dtype='float32')
            running_sum += coefs
            word2coef_dict[word] = coefs
        average_emb = running_sum / (idx + 1)

    with open( os.path.join(glove_outname, 'glove_dict.pkl'), 'wb') as f:
        pickle.dump( word2coef_dict, f)
    np.save( os.path.join(glove_outname, 'default_value'), average_emb)
    print(f'Glove dictionary saved at {glove_outname}')


# Copyright (c) 2020-2021, Matthew Broadway
# License: MIT License
import math
from abc import ABCMeta
from typing import (
    Iterable,
    TYPE_CHECKING,
    Optional,
    Dict,
    Any,
    Tuple,
)
from collections import defaultdict
from functools import lru_cache

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.font_manager import FontProperties
from matplotlib.lines import Line2D
from matplotlib.patches import PathPatch
from matplotlib.path import Path
from matplotlib.textpath import TextPath
import numpy as np

from ezdxf.addons.drawing.backend import Backend, prepare_string_for_rendering
from ezdxf.addons.drawing.properties import Properties, LayoutProperties
from ezdxf.addons.drawing.type_hints import FilterFunc
from ezdxf.tools.fonts import FontMeasurements
from ezdxf.addons.drawing.type_hints import Color
from ezdxf.tools import fonts
from ezdxf.math import Vec3, Matrix44
from ezdxf.path import Command
from ezdxf.render.linetypes import LineTypeRenderer as EzdxfLineTypeRenderer
from .config import Configuration, LinePolicy, HatchPolicy
from .matplotlib_hatch import HATCH_NAME_MAPPING
from .line_renderer import AbstractLineRenderer

if TYPE_CHECKING:
    from ezdxf.eztypes import Layout

# matplotlib docs: https://matplotlib.org/index.html

# line style:
# https://matplotlib.org/api/_as_gen/matplotlib.lines.Line2D.html#matplotlib.lines.Line2D.set_linestyle
# https://matplotlib.org/gallery/lines_bars_and_markers/linestyles.html

# line width:
# https://matplotlib.org/api/_as_gen/matplotlib.lines.Line2D.html#matplotlib.lines.Line2D.set_linewidth
# points unit (pt), 1pt = 1/72 inch, 1pt = 0.3527mm
POINTS = 1.0 / 0.3527  # mm -> points
CURVE4x3 = (Path.CURVE4, Path.CURVE4, Path.CURVE4)

class MatplotlibBackend(Backend):
    def __init__(
        self,
        ax: plt.Axes,
        *,
        adjust_figure: bool = True,
        font: FontProperties = FontProperties(),
        use_text_cache: bool = True,
    ):
        super().__init__()
        self.ax = ax
        self._adjust_figure = adjust_figure
        self._scale_dashes_backup = plt.rcParams["lines.scale_dashes"]
        # Disable internal line style scaling by matplotlib
        plt.rcParams["lines.scale_dashes"] = False

        # like set_axis_off, except that the face_color can still be set
        self.ax.xaxis.set_visible(False)
        self.ax.yaxis.set_visible(False)
        for s in self.ax.spines.values():
            s.set_visible(False)

        self.ax.autoscale(False)
        self.ax.set_aspect("equal", "datalim")

        self._current_z = 0
        self._text_renderer = TextRenderer(font, use_text_cache)
        self._line_renderer: MatplotlibLineRenderer

    def configure(self, config: Configuration) -> None:
        if config.min_lineweight is None:
            # If not set by user, use ~1 pixel
            config = config.with_changes(
                min_lineweight=72 / self.ax.get_figure().dpi
            )
        super().configure(config)
        if config.line_policy == LinePolicy.SOLID:
            self._line_renderer = InternalLineRenderer(
                config, self.ax, solid_only=True
            )
        elif config.line_policy == LinePolicy.APPROXIMATE:
            self._line_renderer = InternalLineRenderer(
                config, self.ax, solid_only=False
            )
        elif config.line_policy == LinePolicy.ACCURATE:
            # This linetype renderer should only be used by "hardcopy" backends!
            # It is just too slow for interactive backends, and the result of
            # the matplotlib line rendering is optimized for displays.
            self._line_renderer = EzdxfLineRenderer(config, self.ax)
        else:
            raise ValueError(config.line_policy)

    def clear_text_cache(self):
        self._text_renderer.clear_cache()

    def _get_z(self) -> int:
        z = self._current_z
        self._current_z += 1
        return z

    def set_background(self, color: Color):
        self.ax.set_facecolor(color)

    def draw_point(self, pos: Vec3, properties: Properties):
        """Draw a real dimensionless point."""
        color = properties.color
        self.ax.scatter([pos.x], [pos.y], s=0.1, c=color, zorder=self._get_z())

    def draw_line(self, start: Vec3, end: Vec3, properties: Properties):
        # matplotlib draws nothing for a zero-length line:
        if start.isclose(end):
            self.draw_point(start, properties)
        else:
            self._line_renderer.draw_line(start, end, properties, self._get_z())

    def draw_path(self, path, properties: Properties):
        self._line_renderer.draw_path(path, properties, self._get_z())

    def draw_filled_paths(
        self,
        paths: Iterable[Path],
        holes: Iterable[Path],
        properties: Properties,
    ):
        fill, hatch = self._get_filling(properties)
        if fill is False and hatch is None:
            return
        if hatch:
            linewidth = self._line_renderer.lineweight(properties)
        else:
            linewidth = 0
        vertices = []
        codes = []
        for path in paths:
            try:
                path = path.counter_clockwise()
            except ValueError:  # cannot detect path orientation
                continue
            v1, c1 = _get_path_patch_data(path)
            vertices.extend(v1)
            codes.extend(c1)

        for hole in holes:
            try:
                hole = hole.clockwise()
            except ValueError:  # cannot detect path orientation
                continue
            v1, c1 = _get_path_patch_data(hole)
            vertices.extend(v1)
            codes.extend(c1)

        patch = PathPatch(
            Path(vertices, codes),
            color=properties.color,
            linewidth=linewidth,
            fill=fill,
            hatch=hatch,
            zorder=self._get_z(),
        )
        self.ax.add_patch(patch)

    def draw_filled_polygon(
        self, points: Iterable[Vec3], properties: Properties
    ):
        self.ax.fill(
            *zip(*((p.x, p.y) for p in points)),
            color=properties.color,
            linewidth=0,
            zorder=self._get_z(),
        )

    def draw_text(
        self,
        text: str,
        transform: Matrix44,
        properties: Properties,
        cap_height: float,
    ):
        if not text.strip():
            return  # no point rendering empty strings
        font_properties = self.get_font_properties(properties.font)
        assert self.current_entity is not None
        text = prepare_string_for_rendering(text, self.current_entity.dxftype())
        transformed_path = _transform_path(
            self._text_renderer.get_text_path(text, font_properties),
            Matrix44.scale(
                self._text_renderer.get_scale(cap_height, font_properties)
            )
            @ transform,
        )
        self.ax.add_patch(
            PathPatch(
                transformed_path,
                facecolor=properties.color,
                linewidth=0,
                zorder=self._get_z(),
            )
        )

    def get_font_properties(
        self, font: Optional[fonts.FontFace]
    ) -> FontProperties:
        if font is None:
            return self._text_renderer.default_font
        font_properties = _get_font_properties(font)
        if font_properties is None:
            return self._text_renderer.default_font
        return font_properties

    def get_font_measurements(
        self, cap_height: float, font: fonts.FontFace = None
    ) -> FontMeasurements:
        return self._text_renderer.get_font_measurements(
            self.get_font_properties(font)
        ).scale_from_baseline(desired_cap_height=cap_height)

    def get_text_line_width(
        self, text: str, cap_height: float, font: fonts.FontFace = None
    ) -> float:
        if not text.strip():
            return 0
        dxftype = (
            self.current_entity.dxftype() if self.current_entity else "TEXT"
        )
        text = prepare_string_for_rendering(text, dxftype)
        font_properties = self.get_font_properties(font)
        path = self._text_renderer.get_text_path(text, font_properties)
        return max(x for x, y in path.vertices) * self._text_renderer.get_scale(
            cap_height, font_properties
        )

    def clear(self):
        self.ax.clear()

    def finalize(self):
        super().finalize()
        self.ax.autoscale(True)
        if self._adjust_figure:
            minx, maxx = self.ax.get_xlim()
            miny, maxy = self.ax.get_ylim()
            data_width, data_height = maxx - minx, maxy - miny
            if not math.isclose(data_width, 0):
                width, height = plt.figaspect(data_height / data_width)
                self.ax.get_figure().set_size_inches(
                    width, height, forward=True
                )
        plt.rcParams["lines.scale_dashes"] = self._scale_dashes_backup

    def _get_filling(self, properties: Properties):
        fill = True
        hatch: Any = None
        assert properties.filling is not None
        name = properties.filling.name.upper()
        if properties.filling.type == 1 and name != "SOLID":
            if self.config.hatch_policy == HatchPolicy.SHOW_OUTLINE:
                fill = False
                hatch = False
            elif (
                self.config.hatch_policy == HatchPolicy.SHOW_APPROXIMATE_PATTERN
            ):
                # Use predefined hatch pattern by name matching:
                fill = False
                hatch = HATCH_NAME_MAPPING.get(name, r"\\\\")
            elif self.config.hatch_policy == HatchPolicy.SHOW_SOLID:
                fill = True
                hatch = False
            else:
                raise ValueError(self.config.hatch_policy)
        return fill, hatch

def _transform_path(path: Path, transform: Matrix44) -> Path:
    vertices = transform.transform_vertices(
        [Vec3(x, y) for x, y in path.vertices]
    )
    return Path([(v.x, v.y) for v in vertices], path.codes)

@lru_cache(maxsize=256)  # fonts.Font is a named tuple
def _get_font_properties(font: fonts.FontFace) -> Optional[FontProperties]:
    # Font-definitions are created by the matplotlib FontManger(),
    # but stored as json file and could be altered by an user:
    font_properties = None
    try:
        font_properties = FontProperties(
            family=font.family,
            style=font.style,
            stretch=font.stretch,
            weight=font.weight,
        )
    except ValueError:
        pass
    return font_properties

class TextRenderer:
    def __init__(self, font: FontProperties, use_cache: bool):
        self._default_font = font
        self._use_cache = use_cache

        # Each font has its own text path cache
        # key is hash(FontProperties)
        self._text_path_cache: Dict[int, Dict[str, TextPath]] = defaultdict(
            dict
        )

        # Each font has its own font measurements cache
        # key is hash(FontProperties)
        self._font_measurement_cache: Dict[int, FontMeasurements] = {}

    @property
    def default_font(self) -> FontProperties:
        return self._default_font

    def clear_cache(self):
        self._text_path_cache.clear()

    def get_scale(
        self, desired_cap_height: float, font: FontProperties
    ) -> float:
        return desired_cap_height / self.get_font_measurements(font).cap_height

    def get_font_measurements(self, font: FontProperties) -> FontMeasurements:
        # None is the default font.
        key = hash(font)
        measurements = self._font_measurement_cache.get(key)
        if measurements is None:
            upper_x = self.get_text_path("X", font).vertices[:, 1].tolist()
            lower_x = self.get_text_path("x", font).vertices[:, 1].tolist()
            lower_p = self.get_text_path("p", font).vertices[:, 1].tolist()
            baseline = min(lower_x)
            measurements = FontMeasurements(
                baseline=baseline,
                cap_height=max(upper_x) - baseline,
                x_height=max(lower_x) - baseline,
                descender_height=baseline - min(lower_p),
            )
            self._font_measurement_cache[key] = measurements
        return measurements

    def get_text_path(self, text: str, font: FontProperties) -> TextPath:
        # None is the default font
        cache = self._text_path_cache[hash(font)]  # defaultdict(dict)
        path = cache.get(text, None)
        if path is None:
            if font is None:
                font = self._default_font
            # must replace $ with \$ to avoid matplotlib interpreting it as math text
            path = TextPath(
                (0, 0),
                text.replace("$", "\\$"),
                size=1,
                prop=font,
                usetex=False,
            )
            if self._use_cache:
                cache[text] = path
        return path

def _get_path_patch_data(path):
    codes = [Path.MOVETO]
    vertices = [path.start]
    for cmd in path:
        if cmd.type == Command.LINE_TO:
            codes.append(Path.LINETO)
            vertices.append(cmd.end)
        elif cmd.type == Command.CURVE4_TO:
            codes.extend(CURVE4x3)
            vertices.extend((cmd.ctrl1, cmd.ctrl2, cmd.end))
        else:
            raise ValueError(f"Invalid command: {cmd.type}")
    return [(p.x, p.y) for p in vertices], codes

def _get_aspect_ratio(ax: plt.Axes) -> float:
    minx, maxx = ax.get_xlim()
    miny, maxy = ax.get_ylim()
    data_width, data_height = maxx - minx, maxy - miny
    if abs(data_height) > 1e-9:
        return data_width / data_height
    return 1.0

def _get_width_height(
    ratio: float, width: float, height: float
) -> Tuple[float, float]:
    if width == 0.0 and height == 0.0:
        raise ValueError("invalid (width, height) values")
    if width == 0.0:
        width = height * ratio
    elif height == 0.0:
        height = width / ratio
    return width, height

def qsave(
    layout: "Layout",
    filename: str,
    *,
    bg: Optional[Color] = None,
    fg: Optional[Color] = None,
    dpi: int = 300,
    backend: str = "agg",
    config: Configuration = None,
    filter_func: FilterFunc = None,
    size_inches: Optional[Tuple[float, float]] = None,
) -> None:
    """Quick and simplified render export by matplotlib.

    Args:
        layout: modelspace or paperspace layout to export
        filename: export filename, file extension determines the format e.g.
            "image.png" to save in PNG format.
        bg: override default background color in hex format #RRGGBB or #RRGGBBAA,
            e.g. use bg="#FFFFFF00" to get a transparent background and a black
            foreground color (ACI=7), because a white background #FFFFFF gets a
            black foreground color or vice versa bg="#00000000" for a transparent
            (black) background and a white foreground color.
        fg: override default foreground color in hex format #RRGGBB or #RRGGBBAA,
            requires also `bg` argument. There is no explicit foreground color
            in DXF defined (also not a background color), but the ACI color 7
            has already a variable color value, black on a light background and
            white on a dark background, this argument overrides this (ACI=7)
            default color value.
        dpi: image resolution (dots per inches).
        size_inches: paper size in inch as `(width, height)` tuple, which also
            defines the size in pixels = (`width` * `dpi`) x (`height` * `dpi`).
            If `width` or `height` is 0.0 the value is calculated by the aspect
            ratio of the drawing.
        backend: the matplotlib rendering backend to use (agg, cairo, svg etc)
            (see documentation for `matplotlib.use() <https://matplotlib.org/3.1.1/api/matplotlib_configuration_api.html?highlight=matplotlib%20use#matplotlib.use>`_
            for a complete list of backends)
        config: drawing parameters
        filter_func: filter function which takes a DXFGraphic object as input
            and returns ``True`` if the entity should be drawn or ``False`` if
            the entity should be ignored

    .. versionadded:: 0.14

    .. versionadded:: 0.15
        added argument `params` to pass parameters to the matplotlib backend

    .. versionchanged:: 0.16
        removed arguments `ltype` and `lineweight_scaling`

    .. versionadded:: 0.17
        added argument `filter_func` to filter DXF entities

    .. versionchanged:: 0.17
        `params` argument replaced by `config` argument

    .. versionchanged:: 0.18
        added argument `size_inches`

    """
    from .properties import RenderContext
    from .frontend import Frontend
    import matplotlib

    # Set the backend to prevent warnings about GUIs being opened from a thread
    # other than the main thread.
    old_backend = matplotlib.get_backend()
    matplotlib.use(backend)
    if config is None:
        config = Configuration.defaults()

    try:
        fig: plt.Figure = plt.figure(dpi=dpi)
        ax: plt.Axes = fig.add_axes((0, 0, 1, 1))
        ctx = RenderContext(layout.doc)
        layout_properties = LayoutProperties.from_layout(layout)
        if bg is not None:
            layout_properties.set_colors(bg, fg)
        out = MatplotlibBackend(ax)
        Frontend(ctx, out, config).draw_layout(
            layout,
            finalize=True,
            filter_func=filter_func,
            layout_properties=layout_properties,
        )
        # transparent=True sets the axes color to fully transparent
        # facecolor sets the figure color
        # (semi-)transparent axes colors do not produce transparent outputs
        # but (semi-)transparent figure colors do.
        if size_inches is not None:
            ratio = _get_aspect_ratio(ax)
            w, h = _get_width_height(ratio, size_inches[0], size_inches[1])
            fig.set_size_inches(w, h, True)
        fig.savefig(
            filename, dpi=dpi, facecolor=ax.get_facecolor(), transparent=True
        )
        plt.close(fig)
    finally:
        matplotlib.use(old_backend)

class MatplotlibLineRenderer(AbstractLineRenderer, metaclass=ABCMeta):
    def __init__(self, config: Configuration, ax: plt.Axes):
        super().__init__(config)
        self.ax = ax

    @property
    def lineweight_scaling(self) -> float:
        return self._config.lineweight_scaling * POINTS

# Scaling factor for internal renderer, just guessing here:
ISO_LIN_PATTERN_FACTOR = 3.0 * POINTS
ANSI_LIN_PATTERN_FACTOR = ISO_LIN_PATTERN_FACTOR * 2.54

class InternalLineRenderer(MatplotlibLineRenderer):
    """matplotlib internal linetype rendering, which is oriented on the output
    medium and dpi: This method is simpler and faster but may not replicate the
    results of CAD applications.
    """

    def __init__(self, config: Configuration, ax: plt.Axes, solid_only: bool):
        super().__init__(config, ax)
        self._solid_only = solid_only

    def draw_line(
        self, start: Vec3, end: Vec3, properties: Properties, z: float
    ):
        self.ax.add_line(
            Line2D(
                (start.x, end.x),
                (start.y, end.y),
                linewidth=self.lineweight(properties),
                linestyle="solid"
                if self._solid_only
                else self.linetype(properties),
                color=properties.color,
                zorder=z,
            )
        )

    def draw_path(self, path, properties: Properties, z: float):
        vertices, codes = _get_path_patch_data(path)
        patch = PathPatch(
            Path(vertices, codes),
            linewidth="solid"
            if self._solid_only
            else self.lineweight(properties),
            linestyle=self.linetype(properties),
            fill=False,
            color=properties.color,
            zorder=z,
        )
        self.ax.add_patch(patch)

    @property
    def measurement_scale(self) -> float:
        return (
            ISO_LIN_PATTERN_FACTOR
            if self._config.measurement
            else ANSI_LIN_PATTERN_FACTOR
        )

    def create_pattern(self, properties: Properties, scale: float):
        """Return matplotlib line style tuple: (offset, on_off_sequence) or
        "solid".

        See examples: https://matplotlib.org/gallery/lines_bars_and_markers/linestyles.html

        """
        if len(properties.linetype_pattern) < 2:
            return "solid"
        else:
            pattern = np.round(np.array(properties.linetype_pattern) * scale)
            pattern = [max(element, 1) for element in pattern]
            if len(pattern) % 2:
                pattern.pop()
            return 0, pattern

class EzdxfLineRenderer(MatplotlibLineRenderer):
    """Replicate AutoCAD linetype rendering oriented on drawing units and
    various ltscale factors. This rendering method break lines into small
    segments which causes a longer rendering time!
    """

    def draw_line(
        self, start: Vec3, end: Vec3, properties: Properties, z: float
    ):
        pattern = self.pattern(properties)
        lineweight = self.lineweight(properties)
        color = properties.color
        if len(pattern) < 2:
            self.ax.add_line(
                Line2D(
                    (start.x, end.x),
                    (start.y, end.y),
                    linewidth=lineweight,
                    color=color,
                    zorder=z,
                )
            )
        else:
            renderer = EzdxfLineTypeRenderer(pattern)
            lines = LineCollection(
                [
                    ((s.x, s.y), (e.x, e.y))
                    for s, e in renderer.line_segment(start, end)
                ],
                linewidths=lineweight,
                color=color,
                zorder=z,
            )
            lines.set_capstyle("butt")
            self.ax.add_collection(lines)

    def draw_path(self, path, properties: Properties, z: float):
        pattern = self.pattern(properties)
        lineweight = self.lineweight(properties)
        color = properties.color
        if len(pattern) < 2:
            vertices, codes = _get_path_patch_data(path)
            patch = PathPatch(
                Path(vertices, codes),
                linewidth=lineweight,
                color=color,
                fill=False,
                zorder=z,
            )
            self.ax.add_patch(patch)
        else:
            renderer = EzdxfLineTypeRenderer(pattern)
            segments = renderer.line_segments(
                path.flattening(
                    self._config.max_flattening_distance, segments=16
                )
            )
            lines = LineCollection(
                [((s.x, s.y), (e.x, e.y)) for s, e in segments],
                linewidths=lineweight,
                color=color,
                zorder=z,
            )
            lines.set_capstyle("butt")
            self.ax.add_collection(lines)


# -*- coding: utf-8 -*-
#training the model.
#process--->1.load data(X:list of lint,y:int). 2.create session. 3.feed data. 4.training (5.validation) ,(6.prediction)
import sys
reload(sys)
sys.setdefaultencoding('utf8')
import tensorflow as tf
import numpy as np
import os
from joint_intent_slots_knowledge_domain_model import joint_knowledge_domain_model
from Data_util import generate_training_data
import jieba

#configuration
FLAGS=tf.app.flags.FLAGS
tf.app.flags.DEFINE_float("learning_rate",0.001,"learning rate")
tf.app.flags.DEFINE_integer("batch_size", 64, "Batch size for training/evaluating.") #批处理的大小 32-->128
tf.app.flags.DEFINE_integer("decay_steps", 1000, "how many steps before decay learning rate.") #6000批处理的大小 32-->128
tf.app.flags.DEFINE_float("decay_rate", 0.99, "Rate of decay for learning rate.") #0.87一次衰减多少
tf.app.flags.DEFINE_string("ckpt_dir","checkpoint_67800/","checkpoint location for the model")
tf.app.flags.DEFINE_integer("sequence_length",25,"max sentence length") #15
tf.app.flags.DEFINE_integer("embed_size",128,"embedding size")
tf.app.flags.DEFINE_boolean("is_training",True,"is traning.true:tranining,false:testing/inference")
tf.app.flags.DEFINE_integer("num_epochs",10,"number of epochs to run.")
tf.app.flags.DEFINE_integer("validate_step", 1000, "how many step to validate.") #1000做一次检验
tf.app.flags.DEFINE_integer("hidden_size",128,"hidden size")
tf.app.flags.DEFINE_float("l2_lambda", 0.0001, "l2 regularization")

tf.app.flags.DEFINE_string("data_source","knowledge_67800/training_data _10w.txt","file for data source") #knowledge/sht_20171125.txt
tf.app.flags.DEFINE_string("knowledge_path","knowledge_67800","file for data source") #skill3_train_20171114.txt
tf.app.flags.DEFINE_boolean("test_mode",False,"whether use test mode. if true, only use a small amount of data")

def main(_):
    #1. load data
    traing_data, valid_data, test_data, vocab_size, intent_num_classes, slots_num_classes,domain_num_classes = generate_training_data(FLAGS.data_source,FLAGS.knowledge_path,FLAGS.test_mode,sequence_length=FLAGS.sequence_length)
    x_train, y_intent_train, y_slots_train,y_domain_train=traing_data
    x_valid, y_intent_valid, y_slots_valid,y_domain_valid = valid_data
    x_test, y_intent_test, y_slots_test,y_domain_test = test_data

    # 2.create session.
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        # Instantiate Model
        sequence_length_batch=[FLAGS.sequence_length]*FLAGS.batch_size
        model=joint_knowledge_domain_model(intent_num_classes, FLAGS.learning_rate, FLAGS.decay_steps, FLAGS.decay_rate, FLAGS.sequence_length,
                 vocab_size, FLAGS.embed_size,FLAGS.hidden_size, sequence_length_batch,slots_num_classes,FLAGS.is_training,domain_num_classes)
        # Initialize Save
        saver = tf.train.Saver()
        if os.path.exists(FLAGS.ckpt_dir + "checkpoint"):
            print("Restoring Variables from Checkpoint")
            saver.restore(sess, tf.train.latest_checkpoint(FLAGS.ckpt_dir))
        else:
            print('Initializing Variables')
            sess.run(tf.global_variables_initializer())
        curr_epoch = sess.run(model.epoch_step)
        # 3.feed data & training
        number_of_training_data = len(x_train)
        print("number_of_training_data:", number_of_training_data)
        previous_eval_loss = 10000
        best_eval_loss = 10000
        batch_size = FLAGS.batch_size
        for epoch in range(curr_epoch, FLAGS.num_epochs):
            loss, acc_intent,acc_slot, acc_domain,counter = 0.0,0.0, 0.0, 0.0,0
            for start, end in zip(range(0, number_of_training_data, batch_size),
                                  range(batch_size, number_of_training_data, batch_size)):
                if epoch == 0 and counter == 0:#print sample to have a look
                    print("trainX[start:end]:", x_train[0:3])
                    print("y_intent_train[start:end]:", y_intent_train[0:3])
                    print("y_domain_train[start:end]:", y_domain_train[0:3])
                    print("y_slots_train[start:end]:", y_slots_train[0:3])
                feed_dict = {model.x: x_train[start:end], model.dropout_keep_prob: 0.5}
                feed_dict[model.y_intent] = y_intent_train[start:end]
                feed_dict[model.y_domain] = y_domain_train[start:end]
                feed_dict[model.y_slots] = y_slots_train[start:end]
                curr_loss, curr_acc_intent,curr_acc_slot,curr_acc_domain, _ = sess.run([model.loss_val, model.accuracy_intent,model.accuracy_slot,model.accuracy_domain, model.train_op],feed_dict)
                loss, counter, acc_intent,acc_slot,acc_domain = loss + curr_loss, counter + 1, acc_intent + curr_acc_intent,acc_slot+curr_acc_slot,acc_domain+curr_acc_domain
                if counter % 50 == 0:
                    print("joint_intent_slots_knowledge_model==>Epoch %d\tBatch %d\tTrain Loss:%.3f\tTrain Accuracy_intent:%.3f\tTrain Accuracy_slot:%.3f\tTrain Accuracy_domain:%.3f"
                          %(epoch, counter, loss / float(counter),acc_intent / float(counter),acc_slot / float(counter),acc_domain/float(counter)))

                if start==0 or start % (FLAGS.validate_step * FLAGS.batch_size) == 0:#evaluation.
                    eval_loss,eval_acc_intent,eval_acc_slot,eval_acc_domain = do_eval(sess, model, x_valid, y_intent_valid, y_slots_valid,y_domain_valid, batch_size)
                    print("joint_intent_slots_knowledge_model.validation.part. previous_eval_loss:",previous_eval_loss ,";current_eval_loss:",eval_loss,";acc_intent:",
                          eval_acc_intent,";acc_slot:",eval_acc_slot,";acc_domain:",eval_acc_domain)
                    if eval_loss > previous_eval_loss:  # if loss is not decreasing
                        # reduce the learning rate by a factor of 0.5
                        print("joint_intent_slots_knowledge_model==>validation.part.going to reduce the learning rate.")
                        learning_rate1 = sess.run(model.learning_rate)
                        lrr = sess.run([model.learning_rate_decay_half_op])
                        learning_rate2 = sess.run(model.learning_rate)
                        print("joint_intent_slots_knowledge_model==>validation.part.learning_rate1:", learning_rate1," ;learning_rate2:", learning_rate2)
                    else:  # loss is decreasing
                        if eval_loss < best_eval_loss:
                            print("joint_intent_slots_knowledge_model==>going to save the model.eval_loss:",eval_loss, ";best_eval_loss:",best_eval_loss)
                            # save model to checkpoint
                            save_path = FLAGS.ckpt_dir + "model.ckpt"
                            saver.save(sess, save_path, global_step=epoch)
                            best_eval_loss = eval_loss
                    previous_eval_loss = eval_loss

            # epoch increment
            print("going to increment epoch counter....")
            sess.run(model.epoch_increment)

        # 5.最后在测试集上做测试，并报告测试准确率 Test
        test_loss, test_acc_intent,test_acc_slot,test_acc_domain = do_eval(sess, model, x_test, y_intent_test, y_slots_test,y_domain_test, batch_size)
        print("test_loss:",test_loss,";test_acc_intent:",test_acc_intent,";test_acc_slot:",test_acc_slot,";test_acc_domain:",test_acc_domain)
    pass

# 在验证集上做验证，报告损失、精确度
def do_eval(sess, model, x_valid, y_intent_valid, y_slots_valid,y_domain_valid, batch_size):
    number_examples = len(x_valid)
    eval_loss, eval_acc_intent,eval_acc_slot, eval_acc_domain,eval_counter = 0.0,0.0,0.0, 0.0, 0
    for start, end in zip(range(0, number_examples, batch_size), range(batch_size, number_examples, batch_size)):
        feed_dict = {model.x: x_valid[start:end], model.dropout_keep_prob: 1.0}
        feed_dict[model.y_intent] = y_intent_valid[start:end]
        feed_dict[model.y_domain] = y_domain_valid[start:end]
        feed_dict[model.y_slots] = y_slots_valid[start:end]
        curr_eval_loss, curr_eval_acc_intent,curr_eval_acc_slot,curr_eval_acc_domain = sess.run([model.loss_val, model.accuracy_intent,model.accuracy_slot,model.accuracy_domain],feed_dict)
        eval_loss, eval_acc_intent,eval_acc_slot,eval_acc_domain, eval_counter = eval_loss + curr_eval_loss, eval_acc_intent +curr_eval_acc_intent , \
                                                                                eval_acc_slot+curr_eval_acc_slot,eval_acc_domain+curr_eval_acc_domain,eval_counter + 1
    return eval_loss / float(eval_counter), eval_acc_intent / float(eval_counter),eval_acc_slot / float(eval_counter),eval_acc_domain/float(eval_counter)

if __name__ == "__main__":
    tf.app.run()