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

###################################################
#
# add_ids_to_xml.py
#
# In order to tell visualization systems, "highlight this
# thing at this time", the document has to be able to identify
# particular elements.  If the original document does NOT have
# id tags on its elements, this module adds some.
#
# The auto-generated IDs have formats like "s0w2m1" meaning
# "sentence 0, word 2, morpheme 1".  But it's flexible if some elements
# already have ids, or if the markup uses different tags than a TEI document.
#
###################################################

from __future__ import print_function, unicode_literals
from __future__ import division, absolute_import

import argparse, os
import logging
from glob import glob
from collections import defaultdict
import numpy as np
from unicodedata import normalize

from readalongs.g2p.util import (load_xml, save_xml, load_json,
    get_lang_attrib, set_lang_attrib, iterate_over_text)
from readalongs.g2p.create_inv_from_map import create_inventory_from_mapping
from .. import lang

class LanguageIdentifier:

    def __init__(self, lang_dir, seen_factor=1.0, unseen_factor=0.01):
        self.seen_factor = seen_factor
        self.unseen_factor = unseen_factor
        self.langs = {}
        self.chars = {}
        self.inventories = {}  # mapping from lang to character inventories
        self.lang_dir = lang_dir
        self.load_inventories()
        self.calculate_prior_probs()

    def load_inventories(self):
        for _, subdir in lang.lang_dirs(self.lang_dir):
            for lang_filename in glob(os.path.join(subdir, "*.json")):
                try:
                    inv = load_json(lang_filename)
                except:
                    logging.error("Invalid JSON in file %s", lang_filename)
                    continue
                if not isinstance(inv, dict):
                    logging.error("File %s is not a JSON dictionary",
                                  lang_filename)
                    continue
                if ("type" not in inv):
                    logging.error("File %s is not a supported "
                                  "conversion format", lang_filename)
                    continue
                if inv["type"] == "inventory":
                    self.add_inventory(inv)
                    continue
                if inv["type"] == "mapping":
                    inv = create_inventory_from_mapping(inv, "in")
                    self.add_inventory(inv)
                    continue

    def add_inventory(self, inv):
        chars = set()
        lang = inv["metadata"]["lang"]
        if not inv["metadata"].get("display", False):
            return      # not an input orthography
        if lang not in self.langs:
            self.langs[lang] = len(self.langs) # make an index for it
        for s in inv["inventory"]:
            s = normalize("NFD", s)
            for c in s:
                if c not in self.chars:  # not yet seen in any lang
                    self.chars[c] = len(self.chars) # make an index for it
                chars.add(s)
        self.inventories[lang] = chars

    def calculate_prior_probs(self):
        probs = np.full((len(self.langs), len(self.chars)), self.unseen_factor)
        for lang, lang_index in self.langs.items():
            seen_indices = [ self.chars[c] for s in self.inventories[lang]
                                           for c in s ]
            probs[lang_index,seen_indices] = self.seen_factor
        probs /= probs.sum(axis=1, keepdims=True)
        self.logprobs = np.log(probs)

    def identify_text(self, s):
        # get the probabilities for each language
        logprobs = np.zeros((len(self.langs),))
        for c in s:
            if c not in self.chars:
                continue            # not a character we use for language id
            char_index = self.chars[c]
            logprobs += self.logprobs[:, char_index]

        logprobs -= np.max(logprobs)  # logprob scaling trick for normalization
        probs = np.exp(logprobs)      # turn back into probs
        probs /= probs.sum()          # normalize

        # return them in a more convenient format
        result = [ (lang, probs[i]) for lang, i in self.langs.items() ]
        return sorted(result, key=lambda x:x[1], reverse=True)

def add_lang_ids(xml, mapping_dir, unit="p"):
    lang_identifier = LanguageIdentifier(mapping_dir)

    elements = xml.xpath(".//" + unit)
    if not elements:
        elements = [xml]
    for element in elements:
        add_lang_ids_to_element(element, lang_identifier)
    return xml

def add_lang_ids_to_element(element, lang_identifier):
    if get_lang_attrib(element):
        return
    text = [ text for lang, text in iterate_over_text(element) ]
    text = "".join(text)
    text = normalize("NFD", text)
    lang_ids = lang_identifier.identify_text(text)
    lang_id = lang_ids[0][0]  # for now just take the most likely
    set_lang_attrib(element, lang_id)

def main(input_path, output_path, mapping_dir, unit="p"):
    xml = load_xml(input_path)
    add_lang_ids(xml, mapping_dir, unit)
    save_xml(output_path, xml)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Convert XML to another orthography while preserving tags')
    parser.add_argument('input', type=str, help='Input XML')
    parser.add_argument('output', type=str, help='Output XML')
    parser.add_argument('--mapping-dir', type=str,
            help="Alternate directory containing orthography mappings")
    parser.add_argument('--unit', type=str, default="p",
            help="Unit at which to do language identification")
    args = parser.parse_args()
    main(args.input, args.output, args.mapping_dir, args.unit)


#! /usr/bin/env python

import mdtraj as md
import argparse
import numpy as np
from LLC_Membranes.llclib import physical, topology
from LLC_Membranes.analysis import p2p
import matplotlib.pyplot as plt
import tqdm
import matplotlib as mpl
import os.path as path

def initialize():

    parser = argparse.ArgumentParser(description='Calculate the number density of selected groups along axis')

    parser.add_argument('-g', '--gro', default='wiggle.gro', help='Name of coordinate file')
    parser.add_argument('-t', '--traj', default='traj_whole.xtc', help='Trajectory file (.trr, .xtc should work)')
    parser.add_argument('-b', '--begin', default=0, type=int, help='Frame to begin calculations')
    parser.add_argument('-c', '--components', nargs='+', help="Region(s) or atoms to include in calculation. Special"
                                                              "groups include 'head groups' and 'tails'")
    parser.add_argument('-m', '--monomer', default='NaGA3C11', help='Name of monomer (no file extension)')
    parser.add_argument('-e', '--end', default=-1, type=int, help='Frame to stop doing calculations')
    parser.add_argument('-bins', default=100, type=int, help='Number of bins to use')
    parser.add_argument('-s', '--solvate', action="store_true",
                        help='If the system is solvated, plot the number density of water as well')
    parser.add_argument('-l', '--load', help='Name of compressed .npz to load')

    args = parser.parse_args()

    return args

def grps(name):
    """ Return names of atoms making up specialized groups. This is specifically for NaGA3C11.

    TODO: incorporate these groups as annotations (maybe)

    :param name: specialized group names. Valid options are 'head groups' and 'tails'

    :type name: str

    :return: atom names that constitute specialized groups
    """

    if name == 'head groups':

        return ['C', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'O3', 'O4']

    elif name == 'tails':

        return ['C7', 'C8', 'C9', 'C10', 'C11', 'C12', 'C13', 'C14', 'C15', 'C16', 'C17', 'C18', 'C19', 'C20', 'C21',
                'C22', 'C23', 'C24', 'C25', 'C26', 'C27', 'C28', 'C29', 'C30', 'C31', 'C32', 'C33', 'C34', 'C35',
                'C36', 'C37', 'C38', 'C39', 'C40', 'C41', 'C42', 'C43', 'C44', 'C45', 'C46', 'C47', 'C48']

class Region(object):

    def __init__(self, traj, gro, components, monomer, begin=0, end=-1):
        """ Determine the density of components. Similar to a radial density, but instead of center of mass, the radial
        density of each atom in the region is averaged together.

        :param traj: name of trajectory to analyze (.xtc or .trr)
        :param gro: name of GROMACS coordinate file with same topology as traj
        :param component: list of atoms or name of groups of atoms
        :param begin: first frame to analyze
        :param end: last grame to analyze

        :type traj: str
        :type gro: str
        :type component: list
        :type begin: int
        :type end: int
        """

        special_groups = ['head groups', 'tails']

        if components[0].lower() in special_groups:
            self.components = grps(components[0].lower())
        else:
            self.components = components

        print('Loading trajectory...', end="")
        self.t = md.load('%s' % traj, top='%s' % gro)[begin:end]
        print('done')

        self.monomer = topology.LC('%s.gro' % monomer)

        self.r = None
        self.density = None

    def component_density(self, bins=50, spline=True, progress=True, npts_spline=10):

        self.r = np.zeros([bins])
        self.density = np.zeros([self.t.n_frames, bins])

        pore_defining_atoms = [a.index for a in self.t.topology.atoms if a.name in self.monomer.pore_defining_atoms
                               and a.residue.name in self.monomer.residues]

        pore_centers = physical.avg_pore_loc(4, self.t.xyz[:, pore_defining_atoms, :], self.t.unitcell_vectors,
                                             spline=spline, progress=progress, npts=npts_spline)

        print('Calculating component density')
        self.r, self.density = physical.compdensity(self.com, pore_centers, self.t.unitcell_vectors,
                                                    nbins=bins, spline=spline, cut=cut)

def duplicate(pos, box):
    """
    Duplicate a set of positions periodically once in the +/- xy directions
    :param pos: xyz positions of a set of coordinates to be duplicated periodically
    :param box: box vectors in mdtraj format (t.unitcell_vectors : [nframes, 3, 3]) for every frame
    :return: Periodically duplicated system
    """

    n = pos.shape[1]  # number atoms in original unit cell

    p = np.zeros([pos.shape[0], n*9, 3])  # will hold periodically duplicated system
    p[:, :n, :] = pos

    # x-direction
    for t in range(pos.shape[0]):
        p[t, n:2*n, :] = pos[t, :, :] + box[t, 0, :]
        p[t, 2*n:3*n, :] = pos[t, :, :] - box[t, 0, :]

    # y-direction
    n *= 3
    for t in range(pos.shape[0]):
        p[t, n:2*n, :] = p[t, :n, :] + box[t, 1, :]
        p[t, 2*n:3*n, :] = p[t, :n, :] - box[t, 1, :]

    return p

