import os
import cv2
import glob
import torch
import argparse
import numpy as np
import torchvision.transforms as transforms
from net import vgg,FPnet,Decoder,testdecoder
from PIL import Image, ImageFile
from function import coral,change_color
from torchvision.utils import save_image

parser = argparse.ArgumentParser()
# Basic options
parser.add_argument('-c','--content',type=str,
                    help='path of your content image')
parser.add_argument('-s','--style',type=str,
                    help='path of your style image')

args = parser.parse_args()

#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")#训练设备

def test_transform(size):
    transform_list = []
    transform_list.append(transforms.Resize(size))
    transform_list.append(transforms.ToTensor())
    transform = transforms.Compose(transform_list)
    return transform

def transfer(contentimg,fbnet,converted="result.jpg"):#20200521decoder10000_1.pth
    '''一次前传得到风格化图像'''
    #mytransfer=test_transform(pixel)

    #contentimg = Image.open(str(contentpath)).convert('RGB')
    #styleimg = Image.open(str(stylepath)).convert('RGB')

    #contentimg=mytransfer(contentimg).unsqueeze(0)
    #styleimg=mytransfer(styleimg).unsqueeze(0)

    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    #decoder=Decoder().to(device).eval()
    #decoder=testdecoder
    #decoder.load_state_dict(torch.load(model_path,map_location="cpu"))

    # decoder = decoder.module
    # decoder.load_state_dict(torch.load(model_path))

    output=fbnet(contentimg,alpha=1.0,lamda=1.0,require_loss=False)

    #save_image(output.cpu(),converted)
    contentimg.detach()
    output.detach()
    return True

if __name__ == "__main__":
    decoder=testdecoder
    decoder.load_state_dict(torch.load('model/decoder.pth',map_location="cpu"))
    encoder=vgg
    encoder.load_state_dict(torch.load('model/vgg_normalised.pth',map_location="cpu"))
    transfer(args.content,args.style,encoder,decoder)


#! /usr/bin env python

"""
Code to generate WFSS dispersed seed images. Starting with a set of imaging seed images,
potentially with padding, given a JWST WFSS GRISMCONF configuration file
"""

import os
from astropy.io import fits
import numpy as np
from .observations.observations \
import observation as Gsim_observation
from multiprocessing import cpu_count

class Grism_seed():
    def __init__(self,image_seeds,cross_filter,mode,config_path=".",extrapolate_SED=False,SED_file=None,instrument="NIRCAM",max_cpu=None, SBE_save=None, renormalize=True, resample=False):
        """A class for a grism simulation

        Attributes
        ----------
        image_seeds: list
            A list of image image_seeds
        cross_filter: str
            A string containing the name of a direct filter
        mode: str
            A string containing either R or C
        config_path: str
            As string pointing to a GRISMCONF configuration file
        extrapolate_SED: bol
            If set to True, the objects' SED will be extrapolated if needed
        SED_file: str
            A string containing the name of an HDF5 file containing the spectra of each source
        instrument: str
            A string containing the name of the instrument (e.g. NIRCAM)
        max_cpu: int
            An integer containing the number of CPU to use in the multiheaded pool when dispersing
        SBE_save: str
            A string containing the name of an output HDF5 file which will contain simulated 2D stamps of each source
        renormalize: bol
            Whether to renormalize the input data to unity over segmentation map area when using an input spectrum.
        reaample: bol
            If true, the disperser will first smooth and resample input spectra appropriately to match the disperser resolution.
        Methods
        -------
        observation(self,orders=["+1","+2"],max_split=-1000,ID=0)

        finalize(self,tofits=None,Back=False,BackLevel=None)

        saveSingleFits(self,name)

        """

        config = os.path.join(config_path,"%s_%s_%s.conf" % (instrument,cross_filter,mode))
        self.config = config

        self.image_seeds = image_seeds
        self.cross_filter = cross_filter
        self.mode = mode
        self.config_path = config_path
        if max_cpu is None:
            max_cpu = cpu_count() - 1

        self.max_cpu = max_cpu
        self.SBE_save = SBE_save

        if self.SBE_save != None:
            print("Will output to ", self.SBE_save)
            if os.path.isfile(self.SBE_save):
                os.unlink(self.SBE_save)

        # Get information about input padding. We use the first image seed for this, just like for the segmentation info.
        h = fits.open(image_seeds[0])[0].header
        self.xstart = int(h["NOMXSTRT"])
        self.xend = int(h["NOMXEND"])
        self.ystart = int(h["NOMYSTRT"])
        self.yend = int(h["NOMYEND"])

        # Get segmentation info, from the first image seed.
        self.seg_data = fits.open(image_seeds[0])[2].data

        self.extrapolate_SED = extrapolate_SED
        self.SED_file = SED_file
        self.renormalize = renormalize
        self.resample = resample

    def observation(self,orders=None,max_split=-1000,ID=0):
        """Sets up an observation.

        Parameters
        ----------
        orders : list
            A list of string containing the name of the orders to disperse
        max_split : int
            Maximum number of pixels to disperse at once (not currently used)
        ID : int
            Specific object ID to dispersed. Set to 0 (default) to disperse all of available objects
        """

        self.this_one = {}

        # If orders are not passed, we get them from the config file

        if orders==None:
            import grismconf
            C = grismconf.Config(self.config)
            print("orders:",C.orders)
            self.orders = C.orders
        else:
            self.orders = orders

        for order in self.orders:
            boundaries = [self.xstart,self.xend,self.ystart,self.yend]
            self.this_one[order] = Gsim_observation(self.image_seeds,self.seg_data,self.config,order=order,max_split=max_split,extrapolate_SED=self.extrapolate_SED,SED_file=self.SED_file,max_cpu=self.max_cpu,ID=ID, SBE_save=self.SBE_save,boundaries=boundaries,renormalize=self.renormalize,resample=self.resample)
            #self.this_one[order].disperse_all()

    def disperse(self,orders=None,cache=False,trans=None):
        """Run the disperser.

        Parameters
        ----------
        orders: list
            Optional list containing the name of the orders to disperse
        cache: bool
            If set to True, the dispersion tables are cached and will be used on subsequent calls.
        """
        if orders==None:
            orders = self.orders

        print("Dispersing orders ", orders)
        for order in orders:
            #print("Dispersing order ",order)
            if self.this_one[order].cache:
                self.this_one[order].disperse_all_from_cache(trans=trans)
            else:
                self.this_one[order].disperse_all(cache=cache)

    def disperse_background_1D(self,background):
        """Produces a dispersed 2D image of the background spectrum contained in the fits file background, meant to be the
        output of thr jwst_background module. This background is dispersed either in the row or column direction, depending on the
        dispersion, and the result is tiled to produce a full 2D image. All orders are generated and added up.

        Parameters
        ----------
        background: 2D numpy array [lambda values,flux values]
            A 2D numpy array containing the spectrum of the background to disperse.
            The wavelength should be in (micron) and the flux (in Mjy/sr), as is produced by
            the jwst_background package
        output: numpy 2D array
            A 2D array containing the model background which can be fed back into finalize()
        """

        bck = 0.
        for order in self.orders:
            print("Computing dispersed background for order ",order)
            bck += self.this_one[order].disperse_background_1D(background)

#        fits.writeto("WFSS_background.fits",bck,overwrite=True)
        return bck

    def finalize(self,Back=None,BackLevel=None,tofits=None):
        """ Produces a 2D dispersed image and add the appropriate background

        Parameters
        ----------
        tofits: str
            Name of a fits file to write the simulation to. Default is set to None
        Back: str
            Name of a fits file containing an image of the background in e-/s in extension 0
            If None, the file listed in the config file is used.
        BackLevel: float
            Renormalization factor for the background image. Renormalization is done by multiplying by BackLevel/np.median(Back)
        """

        # Initialize final image with the background estimate
        final = 0.
        if (Back is None) and (BackLevel is not None):
            # Use pre-computed background from config file, scaled by BackLevel
            import grismconf
            bck_file = grismconf.Config(self.config).BCK
            print("Adding pre-computed 2D dispersed background ",bck_file,"scaled to",BackLevel)
            import grismconf
            final = fits.open(bck_file)[1].data * BackLevel
        if (Back is None) and (BackLevel is None):
            # Use no background
            print("No background added")
            final = 0.
        if (type(Back)==np.ndarray) and (BackLevel is not None):
            # Use a passed background and scale its median to BackLevel
            print("adding passed background array scaled to",BackLevel)
            final = Back/np.median(Back) * BackLevel
        if (type(Back)==np.ndarray) and (BackLevel is None):
            # Use a passed background as is
            print("adding passed background array as is")
            final = Back

        if not ((Back is None) and (BackLevel is None)) and (tofits!=None):
            # Save the background image to a fits file
            hprime = fits.PrimaryHDU()
            himg = fits.ImageHDU(final)
            himg.header['EXTNAME'] = 'BACKGRND'
            himg.header['UNITS'] = 'e/s'
            if BackLevel is not None:
                himg.header['BackLevel'] = BackLevel
            hlist = fits.HDUList([hprime, himg])
            hlist.writeto(tofits, overwrite=True)

        for order in self.orders:
            print("Adding contribution from order ",order)
            try:
                sim = self.this_one[order].simulated_image[self.ystart:self.yend+1,self.xstart:self.xend+1]
            except AttributeError:
                print("Contribution from order",order,"is missing. Skipping it.")
                continue
            final = final + sim
        self.final = final

        if (Back is None) and (BackLevel is None) and (tofits!=None):
            # Save the background image to a fits file
            hprime = fits.PrimaryHDU()
            himg = fits.ImageHDU(final)
            himg.header['EXTNAME'] = 'BACKGRND'
            himg.header['UNITS'] = 'e/s'
            if BackLevel is not None:
                himg.header['BackLevel'] = BackLevel
            hlist = fits.HDUList([hprime, himg])
            hlist.writeto(tofits, overwrite=True)

    def saveSingleFits(self,name):
        """A helper function to write the 2D simulated imae into a fits file

        Parameters
        ----------
        fname: str
            The file name to use for the output
        """

        #save an array into the first extension of a fits file
        h0 = fits.PrimaryHDU()
        h1 = fits.ImageHDU(self.final,name='DATA')

        hdulist = fits.HDUList([h0,h1])
        hdulist.writeto(name,overwrite=True)

