import torch
from torch.utils import data
from torch import nn
from torch.optim import lr_scheduler
from dataset.SynthText import SynthText
from dataset.ICDAR15 import ICDAR15
from network.model import EAST
from network.loss import Loss
from network.loss_target import Loss_target
import os
from shapely.geometry import *
import time
import numpy as np
from PIL import Image, ImageDraw
from tqdm import tqdm
from lib.detect import detect
from evaluate.script import getresult
import argparse
import os
from lib.pseduo_label import generate_pseduo

parser = argparse.ArgumentParser(description='EAST reimplementation')

# Model path
parser.add_argument('--resume', default="./SynthText.pth", type=str,
                    help='Checkpoint state_dict file to resume training from')
parser.add_argument('--target_pseudo_negative', default="./pseudo_label_negative/", type=str,
                    help='the pseduo label of target domain')
parser.add_argument('--target_pseudo_positive', default="./pseudo_label_positive/", type=str,
                    help='the positive sample of pseduo label in target domain')
parser.add_argument('--target_image', default="./Image/", type=str,
                    help='the image of target domain')
parser.add_argument('--path_save', default="./Model_save/", type=str,
                    help='save model')

# Training strategy
parser.add_argument('--epoch_iter', default=8000, type = int,
                    help='the max epoch iter')
parser.add_argument('--batch_size', default=6, type = int,
                    help='batch size of training')
# parser.add_argument('--cdua', default=True, type=str2bool,
#                     help='Use CUDA to train model')
parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,
                    help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
                    help='Momentum value for optim')
parser.add_argument('--weight_decay', default=5e-4, type=float,
                    help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float,
                    help='Gamma update for SGD')
parser.add_argument('--num_workers', default=16, type=int,
                    help='Number of workers used in dataloading')

args = parser.parse_args()