def compdensity(component, pore_centers, start, box, cut=1.5, pores=4, nbins=50, rmax=3.5, buffer=0.0):
    """ Measure the density of a component as a function of the distance from the pore centers

    :param component: the coordinates of the component(s) which you want a radial distribution of at each frame
    :param pore_centers: a numpy array of the locations of each pore center at each trajectory frame
    :param start: the frame number at which to start calculations (should be after equilibration)
    :param cut: cutoff distance for distance calculations. Will not count anything further than cut from the pore center
    :param pores: number of pores (int) default=4
    :param rmax: maximum distance from pore center to calculate density for, default = 3.5 nm
    :param buffer: percentage used to define the location of z planes between which component density will be computed,
           float, default = 0 (i.e. no buffer). Should be between 0 and 1. e.g. for 1 percent, use 0.01 as the buffer

    :type component: numpy.ndarray
    :type pore_centers: numpy.ndarray
    :type start: int
    :type cut: float
    :type pores: int
    :type rmax: float
    :type buffer: float

    :return: the density of "component" as a function the distance from the pore center. Also
             returns the calculated bin width for plotting
    """

    # Extract basic system information. It's important to follow the format of the component array to get it right

    tot_atoms = component.shape[1]  # the total number of components in a single frame
    n_ppore = tot_atoms // pores  # the total number of components in each pore

    nT = component.shape[0]
    zbox = np.mean(box[:, 2, 2])
    # Find the approximate max and minimum z values of the components based on the last frame

    zmax = np.max(component[-1, :, 2])
    zmin = np.min(component[-1, :, 2])
    thickness = zmax - zmin  # approximate membrane thickness

    # now find the maximum and minimum permissible z dimensions based on the buffer
    zmax_buff = zmax - thickness * buffer  # Could use buffer/2 depending on how you interpret what buffer % means
    zmin_buff = zmin + thickness * buffer

    density = np.zeros([nbins])  # number / nm^3
    for t in tqdm.tqdm(range(start, nT)):
        for p in range(pores):
            # narrow down the positions to those that are within 'cut' of at least one pore
            distances = np.linalg.norm(component[t, :, :2] - pore_centers[t, p, :], axis=1)
            d_sorted = np.sort(distances)
            # find where the distances exceed the cutoff
            stop = 0
            while d_sorted[stop] < cut:
                stop += 1

            hist, bin_edges = np.histogram(d_sorted[:stop], bins=nbins, range=(0, cut))  # the range option is necessary
            #  to make sure we have equal sized bins on every iteration

            density += hist

    density /= zbox * (nT - start)  # take average
    bin_width = cut / nbins

    # normalize based on area of anulus where bin is located
    r = np.zeros([nbins])
    normalization = []
    for i in range(nbins):
        normalization.append(np.pi * (bin_edges[i + 1] ** 2 - bin_edges[i] ** 2))
        density[i] /= np.pi * (bin_edges[i + 1] ** 2 - bin_edges[i] ** 2)
        r[i] = (bin_edges[i + 1] + bin_edges[i]) / 2

    return density, r, bin_width

if __name__ == "__main__":

    args = initialize()  # parse the args

    # run : regional_density.py - m offset_disordered_regional_density.npz offset_regional_density.npz
    # layered_disordered_regional_density.npz layered_regional_density.npz solvated_regional_density.npz

    # mpl.style.use('seaborn')

    # regions = ['Tails', 'Head Groups', 'Sodium']
    regions = ['Tails', 'Head Groups', 'Sodium']

    # if args.multi:
    if False:

        # colors = ['blue', 'red', 'green', 'xkcd:orange']
        # colors = ['xkcd:red', 'xkcd:green', 'blue', 'xkcd:yellow']
        colors = ['xkcd:blue', 'xkcd:olive', 'xkcd:orangered', 'xkcd:magenta', 'xkcd:gold']
        names = ['Ordered Parallel Displaced', 'Ordered Sandwiched', 'Disordered Sandwiched',
                 'Disordered Parallel Displaced', 'Solvated Parallel Displaced']

        # import pylab
        #
        # fig = pylab.figure()
        # figlegend = pylab.figure(figsize=(7.75, 0.6))
        # ax = fig.add_subplot(111)
        # for i in range(5):
        #     ax.plot(range(10), pylab.randn(10), color=colors[i], label=names[i])
        #
        # #lines = ax.plot(range(10), pylab.randn(10), range(10), pylab.randn(10),range(10), pylab.randn(10), range(10), pylab.randn(10), range(10), pylab.randn(10))
        # figlegend.legend(*ax.get_legend_handles_labels(), 'best', ncol=3)
        # fig.show()
        # figlegend.show()
        # fig.tight_layout()
        # figlegend.savefig('legend.pdf')
        #
        # plt.show()
        # exit()

        # names = ['Parallel Displaced (d=3.7)', 'Sandwiched (d=3.7)', 'Sandwiched (d=5.0)', 'Parallel Displaced (d=5.0)',
        #          'Solvated Parallel Displaced']
        # names = ['Dry', 'Solvated']
        # colors = ['xkcd:orange', 'xkcd:blue', 'xkcd:orange']

        n = len(args.multi)
        system = np.load(args.multi[0])

        # It is assumed that all of the data uses the same number of bins with the same bin width
        r = system['r']
        bin_width = system['bw']

        results = np.zeros([n, len(regions), len(r)])
        # results = np.zeros([n, 4, len(r)])

        # results[0, :3, :] = system['results']

        for i in range(n - 1):
            system = np.load(args.multi[i])
            results[i, :, :] = system['results']

        system = np.load(args.multi[-1])
        results[-1, :, :] = system['results'][:3, :]

        outline = np.zeros([4, n, r.shape[0]*2 + 2, 2])
        for i in range(len(regions)):
            print(i)
            fig = plt.figure(i)
            for j in range(n):
                #plt.bar(r, results[j, i, :], bin_width, color=colors[j], alpha=1, label=names[j])

                half_width = bin_width / 2
                # outline = np.zeros([r.shape[0]*2 + 2, 2])
                outline[i, j, 0, 0] = r[0] - half_width
                outline[i, j, -1, 0] = r[-1] + half_width
                outline[i, j, 1:-1:2, 0] = r - half_width
                outline[i, j, 2:-1:2, 0] = r + half_width
                outline[i, j, 1:-1:2, 1] = results[j, i, :]
                outline[i, j, 2:-1:2, 1] = results[j, i, :]
                if i == 0:
                    plt.plot(outline[i, j, :-1, 0], outline[i, j, :-1, 1], color=colors[j], linewidth=2, label=names[j])
                else:
                    plt.plot(outline[i, j, 1:, 0], outline[i, j, 1:, 1], color=colors[j], linewidth=2,
                             label=names[j])
            # plt.title(regions[i], fontsize=14)

            #plt.legend()#fontsize=12)

            plt.ylabel('Component Number Density \n (number/nm$^3$)', fontsize=18)
            plt.xlabel('Distance from pore center, r (nm)', fontsize=18)
            plt.axes().tick_params(labelsize=14)
            # plt.ylim([0, 0.6])
            plt.tight_layout()
            plt.savefig("%s_density.pdf" % regions[i])

        plt.show()

    if args.solvate:
        colors = ['xkcd:red', 'xkcd:green', 'blue', 'xkcd:yellow']
    else:
        colors = ['red', 'green', 'blue']

    if not args.load:

        R = Region(args.traj, args.gro, args.components, begin=args.begin, end=args.end)

        print('Loading trajectory...', end="")
        t = md.load('%s' % args.traj, top='%s' % args.gro)[args.begin:args.end]
        print('done')

        box = t.unitcell_vectors
        nT = t.n_frames
        npores = 4
        r_max = 0

        if args.solvate:
            results = np.zeros([len(regions) + 1, args.bins])
        else:
            results = np.zeros([len(regions), args.bins])

        #keep = [a.index for a in t.topology.atoms if a.residue.name != 'HOH']  # everything kept if system not solvated
        components = ['C', 'C1', 'C2', 'C3', 'C4', 'C5']
        comp = [a.index for a in t.topology.atoms if a.name in components]

        p_centers = physical.avg_pore_loc(npores, t.xyz[:, comp, :], box)

        for i, reg in enumerate(regions):

            print('Calculating number density of %s region' % reg)

            pos = p2p.restrict_atoms(t, reg)  # restrict trajectory to region

            p = duplicate(pos, t.unitcell_vectors)  # duplicate things periodically

            equil = 0
            density, r, bin_width = compdensity(pos, p_centers, equil, box, pores=npores, nbins=args.bins)

            results[i, :] = density

            plt.bar(r, density, bin_width, color=colors[i], alpha=0.6, label=reg)

        np.savez_compressed("regional_density", results=results, r=r, bw=bin_width, box=t.unitcell_vectors)
        print('Arrays saved as regional_density.npz')

        if args.solvate:

            print('Calculating number density of solvent')
            keep = [a.index for a in t.topology.atoms if a.residue.name == 'HOH' and a.name == 'O']
            pos = t.xyz[:, keep, :]
            p = duplicate(pos, t.unitcell_vectors)
            equil = 0
            density, r, bin_width = Structure_char.compdensity(pos, p_centers, equil, t.unitcell_vectors, pores=npores, buffer=0, nbins=args.bins)
            results[-1, :] = density
            plt.bar(r, density, bin_width, color='xkcd:gold', alpha=0.75, label='Water')

    else:

        d = np.load(args.load)

        bw = d['bw']
        r = d['r']
        results = d['results']
        box = d['box']

        if args.solvate:
            regions.append('Water')

        for i in range(results.shape[0]):
            plt.bar(r, results[i, :], bw, color=colors[i], alpha=0.75, label=regions[i])

    un_normalize = np.zeros_like(results)

    annulus_area = []
    dr = r[-1] - r[-2]  # width of annulus
    # r from above is calculated as the average of the edges of the bins. So I need to reverse that here
    for i in r:
        inner_r = i - (dr/2)
        outer_r = i + (dr/2)
        annulus_area.append(np.pi*(outer_r**2 - inner_r**2))

    for i in range(results.shape[0]):
        un_normalize[i, :] = results[i, :] * annulus_area

    r_cut = 0.6
    cut_index = 0
    while r[cut_index] < r_cut:
        cut_index += 1

    print('Percentage of sodium within %s nm of pore center: %2.2f %%' % (r_cut, 100*(sum(un_normalize[2, :cut_index])/sum(un_normalize[2, :]))))

    r_cut = 0.6
    cut_index = len(r) - 1
    while r[cut_index] > r_cut:
        cut_index -= 1

    print('Percentage of tails within %s nm of pore center: %2.2f %%' % (r_cut, 100*(sum(un_normalize[0, :cut_index])/sum(un_normalize[0, :]))))

    r_cut_pore = r_cut
    pore_cut_index = cut_index

    r_cut = 1
    cut_index = len(r) - 1
    while r[cut_index] > r_cut:
        cut_index -= 1

    print('Percentage of tails between %s and %s nm of pore center: %2.2f %%' % (r_cut_pore, r_cut,
        100*(sum(un_normalize[0, pore_cut_index:cut_index])/sum(un_normalize[0, :]))))

    print('Maximum of Head group region: r = %s' % r[np.argmax(results[1, :])])

    plt.legend(prop={'size':14}, loc=1)
    plt.ylabel('Component number density (count/nm$^3$)', fontsize=14)
    plt.xlabel('Distance from pore center (nm)', fontsize=14)
    plt.axes().tick_params(labelsize=14)
    plt.xlim(0, 1.5)
    plt.tight_layout()
    plt.savefig("regional_density.png")
    plt.show()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