if __name__ == '__main__':
    import glob

    image_seeds = glob.glob("V4*.fits")
    seed = Grism_seed(image_seeds,"F444W","modA_R","/Users/npirzkal/Dropbox/GRISMDATA/NIRCAM/")


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

import random
import numpy as np
from utils.loader import Loader

import torch
import pickle
from functools import partial
import sys
import time
from PIL import Image, ImageDraw
sys.path.append(sys.path[0] + "/../yolov3")
# print(sys.path)
from yolov3.utils import *
from yolov3.image import letterbox_image, correct_yolo_boxes
from yolov3.darknet import Darknet
from ROI import *

import h5py, tqdm, spacy, misc.verb, re
from misc.refer import REFER

class GtLoader(Loader):

    def __init__(self, data_json, opt, tree_pth=None):
        # parent loader instance, see loader.py
        Loader.__init__(self, data_json)
        # print('data_json, visual_feats_dir, opt, tree_pth', data_json, visual_feats_dir, opt, tree_pth)
        self.opt = opt
        self.batch_size = opt.batch_size
        # self.batch_size = opt["batch_size"]
        # self.vis_dim = 2048 + 512 + 512
        self.vis_dim = 4096

        # img_iterators for each split
        self.split_ix = {}
        self.iterators = {}

        # load ann feats
        # print('loading visual feats ...')
        # pickle.load = partial(pickle.load, encoding="latin1")
        # pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
        # self.visual_feats = torch.load(visual_feats_dir,
            # map_location=lambda storage, loc: storage, pickle_module=pickle)

        # for ke in self.visual_feats:
            # print(ke, np.shape(self.visual_feats[ke]))
        # print('visual feats loaded')
        # if tree_pth:
            # # self.trees = torch.load(tree_pth, 'r')
            # ## changed
            # with open(tree_pth, 'rb') as fp:
                # self.trees = pickle.load(fp)
        # else:
            # self.trees = None

        # for image_id, image in self.Images.items():
            # split = self.Refs[image['ref_ids'][0]]['split']

            # if split not in self.split_ix:
                # self.split_ix[split] = []
                # self.iterators[split] = 0

            # # add sentences to each subsets
            # sent_ids = []
            # for ref_id in self.Images[image_id]['ref_ids']:
                # sent_ids += self.Refs[ref_id]['sent_ids']
            # self.split_ix[split].append(image_id)

        # for k, v in self.split_ix.items():
            # print('assigned %d images to split %s' % (len(v), k))

        print("Loading Yolov3 model")
        self.cfgfile = "yolov3/cfg/yolo_v3.cfg"
        self.weightfile = "yolov3/yolov3.weights"
        self.yolo = Darknet(self.cfgfile)
        # # self.yolo.print_network()
        self.yolo.load_weights(self.weightfile)
        print('Loading weights from %s... Done!' % (self.weightfile))
        self.yolo.cuda()

        print("Loading ROI Pool")
        self.roi_pool = TorchROIPool(2, 1.0)

    # shuffle split
    def shuffle(self, split):
        random.shuffle(self.split_ix[split])

    # reset iterator
    def reset_iterator(self, split):
        self.iterators[split] = 0

    # get reverse coordinates for yolov3 model
    def get_coordinates(self, boxes, im_w, im_h):

        # im_w, im_h = float(im_w), float(im_h)
        # net_w, net_h = float(net_w), float(net_h)
        # if net_w/im_w < net_h/im_h:
            # new_w = net_w
            # new_h = (im_h * net_w)/im_w
        # else:
            # new_w = (im_w * net_h)/im_h
            # new_h = net_h

        # xo, xs = (net_w - new_w)/(2*net_w), net_w/new_w
        # yo, ys = (net_h - new_h)/(2*net_h), net_h/new_h

        b = boxes
        # from x1, y1, w, h to x1, y1, x2, y2
        # b[2] = b[0] + b[2]
        # b[3] = b[1] + b[3]

        target_size = 13
        x_scale = target_size / im_w
        y_scale = target_size / im_h

        b[0] = int(np.round(b[0] * x_scale))
        b[1] = int(np.round(b[1] * y_scale))
        b[2] = int(np.round(b[2] * x_scale))
        b[3] = int(np.round(b[3] * y_scale))
        # from x1, y1, x2, y2 to correct yolo coord
        # b[0] = (b[2] + b[0]) / (im_w * 2.0)
        # b[2] = (b[2] - b[0]) / im_w
        # b[1] = (b[3] + b[1]) / (im_h * 2.0)
        # b[3] = (b[3] - b[1]) / im_h
        # from correct yolo coord to mode layers coord
        # b[0] = b[0] / xs + xo
        # b[1] = b[1] / ys + yo
        # b[2] /= xs
        # b[3] /= ys
        return b

    def doc_prune(self, doc):
        new = []
        old = []

        for token in doc:
            if token.pos_ in ['DET', 'SYM', 'SPACE']:
                continue
            else:
                old.append(token.lower_)
                new.append(token.lower_)

        new_sent = " ".join(new)
        new_sent = re.sub(' +',' ', new_sent)

        old_sent = " ".join(old)
        old_sent = re.sub(' +',' ', old_sent)

        return new_sent.rstrip(), new, old_sent.rstrip(), old

    def doc_to_tree(self, doc):
        def traversal(node, sent_list=[]):
            txt = {}
            txt[(node.i, node.text, node.pos_, node.tag_, node.dep_)] = []

            sent_list.append((node.i, node.text, node.pos_, node.tag_, node.dep_))
            for child in node.children:
                tree, sent_list = traversal(child, sent_list)
                txt[(node.i, node.text, node.pos_, node.tag_, node.dep_)].append(tree)
            return txt, sent_list

        # traversal root's every children
        tree = {}
        try:
            root = next(doc.sents).root
        except:
            tree[(0, "UNK", "UNK", "UNK", "UNK")] = []
            print("a NULL sentence: {}".format(doc))
            return tree, [(0, "UNK", "UNK", "UNK", "UNK")]  # cause there a blank sentence in data

        # start traversal
        tree[(root.i, root.text, root.pos_, root.tag_, root.dep_)] = []
        sent_list = []
        sent_list.append((root.i, root.text, root.pos_, root.tag_, root.dep_))
        for child in root.children:
            t, sent_list = traversal(child, sent_list)
            tree[(root.i, root.text, root.pos_, root.tag_, root.dep_)].append(t)

        return tree, sent_list

    def load_sent(self, sent):
        nlp = spacy.load('en_core_web_sm', disable=['ner'])
        print(nlp)

        new_sents = []
        for idx, doc in enumerate(nlp.pipe([sent], disable = ['parser'])):
            new_sent, new, old_sent, old = self.doc_prune(doc)
            new_sents.append(new_sent)

        tree_list = []
        token_list = []
        tag_list = []
        dep_list = []
        for idx, doc in enumerate(nlp.pipe(new_sents)):
            t, s = self.doc_to_tree(doc)
            tree_list.append(t)
            s = sorted(s, key = lambda x: x[0])
            token = [x[1] for x in s]
            tag = [x[3] for x in s]
            dep = [x[4] for x in s]
            token_list.append(token)
            tag_list.append(tag)
            dep_list.append(dep)

        tree = {'tree': tree_list[0], 'tokens': token_list[0], 'tags': tag_list[0], 'deps': dep_list[0]}
        return tree

    # get one of data
    def get_data(self, img_path, sent):
        # options
        # split_ix = self.split_ix[split]
        # max_index = len(split_ix) - 1  # don't forget to -1
        wrapped = False

        # get information about this batch image
        # image_id = split_ix[self.iterators[split]]
        # ann_ids = self.Images[image_id]['ann_ids']

        # fetch sentences
        sent_ids = []
        gt_ixs = []
        sents = []
        trees = []
        vis = []

        # print("\n\nimage_id", image_id, self.Images[image_id])
        # img = Image.open(img_path).convert('RGB')
        img = Image.open(img_path)
        sized = letterbox_image(img, self.yolo.width, self.yolo.height)
        use_cuda = True

        boxes, layer74 = do_detect(self.yolo,  sized, 0.1, 0.4, use_cuda)
        print("Layer 74", np.shape(layer74))
        correct_yolo_boxes(boxes, img.width, img.height, self.yolo.width, self.yolo.height)
        print("boxes ", boxes)

        width = img.width
        height = img.height
        for box in boxes:
            box[0], box[1], box[2], box[3] = (box[0] - box[2]/2.0) * width, (box[1] - box[3]/2.0) * height, (box[0] + box[2]/2.0) * width, (box[1] + box[3]/2.0) * height
            # print(x1, y1, x2, y2)
            # box = [x1, y1, x2, y2]
        print(boxes)

        for box in boxes:
            box = self.get_coordinates(box, img.width, img.height)
            if box[0] == box[2]:
                if box[2] != 13:
                    box[2] += 1
                else:
                    box[0] -= 1
            if box[1] == box[3]:
                if box[3] != 13:
                    box[3] += 1
                else:
                    box[1] -= 1
        print(boxes)
        ann_boxes = torch.Tensor(boxes)
        # print(ann_boxes, "********")
        v = self.roi_pool(layer74, ann_boxes)
        # print(np.shape(v))
        v = v.reshape(1, v.size(0), v.size(1) * v.size(2) * v.size(3))
        print(np.shape(v))
        vis = v.detach()
        # convert to tensor
        vis = torch.from_numpy(np.asarray(vis)).cuda()
        # gt_ixs = torch.from_numpy(np.asarray(gt_ixs)).cuda()

        # update iter status
        # if self.iterators[split] + 1 > max_index:
            # self.iterators[split] = 0
            # wrapped = True
        # else:
            # self.iterators[split] += 1

        tree = self.load_sent(sent)

        # # return
        data = {}
        data['vis'] = vis.float()  # (num_bboxs, fc7_dim)
        data['sents'] = [tree]
        data['trees'] = [tree['tree']]

        # data['gts'] = gt_ixs.long()   # (num_sents, )
        # data['sent_ids'] = sent_ids
        # data['ann_ids'] = ann_ids
        # data['bounds'] = {'it_pos_now': self.iterators[split],
                          # 'it_max': max_index,
                          # 'wrapped': wrapped}
        print(np.shape(data['vis']), data.keys(), data)
        return data