def train(epoch,  model, optimizer, train_loader_source,train_loader_target):
    model.train()
    scheduler.step()
    epoch_loss = 0
    epoch_time = time.time()

    for i, (source,target) in enumerate(zip(train_loader_source,train_loader_target)):
        start_time = time.time()
        img, gt_score, gt_geo, ignored_map = source
        img_target, gt_score_target, gt_geo_target, ignored_map_target = target

        # target domain training
        img_target, gt_score_target, gt_geo_target, ignored_map_target = \
        img_target.to(device), gt_score_target.to(device), gt_geo_target.to(device), ignored_map_target.to(device)

        (pred_score_target, pred_geo_target),img_class_t = model(img_target)

        geo_loss_target, classify_loss_target,doamin_loss_target = \
        criterion_target(gt_score_target, pred_score_target, gt_geo_target, pred_geo_target, ignored_map_target,img_class_t)

        # source domain training
        img, gt_score, gt_geo, ignored_map  = img.to(device), gt_score.to(device), gt_geo.to(device), ignored_map.to(device)

        (pred_score, pred_geo), img_class_s = model(img)

        geo_loss_source,classify_loss_source , doamin_loss_source  = \
        criterion(gt_score, pred_score, gt_geo, pred_geo, ignored_map, img_class_s)
        #
        # if epoch<5:
        #     k=0
        # else:
        #     k=epoch%5 * 0.1
        loss =  (geo_loss_source + classify_loss_source + 0.02*classify_loss_target) + (doamin_loss_target + doamin_loss_source)*0.1
        epoch_loss += loss.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print('Epoch is [{}/{}], mini-batch is [{}/{}], time consumption is {:.8f}, batch_loss is {:.8f}'.format( \
            epoch + 1, args.epoch_iter, i + 1, int(1000 / args.batch_size), time.time() - start_time, loss.item()))

    print('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(epoch_loss / int(1000 / args.batch_size),time.time() - epoch_time))
    print(time.asctime(time.localtime(time.time())))

if __name__ == '__main__':
    train_img_path = os.path.abspath('./SynthText/')
    train_gt_path = os.path.abspath('./SynthText/gt.mat')

    # source domain
    trainset_ = SynthText(train_img_path, train_gt_path)
    train_loader_source = data.DataLoader(trainset_, batch_size=args.batch_size,
                                   shuffle=True, num_workers=args.num_workers, drop_last=True)

    # target domain
    trainset = ICDAR15(args.target_image, args.target_pseudo_positive,args.target_pseudo_negative)
    train_loader_target = data.DataLoader(trainset, batch_size=args.batch_size,
                                          shuffle=True, num_workers=args.num_workers, drop_last=True)

    criterion = Loss()
    criterion_target = Loss_target()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = EAST()
    data_parallel = False
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
        data_parallel = True
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[args.epoch_iter // 2], gamma=0.1)

    # generating pseudo-label
    print("loading pretrained model from ",args.resume)
    model.load_state_dict(torch.load(args.resume))
    generate_pseduo(model, args.target_image, args.target_pseudo_positive, args.target_pseudo_negative, device)

    for epoch in range(args.epoch_iter):

        train( epoch, model, optimizer, train_loader_source,train_loader_target)
        # 进行target domain的eval，看看指标。


# Template: https://www.pygame.org/docs/tut/tom_games2.html
import pygame,math
import numpy as np
import sys,os
from labcommon import Gondola,ButtonBox # @UnresolvedImport
from version import SIMULATOR_VERSION
from _operator import pos

INFOA = "LITEC Lab 6 Simulator"
INFOB = INFOA+" Version "+SIMULATOR_VERSION

SIMSTEP = 0.01 # seconds between updates

class Simulation():
    def __init__(self,controlmodel,runctl,asset_path=None):
        self.ctlmod = controlmodel
        self.runctl = runctl
        # Initialize screen
        pygame.init()
        pygame.font.init()
        self.clock = pygame.time.Clock()
        self.screen = pygame.display.set_mode((1000,800))
        pygame.display.set_caption(INFOA)

        # Set background
        self.background = pygame.Surface(self.screen.get_size()).convert()
        self.background.fill((225,225,225))

        # Get file locations
        if asset_path is None:
            try:
                base_path = sys._MEIPASS #@UndefinedVariable
            except Exception:
                base_path = os.path.abspath(".")
            asset_path = base_path+"/assets/"

        pygame.display.set_icon(pygame.image.load(asset_path+"icon.png"))

        self.gond = Gondola(asset_path,center=(600,400),simstep=SIMSTEP,sym=False)

        self.compassrose = pygame.image.load(asset_path+"compassrose.png").convert_alpha()
        self.compassrose_rect = self.compassrose.get_rect()
        self.compassrose_rect.center = self.gond.rect_gond.center

        self.target_pnts_orig = np.array([[0,0],[-352,-375]])
        self.target_pnts = self.target_pnts_orig
        self.heading_pnts_orig = np.array([[0,0],[-352,-352+23]])
        self.heading_pnts = self.heading_pnts_orig

        self.font = pygame.font.SysFont('Serif', 14)
        self.info = self.font.render(INFOB,True,(0,0,0))
        self.info_rect = self.info.get_rect()
        self.info_rect.bottomleft = (5,800)

        # Control boxes
        self.info_font = pygame.font.SysFont('Serif', 20)
        self.boxes = []
        reset_box = ButtonBox()
        reset_box.text = self.info_font.render('RESTART',True,(0,0,0))
        reset_box.rect = pygame.Rect(5,5,250,40)
        reset_box.func = self.reset
        self.boxes.append(reset_box)
        start_box = ButtonBox()
        start_box.text = self.info_font.render('START ANGLE: {} deg'.format(self.gond.theta_orig/10),True,(0,0,0))
        start_box.rect = pygame.Rect(5,5+40+5,250,40)
        start_box.func = self.setangle
        self.boxes.append(start_box)
        spin_label = ButtonBox()
        spin_label.text = self.info_font.render('FORCE SPIN:',True,(0,0,0))
        spin_label.rect = pygame.Rect(5,5+40+5+40+5,205,30)
        spin_label.colordown = spin_label.color
        self.boxes.append(spin_label)
        spin_box = ButtonBox()
        spin_box.rect = pygame.Rect(5,5+40+5+40+5+30+5,100,100)
        spin_box.ico = pygame.transform.flip(pygame.image.load(asset_path+"Gondola/gondola_spin_ico.png").convert_alpha(),1,0)
        spin_box.ico_offset = np.array([0,-5])
        spin_box.func = self.spinleft
        self.boxes.append(spin_box)
        spin_box = ButtonBox()
        spin_box.rect = pygame.Rect(110,5+40+5+40+5+30+5,100,100)
        spin_box.ico = pygame.image.load(asset_path+"Gondola/gondola_spin_ico.png").convert_alpha()
        spin_box.ico_offset = np.array([0,-5])
        spin_box.func = self.spinright
        self.boxes.append(spin_box)

        self.reqblock = pygame.Surface((200,200)).convert()
        self.reqblock.fill((225,225,225))
        self.reqblock_rect = self.reqblock.get_rect()
        self.reqblock_rect.topleft = (5,500)
        self.reqblock.blit(self.info_font.render('Use These Fans:',True,(0,0,0)),(5,5))

        self.cfgdone = False

        self.des_heading = None

        self.reset()

    def reset(self,box=None):
        self.gond.reset()

    def setconfig(self,cfg):
        self.cfgdone = True
        # Things to generate:
        # a    Fans to use
        # b    Inertia of the gondola
        # c    Offset moment
        # d    desired heading
        #
        # x x x x x x x x
        # a a a b b b c c

        offset = 5+25

        self.reqblock = pygame.Surface((200,200)).convert()
        self.reqblock.fill((225,225,225))
        self.reqblock_rect = self.reqblock.get_rect()
        self.reqblock_rect.topleft = (5,500)
        self.reqblock.blit(self.info_font.render('Use These Fans:',True,(0,0,0)),(5,5))

        # get fans
        if cfg & 0xE0:
            if cfg & 0x80:
                self.reqblock.blit(self.info_font.render('RIGHT',True,(127,0,0)),(20,offset))
                offset += 25
            if cfg & 0x40:
                self.reqblock.blit(self.info_font.render('LEFT',True,(127,0,0)),(20,offset))
                offset += 25
            if cfg & 0x20:
                self.reqblock.blit(self.info_font.render('TAIL',True,(127,0,0)),(20,offset))
                offset += 25
        else:
            self.reqblock.blit(self.info_font.render('TAIL',True,(127,0,0)),(20,offset))
            offset += 25

        # get inertia
        scale = ((cfg >> 2) & 0x03) * (1-2*((cfg & 0x10) > 0))
        self.gond.I = 1+scale/10

        # get moment sign
        self.gond.Mo *= (1-2*((cfg & 0x01) > 0))

        # Desired heading
        tmp = ((cfg >> 4) ^ (cfg & 0x0F)) & 0x07
        self.des_heading = 45*tmp
        self.reqblock.blit(self.info_font.render('Use Desired Heading:',True,(0,0,0)),(5,offset))
        offset += 25
        self.reqblock.blit(self.info_font.render('{}\u00b0'.format(self.des_heading),True,(0,0,0)),(20,offset))

        anglesin = np.sin(np.radians(-self.des_heading))
        anglecos = np.cos(np.radians(-self.des_heading))
        txmatrix = np.array([[anglecos,anglesin],[-anglesin,anglecos]])
        self.target_pnts = txmatrix.dot(self.target_pnts_orig)

    def setangle(self,box):
        self.gond.theta_orig += 450
        self.gond.theta_orig %= 3600
        box.text = self.info_font.render('START ANGLE: {} deg'.format(self.gond.theta_orig/10),True,(0,0,0))

    def spinleft(self,box=None):
        self.gond.theta_t = -5000

    def spinright(self,box=None):
        self.gond.theta_t = 5000

    def update(self):
        self.ctlmod.timestep(SIMSTEP)

        for fan in self.gond.fans:
            fan.setdc(self.ctlmod.xbr.getpin(fan.port,fan.pin,'CCM'),self.ctlmod.pca0.Tperiod)

        self.gond.update()

        # Emit compass reading
        self.ctlmod.compass.setdirection(np.radians(self.gond.theta/10))

    def draw(self):
        self.screen.blit(self.background,(0,0))
        self.gond.draw(self.screen)
        for box in self.boxes:
            box.draw(self.screen)

        self.screen.blit(self.compassrose,self.compassrose_rect)

        pygame.draw.rect(self.screen,(245,245,245),self.info_rect.inflate(20,5))
        self.screen.blit(self.info,self.info_rect)

        font = pygame.font.SysFont('Serif',20)
        if self.des_heading is not None:
            error = self.des_heading - self.gond.theta/10
            if abs(error) <= 5:
                err_color = (0,150,0)
            else:
                err_color = (200,0,0)
            self.screen.blit(font.render("Error:",True,err_color),(5,700))
            txt = font.render("{:5.1f}\u00b0".format(error),True,err_color)
            self.screen.blit(txt,(175-txt.get_width(),700))
        self.screen.blit(font.render("Heading:",True,(0,0,0)),(5,725))
        txt = font.render("{:5.1f}\u00b0".format(self.gond.theta/10),True,(0,0,0))
        self.screen.blit(txt,(175-txt.get_width(),725))
        self.screen.blit(font.render("Velocity:",True,(0,0,0)),(5,750))
        txt = font.render("{:6.1f}\u00b0/s".format(self.gond.theta_t/10),True,(0,0,0))
        self.screen.blit(txt,(175-txt.get_width(),750))

        pygame.draw.line(self.screen,(0,100,255),self.target_pnts[:,0]+np.array(self.gond.rect_gond.center),self.target_pnts[:,1]+np.array(self.gond.rect_gond.center),width=5)
        anglesin = np.sin(np.radians(-self.gond.theta/10))
        anglecos = np.cos(np.radians(-self.gond.theta/10))
        txmatrix = np.array([[anglecos,anglesin],[-anglesin,anglecos]])
        self.heading_pnts = txmatrix.dot(self.heading_pnts_orig)
        pygame.draw.line(self.screen,(255,0,0),self.heading_pnts[:,0]+np.array(self.gond.rect_gond.center),self.heading_pnts[:,1]+np.array(self.gond.rect_gond.center),width=5)

        if self.reqblock is not None:
            self.screen.blit(self.reqblock,self.reqblock_rect)

        if not self.cfgdone:
            font = pygame.font.SysFont('Serif', 30,bold=True)
            endrect = pygame.Rect((0,0),(400,100))
            endrect.center = (600,400)
            endtext1 = font.render('RIN NOT PROVIDED',True,(0,0,0))
            endtext1_rect = endtext1.get_rect(center=endrect.center)
            endtext1_rect.bottom = endrect.centery-3
            endtext2 = font.render('#define RIN xxxxxxxxx',True,(255,0,0))
            endtext2_rect = endtext2.get_rect(center=endrect.center)
            endtext2_rect.top = endrect.centery+3
            pygame.draw.rect(self.screen,(0,0,0),endrect,width=5)
            pygame.draw.rect(self.screen,(255,255,255),endrect)
            self.screen.blit(endtext1,endtext1_rect)
            self.screen.blit(endtext2,endtext2_rect)

    def run(self):
        while self.runctl > 0:
            if self.runctl == 2:
                self.runctl.run = 1
                self.reset()
            if self.runctl > 2:
                self.setconfig(self.runctl.run)
                self.runctl.run = 1
                self.reset()

            self.update()

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    return
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    pos = pygame.mouse.get_pos()
                    for box in self.boxes:
                        if box.rect.collidepoint(pos):
                            box.hit = True
                elif event.type == pygame.MOUSEBUTTONUP:
                    pos = pygame.mouse.get_pos()
                    for box in self.boxes:
                        if box.hit:
                            if box.rect.collidepoint(pos):
                                box.func(box)
                        box.hit = False

            self.draw()
            pygame.display.flip()

            self.clock.tick(1/SIMSTEP)

if __name__ == "__main__":
    sim = Simulation(0,1,'../assets/')
    sim.run()


"""Molecular database (MDB) class

   * MdbExomol is the MDB for ExoMol
   * MdbHit is the MDB for HITRAN or HITEMP

"""
import numpy as np
import jax.numpy as jnp
import pathlib
from exojax.spec import hapi, exomolapi, exomol
from exojax.spec.hitran import gamma_natural as gn
import pandas as pd

__all__ = ['MdbExomol','MdbHit']

class MdbExomol(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        jlower (jnp array): J_lower
        jupper (jnp array): J_upper
        n_Tref (jnp array): temperature exponent
        alpha_ref (jnp array): alpha_ref (gamma0)
        n_Tref_def: default temperature exponent in .def file, used for jlower not given in .broad
        alpha_ref_def: default alpha_ref (gamma0) in .def file, used for jlower not given in .broad

    """
    def __init__(self,path,nurange=[-np.inf,np.inf],margin=1.0,crit=-np.inf, bkgdatm="H2", broadf=True):
        """Molecular database for Exomol form

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction
           bkgdatm: background atmosphere for broadening. e.g. H2, He,
           broadf: if False, the default broadening parameters in .def file is used

        Note:
           The trans/states files can be very large. For the first time to read it, we convert it to the feather-format. After the second-time, we use the feather format instead.

        """
        explanation="Note: Couldn't find the feather format. We convert data to the feather format. After the second time, it will become much faster."

        self.path = pathlib.Path(path)
        t0=self.path.parents[0].stem
        molec=t0+"__"+str(self.path.stem)
        self.bkgdatm=bkgdatm
        print("Background atmosphere: ",self.bkgdatm)
        molecbroad=t0+"__"+self.bkgdatm

        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]
        self.broadf=broadf
        #Where exomol files are
        self.states_file = self.path/pathlib.Path(molec+".states.bz2")
        self.pf_file = self.path/pathlib.Path(molec+".pf")
        self.def_file = self.path/pathlib.Path(molec+".def")
        self.broad_file = self.path/pathlib.Path(molecbroad+".broad")

        if not self.def_file.exists():
                self.download(molec,extension=[".def"])
        if not self.pf_file.exists():
                self.download(molec,extension=[".pf"])
        if not self.states_file.exists():
                self.download(molec,extension=[".states.bz2"])
        if not self.broad_file.exists():
                self.download(molec,extension=[".broad"])

        #load def
        self.n_Texp_def, self.alpha_ref_def, self.molmass, numinf, numtag=exomolapi.read_def(self.def_file)
        #  default n_Texp value if not given
        if self.n_Texp_def is None:
            self.n_Texp_def=0.5
        #  default alpha_ref value if not given
        if self.alpha_ref_def is None:
            self.alpha_ref_def=0.07

        #load states
        if self.states_file.with_suffix(".feather").exists():
            states=pd.read_feather(self.states_file.with_suffix(".feather"))
        else:
            print(explanation)
            states=exomolapi.read_states(self.states_file)
            states.to_feather(self.states_file.with_suffix(".feather"))
        #load pf
        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

        #trans file(s)
        print("Reading transition file")
        if numinf is None:
            self.trans_file = self.path/pathlib.Path(molec+".trans.bz2")
            if not self.trans_file.exists():
                self.download(molec,[".trans.bz2"])

            if self.trans_file.with_suffix(".feather").exists():
                trans=pd.read_feather(self.trans_file.with_suffix(".feather"))
                ndtrans=trans.to_numpy()
                del trans
            else:
                print(explanation)
                trans=exomolapi.read_trans(self.trans_file)
                trans.to_feather(self.trans_file.with_suffix(".feather"))
                ndtrans=trans.to_numpy()
                del trans
            #compute gup and elower
            self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper=exomolapi.pickup_gE(states,ndtrans)
        else:
            imin=np.searchsorted(numinf,self.nurange[0],side="right")-1 #left side
            imax=np.searchsorted(numinf,self.nurange[1],side="right")-1 #left side
            self.trans_file=[]
            for k,i in enumerate(range(imin,imax+1)):
                trans_file = self.path/pathlib.Path(molec+"__"+numtag[i]+".trans.bz2")
                if not trans_file.exists():
                    self.download(molec,extension=[".trans.bz2"],numtag=numtag[i])
                if trans_file.with_suffix(".feather").exists():
                    trans=pd.read_feather(trans_file.with_suffix(".feather"))
                    ndtrans=trans.to_numpy()
                    del trans
                else:
                    print(explanation)
                    trans=exomolapi.read_trans(trans_file)
                    trans.to_feather(trans_file.with_suffix(".feather"))
                    ndtrans=trans.to_numpy()
                    del trans
                self.trans_file.append(trans_file)
                #compute gup and elower
                if k==0:
                    self._A, self.nu_lines, self._elower, self._gpp, self._jlower, self._jupper=exomolapi.pickup_gE(states,ndtrans)
                else:
                    Ax, nulx, elowerx, gppx, jlowerx, jupperx=exomolapi.pickup_gE(states,ndtrans)
                    self._A=np.hstack([self._A,Ax])
                    self.nu_lines=np.hstack([self.nu_lines,nulx])
                    self._elower=np.hstack([self._elower,elowerx])
                    self._gpp=np.hstack([self._gpp,gppx])
                    self._jlower=np.hstack([self._jlower,jlowerx])
                    self._jupper=np.hstack([self._jupper,jupperx])

        self.Tref=296.0
        self.QTref=np.array(self.QT_interp(self.Tref))

        ##Line strength: input should be ndarray not jnp array
        self.Sij0=exomol.Sij0(self._A,self._gpp,self.nu_lines,self._elower,self.QTref)

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied. self.mask is updated.

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """
        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self._A=self._A[mask]
        self._elower=self._elower[mask]
        self._gpp=self._gpp[mask]
        self._jlower=self._jlower[mask]
        self._jupper=self._jupper[mask]

        #jnp arrays
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.gamma_natural=gn(self.A)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.jlower=jnp.array(self._jlower,dtype=int)
        self.jupper=jnp.array(self._jupper,dtype=int)
        ##Broadening parameters
        self.set_broadening()

    def set_broadening(self,alpha_ref_def=None,n_Texp_def=None):
        """setting broadening parameters

        Args:
           alpha_ref: set default alpha_ref and apply it. None=use self.alpha_ref_def
           n_Texp_def: set default n_Texp and apply it. None=use self.n_Texp_def
        """
        if alpha_ref_def:
            self.alpha_ref_def = alpha_ref_def
        if n_Texp_def:
            self.n_Texp_def = n_Texp_def

        if self.broadf:
            try:
                print(".broad is used.")
                bdat=exomolapi.read_broad(self.broad_file)
                codelv=exomolapi.check_bdat(bdat)
                print("Broadening code level=",codelv)
                if codelv=="a0":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    self.alpha_ref=jnp.array(j2alpha_ref[self._jlower])
                    self.n_Texp=jnp.array(j2n_Texp[self._jlower])
                elif codelv=="a1":
                    j2alpha_ref, j2n_Texp = exomolapi.make_j2b(bdat,\
                        alpha_ref_default=self.alpha_ref_def,\
                        n_Texp_default=self.n_Texp_def,\
                        jlower_max=np.max(self._jlower))
                    jj2alpha_ref, jj2n_Texp=exomolapi.make_jj2b(bdat,\
                        j2alpha_ref_def=j2alpha_ref,j2n_Texp_def=j2n_Texp,\
                        jupper_max=np.max(self._jupper))
                    self.alpha_ref=jnp.array(jj2alpha_ref[self._jlower,self._jupper])
                    self.n_Texp=jnp.array(jj2n_Texp[self._jlower,self._jupper])
            except:
                print("Warning: Cannot load .broad. The default broadening parameters are used.")
                self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
                self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

        else:
            print("The default broadening parameters are used.")
            self.alpha_ref=jnp.array(self.alpha_ref_def*np.ones_like(self._jlower))
            self.n_Texp=jnp.array(self.n_Texp_def*np.ones_like(self._jlower))

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def download(self,molec,extension,numtag=None):
        """Downloading Exomol files

        Args:
           molec: like "12C-16O__Li2015"
           extension: extension list e.g. [".pf",".def",".trans.bz2",".states.bz2",".broad"]
           numtag: number tag of transition file if exists. e.g. "11100-11200"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])

        for ext in extension:
            if ext==".trans.bz2" and numtag is not None:
                ext="__"+numtag+ext

            if ext==".broad":
                pfname_arr=[tag[0]+"__H2"+ext,tag[0]+"__He"+ext,tag[0]+"__air"+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"
            else:
                pfname_arr=[molec+ext]
                url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

            for pfname in pfname_arr:
                pfpath=url+pfname
                os.makedirs(str(self.path), exist_ok=True)
                print("Downloading "+pfpath)
                try:
                    urllib.request.urlretrieve(pfpath,str(self.path/pfname))
                except:
                    print("Error: Couldn't download "+ext+" file and save.")

class MdbHit(object):
    """ molecular database of ExoMol

    MdbExomol is a class for ExoMol.

    Attributes:
        nurange: nu range [min,max] (cm-1)
        nu_lines (nd array): line center (cm-1)
        Sij0 (nd array): line strength at T=Tref (cm)
        dev_nu_lines (jnp array): line center in device (cm-1)
        logsij0 (jnp array): log line strength at T=Tref
        A (jnp array): Einstein A coeeficient
        gamma_natural (jnp array): gamma factor of the natural broadening
        gamma_air (jnp array): gamma factor of air pressure broadening
        gamma_self (jnp array): gamma factor of self pressure broadening
        elower (jnp array): the lower state energy (cm-1)
        gpp (jnp array): statistical weight
        n_air (jnp array): air temperature exponent

    """

    def __init__(self,path,nurange=[-np.inf,np.inf],margin=250.0,crit=-np.inf):
        """Molecular database for HITRAN/HITEMP form

        Args:
           path: path for HITRAN/HITEMP par file
           nurange: wavenumber range list (cm-1) [min,max] or wavenumber grid
           margin: margin for nurange (cm-1)
           crit: line strength lower limit for extraction

        """
        #downloading
        self.path = pathlib.Path(path)
        molec=str(self.path.stem)
        if not self.path.exists():
            self.download()

        #bunzip2 if suffix is .bz2
        if self.path.suffix==".bz2":
            import bz2,shutil
            if self.path.with_suffix('').exists():
                import os
                os.remove(self.path.with_suffix(''))
            print("bunziping")
            with bz2.BZ2File(str(self.path)) as fr:
                with open(str(self.path.with_suffix('')),"wb") as fw:
                    shutil.copyfileobj(fr,fw)
            self.path=self.path.with_suffix('')

        molec=str(self.path.stem)

        hapi.db_begin(str(self.path.parent))
        self.Tref=296.0
        self.molecid = search_molecid(molec)
        self.crit = crit
        self.margin = margin
        self.nurange=[np.min(nurange),np.max(nurange)]

        #nd arrays using DRAM (not jnp, not in GPU)
        self.nu_lines = hapi.getColumn(molec, 'nu')
        self.Sij0 = hapi.getColumn(molec, 'sw')
        self.delta_air = hapi.getColumn(molec, 'delta_air')
        self.isoid = hapi.getColumn(molec,'local_iso_id')
        self.uniqiso=np.unique(self.isoid)

        self._A=hapi.getColumn(molec, 'a')
        self._n_air = hapi.getColumn(molec, 'n_air')
        self._gamma_air = hapi.getColumn(molec, 'gamma_air')
        self._gamma_self =hapi.getColumn(molec, 'gamma_self')
        self._elower = hapi.getColumn(molec, 'elower')
        self._gpp = hapi.getColumn(molec, 'gpp')

        ### MASKING ###
        mask=(self.nu_lines>self.nurange[0]-self.margin)\
        *(self.nu_lines<self.nurange[1]+self.margin)\
        *(self.Sij0>self.crit)

        self.masking(mask)

    def masking(self,mask):
        """applying mask and (re)generate jnp.arrays

        Args:
           mask: mask to be applied

        Note:
           We have nd arrays and jnp arrays. We apply the mask to nd arrays and generate jnp array from the corresponding nd array. For instance, self._A is nd array and self.A is jnp array.

        """

        #numpy float 64 Do not convert them jnp array
        self.nu_lines = self.nu_lines[mask]
        self.Sij0 = self.Sij0[mask]
        self.delta_air=self.delta_air[mask]
        self.isoid = self.isoid[mask]
        self.uniqiso=np.unique(self.isoid)

        ##numpy float 64 copy source for jnp
        self._A=self._A[mask]
        self._n_air = self._n_air[mask]
        self._gamma_air = self._gamma_air[mask]
        self._gamma_self = self._gamma_self[mask]
        self._elower = self._elower[mask]
        self._gpp = self._gpp[mask]

        #jnp.array copy from the copy sources
        self.dev_nu_lines=jnp.array(self.nu_lines)
        self.logsij0=jnp.array(np.log(self.Sij0))
        self.A=jnp.array(self._A)
        self.n_air=jnp.array(self._n_air)
        self.gamma_air = jnp.array(self._gamma_air)
        self.gamma_self = jnp.array(self._gamma_self)
        self.elower=jnp.array(self._elower)
        self.gpp=jnp.array(self._gpp)
        self.gamma_natural=gn(self.A)

    def download(self):
        """Downloading HITRAN/HITEMP par file

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.url import url_HITRAN12
        from exojax.utils.url import url_HITEMP

        try:
            url = url_HITRAN12()+self.path.name
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print(url)
            print("HITRAN download failed")
        try:
            url = url_HITEMP()+self.path.name
            print(url)
            urllib.request.urlretrieve(url,str(self.path))
        except:
            print("HITEMP download failed")

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

    def ExomolQT(self,path):
        """use a partition function from ExoMol

        Args:
           path: path for Exomol data directory/tag. For instance, "/home/CO/12C-16O/Li2015"

        """
        #load pf

        self.empath = pathlib.Path(path)
        t0=self.empath.parents[0].stem
        molec=t0+"__"+str(self.empath.stem)
        self.pf_file = self.empath/pathlib.Path(molec+".pf")
        if not self.pf_file.exists():
                self.exomol_pf_download(molec)

        pf=exomolapi.read_pf(self.pf_file)
        self.gQT=jnp.array(pf["QT"].to_numpy()) #grid QT
        self.T_gQT=jnp.array(pf["T"].to_numpy()) #T forgrid QT

    def exomol_pf_download(self,molec):
        """Downloading Exomol pf files

        Args:
           molec: like "12C-16O__Li2015"

        Note:
           The download URL is written in exojax.utils.url.

        """
        import urllib.request
        from exojax.utils.molname import e2s
        import os
        from exojax.utils.url import url_ExoMol

        tag=molec.split("__")
        molname_simple=e2s(tag[0])
        url = url_ExoMol()+molname_simple+"/"+tag[0]+"/"+tag[1]+"/"

        ext=".pf"
        pfname=molec+ext
        pfpath=url+pfname
        os.makedirs(str(self.empath), exist_ok=True)
        print("Downloading "+pfpath)
        try:
            urllib.request.urlretrieve(pfpath,str(self.empath/pfname))
        except:
            print("Error: Couldn't download "+ext+" file and save.")

    def QT_interp(self,T):
        """interpolated partition function

        Args:
           T: temperature

        Returns:
           Q(T) interpolated in jnp.array

        """
        return jnp.interp(T,self.T_gQT,self.gQT)

    def qr_interp(self,T):
        """interpolated partition function ratio

        Args:
           T: temperature

        Returns:
           qr(T)=Q(T)/Q(Tref) interpolated in jnp.array

        """
        return self.QT_interp(T)/self.QT_interp(self.Tref)

    def Qr_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr = partition function ratio array [N_Tarr x N_iso]

        Note:
           N_Tarr = len(Tarr), N_iso = len(self.uniqiso)

        """
        allT=list(np.concatenate([[self.Tref],Tarr]))
        Qrx=[]
        for iso in self.uniqiso:
            Qrx.append(hapi.partitionSum(self.molecid,iso, allT))
        Qrx=np.array(Qrx)
        qr=Qrx[:,1:].T/Qrx[:,0] #Q(T)/Q(Tref)
        return qr

    def Qr_line_HAPI(self,T):
        """Partition Function ratio using HAPI partition sum

        Args:
           T: temperature (K)

        Returns:
           Qr_line, partition function ratio array for lines [Nlines]

        Note:
           Nlines=len(self.nu_lines)

        """
        qr_line=np.ones_like(self.isoid,dtype=np.float64)
        qrx=self.Qr_HAPI([T])
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            qr_line[mask]=qrx[0,idx]
        return qr_line

    def Qr_layer_HAPI(self,Tarr):
        """Partition Function ratio using HAPI partition sum

        Args:
           Tarr: temperature array (K)

        Returns:
           Qr_layer, partition function ratio array for lines [N_Tarr x Nlines]

        Note:
           Nlines=len(self.nu_lines)
           N_Tarr=len(Tarr)

        """
        NP=len(Tarr)
        qt=np.zeros((NP,len(self.isoid)))
        qr=self.Qr_HAPI(Tarr)
        for idx,iso in enumerate(self.uniqiso):
            mask=self.isoid==iso
            for ilayer in range(NP):
                qt[ilayer,mask]=qr[ilayer,idx]
        return qt

def search_molecid(molec):
    """molec id from molec (source table name) of HITRAN/HITEMP

    Args:
       molec: source table name

    Return:
       int: molecid (HITRAN molecular id)

    """
    try:
        hitf=molec.split("_")
        molecid=int(hitf[0])
        return molecid

    except:
        print("Warning: Define molecid by yourself.")
        return None

if __name__ == "__main__":
    #mdb=MdbExomol("/home/kawahara/exojax/data/CO/12C-16O/Li2015/")
    #mdb=MdbExomol("/home/kawahara/exojax/data/CH4/12C-1H4/YT34to10/",nurange=[6050.0,6150.0])
    mdb=MdbExomol('.database/H2O/1H2-16O/POKAZATEL',[4310.0,4320.0],crit=1.e-45)

#    mask=mdb.A>1.e-42
#    mdb.masking(mask)
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NH3/14N-1H3/CoYuTe/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/H2S/1H2-32S/AYT2/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/FeH/56Fe-1H/MoLLIST/",nurange=[6050.0,6150.0])
#    mdb=MdbExomol("/home/kawahara/exojax/data/exomol/NO/14N-16O/NOname/14N-16O__NOname")


"""
Copyright (c) 2020, salesforce.com, inc.
All rights reserved.
SPDX-License-Identifier: BSD-3-Clause
For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import json
import pathlib

import numpy as np

np.random.seed(0)
import seaborn as sns

sns.set()
import pickle

from protein_attention.utils import get_reports_path, get_cache_path
from matplotlib.ticker import FuncFormatter

import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
from mpl_toolkits.axes_grid1.colorbar import colorbar

secondary_names = {
    '0': 'Helix',
    '1': 'Beta',
    '2': 'Turn/Bend',
    '3': 'Blank'
}

def to_filename(s, extension):
    return "".join(x if (x.isalnum() or x in "._-") else '_' for x in s) + "." + extension

def create_figure(feature_name, weighted_sum, weight_total, report_dir, min_total, filetype):
    assert filetype in ('png', 'pdf')

    mean_by_head = weighted_sum / weight_total
    exclude_mask = np.array(weight_total) < min_total

    masked_mean_by_head = np.ma.masked_array(mean_by_head, mask=exclude_mask)
    layer_max = masked_mean_by_head.max(-1)

    n_layers, n_heads = mean_by_head.shape
    if n_layers == 12 and n_heads == 12:
        plt.figure(figsize=(3, 2.2))
        ax1 = plt.subplot2grid((100, 85), (0, 0), colspan=65, rowspan=99)  # Heatmap
        ax2 = plt.subplot2grid((100, 85), (12, 70), colspan=15, rowspan=75)  # Barchart
    elif n_layers == 30 and n_heads == 16:
        plt.figure(figsize=(3, 2.2))
        ax1 = plt.subplot2grid((100, 85), (0, 5), colspan=55, rowspan=96)
        ax2 = plt.subplot2grid((100, 85), (0, 62), colspan=17, rowspan=97)
    elif n_layers == 12 and n_heads == 64:
        plt.figure(figsize=(8.5, 2.2))
        ax1 = plt.subplot2grid((100, 160), (0, 5), colspan=135, rowspan=96)
        ax2 = plt.subplot2grid((100, 160), (22, 144), colspan=10, rowspan=53)
        plt.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95, wspace=0.01, hspace=0.01)
    else:
        raise NotImplementedError

    xtick_labels = [str(i) if i % 2 == 0 else '' for i in range(1, n_heads + 1)]
    ytick_labels = [str(i) if i % 2 == 0 else '' for i in range(1, n_layers + 1)]
    heatmap = sns.heatmap((mean_by_head * 100).tolist(), center=0.0, ax=ax1,
                          square=True, cbar=False, linewidth=0.1, linecolor='#D0D0D0',
                          cmap=LinearSegmentedColormap.from_list('rg', ["#F14100", "white", "#3D4FC4"], N=256),
                          mask=exclude_mask,
                          xticklabels=xtick_labels,
                          yticklabels=ytick_labels)
    for _, spine in heatmap.spines.items():
        spine.set_visible(True)
        spine.set_edgecolor('#D0D0D0')
        spine.set_linewidth(0.1)
    plt.setp(heatmap.get_yticklabels(), fontsize=7)
    plt.setp(heatmap.get_xticklabels(), fontsize=7)
    heatmap.tick_params(axis='x', pad=1, length=2)
    heatmap.tick_params(axis='y', pad=.5, length=2)
    heatmap.yaxis.labelpad = 3
    heatmap.invert_yaxis()
    heatmap.set_facecolor('#E7E6E6')
    # split axes of heatmap to put colorbar
    ax_divider = make_axes_locatable(ax1)
    if n_layers == 12 and n_heads == 12:
        cax = ax_divider.append_axes('left', size='7%', pad='33%')
    elif n_layers == 30 and n_heads == 16:
        cax = ax_divider.append_axes('left', size='7%', pad='45%')
    elif n_layers == 12 and n_heads == 64:
        cax = ax_divider.append_axes('left', size='1.5%', pad='7%')
    else:
        raise NotImplementedError
    # # make colorbar for heatmap.
    # # Heatmap returns an axes obj but you need to get a mappable obj (get_children)
    cbar = colorbar(ax1.get_children()[0], cax=cax, orientation='vertical', format='%.0f%%')
    cax.yaxis.set_ticks_position('left')
    cbar.solids.set_edgecolor("face")
    cbar.ax.tick_params(labelsize=7, length=4, pad=2)
    ax1.set_title('% Attention', size=9)
    ax1.set_xlabel('Head', size=8)
    ax1.set_ylabel('Layer', size=8)
    for _, spine in ax1.spines.items():
        spine.set_visible(True)
    ax2.set_title('Max', size=9)
    bp = sns.barplot(x=layer_max * 100, ax=ax2, y=list(range(layer_max.shape[0])), color="#3D4FC4", orient="h",
                     edgecolor="none")
    formatter = FuncFormatter(lambda y, pos: '0' if (y == 0) else "%d%%" % (y))
    ax2.xaxis.set_major_formatter(formatter)
    plt.setp(bp.get_xticklabels(), fontsize=7)
    bp.tick_params(axis='x', pad=1, length=3)
    ax2.invert_yaxis()
    ax2.set_yticklabels([])
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.spines['left'].set_visible(False)
    ax2.xaxis.set_ticks_position('bottom')
    ax2.axvline(0, linewidth=.85, color='black')
    fname = report_dir / to_filename(feature_name, filetype)
    print('Saving', fname)
    plt.savefig(fname, format=filetype)
    plt.close()

if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('exp_name', help='Name of experiment')
    args = parser.parse_args()
    print(args)
    min_total = 100
    filetype = 'pdf'

    cache_path = get_cache_path() / f'{args.exp_name}.pickle'
    report_dir = get_reports_path() / 'attention_analysis' / args.exp_name
    pathlib.Path(report_dir).mkdir(parents=True, exist_ok=True)

    args, feature_to_weighted_sum, weight_total = pickle.load(open(cache_path, "rb"))
    with open(report_dir / 'args.json', 'w') as f:
        json.dump(vars(args), f)
    print(args)
    for feature_name, weighted_sum in feature_to_weighted_sum.items():
        create_figure(feature_name, weighted_sum, weight_total, report_dir, min_total=min_total, filetype=filetype)


from matplotlib import pyplot as plt
import numpy as np
import os
import tensorflow as tf
import time

import model
from inputs.detection.construct import construct_network_input_nodes

def visual_test(tfrecords, bbox_priors, checkpoint_dir, specific_model_path, cfg):

  graph = tf.Graph()

  sess_config = tf.ConfigProto(
    log_device_placement=False,
    #device_filters = device_filters,
    allow_soft_placement = True,
    gpu_options = tf.GPUOptions(
        per_process_gpu_memory_fraction=cfg.SESSION_CONFIG.PER_PROCESS_GPU_MEMORY_FRACTION
    )
  )
  sess = tf.Session(graph=graph, config=sess_config)

  with graph.as_default(), sess.as_default():

    images, batched_bboxes, batched_num_bboxes, paths = construct_network_input_nodes(
      tfrecords=tfrecords,
      max_num_bboxes=cfg.MAX_NUM_BBOXES,
      num_epochs=None,
      batch_size=cfg.BATCH_SIZE,
      num_threads=cfg.NUM_INPUT_THREADS,
      add_summaries = False,
      augment=cfg.AUGMENT_IMAGE,
      shuffle_batch=False,
      cfg=cfg
    )
    features = model.build(graph, images, cfg)

    locations, confidences = model.add_detection_heads(graph, features, num_bboxes_per_cell=5, batch_size=cfg.BATCH_SIZE, cfg=cfg)

    # Restore the moving average variables for the conv filters, beta and gamma for
    # batch normalization and the softmax params
    ema = tf.train.ExponentialMovingAverage(decay=cfg.MOVING_AVERAGE_DECAY)
    shadow_vars = {
      ema.average_name(var) : var
      for var in graph.get_collection('conv_params')
    }
    shadow_vars.update({
      ema.average_name(var) : var
      for var in graph.get_collection('batchnorm_params')
    })
    shadow_vars.update({
      ema.average_name(var) : var
      for var in graph.get_collection('softmax_params')
    })
    shadow_vars.update({
      ema.average_name(var) : var
      for var in graph.get_collection('batchnorm_mean_var')
    })

    # Restore the parameters
    saver = tf.train.Saver(shadow_vars, reshape=True)

    fetches = [locations, confidences, images, batched_bboxes, batched_num_bboxes]

    coord = tf.train.Coordinator()

    tf.initialize_all_variables().run()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    plt.ion()

    try:

      if specific_model_path == None:
        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
          specific_model_path = ckpt.model_checkpoint_path
        else:
          print('No checkpoint file found')
          return

      # Restores from checkpoint
      saver.restore(sess, specific_model_path)
      # Assuming model_checkpoint_path looks something like:
      #   /my-favorite-path/cifar10_train/model.ckpt-0,
      # extract global_step from it.
      global_step = int(specific_model_path.split('/')[-1].split('-')[-1])
      print "Found model for global step: %d" % (global_step,)

      total_sample_count = 0
      step = 0
      done = False
      while not coord.should_stop() and not done:

        t = time.time()
        outputs = sess.run(fetches)
        dt = time.time()-t

        locs = outputs[0]
        confs = outputs[1]
        imgs = outputs[2]
        gt_bboxes = outputs[3]
        gt_num_bboxes = outputs[4]

        print locs.shape
        print confs.shape

        for b in range(cfg.BATCH_SIZE):

          # Show the image
          image = imgs[b]
          plt.imshow((image * cfg.IMAGE_STD + cfg.IMAGE_MEAN).astype(np.uint8))

           # Draw the GT Boxes in blue
          for i in range(gt_num_bboxes[b]):
            gt_bbox = gt_bboxes[b][i]
            xmin, ymin, xmax, ymax = gt_bbox * cfg.INPUT_SIZE
            plt.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin], 'b-')

          # Draw the most confident boxes in red
          indices = np.argsort(confs[b].ravel())[::-1]
          for i, index in enumerate(indices[0:gt_num_bboxes[b]]):

            loc = locs[b][index].ravel()
            conf = confs[b][index]
            prior = np.array(bbox_priors[index])

            xmin, ymin, xmax, ymax = (prior + loc) * cfg.INPUT_SIZE
            plt.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin], 'r-')

            # Draw the prior box that was chosen
            xmin, ymin, xmax, ymax = prior * cfg.INPUT_SIZE
            plt.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin], 'g-')

            print "Pred Confidence for box %d: %f" % (i, conf)

          plt.show()

          t = raw_input("push button")
          if t != '':
            done = True
            break

          plt.clf()

    except tf.errors.OutOfRangeError as e:
      pass

    coord.request_stop()
    coord.join(threads)


