# coding=utf-8
from __future__ import print_function

from six import string_types
import importlib
import inspect

from fancytools.fcollections.NestedOrderedDict import NestedOrderedDict

class GetCallablesInPackage(NestedOrderedDict):
    """
    Return a NestedOrderedDict of all functions and classes within a given package.
    excluding 'private' modules beginning with '_'

    =========  ======================================
    Arguments  Description
    =========  ======================================
    package    package or string of path to package
               to get all callables from
    =========  ======================================

    =====================  =======  ===============================
    Optional               Default  Description
    =====================  =======  ===============================
    modules_in_structure   False    include the modules of the
                                    callables in the output
    include_classes        True     include classes in the output
    include_functions      True     include functions in the output
    min_level              0        minimum depth level in the
                                    package structure
    max_level              None     maximum depth level in the
                                    package structure
    exclude_empty_pck      True     exclude sub-packages without
                                    any callables from the output
    =====================  =======  ===============================

    """

    def __init__(self, package, modules_in_structure=False,
                 include_classes=True, include_functions=False,
                 min_level=0, max_level=None, exclude_empty_pck=True):
        """
        """

        NestedOrderedDict.__init__(self)

        if isinstance(package, string_types):
            package = importlib.import_module(package)

        self._package = package
        self._include_classes = include_classes
        self._include_functions = include_functions
        self._modules_in_structure = modules_in_structure
        self._min_level = min_level
        self._max_level = max_level

        self._objects = [package]

        self._buildRecursive(package, self, 0)

        if exclude_empty_pck:
            self._cleanRecursive(self)

    def _cleanRecursive(self, subSelf):
        """
        Delete all NestedOrderedDict that haven't any entries.
        """
        for key, item in list(subSelf.items()):
            if self.isNestedDict(item):
                if not item:
                    subSelf.pop(key)
                else:
                    self._cleanRecursive(item)

    def _buildRecursive(self, module, subSelf, level):

        for (name, obj) in inspect.getmembers(module):
            # take only 'public' modules
            if name[0] == '_':
                continue
            # print name, level
            if inspect.ismodule(obj):
                # don't handle modules thats have been used before:
                if obj in self._objects:
                    continue
                self._objects.append(obj)

                if obj.__name__.startswith(self._package.__name__):
                    # either use module for submenus or check whether module is
                    # actually a package:
                    if self._modules_in_structure or obj.__file__.endswith(
                            '__init__.pyc'):
                        # limit max recursion level
                        if not self._max_level or level <= self._max_level:
                            # if self._max_level:
                            l = subSelf[name] = NestedOrderedDict()
                            self._buildRecursive(obj, l, level + 1)
                    else:
                        # use same menu for all classes of all files on one
                        # folder/module
                        self._buildRecursive(obj, subSelf, level)

            elif level >= self._min_level:
                    # if object is a class
                if (
                    # ( not self._include_classes or ( self._include_classes and inspect.isclass(obj) ) or
                    (self._include_classes and inspect.isclass(obj)) or
                    # and/or a function
                   # ( not self._include_functions or ( self._include_functions and inspect.isfunction(obj)) )
                   (self._include_functions and inspect.isfunction(obj))
                   ):
                        # if object belongs to its parent module (mod is not imported)
                        # OR module paths are the same (in case of from mod
                        # import cls and cls.name == mod.name)
                    try:
                        if self.belongsToModule(obj, module):
                            subSelf[name] = obj
                    except AttributeError:
                        pass  # obj has no __module__ is e.g. a string

    @staticmethod
    def belongsToModule(obj, module):
        """Returns True is an object belongs to a module."""
        return obj.__module__ == module.__name__ or obj.__module__.startswith(
            module.__name__)

    def buildHirarchy(self, horizontal_operation=None,
                      vertical_operation=None):
        """Walk through the nested dict structure and executes
        horizontal_operation(name, callable) resp.
        vertical_operation(name, callable) if defined.
        """
        def buildRecursive(pkey, pval):
            if self.isNestedDict(pval):
                if vertical_operation:
                    vertical_operation(pkey, pval)
                for key, val in list(pval.items()):
                    # if not isinstance(val, NestedOrderedDict):
                    buildRecursive(key, val)
            else:
                if horizontal_operation:
                    horizontal_operation(pkey, pval)
        buildRecursive(self._package.__name__, self)

    @staticmethod
    def isNestedDict(instance):
        """convenience function for

        >>> isinstance(instance, NestedOrderedDict)"""
        return isinstance(instance, NestedOrderedDict)

if __name__ == '__main__':
    import sys
    import numpy
    g = GetCallablesInPackage(numpy, include_functions=True)
    print(g)

    class printClassStructure(object):

        def __init__(self):
            self.indent = ''
            self._hor_as_last_called = True

        def hor(self, key, item):
            sys.stdout.write("%s%s - " % (self.indent, key))
            self.indent = ''
            self._hor_as_last_called = True

        def vert(self, key, item):
            if self._hor_as_last_called:
                print('')
            print("%s-> %s" % (self.indent, key))
            self.indent += '\t'
            self._hor_as_last_called = False

    p = printClassStructure()

    g.buildHirarchy(horizontal_operation=p.hor,
                    vertical_operation=p.vert)


import uuid
from rest_framework.response import Response

class ResponseHelper():
    def __init__(self, status, message, data=None):
        self.status = status
        self.message = message
        self.data = data

    def helper_response(self):
        return Response({
            'status_code': self.status,
            'message': self.message,
            'data': self.data
        })

    def helper_response_without_data(self):
        return Response({
            'status_code': self.status,
            'message': self.message,
        })

class GetObjectHelper():
    def __init__(self, model, id):
        self.model = model
        self.id = id

    def get_object(self):
        '''
        Helper method to get the object with the given uuid
        '''
        try:
            uuid.UUID(self.id)
        except Exception:
            return None

        try:
            return self.model.objects.get(uuid=self.id)
        except:
            return None

class BodyItemHelper():
    def __init__(self, data):
        self.data = data

    def parser_data(self):
        '''
        method helper to parse data from request body
        to be used in create_item method in views.py
        '''

        name = self.data['name']
        type = self.data['type']
        prices = self.data['prices']

        regular_price = None
        vip_price = None
        wholesale_price = None

        for p in prices:
            if p['priceFor'] == 'regular':
                regular_price = p['price']
            elif p['priceFor'] == 'vip':
                vip_price = p['price']
            elif p['priceFor'] == 'wholesale':
                wholesale_price = p['price']
            else:
                print('error')

        try:
            if vip_price is None:
                vip_price = regular_price
        except:
            pass

        try:
            if wholesale_price is None:
                wholesale_price = regular_price
        except:
            pass

        return {
            'name': name,
            'type': type,
            'regular_price': regular_price,
            'vip_price': vip_price,
            'wholesale_price': wholesale_price
        }