last mod 8/25/19
testing pieces of lidar processing code
this piece includes separation of measurements into full-car and part-of-car
and conversion into xyalw format with normal distribution
also includes monogrnet for comparison
"""

import numpy as np
from imageio import imread
from cv2 import imshow, waitKey, destroyWindow

from segment import lasers2use, segmentPoints, makeMeasurement, uv2xy, laser_angles
from segment import getGroundForLaser, makeOcclusionMap, boundMeasurement
from segment import getGrndHeight, msmtBound2msmtNormal

from presavedSensor import getMsmts as getCamMsmts

from plotStuff import base_image, plotRectangle, plotPoints, drawLine
from plotStuff import plotImgKitti, addRect2KittiImg, hsv2Rgba, plotRectangleEdges
from plotStuff import drawOcclusion#fillArc
from calibs import calib_extrinsics, calib_projections, view_by_day
from trackinginfo import sceneranges
from trackinginfo import calib_map_training as calib_map
from analyzeGT import readGroundTruthFileTracking

lidar_files = '/home/m2/Data/kitti/tracking_velodyne/training/{:04d}/{:06d}.bin'
img_files = '/home/m2/Data/kitti/tracking_image/training/{:04d}/{:06d}.png'
gt_files = '/home/m2/Data/kitti/tracking_gt/{:04d}.txt'
ground_files = '/home/m2/Data/kitti/tracking_ground/training/{:02d}f{:06d}.npy'
scene_idx = 6
startfileidx, endfileidx = sceneranges[scene_idx]
startfileidx = 90#135
endfileidx = 150#136

def clear(): destroyWindow('a')

calib_idx = calib_map[scene_idx]
calib_extrinsic = calib_extrinsics[calib_idx].copy()
calib_projection = calib_projections[calib_idx]
calib_intrinsic = calib_projection.dot(np.linalg.inv(calib_extrinsic))
calib_extrinsic[2,3] += 1.65
view_angle = view_by_day[calib_idx]
view = view_angle + .1
bevTM = np.array(((-320./30,0,640), (0,-320./30,320)))
bevTM = bevTM.dot(calib_extrinsic[[0,1,3],:]) # lidar xyz to pixels

with open(gt_files.format(scene_idx), 'r') as fd: gtfilestr = fd.read()
gt_all, gtdc = readGroundTruthFileTracking(gtfilestr, ('Car', 'Van'))

for fileidx in range(startfileidx, endfileidx):
    img = imread(img_files.format(scene_idx, fileidx))[:,:,::-1]
    data = np.fromfile(lidar_files.format(scene_idx, fileidx),
                   dtype=np.float32).reshape((-1,4))[:,:3]
    gt = gt_all[fileidx]
    ground = np.load(ground_files.format(scene_idx, fileidx))

    starts = np.where(np.diff(np.sign(data[:,1])) > 0)[0]
    starts = np.concatenate(([0], starts+1, [len(data)]))
    true_starts = np.append(np.diff(starts) > 2, [True])
    starts = starts[true_starts]
    assert starts.shape[0] > lasers2use[-1]+1
    lidar = []
    for laser in lasers2use:
        pts = data[starts[laser]:starts[laser+1]]
        include = pts[:,0] > 0
        include &= abs(pts[:,1]) < pts[:,0]*view + 2.
        include &= pts.dot(calib_extrinsic[2,:3])+calib_extrinsic[2,3] > -.3
        pts = pts[include]
        swap_idx = np.where(np.diff(np.arctan2(pts[:,1],pts[:,0]))<-.05)[0]
        assert len(swap_idx) <= 1
        if len(swap_idx) == 1:
            swap_idx = swap_idx[0] + 1
            pts = np.append(pts[swap_idx:], pts[:swap_idx], axis=0)
        lidar.append(pts.copy())

    plotimg1 = plotImgKitti(view_angle)
    plotimg2 = plotImgKitti(view_angle)
    # draw object
    for gtobj in gt:
        box = np.array(gtobj['box'])
        addRect2KittiImg(plotimg1, box, (200.,200.,200,.9))
        addRect2KittiImg(plotimg2, box, (200.,200.,200,.9))

    cammsmts = getCamMsmts(scene_idx, fileidx)
    for msmt in cammsmts:
        if msmt[5] > .05:
            # high scored msmts are blue, low are purple
            plotRectangleEdges(plotimg2, msmt[:5], hsv2Rgba((1-msmt[5])*.9/.95, 1,1,.8))

    # draw laser points
    for laseridx, pts in enumerate(lidar):
        lasercolor = hsv2Rgba(laseridx / (len(lidar)+2.), 1., 1., .75)
        segs = segmentPoints(pts)
        # get rectangles from point segments

        msmts = []
        segsareground = []
        segsinclude = []
        for segidx, seg in enumerate(segs):
            segmiddle = np.mean(seg,axis=0)
            segmiddle = calib_extrinsic[:3,:3].dot(segmiddle) + calib_extrinsic[:3,3]
            seggroundelev = getGrndHeight(ground, segmiddle[0], segmiddle[1])
            heights = seg.dot(calib_extrinsic[2,:3])+calib_extrinsic[2,3]
            segisground = max(heights)-seggroundelev < .3
            seginclude = (not segisground) and segmiddle[2]-seggroundelev < 2.
            segsareground.append(segisground)
            segsinclude.append(seginclude)
            msmts.append(makeMeasurement(seg))

        ground_present = getGroundForLaser(ground, laser_angles[laser])
        starting_angle = min(-view, np.arctan2(pts[0,1],pts[0,0]))
        ending_angle = max(view, np.arctan2(pts[-1,1],pts[-1,0]))
        occlusionmap = makeOcclusionMap(segs, segsareground, msmts,
                                        starting_angle, ending_angle, True,
                                        ground_present, None)

        outputmsmts = []
        for segidx in range(len(segs)):
            if not segsinclude[segidx]:
                continue
            msmt = msmts[segidx]
            msmt = boundMeasurement(segidx, msmt, occlusionmap)
            if msmt[3]-msmt[2] < 1. and msmt[5]-msmt[4] < 1.:
                msmt1,msmt2 = None,None
            else:
                msmt1, msmt2 = msmtBound2msmtNormal(msmt)

            if msmt1 is None and msmt2 is None:
                rec = uv2xy(msmt[:6])[:2]
                rec = calib_extrinsic[:2,:2].dot(rec) + calib_extrinsic[:2,2]
                nottoobig = ((msmt[3]-msmt[2] < 5.88) and
                             (msmt[5]-msmt[4] < 2.28)) or (
                             (msmt[3]-msmt[2] < 2.28) and
                             (msmt[5]-msmt[4] < 5.88))
                nottoosmall = max(msmt[5]-msmt[4],msmt[3]-msmt[2]) > .35
                if not nottoobig:
                    print("toobig dude at {:.1f}  {:.1f}".format(rec[0],rec[1]))
                    #continue #ABCDEEEE
                #if not (nottoobig and nottoosmall): continue
                if not nottoobig: continue
                if nottoosmall: lasercolor = hsv2Rgba(.333, 1., 1., 1.)
                else: lasercolor = hsv2Rgba(0., 1., 1., 1.)
                if rec[0] < .2 or rec[0] > 59 or abs(rec[1]) > 29.5: continue
                rec = bevTM.dot((rec[0],rec[1],0,1)).astype(int)
                scattershape = np.mgrid[-2:3,-2:3].transpose((1,2,0)).reshape((25,2))
                scattershape = scattershape[np.hypot(scattershape[:,0],
                                                     scattershape[:,1])<2.25]
                assignlasercolor = np.append(lasercolor[:3]/lasercolor[3], (1.,))
                plotPoints(plotimg1, rec[None,0], rec[None,1], scattershape,
                           assignlasercolor)

            if msmt1 is not None:
                rec, cov = msmt1
                assert cov[0,0] < 1000 and cov[1,1] < 1000
                rec[:2] = calib_extrinsic[:2,:2].dot(rec[:2]) +\
                                    calib_extrinsic[:2,2]
                vardiff = (cov[0,0]-cov[1,1])/2
                varmean = (cov[0,0]+cov[1,1])/2
                stdinmaindirection = (varmean + np.hypot(vardiff, cov[0,1]))**.5
                color = hsv2Rgba(min(stdinmaindirection/4.5,.9), 1., 1., 1.)
                plotRectangleEdges(plotimg1, rec, color)
                outputmsmts.append(msmt1)
            if msmt2 is not None:
                rec, cov = msmt2
                assert cov[0,0] < 1000 and cov[1,1] < 1000
                rec[:2] = calib_extrinsic[:2,:2].dot(rec[:2]) +\
                                    calib_extrinsic[:2,2]
                vardiff = (cov[0,0]-cov[1,1])/2
                varmean = (cov[0,0]+cov[1,1])/2
                stdinmaindirection = (varmean + np.hypot(vardiff, cov[0,1]))**.5
                color = hsv2Rgba(min(stdinmaindirection/4.5,.9), 1., 1., 1.)
                plotRectangleEdges(plotimg1, rec, color)
                outputmsmts.append(msmt2)

    plotimg1 = np.minimum((plotimg1[:,:,:3]/plotimg1[:,:,3:]),255.).astype(np.uint8)
    plotimg2 = np.minimum((plotimg2[:,:,:3]/plotimg2[:,:,3:]),255.).astype(np.uint8)
    img = img[:368]
    display_img = np.zeros((640+img.shape[0], 1280, 3), dtype=np.uint8)
    display_img[:640, :640] = plotimg1
    display_img[:640, 640:] = plotimg2
    display_img[640:, (1280-img.shape[1])//2:(1280+img.shape[1])//2] = img
    imshow('a', display_img);
    if waitKey(500) == ord('q'):
        break


# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from typing import List, Union

import numpy as np
import paddle
import soundfile as sf
import yaml
from yacs.config import CfgNode

from paddlehub.env import MODULE_HOME
from paddlehub.module.module import moduleinfo, serving
from paddlehub.utils.log import logger
from paddlespeech.t2s.frontend.zh_frontend import Frontend
from paddlespeech.t2s.models.fastspeech2 import FastSpeech2
from paddlespeech.t2s.models.fastspeech2 import FastSpeech2Inference
from paddlespeech.t2s.models.parallel_wavegan import PWGGenerator
from paddlespeech.t2s.models.parallel_wavegan import PWGInference
from paddlespeech.t2s.modules.normalizer import ZScore
from paddlespeech.vector.exps.ge2e.audio_processor import SpeakerVerificationPreprocessor
from paddlespeech.vector.models.lstm_speaker_encoder import LSTMSpeakerEncoder

@moduleinfo(
    name="ge2e_fastspeech2_pwgan",
    version="1.0.0",
    summary="",
    author="paddlepaddle",
    author_email="",
    type="audio/voice_cloning",
)
class VoiceCloner(paddle.nn.Layer):
    def __init__(self, speaker_audio: str = None, output_dir: str = './'):
        super(VoiceCloner, self).__init__()

        speaker_encoder_ckpt = os.path.join(MODULE_HOME, 'ge2e_fastspeech2_pwgan', 'assets',
                                            'ge2e_ckpt_0.3/step-3000000.pdparams')
        synthesizer_res_dir = os.path.join(MODULE_HOME, 'ge2e_fastspeech2_pwgan', 'assets',
                                           'fastspeech2_nosil_aishell3_vc1_ckpt_0.5')
        vocoder_res_dir = os.path.join(MODULE_HOME, 'ge2e_fastspeech2_pwgan', 'assets', 'pwg_aishell3_ckpt_0.5')

        # Speaker encoder
        self.speaker_processor = SpeakerVerificationPreprocessor(
            sampling_rate=16000,
            audio_norm_target_dBFS=-30,
            vad_window_length=30,
            vad_moving_average_width=8,
            vad_max_silence_length=6,
            mel_window_length=25,
            mel_window_step=10,
            n_mels=40,
            partial_n_frames=160,
            min_pad_coverage=0.75,
            partial_overlap_ratio=0.5)
        self.speaker_encoder = LSTMSpeakerEncoder(n_mels=40, num_layers=3, hidden_size=256, output_size=256)
        self.speaker_encoder.set_state_dict(paddle.load(speaker_encoder_ckpt))
        self.speaker_encoder.eval()

        # Voice synthesizer
        with open(os.path.join(synthesizer_res_dir, 'default.yaml'), 'r') as f:
            fastspeech2_config = CfgNode(yaml.safe_load(f))
        with open(os.path.join(synthesizer_res_dir, 'phone_id_map.txt'), 'r') as f:
            phn_id = [line.strip().split() for line in f.readlines()]

        model = FastSpeech2(idim=len(phn_id), odim=fastspeech2_config.n_mels, **fastspeech2_config["model"])
        model.set_state_dict(paddle.load(os.path.join(synthesizer_res_dir, 'snapshot_iter_96400.pdz'))["main_params"])
        model.eval()

        stat = np.load(os.path.join(synthesizer_res_dir, 'speech_stats.npy'))
        mu, std = stat
        mu = paddle.to_tensor(mu)
        std = paddle.to_tensor(std)
        fastspeech2_normalizer = ZScore(mu, std)
        self.sample_rate = fastspeech2_config.fs

        self.fastspeech2_inference = FastSpeech2Inference(fastspeech2_normalizer, model)
        self.fastspeech2_inference.eval()

        # Vocoder
        with open(os.path.join(vocoder_res_dir, 'default.yaml')) as f:
            pwg_config = CfgNode(yaml.safe_load(f))

        vocoder = PWGGenerator(**pwg_config["generator_params"])
        vocoder.set_state_dict(
            paddle.load(os.path.join(vocoder_res_dir, 'snapshot_iter_1000000.pdz'))["generator_params"])
        vocoder.remove_weight_norm()
        vocoder.eval()

        stat = np.load(os.path.join(vocoder_res_dir, 'feats_stats.npy'))
        mu, std = stat
        mu = paddle.to_tensor(mu)
        std = paddle.to_tensor(std)
        pwg_normalizer = ZScore(mu, std)

        self.pwg_inference = PWGInference(pwg_normalizer, vocoder)
        self.pwg_inference.eval()

        # Text frontend
        self.frontend = Frontend(phone_vocab_path=os.path.join(synthesizer_res_dir, 'phone_id_map.txt'))

        # Speaking embedding
        self._speaker_embedding = None
        if speaker_audio is None or not os.path.isfile(speaker_audio):
            speaker_audio = os.path.join(MODULE_HOME, 'ge2e_fastspeech2_pwgan', 'assets', 'voice_cloning.wav')
            logger.warning(f'Due to no speaker audio is specified, speaker encoder will use defult '
                           f'waveform({speaker_audio}) to extract speaker embedding. You can use '
                           '"set_speaker_embedding()" method to reset a speaker audio for voice cloning.')
        self.set_speaker_embedding(speaker_audio)

        self.output_dir = os.path.abspath(output_dir)
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def get_speaker_embedding(self):
        return self._speaker_embedding.numpy()

    @paddle.no_grad()
    def set_speaker_embedding(self, speaker_audio: str):
        assert os.path.exists(speaker_audio), f'Speaker audio file: {speaker_audio} does not exists.'
        mel_sequences = self.speaker_processor.extract_mel_partials(
            self.speaker_processor.preprocess_wav(speaker_audio))
        self._speaker_embedding = self.speaker_encoder.embed_utterance(paddle.to_tensor(mel_sequences))

        logger.info(f'Speaker embedding has been set from file: {speaker_audio}')

    @paddle.no_grad()
    def generate(self, data: Union[str, List[str]], use_gpu: bool = False):
        assert self._speaker_embedding is not None, f'Set speaker embedding before voice cloning.'

        if isinstance(data, str):
            data = [data]
        elif isinstance(data, list):
            assert len(data) > 0 and isinstance(data[0],
                                                str) and len(data[0]) > 0, f'Input data should be str of List[str].'
        else:
            raise Exception(f'Input data should be str of List[str].')

        paddle.set_device('gpu') if use_gpu else paddle.set_device('cpu')
        files = []
        for idx, text in enumerate(data):
            phone_ids = self.frontend.get_input_ids(text, merge_sentences=True)["phone_ids"][0]
            wav = self.pwg_inference(self.fastspeech2_inference(phone_ids, spk_emb=self._speaker_embedding))
            output_wav = os.path.join(self.output_dir, f'{idx+1}.wav')
            sf.write(output_wav, wav.numpy(), samplerate=self.sample_rate)
            files.append(output_wav)

        return files


import cv2
import numpy as np

import torch
import torch.nn as nn

import random
from BPTSN import BPTSN
from tsn import TSN
import pdb
from tqdm import tqdm
from confusion_matrix_figure import draw_confusion_matrix
from multi_label_loss import MLL
from multi_label_loss import MLLSampler
from collections.abc import Iterable
from R2Plus1D import R2Plus1DClassifier
from W3D import W3D
from BPI3D import BPI3D
from BPC3D import BPC3D
from DI3D import DI3D, MBI3D

from torch.utils.tensorboard import SummaryWriter
from spatial_transform import (
    Compose, Normalize, Scale, CenterCrop, CornerCrop, MultiScaleCornerCrop,
    MultiScaleRandomCrop, RandomHorizontalFlip, ToTensor)
from target_transform import ClassLabel, VideoID
from temporal_transform import TemporalRandomCrop, TemporalCenterCrop, TemporalBeginCrop, LoopPadding, RepeatPadding
from pev import PEV
from charades_dataset import Charades as Dataset
from pytorch_i3d import InceptionI3d
import numpy as np
import videotransforms
from torchvision import datasets, transforms
import torchvision
from torch.autograd import Variable
from torch.optim import lr_scheduler
import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn
import torch
import argparse
import sys
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"]='0,1,2,3'
#from mi3d import MInceptionI3d

parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, help='rgb or flow')
parser.add_argument('--save_model', type=str)
parser.add_argument('--root', type=str)
parser.add_argument('--fuse', type=str, default='cat',
                    choices=['cat', 'cbp', 'cbp+cat'])
parser.add_argument('--eval', action='store_true')
parser.add_argument('--visualize', action='store_true')
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--epochs', type=int, default=80)
parser.add_argument('--model', type=str,
                    choices=['i3d', 'r2plus1d', 'w3d', 'tsn', 'bptsn', 'bpi3d', 'bpc3d', 'di3d', 'mbi3d'])
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--sample_freq', type=int, default=1)
parser.add_argument('--sample_step', type=int, default=1)
parser.add_argument('--n_samples', type=int, default=6,
                    help='num of samples for each video')
parser.add_argument('--clip_len', type=int, default=32)
parser.add_argument('--resume', type=str, default=None)
parser.add_argument('--print_loss', action='store_true')
parser.add_argument('--view', type=str, choices=['s', 'f', 'fs'])
parser.add_argument('--split_idx', type=int, choices=[1, 2, 3])

args = parser.parse_args()

split_idx = args.split_idx

data_root = '/home/lizhongguo/dataset/pev_of'
torch.manual_seed(123)
np.random.seed(123)

# integrated gradients

def integrated_gradients(inputs, model, target_label_idx, predict_and_gradients, baseline=None, steps=50, cuda=True):
    if baseline is None:
        baseline = [0 * i for i in inputs]
    # scale inputs and compute gradients
    scaled_inputs = [[bl + (float(i) / steps) * (sinput - bl)
                      for i in range(0, steps + 1)] for sinput, bl in zip(inputs, baseline)]
    grads, _ = predict_and_gradients(
        scaled_inputs, model, target_label_idx, cuda)
    avg_grads = np.average(grads[:-1], axis=0)
    print(avg_grads.shape)
    integrated_grad = [(i.cpu().numpy() - b.cpu().numpy())
                       * ag for i, b, ag in zip(inputs, baseline, avg_grads)]
    return integrated_grad

def calculate_outputs_and_gradients(inputs, model, target_label_idx, cuda=False):
    # do the pre-processing
    gradients = []
    for input_f, input_s in zip(inputs[0], inputs[1]):
        input_f.requires_grad = True
        input_s.requires_grad = True
        output = model(input_f, input_s)
        loss = F.cross_entropy(output, target_label_idx)
        # clear grad
        model.zero_grad()
        loss.backward()
        # print(input_f.grad)
        gradient = [input_f.grad.detach().cpu().numpy(),
                    input_s.grad.detach().cpu().numpy()]
        gradients.append(gradient)
        model.zero_grad()

    return np.array(gradients), target_label_idx

def convert_to_gray_scale(attributions):
    return np.average(attributions, axis=2)

def linear_transform(attributions, clip_above_percentile=99.9, clip_below_percentile=70.0, low=0.2, plot_distribution=False):
    m = compute_threshold_by_top_percentage(
        attributions, percentage=100-clip_above_percentile, plot_distribution=plot_distribution)
    e = compute_threshold_by_top_percentage(
        attributions, percentage=100-clip_below_percentile, plot_distribution=plot_distribution)
    transformed = (1 - low) * (np.abs(attributions) - e) / (m - e) + low
    transformed *= np.sign(attributions)
    transformed *= (transformed >= low)
    transformed = np.clip(transformed, 0.0, 1.0)
    return transformed

def compute_threshold_by_top_percentage(attributions, percentage=60, plot_distribution=True):
    if percentage < 0 or percentage > 100:
        raise ValueError('percentage must be in [0, 100]')
    if percentage == 100:
        return np.min(attributions)
    flat_attributions = attributions.flatten()
    attribution_sum = np.sum(flat_attributions)
    sorted_attributions = np.sort(np.abs(flat_attributions))[::-1]
    cum_sum = 100.0 * np.cumsum(sorted_attributions) / attribution_sum
    threshold_idx = np.where(cum_sum >= percentage)[0][0]
    threshold = sorted_attributions[threshold_idx]
    if plot_distribution:
        raise NotImplementedError
    return threshold

def polarity_function(attributions, polarity):
    if polarity == 'positive':
        return np.clip(attributions, 0, 1)
    elif polarity == 'negative':
        return np.clip(attributions, -1, 0)
    else:
        raise NotImplementedError

def overlay_function(attributions, image):
    return np.clip(0.7 * image + 0.5 * attributions, 0, 1)

def visualize(attributions, image, positive_channel=[255, 255, 0], negative_channel=[255, 0, 0], polarity='positive',
              clip_above_percentile=99.9, clip_below_percentile=0, morphological_cleanup=False,
              structure=np.ones((3, 3)), outlines=False, outlines_component_percentage=90, overlay=True,
              mask_mode=False, plot_distribution=False):
    if polarity == 'both':
        raise NotImplementedError

    elif polarity == 'positive':
        attributions = polarity_function(attributions, polarity=polarity)
        channel = positive_channel

    # convert the attributions to the gray scale
    attributions = convert_to_gray_scale(attributions)
    attributions = linear_transform(attributions, clip_above_percentile,
                                    clip_below_percentile, 0.0, plot_distribution=plot_distribution)
    attributions_mask = attributions.copy()
    if morphological_cleanup:
        raise NotImplementedError
    if outlines:
        raise NotImplementedError
    attributions = np.expand_dims(attributions, 2) * channel
    if overlay:
        if mask_mode == False:
            attributions = overlay_function(attributions, image)
        else:
            attributions = np.expand_dims(attributions_mask, 2)
            attributions = np.clip(attributions * image, 0, 1)
    return attributions

def evaluate(max_steps=320):

    if args.model == 'i3d' or args.model == 'bpi3d' or args.model == 'di3d' or args.model == 'mbi3d':
        scale_size = 224
    elif args.model == 'r2plus1d' or args.model == 'w3d' or args.model == 'bpc3d':
        scale_size = 112
    elif args.model == 'tsn' or args.model == 'bptsn':
        scale_size = 224  # 299
    else:
        raise Exception('Model %s not implemented' % args.model)

    if args.model == 'i3d' or args.model == 'r2plus1d' or args.model == 'bpi3d' or args.model == 'bpc3d' \
            or args.model == 'di3d' or args.model == 'mbi3d':
        if args.mode == 'rgb':
            mean = [0.5, 0.5, 0.5]
            std = [0.5, 0.5, 0.5]
        elif args.mode == 'flow':
            mean = [0.5, 0.5]
            std = [0.03, 0.03]
        elif args.mode == 'rgb+flow':
            mean = [0.5, 0.5, 0.5, 0.5, 0.5]
            std = [0.5, 0.5, 0.5, 0.03, 0.03]

    elif args.model == 'tsn' or args.model == 'bptsn':
        if args.mode == 'rgb':
            mean = [104./255., 117./255., 128./255.]
            std = [1., 1., 1.]
        else:
            mean = [0.5, 0.5, 0.5]
            std = [1., 1., 1.]

    else:
        raise Exception('Model %s not implemented' % args.model)

    train_transforms = Compose([MultiScaleRandomCrop([1.0, 0.95, 0.95*0.95], scale_size),
                                RandomHorizontalFlip(),
                                ToTensor(255.0),
                                Normalize(mean, std)
                                ])

    test_transforms = Compose([CenterCrop(scale_size),
                               ToTensor(255.0),
                               Normalize(mean, std)
                               ])

    clip_len = args.clip_len // args.sample_step
    temporal_transforms = Compose([TemporalRandomCrop(clip_len),
                                   RepeatPadding(clip_len)])

    val_temporal_transforms = Compose([TemporalBeginCrop(clip_len),
                                       RepeatPadding(clip_len)])
    val_dataset = PEV(
        data_root,
        '/home/lizhongguo/dataset/pev_split/val_split_%d.txt' % split_idx,
        'evaluation',
        args.n_samples,
        spatial_transform=test_transforms,
        temporal_transform=val_temporal_transforms,
        target_transform=ClassLabel(),
        sample_duration=clip_len, sample_freq=args.sample_freq, mode=args.mode, sample_step=args.sample_step, view=args.view)

    val_sampler = None

    val_dataloader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=(val_sampler is None), num_workers=2, pin_memory=True, sampler=val_sampler, drop_last=False)

    # setup the model
    if args.model == 'mbi3d':
        model = MBI3D(7, args.fuse, ['f%s' % args.mode, 's%s' % args.mode])
        checkpoint = torch.load(
            'pev_split_%d_%s_%s_best%sfs.pt' % (args.split_idx, args.model, args.mode,
                                                '_' if args.fuse == 'cbp' else '_'+args.fuse+'_'), map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'])
        model.cuda()
        model.train(False)
        model.eval()

    for data in val_dataloader:
        input_f, input_s, labels, index = data
        input_f, input_s = input_f.cuda(), input_s.cuda()
        labels = labels.cuda(non_blocking=True)
        #baseline_f = torch.zeros_like(input_f)
        #baseline_s = torch.zeros_like(input_s)
        #output = model(input_f, input_s)
        ig = integrated_gradients(
            (input_f, input_s), model, labels, calculate_outputs_and_gradients)
        ig = [np.transpose(ig[0],(0, 2, 3, 4, 1)), np.transpose(ig[1],(0, 2, 3, 4, 1))]

        m = 0.5
        v = 0.03 if args.mode == 'flow' else 0.5
        input_f.div_(1/v).sub_(-m)
        input_s.div_(1/v).sub_(-m)
        input_f = input_f.permute(0, 2, 3, 4, 1)
        input_s = input_s.permute(0, 2, 3, 4, 1)

        video_id = val_dataset.data[index]['video_id']

        for attr_f, attr_s, input_f, input_s in zip(ig[0], ig[1], input_f, input_s):
            # N
            idx = 0
            for a_f, a_s, i_f, i_s in zip(attr_f, attr_s, input_f, input_s):
                idx += 1
                img_integrated_gradient_overlay = visualize(a_f, i_f.cpu().numpy(), clip_above_percentile=99, clip_below_percentile=90,
                                                            overlay=True, mask_mode=False)
                cv2.imwrite('vis/ig_%05d_f_%d.jpg' %
                            (video_id, idx), 255*img_integrated_gradient_overlay[:,:,(2,1,0)])
                img_integrated_gradient_overlay = visualize(a_s, i_s.cpu().numpy(), clip_above_percentile=99, clip_below_percentile=90,
                                                            overlay=True, mask_mode=False)
                cv2.imwrite('vis/ig_%05d_s_%d.jpg' %
                            (video_id ,idx), 255*img_integrated_gradient_overlay[:,:,(2,1,0)])

if __name__ == '__main__':
    evaluate()


#!/usr/bin/env python3

import csv
import os
import random
import time

import easyargs
import matplotlib
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torchvision
from torch import nn, optim
from torch.optim import lr_scheduler

matplotlib.use('Agg')
# from create_dataset_cifar100 import myDataset
PACKAGE_PARENT = '..'
SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
os.sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
from glo.utils import save_checkpoint, get_cub_param, get_loader_with_idx, get_cifar100_small, get_classifier, \
	get_test_stl_param, get_train_n_unlabled_stl_param
from cub2011 import Cub2011

from glo.evaluation import fewshot_setting, accuracy
from logger import Logger

dim = 512

def train_linear_classifier(model, criterion, optimizer, train_data, num_epochs, seed, name, offset_idx=0,
                            offset_label=0, batch_size=128, fewshot=False, augment=False, autoaugment=False,
                            start_epoch=0, aug_param=None, test_data=None):
	if aug_param is None:
		aug_param = {'std': None, 'mean': None, 'rand_crop': 32, 'image_size': 32}
	model.train()
	sampler = None
	milestones = [60, 120, 160]
	if fewshot:
		num_epoch_repetition = fewshot_setting(train_data)
		num_epochs *= num_epoch_repetition
		milestones = list(map(lambda x: num_epoch_repetition * x, milestones))
	train_data_loader = get_loader_with_idx(train_data, batch_size, num_workers=8,
	                                        augment=augment,
	                                        offset_idx=offset_idx, offset_label=offset_label, sampler=sampler,
	                                        autoaugment=autoaugment, **aug_param)
	# train_data_loader = DataLoader(train_data, batch_size=128, shuffle=True, num_workers=4)
	print(f" train_len:{len(train_data)}")
	if aug_param['image_size'] > 32:
		scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.9)
	else:
		scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=0.2)
	train_loss = 0
	correct = 0
	total = 0
	use_cuda = torch.cuda.is_available()
	for epoch in range(start_epoch, num_epochs):
		model.train()
		end_epoch_time = time.time()
		for i, (_, inputs, targets) in enumerate(train_data_loader):
			if use_cuda:
				inputs, targets = inputs.cuda(), targets.cuda(non_blocking=True)
			optimizer.zero_grad()
			output = model(inputs)
			# _, preds = torch.max(output.data, 1)
			loss = criterion(output, targets)
			loss.backward()
			optimizer.step()

			train_loss += loss.item()
			_, predicted = torch.max(output.data, 1)
			total += targets.size(0)
			correct += predicted.eq(targets.data).sum()
			acc = 100. * correct / total
			# sys.stdout.write('\r')
			# sys.stdout.write('| Epoch [%3d/%3d] Iter[%3d/%3d]\t\tLoss: %.4f Acc@1: %.3f%%'
			#                  % (epoch, num_epochs, i + 1,
			#                     (len(train_data) // batch_size) + 1, loss.item(), acc))
			# sys.stdout.flush()
		print(f" =>Epoch[{epoch}] acc:{acc} lr:{scheduler.get_lr()[0]}  Time: {(time.time() - end_epoch_time) / 60:8.2f}m")
		if epoch % 10 == 0 and test_data is not None:
			fewshot_acc_1_test, fewshot_acc_5_test = accuracy(model, test_data, batch_size=batch_size, aug_param=aug_param)
			print(f"\n=> Mid accuracy@1: {fewshot_acc_1_test} | accuracy@5: {fewshot_acc_5_test}\n")
			model.train()
		# if epoch % 5 == 0:
		# 	checkpoint(acc, epoch, model, seed, name)
		scheduler.step(epoch)
		# save_checkpoint(f'baseline/{name}', epoch, model)

# def train_linear_classifier_2_datasets(model, criterion, optimizer, train_data1, train_data2, test_data, num_epochs,
#                                        resolution=84, offset_idx=0, offset_label=0, batch_size=128):
# 	model.train()
# 	dataloaders_large = get_loader_with_idx(train_data1, batch_size, image_size=resolution, num_workers=8,
# 	                                        augment=False)
# 	dataloaders_small = get_loader_with_idx(train_data2, batch_size, image_size=resolution, num_workers=8,
# 	                                        augment=False,
# 	                                        offset_idx=offset_idx, offset_label=offset_label)
#
# 	# train_data_loader = DataLoader(train_data,batch_size=batch_size,shuffle=True, num_workers=2, pin_memory=False)
#
# 	for epoch in range(num_epochs):
# 		for i, batch in enumerate(dataloaders_large):
# 			# dataloader_iterator = iter(dataloaders_small)
# 			# idx = batch[0]
# 			imgs = maybe_cuda(batch[1])
# 			targets1 = maybe_cuda(batch[2])
# 			optimizer.zero_grad()
# 			if i <= len(train_data2):
# 				data2 = next(dataloader_iterator)
# 			else:
# 				dataloader_iterator = iter(dataloaders_small)
# 				data2 = next(dataloader_iterator)
# 			imgs2 = maybe_cuda(data2[1])
# 			targets2 = maybe_cuda(data2[2])
# 			targets = torch.cat([targets1, targets2])
# 			output = model(torch.cat([imgs, imgs2]).squeeze())
# 			# output = model(imgs).squeeze()
# 			# _, preds = torch.max(output.data, 1)
# 			loss = criterion(output, targets.view(-1, 1).squeeze()).mean()
# 			loss.backward()
# 			optimizer.step()
# 		print(f"Epoch:{epoch}, loss:{loss}")
# 		acc_1_valid, acc_5_valid = accuracy(model, train_data1)
# 		fewshot_acc_1_valid, fewshot_acc_5_valid = accuracy(model, train_data2)
# 		test_acc_1_valid, test_acc_5_valid = accuracy(model, test_data)
# 		print('acc_1_valid accuracy@1', acc_1_valid)
# 		print('acc_5_valid accuracy@5', acc_5_valid)
# 		print('fewshot_acc_1_valid accuracy@1', fewshot_acc_1_valid)
# 		print('fewshot_acc_5_valid accuracy@1', fewshot_acc_5_valid)
# 		print('test_acc_1_valid accuracy@1', test_acc_1_valid)
# 		print('test_acc_5_valid accuracy@1', test_acc_5_valid)

# for i, batch in enumerate(dataloaders_small):
# 	idx = batch[0]
# 	imgs = maybe_cuda(batch[1])
# 	targets1 = maybe_cuda(batch[2])
# 	optimizer.zero_grad()
# 	output = model(imgs).squeeze()
# 	loss = criterion(output, targets1.view(-1, 1).squeeze())#.mean()
# 	loss.backward()
# 	optimizer.step()
# 	if loss < best_loss:
# 		best_loss = loss
# print(f"Epoch:{epoch}, loss:{loss}")
# acc_1_valid, acc_5_valid = accuracy(model, train_data1)
# fewshot_acc_1_valid, fewshot_acc_5_valid = accuracy(model, train_data2)
# test_acc_1_valid, test_acc_5_valid = accuracy(model, test_data)
# print('acc_1_valid accuracy@1', acc_1_valid)
# print('acc_5_valid accuracy@5', acc_5_valid)
# print('fewshot_acc_1_valid accuracy@1', fewshot_acc_1_valid)
# print('fewshot_acc_5_valid accuracy@1', fewshot_acc_5_valid)
# print('test_acc_1_valid accuracy@1', test_acc_1_valid)
# print('test_acc_5_valid accuracy@1', test_acc_5_valid)

# @easyargs
# def run_eval(seed=0, debug=False):
# 	image_path_train = '/cs/labs/daphna/daphna/data/miniimagenet/train'
# 	image_path_test = '/cs/labs/daphna/daphna/data/miniimagenet/test'
# 	PATH = "/cs/labs/daphna/idan.azuri/myglo/glo/"
# 	if debug:
# 		image_path_test = image_path_train = '/Users/idan.a/repos/data/miniimagenet/test'
# 		PATH = "/Users/idan.a/repos/myglo/glo/"
# 	train_data = DatasetFolderFromList(image_path_train)
# 	offset_idx = max(train_data.path2idx.values()) + 1
# 	offset_labels = max(train_data.targets) + 1
# 	fewshot_train_data = DatasetFolderFromList(image_path_test, class_cap=500, offset_idx=offset_idx)
# 	updated_offset_idx = max(fewshot_train_data.path2idx.values()) + 1
# 	fewshot_test_data = DatasetFolderFromList(image_path_test, offset_idx=updated_offset_idx,
# 	                                          ignore_samples_by_path=fewshot_train_data.path2idx)
# 	random.seed(seed)
# 	torch.manual_seed(seed)
# 	np.random.seed(seed)
# 	print(f"train_data: {len(train_data)}")
# 	print(f"fewshot_train_data: {len(fewshot_train_data)}")
# 	print(f"test data: {len(fewshot_test_data)}")
# 	print(f"train_data.classes: {len(train_data.classes)}")
# 	print(f"fewshot_train_data.classes: {len(fewshot_train_data.classes)}")
# 	total_classes = len(train_data.classes) + len(fewshot_train_data.classes)
# 	print(f"total classes: {total_classes}")
#
# 	# data_train, valid_data_test = train_valid_split(train_data, split_fold=10, random_seed=seed)
# 	# fewshot_data_for_train, fewshot_data_for_test = train_valid_split(fewshot_train_data, split_fold=10, random_seed=seed)
#
# 	# classifier = cnn(num_classes=total_classes, im_size=resolution)  # Pixel space
# 	# classifier = models.resnet152(pretrained=False, num_classes=100)
# 	# classifier = WideResNet(depth=28, num_classes=100, widen_factor=10, dropRate=0.3)
# 	classifier = cnn(num_classes=100, im_size=84)  # Pixel space
# 	criterion = nn.CrossEntropyLoss()
# 	num_gpus = torch.cuda.device_count()
# 	if num_gpus > 1:
# 		print(f"=> Using {num_gpus} GPUs")
# 		classifier = nn.DataParallel(classifier.cuda(), device_ids=list(range(num_gpus)))
# 		# criterion = nn.DataParallel(classifier.cuda(), device_ids=list(range(num_gpus)))
# 		criterion = criterion.cuda()
# 	else:
# 		classifier = maybe_cuda(classifier)
# 		criterion = maybe_cuda(criterion)
# 	optimizer = optim.Adam(classifier.parameters(), lr=1e-3)
# 	train_linear_classifier_2_datasets(classifier, criterion, optimizer, train_data, fewshot_train_data,
# 	                                   test_data=fewshot_test_data, num_epochs=30, offset_idx=offset_idx,
# 	                                   offset_label=offset_labels)
# 	# train_linear_classifier(classifier,criterion, optimizer, fewshot_train_data, num_epochs=30, offset_idx=offset_idx, offset_label=offset_labels)
# 	# train_linear_classifier(classifier,criterion, optimizer, train_data, num_epochs=20)
#
# 	acc_1_valid, acc_5_valid = accuracy(classifier, train_data)
# 	fewshot_acc_1_valid, fewshot_acc_5_valid = accuracy(classifier, fewshot_train_data)
# 	fewshot_acc_1_test, fewshot_acc_5_test = accuracy(classifier, fewshot_test_data)
#
# 	print('acc_1_valid accuracy@1', acc_1_valid)
# 	print('acc_5_valid accuracy@5', acc_5_valid)
# 	print('fewshot_acc_1_valid accuracy@1', fewshot_acc_1_valid)
# 	print('fewshot_acc_5_valid accuracy@1', fewshot_acc_5_valid)
# 	print('fewshot_acc_1_test accuracy@1', fewshot_acc_1_test)
# 	print('fewshot_acc_5_test accuracy@5', fewshot_acc_5_test)
# 	# name = rn.split("_")[1:-5]
# 	# name='_'.join(name[:-1])
# 	# training data
#
# 	w = csv.writer(open(f"CNN_test_acc1.csv", "w+"))
# 	w.writerow(fewshot_acc_1_test)
#
# 	w = csv.writer(open(f"CNN_test_acc5.csv", "w+"))
# 	w.writerow(fewshot_acc_5_test)

@easyargs
def run_eval_generic(epochs=200, d="", fewshot=False, augment=False, autoaugment=False, data="", shot=None,
                     resume=False, pretrained=False,batch_size=128):
	global classifier
	dir = "baselines_aug"
	if not os.path.isdir(dir):
		os.mkdir(dir)
	name = f"{data}_{d}_baseline_"
	if not shot is None:
		name = name + f"shot{shot}"
	if not autoaugment is None:
		name = name + "_aug"

	logger = Logger(os.path.join(f"{dir}/", f'{d}_log.txt'), title=name)
	logger.set_names(["valid_acc@1", "valid_acc@5", "test_acc@1", "test_acc@5"])
	data_dir = '../../data'
	cifar_dir_small="../../data/cifar-100"
	aug_param = aug_param_test=None

	if data == "cifar":
		# batch_size = 128
		lr = 0.1
		classes = 100
		WD = 5e-4
		if not fewshot:
			train_data = torchvision.datasets.CIFAR100(root=data_dir, train=True, download=True)
			transductive_train_data = torchvision.datasets.CIFAR100(root=data_dir, train=False, download=True)
			test_data = transductive_train_data
		else:
			print("=> Fewshot")
			train_data, transductive_train_data = get_cifar100_small(cifar_dir_small, shot)
			test_data = torchvision.datasets.CIFAR100(root=data_dir, train=False, download=True)
		print(f"train_data size:{len(train_data)}")
		print(f"test_data size:{len(test_data)}")
	if data == "cub":
		print("CUB-200")
		classes = 200
		# batch_size = 16
		lr = 0.001
		WD = 1e-5
		aug_param = aug_param_test=get_cub_param()
		train_data = Cub2011(root=f"../../data/{data}", train=True)
		test_data = Cub2011(root=f"../../data/{data}", train=False)
		print(f"train_data size:{len(train_data)},{train_data.data.shape}")
		print(f"test_data size:{len(test_data)},{test_data.data.shape}")
	if data == "stl":
		print("STL-10")
		classes = 10
		WD=4e-4
		# batch_size = 128
		lr=2e-3
		aug_param = get_train_n_unlabled_stl_param()
		aug_param_test = get_test_stl_param()
		train_data = torchvision.datasets.STL10(root=f"../../data/{data}", split='train', download=True)
		# transductive_train_data = torchvision.datasets.STL10(root=f"../../data/{data}", split='unlabeled', download=True)
		test_data = torchvision.datasets.STL10(root=f"../../data/{data}", split='test', download=True)
		print(f"train_data size:{len(train_data)},{train_data.data.shape}")
		print(f"test_data size:{len(test_data)},{test_data.data.shape}")
	scores_test_acc1_fewshot = dict()
	scores_test_acc5_fewshot = dict()
	for seed in [1, 17, 32]:
		# offset_idx = len(train_data)
		print(f"{data} num classes: {classes}")
		random.seed(seed)
		torch.manual_seed(seed)
		np.random.seed(seed)
		start_epoch = 0
		if resume:
			# Load checkpoint.
			print('==> Resuming from checkpoint..')
			assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
			checkpoint = torch.load(f'./checkpoint/{name}_{seed}_ckpt.t7')
			classifier = checkpoint['net']
			best_acc = checkpoint['acc']
			start_epoch = checkpoint['epoch'] + 1
			rng_state = checkpoint['rng_state']
			torch.set_rng_state(rng_state)
			print(f"=> Model loaded start_epoch{start_epoch}, acc={best_acc}")

		else:
			classifier = get_classifier(classes, d, pretrained)
		criterion = nn.CrossEntropyLoss().cuda()
		num_gpus = torch.cuda.device_count()
		if num_gpus > 1:
			classifier = nn.DataParallel(classifier).cuda()
		else:
			classifier = classifier.cuda()
		optimizer = optim.SGD(classifier.parameters(), lr, momentum=0.9, weight_decay=WD, nesterov=True)
		cudnn.benchmark = True
		print(' => Total params: %.2fM' % (sum(p.numel() for p in classifier.parameters()) / 1000000.0))
		print(f"=> {d}  Training model")
		print(f"=> Training Epochs = {str(epochs)}")
		print(f"=> Initial Learning Rate = {str(lr)}")

		# criterion = maybe_cuda(criterion)
		# optimizer = optim.Adam(classifier.parameters(), lr=1e-3)
		train_linear_classifier(classifier, criterion, optimizer, seed=seed, name=name, train_data=train_data,
		                        batch_size=batch_size,
		                        num_epochs=epochs,
		                        fewshot=fewshot,
		                        augment=augment, autoaugment=autoaugment, aug_param=aug_param, test_data=test_data,
		                        start_epoch=start_epoch)

		acc_1_valid, acc_5_valid = accuracy(classifier, train_data, batch_size=batch_size, aug_param=aug_param)
		fewshot_acc_1_test, fewshot_acc_5_test = accuracy(classifier, test_data, batch_size=batch_size, aug_param=aug_param_test)

		print('fewshot_acc_1_valid accuracy@1', acc_1_valid)
		print('fewshot_acc_5_valid accuracy@1', acc_5_valid)

		print('fewshot_acc_1_test accuracy@1', fewshot_acc_1_test)
		print('fewshot_acc_5_test accuracy@5', fewshot_acc_5_test)
		logger.append([acc_1_valid, acc_5_valid, fewshot_acc_1_test, fewshot_acc_5_test])
		scores_test_acc1_fewshot[seed] = fewshot_acc_1_test
		scores_test_acc5_fewshot[seed] = fewshot_acc_5_test

	w = csv.writer(open(f"baseline_shot_{shot}_acc1.csv", "w+"))
	for key, val in scores_test_acc1_fewshot.items():
		w.writerow([key, val])
	w = csv.writer(open(f"baseline_shot_{shot}_acc5.csv", "w+"))
	for key, val in scores_test_acc5_fewshot.items():
		w.writerow([key, val])
	logger.close()

def checkpoint(acc, epoch, net, seed, name):
	# Save checkpoint.
	print('Saving..')
	state = {
		'net': net,
		'acc': acc,
		'epoch': epoch,
		'rng_state': torch.get_rng_state()
	}
	if not os.path.isdir('checkpoint'):
		os.mkdir('checkpoint')
	torch.save(state, f'./checkpoint/{name}_{seed}_ckpt.t7')

if __name__ == '__main__':
	run_eval_generic()
	print("Done!")


"""
Module to run tests on SlitTraceSet
"""
import os
import pytest

import numpy as np

from pypeit.images import detector_container
from pypeit import io

def data_path(filename):
    data_dir = os.path.join(os.path.dirname(__file__), 'files')
    return os.path.join(data_dir, filename)

# Example (shane_kast_blue)
def_det = dict(
    dataext=0,
    specaxis=1,
    specflip=False,
    spatflip=False,
    platescale=0.43,
    saturation=65535.,
    mincounts=-1e10,
    nonlinear=0.76,
    numamplifiers=2,
    gain=np.asarray([1.2, 1.2]),
    ronoise=np.asarray([3.7, 3.7]),
    det=1,
    xgap=0.,
    ygap=0.,
    ysize=1.,
    darkcurr=0.0,
    binning='1,1',
    datasec=np.asarray(['[:, 1:1024]', '[:, 1025:2048]']),  # These are rows, columns on the raw frame, 1-indexed
    oscansec=np.asarray(['[:, 2050:2080]', '[:, 2081:2111]']))

def test_init():
    detector = detector_container.DetectorContainer(**def_det)

def test_bundle():
    detector = detector_container.DetectorContainer(**def_det)
    data = detector._bundle()
    assert len(data) == 1

def test_io():
    detector = detector_container.DetectorContainer(**def_det)
    detector.to_file(data_path('tmp_detector.fits'), overwrite=True)

    _new_detector = detector.from_file(data_path('tmp_detector.fits'))

    os.remove(data_path('tmp_detector.fits'))


#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------
"""
Convert a Bert large model exported by Keras2Onnx to a tiny model for test purpose.
The input model is generated like the following (need install keras2onnx from source):

    import numpy
    import keras2onnx
    from transformers import (TFBertForQuestionAnswering, BertTokenizer)

    tokenizer = BertTokenizer.from_pretrained'bert-large-uncased-whole-word-masking-finetuned-squad', do_lower_case=True, cache_dir=cache_dir)
    model = TFBertForQuestionAnswering.from_pretrained(model_name_or_path, cache_dir=cache_dir)

    question, text = "What is ONNX Runtime?", "ONNX Runtime is a performance-focused inference engine for ONNX models."
    inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors='tf')

    output_model_path =  os.path.join(output_dir, 'keras_{}.onnx'.format(model_name_or_path))
    if not os.path.exists(output_model_path):
        model.predict(inputs)
        onnx_model = keras2onnx.convert_keras(model, model.name)
        keras2onnx.save_model(onnx_model, output_model_path)