# encoding=utf-8

import tensorflow as tf
import numpy as np
import  cv2
from config import cfg
from datasets.pascal_voc import pascal_voc

__sets = {}

for year in ['2007', '2012']:
    for split in ['train', 'val', 'trainval', 'test']:
        name = 'voc_{}_{}'.format(year, split)
        # 不使用标记为difficulty的图像
        __sets[name] = (lambda split=split, year=year: pascal_voc(split, year, use_diff=False))

'''
def image_preprocess(image, bgr=True):
    # 将图像减去均值，再除以标准差
    with tf.name_scope('image_preprocess'):
        if image.dtype.base_dtype != tf.float32:
            image = tf.cast(image, tf.float32)

        mean = cfg.PREPROC_PIXEL_MEAN
        std = np.asarray(cfg.PREPROC_PIXEL_STD)
        if bgr:
            mean = mean[::-1]
            std = std[::-1]
        image_mean = tf.constant(mean, dtype=tf.float32)
        image_invstd = tf.constant(1.0 / std, dtype=tf.float32)
        image = (image - image_mean) * image_invstd
        return image
'''

def get_voc_dataset(dataset_name='voc_2007_trainval'):
    return __sets[dataset_name]()

def im_list_to_blob(ims):
    """Convert a list of images into a network input.

    Assumes images are already prepared (means subtracted, BGR order, ...).
    """
    max_shape = np.array([im.shape for im in ims]).max(axis=0)
    num_images = len(ims)
    blob = np.zeros((num_images, max_shape[0], max_shape[1], 3),
                    dtype=np.float32)
    for i in range(num_images):
        im = ims[i]
        blob[i, 0:im.shape[0], 0:im.shape[1], :] = im

    return blob