"""
Import necessary libraries to create a variational autoencoder
The code is mainly developed using the PyTorch library
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
import os
import numpy as np
import copy
import pickle
from torch.utils.tensorboard import SummaryWriter
from progress.bar import Bar
import torchvision
import socket
from mpi4py import MPI
import time

model_name = 'VAE_MNIST'
"""
Initialize Hyperparameters
"""
batch_size = 2#2
learning_rate = 1e-4
num_epochs = 50
input_size = 20#82
imgChannels = 1#3
n_filters = 5
imsize2 = input_size - (n_filters-1) * 3
convdim1 = 8
convdim2 = 16
convdim3 = 32
zDim = 300

writer = SummaryWriter(log_dir='/local/scratch/jrs596/VAE/logs')

"""
Create dataloaders
"""

data_transforms = transforms.Compose([
        transforms.Resize((input_size,input_size)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor()])

#train_dir = '/local/scratch/jrs596/dat/ResNetFung50+_images_unorganised'
#train_dir = '/local/scratch/jrs596/dat/ResNetFung50+_images_unorganised_test'
#train_dataset = datasets.ImageFolder(train_dir, data_transforms['train'])
#train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

mnist_data = torchvision.datasets.MNIST('/local/scratch/jrs596/dat/mnist', transform=data_transforms)
train_loader = torch.utils.data.DataLoader(mnist_data,
                                          batch_size=2,
                                          shuffle=True,
                                          num_workers=4)

n = len(train_loader.dataset)
if n%batch_size != 0:
    print('Total N samples must be divisable by batch size')
    print('N = ' + str(n))
    exit(0)

class VAE(nn.Module):
    def __init__(self, imgChannels=imgChannels, featureDim=int((batch_size*convdim3*imsize2*imsize2)/4),
        zDim=zDim):
        super(VAE, self).__init__()

        # Initializing the 2 convolutional layers and 2 full-connected layers for the encoder
        layer2=int(featureDim/10)
        self.encConv1 = nn.Conv2d(imgChannels, convdim1, n_filters).to('cuda:0')
        self.encConv2 = nn.Conv2d(convdim1, convdim2, n_filters).to('cuda:0')
        self.encConv3 = nn.Conv2d(convdim2, convdim3, n_filters).to('cuda:0')
        self.pool = nn.MaxPool2d(2, stride=2, return_indices=True).to('cuda:0')
        self.unPool = nn.MaxUnpool2d(2, stride=2).to('cuda:0')
        self.batchNorm = nn.BatchNorm2d(16).to('cuda:0')
        self.encFC1 = nn.Linear(featureDim, layer2).to('cuda:0')
        self.encFC1_1 = nn.Linear(layer2, layer2).to('cuda:0')
        self.encFC1_2 = nn.Linear(layer2, zDim).to('cuda:0')
        self.encFC2 = nn.Linear(featureDim, layer2).to('cuda:0')
        self.encFC2_1 = nn.Linear(layer2, layer2).to('cuda:0')
        self.encFC2_2 = nn.Linear(layer2, zDim).to('cuda:0')

        # Initializing the fully-connected layer and 2 convolutional layers for decoder
        self.decFC1_1 = nn.Linear(zDim, layer2).to('cuda:0')
        self.decFC1_2 = nn.Linear(layer2, layer2).to('cuda:0')
        self.decFC1 = nn.Linear(layer2, featureDim).to('cuda:1')
        self.decConv1 = nn.ConvTranspose2d(convdim3, convdim2, n_filters).to('cuda:0')
        self.decConv2 = nn.ConvTranspose2d(convdim2, convdim1, n_filters).to('cuda:0')
        self.decConv3 = nn.ConvTranspose2d(convdim1, imgChannels, n_filters).to('cuda:0')

    def encoder(self, x):

        # Input is fed into 2 convolutional layers sequentially
        # The output feature map are fed into 2 fully-connected layers to predict mean (mu) and variance (logVar)
        # Mu and logVar are used for generating middle representation z and KL divergence loss
        x = F.relu(self.encConv1(x.to('cuda:0')))
        #x = self.pool(x)
        x = F.relu(self.encConv2(x))
        #x = self.pool(x)
        x = self.batchNorm(x)
        x = F.relu(self.encConv3(x))
        x, indices = self.pool(x)
        x_dim = np.prod(list(x.shape))
        x = x.view(-1, x_dim)
        mu = self.encFC1(x)
        mu = self.encFC1_1(mu)
        mu = self.encFC1_2(mu)
        logVar = self.encFC2(x)
        logVar = self.encFC2_1(logVar)
        logVar = self.encFC2_2(logVar)
        return mu, logVar, indices

    def reparameterize(self, mu, logVar):

        #Reparameterization takes in the input mu and logVar and sample the mu + std * eps
        std = torch.exp(logVar/2)
        eps = torch.randn_like(std)
        return mu + std * eps

    def decoder(self, z, indices):

        # z is fed back into a fully-connected layers and then into two transpose convolutional layers
        # The generated output is the same size of the original input
        x = F.relu(self.decFC1_1(z))
        x = F.relu(self.decFC1_2(x))
        x = F.relu(self.decFC1(x.to('cuda:1')))
        x = x.view(batch_size, convdim3, int(imsize2/2) ,int(imsize2/2))
        x = self.unPool(x.to('cuda:0'), indices, output_size=([batch_size,convdim3,imsize2,imsize2]))
        x = F.relu(self.decConv1(x))
        x = F.relu(self.decConv2(x))
        x = torch.sigmoid(self.decConv3(x))
        return x

    def forward(self, x):

        # The entire pipeline of the VAE: encoder -> reparameterization -> decoder
        # output, mu, and logVar are returned for loss computation
        mu, logVar, indices = self.encoder(x)
        z = self.reparameterize(mu, logVar)
        out = self.decoder(z, indices)
        return out, mu, logVar

model = VAE()

"""
Training
"""
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

for epoch in range(num_epochs):
    running_ELBO = 0.0
    running_CE = 0.0
    running_KL = 0.0

    with Bar('Learning...', max=n/batch_size) as bar:
          for idx, data in enumerate(train_loader):

            imgs = data[0]
            labels = imgs.to('cuda:0')
            # Feeding a batch of images into the network to obtain the output image, mu, and logVar

            out, mu, logVar = model(imgs)

            kl_divergence = 0.5 * torch.sum(torch.exp(logVar) + mu**2 - 1. - logVar)

            binary_cross_entropy = F.binary_cross_entropy(out, labels)
            loss = binary_cross_entropy + kl_divergence

            running_ELBO += loss.item() * imgs.shape[0]
            running_CE += binary_cross_entropy.item() * imgs.shape[0]
            running_KL += kl_divergence.item() * imgs.shape[0]

            # Backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            bar.next()

    epoch_loss = round(float(running_ELBO / n),4)
    epoch_CE = round(float(running_CE / n),4)
    epoch_KL = round(float(running_KL / n),4)

    writer.add_scalar("ELBO/train", epoch_loss, epoch)
    writer.add_scalar("Cross entropy/train", epoch_CE, epoch)
    writer.add_scalar("KL divergence/train", epoch_KL, epoch)

    if epoch%10 == 0:
        print('Saving checkpoint')
        PATH = '/local/scratch/jrs596/VAE/models'
        torch.save(model.state_dict(), os.path.join(PATH, model_name + '.pth'))

    print('Epoch {}: ELBO loss {}'.format(epoch, epoch_loss))
    print('KL divergence {}: Binary Cross Entropy {}'.format(epoch_KL, epoch_CE))

writer.flush()
writer.close()


# import ingest
import os
try:
    from secrets import finn_hub_key
except ImportError:
    print('Failed to import "finn_hub_key" from "secrets.py". Aborting fast-forward.')
    exit(0)

from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from time import sleep

import psycopg2
from pgcopy import CopyManager
import finnhub

# CONNECTION = "postgres://root:password@timedb:5432/quantx"
CONNECTION: str = 'postgresql://{}:{}@{}:{}/{}'.format(
    os.environ.get("POSTGRES_USER", "root"),
    os.environ.get("POSTGRES_PASSWORD", "password"),
    os.environ.get("POSTGRES_HOST", "localhost"),
    os.environ.get("POSTGRES_PORT", 5432),
    os.environ.get("POSTGRES_DB", "quantx")
)
symbols = []
# fast forward till the end of the past hour
end_time = datetime.now()
end_time = end_time.replace(second=0, microsecond=0, minute=0, hour=end_time.hour);
end_time -= timedelta(minutes=1)

def main():
    with psycopg2.connect(CONNECTION) as conn:
        with conn.cursor() as curs:
            curs.execute("""select quote.symbol, MAX(time) from
                        symbol LEFT JOIN quote ON symbol.symbol = quote.symbol
                        GROUP BY quote.symbol
                        ;""")
            symbol_time = curs.fetchall()

    # convert list of tuple to list of list:
    symbol_time = [list(tup) for tup in symbol_time]
    for entry in symbol_time:
        if entry[1] is None:
            # by default, retrieve one month of data
            entry[1] = datetime.now() - timedelta(days=30)
        else:
            entry[1] += timedelta(minutes=1) # start grabbing data from 1 min after the last data point
    if not symbol_time:
        raise RuntimeWarning("No Symbols were retrieved from db!")

    for entry in symbol_time:
        start_time = entry[1]
    current_time = start_time
    while current_time < end_time:
        current_time = advance_one_month(current_time, symbol=entry[0])

def advance_one_month(current_start, symbol):
    '''Advance one month per api call because of FinnHub restrictions'''
    current_end = current_start + relativedelta(months=+1, minutes=-1)
    current_end = min(current_end, end_time)
    print(f"Retreiving data from {current_start} to {current_end}")
    responses = get_candles_for_symbol(current_start, current_end, symbol)
    write_db(responses, symbol)
    return current_end + timedelta(minutes = 1)

def write_db(responses, symbol):
    with psycopg2.connect(CONNECTION) as conn:
        cols=['time', 'symbol', 'price_open', 'price_high', 'price_low', 'price_close']
        mgr = CopyManager(conn, 'quote', cols)
        for response in responses:
            if response['s'] != 'ok': # out of market time, etc.
                continue
            datetimes = [datetime.fromtimestamp(x) for x in response['t']]
            data = [list(a) for a in zip(
                datetimes,
                [symbol]*len(datetimes),
                response['o'], response['h'], response['l'], response['c'])]
            mgr.copy(data)

def get_candles_for_symbol(start: datetime, end: datetime, symbol: str):
    start_time = int(start.timestamp())
    end_time = int(end.timestamp())

    responses = []
    finnhub_client = finnhub.Client(api_key=finn_hub_key)

    # make the call

    for _ in range(10):
        try:
            response = finnhub_client.stock_candles(symbol, '1', start_time, end_time)
        except finnhub.FinnhubAPIException as e:
            if  e.status_code == 429: # out of API calls
                sleep(60)
                continue
            else:
                raise e
        break
    else:
        raise RuntimeWarning(f"API Call Failed for symbol {symbol}, start time = {start_time}, end time = {end_time}.")
    #print(response)
    responses.append(response)
    return responses

if __name__ == "__main__":
    main()


from PySide6.QtWidgets import QDialog, QWidget, QGroupBox, QScrollArea, QVBoxLayout, QTableWidget, QAbstractItemView
from PySide6.QtCharts import QChart, QLineSeries, QDateTimeAxis, QValueAxis, QChartView, QBarSeries
from PySide6.QtCore import Qt
from PySide6.QtGui import QFont

class StockInfoForecasts(QDialog):
    def __init__(self):
        super().__init__()
        self.setStyleSheet('background-color: deepskyblue')
        self.chart_groupbox = QGroupBox(self)
        self.chart_groupbox.setTitle("Charts and Graphs")
        self.chart_groupbox.setGeometry(0, 0, 1300, 600)
        self.chart_groupbox.content_container = QScrollArea(self)
        self.prediction_chart_widget = QWidget()
        self.prediction_chart_widget.resize(1300, 2000)
        self.prediction_chart_widget.setLayout(QVBoxLayout())
        self.chart_groupbox.content_container.setWidget(self.prediction_chart_widget)
        self.chart_groupbox.content_container.setGeometry(5, 15, 1290, 650)
        self.ptchart = QChart()
        ptlineseries = QLineSeries()
        ptlineseries.setName("stock")
        self.ptchart.addSeries(ptlineseries)
        x_axis = QDateTimeAxis()
        x_axis.setTickCount(7)
        x_axis.setFormat("MM-dd-yyyy")
        x_axis.setTitleText("Date")
        x_axis.setVisible(True)
        self.ptchart.addAxis(x_axis, Qt.AlignmentFlag.AlignBottom)
        ptlineseries.attachAxis(x_axis)
        y_axis = QValueAxis()
        self.ptchart.addAxis(y_axis, Qt.AlignmentFlag.AlignLeft)
        ptlineseries.attachAxis(y_axis)
        self.ptchartview = QChartView(self.ptchart)
        self.pt_groupbox = QGroupBox(self.prediction_chart_widget)
        self.pt_groupbox.setTitle("Analyst Price Targets")
        self.pt_groupbox.setGeometry(10, 10, 1200, 350)
        self.ptchartview.setParent(self.pt_groupbox)
        self.ptchartview.setGeometry(10, 15, 800, 300)
        self.ptchartview.setVisible(True)
        self.pt_label_container = QWidget(self.pt_groupbox)
        self.pt_label_container.setStyleSheet('background-color: white')
        self.pt_label_container.setGeometry(820, 22, 350, 285)
        self.pt_label_container.setLayout(QVBoxLayout())
        self.qtr_earnings_chart = QChart()
        self.qtr_earnings_chart.setTitle("Quarterly EPS")
        self.qtr_earnings_barseries = QBarSeries()
        self.qtr_earnings_barseries.setName("Earnings")
        self.qtr_earnings_chart.addSeries(self.qtr_earnings_barseries)
        self.qtr_earnings_groupbox = QGroupBox(self.prediction_chart_widget)
        self.qtr_earnings_groupbox.setTitle("Earnings History and Projections")
        self.qtr_earnings_groupbox.setGeometry(10, 370, 1200, 350)
        self.qtr_earnings_chartview = QChartView(self.qtr_earnings_chart)
        self.qtr_earnings_chartview.setParent(self.qtr_earnings_groupbox)
        self.qtr_earnings_chartview.setVisible(True)
        self.qtr_earnings_chartview.setGeometry(10, 15, 800, 300)
        self.qtr_earnings_table = QTableWidget(self.qtr_earnings_groupbox)
        self.qtr_earnings_table.setGeometry(820, 20, 350, 290)
        self.qtr_earnings_table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.qtr_earnings_table.setFont(QFont('arial', 10))
        self.qtr_earnings_table.setStyleSheet('background-color: white;')
        self.qtr_revtrend_chart = QChart()
        self.qtr_revtrend_chart.setTitle("Quarterly Revenue Trend")
        self.qtr_revtrend_barseries = QBarSeries()
        self.qtr_revtrend_chart.addSeries(self.qtr_revtrend_barseries)
        self.qtr_revtrend_gb = QGroupBox(self.prediction_chart_widget)
        self.qtr_revtrend_gb.setTitle("Revenue History and Projections")
        self.qtr_revtrend_gb.setGeometry(10, 730, 1200, 350)
        self.qtr_revtrend_chartview = QChartView(self.qtr_revtrend_chart)
        self.qtr_revtrend_chartview.setParent(self.qtr_revtrend_gb)
        self.qtr_revtrend_chartview.setVisible(True)
        self.qtr_revtrend_chartview.setGeometry(10, 15, 800, 300)
        self.qtr_revtrend_label_container = QWidget(self.qtr_revtrend_gb)
        self.qtr_revtrend_label_container.setGeometry(820, 20, 350, 290)
        self.qtr_revtrend_label_container.setStyleSheet('background-color: white;')
        self.qtr_revtrend_label_container.setLayout(QVBoxLayout())
        self.yr_earnings_chart = QChart()
        self.yr_earnings_chart.setTitle("Yearly EPS")
        self.yr_earnings_barseries = QBarSeries()
        self.yr_earnings_barseries.setName("Earnings")
        self.yr_earnings_chart.addSeries(self.yr_earnings_barseries)
        self.yr_earnings_groupbox = QGroupBox(self.prediction_chart_widget)
        self.yr_earnings_groupbox.setTitle("Earnings History and Projections")
        self.yr_earnings_groupbox.setGeometry(10, 1090, 1200, 350)
        self.yr_earnings_chartview = QChartView(self.yr_earnings_chart)
        self.yr_earnings_chartview.setParent(self.yr_earnings_groupbox)
        self.yr_earnings_chartview.setVisible(True)
        self.yr_earnings_chartview.setGeometry(10, 15, 800, 300)
        self.yr_earnings_label_container = QWidget(self.yr_earnings_groupbox)
        self.yr_earnings_label_container.setGeometry(820, 20, 350, 290)
        self.yr_earnings_label_container.setStyleSheet('background-color: white;')
        self.yr_earnings_label_container.setLayout(QVBoxLayout())
        self.yr_revtrend_chart = QChart()
        self.yr_revtrend_chart.setTitle("Yearly Revenue Trend")
        self.yr_revtrend_barseries = QBarSeries()
        self.yr_revtrend_chart.addSeries(self.yr_revtrend_barseries)
        self.yr_revtrend_gb = QGroupBox(self.prediction_chart_widget)
        self.yr_revtrend_gb.setTitle("Revenue History and Projections")
        self.yr_revtrend_gb.setGeometry(10, 1450, 1200, 360)
        self.yr_revtrend_chartview = QChartView(self.yr_revtrend_chart)
        self.yr_revtrend_chartview.setParent(self.yr_revtrend_gb)
        self.yr_revtrend_chartview.setVisible(True)
        self.yr_revtrend_chartview.setGeometry(10, 15, 800, 300)
        self.yr_revtrend_label_container = QWidget(self.yr_revtrend_gb)
        self.yr_revtrend_label_container.setGeometry(820, 20, 350, 290)
        self.yr_revtrend_label_container.setStyleSheet('background-color: white;')
        self.yr_revtrend_label_container.setLayout(QVBoxLayout())


from collections import OrderedDict

LEFT_SIDEBAR_MENU_ITEMS = OrderedDict(
    [ ('LIBRARY', ['Songs']),
    ('PLAYLISTS', [])])

DEFAULT_VIEWS = LEFT_SIDEBAR_MENU_ITEMS['LIBRARY']

DEFAULT_VIEWS_COUNT = len(DEFAULT_VIEWS)

def seconds_to_hms(seconds):
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)

    if h or (h and not m):
        return '{:02}:{:02}:{:02}'.format(h, m, s)
    else:
        return '{:02}:{:02}'.format(m, s)


# coding:utf-8
'''
@author: 猫大白
Project:Python的unittest框架
'''
from selenium import webdriver
from time import sleep
# driver = webdriver.Chrome()
# driver.implicitly_wait(30)
# driver.get('https://www.baidu.com')

#
# driver.find_element('id','kw').send_keys('临渊')
# driver.find_element('id','su').click()
# sleep(0.5)
# assert '临渊' in driver.title
# driver.quit()

# unittest形式
import unittest

class TestBaiDu(unittest.TestCase):
    def setUp(self) -> None:
        self.driver = webdriver.Chrome()
        self.driver.implicitly_wait(30)

    def tearDown(self) -> None:
        self.driver.quit()

    def test_baidu_search(self):
        self.driver.get('https://www.baidu.com')
        self.driver.find_element('id','kw').send_keys('临渊')
        self.driver.find_element('id','su').click()
        sleep(0.5)
        assert '临渊' in self.driver.title


import sys

from PyQt5.QtCore import QFile, QFileInfo, Qt, QTextCodec
from PyQt5.QtGui import (QFont, QFontDatabase, QFontInfo, QIcon, QKeySequence,
        QPixmap, QTextBlockFormat, QTextCharFormat, QTextCursor,
        QTextDocumentWriter, QTextListFormat)
from PyQt5.QtWidgets import (QAction, QActionGroup, QApplication, QColorDialog,
        QComboBox, QFileDialog, QFontComboBox, QMainWindow, QMenu, QMessageBox,
        QTextEdit, QToolBar)
from PyQt5.QtPrintSupport import QPrintDialog, QPrinter, QPrintPreviewDialog

import textedit_rc

if sys.platform.startswith('darwin'):
    rsrcPath = ":/images/mac"
else:
    rsrcPath = ":/images/win"

class TextEdit(QMainWindow):
    def __init__(self, fileName=None, parent=None):
        super(TextEdit, self).__init__(parent)

        self.setWindowIcon(QIcon(':/images/logo.png'))
        self.setToolButtonStyle(Qt.ToolButtonFollowStyle)
        self.setupFileActions()
        self.setupEditActions()
        self.setupTextActions()

        helpMenu = QMenu("Help", self)
        self.menuBar().addMenu(helpMenu)
        helpMenu.addAction("About", self.about)
        helpMenu.addAction("About &Qt", QApplication.instance().aboutQt)

        self.textEdit = QTextEdit(self)
        self.textEdit.currentCharFormatChanged.connect(
                self.currentCharFormatChanged)
        self.textEdit.cursorPositionChanged.connect(self.cursorPositionChanged)
        self.setCentralWidget(self.textEdit)
        self.textEdit.setFocus()
        self.setCurrentFileName()
        self.fontChanged(self.textEdit.font())
        self.colorChanged(self.textEdit.textColor())
        self.alignmentChanged(self.textEdit.alignment())
        self.textEdit.document().modificationChanged.connect(
                self.actionSave.setEnabled)
        self.textEdit.document().modificationChanged.connect(
                self.setWindowModified)
        self.textEdit.document().undoAvailable.connect(
                self.actionUndo.setEnabled)
        self.textEdit.document().redoAvailable.connect(
                self.actionRedo.setEnabled)
        self.setWindowModified(self.textEdit.document().isModified())
        self.actionSave.setEnabled(self.textEdit.document().isModified())
        self.actionUndo.setEnabled(self.textEdit.document().isUndoAvailable())
        self.actionRedo.setEnabled(self.textEdit.document().isRedoAvailable())
        self.actionUndo.triggered.connect(self.textEdit.undo)
        self.actionRedo.triggered.connect(self.textEdit.redo)
        self.actionCut.setEnabled(False)
        self.actionCopy.setEnabled(False)
        self.actionCut.triggered.connect(self.textEdit.cut)
        self.actionCopy.triggered.connect(self.textEdit.copy)
        self.actionPaste.triggered.connect(self.textEdit.paste)
        self.textEdit.copyAvailable.connect(self.actionCut.setEnabled)
        self.textEdit.copyAvailable.connect(self.actionCopy.setEnabled)
        QApplication.clipboard().dataChanged.connect(self.clipboardDataChanged)

        if fileName is None:
            fileName = ':/example.html'

        if not self.load(fileName):
            self.fileNew()

    def closeEvent(self, e):
        if self.maybeSave():
            e.accept()
        else:
            e.ignore()

    def setupFileActions(self):
        tb = QToolBar(self)
        tb.setWindowTitle("File Actions")
        self.addToolBar(tb)

        menu = QMenu("&File", self)
        self.menuBar().addMenu(menu)

        self.actionNew = QAction(
                QIcon.fromTheme('document-new',
                        QIcon(rsrcPath + '/filenew.png')),
                "&New", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.New, triggered=self.fileNew)
        tb.addAction(self.actionNew)
        menu.addAction(self.actionNew)

        self.actionOpen = QAction(
                QIcon.fromTheme('document-open',
                        QIcon(rsrcPath + '/fileopen.png')),
                "&Open...", self, shortcut=QKeySequence.Open,
                triggered=self.fileOpen)
        tb.addAction(self.actionOpen)
        menu.addAction(self.actionOpen)
        menu.addSeparator()

        self.actionSave = QAction(
                QIcon.fromTheme('document-save',
                        QIcon(rsrcPath + '/filesave.png')),
                "&Save", self, shortcut=QKeySequence.Save,
                triggered=self.fileSave, enabled=False)
        tb.addAction(self.actionSave)
        menu.addAction(self.actionSave)

        self.actionSaveAs = QAction("Save &As...", self,
                priority=QAction.LowPriority,
                shortcut=Qt.CTRL + Qt.SHIFT + Qt.Key_S,
                triggered=self.fileSaveAs)
        menu.addAction(self.actionSaveAs)
        menu.addSeparator()

        self.actionPrint = QAction(
                QIcon.fromTheme('document-print',
                        QIcon(rsrcPath + '/fileprint.png')),
                "&Print...", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.Print, triggered=self.filePrint)
        tb.addAction(self.actionPrint)
        menu.addAction(self.actionPrint)

        self.actionPrintPreview = QAction(
                QIcon.fromTheme('fileprint',
                        QIcon(rsrcPath + '/fileprint.png')),
                "Print Preview...", self,
                shortcut=Qt.CTRL + Qt.SHIFT + Qt.Key_P,
                triggered=self.filePrintPreview)
        menu.addAction(self.actionPrintPreview)

        self.actionPrintPdf = QAction(
                QIcon.fromTheme('exportpdf',
                        QIcon(rsrcPath + '/exportpdf.png')),
                "&Export PDF...", self, priority=QAction.LowPriority,
                shortcut=Qt.CTRL + Qt.Key_D,
                triggered=self.filePrintPdf)
        tb.addAction(self.actionPrintPdf)
        menu.addAction(self.actionPrintPdf)
        menu.addSeparator()

        self.actionQuit = QAction("&Quit", self, shortcut=QKeySequence.Quit,
                triggered=self.close)
        menu.addAction(self.actionQuit)

    def setupEditActions(self):
        tb = QToolBar(self)
        tb.setWindowTitle("Edit Actions")
        self.addToolBar(tb)

        menu = QMenu("&Edit", self)
        self.menuBar().addMenu(menu)

        self.actionUndo = QAction(
                QIcon.fromTheme('edit-undo',
                        QIcon(rsrcPath + '/editundo.png')),
                "&Undo", self, shortcut=QKeySequence.Undo)
        tb.addAction(self.actionUndo)
        menu.addAction(self.actionUndo)

        self.actionRedo = QAction(
                QIcon.fromTheme('edit-redo',
                        QIcon(rsrcPath + '/editredo.png')),
                "&Redo", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.Redo)
        tb.addAction(self.actionRedo)
        menu.addAction(self.actionRedo)
        menu.addSeparator()

        self.actionCut = QAction(
                QIcon.fromTheme('edit-cut', QIcon(rsrcPath + '/editcut.png')),
                "Cu&t", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.Cut)
        tb.addAction(self.actionCut)
        menu.addAction(self.actionCut)

        self.actionCopy = QAction(
                QIcon.fromTheme('edit-copy',
                        QIcon(rsrcPath + '/editcopy.png')),
                "&Copy", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.Copy)
        tb.addAction(self.actionCopy)
        menu.addAction(self.actionCopy)

        self.actionPaste = QAction(
                QIcon.fromTheme('edit-paste',
                        QIcon(rsrcPath + '/editpaste.png')),
                "&Paste", self, priority=QAction.LowPriority,
                shortcut=QKeySequence.Paste,
                enabled=(len(QApplication.clipboard().text()) != 0))
        tb.addAction(self.actionPaste)
        menu.addAction(self.actionPaste)

    def setupTextActions(self):
        tb = QToolBar(self)
        tb.setWindowTitle("Format Actions")
        self.addToolBar(tb)

        menu = QMenu("F&ormat", self)
        self.menuBar().addMenu(menu)

        self.actionTextBold = QAction(
                QIcon.fromTheme('format-text-bold',
                        QIcon(rsrcPath + '/textbold.png')),
                "&Bold", self, priority=QAction.LowPriority,
                shortcut=Qt.CTRL + Qt.Key_B, triggered=self.textBold,
                checkable=True)
        bold = QFont()
        bold.setBold(True)
        self.actionTextBold.setFont(bold)
        tb.addAction(self.actionTextBold)
        menu.addAction(self.actionTextBold)

        self.actionTextItalic = QAction(
                QIcon.fromTheme('format-text-italic',
                        QIcon(rsrcPath + '/textitalic.png')),
                "&Italic", self, priority=QAction.LowPriority,
                shortcut=Qt.CTRL + Qt.Key_I, triggered=self.textItalic,
                checkable=True)
        italic = QFont()
        italic.setItalic(True)
        self.actionTextItalic.setFont(italic)
        tb.addAction(self.actionTextItalic)
        menu.addAction(self.actionTextItalic)

        self.actionTextUnderline = QAction(
                QIcon.fromTheme('format-text-underline',
                        QIcon(rsrcPath + '/textunder.png')),
                "&Underline", self, priority=QAction.LowPriority,
                shortcut=Qt.CTRL + Qt.Key_U, triggered=self.textUnderline,
                checkable=True)
        underline = QFont()
        underline.setUnderline(True)
        self.actionTextUnderline.setFont(underline)
        tb.addAction(self.actionTextUnderline)
        menu.addAction(self.actionTextUnderline)

        menu.addSeparator()

        grp = QActionGroup(self, triggered=self.textAlign)

        # Make sure the alignLeft is always left of the alignRight.
        if QApplication.isLeftToRight():
            self.actionAlignLeft = QAction(
                    QIcon.fromTheme('format-justify-left',
                            QIcon(rsrcPath + '/textleft.png')),
                    "&Left", grp)
            self.actionAlignCenter = QAction(
                    QIcon.fromTheme('format-justify-center',
                            QIcon(rsrcPath + '/textcenter.png')),
                    "C&enter", grp)
            self.actionAlignRight = QAction(
                    QIcon.fromTheme('format-justify-right',
                            QIcon(rsrcPath + '/textright.png')),
                    "&Right", grp)
        else:
            self.actionAlignRight = QAction(
                    QIcon.fromTheme('format-justify-right',
                            QIcon(rsrcPath + '/textright.png')),
                    "&Right", grp)
            self.actionAlignCenter = QAction(
                    QIcon.fromTheme('format-justify-center',
                            QIcon(rsrcPath + '/textcenter.png')),
                    "C&enter", grp)
            self.actionAlignLeft = QAction(
                    QIcon.fromTheme('format-justify-left',
                            QIcon(rsrcPath + '/textleft.png')),
                    "&Left", grp)

        self.actionAlignJustify = QAction(
                QIcon.fromTheme('format-justify-fill',
                        QIcon(rsrcPath + '/textjustify.png')),
                "&Justify", grp)

        self.actionAlignLeft.setShortcut(Qt.CTRL + Qt.Key_L)
        self.actionAlignLeft.setCheckable(True)
        self.actionAlignLeft.setPriority(QAction.LowPriority)

        self.actionAlignCenter.setShortcut(Qt.CTRL + Qt.Key_E)
        self.actionAlignCenter.setCheckable(True)
        self.actionAlignCenter.setPriority(QAction.LowPriority)

        self.actionAlignRight.setShortcut(Qt.CTRL + Qt.Key_R)
        self.actionAlignRight.setCheckable(True)
        self.actionAlignRight.setPriority(QAction.LowPriority)

        self.actionAlignJustify.setShortcut(Qt.CTRL + Qt.Key_J)
        self.actionAlignJustify.setCheckable(True)
        self.actionAlignJustify.setPriority(QAction.LowPriority)

        tb.addActions(grp.actions())
        menu.addActions(grp.actions())
        menu.addSeparator()

        pix = QPixmap(16, 16)
        pix.fill(Qt.black)
        self.actionTextColor = QAction(QIcon(pix), "&Color...", self,
                triggered=self.textColor)
        tb.addAction(self.actionTextColor)
        menu.addAction(self.actionTextColor)

        tb = QToolBar(self)
        tb.setAllowedAreas(Qt.TopToolBarArea | Qt.BottomToolBarArea)
        tb.setWindowTitle("Format Actions")
        self.addToolBarBreak(Qt.TopToolBarArea)
        self.addToolBar(tb)

        comboStyle = QComboBox(tb)
        tb.addWidget(comboStyle)
        comboStyle.addItem("Standard")
        comboStyle.addItem("Bullet List (Disc)")
        comboStyle.addItem("Bullet List (Circle)")
        comboStyle.addItem("Bullet List (Square)")
        comboStyle.addItem("Ordered List (Decimal)")
        comboStyle.addItem("Ordered List (Alpha lower)")
        comboStyle.addItem("Ordered List (Alpha upper)")
        comboStyle.addItem("Ordered List (Roman lower)")
        comboStyle.addItem("Ordered List (Roman upper)")
        comboStyle.activated.connect(self.textStyle)

        self.comboFont = QFontComboBox(tb)
        tb.addWidget(self.comboFont)
        self.comboFont.activated[str].connect(self.textFamily)

        self.comboSize = QComboBox(tb)
        self.comboSize.setObjectName("comboSize")
        tb.addWidget(self.comboSize)
        self.comboSize.setEditable(True)

        db = QFontDatabase()
        for size in db.standardSizes():
            self.comboSize.addItem("%s" % (size))

        self.comboSize.activated[str].connect(self.textSize)
        self.comboSize.setCurrentIndex(
                self.comboSize.findText(
                        "%s" % (QApplication.font().pointSize())))

    def load(self, f):
        if not QFile.exists(f):
            return False

        fh = QFile(f)
        if not fh.open(QFile.ReadOnly):
            return False

        data = fh.readAll()
        codec = QTextCodec.codecForHtml(data)
        unistr = codec.toUnicode(data)

        if Qt.mightBeRichText(unistr):
            self.textEdit.setHtml(unistr)
        else:
            self.textEdit.setPlainText(unistr)

        self.setCurrentFileName(f)
        return True

    def maybeSave(self):
        if not self.textEdit.document().isModified():
            return True

        if self.fileName.startswith(':/'):
            return True

        ret = QMessageBox.warning(self, "Application",
                "The document has been modified.\n"
                "Do you want to save your changes?",
                QMessageBox.Save | QMessageBox.Discard | QMessageBox.Cancel)

        if ret == QMessageBox.Save:
            return self.fileSave()

        if ret == QMessageBox.Cancel:
            return False

        return True

    def setCurrentFileName(self, fileName=''):
        self.fileName = fileName
        self.textEdit.document().setModified(False)

        if not fileName:
            shownName = 'untitled.txt'
        else:
            shownName = QFileInfo(fileName).fileName()

        self.setWindowTitle(self.tr("%s[*] - %s" % (shownName, "Rich Text")))
        self.setWindowModified(False)

    def fileNew(self):
        if self.maybeSave():
            self.textEdit.clear()
            self.setCurrentFileName()

    def fileOpen(self):
        fn, _ = QFileDialog.getOpenFileName(self, "Open File...", None,
                "HTML-Files (*.htm *.html);;All Files (*)")

        if fn:
            self.load(fn)

    def fileSave(self):
        if not self.fileName:
            return self.fileSaveAs()

        writer = QTextDocumentWriter(self.fileName)
        success = writer.write(self.textEdit.document())
        if success:
            self.textEdit.document().setModified(False)

        return success

    def fileSaveAs(self):
        fn, _ = QFileDialog.getSaveFileName(self, "Save as...", None,
                "ODF files (*.odt);;HTML-Files (*.htm *.html);;All Files (*)")

        if not fn:
            return False

        lfn = fn.lower()
        if not lfn.endswith(('.odt', '.htm', '.html')):
            # The default.
            fn += '.odt'

        self.setCurrentFileName(fn)
        return self.fileSave()

    def filePrint(self):
        printer = QPrinter(QPrinter.HighResolution)
        dlg = QPrintDialog(printer, self)

        if self.textEdit.textCursor().hasSelection():
            dlg.addEnabledOption(QPrintDialog.PrintSelection)

        dlg.setWindowTitle("Print Document")

        if dlg.exec_() == QPrintDialog.Accepted:
            self.textEdit.print_(printer)

        del dlg

    def filePrintPreview(self):
        printer = QPrinter(QPrinter.HighResolution)
        preview = QPrintPreviewDialog(printer, self)
        preview.paintRequested.connect(self.printPreview)
        preview.exec_()

    def printPreview(self, printer):
        self.textEdit.print_(printer)

    def filePrintPdf(self):
        fn, _ = QFileDialog.getSaveFileName(self, "Export PDF", None,
                "PDF files (*.pdf);;All Files (*)")

        if fn:
            if QFileInfo(fn).suffix().isEmpty():
                fn += '.pdf'

            printer = QPrinter(QPrinter.HighResolution)
            printer.setOutputFormat(QPrinter.PdfFormat)
            printer.setOutputFileName(fn)
            self.textEdit.document().print_(printer)

    def textBold(self):
        fmt = QTextCharFormat()
        fmt.setFontWeight(self.actionTextBold.isChecked() and QFont.Bold or QFont.Normal)
        self.mergeFormatOnWordOrSelection(fmt)

    def textUnderline(self):
        fmt = QTextCharFormat()
        fmt.setFontUnderline(self.actionTextUnderline.isChecked())
        self.mergeFormatOnWordOrSelection(fmt)

    def textItalic(self):
        fmt = QTextCharFormat()
        fmt.setFontItalic(self.actionTextItalic.isChecked())
        self.mergeFormatOnWordOrSelection(fmt)

    def textFamily(self, family):
        fmt = QTextCharFormat()
        fmt.setFontFamily(family)
        self.mergeFormatOnWordOrSelection(fmt)

    def textSize(self, pointSize):
        pointSize = float(pointSize)
        if pointSize > 0:
            fmt = QTextCharFormat()
            fmt.setFontPointSize(pointSize)
            self.mergeFormatOnWordOrSelection(fmt)

    def textStyle(self, styleIndex):
        cursor = self.textEdit.textCursor()
        if styleIndex:
            styleDict = {
                1: QTextListFormat.ListDisc,
                2: QTextListFormat.ListCircle,
                3: QTextListFormat.ListSquare,
                4: QTextListFormat.ListDecimal,
                5: QTextListFormat.ListLowerAlpha,
                6: QTextListFormat.ListUpperAlpha,
                7: QTextListFormat.ListLowerRoman,
                8: QTextListFormat.ListUpperRoman,
            }

            style = styleDict.get(styleIndex, QTextListFormat.ListDisc)
            cursor.beginEditBlock()
            blockFmt = cursor.blockFormat()
            listFmt = QTextListFormat()

            if cursor.currentList():
                listFmt = cursor.currentList().format()
            else:
                listFmt.setIndent(blockFmt.indent() + 1)
                blockFmt.setIndent(0)
                cursor.setBlockFormat(blockFmt)

            listFmt.setStyle(style)
            cursor.createList(listFmt)
            cursor.endEditBlock()
        else:
            bfmt = QTextBlockFormat()
            bfmt.setObjectIndex(-1)
            cursor.mergeBlockFormat(bfmt)

    def textColor(self):
        col = QColorDialog.getColor(self.textEdit.textColor(), self)
        if not col.isValid():
            return

        fmt = QTextCharFormat()
        fmt.setForeground(col)
        self.mergeFormatOnWordOrSelection(fmt)
        self.colorChanged(col)

    def textAlign(self, action):
        if action == self.actionAlignLeft:
            self.textEdit.setAlignment(Qt.AlignLeft | Qt.AlignAbsolute)
        elif action == self.actionAlignCenter:
            self.textEdit.setAlignment(Qt.AlignHCenter)
        elif action == self.actionAlignRight:
            self.textEdit.setAlignment(Qt.AlignRight | Qt.AlignAbsolute)
        elif action == self.actionAlignJustify:
            self.textEdit.setAlignment(Qt.AlignJustify)

    def currentCharFormatChanged(self, format):
        self.fontChanged(format.font())
        self.colorChanged(format.foreground().color())

    def cursorPositionChanged(self):
        self.alignmentChanged(self.textEdit.alignment())

    def clipboardDataChanged(self):
        self.actionPaste.setEnabled(len(QApplication.clipboard().text()) != 0)

    def about(self):
        QMessageBox.about(self, "About",
                "This example demonstrates Qt's rich text editing facilities "
                "in action, providing an example document for you to "
                "experiment with.")

    def mergeFormatOnWordOrSelection(self, format):
        cursor = self.textEdit.textCursor()
        if not cursor.hasSelection():
            cursor.select(QTextCursor.WordUnderCursor)

        cursor.mergeCharFormat(format)
        self.textEdit.mergeCurrentCharFormat(format)

    def fontChanged(self, font):
        self.comboFont.setCurrentIndex(
                self.comboFont.findText(QFontInfo(font).family()))
        self.comboSize.setCurrentIndex(
                self.comboSize.findText("%s" % font.pointSize()))
        self.actionTextBold.setChecked(font.bold())
        self.actionTextItalic.setChecked(font.italic())
        self.actionTextUnderline.setChecked(font.underline())

    def colorChanged(self, color):
        pix = QPixmap(16, 16)
        pix.fill(color)
        self.actionTextColor.setIcon(QIcon(pix))

    def alignmentChanged(self, alignment):
        if alignment & Qt.AlignLeft:
            self.actionAlignLeft.setChecked(True)
        elif alignment & Qt.AlignHCenter:
            self.actionAlignCenter.setChecked(True)
        elif alignment & Qt.AlignRight:
            self.actionAlignRight.setChecked(True)
        elif alignment & Qt.AlignJustify:
            self.actionAlignJustify.setChecked(True)

if __name__ == '__main__':
    app = QApplication(sys.argv)

    mainWindows = []
    for fn in sys.argv[1:] or [None]:
        textEdit = TextEdit(fn)
        textEdit.resize(700, 800)
        textEdit.show()
        mainWindows.append(textEdit)

    sys.exit(app.exec_())


import logging
import torch
import argparse
import time
import multiprocessing
import os
import numpy as np
import math

from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
from torch.utils.data.distributed import DistributedSampler
from transformers import AdamW, get_linear_schedule_with_warmup
from tqdm import tqdm

from configs import add_args, set_dist, set_seed, set_hyperparas
from models import bulid_or_load_gen_model,bulid_or_load_cls_model
from utils import get_filenames, get_elapse_time, load_and_cache_gen_data, load_and_cache_defect_data,load_and_cache_clone_data, get_lang_by_task
from evaluator import smooth_bleu
from evaluator.CodeBLEU import calc_code_bleu
from evaluator.bleu import _bleu
import sys
from sklearn.metrics import recall_score, precision_score, f1_score
from tree_sitter import Language, Parser
from utils import retrieve2file
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

def evaluate_cls(args, model, eval_examples, eval_data, write_to_pred=False):
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.dev_batch_size)

    # Eval!
    if write_to_pred == False:
        logger.info("***** Running evaluation *****")
        logger.info("  Num examples = %d", len(eval_examples))
        logger.info("  Num batches = %d", len(eval_dataloader))
        logger.info("  Batch size = %d", args.dev_batch_size)
    eval_loss = 0.0
    nb_eval_steps = 0
    model.eval()
    logits = []
    labels = []
    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Evaluating"):
        if args.sub_task == "POJ":
            inputs = batch[0].to(args.device)
            p_inputs = batch[1].to(args.device)
            n_inputs = batch[2].to(args.device)
            label = batch[3].to(args.device)
        else:
            inputs = batch[0].to(args.device) # inputs shape: [batch_size, args.max_source_length+args.max_target_length]
            label = batch[1].to(args.device) # label shape: [batch_size]
        with torch.no_grad():
            if args.sub_task == "POJ":
                lm_loss, logit = model(inputs,p_inputs,n_inputs,label)
            else:
                lm_loss, logit = model(inputs, label) # logit shape:[batch_size]

            eval_loss += lm_loss.mean().item()
            logits.append(logit.cpu().numpy())#logit shape: [batch_size]
            labels.append(label.cpu().numpy())#label shape: [batch_size]
        nb_eval_steps += 1
    logits = np.concatenate(logits, 0)
    labels = np.concatenate(labels, 0)
    if args.model_name == "unixcoder" and args.task == "clone":
        preds = logits > 0.5
    else:
        preds = logits[:, 1] > 0.5
    if args.task == 'defect':
        eval_acc = np.mean(labels == preds)
        eval_loss = eval_loss / nb_eval_steps
        perplexity = torch.tensor(eval_loss)

        result = {
            "eval_loss": float(perplexity),
            "eval_acc": round(eval_acc, 4) * 100,
        }
    elif args.task == 'clone':
        recall = recall_score(labels, preds)
        precision = precision_score(labels, preds)
        f1 = f1_score(labels, preds)
        result = {
            "eval_recall": float(recall) * 100,
            "eval_precision": float(precision) * 100,
            "eval_f1": float(f1) * 100,
            "eval_threshold": 0.5,
        }
    if write_to_pred == False:
        logger.info("***** Eval results *****")
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(round(result[key], 4)))

    if write_to_pred:
        with open(os.path.join(args.output_dir, "predictions.txt"), 'w') as f:
            for example, pred in zip(eval_examples, preds):
                if args.task == 'defect':
                    if args.model_name == "unixcoder":
                        if pred:
                            f.write(str(example.idx) + '\t1\n')
                        else:
                            f.write(str(example.idx) + '\t0\n')
                    else:
                        if pred:
                            f.write(str(example.idx) + '\t1\n')
                        else:
                            f.write(str(example.idx) + '\t0\n')
                elif args.task == 'clone':
                    if pred:
                        f.write(example.url1 + '\t' + example.url2 + '\t' + '1' + '\n')
                    else:
                        f.write(example.url1 + '\t' + example.url2 + '\t' + '0' + '\n')

    return result

def eval_ppl_epoch(args, eval_data, eval_examples, model, tokenizer):
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.dev_batch_size,
                                 num_workers=4, pin_memory=True)
    # Start evaluating model
    logger.info("  " + "***** Running ppl evaluation *****")
    logger.info("  Num examples = %d", len(eval_examples))
    logger.info("  Batch size = %d", args.dev_batch_size)

    model.eval()
    eval_loss, batch_num = 0, 0
    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval ppl"):
        batch = tuple(t.to(args.device) for t in batch)
        source_ids, target_ids = batch
        source_mask = source_ids.ne(tokenizer.pad_token_id)
        target_mask = target_ids.ne(tokenizer.pad_token_id)

        with torch.no_grad():
            if args.model_name in ['roberta', 'codebert', 'graphcodebert']:
#                     loss, _, _, attention = model(source_ids=source_ids, source_mask=source_mask,
#                                           target_ids=target_ids, target_mask=target_mask)
                loss, _, _, _ = model(source_ids=source_ids, source_mask=source_mask,
                                    target_ids=target_ids, target_mask=target_mask)
                if args.n_gpu > 1:
                    loss = loss.mean()
                eval_loss += loss.item()
                batch_num += 1
            elif args.model_name in ['unixcoder']:
                _,loss,num = model(source_ids=source_ids,target_ids=target_ids)
                if args.n_gpu > 1:
                    loss = loss.mean()
                eval_loss += loss.sum().item()
                batch_num += num.sum().item()
            else:
                outputs = model(input_ids=source_ids, attention_mask=source_mask,
                                labels=target_ids, decoder_attention_mask=target_mask)
                if isinstance(outputs,dict):
                    loss=outputs['loss']
                else:
                    loss = outputs.loss
                if args.n_gpu > 1:
                    loss = loss.mean()
                eval_loss += loss.item()
                batch_num += 1

    eval_loss = eval_loss / batch_num
    eval_ppl = round(np.exp(eval_loss), 5)
    return eval_ppl

def eval_bleu_epoch(args, eval_data, eval_examples, model, tokenizer, split_tag, criteria):
    logger.info(
        "  ***** Running bleu evaluation on {} data*****".format(split_tag))
    logger.info("  Num examples = %d", len(eval_examples))
    if split_tag == 'dev':
        batch_size = args.dev_batch_size
    elif split_tag == 'test':
        batch_size = args.test_batch_size
    else:
        batch_size = args.batch_size
    logger.info("  Batch size = %d", batch_size)
    eval_sampler = SequentialSampler(eval_data)
    if args.data_num == -1:
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=batch_size,
                                     num_workers=4, pin_memory=True)
    else:
        eval_dataloader = DataLoader(
            eval_data, sampler=eval_sampler, batch_size=batch_size)

    model.eval()
    pred_ids = []
    bleu, codebleu = 0.0, 0.0
    for batch in tqdm(eval_dataloader, total=len(eval_dataloader), desc="Eval bleu for {} set".format(split_tag)):
        source_ids = batch[0].to(args.device) #shape: (batch_size, max_source_len)
        source_mask = source_ids.ne(tokenizer.pad_token_id)
        if hasattr(model, 'module'):
            model = model.module # extract the model from the DataParallel wrapper
        with torch.no_grad():
            if args.model_name in ['roberta', 'codebert', 'graphcodebert']:
                preds, _ = model(source_ids=source_ids,
                                 source_mask=source_mask)
                top_preds = [pred[0].cpu().numpy() for pred in preds]
            elif args.model_name in ['unixcoder']:
                preds = model(source_ids=source_ids)  # preds shape: [batch_size, self.beam_size, max_target_len]
                top_preds = [pred[0].cpu().numpy() for pred in preds]# top_preds shape: batch_size * [max_target_len]
            else:
                preds = model.generate(source_ids,
                                       attention_mask=source_mask,
                                       use_cache=True,
                                       num_beams=args.beam_size,
                                       early_stopping=args.task == 'summarize',
                                       max_length=args.max_target_length)
                top_preds = list(preds.cpu().numpy())
            pred_ids.extend(top_preds)
    # pdb.set_trace()
    pred_nls = [tokenizer.decode(
        id, skip_special_tokens=True, clean_up_tokenization_spaces=False) for id in pred_ids]
    # unixcoder in fewshot will generate '\n' in small batch, and gradually disappear
    if args.model_name in ['unixcoder']:
        pred_nls = [id.replace('\n',' ').replace("        "," ").replace("    "," ").replace("\t"," ") for id in pred_nls]
    output_fn = os.path.join(args.res_dir, "test_{}.output".format(criteria))
    gold_fn = os.path.join(args.res_dir, "test_{}.gold".format(criteria))
    src_fn = os.path.join(args.res_dir, "test_{}.src".format(criteria))

    if args.task in ['defect']:
        target_dict = {0: 'false', 1: 'true'}
        golds = [target_dict[ex.target] for ex in eval_examples]
        eval_acc = np.mean([int(p == g) for p, g in zip(pred_nls, golds)])
        result = {'em': eval_acc, 'bleu': 0, 'codebleu': 0}

        with open(output_fn, 'w',encoding='utf-8') as f, open(gold_fn, 'w',encoding='utf-8') as f1, open(src_fn, 'w',encoding='utf-8') as f2:
            for pred_nl, gold in zip(pred_nls, eval_examples):
                f.write(pred_nl.strip() + '\n')
                f1.write(target_dict[gold.target] + '\n')
                f2.write(gold.source.strip() + '\n')
            logger.info("Save the predictions into %s", output_fn)
    else:
        dev_accs, predictions = [], []
        with open(output_fn, 'w',encoding='utf-8') as f, open(gold_fn, 'w',encoding='utf-8') as f1, open(src_fn, 'w',encoding='utf-8') as f2:
            for pred_nl, gold in zip(pred_nls, eval_examples):
                dev_accs.append(pred_nl.strip() == gold.target.strip())
                if args.task in ['summarize']:
                    predictions.append(str(gold.idx) + '\t' + pred_nl)
                    f.write(str(gold.idx) + '\t' +
                            pred_nl.strip().encode('utf8').decode() + '\n')
                    f1.write(str(gold.idx) + '\t' +
                             gold.target.strip().encode('utf8').decode() + '\n')
                    f2.write(str(gold.idx) + '\t' +
                             gold.source.strip().encode('utf8').decode() + '\n')
                else:
                    f.write(pred_nl.strip().encode('utf8').decode() + '\n')
                    f1.write(gold.target.strip().encode(
                        'utf8').decode() + '\n')
                    f2.write(gold.source.strip().encode(
                        'utf8').decode() + '\n')
        if args.task in ['summarize']:
            (goldMap, predictionMap) = smooth_bleu.computeMaps(predictions, gold_fn)
            bleu = round(smooth_bleu.bleuFromMaps(
                goldMap, predictionMap)[0], 2)
        else:
            bleu = round(_bleu(gold_fn, output_fn), 2)
            if split_tag == 'test' and args.task in ['refine', 'translate', 'generate' , 'clone']:
                args.lang = get_lang_by_task(args.task, args.sub_task)
                codebleu = calc_code_bleu.get_codebleu(
                    gold_fn, output_fn, args.lang,args=args)
        # except:
        #     bleu = 0.0
        #     codebleu = 0.0

        em = np.mean(dev_accs) * 100
        result = {'em': em, 'bleu': bleu}
        if not args.task == 'summarize' and split_tag == 'test':
            result['codebleu'] = codebleu * 100

    logger.info("***** Eval results *****")
    for key in sorted(result.keys()):
        logger.info("  %s = %s", key, str(round(result[key], 4)))

    return result

def main():
    parser = argparse.ArgumentParser()
    args = add_args(parser)
    t0 = time.time()

    set_dist(args)
    set_seed(args)
    set_hyperparas(args)

    logger.info(args)
    logger.info("************* args: *************")
    logger.info("args.model_name: "+str(args.model_name))
    logger.info("args.few_shot: "+str(args.few_shot))
    logger.info("args.task: "+str(args.task))
    logger.info("args.sub_task: "+str(args.sub_task))
    logger.info("*********************************")

    if args.task in ['summarize', 'translate', 'refine', 'generate','complete']:
        config, model, tokenizer = bulid_or_load_gen_model(args)
    elif args.task in ['defect','clone']:
        config, model, tokenizer = bulid_or_load_cls_model(args)
    model.to(args.device)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)
    pool = multiprocessing.Pool(args.cpu_count)
    args.train_filename, args.dev_filename, args.test_filename = get_filenames(
        args.data_dir, args.task, args.sub_task)
    fa = open(os.path.join(args.output_dir, 'summary.log'), 'a+',encoding='utf-8')

    if args.do_train:
        if args.local_rank in [-1, 0] and args.data_num == -1:
            summary_fn = '{}/{}'.format(args.summary_dir,
                                        '/'.join(args.output_dir.split('/')[1:]))
            tb_writer = SummaryWriter(summary_fn)
        # Prepare training data loader
        if args.task in ['summarize', 'translate', 'refine', 'generate','complete']:
            train_examples, train_data = load_and_cache_gen_data(
                args, args.train_filename, pool, tokenizer, 'train')
        elif args.task in ['defect']:
            train_examples, train_data = load_and_cache_defect_data(args, args.train_filename, pool, tokenizer, 'train')
        elif args.task in ['clone']:
            train_examples, train_data = load_and_cache_clone_data(args, args.train_filename, pool, tokenizer, 'train')
        if 0 and args.prefix_tuning:
            if args.retriever_mode in ['random', 'retrieve']:
                retrieve2file(args, train_examples, train_data)
        elif args.retriever_mode == 'old':
            args.prefix_dir = args.old_prefix_dir

        train_sampler = RandomSampler(
            train_data) if args.local_rank == -1 else DistributedSampler(train_data)
        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.batch_size,
                                      num_workers=4, pin_memory=True)
        # Prepare optimizer and schedule (linear warmup and decay)
        no_decay = ['bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
             'weight_decay': args.weight_decay},
            {'params': [p for n, p in model.named_parameters() if any(
                nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters,
                          lr=args.lr, eps=args.adam_epsilon)
        num_train_optimization_steps = args.num_train_epochs * \
            len(train_dataloader)
        save_steps = max(len(train_dataloader), 1)
        if not args.is_clone_sample and args.task in ['clone']:
            save_steps = save_steps//20
        if args.task in ['defect']:
            save_steps = save_steps//1#4
        scheduler = get_linear_schedule_with_warmup(optimizer,
                                                    num_warmup_steps=int(args.warmup_steps) if args.warmup_steps >= 1 else num_train_optimization_steps * args.warmup_steps,
                                                    num_training_steps=num_train_optimization_steps)

        # Start training
        train_example_num = len(train_data)
        logger.info("***** Running training *****")
        logger.info("  Num examples = %d", train_example_num)
        logger.info("  Batch size = %d", args.batch_size)
        logger.info("  Batch num = %d", math.ceil(
            train_example_num / args.batch_size))
        logger.info("  Num epoch = %d", args.num_train_epochs)

        if args.task in ['summarize', 'translate', 'refine', 'generate','complete']:
            dev_dataset = {}
            global_step, best_bleu_em, best_ppl = 0, -1, 1e6
            not_loss_dec_cnt, not_bleu_em_inc_cnt = 0, 0 if args.do_eval_bleu else 1e6

            for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
                bar = tqdm(train_dataloader, total=len(
                    train_dataloader), desc="Training")
                nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
                model.train()
                # kkk=0
                for step, batch in enumerate(bar):
                    # if kkk<2:
                    #     kkk+=1
                    # else:
                    #     break
                    batch = tuple(t.to(args.device) for t in batch)

                    source_ids, target_ids = batch
                    source_mask = source_ids.ne(tokenizer.pad_token_id)
                    target_mask = target_ids.ne(tokenizer.pad_token_id)

                    if args.model_name in ['roberta', 'codebert', 'graphcodebert']:
    #                     loss, _, _, attention = model(source_ids=source_ids, source_mask=source_mask,
    #                                           target_ids=target_ids, target_mask=target_mask)
                        loss, _, _, _ = model(source_ids=source_ids, source_mask=source_mask,
                                            target_ids=target_ids, target_mask=target_mask)
                    elif args.model_name in ['unixcoder']:
                        loss,_,_ = model(source_ids=source_ids,target_ids=target_ids)#mask strategy (e2d) written in Seq2seq
                    else:
                        outputs = model(input_ids=source_ids, attention_mask=source_mask,
                                        labels=target_ids, decoder_attention_mask=target_mask)
                        #attention =outputs.encoder_attentions
                        if isinstance(outputs,dict):
                            loss=outputs['loss']
                        else:
                            loss = outputs.loss

                    if args.n_gpu > 1:
                        loss = loss.mean()  # mean() to average on multi-gpu.
                    if args.gradient_accumulation_steps > 1:
                        loss = loss / args.gradient_accumulation_steps
                    tr_loss += loss.item()

                    nb_tr_examples += source_ids.size(0)
                    nb_tr_steps += 1
                    loss.backward()

                    if nb_tr_steps % args.gradient_accumulation_steps == 0:
                        # Update parameters
                        optimizer.step()
                        optimizer.zero_grad()
                        scheduler.step()
                        global_step += 1
                        train_loss = round(
                            tr_loss * args.gradient_accumulation_steps / (nb_tr_steps + 1), 4)
                        bar.set_description("[{}] Train loss {}".format(
                            cur_epoch, round(train_loss, 3)))
                        if args.data_num == -1:
                            tb_writer.add_scalar('train_loss_per_epoch', round(train_loss, 3), cur_epoch)
                            tb_writer.add_scalar('train_loss_per_step', round(train_loss, 3), global_step)
                    # result = {'epoch': cur_epoch,
                    #         'global_step': global_step, 'train_loss': round(train_loss, 3)}
                    # for key in sorted(result.keys()):
                    #     logger.info("  %s = %s", key, str(result[key]))

                if args.do_eval:
                    # Eval model with dev dataset
                    if 'dev_loss' in dev_dataset:
                        eval_examples, eval_data = dev_dataset['dev_loss']
                    else:
                        eval_examples, eval_data = load_and_cache_gen_data(
                            args, args.dev_filename, pool, tokenizer, 'dev')
                        dev_dataset['dev_loss'] = eval_examples, eval_data

                    eval_ppl = eval_ppl_epoch(
                        args, eval_data, eval_examples, model, tokenizer)
                    result = {'epoch': cur_epoch,
                            'global_step': global_step, 'eval_ppl': eval_ppl}
                    for key in sorted(result.keys()):
                        logger.info("  %s = %s", key, str(result[key]))
                    logger.info("  " + "*" * 20)
                    if args.data_num == -1:
                        tb_writer.add_scalar('dev_ppl', eval_ppl, cur_epoch)

                    # save last checkpoint
                    if args.save_last_checkpoints:
                        last_output_dir = os.path.join(
                            args.output_dir, 'checkpoint-last')
                        if not os.path.exists(last_output_dir):
                            os.makedirs(last_output_dir)
                        model_to_save = model.module if hasattr(
                            model, 'module') else model
                        output_model_file = os.path.join(
                            last_output_dir, "pytorch_model.bin")
                        torch.save(model_to_save.state_dict(), output_model_file)
                        logger.info("Save the last model into %s",
                                    output_model_file)

                    if eval_ppl < best_ppl:
                        not_loss_dec_cnt = 0
                        logger.info("  Best ppl:%s", eval_ppl)
                        logger.info("  " + "*" * 20)
                        fa.write("[%d] Best ppl changed into %.4f\n" %
                                (cur_epoch, eval_ppl))
                        best_ppl = eval_ppl

                        # Save best checkpoint for best ppl
                        output_dir = os.path.join(
                            args.output_dir, 'checkpoint-best-ppl')
                        if not os.path.exists(output_dir):
                            os.makedirs(output_dir)
                        if args.always_save_model:
                            model_to_save = model.module if hasattr(
                                model, 'module') else model
                            output_model_file = os.path.join(
                                output_dir, "pytorch_model.bin")
                            torch.save(model_to_save.state_dict(),
                                    output_model_file)
                            logger.info(
                                "Save the best ppl model into %s", output_model_file)
                    else:
                        not_loss_dec_cnt += 1
                        logger.info(
                            "Ppl does not decrease for %d epochs", not_loss_dec_cnt)
                        if all([x > args.patience for x in [not_bleu_em_inc_cnt, not_loss_dec_cnt]]):
                            early_stop_str = "[%d] Early stop as not_bleu_em_inc_cnt=%d, and not_loss_dec_cnt=%d\n" % (
                                cur_epoch, not_bleu_em_inc_cnt, not_loss_dec_cnt)
                            logger.info(early_stop_str)
                            fa.write(early_stop_str)
                            break
                    logger.info("***** CUDA.empty_cache() *****")
                    torch.cuda.empty_cache()
                    if args.do_eval_bleu:
                        eval_examples, eval_data = load_and_cache_gen_data(args, args.dev_filename, pool, tokenizer, 'dev',
                                                                        only_src=True, is_sample=True)

                        result = eval_bleu_epoch(
                            args, eval_data, eval_examples, model, tokenizer, 'dev', 'e%d' % cur_epoch)
                        dev_bleu, dev_em = result['bleu'], result['em']
                        if args.task in ['summarize']:
                            dev_bleu_em = dev_bleu
                        elif args.task in ['defect']:
                            dev_bleu_em = dev_em
                        else:
                            dev_bleu_em = dev_bleu + dev_em
                        if args.data_num == -1:
                            tb_writer.add_scalar(
                                'dev_bleu_em', dev_bleu_em, cur_epoch)
                            # tb_writer.add_scalar('dev_em', dev_em, cur_epoch)
                        if dev_bleu_em > best_bleu_em:
                            not_bleu_em_inc_cnt = 0
                            logger.info("  [%d] Best bleu+em: %.2f (bleu: %.2f, em: %.2f)",
                                        cur_epoch, dev_bleu_em, dev_bleu, dev_em)
                            logger.info("  " + "*" * 20)
                            best_bleu_em = dev_bleu_em
                            fa.write("[%d] Best bleu+em changed into %.2f (bleu: %.2f, em: %.2f)\n" % (
                                cur_epoch, best_bleu_em, dev_bleu, dev_em))
                            # Save best checkpoint for best bleu
                            output_dir = os.path.join(
                                args.output_dir, 'checkpoint-best-bleu')
                            if not os.path.exists(output_dir):
                                os.makedirs(output_dir)
                            if args.data_num == -1 or args.always_save_model:
                                model_to_save = model.module if hasattr(
                                    model, 'module') else model
                                output_model_file = os.path.join(
                                    output_dir, "pytorch_model.bin")
                                torch.save(model_to_save.state_dict(),
                                        output_model_file)
                                logger.info(
                                    "Save the best bleu model into %s", output_model_file)
                        else:
                            not_bleu_em_inc_cnt += 1
                            logger.info(
                                "Bleu_em does not increase for %d epochs", not_bleu_em_inc_cnt)
                            if all([x > args.patience for x in [not_bleu_em_inc_cnt, not_loss_dec_cnt]]):
                                stop_early_str = "[%d] Early stop as not_bleu_em_inc_cnt=%d, and not_loss_dec_cnt=%d\n" % (
                                    cur_epoch, not_bleu_em_inc_cnt, not_loss_dec_cnt)
                                logger.info(stop_early_str)
                                fa.write(stop_early_str)
                                break
                logger.info("***** CUDA.empty_cache() *****")
                torch.cuda.empty_cache()

        elif args.task in ['defect']:
            global_step, best_acc = 0, 0
            not_acc_inc_cnt = 0
            is_early_stop = False
            for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
                bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
                nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
                model.train()
                for step, batch in enumerate(bar):
                    if args.sub_task == "POJ":
                        source_ids = batch[0].to(args.device)
                        p_inputs = batch[1].to(args.device)
                        n_inputs = batch[2].to(args.device)
                        labels = batch[3].to(args.device)
                        model.train()
                        loss,vec = model(source_ids,p_inputs,n_inputs,labels)
                    else:
                        if args.model_name == 'unixcoder':
                            source_ids = batch[0].to(args.device)
                            labels = batch[1].to(args.device)
                            loss, logits = model(source_ids, labels)
                        elif args.model_name in ['t5','codet5','bart','plbart']:
                            batch = tuple(t.to(args.device) for t in batch)
                            source_ids, labels = batch
                            outputs = model(source_ids, labels)
                            if isinstance(outputs,dict):
                                loss=outputs['loss']
                            else:
                                loss, logits = outputs
                        else:
                            batch = tuple(t.to(args.device) for t in batch)
                            source_ids, labels = batch
                            loss, logits = model(source_ids, labels)

                    if args.n_gpu > 1:
                        loss = loss.mean()  # mean() to average on multi-gpu.
                    if args.gradient_accumulation_steps > 1:
                        loss = loss / args.gradient_accumulation_steps
                    tr_loss += loss.item()

                    nb_tr_examples += source_ids.size(0)
                    nb_tr_steps += 1
                    loss.backward()
                    # torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                    if nb_tr_steps % args.gradient_accumulation_steps == 0:
                        # Update parameters
                        optimizer.step()
                        optimizer.zero_grad()
                        scheduler.step()
                        global_step += 1
                        train_loss = round(tr_loss * args.gradient_accumulation_steps / nb_tr_steps, 4)
                        bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
                        if args.data_num == -1:
                            tb_writer.add_scalar('train_loss_per_epoch', round(train_loss, 3), cur_epoch)
                            tb_writer.add_scalar('train_loss_per_step', round(train_loss, 3), global_step)

                    if (step + 1) % save_steps == 0 and args.do_eval:
                        logger.info("***** CUDA.empty_cache() *****")
                        torch.cuda.empty_cache()

                        eval_examples, eval_data = load_and_cache_defect_data(args, args.dev_filename, pool, tokenizer,
                                                                            'dev',is_sample=True)

                        result = evaluate_cls(args, model, eval_examples, eval_data)
                        eval_acc = result['eval_acc']

                        if args.data_num == -1:
                            tb_writer.add_scalar('dev_acc', round(eval_acc, 4), len(train_dataloader)//save_steps*cur_epoch+(step + 1) // save_steps-1)

                        # save last checkpoint
                        last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
                        if not os.path.exists(last_output_dir):
                            os.makedirs(last_output_dir)

                        if True or args.data_num == -1 and args.save_last_checkpoints:
                            model_to_save = model.module if hasattr(model, 'module') else model
                            output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
                            torch.save(model_to_save.state_dict(), output_model_file)
                            logger.info("Save the last model into %s", output_model_file)

                        if eval_acc > best_acc:
                            not_acc_inc_cnt = 0
                            logger.info("  Best acc: %s", round(eval_acc, 4))
                            logger.info("  " + "*" * 20)
                            fa.write("[%d] Best acc changed into %.4f\n" % (cur_epoch, round(eval_acc, 4)))
                            best_acc = eval_acc
                            # Save best checkpoint for best ppl
                            output_dir = os.path.join(args.output_dir, 'checkpoint-best-acc')
                            if not os.path.exists(output_dir):
                                os.makedirs(output_dir)
                            if args.data_num == -1 or True:
                                model_to_save = model.module if hasattr(model, 'module') else model
                                output_model_file = os.path.join(output_dir, "pytorch_model.bin")
                                torch.save(model_to_save.state_dict(), output_model_file)
                                logger.info("Save the best ppl model into %s", output_model_file)
                        else:
                            not_acc_inc_cnt += 1
                            logger.info("acc does not increase for %d epochs", not_acc_inc_cnt)
                            if not_acc_inc_cnt > args.patience*len(train_dataloader)//save_steps:
                                logger.info("Early stop as acc do not increase for %d times", not_acc_inc_cnt)
                                fa.write("[%d] Early stop as not_acc_inc_cnt=%d\n" % (cur_epoch, not_acc_inc_cnt))
                                is_early_stop = True
                                break

                    model.train()
                if is_early_stop:
                    break

                logger.info("***** CUDA.empty_cache() *****")
                torch.cuda.empty_cache()
        elif args.task in ['clone']:
            global_step, best_f1 = 0, 0
            not_f1_inc_cnt = 0
            is_early_stop = False
            for cur_epoch in range(args.start_epoch, int(args.num_train_epochs)):
                bar = tqdm(train_dataloader, total=len(train_dataloader), desc="Training")
                nb_tr_examples, nb_tr_steps, tr_loss = 0, 0, 0
                model.train()
                for step, batch in enumerate(bar):
                    if args.sub_task == "POJ":
                        source_ids = batch[0].to(args.device)
                        p_inputs = batch[1].to(args.device)
                        n_inputs = batch[2].to(args.device)
                        labels = batch[3].to(args.device)
                        loss,vec = model(source_ids,p_inputs,n_inputs,labels)
                    else:
                        if args.model_name == 'unixcoder':
                            source_ids = batch[0].to(args.device)
                            labels = batch[1].to(args.device)
                            loss, logits = model(source_ids, labels)
                        elif args.model_name in ['t5','codet5','bart','plbart']:
                            batch = tuple(t.to(args.device) for t in batch)
                            source_ids, labels = batch
                            outputs = model(source_ids, labels=labels)
                            if isinstance(outputs,dict):
                                loss=outputs['loss']
                            else:
                                loss, logits = outputs
                        else:
                            batch = tuple(t.to(args.device) for t in batch)#shape: (2*(batch_size, seq_len))
                            source_ids, labels = batch #[batch,1024] [batch]
                            loss, logits = model(source_ids, labels)

                    if args.n_gpu > 1:
                        loss = loss.mean()  # mean() to average on multi-gpu.
                    if args.gradient_accumulation_steps > 1:
                        loss = loss / args.gradient_accumulation_steps
                    tr_loss += loss.item()

                    nb_tr_examples += source_ids.size(0)
                    nb_tr_steps += 1
                    loss.backward()
                    # torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                    if nb_tr_steps % args.gradient_accumulation_steps == 0:
                        # Update parameters
                        optimizer.step()
                        optimizer.zero_grad()
                        scheduler.step()
                        global_step += 1
                        train_loss = round(tr_loss * args.gradient_accumulation_steps / nb_tr_steps, 4)
                        bar.set_description("[{}] Train loss {}".format(cur_epoch, round(train_loss, 3)))
                        if args.data_num == -1:
                            tb_writer.add_scalar('train_loss_per_epoch', round(train_loss, 3), cur_epoch)
                            tb_writer.add_scalar('train_loss_per_step', round(train_loss, 3), global_step)

                    if (step + 1) % save_steps == 0 and args.do_eval:
                        logger.info("***** CUDA.empty_cache() *****")
                        torch.cuda.empty_cache()

                        eval_examples, eval_data = load_and_cache_clone_data(args, args.dev_filename, pool, tokenizer,
                                                                            'dev', is_sample=True)

                        result = evaluate_cls(args, model, eval_examples, eval_data)
                        eval_f1 = result['eval_f1']

                        if args.data_num == -1:
                            tb_writer.add_scalar('dev_f1', round(eval_f1, 4), len(train_dataloader)//save_steps*cur_epoch+(step + 1) // save_steps-1)

                        # save last checkpoint
                        last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
                        if not os.path.exists(last_output_dir):
                            os.makedirs(last_output_dir)

                        if True or args.data_num == -1 and args.save_last_checkpoints:
                            model_to_save = model.module if hasattr(model, 'module') else model
                            output_model_file = os.path.join(last_output_dir, "pytorch_model.bin")
                            torch.save(model_to_save.state_dict(), output_model_file)
                            logger.info("Save the last model into %s", output_model_file)

                        if eval_f1 > best_f1:
                            not_f1_inc_cnt = 0
                            logger.info("  Best f1: %s", round(eval_f1, 4))
                            logger.info("  " + "*" * 20)
                            fa.write("[%d] Best f1 changed into %.4f\n" % (cur_epoch, round(eval_f1, 4)))
                            best_f1 = eval_f1
                            # Save best checkpoint for best ppl
                            output_dir = os.path.join(args.output_dir, 'checkpoint-best-f1')
                            if not os.path.exists(output_dir):
                                os.makedirs(output_dir)
                            if args.data_num == -1 or True:
                                model_to_save = model.module if hasattr(model, 'module') else model
                                output_model_file = os.path.join(output_dir, "pytorch_model.bin")
                                torch.save(model_to_save.state_dict(), output_model_file)
                                logger.info("Save the best ppl model into %s", output_model_file)
                        else:
                            not_f1_inc_cnt += 1
                            logger.info("F1 does not increase for %d epochs", not_f1_inc_cnt)
                            if not_f1_inc_cnt > args.patience*len(train_dataloader)//save_steps:
                                logger.info("Early stop as f1 do not increase for %d times", not_f1_inc_cnt)
                                fa.write("[%d] Early stop as not_f1_inc_cnt=%d\n" % (cur_epoch, not_f1_inc_cnt))
                                is_early_stop = True
                                break

                    model.train()
                if is_early_stop:
                    break

                logger.info("***** CUDA.empty_cache() *****")
                torch.cuda.empty_cache()

        if args.local_rank in [-1, 0] and args.data_num == -1:
            tb_writer.close()
        logger.info("Finish training and take %s", get_elapse_time(t0))

    if args.do_test:

        logger.info("  " + "***** Testing *****")
        logger.info("  Batch size = %d", args.test_batch_size)
        if args.task in ['summarize', 'translate', 'refine', 'generate','complete']:
            for criteria in ['best-bleu', 'best-ppl']:  # 'best-bleu', 'best-ppl', 'last'
                file = os.path.join(
                    args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
                logger.info("Reload model from {}".format(file))

                model = model.module if hasattr(model, 'module') else model
                model.load_state_dict(torch.load(file))
                # if args.n_gpu > 1:
                #     # multi-gpu training
                #     model = torch.nn.DataParallel(model)
                # model.load_state_dict(torch.load(file))
                eval_examples, eval_data = load_and_cache_gen_data(args, args.test_filename, pool, tokenizer, 'test',
                                                                only_src=True, is_sample=False)
                result = eval_bleu_epoch(
                    args, eval_data, eval_examples, model, tokenizer, 'test', criteria)
                test_bleu, test_em = result['bleu'], result['em']
                test_codebleu = result['codebleu'] if 'codebleu' in result else 0
                result_str = "[%s] bleu-4: %.2f, em: %.4f, codebleu: %.4f\n" % (
                    criteria, test_bleu, test_em, test_codebleu)
                logger.info("*********************************")
                logger.info(result_str)
                logger.info("*********************************")
                fa.write(result_str)
                if args.res_fn:
                    with open(args.res_fn, 'a+',encoding='utf-8') as f:
                        f.write('[Time: {}] {}\n'.format(
                            get_elapse_time(t0), file))
                        f.write(result_str)
        elif args.task in ['defect']:
            for criteria in ['best-acc']:
                file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
                logger.info("Reload model from {}".format(file))

                model = model.module if hasattr(model, 'module') else model
                model.load_state_dict(torch.load(file))
                # if args.n_gpu > 1:
                #     # multi-gpu training
                #     model = torch.nn.DataParallel(model)

                eval_examples, eval_data = load_and_cache_defect_data(args, args.test_filename, pool, tokenizer, 'test',
                                                                    False)
                result = evaluate_cls(args, model, eval_examples, eval_data, write_to_pred=True)
                logger.info(" test_acc = %.4f" % ( result['eval_acc']))
                logger.info("*********************************")
                logger.info("[%s]  test_acc = %.4f" % (criteria, result['eval_acc']))
                logger.info("*********************************")

                fa.write("[%s] test-acc: %.4f\n" % (criteria, result['eval_acc']))
                if args.res_fn:
                    with open(args.res_fn, 'a+') as f:
                        f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
                        f.write("[%s] acc: %.4f\n\n" % (
                            criteria, result['eval_acc']))
        elif args.task in ['clone']:
            for criteria in ['best-f1']:
                file = os.path.join(args.output_dir, 'checkpoint-{}/pytorch_model.bin'.format(criteria))
                logger.info("Reload model from {}".format(file))

                model = model.module if hasattr(model, 'module') else model
                model.load_state_dict(torch.load(file))
                # if args.n_gpu > 1:
                #     # multi-gpu training
                #     model = torch.nn.DataParallel(model)
                # model.load_state_dict(torch.load(file))

                eval_examples, eval_data = load_and_cache_clone_data(args, args.test_filename, pool, tokenizer, 'test',
                                                                    False)

                result = evaluate_cls(args, model, eval_examples, eval_data, write_to_pred=True)

                logger.info("  test_f1 = %.4f", result['eval_f1'])
                logger.info("  test_prec = %.4f", result['eval_precision'])
                logger.info("  test_rec = %.4f", result['eval_recall'])
                logger.info("*********************************")
                logger.info("[%s] test-f1: %.4f, precision: %.4f, recall: %.4f\n" % (
                    criteria, result['eval_f1'], result['eval_precision'], result['eval_recall']))
                logger.info("*********************************")

                fa.write("[%s] test-f1: %.4f, precision: %.4f, recall: %.4f\n" % (
                    criteria, result['eval_f1'], result['eval_precision'], result['eval_recall']))
                if args.res_fn:
                    with open(args.res_fn, 'a+') as f:
                        f.write('[Time: {}] {}\n'.format(get_elapse_time(t0), file))
                        f.write("[%s] f1: %.4f, precision: %.4f, recall: %.4f\n\n" % (
                            criteria, result['eval_f1'], result['eval_precision'], result['eval_recall']))

    logger.info("Finish and take {}".format(get_elapse_time(t0)))
    fa.write("Finish and take {}".format(get_elapse_time(t0)))
    fa.close()

if __name__ == "__main__":
    main()