# Assuming directory structure:
#    ./dataset/vcc2016/wav/Training Set/SF1/100001.wav

import os
from os.path import join

import librosa
import numpy as np
import pyworld as pw
import tensorflow as tf

args = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string('dir_to_wav', './data_multi/wav', 'Dir to *.wav')
tf.app.flags .DEFINE_string('dir_to_bin', './data_multi/bin', 'Dir to output *.bin')
tf.app.flags.DEFINE_integer('fs', 16000, 'Global sampling frequency')
tf.app.flags.DEFINE_float('f0_ceil', 500, 'Global f0 ceiling')

FFT_SIZE = 1024
SP_DIM = FFT_SIZE // 2 + 1
FEAT_DIM = SP_DIM + SP_DIM + 1 + 1 + 1  # [sp, ap, f0, en, s]
RECORD_BYTES = FEAT_DIM * 4  # all features saved in `float32`

EPSILON = 1e-10

def wav2pw(x, fs=16000, fft_size=FFT_SIZE):
    ''' Extract WORLD feature from waveform '''
    _f0, t = pw.dio(x, fs, f0_ceil=args.f0_ceil)            # raw pitch extractor
    f0 = pw.stonemask(x, _f0, t, fs)  # pitch refinement
    sp = pw.cheaptrick(x, f0, t, fs, fft_size=fft_size)
    ap = pw.d4c(x, f0, t, fs, fft_size=fft_size) # extract aperiodicity
    return {
        'f0': f0,
        'sp': sp,
        'ap': ap,
    }

def list_dir(path):
    ''' retrieve the 'short name' of the dirs '''
    return sorted([f for f in os.listdir(path) if os.path.isdir(join(path, f))])

def list_full_filenames(path):
    ''' return a generator of full filenames '''
    return (
        join(path, f)
            for f in os.listdir(path)
                if not os.path.isdir(join(path, f)))

def extract(filename, fft_size=FFT_SIZE, dtype=np.float32):
    ''' Basic (WORLD) feature extraction '''
    x, _ = librosa.load(filename, sr=args.fs, mono=True, dtype=np.float64)
    features = wav2pw(x, args.fs, fft_size=fft_size)
    ap = features['ap']
    f0 = features['f0'].reshape([-1, 1])
    sp = features['sp']
    en = np.sum(sp + EPSILON, axis=1, keepdims=True)
    sp = np.log10(sp / en)
    return np.concatenate([sp, ap, f0, en], axis=1).astype(dtype)

def extract_and_save_bin_to(dir_to_wav, dir_to_bin, speakers):
    '''
    NOTE: the directory structure must be [args.dir_to_wav]/[Set]/[speakers]
    '''
    counter = 1
    N = len(tf.gfile.Glob(join(dir_to_wav, '*', '*', '*.wav')))
    for d in list_dir(dir_to_wav):  # ['Training Set', 'Testing Set']
        path = join(dir_to_wav, d)
        for s in list_dir(path):  # ['SF1', ..., 'TM3']
            path = join(dir_to_wav, d, s)
            output_dir = join(dir_to_bin, d, s)
            tf.gfile.MakeDirs(output_dir)
            for f in list_full_filenames(path):  # ['10001.wav', ...]
                print('\rFile {}/{}: {:50}'.format(counter, N, f), end='')
                features = extract(f)
                labels = speakers.index(s) * np.ones(
                    [features.shape[0], 1],
                    np.float32,
                )
                b = os.path.splitext(f)[0]
                _, b = os.path.split(b)
                features = np.concatenate([features, labels], 1)
                with open(join(output_dir, '{}.bin'.format(b)), 'wb') as fp:
                    fp.write(features.tostring())
                counter += 1
        print()

class Tanhize(object):
    ''' Normalizing `x` to [-1, 1] '''
    def __init__(self, xmin, xmax):
        self.xmin = xmin
        self.xmax = xmax
        self.xscale = xmax - xmin

    def forward_process(self, x):
        x = (x - self.xmin) / self.xscale
        return tf.clip_by_value(x, 0., 1.) * 2. - 1.

    def backward_process(self, x):
        return (x * .5 + .5) * self.xscale + self.xmin

def read(
    file_pattern,
    batch_size,
    record_bytes=RECORD_BYTES,
    capacity=2048,
    min_after_dequeue=1536,
    num_threads=8,
    data_format='NCHW',
    normalizer=None,
    ):
    '''
    Read only `sp` and `speaker`
    Return:
        `feature`: [b, c]
        `speaker`: [b,]
    '''
    with tf.device('cpu'):
        with tf.name_scope('InputSpectralFrame'):
            files = tf.gfile.Glob(file_pattern)
            filename_queue = tf.train.string_input_producer(files)

            reader = tf.FixedLengthRecordReader(record_bytes)
            _, value = reader.read(filename_queue)
            value = tf.decode_raw(value, tf.float32)

            value = tf.reshape(value, [FEAT_DIM,])
            feature = value[:SP_DIM]   # NCHW format

            if normalizer is not None:
                feature = normalizer.forward_process(feature)

            if data_format == 'NCHW':
                feature = tf.reshape(feature, [1, SP_DIM, 1])
            elif data_format == 'NHWC':
                feature = tf.reshape(feature, [SP_DIM, 1, 1])
            else:
                pass

            speaker = tf.cast(value[-1], tf.int64)
            #zk
            f0 = tf.cast(value[-3], tf.int64)

            return tf.train.shuffle_batch(
                [feature, speaker, f0],
                batch_size,
                capacity=capacity,
                min_after_dequeue=min_after_dequeue,
                num_threads=num_threads,
                # enqueue_many=True,
            )