def _get_image_blob(im):
    """Converts an image into a network input.
    Arguments:
      im (ndarray): a color image in BGR order
    Returns:
      blob (ndarray): a data blob holding an image pyramid
      im_scale_factors (list): list of image scales (relative to im) used
        in the image pyramid
    """
    im_orig = im.astype(np.float32, copy=True)
    im_orig -= cfg.PIXEL_MEANS

    im_shape = im_orig.shape
    im_size_min = np.min(im_shape[0:2])
    im_size_max = np.max(im_shape[0:2])

    processed_ims = []
    im_scale_factors = []

    target_size = cfg.IMAGE_SCALES
    im_scale = float(target_size) / float(im_size_min)
    # Prevent the biggest axis from being more than MAX_SIZE
    if np.round(im_scale * im_size_max) > cfg.IMAGE_MAX_SIZE:
        im_scale = float(cfg.IMAGE_MAX_SIZE) / float(im_size_max)
    im = cv2.resize(im_orig, None, None, fx=im_scale, fy=im_scale,
                    interpolation=cv2.INTER_LINEAR)
    im_scale_factors.append(im_scale)
    processed_ims.append(im)

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

    return blob, np.array(im_scale_factors)

def _get_blobs(im):
    """Convert an image and RoIs within that image into network inputs."""
    blobs = {}
    blobs['data'], im_scale_factors = _get_image_blob(im)

    return blobs, im_scale_factors


