import os
import pyautogui
import time
import math
import cv2
import numpy as np
import ffmpeg
import argparse
import logging

from datetime import datetime, timedelta
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support import ui
from skimage.metrics import structural_similarity as ssim

pyautogui.FAILSAFE = False

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SKYROOM_ICON = (1272, 43)
START_RECORDING_ICON = (1136, 178)
BETWEEN_PAUSE_AND_STOP_ICON = (1196, 185)
RECORDING_TAB_ICON = (209, 17)
STOP_RECORDING_ICON = (1217, 180)
CENTER_OF_DOWNLOAD_BAR = (825, 861)
CLOSE_NOTIFICATION = (413, 83)

FAILURE_TEST_INTERVAL = timedelta(minutes=5)

def split_to_100bulks(arr):
    result = []
    for i in range(math.ceil(len(arr) / 100)):
        result.append(arr[i * 100:(i + 1) * 100])
    return result

def goto_class(driver):
    driver.find_element_by_id('btn_guest').click()

def is_skyroom_extension_open():
    return pyautogui.pixelMatchesColor(*SKYROOM_ICON, (225, 225, 225))

def is_tab_in_recording():
    return pyautogui.pixelMatchesColor(*RECORDING_TAB_ICON, (26, 115, 232))

def close_chrome_notification():
    pyautogui.click(*CLOSE_NOTIFICATION)

def open_skyroom_popup():
    for repeat_number in range(10):
        try:
            if not is_skyroom_extension_open():
                pyautogui.click(*SKYROOM_ICON)
                time.sleep(1)
            if not is_skyroom_extension_open():
                raise Exception(
                    'I clicked on skyroom icon but popup is not open! :|')

            break
        except Exception as e:
            logger.exception(e)

def close_skyroom_popup():
    for repeat_number in range(10):
        try:
            if is_skyroom_extension_open():
                pyautogui.click(*SKYROOM_ICON)
                time.sleep(1)
            if is_skyroom_extension_open():
                raise Exception(
                    'I clicked on skyroom icon but popup is not closed! :|')

            break
        except Exception as e:
            logger.exception(e)
    pyautogui.click(0, 0)

def force_refresh(driver):
    driver.execute_script('window.onbeforeunload = function() {};')
    time.sleep(0.5)
    driver.refresh()

def is_need_recording():
    return not os.path.exists('./force-stop-recording')

def main():
    parser = argparse.ArgumentParser(
        description='This command will record a VClass page')
    parser.add_argument('-u', '--url', type=str,
                        default='', help='URL of vclass')
    parser.add_argument('-d', '--duration', type=float,
                        default=0, help='Duration of class in minutes')
    parser.add_argument('-n', '--name', type=str,
                        default='', help='Name of downloads folder')
    parser.add_argument('-a', '--username', type=str,
                        default='ضبط کننده', help='Username of skyroom user')
    parser.add_argument('-e', '--encoding', type=str,
                        default='no-encode',
                        help='Encoding quality, see readme.md')
    parser.add_argument('-v', '--video', type=str,
                        default='',
                        help='Video to encode')

    args = parser.parse_args()

    SOURCE_DIR = os.path.abspath(os.path.dirname(__file__))
    BASE_DIR = os.path.dirname(SOURCE_DIR)

    if args.url != '' and args.video != '':
        print('You cant define the -v argument for recording')
        exit()

    if args.url != '':
        try:
            os.mkdir(os.path.join(
                BASE_DIR,
                'downloads',
                args.name))
        except FileExistsError:
            pass
        download_path = os.path.join(
            BASE_DIR,
            'downloads',
            args.name,
            datetime.now().strftime("%Y-%m-%d--%H-%M"))
        os.mkdir(download_path)
        chrome_options = Options()
        chrome_options.add_extension(os.path.join(SOURCE_DIR, 'skyroom.crx'))
        chrome_options.add_experimental_option("prefs", {
            "download.default_directory": download_path,
            "download.prompt_for_download": False,
            "download.directory_upgrade": True,
            "safebrowsing.enabled": True
        })
        chrome_options.add_argument("--no-sandbox")
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument("disable-infobars")

        logger.info('Opening google chrome')
        driver = None
        for retry_number in range(10):
            try:
                if driver:
                    logger.info('Driver is not none, close it.')
                    driver.close()
            except Exception as e:
                logger.exception(e)
            try:
                driver = webdriver.Chrome(options=chrome_options)

                driver.implicitly_wait(10)
                time.sleep(2)

                driver.maximize_window()
                time.sleep(1)
                driver.maximize_window()

                break
            except Exception as e:
                logger.exception(e)

        logger.info('Create new window for not closing browser')
        for retry_number in range(10):
            try:
                if len(driver.window_handles) < 2:
                    driver.execute_script(
                        "window.open('about:blank', 'tab2');")
                time.sleep(2 * (retry_number + 1))

                break
            except Exception as e:
                logger.exception(e)

        logger.info('Configure skyroom extension')
        for retry_number in range(10):
            try:
                if len(driver.window_handles) > 1:
                    driver.switch_to.window(driver.window_handles[0])
                    driver.get(
                        "chrome-extension://"
                        "pejdnafppnpfimpnipdkiidjancinenc/options.html")
                    ui.Select(driver.find_element_by_id(
                        'video_bitrate')).select_by_value('1000')
                    ui.Select(driver.find_element_by_id(
                        'audio_bitrate')).select_by_value('32')

                    driver.find_element_by_id('btnSave').click()
                time.sleep(2 * (retry_number + 1))

                break
            except Exception as e:
                logger.exception(e)

        logger.info('Open vclass')
        for retry_number in range(10):
            try:
                driver.switch_to.window(driver.window_handles[0])
                driver.get(args.url)
                time.sleep(2 * (retry_number + 1))

                break
            except Exception as e:
                logger.exception(e)

        logger.info('Login as guest')
        for retry_number in range(10):
            try:
                force_refresh(driver)
                time.sleep(5 * retry_number)
                goto_class(driver)
                driver.find_element_by_xpath("//input[@class='full-width']")
                driver.execute_script(
                    "document.querySelector('.dlg-nickname .full-width').value"
                    f" = '{args.username}';")
                driver.execute_script(
                    "document.querySelector('.dlg-nickname .btn').click();")

                break
            except Exception as e:
                logger.exception(e)

        logger.info('Start record')
        for retry_number in range(10):
            try:
                open_skyroom_popup()

                if not is_tab_in_recording():
                    if not pyautogui.pixelMatchesColor(
                            *START_RECORDING_ICON, (255, 0, 0)):
                        close_skyroom_popup()
                        raise Exception('Recording red color not found :|')

                    pyautogui.click(*START_RECORDING_ICON)
                    time.sleep(5)

                if not pyautogui.pixelMatchesColor(
                        *BETWEEN_PAUSE_AND_STOP_ICON, (245, 245, 245)):
                    raise Exception('I can not see pause and stop icon!')
                if not is_tab_in_recording():
                    raise Exception('Tab blue recording icon can not be seen!')

                close_skyroom_popup()

                break
            except Exception as e:
                logger.exception(e)

        close_chrome_notification()

        logger.info('Recording is started, watch for freeze detection!')
        end_time = datetime.now() + timedelta(minutes=args.duration)
        old_screenshot = cv2.imdecode(np.frombuffer(
            driver.get_screenshot_as_png(), np.uint8), -1)
        next_failure_test = datetime.now() + FAILURE_TEST_INTERVAL
        while datetime.now() < end_time and is_need_recording():
            if datetime.now() >= next_failure_test:
                for retry_number in range(10):
                    try:
                        cur_screenshot = cv2.imdecode(np.frombuffer(
                            driver.get_screenshot_as_png(), np.uint8), -1)
                        similarity = ssim(
                            old_screenshot, cur_screenshot, multichannel=True)
                        if similarity > 0.98:
                            logger.info(
                                'Screenshots are too similar, refresh!')
                            force_refresh(driver)
                            goto_class(driver)
                        old_screenshot = cur_screenshot
                        break
                    except Exception as e:
                        logger.exception(e)
                next_failure_test = datetime.now() + FAILURE_TEST_INTERVAL

            time.sleep(max(0, min(
                10,
                (end_time - datetime.now()).total_seconds()
            )))

        logger.info('Time is over, stop recording')
        for retry_number in range(10):
            try:
                time.sleep(5 * retry_number)
                if is_tab_in_recording():
                    open_skyroom_popup()

                    if not pyautogui.pixelMatchesColor(
                            *STOP_RECORDING_ICON, (0, 0, 0)):
                        raise Exception(
                            'I can not see stop recording black icon')
                    pyautogui.click(*STOP_RECORDING_ICON)
                    time.sleep(5)
                    close_skyroom_popup()

                if is_tab_in_recording():
                    raise Exception(
                        'I stop recording but tab is in recording :|')
                if len(driver.window_handles) < 2:
                    raise Exception('Recording window is not open :(')

                break
            except Exception as e:
                logger.exception(e)

        # download file
        for retry_number in range(10):
            try:
                time.sleep(5)

                driver.switch_to.window(driver.window_handles[1])
                time.sleep(0.5)

                driver.find_element_by_id('download').click()
                time.sleep(5 * (1 + retry_number))

                if not pyautogui.pixelMatchesColor(
                        *CENTER_OF_DOWNLOAD_BAR, (255, 255, 255)):
                    raise Exception('I can not see download bar')

                break
            except Exception as e:
                logger.exception(e)

        # close windows
        time.sleep(60 + 3 * args.duration)
        for retry_number in range(100):
            try:
                if len(os.listdir(download_path)) == 0 or not os.listdir(
                        download_path)[0].endswith('.webm'):
                    raise Exception('Downloaded file can not be found!')
                    time.sleep(10 * (1 + retry_number))
                break
            except Exception as e:
                logger.exception(e)

        for retry_number in range(100):
            try:
                driver.switch_to.window(driver.window_handles[0])
                driver.close()
                time.sleep(3)
                driver.switch_to.window(driver.window_handles[0])
                driver.close()
                break
            except Exception as e:
                logger.exception(e)

        webm_file = os.path.join(
            download_path,
            os.listdir(download_path)[0]
        )
        new_webm_file = os.path.join(download_path, 'video.webm')
        os.rename(webm_file, new_webm_file)
        output_mp4 = os.path.join(download_path, 'video.mp4')
    else:
        new_webm_file = args.video
        output_mp4 = os.path.join(os.path.dirname(args.video), 'video.mp4')

    ffmpeg_pipe = ffmpeg.input(new_webm_file)
    if args.encoding == "high":
        ffmpeg_pipe.output(
            output_mp4,
            **{'vcodec': 'libx264', 'acodec': 'aac', 'b:a': '128k', },
            crf='36',
            tune='stillimage',
            preset='veryfast',
            r='25',
            movflags='+faststart'
        ).run()
    elif args.encoding == "medium":
        ffmpeg_pipe.output(
            output_mp4,
            **{'vcodec': 'libx264', 'acodec': 'aac', 'b:a': '60k', },
            crf='36',
            tune='stillimage',
            preset='veryfast',
            s='852x480',
            r='20',
            movflags='+faststart'
        ).run()
    elif args.encoding == "low":
        ffmpeg_pipe.output(
            output_mp4,
            **{'vcodec': 'libx264', 'acodec': 'aac', 'b:a': '60k', },
            crf='36',
            tune='stillimage',
            preset='veryfast',
            s='640x360',
            r='15',
            movflags='+faststart'
        ).run()