def read_whole_features(file_pattern, num_epochs=1):
    '''
    Return
        `feature`: `dict` whose keys are `sp`, `ap`, `f0`, `en`, `speaker`
    '''
    with tf.device('cpu'):
        with tf.name_scope('InputPipline'):
            files = tf.gfile.Glob(file_pattern)
            print('{} files found'.format(len(files)))
            filename_queue = tf.train.string_input_producer(files, num_epochs=num_epochs)
            reader = tf.WholeFileReader()
            key, value = reader.read(filename_queue)
            print("Processing {}".format(key), flush=True)
            value = tf.decode_raw(value, tf.float32)
            value = tf.reshape(value, [-1, FEAT_DIM])
            return {
                'sp': value[:, :SP_DIM],
                'ap': value[:, SP_DIM : 2*SP_DIM],
                'f0': value[:, SP_DIM * 2],
                'en': value[:, SP_DIM * 2 + 1],
                'speaker': tf.cast(value[:, SP_DIM * 2 + 2], tf.int64),
                'filename': key,
            }

def pw2wav(features, feat_dim=513, fs=16000):
    ''' NOTE: Use `order='C'` to ensure Cython compatibility '''
    en = np.reshape(features['en'], [-1, 1])
    sp = np.power(10., features['sp'])
    sp = en * sp
    if isinstance(features, dict):
        return pw.synthesize(
            features['f0'].astype(np.float64).copy(order='C'),
            sp.astype(np.float64).copy(order='C'),
            features['ap'].astype(np.float64).copy(order='C'),
            fs,
        )
    features = features.astype(np.float64)
    sp = features[:, :feat_dim]
    ap = features[:, feat_dim:feat_dim*2]
    f0 = features[:, feat_dim*2]
    en = features[:, feat_dim*2 + 1]
    en = np.reshape(en, [-1, 1])
    sp = np.power(10., sp)
    sp = en * sp
    return pw.synthesize(
        f0.copy(order='C'),
        sp.copy(order='C'),
        ap.copy(order='C'),
        fs
    )

def make_speaker_tsv(path):
    speakers = []
    for d in list_dir(path):
        speakers += list_dir(join(path, d))
    speakers = sorted(set(speakers))
    with open('./etc/speakers.tsv', 'w') as fp:
        for s in speakers:
            fp.write('{}\n'.format(s))
    return speakers

if __name__ == '__main__':
    speakers = make_speaker_tsv(args.dir_to_wav)
    extract_and_save_bin_to(
        args.dir_to_wav,
        args.dir_to_bin,
        speakers,
    )


import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='1'
from os import listdir

import sys
import time
import tools.ops

import numpy as np

import tensorflow as tf
import scipy.misc as sm

from models.mfb_net_cross import *
from tools.utilities import *
from tools.ops import *

flags = tf.app.flags
flags.DEFINE_integer('batch_size', 1, 'Batch size.')
flags.DEFINE_integer('num_epochs', 1, 'Number of epochs.')	# ~13 min per epoch
flags.DEFINE_integer('num_gpus', 1, 'Number of GPUs.')
flags.DEFINE_integer('seq_length', 16, 'Length of each video clip.')
flags.DEFINE_integer('height', 128, 'Height of video frame.')
flags.DEFINE_integer('width', 128, 'Width of video frame.')
flags.DEFINE_integer('channel', 3, 'Number of channels for each frame.')

FLAGS = flags.FLAGS

prefix          = 'mfb_cross'
model_save_dir  = './ckpt/' + prefix
loss_save_dir   = './loss'
test_save_dir   = './test/' + prefix
test_list_path  = './dataset/testlist.txt'
dataset_path    = './dataset/UCF-101-tf-records'

use_pretrained_model = True
save_predictions     = True

def run_testing():

	# Create model directory
	if not os.path.exists(model_save_dir):
		os.makedirs(model_save_dir)
	model_filename = "./mfb_baseline_ucf24.model"

	tower_ffg_losses, tower_fbg_losses, tower_lfg_losses, tower_feat_losses = [], [], [], []
	tower_ffg_m_losses, tower_fbg_m_losses, tower_lfg_m_losses = [], [], []

	global_step = tf.get_variable(
                'global_step',
                [],
                initializer=tf.constant_initializer(0),
                trainable=False
                )
	starter_learning_rate = 1e-4
	learning_rate = tf.train.exponential_decay(starter_learning_rate, global_step,
                                           100000000, 0.5, staircase=True)
	opt = tf.train.AdamOptimizer(learning_rate)

	# Create a session for running Ops on the Graph.
	config = tf.ConfigProto(allow_soft_placement=True)
	sess = tf.Session(config=config)
	coord = tf.train.Coordinator()
	threads = None

	test_list_file = open(test_list_path, 'r')
	test_list = test_list_file.read().splitlines()
	for i, line in enumerate(test_list):
		test_list[i] = os.path.join(dataset_path, test_list[i])

	assert(len(test_list) % FLAGS.num_gpus == 0)
	num_for_each_gpu = len(test_list) // FLAGS.num_gpus

	clips_list, img_masks_list, loss_masks_list = [], [], []
	with sess.as_default():
		for i in range(FLAGS.num_gpus):
			clips, img_masks, loss_masks = input_pipeline(test_list[i*num_for_each_gpu:(i+1)*num_for_each_gpu], \
											FLAGS.batch_size, read_threads=1, num_epochs=FLAGS.num_epochs, is_training=False)
			clips_list.append(clips)
			img_masks_list.append(img_masks)
			loss_masks_list.append(loss_masks)

	mfb_list = []
	with tf.variable_scope('vars') as var_scope:
		for gpu_index in range(FLAGS.num_gpus):
			with tf.device('/gpu:%d' % (gpu_index)):
				with tf.name_scope('%s_%d' % ('tower', gpu_index)) as scope:

					# construct model
					mfb = mfb_net(clips_list[gpu_index], FLAGS.height, FLAGS.width, FLAGS.seq_length, \
									FLAGS.channel, FLAGS.batch_size, is_training=False)
					mfb_list.append(mfb)
					_, first_fg_loss, first_bg_loss, last_fg_loss, feat_loss, _ = \
						tower_loss(scope, mfb, clips_list[gpu_index], img_masks_list[gpu_index], loss_masks_list[gpu_index])

					var_scope.reuse_variables()

					tower_ffg_losses.append(first_fg_loss)
					tower_fbg_losses.append(first_bg_loss)
					tower_lfg_losses.append(last_fg_loss)
					tower_feat_losses.append(feat_loss)

	# concatenate the losses of all towers
	ffg_loss_op  = tf.reduce_mean(tower_ffg_losses)
	fbg_loss_op  = tf.reduce_mean(tower_fbg_losses)
	lfg_loss_op  = tf.reduce_mean(tower_lfg_losses)
	feat_loss_op = tf.reduce_mean(tower_feat_losses)

	# saver for saving checkpoints
	saver = tf.train.Saver()
	init = tf.initialize_all_variables()

	sess.run(init)
	if not os.path.exists(model_save_dir):
		os.makedirs(model_save_dir)
	if use_pretrained_model:
		print('[*] Loading checkpoint ...')
		model = tf.train.latest_checkpoint(model_save_dir)
		if model is not None:
			saver.restore(sess, model)
			print('[*] Loading success: %s!'%model)
		else:
			print('[*] Loading failed ...')

	# Create loss output folder
	if not os.path.exists(loss_save_dir):
		os.makedirs(loss_save_dir)
	loss_file = open(os.path.join(loss_save_dir, prefix+'_test.txt'), 'a+')

	# Create test output folder
	if not os.path.exists(test_save_dir):
		os.makedirs(test_save_dir)

	# start queue runner
	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)

	ffg_loss_list, fbg_loss_list, lfg_loss_list, feat_loss_list = [], [], [], []
	try:
		with sess.as_default():
			print('\n\n\n*********** start testing ***********\n\n\n')
			step = global_step.eval()
			print('[step = %d]'%step)
			cnt = 0
			while not coord.should_stop():
				# Run training steps or whatever
				mfb = mfb_list[0]
				ffg, fbg, lfg, gt_ffg, gt_fbg, gt_lfg, ffg_loss, fbg_loss, lfg_loss, feat_loss = sess.run([
					mfb.first_fg_rec, mfb.first_bg_rec, mfb.last_fg_rec, \
					mfb.gt_ffg, mfb.gt_fbg, mfb.gt_lfg, ffg_loss_op, fbg_loss_op, lfg_loss_op, feat_loss_op])

				ffg, fbg, lfg, gt_ffg, gt_fbg, gt_lfg = \
					ffg[0], fbg[0], lfg[0], gt_ffg[0], gt_fbg[0], gt_lfg[0]

				ffg, fbg, lfg = (ffg+1)/2*255.0, (fbg+1)/2*255.0, (lfg+1)/2*255.0
				gt_ffg, gt_fbg, gt_lfg = (gt_ffg+1)/2*255.0, (gt_fbg+1)/2*255.0, (gt_lfg+1)/2*255.0

				img = gen_pred_img(ffg, fbg, lfg)
				gt  = gen_pred_img(gt_ffg, gt_fbg, gt_lfg)
				save_img = np.concatenate((img, gt))
				sm.imsave(os.path.join(test_save_dir, '%05d.jpg'%cnt), save_img)

				ffg_loss_list.append(ffg_loss)
				fbg_loss_list.append(fbg_loss)
				lfg_loss_list.append(lfg_loss)
				feat_loss_list.append(feat_loss)

				line = '%05d: ffg_loss=%.8f, fbg_loss=%.8f, lfg_loss=%.8f, feat_loss=%.8f' \
						%(cnt, ffg_loss, fbg_loss, lfg_loss, feat_loss)
				loss_file.write(line + '\n')
				loss_file.flush()
				print(line)
				cnt += 1

	except tf.errors.OutOfRangeError:
		print('Done training -- epoch limit reached')
	finally:
		# When done, ask the threads to stop.
		coord.request_stop()

	# Wait for threads to finish.
	coord.join(threads)
	sess.close()

	mean_ffg = np.mean(np.asarray(ffg_loss_list))
	mean_fbg = np.mean(np.asarray(fbg_loss_list))
	mean_lfg = np.mean(np.asarray(lfg_loss_list))
	mean_feat = np.mean(np.asarray(feat_loss_list))

	line = '[step=%d] ffg_loss=%.8f, fbg_loss=%.8f, lfg_loss=%.8f, feat_loss=%.8f' \
			%(step, mean_ffg, mean_fbg, mean_lfg, mean_feat)
	print(line)
	loss_file.write(line + '\n')