import os
import random
from contextlib import redirect_stdout
from PIL import Image
import cv2
import torch
import torch.nn.functional as F
from torch.utils import data
from pycocotools.coco import COCO
import numpy as np

class CocoDataset(data.dataset.Dataset):
    'Dataset looping through a set of images'

    def __init__(self, path, resize, max_size, stride, annotations=None, training=False):
        super().__init__()

        self.path = os.path.expanduser(path)
        self.resize = resize
        self.max_size = max_size
        self.stride = stride
        self.mean = [102.9801, 115.9465, 122.7717]
        self.std = [1., 1., 1.]
        self.training = training

        with redirect_stdout(None):
            self.coco = COCO(annotations)
        self.ids = list(self.coco.imgs.keys())
        self.ids = sorted(self.ids)
        if 'categories' in self.coco.dataset:
            self.categories_inv = { k: i for i, k in enumerate(self.coco.getCatIds()) }

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

    def __getitem__(self, index):
        ' Get sample'

        # Load image
        id = self.ids[index]
        if self.coco:
            image = self.coco.loadImgs(id)[0]['file_name']
        im = Image.open('{}/{}'.format(self.path, image)).convert("RGB")

        # Randomly sample scale for resize during training
        resize = self.resize
        if isinstance(resize, list):
            resize = random.randint(self.resize[0], self.resize[-1])

        ratio = resize / min(im.size)
        if ratio * max(im.size) > self.max_size:
            ratio = self.max_size / max(im.size)
        im = im.resize((int(ratio * d) for d in im.size), Image.BILINEAR)

        r, g, b = im.split()
        im = Image.merge("RGB", (b, g, r))

        # Convert to tensor and normalize
        data = torch.ByteTensor(torch.ByteStorage.from_buffer(im.tobytes()))
        data = data.float().view(*im.size[::-1], len(im.mode))

        data = data.permute(2, 0, 1)
        for t, mean, std in zip(data, self.mean, self.std):
            t.sub_(mean).div_(std)

        # Apply padding
        pw, ph = ((self.stride - d % self.stride) % self.stride for d in im.size)
        data = F.pad(data, (0, pw, 0, ph))

        """
        # debug
        if id == 139:
            print('id: {}, shape: {}'.format(id, data.shape))
            np.save(os.path.join('/workspace/retinanet/debug', '{}.npy'.format(id)), data.cpu().numpy())
        """

        return data, id, ratio

    def _get_target(self, id):
        'Get annotations for sample'

        ann_ids = self.coco.getAnnIds(imgIds=id)
        annotations = self.coco.loadAnns(ann_ids)

        boxes, categories = [], []
        for ann in annotations:
            if ann['bbox'][2] < 1 and ann['bbox'][3] < 1:
                continue
            boxes.append(ann['bbox'])
            cat = ann['category_id']
            if 'categories' in self.coco.dataset:
                cat = self.categories_inv[cat]
            categories.append(cat)

        if boxes:
            target = (torch.FloatTensor(boxes),
                torch.FloatTensor(categories).unsqueeze(1))
        else:
            target = (torch.ones([1, 4]), torch.ones([1, 1]) * -1)

        return target

    def collate_fn(self, batch):
        'Create batch from multiple samples'

        if self.training:
            data, targets = zip(*batch)
            max_det = max([t.size()[0] for t in targets])
            targets = [torch.cat([t, torch.ones([max_det - t.size()[0], 5]) * -1]) for t in targets]
            targets = torch.stack(targets, 0)
        else:
            data, indices, ratios = zip(*batch)

        # Pad data to match max batch dimensions
        sizes = [d.size()[-2:] for d in data]
        w, h = (max(dim) for dim in zip(*sizes))

        data_stack = []
        for datum in data:
            pw, ph = w - datum.size()[-2], h - datum.size()[-1]
            data_stack.append(
                F.pad(datum, (0, ph, 0, pw)) if max(ph, pw) > 0 else datum)

        data = torch.stack(data_stack)

        if self.training:
            return data, targets

        ratios = torch.FloatTensor(ratios).view(-1, 1, 1)
        return data, torch.IntTensor(indices), ratios

class DataIterator():
    'Data loader for data parallel'

    def __init__(self, path, resize, max_size, batch_size, stride, world, annotations, training=False):
        self.resize = resize
        self.max_size = max_size

        self.dataset = CocoDataset(path, resize=resize, max_size=max_size,
            stride=stride, annotations=annotations, training=training)
        self.ids = self.dataset.ids
        self.coco = self.dataset.coco

        self.sampler = data.distributed.DistributedSampler(self.dataset) if world > 1 else None
        self.dataloader = data.DataLoader(self.dataset, batch_size=batch_size // world,
            sampler=self.sampler, collate_fn=self.dataset.collate_fn, num_workers=2, pin_memory=True)

    def __repr__(self):
        return '\n'.join([
            '    loader: pytorch',
            '    resize: {}, max: {}'.format(self.resize, self.max_size),
        ])

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

    def __iter__(self):
        for output in self.dataloader:
            if self.dataset.training:
                data, target = output
            else:
                data, ids, ratio = output

            if torch.cuda.is_available():
                data = data.cuda(non_blocking=True)

            if self.dataset.training:
                if torch.cuda.is_available():
                    target = target.cuda(non_blocking=True)
                yield data, target
            else:
                if torch.cuda.is_available():
                    ids = ids.cuda(non_blocking=True)
                    ratio = ratio.cuda(non_blocking=True)
                yield data, ids, ratio


"""
Shared constants and delayed imports.

This module imports shared libraries either directly or with `LazyImport`.

`LazyImport` shall only be used to imported
"""
# Known issues of LazyImport ::
#
#     1) High overhead when called hundreds of thousands times.
#     For example, NumPy must not be imported with LazyImport.

import math
import os
import coloredlogs         # NOQA
import numpy as np         # NOQA

from andes.utils.lazyimport import LazyImport
from distutils.spawn import find_executable