if __name__ == "__main__":
    main()


#!/usr/bin/env python

import numpy as np
import openbabel as ob
import pybel as pb
import copy
from aqml.cheminfo import *
import aqml.cheminfo.openbabel.amon_f as cioaf
import aqml.cheminfo.openbabel.amon as cioa
import aqml.cheminfo.openbabel.obabel as cib
#import networkx as nx

def is_connected(g):
    nv = g.shape[0]
    ne = (g>0).sum()/2
    return ne - nv + 1 >= 0

class RawMol(object):

    def __init__(self, zs, coords, allow_radical=False, allow_charge=False):
        """
        special cases:
           SMILES          xyzfile (initial)     xyzfile updated
        1) C[NH+]=CC(=O)O  CN[C]C([O])O          CN[C]C(=O)O
        2) CC(=O)[O-]      CC([O])([O])          CC(=O)(=O)

        """
        self.zs = zs
        self.coords = coords
        iok = True
        if not allow_radical:
            if sum(zs)%2 != 0: iok = False; print ' **'
        g = cib.perceive_g(zs, coords)
        self.g = g

        can = None; m = None

        if iok:
          ic = is_connected(g)
          if not ic:
            iok = False
            print ' ** dissociated'
          else:
            cns = g.sum(axis=0)
            na = len(zs)
            chgs = np.zeros(na,np.int)
            tvs = -999*np.ones(na,np.int)
            dic = {1:1, 4:2,  5:3, 6:4, 7:3, 8:2, 9:1, \
                    13:3, 14:4, 32:4, 34:2, 50:4 }
            zsf = dic.keys()
            ias = np.arange(na).astype(np.int)
            ias_visited = []
            for ia in range(na):
                #print ' ia, tvs = ', ia, tvs
                zi = zs[ia]
                nb = (g[ia] > 0).sum()
                if zi in zsf:
                    if tvs[ia] < 0:
                        tvs[ia] = dic[zi]
                    if zi == 6 and cns[ia] == 1:
                        # e.g., -N$C (or -[N+]#[C-]
                        chgs[ia] = -1; ibs = ias[ g[ia] > 0 ]
                        ib = ibs[0]
                        assert ibs.shape[0] == 1 and zs[ib] == 7
                        #if ib not in ias_visited:
                        #ias_visited.append(ib)
                        tvs[ib] = 5; chgs[ib] = 1
                    elif zi == 7:
                        if cns[ia] == 4:
                            if not allow_charge:
                                iok = False; print ' *** '
                        elif cns[ia] == 3:
                            ibs = ias[ g[ia] > 0 ]
                            for ib in ibs:
                                if zs[ib] in [8,] and (g[ib]>0).sum() == 1:
                                    # R-N(=O)=O or R-N(=CR)=O
                                    tvs[ia] = 5; break
                        elif cns[ia] == 1:
                            #print '+++++'
                            ibs = ias[ g[ia] > 0 ]
                            ib = ibs[0]
                            if len(ibs) == 1 and zs[ib] == 7:
                                # C=N#N,
                                chgs[ia] = -1;
                                #if ib not in ias_visited:
                                chgs[ib] = 1; tvs[ib] = 5
                                #ias_visited.append( ib )
                                #print 'tvs = ', tvs
                elif zi in [15,33,51]:
                    assert nb in [3,4,5]
                    tvs[ia] = {3:3, 4:5, 5:5}[nb]
                elif zi in [16,34,52]:
                    assert nb in [2,3,4]
                    tvs[ia] = {2:2, 3:4, 4:6}[nb]
                elif zi in [17,35,53]:
                    tvs[ia] = 1 if nb == 1 else 7
                else:
                    raise '#ERROR: new element?'

            #print ' --   zs = ', zs
            #print ' --  tvs = ', tvs
            #print ' -- chgs = ', chgs
            #print 'iok = ', iok
            if iok:
                bosr = []
                cmg = cioaf.MG(bosr, zs, chgs, tvs, g, coords, use_bosr=False)
                #cmg = cioa.MG(bosr, zs, chgs, tvs, g, coords, use_bosr=False)
                cans, ms = cmg.update_m() #debug=True, icon=True)
                nm = len(cans)
                if nm == 0:
                    iok = False
                    print ' ++ found radical or charged species'
                elif nm == 1:
                    can = cans[0]; m = ms[0]
                else:
                    iok = False
                    print ' ++ more than one possiblities: ', cans

        self.iok = iok
        self.can = can
        self.m = m

if __name__ == '__main__':

    import os, sys
    import aqml.cheminfo.rw.xyz as cix

    # Note that file in `fs represent
    fs = sys.argv[1:]
    for f in fs:
        symbols, positions = cix.read_xyz(f)[0]
        numbers = [ atomic_numbers[si] for si in symbols ]
        obj = RawMol( np.array(numbers), np.array(positions) )
        if obj.iok:
            print f[:-4], obj.can
        else:
            print f[:-4], ' ** FAILED'


import torch
import torch.nn as nn
import argparse
import os
import json
import shutil
import numpy as np
from model import *
from dataset import *
from torch.utils.data import DataLoader
import torch.utils.data.sampler as torch_sampler
from evaluate_tDCF_asvspoof19 import compute_eer_and_tdcf
from loss import *
from collections import defaultdict
from tqdm import tqdm, trange
import random
from utils import *
import eval_metrics as em
from ECAPA_TDNN import *

torch.set_default_tensor_type(torch.FloatTensor)

def initParams():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('--seed', type=int, help="random number seed", default=688)

    # Data folder prepare
    parser.add_argument("-a", "--access_type", type=str, help="LA or PA", default='LA')
    parser.add_argument("-d", "--path_to_database", type=str, help="dataset path", default='/data/neil/DS_10283_3336/')
    parser.add_argument("-f", "--path_to_features", type=str, help="features path",
                        default='/data2/neil/ASVspoof2019LA/')

    parser.add_argument("-o", "--out_fold", type=str, help="output folder", required=True, default='./models/try/')

    parser.add_argument("--ratio", type=float, default=0.5,
                        help="ASVspoof ratio in a training batch, the other should be augmented")

    # Dataset prepare
    parser.add_argument("--feat", type=str, help="which feature to use", default='LFCC',
                        choices=["CQCC", "LFCC"])
    parser.add_argument("--feat_len", type=int, help="features length", default=750)
    parser.add_argument('--pad_chop', type=str2bool, nargs='?', const=True, default=True, help="whether pad_chop in the dataset")
    parser.add_argument('--padding', type=str, default='repeat', choices=['zero', 'repeat', 'silence'],
                        help="how to pad short utterance")
    parser.add_argument("--enc_dim", type=int, help="encoding dimension", default=256)

    parser.add_argument('-m', '--model', help='Model arch', default='lcnn',
                        choices=['cnn', 'resnet', 'lcnn', 'res2net', 'ecapa'])

    # Training hyperparameters
    parser.add_argument('--num_epochs', type=int, default=200, help="Number of epochs for training")
    parser.add_argument('--batch_size', type=int, default=64, help="Mini batch size for training")
    parser.add_argument('--lr', type=float, default=0.0005, help="learning rate")
    parser.add_argument('--lr_decay', type=float, default=0.5, help="decay learning rate")
    parser.add_argument('--interval', type=int, default=30, help="interval to decay lr")

    parser.add_argument('--beta_1', type=float, default=0.9, help="bata_1 for Adam")
    parser.add_argument('--beta_2', type=float, default=0.999, help="beta_2 for Adam")
    parser.add_argument('--eps', type=float, default=1e-8, help="epsilon for Adam")
    parser.add_argument("--gpu", type=str, help="GPU index", default="1")
    parser.add_argument('--num_workers', type=int, default=0, help="number of workers")

    parser.add_argument('--base_loss', type=str, default="ce", choices=["ce", "bce"], help="use which loss for basic training")
    parser.add_argument('--add_loss', type=str, default=None,
                        choices=[None, 'isolate', 'ang_iso', 'p2sgrad'], help="add other loss for one-class training")
    parser.add_argument('--weight_loss', type=float, default=1, help="weight for other loss")
    parser.add_argument('--r_real', type=float, default=0.9, help="r_real for isolate loss")
    parser.add_argument('--r_fake', type=float, default=0.2, help="r_fake for isolate loss")
    parser.add_argument('--alpha', type=float, default=20, help="scale factor for angular isolate loss")
    parser.add_argument('--num_centers', type=int, default=3, help="num of centers for multi isolate loss")

    parser.add_argument('--visualize', action='store_true', help="feature visualization")
    parser.add_argument('--test_only', action='store_true', help="test the trained model in case the test crash sometimes or another test method")
    parser.add_argument('--continue_training', action='store_true', help="continue training with trained model")

    parser.add_argument('--ADV_AUG', type=str2bool, nargs='?', const=True, default=False,
                        help="whether to use adversarial augmentation in training")
    parser.add_argument('--LA_aug', type=str2bool, nargs='?', const=True, default=False,
                        help="whether to use LA_augmentation in training")
    parser.add_argument('--DF_aug', type=str2bool, nargs='?', const=True, default=False,
                        help="whether to use DF_augmentation in training")
    parser.add_argument('--LAPA_aug', type=str2bool, nargs='?', const=True, default=False,
                        help="whether to use LAPA_augmentation in training")
    parser.add_argument('--DFPA_aug', type=str2bool, nargs='?', const=True, default=False,
                        help="whether to use DFPA_augmentation in training")
    parser.add_argument('--lambda_', type=float, default=0.05, help="lambda for gradient reversal layer")
    parser.add_argument('--lr_d', type=float, default=0.0001, help="learning rate")

    parser.add_argument('--pre_train', action='store_true', help="whether to pretrain the model")
    parser.add_argument('--test_on_eval', action='store_true',
                        help="whether to run EER on the evaluation set")

    args = parser.parse_args()

    # Check ratio
    assert (args.ratio > 0) and (args.ratio <= 1)

    # Change this to specify GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # Set seeds
    setup_seed(args.seed)

    if args.test_only or args.continue_training:
        pass
    else:
        # Path for output data
        if not os.path.exists(args.out_fold):
            os.makedirs(args.out_fold)
        else:
            shutil.rmtree(args.out_fold)
            os.mkdir(args.out_fold)

        # Folder for intermediate results
        if not os.path.exists(os.path.join(args.out_fold, 'checkpoint')):
            os.makedirs(os.path.join(args.out_fold, 'checkpoint'))
        else:
            shutil.rmtree(os.path.join(args.out_fold, 'checkpoint'))
            os.mkdir(os.path.join(args.out_fold, 'checkpoint'))

        # Path for input data
        # assert os.path.exists(args.path_to_database)
        assert os.path.exists(args.path_to_features)

        # Save training arguments
        with open(os.path.join(args.out_fold, 'args.json'), 'w') as file:
            file.write(json.dumps(vars(args), sort_keys=True, separators=('\n', ':')))

        with open(os.path.join(args.out_fold, 'train_loss.log'), 'w') as file:
            file.write("Start recording training loss ...\n")
        with open(os.path.join(args.out_fold, 'dev_loss.log'), 'w') as file:
            file.write("Start recording validation loss ...\n")
        with open(os.path.join(args.out_fold, 'test_loss.log'), 'w') as file:
            file.write("Start recording test loss ...\n")

    args.cuda = torch.cuda.is_available()
    print('Cuda device available: ', args.cuda)
    args.device = torch.device("cuda" if args.cuda else "cpu")

    return args