def main(_):
	run_testing()

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


import torch
import argparse
import os
import cv2
import numpy as np
import torch.optim as optim
from multiprocessing import cpu_count
from torch.utils.data import DataLoader
from modeling.anime_gan import Generator
from modeling.anime_gan import Discriminator
from modeling.losses import AnimeGanLoss
from modeling.losses import LossSummary
from utils.common import load_checkpoint
from utils.common import save_checkpoint
from utils.common import set_lr
from utils.common import initialize_weights
from utils.image_processing import denormalize_input
from dataset import AnimeDataSet
from tqdm import tqdm

gaussian_mean = torch.tensor(0.0)
gaussian_std = torch.tensor(0.1)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='Hayao')
    parser.add_argument('--data-dir', type=str, default='/content/dataset')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--init-epochs', type=int, default=5)
    parser.add_argument('--batch-size', type=int, default=6)
    parser.add_argument('--checkpoint-dir', type=str, default='/content/checkpoints')
    parser.add_argument('--save-image-dir', type=str, default='/content/images')
    parser.add_argument('--gan-loss', type=str, default='lsgan', help='lsgan / hinge / bce')
    parser.add_argument('--resume', type=str, default='False')
    parser.add_argument('--use_sn', action='store_true')
    parser.add_argument('--save-interval', type=int, default=1)
    parser.add_argument('--debug-samples', type=int, default=0)
    parser.add_argument('--lr-g', type=float, default=2e-4)
    parser.add_argument('--lr-d', type=float, default=4e-4)
    parser.add_argument('--init-lr', type=float, default=1e-3)
    parser.add_argument('--wadvg', type=float, default=10.0, help='Adversarial loss weight for G')
    parser.add_argument('--wadvd', type=float, default=10.0, help='Adversarial loss weight for D')
    parser.add_argument('--wcon', type=float, default=1.5, help='Content loss weight')
    parser.add_argument('--wgra', type=float, default=3.0, help='Gram loss weight')
    parser.add_argument('--wcol', type=float, default=30.0, help='Color loss weight')
    parser.add_argument('--d-layers', type=int, default=3, help='Discriminator conv layers')
    parser.add_argument('--d-noise', action='store_true')

    return parser.parse_args()

def collate_fn(batch):
    img, anime, anime_gray, anime_smt_gray = zip(*batch)
    return (
        torch.stack(img, 0),
        torch.stack(anime, 0),
        torch.stack(anime_gray, 0),
        torch.stack(anime_smt_gray, 0),
    )

def check_params(args):
    data_path = os.path.join(args.data_dir, args.dataset)
    if not os.path.exists(data_path):
        raise FileNotFoundError(f'Dataset not found {data_path}')

    if not os.path.exists(args.save_image_dir):
        print(f'* {args.save_image_dir} does not exist, creating...')
        os.makedirs(args.save_image_dir)

    if not os.path.exists(args.checkpoint_dir):
        print(f'* {args.checkpoint_dir} does not exist, creating...')
        os.makedirs(args.checkpoint_dir)

    assert args.gan_loss in {'lsgan', 'hinge', 'bce'}, f'{args.gan_loss} is not supported'