# Library preference:
# KVXOPT + ipadd > CVXOPT + ipadd > KXVOPT > CVXOPT (+ KLU or not)

try:
    import kvxopt
    from kvxopt import umfpack   # test if shared libs can be found
    from kvxopt import spmatrix as kspmatrix
    KIP_ADD = True
except ImportError:
    kvxopt = None
    kspmatrix = None
    KIP_ADD = False

if KIP_ADD is False:
    from cvxopt import spmatrix as cspmatrix
    if hasattr(cspmatrix, 'ipadd'):
        CIP_ADD = True
    else:
        CIP_ADD = False

if kvxopt is None or (KIP_ADD is False and CIP_ADD is True):
    from cvxopt import umfpack                           # NOQA
    from cvxopt import spmatrix, matrix, sparse, spdiag  # NOQA
    from cvxopt import mul, div                          # NOQA
    from cvxopt.lapack import gesv                       # NOQA
    from cvxopt import printing                          # NOQA
    try:
        from cvxoptklu import klu  # NOQA
    except ImportError:
        klu = None
    IP_ADD = CIP_ADD
else:
    from kvxopt import umfpack, klu                      # NOQA
    from kvxopt import spmatrix, matrix, sparse, spdiag  # NOQA
    from kvxopt import mul, div                          # NOQA
    from kvxopt.lapack import gesv                       # NOQA
    from kvxopt import printing                          # NOQA
    IP_ADD = KIP_ADD

printing.options['dformat'] = '%.1f'
printing.options['width'] = -1

from andes.utils.texttable import Texttable              # NOQA
from andes.utils.paths import get_dot_andes_path         # NOQA

# --- constants ---

deg2rad = math.pi/180
jac_names = ('fx', 'fy', 'gx', 'gy')
jac_types = ('c', '')

jac_full_names = list()
for jname in jac_names:
    for jtype in jac_types:
        jac_full_names.append(jname + jtype)

# --- lazy import packages ---
tqdm = LazyImport('from tqdm import tqdm')

pd = LazyImport('import pandas')
cupy = LazyImport('import cupy')
mpl = LazyImport('import matplotlib')
unittest = LazyImport('import unittest')
yaml = LazyImport('import yaml')

plt = LazyImport('from matplotlib import pyplot')
Pool = LazyImport('from pathos.multiprocessing import Pool')
Process = LazyImport('from multiprocess import Process')

newton_krylov = LazyImport('from scipy.optimize import newton_krylov')
fsolve = LazyImport('from scipy.optimize import fsolve')
solve_ivp = LazyImport('from scipy.integrate import solve_ivp')
odeint = LazyImport('from scipy.integrate import odeint')

# --- Shared functions ---

def set_latex():
    """
    Enables LaTeX for matplotlib based on the `with_latex` option and `dvipng` availability.

    Returns
    -------
    bool
        True for LaTeX on, False for off
    """

    if find_executable('dvipng'):
        mpl.rc('text', usetex=True)

        no_warn_file = os.path.join(get_dot_andes_path(), '.no_warn_latex')
        if not os.path.isfile(no_warn_file):
            print('Using LaTeX for rendering. If an error occurs:')
            print('a) If you are using `andes plot`, disable with option "-d",')
            print('b) If you are using `plot()`, set "latex=False".')

            try:
                with open(os.path.join(get_dot_andes_path(), '.no_warn_latex'), 'w') as f:
                    f.write('0')
            except OSError:
                pass

        return True

    return False


# -*- coding: utf-8 -*-
"""
Created on Sun Jan 28 21:24:59 2018

@author: HP
"""

from numpy import asarray
from numpy import zeros
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers import LSTM, Dropout, GRU
from keras.layers import Embedding
from collections import defaultdict
from nltk.corpus import brown,stopwords
import random
import nltk
import numpy as np
#brown.categories()
#brown.words(categories='news')
#brown.words(fileids=['cg22'])
#brown.sents(categories=['news', 'editorial', 'reviews'])

dataset = [] # 500 samples

for category in brown.categories():
    for fileid in brown.fileids(category):
        dataset.append((brown.words(fileids = fileid),category))

dataset = [([w.lower() for w in text],category) for text,category in dataset]

labels=[]
for sample in dataset:
    labels.append(sample[1])

inputset=[]
for sample in dataset:
    inputset.append(' '.join(sample[0]))

input_train=[]
j=0;
for zz in inputset:
    j=j+1
    if (j%5 is not 0):
        input_train.append(zz)

input_test=[]
j=0;
for zz in inputset:
    j=j+1
    if (j%5 is 0):
        input_test.append(zz)

categ=brown.categories()

label_class=[]
for x in labels:
    label_class.append(categ.index(x))

label_train=[]
j=0;
for zz in label_class:
    j=j+1
    if (j%5 is not 0):
        label_train.append(zz)

label_test=[]
j=0;
for zz in label_class:
    j=j+1
    if (j%5 is 0):
        label_test.append(zz)

len_finder=[]
for dat in inputset:
    len_finder.append(len(dat))

#one hot encoding
i=0
y=np.zeros((len(label_class),max(label_class)+1))
for x in label_class:
    y[i][x]=1
    i=i+1

i=0
y_train=np.zeros((len(label_train),max(label_train)+1))
for x in label_train:
    y_train[i][x]=1
    i=i+1

i=0
y_test=np.zeros((len(label_test),max(label_test)+1))
for x in label_test:
    y_test[i][x]=1
    i=i+1

t = Tokenizer()
t.fit_on_texts(input_train)
vocab_size = len(t.word_index) + 1
# integer encode the documents
encoded_docs = t.texts_to_sequences(input_train)
#print(encoded_docs)
# pad documents to a max length of 4 words
max_length = max(len_finder)
padded_docs = pad_sequences(encoded_docs, maxlen=max_length, padding='post')
#print(padded_docs)
# load the whole embedding into memory
embeddings_index = dict()
f = open("G:\\NLP\\Dataset\\GloVe\\glove.6B.100d.txt",  encoding="utf8")
for line in f:
	values = line.split()
	word = values[0]
	coefs = asarray(values[1:], dtype='float32')
	embeddings_index[word] = coefs
f.close()
print('Loaded %s word vectors.' % len(embeddings_index))
# create a weight matrix for words in training docs
embedding_matrix = zeros((vocab_size, 100))
for word, i in t.word_index.items():
	embedding_vector = embeddings_index.get(word)
	if embedding_vector is not None:
		embedding_matrix[i] = embedding_vector

# define model
model = Sequential()
e = Embedding(vocab_size, 100, weights=[embedding_matrix], input_length=max(len_finder), trainable=False)
model.add(e)
#model.add(Flatten())
model.add(LSTM(100))
model.add(Dropout(0.5))

model1=Sequential()
model1.add(model)

model1.add(Dense(15, activation='softmax'))

# compile the model
model1.compile(optimizer='adam', loss='binary_crossentropy', metrics=['acc'])
# summarize the model
print(model1.summary())
# fit the model
model1.fit(padded_docs, y_train, epochs=1, verbose=0)

#Testing the model
tt = Tokenizer()
tt.fit_on_texts(input_test)
tvocab_size = len(tt.word_index) + 1
# integer encode the documents
tencoded_docs = tt.texts_to_sequences(input_test)
#print(encoded_docs)
# pad documents to a max length of 4 words
tpadded_docs = pad_sequences(tencoded_docs, maxlen=max_length, padding='post')
#print(padded_docs)

# evaluate the model
loss, accuracy = model1.evaluate(tpadded_docs, y_test, verbose=0)
print('Accuracy: %f' % (accuracy*100))

Simple Text classifiers/Text Classification on Brown Corpus/Brown_classifier.py


import os
import codecs
import numpy as np
import tensorflow as tf
from models.attention import multihead_attention_decoder
from models.basemodel import BaseModel, make_custom_embedding_matrix
from utils import gumbel_softmax
from models.attention import attention_decoder