def adjust_learning_rate(args, lr, optimizer, epoch_num):
    lr = lr * (args.lr_decay ** (epoch_num // args.interval))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def shuffle(feat, tags, labels):
    shuffle_index = torch.randperm(labels.shape[0])
    feat = feat[shuffle_index]
    tags = tags[shuffle_index]
    labels = labels[shuffle_index]
    # this_len = this_len[shuffle_index]
    return feat, tags, labels

def train(args):
    torch.set_default_tensor_type(torch.FloatTensor)

    # initialize model
    if args.model == 'resnet':
        node_dict = {"LFCC": 3}
        feat_model = ResNet(node_dict[args.feat], args.enc_dim, resnet_type='18', nclasses=1 if args.base_loss == "bce" else 2).to(args.device)
    elif args.model == 'lcnn':
        feat_model = LCNN(60, args.enc_dim, nclasses=2).to(args.device)
    elif args.model == 'ecapa':
        node_dict = {"LFCC": 60}
        feat_model = Res2Net2(Bottle2neck, C=512, model_scale=8, nOut=2, n_mels=node_dict[args.feat]).to(args.device)
    elif args.model == 'res2net':
        feat_model = Res2Net(SEBottle2neck, [3, 4, 6, 3], baseWidth=26, scale=4, pretrained=False, num_classes=2).to(args.device)

    if args.continue_training:
        feat_model = torch.load(os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt')).to(args.device)
    # feat_model = nn.DataParallel(feat_model, list(range(torch.cuda.device_count())))  # for multiple GPUs
    feat_optimizer = torch.optim.Adam(feat_model.parameters(), lr=args.lr,
                                      betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)

    training_set = ASVspoof2019(args.access_type, args.path_to_features, 'train',
                                args.feat, feat_len=args.feat_len, pad_chop=args.pad_chop, padding=args.padding)
    validation_set = ASVspoof2019(args.access_type, args.path_to_features, 'dev',
                                  args.feat, feat_len=args.feat_len, pad_chop=args.pad_chop, padding=args.padding)
    if args.LA_aug:
        training_set = ASVspoof2021LA_aug(part="train",
                                        feature=args.feat, feat_len=args.feat_len,
                                        pad_chop=args.pad_chop, padding=args.padding)
        validation_set = ASVspoof2021LA_aug(part="dev",
                                        feature=args.feat, feat_len=args.feat_len,
                                        pad_chop=args.pad_chop, padding=args.padding)
    if args.DF_aug:
        training_set = ASVspoof2021DF_aug(part="train",
                                          feature=args.feat, feat_len=args.feat_len,
                                          pad_chop=args.pad_chop, padding=args.padding)
        validation_set = ASVspoof2021DF_aug(part="dev",
                                            feature=args.feat, feat_len=args.feat_len,
                                            pad_chop=args.pad_chop, padding=args.padding)
    if args.LAPA_aug:
        training_set = ASVspoof2021LAPA_aug(part="train",
                                          feature=args.feat, feat_len=args.feat_len,
                                          pad_chop=args.pad_chop, padding=args.padding)
        validation_set = ASVspoof2021LAPA_aug(part="dev",
                                            feature=args.feat, feat_len=args.feat_len,
                                            pad_chop=args.pad_chop, padding=args.padding)
    if args.DFPA_aug:
        training_set = ASVspoof2021DFPA_aug(part="train",
                                          feature=args.feat, feat_len=args.feat_len,
                                          pad_chop=args.pad_chop, padding=args.padding)
        validation_set = ASVspoof2021DFPA_aug(part="dev",
                                            feature=args.feat, feat_len=args.feat_len,
                                            pad_chop=args.pad_chop, padding=args.padding)

    if args.ADV_AUG:
        assert (args.LA_aug or args.DF_aug or args.LAPA_aug or args.DFPA_aug)
        if args.LA_aug or args.DF_aug:
            classifier = ChannelClassifier(args.enc_dim, len(training_set.channel), args.lambda_).to(args.device)
            classifier_optimizer = torch.optim.Adam(classifier.parameters(), lr=args.lr_d,
                                                    betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)
        else:
            classifier1 = ChannelClassifier(args.enc_dim, len(training_set.channel), args.lambda_).to(args.device)
            classifier1_optimizer = torch.optim.Adam(classifier1.parameters(), lr=args.lr_d,
                                                    betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)
            classifier2 = ChannelClassifier(args.enc_dim, len(training_set.devices), args.lambda_).to(args.device)
            classifier2_optimizer = torch.optim.Adam(classifier2.parameters(), lr=args.lr_d,
                                                     betas=(args.beta_1, args.beta_2), eps=args.eps,
                                                     weight_decay=0.0005)

    trainOriDataLoader = DataLoader(training_set, batch_size=int(args.batch_size * args.ratio),
                                    shuffle=False, num_workers=args.num_workers,
                                    sampler=torch_sampler.SubsetRandomSampler(range(25380)))
    trainAugDataLoader = DataLoader(training_set, batch_size=args.batch_size - int(args.batch_size * args.ratio),
                                    shuffle=False, num_workers=args.num_workers,
                                    sampler=torch_sampler.SubsetRandomSampler(range(25380, len(training_set))))
    trainOri_flow = iter(trainOriDataLoader)
    trainAug_flow = iter(trainAugDataLoader)

    valOriDataLoader = DataLoader(validation_set, batch_size=int(args.batch_size * args.ratio),
                                    shuffle=False, num_workers=args.num_workers,
                                    sampler=torch_sampler.SubsetRandomSampler(range(24844)))
    valAugDataLoader = DataLoader(validation_set, batch_size=args.batch_size - int(args.batch_size * args.ratio),
                                    shuffle=False, num_workers=args.num_workers,
                                    sampler=torch_sampler.SubsetRandomSampler(range(24844, len(validation_set))))
    valOri_flow = iter(valOriDataLoader)
    valAug_flow = iter(valAugDataLoader)

    test_set = ASVspoof2019(args.access_type, args.path_to_features, "eval", args.feat, feat_len=args.feat_len, pad_chop=args.pad_chop, padding=args.padding)
    testDataLoader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=test_set.collate_fn)

    feat, _, _, _, _ = training_set[23]
    print("Feature shape", feat.shape)

    if args.base_loss == "ce":
        criterion = nn.CrossEntropyLoss()
    else:
        assert False

    if args.add_loss == "isolate":
        iso_loss = IsolateLoss(2, args.enc_dim, r_real=args.r_real, r_fake=args.r_fake).to(args.device)
        if args.continue_training:
            iso_loss = torch.load(os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt')).to(args.device)
        iso_loss.train()
        iso_optimzer = torch.optim.SGD(iso_loss.parameters(), lr=args.lr)

    if args.add_loss == "iso_sq":
        iso_loss = IsolateSquareLoss(2, args.enc_dim, r_real=args.r_real, r_fake=args.r_fake).to(args.device)
        if args.continue_training:
            iso_loss = torch.load(os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt')).to(args.device)
        iso_loss.train()
        iso_optimzer = torch.optim.SGD(iso_loss.parameters(), lr=args.lr)

    if args.add_loss == "ang_iso":
        ang_iso = AngularIsoLoss(args.enc_dim, r_real=args.r_real, r_fake=args.r_fake, alpha=args.alpha).to(args.device)
        ang_iso.train()
        ang_iso_optimzer = torch.optim.SGD(ang_iso.parameters(), lr=args.lr)

    if args.add_loss == "p2sgrad":
        p2sgrad_loss = P2SGradLoss(in_dim=args.enc_dim, out_dim=2, smooth=0.0).to(args.device)
        p2sgrad_loss.train()
        p2sgrad_optimzer = torch.optim.SGD(p2sgrad_loss.parameters(), lr=args.lr)

    early_stop_cnt = 0
    prev_loss = 1e8
    add_size = args.batch_size - int(args.batch_size * args.ratio)

    if args.add_loss is None:
        monitor_loss = 'base_loss'
    else:
        monitor_loss = args.add_loss

    for epoch_num in tqdm(range(args.num_epochs)):
        genuine_feats, ip1_loader, tag_loader, idx_loader = [], [], [], []
        feat_model.train()
        trainlossDict = defaultdict(list)
        devlossDict = defaultdict(list)
        testlossDict = defaultdict(list)
        adjust_learning_rate(args, args.lr, feat_optimizer, epoch_num)
        if args.add_loss == "isolate":
            adjust_learning_rate(args, args.lr, iso_optimzer, epoch_num)
        if args.add_loss == "ang_iso":
            adjust_learning_rate(args, args.lr, ang_iso_optimzer, epoch_num)
        if args.add_loss == "p2sgrad":
            adjust_learning_rate(args, args.lr, p2sgrad_optimzer, epoch_num)
        if args.ADV_AUG:
            if args.LA_aug or args.DF_aug:
                adjust_learning_rate(args, args.lr_d, classifier_optimizer, epoch_num)
            else:
                adjust_learning_rate(args, args.lr_d, classifier1_optimizer, epoch_num)
                adjust_learning_rate(args, args.lr_d, classifier2_optimizer, epoch_num)
        print('\nEpoch: %d ' % (epoch_num + 1))
        correct_m, total_m, correct_c, total_c, correct_v, total_v = 0, 0, 0, 0, 0, 0

        for i in trange(0, len(trainOriDataLoader), total=len(trainOriDataLoader), initial=0):
            try:
                featOri, audio_fnOri, tagsOri, labelsOri, channelsOri = next(trainOri_flow)
            except StopIteration:
                trainOri_flow = iter(trainOriDataLoader)
                featOri, audio_fnOri, tagsOri, labelsOri, channelsOri = next(trainOri_flow)

            try:
                featAug, audio_fnAug, tagsAug, labelsAug, channelsAug = next(trainAug_flow)
            except StopIteration:
                trainAug_flow = iter(trainAugDataLoader)
                featAug, audio_fnAug, tagsAug, labelsAug, channelsAug = next(trainAug_flow)

            feat = torch.cat((featOri, featAug), 0)
            tags = torch.cat((tagsOri, tagsAug), 0)
            labels = torch.cat((labelsOri, labelsAug), 0)
            # if not args.LAPA_aug:
            channels = torch.cat((channelsOri, channelsAug), 0)
            # else:
            #     channels = torch.cat((np.array(channelsOri), np.array(channelsAug)), 0)

            # count = 0
            # for channel in list(channels):
            #     if channel == "no_channel":
            #         count += 1
            # print(count / 64)

            # if i > 2: break
            feat = feat.transpose(2,3).to(args.device)

            tags = tags.to(args.device)
            labels = labels.to(args.device)
            # this_len = this_len.to(args.device)

            if args.ratio < 1:
                feat, tags, labels = shuffle(feat, tags, labels)

            if args.model == 'ecapa':
                feat = torch.squeeze(feat)

            feats, feat_outputs = feat_model(feat)

            if args.base_loss == "bce":
                feat_loss = criterion(feat_outputs, labels.unsqueeze(1).float())
            else:
                feat_loss = criterion(feat_outputs, labels)

            trainlossDict['base_loss'].append(feat_loss.item())

            if args.add_loss == None:
                feat_optimizer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()

            if args.add_loss in ["isolate", "iso_sq"]:
                isoloss = iso_loss(feats, labels)
                feat_loss = isoloss * args.weight_loss
                feat_optimizer.zero_grad()
                iso_optimzer.zero_grad()
                trainlossDict[args.add_loss].append(isoloss.item())
                feat_loss.backward()
                feat_optimizer.step()
                iso_optimzer.step()

            if args.add_loss == "ang_iso":
                ang_isoloss, _ = ang_iso(feats, labels)
                feat_loss = ang_isoloss * args.weight_loss
                if epoch_num > 0 and args.ADV_AUG:
                    if args.LA_aug or args.DF_aug:
                        channels = channels.to(args.device)
                        # feats = grl(feats)
                        classifier_out = classifier(feats)
                        _, predicted = torch.max(classifier_out.data, 1)
                        total_m += channels.size(0)
                        correct_m += (predicted == channels).sum().item()
                        device_loss = criterion(classifier_out, channels)
                        # print(feat_loss.item())
                        feat_loss += device_loss
                        # print(device_loss.item())
                        trainlossDict["adv_loss"].append(device_loss.item())
                    else:
                        channels = channels.to(args.device)
                        codec = channels[:, 0]
                        devic = channels[:, 1]
                        classifier1_out = classifier1(feats)
                        classifier2_out = classifier2(feats)
                        _, predicted = torch.max(classifier1_out.data, 1)
                        total_m += channels.size(0)
                        correct_m += (predicted == codec).sum().item()
                        codec_loss = criterion(classifier1_out, codec)
                        devic_loss = criterion(classifier2_out, devic)
                        advaug_loss = codec_loss + devic_loss
                        feat_loss += advaug_loss
                        trainlossDict["adv_loss"].append(advaug_loss.item())
                feat_optimizer.zero_grad()
                ang_iso_optimzer.zero_grad()
                trainlossDict[args.add_loss].append(ang_isoloss.item())
                feat_loss.backward()
                feat_optimizer.step()
                ang_iso_optimzer.step()

            if args.add_loss == "p2sgrad":
                feat_loss, _ = p2sgrad_loss(feats, labels)
                trainlossDict[args.add_loss].append(feat_loss.item())
                feat_optimizer.zero_grad()
                p2sgrad_optimzer.zero_grad()
                feat_loss.backward()
                feat_optimizer.step()
                p2sgrad_optimzer.step()

            if args.ADV_AUG:
                if args.LA_aug or args.DF_aug:
                    channels = channels.to(args.device)
                    feats, _ = feat_model(feat)
                    feats = feats.detach()
                    # feats = grl(feats)
                    classifier_out = classifier(feats)
                    _, predicted = torch.max(classifier_out.data, 1)
                    total_c += channels.size(0)
                    correct_c += (predicted == channels).sum().item()
                    device_loss_c = criterion(classifier_out, channels)
                    classifier_optimizer.zero_grad()
                    device_loss_c.backward()
                    classifier_optimizer.step()
                else:
                    channels = channels.to(args.device)
                    codec = channels[:, 0]
                    devic = channels[:, 1]
                    feats, _ = feat_model(feat)
                    feats = feats.detach()
                    # feats = grl(feats)
                    classifier1_out = classifier1(feats)
                    classifier2_out = classifier2(feats)
                    _, predicted = torch.max(classifier1_out.data, 1)
                    total_c += channels.size(0)
                    correct_c += (predicted == codec).sum().item()
                    codec_loss_c = criterion(classifier1_out, codec)
                    classifier1_optimizer.zero_grad()
                    codec_loss_c.backward()
                    classifier1_optimizer.step()
                    devic_loss_c = criterion(classifier2_out, devic)
                    classifier2_optimizer.zero_grad()
                    devic_loss_c.backward()
                    classifier2_optimizer.step()

            # genuine_feats.append(feats[labels==0])
            ip1_loader.append(feats)
            idx_loader.append((labels))
            tag_loader.append((tags))

            # if epoch_num > 0:
            #     print(100 * correct_m / total_m)
            #     print(100 * correct_c / total_c)

            # desc_str = ''
            # for key in sorted(trainlossDict.keys()):
            #     desc_str += key + ':%.5f' % (np.nanmean(trainlossDict[key])) + ', '
            # t.set_description(desc_str)
            # print(desc_str)

            if epoch_num > 0 and args.ADV_AUG:
                with open(os.path.join(args.out_fold, "train_loss.log"), "a") as log:
                    log.write(str(epoch_num) + "\t" + str(i) + "\t" +
                              str(trainlossDict["adv_loss"][-1]) + "\t" +
                              str(100 * correct_m / total_m) + "\t" +
                              str(100 * correct_c / total_c) + "\t" +
                              str(trainlossDict[monitor_loss][-1]) + "\n")
            else:
                with open(os.path.join(args.out_fold, "train_loss.log"), "a") as log:
                    log.write(str(epoch_num) + "\t" + str(i) + "\t" +
                              str(trainlossDict[monitor_loss][-1]) + "\n")

        # print(len(what))
        # print(len(list(set(what))))
        # assert len(what) == len(list(set(what)))
        # Val the model
        # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
        feat_model.eval()
        with torch.no_grad():
            ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
            # with trange(2) as v:
            # with trange(len(valDataLoader)) as v:
            #     for i in v:
            for i in trange(0, len(valOriDataLoader), total=len(valOriDataLoader), initial=0):
                try:
                    featOri, audio_fnOri, tagsOri, labelsOri, channelsOri = next(valOri_flow)
                except StopIteration:
                    valOri_flow = iter(valOriDataLoader)
                    featOri, audio_fnOri, tagsOri, labelsOri, channelsOri = next(valOri_flow)

                try:
                    featAug, audio_fnAug, tagsAug, labelsAug, channelsAug = next(valAug_flow)
                except StopIteration:
                    valAug_flow = iter(valAugDataLoader)
                    featAug, audio_fnAug, tagsAug, labelsAug, channelsAug = next(valAug_flow)

                feat = torch.cat((featOri, featAug), 0)
                tags = torch.cat((tagsOri, tagsAug), 0)
                labels = torch.cat((labelsOri, labelsAug), 0)
                channels = torch.cat((channelsOri, channelsAug), 0)

                # if i > 2: break
                feat = feat.transpose(2,3).to(args.device)

                tags = tags.to(args.device)
                labels = labels.to(args.device)

                feat, tags, labels = shuffle(feat, tags, labels)

                if args.model == 'ecapa':
                    feat = torch.squeeze(feat)

                feats, feat_outputs = feat_model(feat)

                if args.base_loss == "bce":
                    feat_loss = criterion(feat_outputs, labels.unsqueeze(1).float())
                    score = feat_outputs[:, 0]
                else:
                    feat_loss = criterion(feat_outputs, labels)
                    score = F.softmax(feat_outputs, dim=1)[:, 0]

                ip1_loader.append(feats)
                idx_loader.append((labels))
                tag_loader.append((tags))

                if args.add_loss in [None]:
                    devlossDict["base_loss"].append(feat_loss.item())
                elif args.add_loss in ["isolate", "iso_sq"]:
                    isoloss = iso_loss(feats, labels)
                    score = torch.norm(feats - iso_loss.center, p=2, dim=1)
                    devlossDict[args.add_loss].append(isoloss.item())
                elif args.add_loss == "ang_iso":
                    ang_isoloss, score = ang_iso(feats, labels)
                    devlossDict[args.add_loss].append(ang_isoloss.item())
                    if epoch_num > 0 and args.ADV_AUG:
                        if args.LA_aug or args.DF_aug:
                            channels = channels.to(args.device)
                            # feats = grl(feats)
                            classifier_out = classifier(feats)
                            _, predicted = torch.max(classifier_out.data, 1)
                            total_v += channels.size(0)
                            correct_v += (predicted == channels).sum().item()
                            device_loss = criterion(classifier_out, channels)
                            # print(feat_loss.item())
                            feat_loss += device_loss
                            # print(device_loss.item())
                            devlossDict["adv_loss"].append(device_loss.item())
                        else:
                            channels = channels.to(args.device)
                            codec = channels[:, 0]
                            devic = channels[:, 1]
                            classifier1_out = classifier1(feats)
                            classifier2_out = classifier2(feats)
                            _, predicted = torch.max(classifier1_out.data, 1)
                            total_v += channels.size(0)
                            correct_v += (predicted == codec).sum().item()
                            codec_loss = criterion(classifier1_out, codec)
                            devic_loss = criterion(classifier2_out, devic)
                            advaug_loss = codec_loss + devic_loss
                            feat_loss += advaug_loss
                            devlossDict["adv_loss"].append(advaug_loss.item())
                elif args.add_loss == 'p2sgrad':
                    feat_loss, score = p2sgrad_loss(feats, labels)
                    devlossDict[args.add_loss].append(feat_loss.item())

                score_loader.append(score)

                # desc_str = ''
                # for key in sorted(devlossDict.keys()):
                #     desc_str += key + ':%.5f' % (np.nanmean(devlossDict[key])) + ', '
                # # v.set_description(desc_str)
                # print(desc_str)
#             scores = torch.cat(score_loader, 0).data.cpu().numpy()
#             labels = torch.cat(idx_loader, 0).data.cpu().numpy()
#             eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]
#             other_eer = em.compute_eer(-scores[labels == 0], -scores[labels == 1])[0]
#             eer = min(eer, other_eer)

#             if epoch_num > 0 and args.ADV_AUG:
#                 with open(os.path.join(args.out_fold, "dev_loss.log"), "a") as log:
#                     log.write(str(epoch_num) + "\t"+ "\t" +
#                               str(np.nanmean(devlossDict["adv_loss"])) + "\t" +
#                               str(100 * correct_v / total_v) + "\t" +
#                               str(np.nanmean(devlossDict[monitor_loss])) + "\t" +
#                               str(eer) + "\n")
#             else:
#                 with open(os.path.join(args.out_fold, "dev_loss.log"), "a") as log:
#                     log.write(str(epoch_num) + "\t" +
#                               str(np.nanmean(devlossDict[monitor_loss])) + "\t" +
#                               str(eer) +"\n")
#             print("Val EER: {}".format(eer))

            if args.visualize and ((epoch_num+1) % 3 == 1):
                feat = torch.cat(ip1_loader, 0)
                tags = torch.cat(tag_loader, 0)
                if args.add_loss == "isolate":
                    centers = iso_loss.center
                elif args.add_loss == "ang_iso":
                    centers = ang_iso.center
                else:
                    centers = torch.mean(feat[labels==0], dim=0, keepdim=True)
                visualize(args, feat.data.cpu().numpy(), tags.data.cpu().numpy(), labels.data.cpu().numpy(), centers.data.cpu().numpy(),
                          epoch_num + 1, "Dev")

        if args.test_on_eval:
            with torch.no_grad():
                ip1_loader, tag_loader, idx_loader, score_loader = [], [], [], []
                for i, (feat, audio_fn, tags, labels) in enumerate(tqdm(testDataLoader)):
                    # if i > 2: break
                    feat = feat.transpose(2,3).to(args.device)
                    tags = tags.to(args.device)
                    labels = labels.to(args.device)

                    if args.model == 'ecapa':
                        feat = torch.squeeze(feat)
                    feats, feat_outputs = feat_model(feat)

                    if args.base_loss == "bce":
                        feat_loss = criterion(feat_outputs, labels.unsqueeze(1).float())
                        score = feat_outputs[:, 0]
                    else:
                        feat_loss = criterion(feat_outputs, labels)
                        score = F.softmax(feat_outputs, dim=1)[:, 0]

                    ip1_loader.append(feats)
                    idx_loader.append((labels))
                    tag_loader.append((tags))

                    if args.add_loss in [None]:
                        testlossDict["base_loss"].append(feat_loss.item())
                    elif args.add_loss in ["isolate", "iso_sq"]:
                        isoloss = iso_loss(feats, labels)
                        score = torch.norm(feats - iso_loss.center, p=2, dim=1)
                        testlossDict[args.add_loss].append(isoloss.item())
                    elif args.add_loss == "ang_iso":
                        ang_isoloss, score = ang_iso(feats, labels)
                        testlossDict[args.add_loss].append(ang_isoloss.item())
                    elif args.add_loss == 'p2sgrad':
                        p2s_loss, score = p2sgrad_loss(feats, labels)
                        testlossDict[args.add_loss].append(p2s_loss.item())

                    score_loader.append(score)

                    # desc_str = ''
                    # for key in sorted(testlossDict.keys()):
                    #     desc_str += key + ':%.5f' % (np.nanmean(testlossDict[key])) + ', '
                    # # v.set_description(desc_str)
                    # print(desc_str)
                scores = torch.cat(score_loader, 0).data.cpu().numpy()
                labels = torch.cat(idx_loader, 0).data.cpu().numpy()
                eer = em.compute_eer(scores[labels == 0], scores[labels == 1])[0]
                other_eer = em.compute_eer(-scores[labels == 0], -scores[labels == 1])[0]
                eer = min(eer, other_eer)

                with open(os.path.join(args.out_fold, "test_loss.log"), "a") as log:
                    log.write(str(epoch_num) + "\t" + str(np.nanmean(testlossDict[monitor_loss])) + "\t" + str(eer) + "\n")
                print("Test EER: {}".format(eer))

        valLoss = np.nanmean(devlossDict[monitor_loss])
        # if args.add_loss == "isolate":
        #     print("isolate center: ", iso_loss.center.data)
        if (epoch_num + 1) % 1 == 0:
            torch.save(feat_model, os.path.join(args.out_fold, 'checkpoint',
                                                'anti-spoofing_feat_model_%d.pt' % (epoch_num + 1)))
            if args.add_loss in ["isolate", "iso_sq"]:
                loss_model = iso_loss
                torch.save(loss_model, os.path.join(args.out_fold, 'checkpoint',
                                                    'anti-spoofing_loss_model_%d.pt' % (epoch_num + 1)))
            elif args.add_loss == "ang_iso":
                loss_model = ang_iso
                torch.save(loss_model, os.path.join(args.out_fold, 'checkpoint',
                                                    'anti-spoofing_loss_model_%d.pt' % (epoch_num + 1)))
            elif args.add_loss == "p2sgrad":
                loss_model = p2sgrad_loss
                torch.save(p2sgrad_loss, os.path.join(args.out_fold, 'checkpoint',
                                                    'anti-spoofing_loss_model_%d.pt' % (epoch_num + 1)))
            else:
                loss_model = None

        if valLoss < prev_loss:
            # Save the model checkpoint
            torch.save(feat_model, os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt'))
            if args.add_loss in ["isolate", "iso_sq"]:
                loss_model = iso_loss
                torch.save(loss_model, os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))
            elif args.add_loss == "ang_iso":
                loss_model = ang_iso
                torch.save(loss_model, os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))
            elif args.add_loss == "p2sgrad":
                loss_model = p2sgrad_loss
                torch.save(p2sgrad_loss, os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))
            else:
                loss_model = None
            prev_loss = valLoss
            early_stop_cnt = 0
        else:
            early_stop_cnt += 1

        if early_stop_cnt == 500:
            with open(os.path.join(args.out_fold, 'args.json'), 'a') as res_file:
                res_file.write('\nTrained Epochs: %d\n' % (epoch_num - 499))
            break
        # if early_stop_cnt == 1:
        #     torch.save(feat_model, os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt')

            # print('Dev Accuracy of the model on the val features: {} % '.format(100 * feat_correct / total))

    return feat_model, loss_model

if __name__ == "__main__":
    args = initParams()
    if not args.test_only:
        _, _ = train(args)
    # model = torch.load(os.path.join(args.out_fold, 'anti-spoofing_feat_model.pt'))
    # if args.add_loss is None:
    #     loss_model = None
    # else:
    #     loss_model = torch.load(os.path.join(args.out_fold, 'anti-spoofing_loss_model.pt'))
    # # TReer_cm, TRmin_tDCF = test(args, model, loss_model, "train")
    # # VAeer_cm, VAmin_tDCF = test(args, model, loss_model, "dev")
    # TEeer_cm, TEmin_tDCF = test(args, model, loss_model)
    # with open(os.path.join(args.out_fold, 'args.json'), 'a') as res_file:
    #     # res_file.write('\nTrain EER: %8.5f min-tDCF: %8.5f\n' % (TReer_cm, TRmin_tDCF))
    #     # res_file.write('\nVal EER: %8.5f min-tDCF: %8.5f\n' % (VAeer_cm, VAmin_tDCF))
    #     res_file.write('\nTest EER: %8.5f min-tDCF: %8.5f\n' % (TEeer_cm, TEmin_tDCF))
    # plot_loss(args)

    # # Test a checkpoint model
    # args = initParams()
    # model = torch.load(os.path.join(args.out_fold, 'checkpoint', 'anti-spoofing_feat_model_19.pt'))
    # loss_model = torch.load(os.path.join(args.out_fold, 'checkpoint', 'anti-spoofing_loss_model_19.pt'))
    # VAeer_cm, VAmin_tDCF = test(args, model, loss_model, "dev")


import os
import numpy as np
from multiprocessing import cpu_count
import logging

from joblib import Parallel, delayed
import pysam
import click
import click_log

from .utils import get_fast5_read_id_mapping, get_output_filenames
from .model import load_adapter_model
from .signal import get_fast5_fiveprime
from .trim import Trimmer

logger = logging.getLogger(__name__)
click_log.basic_config(logger)

def predict_bam_batch(batch, reads,
                      model, threshold,
                      batch_size, squiggle_size):
    '''
    Make a prediction for a batch of signals
    '''
    read_ids, sig_len, batch = zip(*batch)
    sig_len = np.array(sig_len)
    batch = np.array(batch).reshape(batch_size, squiggle_size, 1)
    batch = batch[~np.isnan(sig_len)]
    read_ids = [r_id for r_id, sl in zip(read_ids, sig_len) if not np.isnan(sl)]
    preds = model.predict(batch).squeeze()
    logger.info('INFO: Made prediction for batch')
    for score, read_id, read in zip(preds, read_ids, reads):
        assert read_id == read.query_name
        read.set_tag('ac', score, value_type='f')
        yield score > threshold, read

def predict_fq_batch(batch, reads,
                     model, trimmer, threshold,
                     batch_size, squiggle_size):
    '''
    Make a prediction for a batch of signals
    '''
    read_ids, _, batch = zip(*batch)
    batch = np.array(batch).reshape(batch_size, squiggle_size, 1)
    preds = model.predict(batch).squeeze()
    reads = trimmer(reads, preds > threshold)
    logger.info('INFO: Made prediction for batch')
    for score, read_id, read in zip(preds, read_ids, reads):
        assert read_id == read.name
        yield score > threshold, read

def bam_batch_iter(bam, read_id_fast5_filemap,
                   batch_size, keep_unmapped):
    batch = []
    batch_reads = []
    for read in bam.fetch(until_eof=keep_unmapped):
        read_id = read.query_name
        try:
            f5_fn = read_id_fast5_filemap[read_id]
        except KeyError:
            logger.warning(
                'WARNING: Fast5 file not found for read {} in bam file'.format(read_id))
            continue
        batch.append((read_id, f5_fn))
        batch_reads.append(read)
        if len(batch) == batch_size:
            logger.info('INFO: Identified batch of read fast5 files')
            yield batch, batch_reads, batch_size
            batch = []
            batch_reads = []
    if len(batch):
        yield batch, batch_reads, len(batch)

def fq_batch_iter(fq, read_id_fast5_filemap, batch_size):
    batch = []
    batch_reads = []
    for read in fq:
        read_id = read.name
        try:
            f5_fn = read_id_fast5_filemap[read_id]
        except KeyError:
            logger.warning(
                'WARNING: Fast5 file not found for read {} in bam file'.format(read_id))
            continue
        batch.append((read_id, f5_fn))
        batch_reads.append(read)
        if len(batch) == batch_size:
            logger.info('INFO: Identified batch of read fast5 files')
            yield batch, batch_reads, batch_size
            batch = []
            batch_reads = []
    if len(batch):
        yield batch, batch_reads, len(batch)

def bam_filter(bam_fn, read_id_fast5_filemap, model,
               pass_output_bam_fn, fail_output_bam_fn, threshold=0.5,
               batch_size=5000, squiggle_size=2000,
               keep_unmapped=False, processes=8):
    '''
    Filter a bam file using a model trained to detect adapters in
    signal from 5' end.

    Parameters:
    ----------

        bam_fn: str, required
          Path to input bam file

        read_id_fast5_filemap: dict, required
          Dict of read_id: fast5 filepath key value pairs

        model: keras.models.Model, required
          Model used to make predictions. Must accept input of
          shape (batch_size, squiggle_size, 1)

        pass_output_bam_fn: str, required
          Bam file path to write reads passing filter to.
          Will be overwritten

        fail_output_bam_fn: str, required
          Bam file path to write reads failing filter to.
          Will be overwritten

        threshold: float, optional, default: 0.5
          Threshold at which to cutoff pass/fail

        batch_size, int, optional, default: 5000
          Number of signals to pass to model prediction
          at a time

        squiggle_size: int, optional, default: 2000
          Length of signal to take from 5' end of RNA fast5 signal,
          to pass to model.

        keep_unmapped: bool, optional, default: False
          Whether to process unmapped reads in the bam file. If false,
          unmapped reads are filtered out (i.e. are not present in
          either pass or fail bam files). This speeds up computation.

        processed: int, optional, default: 8
          Number of parallel processes to use to read signals from fast5

    Returns:
    --------

        None
    '''
    in_bam = pysam.AlignmentFile(bam_fn, mode='rb')
    # create two output bams
    out_bams = {
        True: pysam.AlignmentFile(pass_output_bam_fn,
                                  mode='wb', template=in_bam),
        False: pysam.AlignmentFile(fail_output_bam_fn,
                                   mode='wb', template=in_bam)
    }
    filter_counts = {True: 0, False: 0}
    with Parallel(n_jobs=processes) as pool:
        for batch, reads, bs in bam_batch_iter(in_bam, read_id_fast5_filemap,
                                               batch_size, keep_unmapped):
            # not sure if joblib guarantees the output order is the same as
            # input order, so pass the read_ids as well to be safe
            batch_signals = pool(
                delayed(get_fast5_fiveprime)(*f, squiggle_size) for f in batch)
            logger.info('INFO: Collected Five prime signals for batch')
            for decision, read in predict_bam_batch(batch_signals, reads, model,
                                                    threshold, bs, squiggle_size):
                out_bams[decision].write(read)
                filter_counts[decision] += 1
            logger.info('INFO: Wrote batch to bam files')
            logger.info('INFO: Processed {} reads'.format(sum(filter_counts.values())))
            logger.info('INFO: Of these:')
            logger.info('INFO:   {} passed the threshold'.format(filter_counts[True]))
            logger.info('INFO:   {} failed the threshold'.format(filter_counts[False]))

    in_bam.close()
    for bam in out_bams.values():
        bam.close()

def fastq_filter(fastq_fn, read_id_fast5_filemap, model,
                 pass_output_fq_fn, fail_output_fq_fn, threshold=0.5,
                 batch_size=5000, squiggle_size=2000,
                 trim=False, processes=8):
    '''
    Filter a bam file using a model trained to detect adapters in
    signal from 5' end.

    Parameters:
    ----------

        fastq_fn: str, required
          Path to input fastq file

        read_id_fast5_filemap: dict, required
          Dict of read_id: fast5 filepath key value pairs

        model: keras.models.Model, required
          Model used to make predictions. Must accept input of
          shape (batch_size, squiggle_size, 1)

        pass_output_fq_fn: str, required
          Fastq file path to write reads passing filter to.
          Will be overwritten

        fail_output_fq_fn: str, required
          Fastq file path to write reads failing filter to.
          Will be overwritten

        threshold: float, optional, default: 0.5
          Threshold at which to cutoff pass/fail

        batch_size, int, optional, default: 5000
          Number of signals to pass to model prediction
          at a time

        squiggle_size: int, optional, default: 2000
          Length of signal to take from 5' end of RNA fast5 signal,
          to pass to model.

        trim: bool, optional, default: False
          If True, a second model trained at the sequence level will
          be used to estimate the length of the adapter and cut it off

        processed: int, optional, default: 8
          Number of parallel processes to use to read signals from fast5

    Returns:
    --------

        None
    '''
    in_fq = pysam.FastxFile(fastq_fn)
    # create two output fastqs
    out_fqs = {
        True: open(pass_output_fq_fn, mode='w'),
        False: open(fail_output_fq_fn, mode='w')
    }
    if trim:
        trimmer = Trimmer().trim_adapters
    else:
        trimmer = lambda x, y: x
    filter_counts = {True: 0, False: 0}
    with Parallel(n_jobs=processes) as pool:
        for batch, reads, bs in fq_batch_iter(in_fq, read_id_fast5_filemap,
                                              batch_size):
            # not sure if joblib guarantees the output order is the same as
            # input order, so pass the read_ids as well to be safe
            batch_signals = pool(
                delayed(get_fast5_fiveprime)(*f, squiggle_size) for f in batch)
            logger.info('INFO: Collected Five prime signals for batch')
            for decision, read in predict_fq_batch(batch_signals, reads, model, trimmer,
                                                   threshold, bs, squiggle_size):
                out_fqs[decision].write('{:s}\n'.format(str(read)))
                filter_counts[decision] += 1
            logger.info('INFO: Wrote batch to Fastq files')
            logger.info('INFO: Processed {} reads'.format(sum(filter_counts.values())))
            logger.info('INFO: Of these:')
            logger.info('INFO:   {} passed the threshold'.format(filter_counts[True]))
            logger.info('INFO:   {} failed the threshold'.format(filter_counts[False]))

    in_fq.close()
    for fq in out_fqs.values():
        fq.close()

@click.group()
def cli():
    pass

@cli.command()
@click_log.simple_verbosity_option(logger)
@click.option('-p', '--pos-bam-file', required=True, type=str)
@click.option('-n', '--neg-bam-file', required=True, type=str)
@click.option('-o', '--output-bam-basename', required=True, type=str)
@click.option('-t', '--threshold', required=True, type=float)
@click.option('--keep-unmapped/--no-keep-unmapped', default=False)
def rethreshold_bam(pos_bam_file, neg_bam_file,
                    output_bam_basename,
                    threshold, keep_unmapped):
    logger.info('INFO: Rethresholding bams at {}'.format(threshold))
    input_bams = [
        pysam.AlignmentFile(bam_fn, mode='rb')
        for bam_fn in (pos_bam_file, neg_bam_file)
    ]
    pass_output_bam_fn, fail_output_bam_fn = get_output_filenames(
        output_bam_basename, 'bam'
    )
    out_bams = {
        True: pysam.AlignmentFile(pass_output_bam_fn,
                                  mode='wb', template=input_bams[0]),
        False: pysam.AlignmentFile(fail_output_bam_fn,
                                   mode='wb', template=input_bams[1])
    }
    filter_counts = {True: 0, False: 0}
    for in_bam in input_bams:
        for record in in_bam.fetch(until_eof=keep_unmapped):
            decision = record.get_tag('ac') > threshold
            out_bams[decision].write(record)
            filter_counts[decision] += 1
    logger.info('INFO: Processed {} reads'.format(sum(filter_counts.values())))
    logger.info('INFO: Of these:')
    logger.info('INFO:   {} passed the threshold'.format(filter_counts[True]))
    logger.info('INFO:   {} failed the threshold'.format(filter_counts[False]))

    for bam in [*input_bams, *out_bams.values()]:
        bam.close()

@cli.command()
@click_log.simple_verbosity_option(logger)
@click.option('-b', '--bam-file', required=True, type=str)
@click.option('-s', '--summary-files-basedir', required=True, type=str)
@click.option('-f', '--fast5-files-basedir', required=True, type=str)
@click.option('-o', '--output-bam-basename', required=True, type=str)
@click.option('-m', '--model-file', required=False, default=None)
@click.option('-k', '--rna-kit', required=False, type=click.Choice(['SQK-RNA001', 'SQK-RNA002']))
@click.option('-t', '--threshold', required=False, default=0.5, type=float)
@click.option('-n', '--batch-size', required=False, default=5000, type=int)
@click.option('-l', '--squiggle-size', required=False, default=2000, type=int)
@click.option('--keep-unmapped/--no-keep-unmapped', default=False)
@click.option('-p', '--processes', required=False, default=-1, type=int)
def filter_bam(bam_file,
               summary_files_basedir,
               fast5_files_basedir,
               output_bam_basename,
               model_file,
               rna_kit,
               threshold,
               batch_size,
               squiggle_size,
               keep_unmapped,
               processes):
    if processes == -1:
        processes = cpu_count()
    logger.info('INFO: {} cpus available'.format(processes))
    read_id_f5_filemap = get_fast5_read_id_mapping(
        summary_files_basedir,
        fast5_files_basedir,
    )
    logger.info(
        ('INFO: Found {} read_ids with associated Fast5'
        ' files in sequencing summaries').format(len(read_id_f5_filemap)))
    pass_output_bam_fn, fail_output_bam_fn = get_output_filenames(
        output_bam_basename
    )
    logger.info('INFO: Reads passing threshold will be written to {}'.format(
        pass_output_bam_fn))
    logger.info('INFO: Reads failing threshold will be written to {}'.format(
        fail_output_bam_fn))
    logger.info('INFO: Loading model')
    model = load_adapter_model(model_file, rna_kit=rna_kit)
    logger.debug('DEBUG: Model Summary:')
    model.summary(print_fn=logger.debug)
    bam_filter(
        bam_file, read_id_f5_filemap, model,
        pass_output_bam_fn, fail_output_bam_fn,
        threshold, batch_size, squiggle_size,
        keep_unmapped, processes
    )

@cli.command()
@click_log.simple_verbosity_option(logger)
@click.option('-q', '--fastq-file', required=True, type=str)
@click.option('-s', '--summary-files-basedir', required=True, type=str)
@click.option('-f', '--fast5-files-basedir', required=True, type=str)
@click.option('-o', '--output-fastq-basename', required=True, type=str)
@click.option('-m', '--model-file', required=False, default=None)
@click.option('-t', '--threshold', required=False, default=0.5, type=float)
@click.option('-n', '--batch-size', required=False, default=5000, type=int)
@click.option('-l', '--squiggle-size', required=False, default=3000, type=int)
@click.option('-c', '--trim/--no-trim', required=False, default=False)
@click.option('-p', '--processes', required=False, default=-1, type=int)
def filter_fastq(fastq_file,
                 summary_files_basedir,
                 fast5_files_basedir,
                 output_fastq_basename,
                 model_file,
                 threshold,
                 batch_size,
                 squiggle_size,
                 trim,
                 processes):
    if processes == -1:
        processes = cpu_count()
    logger.info('INFO: {} cpus available'.format(processes))
    read_id_f5_filemap = get_fast5_read_id_mapping(
        summary_files_basedir,
        fast5_files_basedir,
    )
    logger.info(
        ('INFO: Found {} read_ids with associated Fast5'
        ' files in sequencing summaries').format(len(read_id_f5_filemap)))
    pass_output_fq_fn, fail_output_fq_fn = get_output_filenames(
        output_fastq_basename, 'fastq'
    )
    logger.info('INFO: Reads passing threshold will be written to {}'.format(
        pass_output_fq_fn))
    logger.info('INFO: Reads failing threshold will be written to {}'.format(
        fail_output_fq_fn))
    logger.info('INFO: Loading model')
    model = load_adapter_model(model_file)
    logger.debug('DEBUG: Model Summary:')
    model.summary(print_fn=logger.debug)
    fastq_filter(
        fastq_file, read_id_f5_filemap, model,
        pass_output_fq_fn, fail_output_fq_fn,
        threshold, batch_size, squiggle_size,
        trim, processes
    )

if __name__ == '__main__':
    run_bam_filter()


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

import argparse
import math
import os
import random
import time

import numpy as np
import paddle
from visualdl import LogWriter
from modeling import GPTModel, GPTForPretraining, GPTPretrainingCriterion, GPTForPretrainingPipe
from paddlenlp.transformers import GPTTokenizer, GPTChineseTokenizer
from paddlenlp.utils.log import logger

from dataset import create_pretrained_dataset
from args import parse_args
import lr
from paddle.distributed import fleet
from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients
import types
from utils import get_timers, set_timers
from types import MethodType
from paddle import _C_ops
from paddle.fluid import core
from paddle.fluid.dygraph import to_variable
import paddle.distributed as dist
from framework import assign_group_by_size, flatten_dense_tensors, obtain_storage, AdamW
from paddle.incubate.distributed.models import moe

from checkpointing import save_checkpoint, load_checkpoint

MODEL_CLASSES = {
    "gpt": (GPTForPretraining, GPTTokenizer),
    "gpt-cn": (GPTForPretraining, GPTChineseTokenizer),
}

set_timers()

def set_hyrbid_parallel_seed(basic_seed, dp_rank, mp_rank, pp_rank):
    assert args.device != "cpu"

    random.seed(basic_seed + dp_rank)
    np.random.seed(basic_seed + dp_rank)
    paddle.seed(basic_seed + dp_rank)
    from paddle.distributed.fleet import meta_parallel
    meta_parallel.model_parallel_random_seed(basic_seed + dp_rank + 1000 *
                                             mp_rank)

    # local_seed/ global_seed is used to control dropout in ModelParallel
    local_seed = basic_seed + 123 + mp_rank * 10 + pp_rank * 1000
    global_seed = basic_seed + dp_rank
    tracker = get_rng_state_tracker()
    tracker.add('global_seed', global_seed)
    tracker.add('local_seed', local_seed)

@paddle.no_grad()
def run_evaluate(args,
                 data_loader,
                 model,
                 criterion,
                 iter_steps,
                 log_writer,
                 global_step,
                 epoch,
                 task_name="valid"):
    model.eval()
    all_loss = []
    local_time = time.time()
    for eval_step, batch in enumerate(data_loader):
        tokens, loss_mask, labels = batch
        with paddle.amp.auto_cast(
                args.use_pure_fp16,
                custom_black_list=[
                    "reduce_sum",
                    "c_softmax_with_cross_entropy",
                    "elementwise_div",
                ],
                level='O2'):
            preds = model(tokens)
        preds = paddle.cast(preds, dtype="float32")
        loss = criterion(preds, labels, loss_mask)

        all_loss.append(float(loss))
        if eval_step >= iter_steps - 1:
            break

    average_loss = sum(all_loss) / len(all_loss)
    logger.info("%s step %d, epoch: %d, batch: %d, loss: %f, speed: %.2f step/s"
                % (task_name, global_step, epoch, eval_step, average_loss,
                   iter_steps / (time.time() - local_time)))
    log_writer.add_scalar(task_name + "_loss", average_loss, global_step)
    model.train()

def initialize_model_and_expert_group(hcg):
    def get_expert_parallel_world_size(self):
        return self.get_data_parallel_world_size(
        ) * self.get_model_parallel_world_size()

    hcg.get_expert_parallel_world_size = types.MethodType(
        get_expert_parallel_world_size, hcg)

    # need create mp_dp group for expert parallel group in advance
    _, mp_dp_comm_group = hcg._set_check_group(parallel_method="pipe")

    def get_expert_parallel_group(self):
        return mp_dp_comm_group

    hcg.get_expert_parallel_group = types.MethodType(get_expert_parallel_group,
                                                     hcg)

def initialize_mp_dp_parameters(model, hcg):
    mp_group = hcg.get_model_parallel_group()
    mp_src_rank = hcg.get_model_parallel_group_src_rank()

    dp_group = hcg.get_data_parallel_group()
    dp_src_rank = hcg.get_data_parallel_group_src_rank()

    for param in model.parameters():
        if "expert_" in param.name:
            continue
        if not param.is_distributed:
            paddle.distributed.broadcast(
                param.detach(),
                src=mp_src_rank,
                group=mp_group,
                use_calc_stream=True)

        paddle.distributed.broadcast(
            param.detach(),
            src=dp_src_rank,
            group=dp_group,
            use_calc_stream=True)

def unscale_method(self, optimizer):
    if not self._enable:
        return

    if getattr(optimizer, '_param_groups', None) and isinstance(
            optimizer._param_groups[0], dict):
        param_grads_fp16 = []
        param_grads_fp32 = []
        for group in optimizer._param_groups:
            for param in group['params']:
                if param._grad_ivar() is not None:
                    if param._grad_ivar().dtype == core.VarDesc.VarType.FP16:
                        param_grads_fp16.append(param._grad_ivar())
                    else:
                        param_grads_fp32.append(param._grad_ivar())
    else:
        param_grads_fp16 = [
            param._grad_ivar() for param in optimizer._parameter_list
            if (param._grad_ivar() is not None
                ) and (param._grad_ivar().dtype == core.VarDesc.VarType.FP16)
        ]
        param_grads_fp32 = [
            param._grad_ivar() for param in optimizer._parameter_list
            if (param._grad_ivar() is not None
                ) and (param._grad_ivar().dtype == core.VarDesc.VarType.FP32)
        ]
    temp_found_inf_fp16 = to_variable(np.array([0]).astype(np.bool))
    temp_found_inf_fp32 = to_variable(np.array([0]).astype(np.bool))

    if len(param_grads_fp16):
        _C_ops.check_finite_and_unscale(param_grads_fp16, self._scale,
                                        param_grads_fp16, temp_found_inf_fp16)
    if len(param_grads_fp32):
        _C_ops.check_finite_and_unscale(param_grads_fp32, self._scale,
                                        param_grads_fp32, temp_found_inf_fp32)
    self._found_inf = 1 if temp_found_inf_fp16 or temp_found_inf_fp32 else 0

    if dist.get_world_size() > 1:
        is_found_inf = paddle.to_tensor([self._found_inf], dtype="int32")
        paddle.distributed.all_reduce(
            is_found_inf, op=paddle.distributed.ReduceOp.MAX, group=None)
        self._found_inf = is_found_inf.numpy()[0]

def all_reduce_parameters(params, group):
    if group.nranks < 2:
        return
    nranks = group.nranks
    with paddle.framework.no_grad():
        for p in params:
            grad = p.grad.scale_(group.nranks)
            paddle.distributed.all_reduce(grad, use_calc_stream=True)

def parameters_classify(model):
    decay_gate_params = []
    decay_expert_params = []
    decay_other_params = []

    gate_params = []
    expert_params = []
    other_params = []

    for param in model.parameters():
        # param_name = param.name
        if "expert_" in param.name:
            if not any(nd in param.name for nd in ["bias", "norm"]):
                decay_expert_params.append(param)
            else:
                expert_params.append(param)
        elif "gate_" in param.name:
            if not any(nd in param.name for nd in ["bias", "norm"]):
                decay_gate_params.append(param)
            else:
                gate_params.append(param)
        else:
            if not any(nd in param.name for nd in ["bias", "norm"]):
                decay_other_params.append(param)
            else:
                other_params.append(param)

    print("all parameters length:", len(model.parameters()))
    print(
        "decay_gate_params len: {}, decay_expert_params len: {}, decay_other_params len: {}".
        format(
            len(decay_gate_params),
            len(decay_expert_params), len(decay_other_params)))
    print("gate_params len: {}, expert_params len: {}, other_params len: {}".
          format(len(gate_params), len(expert_params), len(other_params)))

    d_gate = obtain_storage(decay_gate_params)
    gate = obtain_storage(gate_params)

    d_expert = obtain_storage(decay_expert_params)
    expert = obtain_storage(expert_params)

    d_other = obtain_storage(decay_other_params)
    other = obtain_storage(other_params)

    opt_fused_tensors = []
    decay_fused_tensors = []
    reduce_fused_tensors = []
    gate_fused_tensors = []

    decay_fused_tensors = d_gate + d_other + d_expert
    opt_fused_tensors = decay_fused_tensors + gate + other + expert
    reduce_fused_tensors = d_other + other
    gate_fused_tensors = d_gate + gate

    expert_fusion_names = []
    for p in d_expert + expert:
        expert_fusion_names.append(p.name)

    return opt_fused_tensors, decay_fused_tensors, reduce_fused_tensors, gate_fused_tensors, expert_fusion_names

def timer_log(log_freq):
    timers = get_timers()
    # Logging
    timers_to_log = []

    def add_to_logging(name):
        if name in timers.timers:
            timers_to_log.append(name)

    add_to_logging('forward-compute')
    add_to_logging('forward-recv')
    add_to_logging('forward-send')
    add_to_logging('forward-send-backward-recv')
    add_to_logging('backward-compute')
    add_to_logging('backward-recv')
    add_to_logging('backward-send')
    add_to_logging('backward-send-forward-recv')
    add_to_logging('backward-params-all-reduce')
    add_to_logging('backward-embedding-all-reduce')
    add_to_logging('optimizer-copy-to-main-grad')
    add_to_logging('optimizer-unscale-and-check-inf')
    add_to_logging('optimizer-clip-main-grad')
    add_to_logging('optimizer-copy-main-to-model-params')
    add_to_logging('optimizer')
    add_to_logging('batch-generator')
    add_to_logging('Prepare Forward')
    add_to_logging('Gate Computation')
    add_to_logging('Limit_By_Capacity')
    add_to_logging('Prune_Gate_By_Cap')
    add_to_logging('Random Routing')
    add_to_logging('Base Operation')
    add_to_logging('AllGather in Limit')
    add_to_logging('MOEScatter')
    add_to_logging('Expert Computation')
    add_to_logging('MOEGather')
    add_to_logging('Score BMM')
    add_to_logging('AllReduce')
    add_to_logging('AllGather')
    add_to_logging('lec reduce')
    add_to_logging('lec reduce2')

    timers.log(timers_to_log, normalizer=log_freq)

def do_train(args):
    paddle.set_device(args.device)
    strategy = fleet.DistributedStrategy()
    strategy.hybrid_configs = {
        "dp_degree": args.dp_degree,
        "mp_degree": args.mp_degree,
        "pp_degree": args.pp_degree
    }

    accumulate_steps = args.local_batch_size // args.micro_batch_size
    strategy.pipeline_configs = {
        "accumulate_steps": accumulate_steps,
        "micro_batch_size": args.micro_batch_size
    }

    fleet.init(is_collective=True, strategy=strategy)

    # obtain rank message of hybrid parallel
    hcg = fleet.get_hybrid_communicate_group()
    global_rank = hcg.get_global_rank()
    mp_rank = hcg.get_model_parallel_rank()
    pp_rank = hcg.get_stage_id()
    dp_rank = hcg.get_data_parallel_rank()
    local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))

    # seed control in hybrid parallel
    set_hyrbid_parallel_seed(args.seed, dp_rank, mp_rank, pp_rank)

    default_global_tokens_num = args.global_batch_size * args.max_seq_len

    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)

    # Define log writer
    log_writer_path = os.path.join(
        args.output_dir, "train_log",
        "{}_globalbsz_{}_pure_fp16_{}_recompute_{}_card_{}".format(
            args.model_name_or_path, args.global_batch_size, args.use_pure_fp16,
            False, global_rank).lower())

    if os.path.exists(log_writer_path):
        import shutil
        shutil.rmtree(log_writer_path)

    log_writer = LogWriter(log_writer_path)

    pretrained_models_list = list(
        model_class.pretrained_init_configuration.keys())

    if args.model_name_or_path in pretrained_models_list:
        model_config = model_class.pretrained_init_configuration[
            args.model_name_or_path]
        model_config["hidden_dropout_prob"] = args.hidden_dropout_prob
        model_config[
            "attention_probs_dropout_prob"] = args.attention_probs_dropout_prob

        model_config['num_partitions'] = args.mp_degree

        # MOE config
        initialize_model_and_expert_group(hcg)

        model_config['expert_mode'] = args.expert_mode
        model_config['hcg'] = hcg
        model_config['num_experts'] = args.num_experts
        model_config['top_k'] = args.top_k
        if args.expert_mode:
            model_config['gate'] = args.gate

        if args.pp_degree == 1:
            model_config["recompute_interval"] = 1 if args.use_recompute else 0
            model_config["recompute_partition"] = args.recompute_partition
            model_config["recompute_offload"] = args.recompute_offload
            if args.use_recompute and args.recompute_partition:
                raise Exception(
                    "when use_recompute is True, recompute_partition must be False in MoE."
                )

            model = GPTForPretraining(GPTModel(**model_config))
        else:
            model_config['topology'] = hcg.topology()
            model_config["recompute_interval"] = 1 if args.use_recompute else 0
            model = GPTForPretrainingPipe(**model_config)
    else:
        model = GPTForPretraining.from_pretrained(
            args.model_name_or_path,
            hidden_dropout_prob=args.hidden_dropout_prob,
            attention_probs_dropout_prob=args.attention_probs_dropout_prob)

    # Create the critrion for the gpt model
    criterion = GPTPretrainingCriterion()

    if args.decay_steps is None:
        args.decay_steps = args.max_steps
    warmup_step = args.warmup_rate * args.decay_steps

    lr_scheduler = None

    if args.lr_decay_style == "none":
        lr_scheduler = None
    elif args.lr_decay_style == "cosine":
        lr_scheduler = lr.CosineAnnealingWithWarmupDecay(
            max_lr=args.max_lr,
            min_lr=args.min_lr,
            warmup_step=warmup_step,
            decay_step=args.decay_steps)