"""

import onnx
import onnx.utils
import sys
import argparse
import numpy as np
from onnx import ModelProto, TensorProto, numpy_helper
from onnxruntime_tools.transformers.onnx_model import OnnxModel
import os
import onnxruntime
import random
from pathlib import Path
import timeit

DICT_SIZE = 20
SEQ_LEN = 7
""" This class creates a tiny bert model for test purpose. """

class TinyBertOnnxModel(OnnxModel):
    def __init__(self, model, verbose):
        super(TinyBertOnnxModel, self).__init__(model, verbose)
        self.resize_model()

    def resize_weight(self, initializer_name, target_shape):
        weight = self.get_initializer(initializer_name)
        w = numpy_helper.to_array(weight)

        target_w = w
        if len(target_shape) == 1:
            target_w = w[:target_shape[0]]
        elif len(target_shape) == 2:
            target_w = w[:target_shape[0], :target_shape[1]]
        elif len(target_shape) == 3:
            target_w = w[:target_shape[0], :target_shape[1], :target_shape[2]]
        else:
            print("at most 3 dimensions")

        tensor = onnx.helper.make_tensor(name=initializer_name + '_resize',
                                         data_type=TensorProto.FLOAT,
                                         dims=target_shape,
                                         vals=target_w.flatten().tolist())

        return tensor

    def resize_model(self):
        graph = self.model.graph
        initializers = graph.initializer

        # parameters of input base model.
        old_parameters = {
            "seq_len": 26,
            "hidden_size": 1024,
            "num_heads": 16,
            "size_per_head": 64,
            "word_dict_size": [28996, 30522],  # list of supported dictionary size.
            "max_word_position": 512
        }

        # parameters of output tiny model.
        new_parameters = {
            "seq_len": SEQ_LEN,
            "hidden_size": 8,
            "num_heads": 2,
            "size_per_head": 4,
            "word_dict_size": DICT_SIZE,
            "max_word_position": 10
        }

        for input in graph.input:
            if (input.type.tensor_type.shape.dim[1].dim_value == old_parameters["seq_len"]):
                print("input", input.name, input.type.tensor_type.shape)
                input.type.tensor_type.shape.dim[1].dim_value = new_parameters["seq_len"]
                print("=>", input.type.tensor_type.shape)

        reshapes = {}
        for initializer in initializers:
            tensor = numpy_helper.to_array(initializer)
            dtype = np.float32 if initializer.data_type == 1 else np.int32
            if len(tensor.shape) == 1 and tensor.shape[0] == 1:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>[", new_parameters["num_heads"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["num_heads"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>[", new_parameters["seq_len"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["seq_len"]], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>[", new_parameters["size_per_head"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["size_per_head"]], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>[", new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>[", 4 * new_parameters["hidden_size"], "]")
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray([4 * new_parameters["hidden_size"]], dtype=dtype),
                                                initializer.name))
            elif len(tensor.shape) == 0:
                if tensor == old_parameters["num_heads"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["num_heads"], "=>", new_parameters["num_heads"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["num_heads"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["seq_len"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["seq_len"], "=>", new_parameters["seq_len"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["seq_len"], dtype=dtype), initializer.name))
                elif tensor == old_parameters["size_per_head"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["size_per_head"], "=>", new_parameters["size_per_head"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["size_per_head"], dtype=dtype),
                                                initializer.name))
                elif tensor == old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          old_parameters["hidden_size"], "=>", new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 4 * old_parameters["hidden_size"]:
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          4 * old_parameters["hidden_size"], "=>", 4 * new_parameters["hidden_size"])
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(4 * new_parameters["hidden_size"], dtype=dtype),
                                                initializer.name))
                elif tensor == 1.0 / np.sqrt(old_parameters["size_per_head"]):
                    print("initializer type={}".format(initializer.data_type), initializer.name,
                          1.0 / np.sqrt(old_parameters["size_per_head"]), "=>",
                          1.0 / np.sqrt(new_parameters["size_per_head"]))
                    initializer.CopyFrom(
                        numpy_helper.from_array(np.asarray(1.0 / np.sqrt(new_parameters["size_per_head"]), dtype=dtype),
                                                initializer.name))

            new_shape = []
            shape_changed = False
            for dim in tensor.shape:
                if (dim == old_parameters["hidden_size"]):
                    new_shape.append(new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim == 4 * old_parameters["hidden_size"]):
                    new_shape.append(4 * new_parameters["hidden_size"])
                    shape_changed = True
                elif (dim in old_parameters["word_dict_size"]):
                    new_shape.append(new_parameters["word_dict_size"])
                    shape_changed = True
                elif (dim == old_parameters["max_word_position"]):
                    new_shape.append(new_parameters["max_word_position"])
                    shape_changed = True
                else:
                    new_shape.append(dim)
            if shape_changed:
                reshapes[initializer.name] = new_shape
                print("initializer", initializer.name, tensor.shape, "=>", new_shape)

        for initializer_name in reshapes:
            self.replace_input_of_all_nodes(initializer_name, initializer_name + '_resize')
            tensor = self.resize_weight(initializer_name, reshapes[initializer_name])
            self.model.graph.initializer.extend([tensor])

        self.use_dynamic_axes()

    def use_dynamic_axes(self, dynamic_batch_dim='batch_size', seq_len=7):
        """
        Update input and output shape to use dynamic axes.
        """
        dynamic_batch_inputs = {}
        for input in self.model.graph.input:
            dim_proto = input.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = input.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len

        for output in self.model.graph.output:
            dim_proto = output.type.tensor_type.shape.dim[0]
            dim_proto.dim_param = dynamic_batch_dim
            dim_proto = output.type.tensor_type.shape.dim[1]
            dim_proto.dim_value = seq_len

def generate_test_data(onnx_file,
                       output_path,
                       batch_size,
                       sequence_length,
                       use_cpu=True,
                       input_tensor_only=False,
                       dictionary_size=DICT_SIZE,
                       test_cases=3):
    input_data_type = np.int32
    for test_case in range(test_cases):
        input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=input_data_type)
        tensor_1 = numpy_helper.from_array(input_1, 'input_ids')

        actual_seq_len = random.randint(sequence_length - 3, sequence_length)
        input_2 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        temp = np.ones((batch_size, actual_seq_len), dtype=input_data_type)
        input_2[:temp.shape[0], :temp.shape[1]] = temp
        tensor_2 = numpy_helper.from_array(input_2, 'attention_mask')

        input_3 = np.zeros((batch_size, sequence_length), dtype=input_data_type)
        tensor_3 = numpy_helper.from_array(input_3, 'token_type_ids')

        path = os.path.join(output_path, 'test_data_set_' + str(test_case))
        try:
            os.mkdir(path)
        except OSError:
            print("Creation of the directory %s failed" % path)
        else:
            print("Successfully created the directory %s " % path)

        if input_tensor_only:
            return

        sess_options = onnxruntime.SessionOptions()
        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
        sess = onnxruntime.InferenceSession(onnx_file, sess_options, providers=['CPUExecutionProvider'])

        input1_name = sess.get_inputs()[0].name
        output_names = [output.name for output in sess.get_outputs()]
        inputs = {'input_ids': input_1, 'attention_mask': input_2, 'token_type_ids': input_3}
        print("inputs", inputs)
        result = sess.run(output_names, inputs)

        with open(os.path.join(path, 'input_{}.pb'.format(0)), 'wb') as f:
            f.write(tensor_1.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(1)), 'wb') as f:
            f.write(tensor_2.SerializeToString())
        with open(os.path.join(path, 'input_{}.pb'.format(2)), 'wb') as f:
            f.write(tensor_3.SerializeToString())

        for i, output_name in enumerate(output_names):
            tensor_result = numpy_helper.from_array(
                np.asarray(result[i]).reshape((batch_size, sequence_length)), output_names[i])
            with open(os.path.join(path, 'output_{}.pb'.format(i)), 'wb') as f:
                f.write(tensor_result.SerializeToString())

        start_time = timeit.default_timer()

        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED

        path_prefix = onnx_file[:-5]  #remove .onnx suffix
        if use_cpu:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_cpu.onnx"
        else:
            sess_options.optimized_model_filepath = path_prefix + "_optimized_gpu.onnx"

        session = onnxruntime.InferenceSession(onnx_file, sess_options=sess_options,
                                               providers=onnxruntime.get_available_providers())
        if use_cpu:
            session.set_providers(['CPUExecutionProvider'])  # use cpu
        else:
            if 'CUDAExecutionProvider' not in session.get_providers():
                print("Warning: GPU not found")
                continue
        outputs = session.run(None, inputs)
        evalTime = timeit.default_timer() - start_time
        if outputs[0].tolist() != result[0].tolist():
            print("Error: not same result after optimization. use_cpu={}, no_opt_output={}, opt_output={}".format(
                use_cpu, result[0].tolist(), outputs[1].tolist()))
        print("** Evaluation done in total {} secs".format(evalTime))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, type=str)
    parser.add_argument('--output', required=True, type=str)
    parser.add_argument('--float16', required=False, action='store_true')
    parser.set_defaults(float16=False)
    args = parser.parse_args()

    model = ModelProto()
    with open(args.input, "rb") as f:
        model.ParseFromString(f.read())

    bert_model = TinyBertOnnxModel(model, False)

    if args.float16:
        bert_model.convert_model_float32_to_float16()

    bert_model.update_graph()
    bert_model.remove_unused_constant()

    print("opset version", bert_model.model.opset_import[0].version)

    with open(args.output, "wb") as out:
        out.write(bert_model.model.SerializeToString())

    p = Path(args.output)
    data_path = p.parent

    batch_size = 1
    sequence_length = SEQ_LEN

    generate_test_data(args.output, data_path, batch_size, sequence_length, use_cpu=not args.float16)

if __name__ == "__main__":
    main()