def save_samples(generator, loader, args, max_imgs=2, subname='gen'):
    '''
    Generate and save images
    '''
    generator.eval()

    max_iter = (max_imgs // args.batch_size) + 1
    fake_imgs = []

    for i, (img, *_) in enumerate(loader):
        with torch.no_grad():
            fake_img = generator(img.cuda())
            fake_img = fake_img.detach().cpu().numpy()
            # Channel first -> channel last
            fake_img  = fake_img.transpose(0, 2, 3, 1)
            fake_imgs.append(denormalize_input(fake_img, dtype=np.int16))

        if i + 1 == max_iter:
            break

    fake_imgs = np.concatenate(fake_imgs, axis=0)

    for i, img in enumerate(fake_imgs):
        save_path = os.path.join(args.save_image_dir, f'{subname}_{i}.jpg')
        cv2.imwrite(save_path, img[..., ::-1])

def gaussian_noise():
    return torch.normal(gaussian_mean, gaussian_std)

def main(args):
    check_params(args)

    print("Init models...")

    G = Generator(args.dataset).cuda()
    D = Discriminator(args).cuda()

    loss_tracker = LossSummary()

    loss_fn = AnimeGanLoss(args)

    # Create DataLoader
    data_loader = DataLoader(
        AnimeDataSet(args),
        batch_size=args.batch_size,
        num_workers=cpu_count(),
        pin_memory=True,
        shuffle=True,
        collate_fn=collate_fn,
    )

    optimizer_g = optim.Adam(G.parameters(), lr=args.lr_g, betas=(0.5, 0.999))
    optimizer_d = optim.Adam(D.parameters(), lr=args.lr_d, betas=(0.5, 0.999))

    start_e = 0
    if args.resume == 'GD':
        # Load G and D
        try:
            start_e = load_checkpoint(G, args.checkpoint_dir)
            print("G weight loaded")
            load_checkpoint(D, args.checkpoint_dir)
            print("D weight loaded")
        except Exception as e:
            print('Could not load checkpoint, train from scratch', e)
    elif args.resume == 'G':
        # Load G only
        try:
            start_e = load_checkpoint(G, args.checkpoint_dir, posfix='_init')
        except Exception as e:
            print('Could not load G init checkpoint, train from scratch', e)

    for e in range(start_e, args.epochs):
        print(f"Epoch {e}/{args.epochs}")
        bar = tqdm(data_loader)
        G.train()

        init_losses = []

        if e < args.init_epochs:
            # Train with content loss only
            set_lr(optimizer_g, args.init_lr)
            for img, *_ in bar:
                img = img.cuda()

                optimizer_g.zero_grad()

                fake_img = G(img)
                loss = loss_fn.content_loss_vgg(img, fake_img)
                loss.backward()
                optimizer_g.step()

                init_losses.append(loss.cpu().detach().numpy())
                avg_content_loss = sum(init_losses) / len(init_losses)
                bar.set_description(f'[Init Training G] content loss: {avg_content_loss:2f}')

            set_lr(optimizer_g, args.lr_g)
            save_checkpoint(G, optimizer_g, e, args, posfix='_init')
            save_samples(G, data_loader, args, subname='initg')
            continue

        loss_tracker.reset()
        for img, anime, anime_gray, anime_smt_gray in bar:
            # To cuda
            img = img.cuda()
            anime = anime.cuda()
            anime_gray = anime_gray.cuda()
            anime_smt_gray = anime_smt_gray.cuda()

            # ---------------- TRAIN D ---------------- #
            optimizer_d.zero_grad()
            fake_img = G(img).detach()

            # Add some Gaussian noise to images before feeding to D
            if args.d_noise:
                fake_img += gaussian_noise()
                anime += gaussian_noise()
                anime_gray += gaussian_noise()
                anime_smt_gray += gaussian_noise()

            fake_d = D(fake_img)
            real_anime_d = D(anime)
            real_anime_gray_d = D(anime_gray)
            real_anime_smt_gray_d = D(anime_smt_gray)

            loss_d = loss_fn.compute_loss_D(
                fake_d, real_anime_d, real_anime_gray_d, real_anime_smt_gray_d)

            loss_d.backward()
            optimizer_d.step()

            loss_tracker.update_loss_D(loss_d)

            # ---------------- TRAIN G ---------------- #
            optimizer_g.zero_grad()

            fake_img = G(img)
            fake_d = D(fake_img)

            adv_loss, con_loss, gra_loss, col_loss = loss_fn.compute_loss_G(
                fake_img, img, fake_d, anime_gray)

            loss_g = adv_loss + con_loss + gra_loss + col_loss

            loss_g.backward()
            optimizer_g.step()

            loss_tracker.update_loss_G(adv_loss, gra_loss, col_loss, con_loss)

            avg_adv, avg_gram, avg_color, avg_content = loss_tracker.avg_loss_G()
            avg_adv_d = loss_tracker.avg_loss_D()
            bar.set_description(f'loss G: adv {avg_adv:2f} con {avg_content:2f} gram {avg_gram:2f} color {avg_color:2f} / loss D: {avg_adv_d:2f}')

        if e % args.save_interval == 0:
            save_checkpoint(G, optimizer_g, e, args)
            save_checkpoint(D, optimizer_d, e, args)
            save_samples(G, data_loader, args)

if __name__ == '__main__':
    args = parse_args()

    print("# ==== Train Config ==== #")
    for arg in vars(args):
        print(arg, getattr(args, arg))
    print("==========================")

    main(args)


# -*- coding: utf-8 -*-
# -*- author: JeremySun -*-
# -*- dating: 20/11/12 -*-

import os
import csv
import copy
import random
import warnings
import numpy as np
import pandas as pd
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.utils.data as Data
from torch.nn.utils import clip_grad_norm_
from sklearn.utils import shuffle as reset
from tqdm import tqdm
import transformers
from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModel
from transformers import get_linear_schedule_with_warmup
from transformers import AdamW

warnings.filterwarnings("ignore")
transformers.logging.set_verbosity_error()

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.enabled = False

classes2idx = {'财经': 0, '房产': 1, '家居': 2, '教育': 3, '科技': 4, '时尚': 5, '时政': 6, '游戏': 7, '娱乐': 8, '体育': 9}
idx2classes = {0: '财经', 1: '房产', 2: '家居', 3: '教育', 4: '科技', 5: '时尚', 6: '时政', 7: '游戏', 8: '娱乐', 9: '体育'}

MODEL_NAME = '../pretrained_model/roberta_base_myself'
OUTPUT_MODEL = '../models/roberta_base_myself.pth'

SEED = 1224
TRAIN = True
BATCH_SIZE = 4
EPOCHS = 5
HIDDEN_DIM = 256
DROPOUT = 0.3
BIDIRECTIONAL = True

best_score = 0.

class CustomDataset(Data.Dataset):
    def __init__(self, data, max_len, with_labels=True, model_name=MODEL_NAME):
        self.data = data
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.max_len = max_len
        self.with_labels = with_labels

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

    def __getitem__(self, index):
        sent = str(self.data.loc[index, 'content'])
        encoded_pair = self.tokenizer(sent, padding='max_length', truncation=True, max_length=self.max_len, return_tensors='pt')
        token_ids = encoded_pair['input_ids'].squeeze(0)  # tensor of token ids
        attn_masks = encoded_pair['attention_mask'].squeeze(0)  # binary tensor with 0 for padded values and 1 for the other values
        token_type_ids = encoded_pair['token_type_ids'].squeeze(0)  # binary tensor with 0 for the 1st sentence tokens and 1 for the 2nd sentence tokens

        if self.with_labels:
            label = self.data.loc[index, 'class_label']
            return token_ids, attn_masks, token_type_ids, label
        else:
            return token_ids, attn_masks, token_type_ids

class CLSModel(nn.Module):
    def __init__(self, model_name, hidden_dim, num_classes, dropout, bidirectional=BIDIRECTIONAL, freeze_bert=False):
        super(CLSModel, self).__init__()
        self.bert = BertModel.from_pretrained(model_name, output_hidden_states=True, return_dict=True)
        embedding_dim = self.bert.config.to_dict()['hidden_size']
        self.rnn = nn.GRU(input_size=embedding_dim * 4, hidden_size=hidden_dim, batch_first=True, dropout=dropout, bidirectional=bidirectional)  # [batch_size, seq_len, input_size]
        # self.linear = nn.Linear(in_features=hidden_dim * 2, out_features=num_classes)
        self.linear = nn.Sequential(
            nn.Linear(in_features=hidden_dim * 2, out_features=512),
            nn.GELU(),
            nn.Dropout(p=dropout),
            nn.Linear(in_features=512, out_features=num_classes)
        )
        self.dropout = nn.Dropout(p=dropout)

        if freeze_bert:
            for p in self.bert.parameters():
                p.requires_grad = False

    def forward(self, input_ids, attn_masks, token_type_ids):
        embedding = self.bert(input_ids, token_type_ids=token_type_ids, attention_mask=attn_masks)  # [batch_size, seq_len, hidden_size]
        hidden_states = torch.cat(tuple([embedding.hidden_states[i] for i in [-1, -2, -3, -4]]), dim=-1)  # [batch_size, seq_len, hidden_size * 4]
        first_hidden_states = hidden_states[:, 0, :]  # [batch_size, hidden_size * 4]
        first_hidden_states = first_hidden_states.unsqueeze(dim=1)  # [batch_size, 1, hidden_size * 4]
        # first_hidden_states.permute(1, 0, 2)  # [1, batch_size, hidden_size * 4]
        _, hidden = self.rnn(first_hidden_states)  # [num_layers * num_directions, batch_size, hidden_dim]
        hidden = self.dropout(torch.cat((hidden[-1, :, :], hidden[-2, :, :]), dim=1) if self.rnn.bidirectional else hidden[-1, :, :])  # torch.cat([batch_size, hidden_dim], [batch_size, hidden_dim]) --> [batch_size, hidden_dim * 2]
        logits = self.linear(hidden)
        return logits

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)

def process_data(filename, classes2idx, with_labels=True):
    data = pd.read_csv(filename, encoding='utf-8')
    print(len(data))
    if with_labels:
        data = data.replace({'class_label': classes2idx})
    return data

def flat_accuracy(preds, labels):
    pred_flat = np.argmax(preds, axis=1).flatten()
    labels_flat = labels.flatten()
    return np.sum(pred_flat == labels_flat) / len(labels_flat)

def save(model, optimizer):
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict()
    }, OUTPUT_MODEL)
    print('\nThe best model has been saved.')

def train_test_split(data_df, test_size=0.2, shuffle=True, random_state=None):
    if shuffle:
        data_df = reset(data_df, random_state=random_state)

    train = data_df[int(len(data_df) * test_size):].reset_index(drop=True)
    test = data_df[: int(len(data_df) * test_size)].reset_index(drop=True)

    return train, test

def train_eval(model, criterion, optimizer, train_loader, val_loader, epochs):
    # checkpoint = torch.load(output_model, map_location='cpu')
    # model.load_state_dict(checkpoint['model_state_dict'])

    model.to(device)
    # total_steps = len(train_loader) * epochs

    print('========== Training ==========')

    for epoch in range(epochs):
        model.train()
        print(f'Epoch {epoch + 1}')
        for i, batch in enumerate(tqdm(train_loader)):
            batch = tuple(t.to(device) for t in batch)
            logits = model(batch[0], batch[1], batch[2])
            loss = criterion(logits, batch[3])
            # if i % 10 == 0:
            # print(i, loss.item())

            optimizer.zero_grad()
            loss.backward()
            clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()
            # scheduler.step()

            if i % 200 == 0:
                evaluate(model, optimizer, val_loader)

def evaluate(model, optimizer, validation_dataloader):
    model.eval()
    eval_loss, eval_accuracy, nb_eval_steps = 0., 0., 0
    for batch in validation_dataloader:
        batch = tuple(t.to(device) for t in batch)
        with torch.no_grad():
            logits = model(batch[0], batch[1], batch[2])
            logits = logits.detach().cpu().numpy()
            label_ids = batch[3].cpu().numpy()
            tmp_eval_accuracy = flat_accuracy(logits, label_ids)
            eval_accuracy += tmp_eval_accuracy
            nb_eval_steps += 1

    print('\nValidation Accuracy: {}'.format(eval_accuracy / nb_eval_steps))

    global best_score

    if best_score < eval_accuracy / nb_eval_steps:
        best_score = eval_accuracy / nb_eval_steps
        save(model, optimizer)

    # model.train()

def test(model, dataloader, with_labels=False):
    checkpoint = torch.load(OUTPUT_MODEL, map_location='cpu')  # load model

    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)

    print('========== Testing ==========')

    pred_label = []
    model.eval()
    for i, batch in enumerate(tqdm(dataloader)):
        batch = tuple(t.to(device) for t in batch)
        with torch.no_grad():
            logits = model(batch[0], batch[1], batch[2])
            logits = logits.detach().cpu().numpy()
            preds = np.argmax(logits, axis=1).flatten()
            pred_label.extend(preds)

    pd.DataFrame(data=pred_label, index=range(len(pred_label))).to_csv('../results/pred.csv', header=['class_label'], encoding='utf-8')

    rel_dict = {'财经': '高风险', '时政': '高风险', '房产': '中风险', '科技': '中风险', '教育': '低风险', '时尚': '低风险', '游戏': '低风险', '家居': '可公开', '体育': '可公开', '娱乐': '可公开'}

    with open('../results/pred.csv', encoding='utf-8') as f:
        rows = [row for row in csv.reader(f)]
        rows = np.array(rows[1:])
        label_list = [label for _, label in rows]  # label list
        final_col = []
        for i in label_list:
            final_col.append(rel_dict[idx2classes[int(i)]])

        data = pd.read_csv('../results/pred.csv')
        data['final'] = final_col
        data = data.replace({'class_label': idx2classes})
        data.to_csv('../results/result.csv', index=False, encoding='utf_8_sig', header=['id', 'class_label', 'rank_label'])

    print('Test completed')