# Generate parameter names needed to perform weight decay.
# All bias and LayerNorm parameters are excluded.
    if args.use_pure_fp16:
        scaler = paddle.amp.GradScaler(init_loss_scaling=args.scale_loss)
        scaler = fleet.distributed_scaler(scaler)
        scaler._unscale = MethodType(unscale_method, scaler)
        model = paddle.amp.decorate(
            models=model, optimizers=None, level='O2', save_dtype='float32')

    opt_fused_tensors, decay_fused_tensors, reduce_fused_tensors, gate_fused_tensors, \
        expert_fusion_names = parameters_classify(model)
    decay_params = [p.name for p in decay_fused_tensors]

    clip = None
    if args.grad_clip > 0:
        is_expert_param_fun = lambda param: param.name in expert_fusion_names
        clip = moe.ClipGradByGlobalNorm(clip_norm=args.grad_clip, \
                                        is_expert_param_func = is_expert_param_fun, \
                                        moe_group = hcg.get_expert_parallel_group())

    optimizer = AdamW(
        learning_rate=lr_scheduler if lr_scheduler is not None else args.max_lr,
        beta1=args.adam_beta1,
        beta2=args.adam_beta2,
        epsilon=args.adam_epsilon,
        parameters=opt_fused_tensors,
        weight_decay=args.weight_decay,
        grad_clip=clip,
        apply_decay_param_fun=lambda x: x in decay_params,  #decay_params,
        multi_precision=args.use_pure_fp16)

    if paddle.distributed.get_world_size() > 1 and args.resume_dir is None:
        print(">> initialize....")
        initialize_mp_dp_parameters(model, hcg)

    #in order to restore reader.
    pass_num = 0
    file_id = 0
    start_epoch = 0
    args.resume_dir = None if len(args.resume_dir) <= 0 else args.resume_dir

    if args.resume_dir is not None:
        global_step, loss_scale, data_meta = load_checkpoint(
            args, model, optimizer, lr_scheduler, tokenizer, dp_rank, mp_rank,
            pp_rank)
        pass_num = data_meta["pass_num"]
        file_id = data_meta["file_id"]
        start_epoch = data_meta["start_epoch"]

    if args.use_pure_fp16:
        scaler = paddle.amp.GradScaler(init_loss_scaling=loss_scale
                                       if args.resume_dir is not None else
                                       args.scale_loss)
        scaler = fleet.distributed_scaler(scaler)
        scaler._unscale = MethodType(unscale_method, scaler)
        model, optimizer = paddle.amp.decorate(
            models=model,
            optimizers=optimizer,
            level='O2',
            save_dtype='float32')

    if args.model_name_or_path not in pretrained_models_list:
        logger.info("Try to load checkpoint from %s " % args.model_name_or_path)
        opt_path = os.path.join(args.model_name_or_path, "model_state.pdopt")
        if os.path.exists(opt_path):
            opt_dict = paddle.load(opt_path)
            optimizer.set_state_dict(opt_dict)
        else:
            logger.warning("No optimizer checkpoint file found in %s." %
                           opt_path)

    global_step = 0 if args.resume_dir is None else global_step
    timers = get_timers()
    tic_train = time.time()
    for epoch in range(start_epoch, args.num_train_epochs):
        files = [
            os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir)
            if (os.path.isfile(os.path.join(args.input_dir, f)) and "npz_"
                not in str(f))
        ]
        files.sort()
        num_files = len(files)
        for f_id in range(file_id, num_files):
            data_file = files[f_id]
            train_data_loader, valid_data_loader, test_data_loader = create_pretrained_dataset(
                args,
                data_file,
                local_rank=local_rank,
                data_world_size=args.dp_degree,
                data_world_rank=dp_rank,
                eos_id=tokenizer.eos_token_id)

            # Bug fix, if not call valid_data_loader, the enumerate will call valid_data_loader
            # many times. and start a new random dataloader.
            valid_data_loader = valid_data_loader()
            test_data_loader = test_data_loader()

            for step, batch in enumerate(train_data_loader()):
                # to remove the train data that has been studyed.
                if step < global_step - pass_num: continue

                global_step += 1
                tokens, loss_mask, labels = batch

                loss_mask.stop_gradient = True
                labels.stop_gradient = True

                loss = 0.0
                for i in range(accumulate_steps):
                    start_index = i * args.micro_batch_size
                    end_index = start_index + args.micro_batch_size
                    timers('forward-compute').start()
                    with paddle.amp.auto_cast(
                            args.use_pure_fp16,
                            custom_black_list=[
                                "reduce_sum",
                                "c_softmax_with_cross_entropy",
                                "elementwise_div",
                            ],
                            level='O2'):
                        preds = model(tokens[start_index:end_index, :])
                        loss_mbs = criterion(
                            preds, labels[start_index:end_index, :],
                            loss_mask[start_index:end_index, :])
                    timers('forward-compute').stop()

                    if args.gate != "naive" and args.balance_loss_weight:
                        aux_loss_list = [
                            l.moe_mlp.gate.get_loss(clear=False)
                            for l in model.gpt.decoder.layers
                            if hasattr(l.moe_mlp, "gate")
                        ]
                        bal_loss = paddle.concat(aux_loss_list)
                        if bal_loss.dtype == paddle.float16:
                            bal_loss = paddle.cast(
                                bal_loss, dtype=paddle.float32)
                        bal_loss = bal_loss.mean()
                        loss_mbs += bal_loss * args.balance_loss_weight
                    loss_mbs = loss_mbs / accumulate_steps

                    timers('backward-compute').start()
                    if args.use_pure_fp16:
                        scaler.scale(loss_mbs).backward()
                    else:
                        loss_mbs.backward()
                    timers('backward-compute').stop()
                    loss = loss + loss_mbs

                timers('backward-params-all-reduce').start()
                all_reduce_parameters(gate_fused_tensors,
                                      hcg.get_expert_parallel_group())
                all_reduce_parameters(reduce_fused_tensors,
                                      hcg.get_data_parallel_group())
                timers('backward-params-all-reduce').stop()

                if args.use_pure_fp16:
                    scaler.minimize(optimizer, loss)
                else:
                    optimizer.step()
                learning_rate = optimizer.get_lr()
                if lr_scheduler is not None:
                    lr_scheduler.step()
                optimizer.clear_grad()

                if global_step % args.logging_freq == 0:
                    avg_loss = loss.numpy()
                    speed = args.logging_freq / (time.time() - tic_train)
                    if args.gate != "naive" and args.balance_loss_weight:
                        bal_loss = bal_loss.numpy()
                        avg_loss -= bal_loss
                    else:
                        bal_loss = -1
                    logger.info(
                        "global step %d, epoch: %d, batch: %d, loss: %.9f, bal_loss: %.9f, speed: %.2f step/s, ips: %.0f tokens/s, learning rate: %.5e"
                        % (global_step, epoch, step, avg_loss, bal_loss, speed,
                           speed * default_global_tokens_num, learning_rate))
                    log_writer.add_scalar("loss", float(loss), global_step)
                    log_writer.add_scalar("learning_rate", learning_rate,
                                          global_step)

                    tic_train = time.time()
                    timer_log(args.logging_freq)

                if (global_step % args.save_steps == 0 or
                        global_step >= args.max_steps):
                    loss_scale = scaler._scale if args.use_pure_fp16 else None
                    save_checkpoint(args, global_step, model, optimizer,
                                    lr_scheduler, tokenizer, loss_scale,
                                    dp_rank, mp_rank, pp_rank, pass_num,
                                    file_id, epoch)
                    print(
                        "save checkpoint for step_{} successfully...loss_scale = {}".
                        format(global_step, loss_scale))

                if global_step % args.eval_freq == 0:
                    # Since the valid data broardcast to all devices, we do evaluate on all device.
                    run_evaluate(args, valid_data_loader, model, criterion,
                                 args.eval_iters, log_writer, global_step,
                                 epoch, "valid")

                if global_step >= args.max_steps:
                    run_evaluate(args, test_data_loader, model, criterion,
                                 args.test_iters, log_writer, global_step,
                                 epoch, "test")
                    logger.info("The training process is complete.")
                    del train_data_loader
                    return

            # to record sum of the length of train_data_loader that has been read.
            pass_num += len(train_data_loader())
            del train_data_loader

if __name__ == "__main__":
    args = parse_args(MODEL_CLASSES)
    do_train(args)