class DiverEmbMin(BaseModel):
    def __init__(self, vocab, hps):
        super(DiverEmbMin, self).__init__(vocab, hps)

    def build_model(self):
        print('models build start')
        with tf.variable_scope('models'):
            self.rand_unif_init = tf.random_uniform_initializer(-self.hps.rand_unif_init_size, self.hps.rand_unif_init_size, seed=777)
            self.trunc_norm_init = tf.truncated_normal_initializer(stddev=self.hps.trunc_norm_init_std)

            with tf.variable_scope('embedding'):
                print('add embedding')
                self.add_embedding()

            with tf.variable_scope('encoder'):
                print("add encoding")
                self.add_encoder()
                self.reduce_state()
                self.add_mechanism()
                self.dec_init_with_mechanism()

            with tf.variable_scope('decoder'):
                print('add decoding')
                self.add_decoder()

            with tf.variable_scope('output_projection'):
                self.add_output_projection()

            if self.hps.mode not in ['decode']:
                self.regularize_by_attn_dists_timestep()
                self.add_loss_calculation()
            else:
                print('#######', len(self.dec_vocab_dists))
                assert len(self.dec_vocab_dists) == 1
                topk_probs_dec, self.topk_ids = tf.nn.top_k(self.dec_vocab_dists[0], self.hps.batch_size * 2)

                self.topk_log_probs = tf.log(topk_probs_dec)

    def reduce_state(self):
        """
        Reduce the bi-directional encoder output state size.
        """
        w_reduce = tf.get_variable('w_reduce', [self.hps.hidden_dim * 2, self.hps.hidden_dim],
                                   dtype=tf.float32, initializer=self.trunc_norm_init)
        b_reduce = tf.get_variable('b_reduce', [self.hps.hidden_dim],
                                   dtype=tf.float32, initializer=self.trunc_norm_init)
        concat_h = tf.concat(axis=1, values=[self.bi_bw_st, self.bi_fw_st])

        new_h = tf.nn.relu(tf.matmul(concat_h, w_reduce) + b_reduce)

        # [batch_size, hidden_dim]
        self.dec_init_state_small = new_h

    def add_mechanism(self):
        self.matrix = tf.get_variable('mech_mat', [self.hps.matrix_num, self.hps.mechanism_dim], dtype=tf.float32,
                                      initializer=self.rand_unif_init)

        W_enc = tf.get_variable('W_enc', [self.hps.hidden_dim, 2 * self.hps.hidden_dim], dtype=tf.float32,
                                initializer=self.trunc_norm_init)

        enc_feature = tf.matmul(self.dec_init_state_small, W_enc)
        maxout_res = tf.contrib.layers.maxout(enc_feature, self.hps.hidden_dim)  # [batch_size, hidden_dim]

        W_attn = tf.get_variable('W_divert_attn', [self.hps.mechanism_dim, self.hps.hidden_dim], dtype=tf.float32,
                                 initializer=self.trunc_norm_init)
        res1 = tf.matmul(self.matrix, W_attn)

        # [batch_size, mechanism_num]
        self.emb_score = tf.nn.softmax(tf.transpose(tf.matmul(res1, maxout_res, transpose_b=True))) + 1e-12

    def regularize_by_attn_dists_timestep(self):
        # attn_dist shape: [max len, Tensor(BatchSize * matrix_num, dim)]
        attn = tf.stack(self.attn_dists[0])
        attn = tf.reduce_sum(attn, axis=0)  # [B * matrix_num, enc_len]

        attnsum = tf.reduce_sum(attn, axis=1, keep_dims=True)
        attn /= attnsum  # normalized
        attn = tf.reshape(attn, [self.hps.batch_size, self.hps.matrix_num, -1])
        aa_t = tf.matmul(attn, tf.transpose(attn, [0, 2, 1]))
        eye = tf.reshape(tf.tile(tf.eye(self.hps.matrix_num), [self.hps.batch_size, 1]),
                         [-1, self.hps.matrix_num, self.hps.matrix_num])
        self.P = tf.reduce_mean(tf.square(tf.norm(aa_t - eye, axis=[-2, -1], ord='fro')))

    def dec_init_with_mechanism(self):
        # [Batch_size * mechanism_num, mechanism_dim] with [mec1, mec2 .... mecn, mec1, ..... mec n]
        tiled_mechanism = tf.tile(self.matrix, [self.hps.batch_size, 1])

        # [batch_size * mechanism_num, hidden_dim] with [state1, state1, ..... state 2, state2 ... state n, state n]
        tiled_dec_init_state = tf.reshape(tf.tile(self.dec_init_state_small, [1, self.hps.matrix_num]),
                                          [self.hps.batch_size * self.hps.matrix_num, self.hps.hidden_dim])

        # [B * mech_num, dec_hidden_dim] with [state1 + mech1, state1 + mech2 ...]
        self.dec_init_state_cloned = tf.concat([tiled_mechanism, tiled_dec_init_state], axis=1)

    def add_loss_calculation(self):
        dec_tgt_batch = tf.tile(self.dec_tgt_batch, [1, self.hps.matrix_num])
        dec_tgt_batch = tf.reshape(dec_tgt_batch, [self.hps.batch_size * self.hps.matrix_num, self.hps.max_dec_len])

        dec_pad_masks = tf.tile(self.dec_pad_masks, [1, self.hps.matrix_num])
        dec_pad_masks = tf.reshape(dec_pad_masks, [self.hps.batch_size * self.hps.matrix_num, self.hps.max_dec_len])

        seq_loss = tf.contrib.seq2seq.sequence_loss(
            tf.stack(self.dec_vocab_scores, axis=1),
            dec_tgt_batch, weights=dec_pad_masks, average_across_batch=False)

        seq_loss = tf.reshape(seq_loss, [self.hps.batch_size, self.hps.matrix_num])
        min_loss = tf.math.reduce_min(seq_loss, axis=1)
        self.seq_loss = tf.reduce_mean(min_loss)
        self.min_loss_index = tf.argmin(seq_loss, axis=1)

        self.loss = self.seq_loss + self.hps.emb_min_coeff * self.P

        tf.summary.scalar('penalty', self.P)
        tf.summary.scalar('loss_seq', self.seq_loss)
        tf.summary.scalar('loss', self.loss)

    def add_decoder(self):
        cell_fw1 = tf.nn.rnn_cell.GRUCell(self.hps.decoder_hidden_dim, kernel_initializer=self.rand_unif_init)
        cell_fw1 = tf.contrib.rnn.DropoutWrapper(cell=cell_fw1, input_keep_prob=self.keep_rate)

        if self.hps.mode == 'decode':
            emb_dec_inp = self.emb_dec_inp
            enc_pad_mask = self.enc_pad_masks
            self.dec_init_state = tf.placeholder(tf.float32, [self.hps.batch_size, self.hps.decoder_hidden_dim],
                                                 name='placeholder')

        else:
            self.dec_init_state = self.dec_init_state_cloned  # [B1 + S1, B1 + S2, B1 + S3 ....]

            emb_dec_inp = tf.stack(self.emb_dec_inp)
            tiled_dec_inp = tf.tile(emb_dec_inp, [1, 1, self.hps.matrix_num])
            reshape_dec_inp = tf.reshape(tiled_dec_inp, [self.hps.max_dec_len, self.hps.batch_size *
                                                         self.hps.matrix_num, self.hps.embed_dim])

            emb_dec_inp = tf.unstack(reshape_dec_inp, axis=0)
            enc_pad_mask = tf.reshape(tf.tile(self.enc_pad_masks, [1, self.hps.matrix_num]),
                                      [-1, tf.shape(self.enc_pad_masks)[1]])

            self.encoder_outputs = tf.reshape(tf.tile(self.encoder_outputs, multiples=[1, self.hps.matrix_num, 1]),
                                              [self.hps.batch_size * self.hps.matrix_num, -1,
                                               2 * self.hps.hidden_dim])

        self.dec_outputs, dec_out_states, attn_dists = attention_decoder(emb_dec_inp, self.dec_init_state,
                                                                         self.encoder_outputs, enc_pad_mask,
                                                                         cell_fw1,
                                                                         initial_state_attention=self.hps.mode == 'decode')
        self.dec_out_states = (dec_out_states, [])
        self.attn_dists = (attn_dists, [], [])

    def add_embedding(self):
        self.embedding_matrix = make_custom_embedding_matrix(self.vocab, self.hps)
        self.emb_enc_inp = tf.nn.embedding_lookup(self.embedding_matrix, self.enc_batch)
        self.emb_dec_inp = [tf.nn.embedding_lookup(self.embedding_matrix, x) for x in
                            tf.unstack(self.dec_inp_batch, axis=1)]

    def run_encoder(self, batch, sess, selected_mechanism=-1):
        assert 0 <= selected_mechanism < self.hps.matrix_num
        feeddict = self.make_feeddict(batch, encoder_only=True)
        feeddict[self.keep_rate] = 1.0
        enc_states, dec_init_state = sess.run([self.encoder_outputs, self.dec_init_state_cloned], feeddict)

        dec_init_state = dec_init_state[selected_mechanism]

        return enc_states, dec_init_state

    def add_output_projection(self):
        with tf.variable_scope('dec_output_projection'):
            w = tf.get_variable('w1', [self.hps.decoder_hidden_dim, len(self.vocab.words)],
                                dtype=tf.float32, initializer=self.trunc_norm_init)
            v = tf.get_variable('v1', [len(self.vocab.words)],
                                dtype=tf.float32, initializer=self.trunc_norm_init)

            dec_flatten = tf.reshape(tf.stack(self.dec_outputs), [-1, self.hps.decoder_hidden_dim])

            dec_vocab_scores = tf.nn.xw_plus_b(dec_flatten, w, v)
            if self.hps.mode == 'decode':
                dec_vocab_scores = tf.reshape(dec_vocab_scores,
                                              [-1, self.hps.batch_size, len(self.vocab.words)])
            else:
                dec_vocab_scores = tf.reshape(dec_vocab_scores,
                                              [-1, self.hps.batch_size * self.hps.matrix_num, len(self.vocab.words)])

            dec_vocab_dists = tf.nn.softmax(dec_vocab_scores)

            self.dec_vocab_scores = tf.unstack(dec_vocab_scores)
            self.dec_vocab_dists = tf.unstack(dec_vocab_dists)

    def run_step(self, batch, sess, is_train=False, freeze_layer=False):
        feeddict = self.make_feeddict(batch)
        to_return = {
            'summaries': self.summaries,
            'loss': self.loss,
            'global_step': self.global_step
        }
        if is_train:
            to_return['selected_emb_idx'] = self.min_loss_index
            if freeze_layer:
                feeddict[self.freeze_layer] = freeze_layer
                to_return['train_op'] = self.train_op
                # to_return['train_op'] = self.train_op_freeze
            else:
                to_return['train_op'] = self.train_op

        return sess.run(to_return, feeddict)


import sys, os, socket
os.environ["CUDA_VISIBLE_DEVICES"]="0"
hostname = socket.gethostname()

if hostname=='tianx-pc':
    homeDir = '/analyse/cdhome/'
    proj0257Dir = '/analyse/Project0257/'
elif hostname[0:7]=='deepnet':
    homeDir = '/home/chrisd/'
    proj0257Dir = '/analyse/Project0257/'

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from PIL import Image
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import ModelCheckpoint, TensorBoard, ProgbarLogger
from io import StringIO
import h5py

sys.path.append(os.path.abspath(homeDir+'dlfaceScripts/SchynsLabDNN/faceNets/'))
from resNetUtils import getActVAE
from vae_models import ResNet10DetEncoder, ResNet10Decoder, Darknet19Encoder, Darknet19Decoder
from vae import AutoEncoder

dataDir = proj0257Dir+'christoph_face_render_withAUs_20190730/'
main_data_txt = dataDir+'images_firstGen_ctrlSim_k1/path/linksToImages.txt'
main_df = pd.read_csv(main_data_txt, delim_whitespace = True, header=None)
main_df.columns = ['filename', 'yID', 'yVector', 'yGender', 'yEthn', 'yAge', 'yEmo', 'yAnglex', 'yAngley', 'yAnglelx', 'yAnglely']
main_df = main_df[['filename']]

maintrain_df = main_df.sample(frac=0.9,random_state=1)
test_df = main_df.drop(maintrain_df.index)
train_df = maintrain_df.sample(frac=0.8,random_state=1)
val_df = maintrain_df.drop(train_df.index)

inputShape = (224, 224, 3)
batchSize = 50
latentSize = 512
epochs = 100