if __name__ == '__main__':
    set_seed(1224)

    data_df = process_data(filename='../data/train/all_labeled_data.csv', classes2idx=classes2idx, with_labels=True)
    train_df, val_df = train_test_split(data_df=data_df, test_size=0.2, shuffle=True, random_state=1)

    print('Reading training data...')
    train_set = CustomDataset(data=train_df, max_len=512, model_name=MODEL_NAME)
    train_loader = Data.DataLoader(dataset=train_set, batch_size=BATCH_SIZE, num_workers=5, shuffle=True)

    print('Reading validation data...')
    val_set = CustomDataset(data=val_df, max_len=512, model_name=MODEL_NAME)
    val_loader = Data.DataLoader(dataset=val_set, batch_size=BATCH_SIZE, num_workers=5, shuffle=True)

    model = CLSModel(model_name=MODEL_NAME, hidden_dim=HIDDEN_DIM, num_classes=len(classes2idx), dropout=DROPOUT, bidirectional=BIDIRECTIONAL, freeze_bert=False)
    criterion = nn.CrossEntropyLoss()

    optimizer = AdamW(model.parameters(), lr=1e-5, weight_decay=1e-2)

    train_eval(model=model, criterion=criterion, optimizer=optimizer, train_loader=train_loader, val_loader=val_loader, epochs=EPOCHS)

    print('Reading test data...')
    test_df = process_data(filename='../data/test_data_cleaned.csv', classes2idx=classes2idx, with_labels=False)
    test_set = CustomDataset(data=test_df, max_len=512, with_labels=False, model_name=MODEL_NAME)
    test_loader = Data.DataLoader(dataset=test_set, batch_size=BATCH_SIZE, num_workers=5, shuffle=False)

    test(model=model, dataloader=test_loader, with_labels=False)


#!/usr/bin/env python
"""Minimises the energy or the gap penalty function of a molecule in a cluster.

Template files are needed as well as an xyz file for the molecule and another
one for the surrounding molecules. Overall the use of subprocess is ugly as it
is repeated 3 or 4 times but it was found to handle memory better than Pool
when interfacing with Gaussian.

Energy units are Hartree inside the program but are printed in eV. Distances are
kept as Angstrom throughout and are converted from Bohr if necessary before
reaching this module

"""
import numpy as np
import subprocess
import os
from datetime import datetime
from scipy.optimize import minimize

from fromage.io import read_file as rf
from fromage.utils import array_operations as ao
from fromage.utils import calc
from fromage.io.parse_config_file import bool_cast

def sequence(in_pos):
    """
    Run Gaussian calculations in parallel and write and return results

    This function is designed to work with the scipy.optimise.minimize function.
    This is why it can only receive one array of floats as input and return two
    arrays of floats. As a result some variables in this function are defined
    elsewhere in the module which is a necessary evil.

    Parameters
    ----------
    in_pos : list of floats
        Input coordinates in array form
    Returns
    -------
    en_out : float
        Combined energy or penalty function value in Hartree
    gr_out : list of floats
        Gradients of en_out in Hartree/Angstrom

    References
    ----------
    Levine, B. G., Coe, J. D. & Martinez, T. J. Optimizing conical intersections
    without derivative coupling vectors: Application to multistate
    multireference second-order perturbation theory (MS-CASPT2).
    J. Phys. Chem. B 112, 405-413 (2008).

    """
    # initialise calculation objects
    rl = calc.setup_calc("rl", low_level)
    ml = calc.setup_calc("ml", low_level)
    mh = calc.setup_calc("mh", high_level)
    if bool_ci:
        mg = calc.setup_calc("mg", high_level)

    # Run the calculations as subprocesses with a maximum of 2 simultameous ones
    # at the same time. This order is optimised for the mh calculation being
    # the longest
    mh_proc = mh.run(ao.array2atom(mol_atoms, in_pos))
    if bool_ci and high_level != "gaussian_cas":
        mg_proc = mg.run(ao.array2atom(mol_atoms, in_pos))
    rl_proc = rl.run(ao.array2atom(mol_atoms, in_pos))
    rl_proc.wait()
    ml_proc = ml.run(ao.array2atom(mol_atoms, in_pos))
    ml_proc.wait()
    mh_proc.wait()
    if bool_ci and high_level != "gaussian_cas":
        mg_proc.wait()

    # read results. Each x_en_gr is a tuple (energy,gradients,scf_energy)
    rl_en_gr = rl.read_out(in_pos, mol_atoms, shell_atoms)
    ml_en_gr = ml.read_out(in_pos)

    if high_level == "gaussian_cas":
        mh_en_gr = mh.read_out(in_pos)[0:3]
        if bool_ci:
            mg_en_gr = (mh.read_out(in_pos)[2], mh.read_out(
                in_pos)[3], mh.read_out(in_pos)[2])
    elif high_level == "molcas":
        mh_en_gr = mh.read_out(in_pos)
        if bool_ci:
            # in molcas the first read out is S_1, gr, S_0 even if the target
            # state is 0
            mg_en_gr = mg.read_out(in_pos)[::-1]
    else:
        mh_en_gr = mh.read_out(in_pos)
        if bool_ci:
            mg_en_gr = mg.read_out(in_pos)

    # combine results
    en_combo = rl_en_gr[0] - ml_en_gr[0] + mh_en_gr[0]
    gr_combo = rl_en_gr[1] - ml_en_gr[1] + mh_en_gr[1]
    scf_combo = rl_en_gr[2] - ml_en_gr[2] + mh_en_gr[2]

    if bool_ci:
        # corresponding ground state energy and gradients
        en_combo_g = rl_en_gr[0] - ml_en_gr[0] + mg_en_gr[0]
        gr_combo_g = rl_en_gr[1] - ml_en_gr[1] + mg_en_gr[1]

        # Penalty function parameters and calculation
        alpha = 0.02
        e_mean = (en_combo + en_combo_g) / 2
        e_diff = en_combo - en_combo_g
        g_ij = e_diff**2 / (e_diff + alpha)
        en_out = e_mean + sigma * g_ij
        gr_out = 0.5 * (gr_combo + gr_combo_g) + sigma * ((e_diff**2 + 2 *
                                                           alpha * e_diff) / (e_diff + alpha)**2) * (gr_combo - gr_combo_g)
    else:
        en_out = en_combo
        gr_out = gr_combo

    # print some updates in the output
    out_file.write("------------------------------\n")
    global iteration
    iteration += 1
    out_file.write("Iteration: " + str(iteration) + "\n")
    out_file.write("Real low energy: {:>30.8f} eV\n".format(
        rl_en_gr[0] * evconv))
    out_file.write("Model low energy: {:>29.8f} eV\n".format(
        ml_en_gr[0] * evconv))
    out_file.write("Model high energy: {:>28.8f} eV\n".format(
        mh_en_gr[0] * evconv))
    out_file.write(
        "ONIOM Total energy: {:>27.8f} eV\n".format(en_combo * evconv))
    out_file.write(
        "ONIOM SCF energy: {:>29.8f} eV\n".format(scf_combo * evconv))
    out_file.write(
        "Energy grad. norm: {:>28.8f} eV/A\n".format(np.linalg.norm(gr_combo * evconv)))
    if bool_ci:
        out_file.write(
            "Penalty function value: {:>23.8f} eV\n".format(en_out * evconv))
        out_file.write("Penalty function grad. norm: {:>18.8f} eV\n".format(
            np.linalg.norm(gr_out * evconv)))

    out_file.write("Gap: {:>42.8f} eV\n".format(
        (en_combo - scf_combo) * evconv))
    out_file.flush()
    return (en_out, gr_out)

if __name__ == '__main__':

    evconv = 27.2114  # Something in Hartree * evconv = Something in eV

    # default settings

    def_inputs = {
        "mol_file": "mol.init.xyz",
        "shell_file": "shell.xyz",
        "out_file": "fromage.out",
        "bool_ci": "0",
        "high_level": "gaussian",
        "low_level": "gaussian",
        "sigma": "3.5",
        "gtol": 1e-5,
        "single_point": "0"}

    inputs = def_inputs.copy()

    # read user inputs
    if os.path.isfile("fromage.in"):
        new_inputs = rf.read_config("fromage.in")
        inputs.update(new_inputs)

    mol_file = inputs["mol_file"]
    shell_file = inputs["shell_file"]
    out_file = inputs["out_file"]
    bool_ci = bool_cast(inputs["bool_ci"])
    high_level = inputs["high_level"]
    low_level = inputs["low_level"]
    gtol = inputs["gtol"]
    single_point = bool_cast(inputs["single_point"])
    # sigma is called lambda in some papers but that is a bad variable name
    # in Python
    sigma = float(inputs["sigma"])
    # output
    out_file = open(out_file, "w", 1)
    # print start time
    start_time = datetime.now()
    out_file.write("STARTING TIME: " + str(start_time) + "\n")

    iteration = 0

    # clean up the last output
    if os.path.exists("geom_mol.xyz"):
        subprocess.call("rm geom_mol.xyz", shell=True)
    if os.path.exists("geom_cluster.xyz"):
        subprocess.call("rm geom_cluster.xyz", shell=True)

    # read initial coordniates
    mol_atoms = rf.read_xyz(mol_file)[0]

    # read shell atoms
    shell_atoms = rf.read_xyz(shell_file)[0]

    # make the initial coordinates into a flat list
    atoms_array = []
    for atom in mol_atoms:
        atoms_array.append(atom.x)
        atoms_array.append(atom.y)
        atoms_array.append(atom.z)

    # make the list into an array
    atoms_array = np.array(atoms_array)
    if single_point:
        sequence(atoms_array)
    else:
        res = minimize(sequence, atoms_array, jac=True,
                       options={'disp': True, 'gtol': gtol})

    out_file.write("DONE\n")
    end_time = datetime.now()
    out_file.write("ELAPSED TIME: " + str(end_time - start_time) + "\n")
    out_file.write("ENDING TIME: " + str(end_time) + "\n")
    out_file.close()


import os, sys
import eval_rgan
import eval_imgsynth
sys.path.insert(0,"scd")
import scd
import tensorflow as tf
import numpy as np
import cv2
import skimage.io

def main():
    config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
    sess=tf.Session(config=config)

    #define tensors and  models
    zb=6
    nf=256
    zdim=1024
    images = tf.placeholder(tf.float32,[2,150,496,3])
    z = tf.placeholder(tf.float32,[zb,zdim])

    scdobj = scd.SCD(input=images)
    imgsynth = eval_imgsynth.ImageSynthesizer(scdobj.end_points["pool5"])

    rgan = eval_rgan.rgan(images=scdobj.end_points["pool5"],z_dim=zdim,nf=nf, training=True)
    imgsynth_from_feature = eval_imgsynth.ImageSynthesizer(rgan.rec_x_out,reuse=True)

    rgan_from_z = eval_rgan.rgan(latent_vec=z,nf=nf,z_dim=zdim,reuse=True, training=True)
    imgsynth_from_z = eval_imgsynth.ImageSynthesizer(rgan_from_z.rec_x_p_out,reuse=True)

    #load weights
    scd_saver = scd.get_saver()
    scd_saver.restore(sess,"/workspace/imgsynth/ssd_300_kitti/ssd_model.ckpt")
    imgsynth_saver = eval_imgsynth.get_saver()
    imgsynth_saver.restore(sess,"/workspace/imgsynth/result_kitti256p_2/model.ckpt-89")
    rgan_saver = eval_rgan.get_saver()
    rgan_saver.restore(sess,"models/model-99999")

    co=[]
    image = cv2.resize(scd.imread_as_jpg("/workspace2/kitti/testing/image_2/007482.png"),(496,150))
    co.append(image)
    image = cv2.resize(scd.imread_as_jpg("/workspace2/kitti/testing/image_2/007481.png"),(496,150))
    co.append(image)
    #reses = imgsynth.generate_image_from_featuremap(np.expand_dims(image,0),images)
    image = np.array(co)

    #simply get detection result
    reses = scdobj.get_image(sess,image)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/detected%d.jpg"%i,a)
    #model:image=>feature, image decoder:feature=>image_hat
    reses = imgsynth.generate_image_from_featuremap(sess,image,images)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/test%d.jpg"%i,a)

    #model:image=>feature, gan:feature=>latent vector=>feature_hat, image decoder:feature_hat=>image_hat
    reses = imgsynth_from_feature.generate_image_from_featuremap(sess,image,images)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/test%d_2.jpg"%i,a)

    #get detection result thru gan and image decoder
    reses = scdobj.get_image(sess,reses)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/gan_detected%d.jpg"%i,a)

    return 0
    #latent vector=>feature_hat, image decoder:feature_hat=>image_hat
    np.random.seed(0)
    zs = []
    zsp = np.random.randn(zb,zdim)+1
    zs = np.array(zsp)
    reses = imgsynth_from_z.generate_image_from_featuremap(sess,zs,z)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/noise%d.jpg"%i,a)

    #noise interpolation
    zsb = []
    for i in range(0,zb):
        #zs.append(zsp[0]*float(i)/(float(zb)-1.)+zsp[1]*(1.-float(i)/(float(zb)-1.)))
        zsb.append(zsp[i]*(1.-i/(zb-1)))
    zsb = np.array(zsb)
    reses = imgsynth_from_z.generate_image_from_featuremap(sess,zsb,z)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/noise_interpolate%d.jpg"%i,a)

    #model:images=>features, gan:features=>latent vectors=>features_hat, interpolate:features_hat=>features_interplated, image decoder:feature_interplated=>image_hat
    zs = []
    zsp = rgan.generate_noise_from_featuremap(sess,image,images)
    for i in range(0,zb):
        zs.append(zsp[0]*float(i)/zb+zsp[1]*(1-float(i)/zb))
    zs = np.array(zs)
    reses = imgsynth_from_z.generate_image_from_featuremap(sess,zs,z)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/interpolate%d.jpg"%i,a)

    zs = []
    zsp = rgan.generate_noise_from_featuremap(sess,image,images)
    for i in range(0,zb):
        zs.append(zsp[0]*float(i)/zb+zsp[1]*(1-float(i)/zb)+3)
    zs = np.array(zs)
    reses = imgsynth_from_z.generate_image_from_featuremap(sess,zs,z)
    for i, a in enumerate(reses):
        print(i)
        skimage.io.imsave("test/interpolate2%d.jpg"%i,a)

if __name__ == "__main__":
    main()


# Copyright 2017 Google Inc. and Skytruth Inc.
#
# 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.

from __future__ import absolute_import

import logging
import numpy as np
import os
import pytz
import subprocess
import tempfile
import tensorflow as tf
import time
import uuid
from datetime import datetime
from datetime import timedelta

class Inferer(object):
    def __init__(self, model, model_checkpoint_path, root_feature_path, parallelism=4):

        self.model = model
        self.estimator = model.make_estimator(model_checkpoint_path)
        self.root_feature_path = root_feature_path
        logging.info('created Inferer with Model, %s, and dims %s', model,
                    model.num_feature_dimensions)
        self.parallelism = 4

    def close(self):
        self.sess.close()

    def _feature_files(self, ids):
        return [
            '%s/%s.tfrecord' % (self.root_feature_path, x)
            for x in ids
        ]

    def _build_time_ranges(self, interval_months, start_date, end_date):
        # TODO: should use min_window_duration here
        window_dur_seconds = self.model.max_window_duration_seconds
        last_viable_date = datetime.now(
            pytz.utc) - timedelta(seconds=window_dur_seconds)
        time_starts = []
        start_year = start_date.year
        month_count = start_date.month - 1
        if start_date.day != 1:
            raise ValueError('start_date must fall on the 1st of the month')
        dt = start_date
        while True:
            year = start_year + month_count // 12
            month = month_count % 12 + 1
            month_count += interval_months
            dt = datetime(year, month, 1, tzinfo=pytz.utc)
            if dt >= end_date:
                break
            time_starts.append(dt)
        delta = timedelta(seconds=self.model.max_window_duration_seconds)
        time_ranges = [(int(time.mktime(dt.timetuple())),
                             int(time.mktime((dt + delta).timetuple())))
                            for dt in time_starts]
        return time_ranges

    def run_inference(self, ids, interval_months, start_date, end_date):
        paths = self._feature_files(ids)

        if self.model.max_window_duration_seconds != 0:
            time_ranges = self._build_time_ranges(interval_months, start_date, end_date)
            input_fn = self.model.make_prediction_input_fn(paths, time_ranges, self.parallelism)
        else:
            input_fn = self.model.make_prediction_input_fn(paths, (start_date, end_date), self.parallelism)

        for result in self.estimator.predict(input_fn=input_fn):

            start_time, end_time = [datetime.utcfromtimestamp(x) for x in result['time_ranges']]
            output = {
                'id': result['id'],
                'start_time': start_time.isoformat(),
                'end_time': end_time.isoformat()
            }
            for k, v in result.items():
                if k in self.model.objective_map:
                    o = self.model.objective_map[k]
                    output[o.metadata_label] = o.build_json_results(v, result['timestamps'])

            yield output