for beta in [1]: #[1,2,5,10,20]:

    # build and compile the autoencoder
    encoder = ResNet10Encoder(inputShape, latentSize=latentSize, latentConstraints='bvae', beta=beta)
    decoder = Darknet19Decoder(inputShape, latentSize=latentSize)
    bvae = AutoEncoder(encoder, decoder)
    bvae.ae.compile(optimizer='adam', loss='mean_absolute_error')

    bvae.ae.load_weights(proj0257Dir+'aeModels/1stGen/001530SimStructk1ColleaguesRN10DN19beta'+str(beta)+'.h5')

    # define source path
    basePth = proj0257Dir+'christoph_face_render_withAUs_20190730/colleaguesRandomJiayu/'
    # define other constants
    genderTxt = ['f','m']
    nBatch = 9

    for gg in range(2):
        for id in range(2):
                for rr in range(2):
                    for cc in range(3):

                        print('beta '+str(beta)+' gg '+str(gg+1)+' id '+str(id+1)+' rr '+str(rr+1)+' cc '+str(cc+1))
                        ths_txt = basePth+genderTxt[gg]+'/id'+str(id+1)+'/linksToImages_array_'+str(rr+1)+'_'+str(cc+1)+'.txt'

                        ths_df = pd.read_csv(ths_txt, delim_whitespace = True, header=None)
                        ths_df.columns = ['filename', 'yID', 'yVector', 'yGender', 'yEthn', 'yAge', 'yEmo', 'yAnglex', 'yAngley', 'yAnglelx', 'yAnglely']
                        ths_df = ths_df[['filename']]
                        test_datagen = ImageDataGenerator(rescale=1./255)
                        ths_generator = test_datagen.flow_from_dataframe(
                            dataframe=ths_df,
                            target_size=(inputShape[0],inputShape[1]),
                            directory='/',
                            x_col='filename',
                            class_mode='input',
                            validate_filenames=False,
                            shuffle=False,
                            batch_size=int(ths_df.shape[0]/nBatch))

                        for bb in range(nBatch):
                            print('loading batch '+str(bb))
                            thsBatch, thsLabels = next(ths_generator)
                            print('evaluating model ... ')
                            latentVec = bvae.encoder.predict_on_batch(thsBatch)

                            thsDestinDir = proj0257Dir+'humanReverseCorrelation/activations/vae/trialsRandom/beta'+str(beta)+'/'+genderTxt[gg]+'/id'+str(id+1)+'/row'+str(rr+1)+'/col'+str(cc+1)+'/'

                            print('saving layer activations')
                            getActVAE(bvae.encoder, thsBatch, thsDestinDir, startLayer=0, batchNr=bb)

                            #print('saving latent vectors')
                            #thsFileName = thsDestinDir+'latent_batch_'+str(bb+1)+'.h5'
                            #hf = h5py.File((thsFileName), 'w')
                            #hf.create_dataset('latentVec', data=latentVec)
                            #hf.close()


from typing import Optional, Sequence
import numpy as np

class MoleculeGraphData:
  """MoleculeGraphData class

  This data class is almost same as `torch_geometric.data.Data
  <https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data>`_.

  Attributes
  ----------
  node_features : np.ndarray
    Node feature matrix with shape [num_nodes, num_node_features]
  edge_index : np.ndarray, dtype int
    Graph connectivity in COO format with shape [2, num_edges]
  targets : np.ndarray
    Graph or node targets with arbitrary shape
  edge_features : np.ndarray, optional (default None)
    Edge feature matrix with shape [num_edges, num_edge_features]
  graph_features : np.ndarray, optional (default None)
    Graph feature vector with shape [num_graph_features,]
  num_nodes : int
    The number of nodes in the graph
  num_node_features : int
    The number of features per node in the graph
  num_edges : int
    The number of edges in the graph
  num_edges_features : int, , optional (default None)
    The number of features per edge in the graph
  """

  def __init__(
      self,
      node_features: np.ndarray,
      edge_index: np.ndarray,
      targets: np.ndarray,
      edge_features: Optional[np.ndarray] = None,
      graph_features: Optional[np.ndarray] = None,
  ):
    """
    Parameters
    ----------
    node_features : np.ndarray
      Node feature matrix with shape [num_nodes, num_node_features]
    edge_index : np.ndarray, dtype int
      Graph connectivity in COO format with shape [2, num_edges]
    targets : np.ndarray
      Graph or node targets with arbitrary shape
    edge_features : np.ndarray, optional (default None)
      Edge feature matrix with shape [num_edges, num_edge_features]
    graph_features : np.ndarray, optional (default None)
      Graph feature vector with shape [num_graph_features,]
    """
    # validate params
    if isinstance(node_features, np.ndarray) is False:
      raise ValueError('node_features must be np.ndarray.')
    if isinstance(edge_index, np.ndarray) is False:
      raise ValueError('edge_index must be np.ndarray.')
    elif edge_index.dtype != np.int:
      raise ValueError('edge_index.dtype must be np.int')
    elif edge_index.shape[0] != 2:
      raise ValueError('The shape of edge_index is [2, num_edges].')
    if isinstance(targets, np.ndarray) is False:
      raise ValueError('y must be np.ndarray.')
    if edge_features is not None:
      if isinstance(edge_features, np.ndarray) is False:
        raise ValueError('edge_features must be np.ndarray or None.')
      elif edge_index.shape[1] != edge_features.shape[0]:
        raise ValueError('The first dimension of edge_features must be the \
                    same as the second dimension of edge_index.')
    if graph_features is not None and isinstance(graph_features,
                                                 np.ndarray) is False:
      raise ValueError('graph_features must be np.ndarray or None.')

    self.node_features = node_features
    self.edge_index = edge_index
    self.edge_features = edge_features
    self.graph_features = graph_features
    self.targets = targets
    self.num_nodes, self.num_node_features = self.node_features.shape
    self.num_edges = edge_index.shape[1]
    if self.node_features is not None:
      self.num_edge_features = self.edge_features.shape[1]

  def to_pyg_data(self):
    """Convert to PyTorch Geometric Data instance

    Returns
    -------
    torch_geometric.data.Data
      Molecule graph data for PyTorch Geometric
    """
    try:
      import torch
      from torch_geometric.data import Data
    except ModuleNotFoundError:
      raise ValueError("This class requires PyTorch Geometric to be installed.")

    return Data(
      x=torch.from_numpy(self.node_features),
      edge_index=torch.from_numpy(self.edge_index),
      edge_attr=None if self.edge_features is None \
        else torch.from_numpy(self.edge_features),
      y=torch.from_numpy(self.targets),
    )

class BatchMoleculeGraphData(MoleculeGraphData):
  """Batch MoleculeGraphData class

  Attributes
  ----------
  graph_index : np.ndarray, dtype int
    This vector indicates which graph the node belongs with shape [num_nodes,]
  """

  def __init__(self, molecule_graphs: Sequence[MoleculeGraphData]):
    """
    Parameters
    ----------
    molecule_graphs : Sequence[MoleculeGraphData]
      List of MoleculeGraphData
    """
    # stack features and targets
    batch_node_features = np.vstack(
        [graph.node_features for graph in molecule_graphs])
    batch_targets = np.vstack([graph.targets for graph in molecule_graphs])

    # before stacking edge_features or graph_features,
    # we should check whether these are None or not
    if molecule_graphs[0].edge_features is not None:
      batch_edge_features = np.vstack(
          [graph.edge_features for graph in molecule_graphs])
    else:
      batch_edge_features = None

    if molecule_graphs[0].graph_features is not None:
      batch_graph_features = np.vstack(
          [graph.graph_features for graph in molecule_graphs])
    else:
      batch_graph_features = None

    # create new edge index
    num_nodes_list = [graph.num_nodes for graph in molecule_graphs]
    batch_edge_index = np.hstack(
      [graph.edge_index + prev_num_node for prev_num_node, graph \
        in zip([0] + num_nodes_list[:-1], molecule_graphs)]
    ).astype(int)

    # graph_index indicates which nodes belong to which graph
    graph_index = []
    for i, num_nodes in enumerate(num_nodes_list):
      graph_index.extend([i] * num_nodes)
    self.graph_index = np.array(graph_index, dtype=int)

    super().__init__(
        node_features=batch_node_features,
        edge_index=batch_edge_index,
        targets=batch_targets,
        edge_features=batch_edge_features,
        graph_features=batch_graph_features,
    )

    @staticmethod  # type: ignore
    def to_pyg_data(molecule_graphs: Sequence[MoleculeGraphData]):
      """Convert to PyTorch Geometric Batch instance

      Parameters
      ----------
      molecule_graphs : Sequence[MoleculeGraphData]
        List of MoleculeGraphData

      Returns
      -------
      torch_geometric.data.Batch
        Batch data of molecule graph for PyTorch Geometric
      """
      try:
        from torch_geometric.data import Batch
      except ModuleNotFoundError:
        raise ValueError(
            "This class requires PyTorch Geometric to be installed.")

      data_list = [mol_graph.to_pyg_data() for mol_graph in molecule_graphs]
      return Batch.from_data_list(data_list=data_list)


# coding: utf-8

from __future__ import division, print_function

import os, sys
import tensorflow as tf
import time
import cv2
import numpy as np

from utils import plot_one_box

os.environ["CUDA_VISIBLE_DEVICES"] = '0'

color_table = [[0,255,0], [255,0,0],[0,0,255]]

classes = ['face', 'mask']

def letterbox_resize(img, new_width, new_height, interp=0):
    '''
    Letterbox resize. keep the original aspect ratio in the resized image.
    '''
    ori_height, ori_width = img.shape[:2]

    resize_ratio = min(new_width / ori_width, new_height / ori_height)

    resize_w = int(resize_ratio * ori_width)
    resize_h = int(resize_ratio * ori_height)

    img = cv2.resize(img, (resize_w, resize_h), interpolation=interp)
    image_padded = np.full((new_height, new_width, 3), 128, np.uint8)

    dw = int((new_width - resize_w) / 2)
    dh = int((new_height - resize_h) / 2)

    image_padded[dh: resize_h + dh, dw: resize_w + dw, :] = img

    return image_padded, resize_ratio, dw, dh

def get_img_list(img_path, exts=['jpg', 'png', 'jpeg', 'JPG']):

    img_list = os.listdir(img_path)
    new_list = []
    for img_name in img_list:
        for ext in exts:
            if img_name.endswith(ext):
                new_list.append(img_name)
                break
    return new_list

def inference(img_dir, out_dir):

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        with tf.gfile.FastGFile('./sur/sur0316/yolov3_tiny_sur.pb', 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name='') #

        sess.run(tf.global_variables_initializer())

        input_data = sess.graph.get_tensor_by_name('input_data:0')
        labels = sess.graph.get_tensor_by_name('labels:0')
        scores = sess.graph.get_tensor_by_name('scores:0')
        boxes = sess.graph.get_tensor_by_name('boxes:0')

        img_names = os.listdir(img_dir)
        count = 0
        times = []
        for img_name in img_names:
            img_ori = cv2.imread(os.path.join(img_dir, img_name))
            #img_ori = cv2.imread(img_name)
            start = time.time()
            img, resize_ratio, dw, dh = letterbox_resize(img_ori, 608, 608)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            "img = img[np.newaxis, :]"
            boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})

            boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
            boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio

            end = time.time()

            for i in range(len(boxes_)):
                x0, y0, x1, y1 = boxes_[i]
                plot_one_box(img_ori, [x0, y0, x1, y1], label=classes[labels_[i]]+',{:.2f}'.format(scores_[i]), color=color_table[labels_[i]])

            img_name = os.path.basename(img_name)
            cv2.imwrite(os.path.join(out_dir, img_name), img_ori)
            count += 1
            print('No.{}, img:{}, time:{:.4f}'.format(count, img_name, end-start))

            if count > 1:
                times.append(end-start)

        print('Total:{}, avg time:{:.4f}'.format(count, np.mean(times)))

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = '1'

    img_dir = './data/test'
    out_dir = './data/test_out'

    inference(img_dir, out_dir)