#-------------------------------- main.py file -----------------------------------------#
"""
Main file - entry point to the code. This file coordinates all other files and
implements all the functionality directly available to the user.
"""

# import statements
import numpy as np
import joblib
from pkg_resources import get_distribution
from boreas.models import makePrediction, RFModelIsotropic, TBNNSModelAnisotropic
from boreas.case import TestCase, TrainingCase
from boreas import process
from boreas import constants

def printInfo():
    """
    Makes sure everything is properly installed.

    We print a welcome message, the version of the package, and attempt to load
    the pre-trained models to make sure the data file is there. Return 1 at the end
    if no exceptions were raised.
    """

    print('Welcome to Boreas - a package for industrial deployment of machine-learned '
          + 'turbulent mixing models for film cooling (formerly known as RaFoFC)!')

    # Get distribution version
    dist = get_distribution('boreas')
    print('Version: {}'.format(dist.version))

    # Try to load the default RF model and print information about it
    print('Attempting to load the default RF model...')
    rf = RFModelIsotropic()
    rf.loadFromDisk()
    print('Default model was found and can be loaded properly.')
    print('\t Description: ', end="", flush=True)
    rf.printDescription()

    # Try to load the default TBNNS model and print information about it
    print('Attempting to load the default TBNN-s model...')
    nn = TBNNSModelAnisotropic()
    nn.loadFromDisk()
    print('Default model was found and can be loaded properly.')
    print('\t Description: ', end="", flush=True)
    nn.printDescription()

    return 1 # return this if everything went ok

def applyMLModel(tecplot_in_path, tecplot_out_path, *,
                 zone = None, deltaT0 = None,
                 use_default_var_names = False, use_default_derivative_names = True,
                 calc_derivatives = True, write_derivatives = True,
                 threshold = None, default_prt = None, clean_features = True,
                 features_load_path = None, features_dump_path = None,
                 ip_file_path = None, csv_file_path = None,
                 variables_to_write = None, outnames_to_write = None,
                 model_path = None, secondary_model_path = None,
                 model_type = "RF", features_type="F2",
                 ensemble_of_models = False, std_ensemble = False):
    """
    Applies ML model on a single test case, given in a Tecplot file.

    Main function of package. Call this to take in a Tecplot file, process it, apply
    the machine learning model, and save results to disk. All optional arguments must
    be used with the identifying keyword (that's what * means)

    Arguments:
    tecplot_in_path -- string containing the path of the input tecplot file. It must be
                       a binary .plt file, resulting from a k-epsilon simulation.
    tecplot_out_path -- string containing the path to which the final tecplot dataset
                        will be saved.
    zone -- optional argument. The zone where the flow field solution is saved in
            Tecplot. By default, it is zone 0. This can be either a string (with the
            zone name) or an integer with the zone index.
    deltaT0 -- optional argument. Temperature scale (Tmax - Tmin) that will be used to
               non-dimensionalize the dataset. If it is not provided (default behavior),
               the user will be prompted to enter an appropriate number.
    use_default_var_names -- optional argument. Boolean flag (True/False) that determines
                             whether default Fluent names will be used to fetch variables
                             in the Tecplot dataset. If the flag is False (default
                             behavior), the user will be prompted to enter names for each
                             variable that is needed.
    use_default_derivative_names -- optional argument. Boolean flag (True/False) that
                                    determine if the user will pick the names for the
                                    derivative quantities in the Tecplot file or whether
                                    default names are used. This flag is only used if the
                                    next flag is False (i.e., if derivatives are
                                    already pre-calculated, then setting this flag to
                                    False allows the user to input the names of each
                                    derivative in the input .plt file). It defaults to
                                    True.
    calc_derivatives -- optional argument. Boolean flag (True/False) that determines
                        whether derivatives need to be calculated in the Tecplot file.
                        Note we need derivatives of U, V, W, and Temperature, with names
                        ddx_{}, ddy_{}, ddz_{}. If such variables were already calculated
                        and exist in the dataset, set this flag to False to speed up the
                        process. By default (True), derivatives are calculated and a new
                        file with derivatives called "derivatives_{}" will be saved to
                        disk.
    write_derivatives -- optional argument. Boolean flag (True/False) that determines
                         whether to write a binary Tecplot file to disk with the newly
                         calculated derivatives. The file will have the same name as the
                         input, except followed by "_derivatives". This is useful because
                         calculating derivatives takes a long time, so you might want to
                         save results to disk as soon as they are calculated.
    threshold -- optional argument. This variable determines the threshold for
                 (non-dimensional) temperature gradient below which we throw away a
                 point. If None, use the value in constants.py (default value is 1e-3).
                 For temperature gradient less than that, we use the Reynolds analogy
                 (with fixed Pr_t). For gradients larger than that, we use the
                 model.
    default_prt -- optional argument, this variable contains the default value of Pr_t to
                   use in regions where gradients are low or features have been cleaned.
                   If this is None (default), then use the value from constants.py.
    clean_features -- optional argument. This determines whether we should remove outlier
                      points from the dataset before applying the model. This is measured
                      by the standard deviation of points around the mean.
    features_load_path -- optional argument. If this is supplied, then the function will
                           try to load the features from disk instead of
                           processing the tecplot file all over again. Since calculating
                           the features can take a while for large datasets, this can be
                           useful to speed up repetitions.
    features_dump_path -- optional argument. If this is provided and we processed the
                           tecplot data from scratch (i.e. we calculated the features),
                           then the function will save the features to disk, so it is
                           much faster to perform the same computations again later.
    ip_file_path -- optional argument. String containing the path to which the
                    interpolation file (which is read by ANSYS Fluent) will be saved. If
                    this argument is None (by default), then no interpolation file is
                    written.
    csv_file_path -- optional argument. String containing the path to which the csv file
                     (which can be read by StarCCM+) will be saved. If this is None
                     (default), then no csv file is written.
    variables_to_write -- optional argument. This is a list of strings containing names
                          of variables in the Tecplot file that we want to write in the
                          Fluent interpolation file/CSV file. By default, it is None,
                          which leads the program to pick only the diffusivity variables
                          just calculated.
    outnames_to_write -- optional argument. This is a list of strings that must have the
                        same length as the previous argument. It contains the names that
                        each of the variables written in the interpolation/csv files will
                        have. By default, this is None, which leads to code to name all
                        variables being written sequentially, starting at "uds-2". Naming
                        them as "user defined scalars x" (uds-x) is an easy way to read
                        them in Fluent.
    model_path -- optional argument. This is the path where the function will look for
                  a pre-trained machine learning model. The file must be a pickled
                  instance of a random forest regressor class or a pickled instance of
                  the TBNN-s class, saved to disk using joblib. If None, the default
                  machine learning model that comes with the package(which is already
                  pre-trained with LES/DNS) is employed.
    secondary_model_path -- optional argument. This is the path where the function will
                            look for a pre-trained random forest model to support the
                            TBNN-s model in the hybrid formulation. The file must be a
                            pickled instance of a random forest regressor class, saved to
                            disk using joblib. By default, the default RF is loaded. This
                            argument is only relevant when model_type = "TBNNS_hybrid".
    model_type -- optional argument. This tells us which type of model we are loading.
                  It must be a string, and the currently supported options are "RF",
                  "TBNNS", and "TBNNS_hybrid". The default option is "RF".
    features_type -- optional argument, string determining the type of features that
                     we are currently extracting. Options are "F1" and "F2". Default
                     value is "F2".
    ensemble_of_models -- optional argument. This is a boolean flag that tells us whether
                          to use a model ensemble instead of a single model instance. If
                          this is true, the model_path parameter must be a list of paths
                          instead of a single path. The default option is "False"
    std_ensemble -- optional argument. This is a boolean flag that instructs the solver
                    to return the standard deviation across the ensemble of models. This
                    can only be True if ensemble_of_models=True; in which case, we only
                    return the standard deviation and not the actual diffusivity. This
                    option only makes sense for the TBNN-s model (since the RF is already
                    an ensemble) and is not supported for the hybrid model. The default
                    option is "False"
    """

    assert model_type == "RF" or model_type == "TBNNS" or model_type == "TBNNS_hybrid", \
            "Invalid model_type received!"
    assert features_type == "F1" or features_type == "F2", \
            "Invalid features_type received!"

    if ensemble_of_models: # check whether model_path is a list if model ensemble
        assert type(model_path) is list, \
            "Error! For a model ensemble, model_path must be a list"
        assert len(model_path) > 0, "Error! model_path is an empty list!"

    # Initialize dataset and get scales for non-dimensionalization. The default behavior
    # is to ask the user for the names and the scales. Passing keyword arguments to this
    # function can be done to go around this behavior
    dataset = TestCase(tecplot_in_path, zone=zone,
                        use_default_names=use_default_var_names)
    dataset.normalize(deltaT0=deltaT0)

    # If this flag is True (default) calculate the derivatives and save the result to
    # disk (since it takes a while to do that...)
    if calc_derivatives:
        dataset.calculateDerivatives()
        if write_derivatives: # write new Tecplot file to disk
            dataset.saveDataset(tecplot_in_path[0:-4] + "_derivatives.plt")
    else:
        print("Derivatives already calculated!")
        dataset.addDerivativeNames(use_default_derivative_names)

    # Here, run the code for applying random forest model ("RF")
    if model_type == "RF":
        # This line processes the dataset and extracts features for the ML step which
        # can take a long time. features_load_path and features_dump_path can be
        # set to make the method load/save the processed quantities from disk.
        x, _ = dataset.extractFeatures(with_tensor_basis=False,
                                       features_type=features_type, threshold=threshold,
                                       features_load_path=features_load_path,
                                       features_dump_path=features_dump_path,
                                       clean_features=clean_features)
        prt_ML = makePrediction("RF", model_path, x, features_type)

        # Adds result to tecplot and sets the default variable names to output
        varname = "Prt_ML"
        dataset.addPrt(prt_ML, varname, default_prt)
        if variables_to_write is None:
            variables_to_write = [varname]
        if outnames_to_write is None:
            outnames_to_write = ["uds-2"]

    # Here, run the code for applying TBNN model ("TBNNS")
    elif model_type == "TBNNS":
        # This line processes the dataset and returns the features and tensor basis
        # at each point in the dataset where gradients are significant.
        x, tb = dataset.extractFeatures(with_tensor_basis=True,
                                        features_type=features_type, threshold=threshold,
                                        features_load_path=features_load_path,
                                        features_dump_path=features_dump_path,
                                        clean_features=clean_features)
        alphaij_ML, g_ML = makePrediction("TBNNS", model_path, x, features_type, tb,
                                          ensemble=ensemble_of_models,
                                          std_flag=std_ensemble)

        # Adds result to tecplot and sets the default variable names to output
        varname = ["Dxx", "Dxy", "Dxz", "Dyx", "Dyy", "Dyz", "Dzx", "Dzy", "Dzz"]
        dataset.addTensorDiff(alphaij_ML, varname, default_prt)
        g_name = ["g1", "g2", "g3", "g4", "g5", "g6"]
        dataset.addG(g_ML, g_name, default_prt)

        if variables_to_write is None:
            variables_to_write = varname
        if outnames_to_write is None:
            outnames_to_write = ["uds-2", "uds-3", "uds-4", "uds-5", "uds-6", "uds-7",
                                 "uds-8", "uds-9", "uds-10"]

    # Here, run the code for applying TBNN-s + random forest model ("TBNNS_hybrid")
    elif model_type == "TBNNS_hybrid":
        # This line processes the dataset and returns the features and tensor basis
        # at each point in the dataset where gradients are significant.
        x, tb = dataset.extractFeatures(with_tensor_basis=True,
                                        features_type=features_type, threshold=threshold,
                                        features_load_path=features_load_path,
                                        features_dump_path=features_dump_path,
                                        clean_features=clean_features)
        alphaij_ML, g_ML = makePrediction("TBNNS", model_path, x, features_type, tb,
                                          ensemble=ensemble_of_models,
                                          std_flag=std_ensemble)

        # Now, get a random forest prediction for the turbulent Prandtl number
        prt_ML = makePrediction("RF", secondary_model_path, x, features_type)

        # Combine alphaij_ML and prt_ML into a single diffusivity tensor
        alphaij_mod = dataset.enforcePrt(alphaij_ML, prt_ML)

        # Adds result to tecplot and sets the default variable names to output
        varname = ["Dxx", "Dxy", "Dxz", "Dyx", "Dyy", "Dyz", "Dzx", "Dzy", "Dzz"]
        dataset.addTensorDiff(alphaij_mod, varname, default_prt)
        g_name = ["g1", "g2", "g3", "g4", "g5", "g6"]
        dataset.addG(g_ML, g_name, default_prt)

        if variables_to_write is None:
            variables_to_write = varname
        if outnames_to_write is None:
            outnames_to_write = ["uds-2", "uds-3", "uds-4", "uds-5", "uds-6", "uds-7",
                                 "uds-8", "uds-9", "uds-10"]

    # Write output: create interp/csv files and produce tecplot file
    if ip_file_path is not None:
        dataset.createInterpFile(ip_file_path, variables_to_write, outnames_to_write)
    if csv_file_path is not None:
        dataset.createCsvFile(csv_file_path, variables_to_write, outnames_to_write)
    dataset.saveDataset(tecplot_out_path)

def produceTrainingFeatures(tecplot_in_path, *, data_path = None,
                            zone = None, deltaT0 = None,
                            use_default_var_names = False,
                            use_default_derivative_names = True,
                            calc_derivatives = True, write_derivatives = True,
                            threshold = None, clean_features = True,
                            features_load_path = None, features_dump_path = None,
                            prt_cap = None, gamma_correction = False,
                            downsample = None, tecplot_out_path = None,
                            model_type = "RF", features_type="F2"):

    """
    Produces features and labels from a single Tecplot file, used for training.

    This function is useful for training your own models. Call it on a single Tecplot
    file (.plt) that contains all mean data including u'c' values, and it will process
    it to generate the features and labels used for training. All optional arguments may
    only be used with the keyword (that's what * means)

    Arguments:
    tecplot_in_path -- string containing the path of the input tecplot file. It must be
                       a binary .plt file, resulting from a k-epsilon simulation.
    data_path -- optional argument. A string containing the path where a joblib file is
                 saved, containing features and labels for training ML models. If None
                 (default), a default name is employed.
    zone -- optional argument. The zone where the flow field solution is saved in
            Tecplot. By default, it is zone 0. This can be either a string (with the
            zone name) or an integer with the zone index.
    deltaT0 -- optional argument. Temperature scale (Tmax - Tmin) that will be used to
               non-dimensionalize the dataset. If it is not provided (default behavior),
               the user will be prompted to enter an appropriate number.
    use_default_var_names -- optional argument. Boolean flag (True/False) that determines
                             whether default Fluent names will be used to fetch variables
                             in the Tecplot dataset. If the flag is False (default
                             behavior), the user will be prompted to enter names for each
                             variable that is needed.
    use_default_derivative_names -- optional argument. Boolean flag (True/False) that
                                    determine if the user will pick the names for the
                                    derivative quantities in the Tecplot file or whether
                                    default names are used. This flag is only used if the
                                    next flag is False (i.e., if derivatives are
                                    already pre-calculated, then setting this flag to
                                    False allows the user to input the names of each
                                    derivative in the input .plt file). It defaults to
                                    True.
    calc_derivatives -- optional argument. Boolean flag (True/False) that determines
                        whether derivatives need to be calculated in the Tecplot file.
                        Note we need derivatives of U, V, W, and Temperature, with names
                        ddx_{}, ddy_{}, ddz_{}. If such variables were already calculated
                        and exist in the dataset, set this flag to False to speed up the
                        process. By default (True), derivatives are calculated and a new
                        file with derivatives called "derivatives_{}" will be saved to
                        disk.
    write_derivatives -- optional argument. Boolean flag (True/False) that determines
                         whether to write a binary Tecplot file to disk with the newly
                         calculated derivatives. The file will have the same name as the
                         input, except followed by "_derivatives". This is useful because
                         calculating derivatives takes a long time, so you might want to
                         save results to disk as soon as they are calculated.
    threshold -- optional argument. This variable determines the threshold for
                 (non-dimensional) temperature gradient below which we throw away a
                 point. If None, use the value in constants.py (default value is 1e-3).
                 For temperature gradient less than that, we use the Reynolds analogy
                 (with fixed Pr_t). For gradients larger than that, we use the
                 model.
    clean_features -- optional argument. This determines whether we should remove outlier
                      points from the dataset before applying the model. This is measured
                      by the standard deviation of points around the mean.
    features_load_path -- optional argument. If this is supplied, then the function will
                          try to load the features from disk instead of
                          processing the tecplot file all over again. Since calculating
                          the features can take a while for large datasets, this can be
                          useful to speed up repetitions.
    features_dump_path -- optional argument. If this is provided and we processed the
                          tecplot data from scratch (i.e. we calculated the features),
                          then the function will save the features to disk, so it is
                          much faster to perform the same computations again later.
    prt_cap -- optional, contains the (symmetric) cap on the value of Pr_t. If None,
               then use the value in constants.py. If this value is 100, for example,
               then 0.01 < Pr_t < 100, and values outside of this range are capped.
    gamma_correction -- optional. If True, use the correction defined in
                        Milani, Ling, Eaton (JTM 2020). That correction only makes
                        sense if training data is on a fixed reference frame (it breaks
                        Galilean invariance), so it is turned off by default. However,
                        it can improve results in some cases.
    downsample -- optional, number that controls how we downsample the data before
                  saving it to disk. If None (default), it will read the number from
                  constants.py. If this number is more than 1, then it represents the
                  number of examples we want to save; if it is less than 1, it represents
                  the ratio of all training examples we want to save.
    tecplot_out_path -- optional, a string containing the path to which the final tecplot
                        dataset will be saved. Useful for sanity checking the results. By
                        default it is None (no .plt file saved)
    model_type -- optional argument. This tells us which type of model we are loading.
                  It must be a string, and the currently supported options are "RF".
                  The default option is "RF".
    features_type -- optional argument, string determining the type of features that
                     we are currently extracting. Options are "F1" and "F2". Default
                     value is "F2".
    """

    assert model_type == "RF" or model_type == "TBNNS" or model_type == "TBNNS_hybrid", \
           "Invalid model_type received!"

    # Initialize dataset and get scales for non-dimensionalization. The default behavior
    # is to ask the user for the names and the scales. Passing keyword arguments to this
    # function can be done to go around this behavior
    dataset = TrainingCase(tecplot_in_path, zone=zone,
                           use_default_names=use_default_var_names)
    dataset.normalize(deltaT0=deltaT0)

    # If this flag is True (default) calculate the derivatives and save the result to
    # disk (since it takes a while to do that...)
    if calc_derivatives:
        dataset.calculateDerivatives()
        if write_derivatives: # write new Tecplot file to disk
            dataset.saveDataset(tecplot_in_path[0:-4] + "_derivatives.plt")
    else:
        print("Derivatives already calculated!")
        dataset.addDerivativeNames(use_default_derivative_names)

    metadata = {}
    metadata["features_type"] = features_type
    if model_type == "RF":
        # This line processes the dataset and extracts features for the ML step which
        # can take a long time. features_load_path and features_dump_path can be
        # set to make the method load/save the processed quantities from disk.
        x, _ = dataset.extractFeatures(with_tensor_basis=False,
                                       features_type=features_type, threshold=threshold,
                                       features_load_path=features_load_path,
                                       features_dump_path=features_dump_path,
                                       clean_features=clean_features)
        gamma = dataset.extractGamma(prt_cap, gamma_correction) # gamma = 1/Prt

        training_list = [x, gamma]  # what is used for training
        metadata["with_tensor_basis"]=False
        metadata["with_gamma"]=True

        # Write the Tecplot data to disk with the extracted Prt_LES for sanity check
        if tecplot_out_path is not None:
            dataset.addPrt(1.0/gamma, "Prt_LES")
            dataset.saveDataset(tecplot_out_path)

    elif model_type == "TBNNS":
        # This line processes the dataset and returns the features and tensor basis
        # at each point in the dataset where gradients are significant.
        x, tb = dataset.extractFeatures(with_tensor_basis=True,
                                        features_type=features_type, threshold=threshold,
                                        features_load_path=features_load_path,
                                        features_dump_path=features_dump_path,
                                        clean_features=clean_features)
        uc, gradT, nut = dataset.extractUc()

        training_list = [x, tb, uc, gradT, nut]  # what is used for training
        metadata["with_tensor_basis"]=True
        metadata["with_gamma"]=False

    elif model_type == "TBNNS_hybrid":
        # This line processes the dataset and returns the features and tensor basis
        # at each point in the dataset where gradients are significant.
        x, tb = dataset.extractFeatures(with_tensor_basis=True,
                                        features_type=features_type, threshold=threshold,
                                        features_load_path=features_load_path,
                                        features_dump_path=features_dump_path,
                                        clean_features=clean_features)

        uc, gradT, nut = dataset.extractUc()
        gamma = dataset.extractGamma(prt_cap, gamma_correction)

        training_list = [x, tb, uc, gradT, nut, gamma]  # what is used for training
        metadata["with_tensor_basis"]=True
        metadata["with_gamma"]=True

    # Saves joblib file to disk with features/labels for this dataset
    # If data_path is None, use default name (appending _trainingdata to the end)
    if data_path is None:
        data_path = tecplot_in_path[0:-4] + "_trainingdata.pckl" # default name

    # Save training features to disk
    process.saveTrainingFeatures(training_list, metadata, data_path, downsample)

def trainRFModel(features_list, description, model_path, *, features_type="F1",
                 downsample=None, n_trees = None, max_depth = None,
                 min_samples_split = None, n_jobs = None):
    """
    Trains a random forest model and saves it to disk.

    Trains a random forest, isotropic model, with features and labels previously
    calculated. Multiple files can be used at the same time (each file in the list
    comes from a given dataset. All optional arguments may only be used with the
    keyword (that's what * means)

    Arguments:
    features_list -- list containing paths to files saved to disk with features
                     and labels for training. These files are produced by the function
                     above (produceTrainingFeatures) from Tecplot files.
    description -- A short, written description of the model being trained. It will be
                   saved to disk together with the model itself.
    model_path -- The path where the trained model will be saved in disk.
    features_type -- optional argument, string determining the type of features that
                     we are currently extracting. Options are "F1" and "F2". Default
                     value is "F1".
    downsample -- optional, number that controls how we downsample the data before
                  using it to train. If None (default), it will read the number from
                  constants.py. If this number is more than 1, then it represents the
                  number of examples we want to save; if it is less than 1, it represents
                  the ratio of all training examples we want to save. Can also be a list
                  of numbers, in which case each number is applied to an element of
                  features_list
    n_trees -- optional. Hyperparameter of the random forest, contains number of
                   trees to use. If None (default), reads value from constants.py
    max_depth -- optional. Hyperparameter of the random forest, contains maximum
                 depth of each tree to use. If None (default), reads value from
                 constants.py
    min_samples_split -- optional. Hyperparameter of the random forest, contains
                         minimum number of samples at a node required to split. Can
                         either be an int (number itself) or a float (ratio of total
                         examples). If None (default), reads value from constants.py
    n_jobs -- optional. Number of processors to use when training the RF (notice that
              training is embarassingly parallel). If None (default behavior), then
              the value is read from constants.py. See manual for
              RandomForestRegressor class; if this is -1, all processors are used.
    """

    # Reads the list of files provided for features/labels
    print("{} file(s) were provided and will be used".format(len(features_list)))
    x_list = []
    y_list = []

    if isinstance(downsample, list): # make sure list is the right size
        assert len(downsample) == len(features_list), \
           "downsample is a list, but it has the wrong number of entries!"

    for i, file in enumerate(features_list):
        if isinstance(downsample, list): # if list, take each element sequentially
            x_features, gamma = process.loadTrainingFeatures(file, "RF", downsample[i],
                                                             features_type)
        else:
            x_features, gamma = process.loadTrainingFeatures(file, "RF", downsample,
                                                             features_type)

        x_list.append(x_features)
        y_list.append(gamma)

    x_total = np.concatenate(x_list, axis=0)
    y_total = np.concatenate(y_list, axis=0)

    # Here, we train and save the model
    rf = RFModelIsotropic()
    rf.train(x_total, y_total, n_trees, max_depth, min_samples_split, n_jobs)
    rf.save(description, model_path)

def trainTBNNSModel(features_list_train, features_list_dev, description, model_path,
                    path_to_saver, *, FLAGS={}, features_type="F2",
                    downsample_train=None, downsample_dev=None,):
    """
    Trains a TBNN-s and saves it to disk.

    Trains a TBNN-s, anisotropic model, with features and labels previously
    calculated. Multiple files can be used at the same time (each file in the list
    comes from a given dataset. All optional arguments may only be used with the
    keyword (that's what * means)

    Arguments:
    features_list_train -- list containing paths to files saved to disk with features
                     and labels for training. These files are produced by the function
                     above (produceTrainingFeatures) from Tecplot files.
    features_list_dev -- list containing paths to files saved to disk with features
                     and labels for the validation set. These files are produced by
                     the function above (produceTrainingFeatures) from Tecplot files.
    description -- A short, written description of the model being trained. It will be
                   saved to disk together with the model itself.
    model_path -- The path where the trained model will be saved in disk.
    path_to_saver -- The path where the model parameters are saved to disk, through the
                    tf.Saver class. Usually, we want to put that in a folder called
                    checkpoints.
    FLAGS -- optional argument, dictionary that controls training parameters for
             the TBNNS model. Check the tbnns package to see what settings can be used.
    features_type -- optional argument, string determining the type of features that
                     we are currently extracting. Options are "F1" and "F2". Default
                     value is "F2".
    downsample_train -- optional, number that controls how we downsample the data before
                  using it to train. If None (default), it will read the number from
                  constants.py. If this number is more than 1, then it represents the
                  number of examples we want to save; if it is less than 1, it represents
                  the ratio of all training examples we want to save. Can also be a list
                  of numbers, in which case each number is applied to an element of
                  features_list.
    downsample_dev -- optional, same as above for the dev set.
    """

    # Reads the list of files provided for features/labels
    print("{} file(s) were provided and will be used for training"\
            .format(len(features_list_train)))
    print("{} file(s) were provided and will be used for validation"\
            .format(len(features_list_dev)))

    # Makes sure downsample list is the right size
    if isinstance(downsample_train, list):
        assert len(downsample_train) == len(features_list_train), \
           "downsample is a list, but it has the wrong number of entries!"
    if isinstance(downsample_dev, list): # make sure list is the right size
        assert len(downsample_dev) == len(features_list_dev), \
           "downsample is a list, but it has the wrong number of entries!"

    # Load training files
    x_list_train=[]; tb_list_train=[]; uc_list_train=[];
    gradT_list_train=[]; nut_list_train=[]
    for i, file in enumerate(features_list_train):
        if isinstance(downsample_train, list): # if list, take each element sequentially
            x, tb, uc, gradT, nut = process.loadTrainingFeatures(file, "TBNNS",
                                                                 downsample_train[i],
                                                                 features_type)
        else:
            x, tb, uc, gradT, nut = process.loadTrainingFeatures(file, "TBNNS",
                                                                 downsample_train,
                                                                 features_type)
        x_list_train.append(x); tb_list_train.append(tb); uc_list_train.append(uc)
        gradT_list_train.append(gradT); nut_list_train.append(nut)
    x_train = np.concatenate(x_list_train, axis=0)
    tb_train = np.concatenate(tb_list_train, axis=0)
    uc_train = np.concatenate(uc_list_train, axis=0)
    gradT_train = np.concatenate(gradT_list_train, axis=0)
    nut_train = np.concatenate(nut_list_train, axis=0)

    # Load dev files
    x_list_dev=[]; tb_list_dev=[]; uc_list_dev=[];
    gradT_list_dev=[]; nut_list_dev=[]
    for i, file in enumerate(features_list_dev):
        if isinstance(downsample_dev, list): # if list, take each element sequentially
            x, tb, uc, gradT, nut = process.loadTrainingFeatures(file, "TBNNS",
                                                                 downsample_dev[i],
                                                                 features_type)
        else:
            x, tb, uc, gradT, nut = process.loadTrainingFeatures(file, "TBNNS",
                                                                 downsample_dev,
                                                                 features_type)
        x_list_dev.append(x); tb_list_dev.append(tb); uc_list_dev.append(uc)
        gradT_list_dev.append(gradT); nut_list_dev.append(nut)
    x_dev = np.concatenate(x_list_dev, axis=0)
    tb_dev = np.concatenate(tb_list_dev, axis=0)
    uc_dev = np.concatenate(uc_list_dev, axis=0)
    gradT_dev = np.concatenate(gradT_list_dev, axis=0)
    nut_dev = np.concatenate(nut_list_dev, axis=0)

    # Edit FLAGS if necessary:
    if features_type=="F1": FLAGS['num_features'] = constants.NUM_FEATURES_F1
    elif features_type=="F2": FLAGS['num_features'] = constants.NUM_FEATURES_F2

    # Here, we train and save the model
    nn = TBNNSModelAnisotropic()
    nn.train(FLAGS, path_to_saver,
             x_train, tb_train, uc_train, gradT_train, nut_train,
             x_dev, tb_dev, uc_dev, gradT_dev, nut_dev)
    nn.save(description, model_path)


# This file is part of LayerModel_lib
#
#     A tool to compute the transmission behaviour of plane electromagnetic waves
#     through human tissue.
#
# Copyright (C) 2018 Jan-Christoph Brumm
#
# Licensed under MIT license.
#
import numpy as np
import os

from LayerModel_lib.voxelmodel import VoxelModel
from LayerModel_lib.voxelmodel_importer import VoxelModelImporter
from LayerModel_lib.coordinate import Coordinate

current_directory = os.path.dirname(__file__)
base_path = os.path.join(current_directory, '..', '..', '..', '..', '..', 'Numerical Human Phantoms', 'Katja')

# path to the AVW File of this model
filename = os.path.join(base_path, 'Katja')
# path to the tissue_mapping file
tissue_file = os.path.join('ImportKatja_tissues.txt')

AVW_Data = VoxelModelImporter(filename, tissue_file, 'AVW')
model_orig = AVW_Data.data['image']
tissue_name_orig = AVW_Data.tissue_names
tissue_mapping = AVW_Data.tissue_mapping

Katja = VoxelModel()
Katja.show_progress_bar = True

# needs to be set manually from README.txt
Katja.set_scale(1.775, 1.775, 4.84)

Katja.name = 'Katja'
Katja.description = 'Katja model from the Helmholtz Zentrum München. ' \
                    'Resolution %.2fmm x %.2fmm x %.2fmm' % (Katja.scaling.x,
                                                             Katja.scaling.y, Katja.scaling.z)

# For some reason Katja needs to be shifted left and back circularly
# first the shift left
model_orig = np.hstack((model_orig[:, 210::, :], model_orig[:, 0:210, :]))
# then shift back
model_orig = np.vstack((model_orig[16::, :, :], model_orig[0:16, :, :]))
# Katja has inverse x coordinates in the model. Correct that by flipping along the first axis
model_orig = np.flip(model_orig, axis=0)

# Calculate the outer_shape of the original and the complete model
outer_shape = AVW_Data.calculate_outer_shape(model_orig, tissue_mapping)

Katja.add_voxel_data(short_name='original',
                     name='Original data from AVW file',
                     model=model_orig,
                     outer_shape=outer_shape,
                     tissue_names=tissue_name_orig)

Katja.add_voxel_data(short_name='complete',
                     name='The \'original\' model converted to our TissueProperties.',
                     model=Katja.models['original'].data,
                     outer_shape=outer_shape,
                     tissue_mapping=tissue_mapping)

# Calculate the trunk model
start_slice = int(177)
end_slice = int(290)

(model_trunk, trunk_mask) = AVW_Data.calculate_trunk_model(Katja, 'complete', z_start=start_slice, z_end=end_slice)
outer_shape_trunk = AVW_Data.calculate_outer_shape(model_trunk)

Katja.add_voxel_data(short_name='trunk',
                     name="The trunk of the 'complete' model. Arms have been removed using "
                          "VoxelModel.remove_arms().",
                     outer_shape=outer_shape_trunk,
                     model=model_trunk,
                     mask=trunk_mask,
                     tissue_mapping=None)

surface = Katja.create_3d_model(model_type='trunk', patch_size=(30, 30))
Katja.models['trunk'].surface_3d = surface

Katja.models['trunk'].endpoints = []
for (i, s) in enumerate(surface):
    Katja.models['trunk'].endpoints.append(Coordinate(np.array(s['centroid'])))

Katja.save_model()


# Script to create Proximity Interaction Graphs.
# Author: Marc Moesser
# This script uses adapted functions from the ECIF script supplied in https://github.com/DIFACQUIM/ECIF

import pandas as pd
import numpy as np
import networkx as nx
from scipy.spatial.distance import cdist
from itertools import product

def GetAtomType(atom):
# Atom types are defined as follows:
# 1) Atom symbol
# 2) Explicit Valence
# 3) # Heavy Atom Neighbors
# 4) # Hydrogen Neighbors
# 5) Boolean: Is atom aromatic?
# 6) Boolean: is atom in a ring?

# This function can be used to identify all unique protein atom types in the dataset.

    AtomType = [atom.GetSymbol(),
                str(atom.GetExplicitValence()),
                str(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() != "H"])),
                str(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() == "H"])),
                str(int(atom.GetIsAromatic())),
                str(int(atom.IsInRing())),
               ]

    return(";".join(AtomType))

def LoadSDFasDF(mol):
# This function converts the input ligand (.MOL file) into a pandas DataFrame with the ligand atom position in 3D (X,Y,Z)

    m = mol

    atoms = []

    for atom in m.GetAtoms():
        if atom.GetSymbol() != "H": # Include only non-hydrogen atoms
            entry = [int(atom.GetIdx())]
            entry.append(str(atom.GetSymbol()))
            pos = m.GetConformer().GetAtomPosition(atom.GetIdx())
            entry.append(float("{0:.4f}".format(pos.x)))
            entry.append(float("{0:.4f}".format(pos.y)))
            entry.append(float("{0:.4f}".format(pos.z)))
            atoms.append(entry)

    df = pd.DataFrame(atoms)
    df.columns = ["ATOM_INDEX","ATOM_TYPE","X","Y","Z"]

    return(df)

def LoadPDBasDF(PDB, Atom_Keys):
# This function converts a protein PDB file into a pandas DataFrame with the protein atom position in 3D (X,Y,Z)

    prot_atoms = []

    f = open(PDB)
    for i in f:
        if i[:4] == "ATOM":
            # Include only non-hydrogen atoms
            if (len(i[12:16].replace(" ","")) < 4 and i[12:16].replace(" ","")[0] != "H") or (len(i[12:16].replace(" ","")) == 4 and i[12:16].replace(" ","")[1] != "H" and i[12:16].replace(" ","")[0] != "H"):
                prot_atoms.append([int(i[6:11]),
                         i[17:20]+"-"+i[12:16].replace(" ",""),
                         float(i[30:38]),
                         float(i[38:46]),
                         float(i[46:54])
                        ])

    f.close()

    df = pd.DataFrame(prot_atoms, columns=["ATOM_INDEX","PDB_ATOM","X","Y","Z"])
    df = df.merge(Atom_Keys, left_on='PDB_ATOM', right_on='PDB_ATOM')[["ATOM_INDEX", "ATOM_TYPE", "X", "Y", "Z"]].sort_values(by="ATOM_INDEX").reset_index(drop=True)
    if list(df["ATOM_TYPE"].isna()).count(True) > 0:
        print("WARNING: Protein contains unsupported atom types. Only supported atom-type pairs are counted.")
    return(df)

def GetAtomContacts(PDB_protein, mol, Atom_Keys, distance_cutoff=6.0):
# This function returns the list of protein atom types the ligand interacts with for a given distance cutoff
# cutoff = 6 Angstrom is standard

    # Protein and ligand structure are loaded as pandas DataFrame
    Target = LoadPDBasDF(PDB_protein, Atom_Keys)
    Ligand = LoadSDFasDF(mol)

    # A cubic box around the ligand is created using the proximity threshold specified (here distance_cutoff = 6 Angstrom by default).
    for i in ["X","Y","Z"]:
        Target = Target[Target[i] < float(Ligand[i].max())+distance_cutoff]
        Target = Target[Target[i] > float(Ligand[i].min())-distance_cutoff]

    # Calculate the possible pairs
    Pairs = list(product(Target["ATOM_TYPE"], Ligand["ATOM_INDEX"]))
    Pairs = [str(x[0])+"-"+str(x[1]) for x in Pairs]
    Pairs = pd.DataFrame(Pairs, columns=["ATOM_PAIR"])

    Distances = cdist(Target[["X","Y","Z"]], Ligand[["X","Y","Z"]], metric="euclidean")
    Distances = Distances.reshape(Distances.shape[0]*Distances.shape[1],1)
    Distances = pd.DataFrame(Distances, columns=["DISTANCE"])

    #Select pairs with distance lower than the cutoff
    Pairs = pd.concat([Pairs,Distances], axis=1)
    Pairs = Pairs[Pairs["DISTANCE"] <= distance_cutoff].reset_index(drop=True)

    contact_pair_list = [i.split("-")[0] for i in Pairs["ATOM_PAIR"]]
    Pairs["PROT_ATOM"] = contact_pair_list
    Pairs["LIG_ATOM"] = [int(i.split("-")[1]) for i in Pairs["ATOM_PAIR"]]

    return Pairs

def atom_features(atom, features=["num_heavy_atoms", "total_num_Hs", "explicit_valence", "is_aromatic", "is_in_ring"]):
    # Computes the ligand atom features for graph node construction
    # The standard features are the following:
    # num_heavy_atoms = # of heavy atom neighbors
    # total_num_Hs = # number of hydrogen atom neighbors
    # explicit_valence = explicit valence of the atom
    # is_aromatic = boolean 1 - aromatic, 0 - not aromatic
    # is_in_ring = boolean 1 - is in ring, 0 - is not in ring

    feature_list = []
    if "num_heavy_atoms" in features:
        feature_list.append(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() != "H"]))
    if "total_num_Hs" in features:
        feature_list.append(len([x.GetSymbol() for x in atom.GetNeighbors() if x.GetSymbol() == "H"]))
    if "explicit_valence" in features: #-NEW ADDITION FOR PLIG
        feature_list.append(atom.GetExplicitValence())
    if "is_aromatic" in features:

        if atom.GetIsAromatic():
            feature_list.append(1)
        else:
            feature_list.append(0)
    if "is_in_ring" in features:
        if atom.IsInRing():
            feature_list.append(1)
        else:
            feature_list.append(0)
    return np.array(feature_list)

def atom_features_PLIG(atom_idx, atom, contact_df, extra_features, Atom_Keys):
    # Generates the protein-ligand interaction features for the PLIG creation

    possible_contacts = list(dict.fromkeys(Atom_Keys["ATOM_TYPE"]))
    feature_list = np.zeros(len(possible_contacts), dtype=int)
    contact_df_slice = contact_df[contact_df["LIG_ATOM"] == atom_idx]

    #count the number of contacts between ligand and protein atoms
    for i,contact in enumerate(possible_contacts):
        for k in contact_df_slice["PROT_ATOM"]:
            if k == contact:
                feature_list[i] +=1

    extra_feature_array = atom_features(atom, extra_features)
    output = np.append(extra_feature_array, feature_list)

    return output

def mol_to_graph(mol, contact_df, Atom_Keys, extra_features=["num_heavy_atoms", "total_num_Hs", "explicit_valence","is_aromatic", "is_in_ring"]):
    #Final function to summarize the generation of PLIGS

    #Extra features are any extra features to be added to the protein-ligand interaction features.
    #In this work, we use the following ligand-based features as the ligand atom node features that are added to the interaction features:

    #num_heavy_atoms
    #total_num_hs
    #explicit_valence
    #is_aromatic
    #is_in_ring

    #However, this is freely customizable! Any additional features can be added here.

    c_size = len([x.GetSymbol() for x in mol.GetAtoms() if x.GetSymbol() != "H"])
    features = []
    heavy_atom_index = []
    idx_to_idx = {}
    counter = 0

    # Generate nodes
    for atom in mol.GetAtoms():
        if atom.GetSymbol() != "H": # Include only non-hydrogen atoms
            idx_to_idx[atom.GetIdx()] = counter
            counter +=1
            heavy_atom_index.append(atom.GetIdx())
            feature = atom_features_PLIG(atom.GetIdx(), atom, contact_df, extra_features, Atom_Keys)
            features.append(feature)

    #Generate edges
    edges = []
    for bond in mol.GetBonds():
        idx1 = bond.GetBeginAtomIdx()
        idx2 = bond.GetEndAtomIdx()
        if idx1 in heavy_atom_index and idx2 in heavy_atom_index:
            edges.append([idx_to_idx[bond.GetBeginAtomIdx()], idx_to_idx[bond.GetEndAtomIdx()]])
    g = nx.Graph(edges).to_directed()
    edge_index = []
    for e1, e2 in g.edges:
        edge_index.append([e1, e2])

    #return molecular graph with its node features and edge indices
    return c_size, features, edge_index


"""Library implementing attention modules.

Authors
 * Anonymous
"""

import torch
import logging
import torch.nn as nn
import numpy as np
from typing import Optional
from speechbrain.dataio.dataio import length_to_mask
import torch.nn.functional as F
import math

logger = logging.getLogger(__name__)

class ContentBasedAttention(nn.Module):
    """ This class implements content-based attention module for seq2seq
    learning.

    Reference: NEURAL MACHINE TRANSLATION BY JOINTLY LEARNING TO ALIGN
    AND TRANSLATE, Bahdanau et.al. https://arxiv.org/pdf/1409.0473.pdf

    Arguments
    ---------
    attn_dim : int
        Size of the attention feature.
    output_dim : int
        Size of the output context vector.
    scaling : float
        The factor controls the sharpening degree (default: 1.0).

    Example
    -------
    >>> enc_tensor = torch.rand([4, 10, 20])
    >>> enc_len = torch.ones([4]) * 10
    >>> dec_tensor = torch.rand([4, 25])
    >>> net = ContentBasedAttention(enc_dim=20, dec_dim=25, attn_dim=30, output_dim=5)
    >>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
    >>> out_tensor.shape
    torch.Size([4, 5])
    """

    def __init__(self, enc_dim, dec_dim, attn_dim, output_dim, scaling=1.0):
        super(ContentBasedAttention, self).__init__()

        self.mlp_enc = nn.Linear(enc_dim, attn_dim)
        self.mlp_dec = nn.Linear(dec_dim, attn_dim)
        self.mlp_attn = nn.Linear(attn_dim, 1, bias=False)
        self.mlp_out = nn.Linear(enc_dim, output_dim)

        self.scaling = scaling

        self.softmax = nn.Softmax(dim=-1)

        # reset the encoder states, lengths and masks
        self.reset()

    def reset(self):
        """Reset the memory in the attention module.
        """
        self.enc_len = None
        self.precomputed_enc_h = None
        self.mask = None

    def forward(self, enc_states, enc_len, dec_states):
        """Returns the output of the attention module.

        Arguments
        ---------
        enc_states : torch.Tensor
            The tensor to be attended.
        enc_len : torch.Tensor
            The real length (without padding) of enc_states for each sentence.
        dec_states : torch.Tensor
            The query tensor.

        """

        if self.precomputed_enc_h is None:

            self.precomputed_enc_h = self.mlp_enc(enc_states)
            self.mask = length_to_mask(
                enc_len, max_len=enc_states.size(1), device=enc_states.device
            )

        dec_h = self.mlp_dec(dec_states.unsqueeze(1))
        attn = self.mlp_attn(
            torch.tanh(self.precomputed_enc_h + dec_h)
        ).squeeze(-1)

        # mask the padded frames
        attn = attn.masked_fill(self.mask == 0, -np.inf)
        attn = self.softmax(attn * self.scaling)

        # compute context vectors
        # [B, 1, L] X [B, L, F]
        context = torch.bmm(attn.unsqueeze(1), enc_states).squeeze(1)
        context = self.mlp_out(context)

        return context, attn

class LocationAwareAttention(nn.Module):
    """This class implements location-aware attention module for seq2seq learning.

    Reference: Attention-Based Models for Speech Recognition, Chorowski et.al.
    https://arxiv.org/pdf/1506.07503.pdf

    Arguments
    ---------
    attn_dim : int
        Size of the attention feature.
    output_dim : int
        Size of the output context vector.
    conv_channels : int
        Number of channel for location feature.
    kernel_size : int
        Kernel size of convolutional layer for location feature.
    scaling : float
        The factor controls the sharpening degree (default: 1.0).

    Example
    -------
    >>> enc_tensor = torch.rand([4, 10, 20])
    >>> enc_len = torch.ones([4]) * 10
    >>> dec_tensor = torch.rand([4, 25])
    >>> net = LocationAwareAttention(
    ...     enc_dim=20,
    ...     dec_dim=25,
    ...     attn_dim=30,
    ...     output_dim=5,
    ...     conv_channels=10,
    ...     kernel_size=100)
    >>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
    >>> out_tensor.shape
    torch.Size([4, 5])
    """

    precomputed_enc_h: Optional[torch.Tensor]

    def __init__(
        self,
        enc_dim,
        dec_dim,
        attn_dim,
        output_dim,
        conv_channels,
        kernel_size,
        scaling=1.0,
    ):
        super(LocationAwareAttention, self).__init__()

        self.mlp_enc = nn.Linear(enc_dim, attn_dim)
        self.mlp_dec = nn.Linear(dec_dim, attn_dim)
        self.mlp_attn = nn.Linear(attn_dim, 1, bias=False)
        self.conv_loc = nn.Conv1d(
            1,
            conv_channels,
            kernel_size=2 * kernel_size + 1,
            padding=kernel_size,
            bias=False,
        )
        self.mlp_loc = nn.Linear(conv_channels, attn_dim)
        self.mlp_attn = nn.Linear(attn_dim, 1, bias=False)
        self.mlp_out = nn.Linear(enc_dim, output_dim)

        self.scaling = scaling

        self.softmax = nn.Softmax(dim=-1)

        # reset the encoder states, lengths and masks
        self.reset()

    def reset(self):
        """Reset the memory in attention module.
        """
        self.enc_len = None
        self.precomputed_enc_h = None
        self.mask = None
        self.prev_attn = None

    def forward(self, enc_states, enc_len, dec_states):
        """Returns the output of the attention module.

        Arguments
        ---------
        enc_states : torch.Tensor
            The tensor to be attended.
        enc_len : torch.Tensor
            The real length (without padding) of enc_states for each sentence.
        dec_states : torch.Tensor
            The query tensor.
        """
        if self.precomputed_enc_h is None:

            self.precomputed_enc_h = self.mlp_enc(enc_states)
            self.mask = length_to_mask(
                enc_len, max_len=enc_states.size(1), device=enc_states.device
            )

            # multiply mask by 1/Ln for each row
            self.prev_attn = self.mask * (1 / enc_len.float()).unsqueeze(1)

        # compute location-aware features
        # [B, 1, L] -> [B, C, L]
        attn_conv = self.conv_loc(self.prev_attn.unsqueeze(1))
        # [B, C, L] -> [B, L, C] -> [B, L, F]
        attn_conv = self.mlp_loc(attn_conv.transpose(1, 2))

        dec_h = self.mlp_dec(dec_states.unsqueeze(1))
        attn = self.mlp_attn(
            torch.tanh(self.precomputed_enc_h + dec_h + attn_conv)
        ).squeeze(-1)

        # mask the padded frames
        attn = attn.masked_fill(self.mask == 0, -np.inf)
        attn = self.softmax(attn * self.scaling)

        # set prev_attn to current attn for the next timestep
        self.prev_attn = attn.detach()

        # compute context vectors
        # [B, 1, L] X [B, L, F]
        context = torch.bmm(attn.unsqueeze(1), enc_states).squeeze(1)
        context = self.mlp_out(context)

        return context, attn

class KeyValueAttention(nn.Module):
    """ This class implements a single-headed key-value attention module for seq2seq
    learning.

    Reference: "Attention Is All You Need" by Vaswani et al., sec. 3.2.1

    Arguments
    ---------
    enc_dim : int
        Size of the encoder feature vectors from which keys and values are computed.
    dec_dim : int
        Size of the decoder feature vectors from which queries are computed.
    attn_dim : int
        Size of the attention feature.
    output_dim : int
        Size of the output context vector.

    Example
    -------
    >>> enc_tensor = torch.rand([4, 10, 20])
    >>> enc_len = torch.ones([4]) * 10
    >>> dec_tensor = torch.rand([4, 25])
    >>> net = KeyValueAttention(enc_dim=20, dec_dim=25, attn_dim=30, output_dim=5)
    >>> out_tensor, out_weight = net(enc_tensor, enc_len, dec_tensor)
    >>> out_tensor.shape
    torch.Size([4, 5])
    """

    def __init__(self, enc_dim, dec_dim, attn_dim, output_dim):
        super(KeyValueAttention, self).__init__()

        self.key_linear = nn.Linear(enc_dim, attn_dim)
        self.query_linear = nn.Linear(dec_dim, attn_dim)
        self.value_linear = nn.Linear(enc_dim, output_dim)
        self.scaling = torch.sqrt(torch.tensor(attn_dim).float())

        # reset the encoder states, lengths and masks
        self.reset()

    def reset(self):
        """Reset the memory in the attention module.
        """
        self.values = None
        self.keys = None
        self.mask = None

    def forward(self, enc_states, enc_len, dec_states):
        """Returns the output of the attention module.

        Arguments
        ---------
        enc_states : torch.Tensor
            The tensor to be attended.
        enc_len : torch.Tensor
            The real length (without padding) of enc_states for each sentence.
        dec_states : torch.Tensor
            The query tensor.
        """

        if self.keys is None:

            self.keys = self.key_linear(enc_states)
            self.values = self.value_linear(enc_states)
            self.mask = length_to_mask(
                enc_len, max_len=enc_states.size(1), device=enc_states.device
            ).unsqueeze(2)

        query = self.query_linear(dec_states).unsqueeze(2)
        scores = torch.matmul(self.keys, query) / self.scaling
        scores = scores.masked_fill(self.mask == 0, -np.inf)
        normalized_scores = scores.softmax(1).transpose(1, 2)
        out = torch.matmul(normalized_scores, self.values).squeeze(1)
        return out, normalized_scores

class RelPosEncXL(nn.Module):
    """

    """

    def __init__(self, emb_dim):
        super().__init__()
        self.emb_dim = emb_dim

        inv_freq = torch.exp(
            torch.arange(0, self.emb_dim, 2, dtype=torch.float32)
            * -(math.log(10000.0) / self.emb_dim)
        )
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x: torch.Tensor):
        """
        Parameters
        ----------
        x : torch.Tensor
        input tensor with shape seq_len, batch_size, embed_dim
        Returns
        -------
        pos_emb : torch.Tensor
        """
        seq_len = x.size(1)
        with torch.no_grad():
            tot_pe = torch.zeros((2, seq_len, self.emb_dim), dtype=x.dtype).to(
                x
            )
            pe_past = tot_pe[0]
            pe_future = tot_pe[1]
            positions = (
                torch.arange(0, seq_len, dtype=x.dtype).to(x).unsqueeze(-1)
            )
            sinusoids = torch.sin(positions * self.inv_freq)
            pe_past[:, 0::2] = sinusoids
            pe_past[:, 1::2] = torch.cos(positions * self.inv_freq)
            pe_future[:, 0::2] = sinusoids  # same for past and future
            pe_future[:, 1::2] = torch.cos(-positions * self.inv_freq)

            pe_past = torch.flip(pe_past, (0,)).unsqueeze(0)
            pe_future = pe_future[1:].unsqueeze(0)
            pe = torch.cat([pe_past, pe_future], dim=1)
            # pe is now 1, 2*seq_len, embed_dim
            return pe

class RelPosMHAXL(nn.Module):
    """ This class implements the relative multihead implementation similar to that in Transformer XL
    https://arxiv.org/pdf/1901.02860.pdf

    Arguments
    ---------
    embed_dim : int
        Size of the encoder feature vectors from which keys and values are computed.
    num_heads: int
        Number of attention heads.
    dropout : float, optional
        Dropout rate.
    vbias: bool, optional
        Whether to use bias for computing value.
    vdim: int, optional
        Size for value. Default is embed_dim (Note each head is embed_dim // num_heads).
    mask_pos_future: bool, optional
        Whether to mask future positional encodings values.
        Must be true for causal applications e.g. decoder.
    Example
    -------
    >>> inputs = torch.rand([6, 60, 512])
    >>> pos_emb = torch.rand([1, 2*60-1, 512])
    >>> net = RelPosMHAXL(num_heads=8, embed_dim=inputs.shape[-1])
    >>> outputs, attn = net(inputs, inputs, inputs, pos_emb)
    >>> outputs.shape
    torch.Size([6, 60, 512])
    """

    def __init__(
        self,
        embed_dim,
        num_heads,
        dropout=0.0,
        vbias=False,
        vdim=None,
        mask_pos_future=False,
    ):
        super(RelPosMHAXL, self).__init__()
        self.embed_dim = embed_dim
        self.vdim = vdim if vdim is not None else embed_dim
        self._qkv_same_embed_dim = self.vdim == embed_dim
        self.mask_pos_future = mask_pos_future
        self.vbias = vbias

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.vhead_dim = self.vdim // num_heads

        assert (
            self.head_dim * num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"
        assert (
            self.vhead_dim * num_heads == self.vdim
        ), "vdim must be divisible by num_heads"

        if self._qkv_same_embed_dim is False:
            self.qk_proj_weight = nn.Parameter(
                torch.empty(2 * embed_dim, embed_dim)
            )
            self.v_proj_weight = nn.Parameter(torch.empty(self.vdim, embed_dim))
        else:
            self.in_proj_weight = nn.Parameter(
                torch.empty(3 * embed_dim, embed_dim)
            )

        if vbias:
            self.value_bias_weight = nn.Parameter(torch.empty(self.vdim))
        else:
            self.vbias = None

        self.dropout_att = nn.Dropout(dropout)
        self.out_proj = nn.Linear(self.vdim, embed_dim)

        self.linear_pos = nn.Linear(embed_dim, embed_dim, bias=False)

        self.pos_bias_u = nn.Parameter(
            torch.empty(self.head_dim, self.num_heads)
        )
        self.pos_bias_v = nn.Parameter(
            torch.empty(self.head_dim, self.num_heads)
        )

        if next(self.parameters()).dtype == torch.float16:
            self.attn_fill_value = -65000
        else:
            self.attn_fill_value = -float("inf")

        self._reset_parameters()
        self.scale = 1 / math.sqrt(self.embed_dim)

    def _reset_parameters(self):
        if self._qkv_same_embed_dim:
            torch.nn.init.xavier_uniform_(self.in_proj_weight)
        else:
            torch.nn.init.xavier_uniform_(self.qk_proj_weight)
            torch.nn.init.xavier_uniform_(self.v_proj_weight)

        if self.vbias is not None:
            torch.nn.init.constant_(self.value_bias_weight, 0.0)

        # positional biases
        torch.nn.init.xavier_uniform_(self.pos_bias_u)
        torch.nn.init.xavier_uniform_(self.pos_bias_v)

    def rel_shift(self, x):
        # batch, head, time1, 2*time1-1.

        zero_pad = torch.zeros(
            (*x.size()[:3], 1), device=x.device, dtype=x.dtype
        )
        x_padded = torch.cat([zero_pad, x], dim=-1)

        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
        x = x_padded[:, :, 1:].view_as(x)[
            :, :, :, : x.size(-1) // 2 + 1
        ]  # only keep the positions from 0 to time2

        if self.mask_pos_future:
            ones = torch.ones((x.size(2), x.size(3)), device=x.device)
            x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]

        return x

    def forward(
        self,
        query,
        key,
        value,
        pos_embs,
        key_padding_mask=None,
        attn_mask=None,
        return_attn_weights=True,
    ):
        """
        Arguments
        ----------
        query : tensor
            (B, L, E) where L is the target sequence length,
            B is the batch size, E is the embedding dimension.
        key : tensor
            (B, S, E) where S is the source sequence length,
            B is the batch size, E is the embedding dimension.
        value : tensor
            (B, S, E) where S is the source sequence length,
            B is the batch size, E is the embedding dimension.
        pos_emb : tensor
            bidirectional sinusoidal positional embedding tensor (1, 2*S-1, E) where S is the max length between source and target sequence lengths,
            and E is the embedding dimension.
        key_padding_mask : tensor
            (B, S) where B is the batch size, S is the source sequence
            length. If a ByteTensor is provided, the non-zero positions will
            be ignored while the position with the zero positions will be
            unchanged. If a BoolTensor is provided, the positions with the
            value of True will be ignored while the position with the value
            of False will be unchanged.
        attn_mask : tensor
            2D mask (L, S) where L is the target sequence length, S is
            the source sequence length.
            3D mask (N*num_heads, L, S) where N is the batch
            size, L is the target sequence length, S is the source sequence
            length. attn_mask ensure that position i is allowed to attend the
            unmasked positions. If a ByteTensor is provided, the non-zero
            positions are not allowed to attend while the zero positions will
            be unchanged. If a BoolTensor is provided, positions with True is
            not allowed to attend while False values will be unchanged. If a
            FloatTensor is provided, it will be added to the attention weight.

        Outputs
        -------
        out : tensor
            (B, L, E) where L is the target sequence length, B is the
            batch size, E is the embedding dimension.
        attn_score : tensor
            (B, L, S) where B is the batch size, L is the target
            sequence length, S is the source sequence length.
        """

        # query, key and value are of shape batch, time, embed_dim
        bsz = query.shape[0]
        klen = key.shape[1]
        qlen = query.shape[1]

        if self._qkv_same_embed_dim:
            # self-attention
            if (query is key or torch.equal(query, key)) and (
                key is value or torch.equal(key, value)
            ):
                query, key, value = (
                    nn.functional.linear(query, self.in_proj_weight)
                    .view(bsz, -1, self.num_heads, self.head_dim * 3)
                    .chunk(3, dim=-1)
                )
            else:
                qweight, kweight, vweight = self.in_proj_weight.chunk(3, dim=0)
                query = nn.functional.linear(query, qweight).view(
                    bsz, -1, self.num_heads, self.head_dim
                )
                key = nn.functional.linear(key, kweight).view(
                    bsz, -1, self.num_heads, self.head_dim
                )
                value = nn.functional.linear(value, vweight).view(
                    bsz, -1, self.num_heads, self.head_dim
                )
        else:
            raise NotImplementedError
            query, key = (
                nn.functional.linear(query, self.qk_proj_weight)
                .view(bsz, -1, self.num_heads, self.head_dim * 2)
                .chunk(2, dim=-1)
            )
            value = nn.functional.linear(value, self.v_proj_weight).view(
                bsz, -1, self.num_heads, self.vhead_dim
            )

        if self.vbias is not None:
            value = value + self.value_bias_weight.view(
                1, 1, self.num_heads, self.vhead_dim
            )

        p_k = self.linear_pos(pos_embs).view(
            1, -1, self.num_heads, self.head_dim
        )
        # (batch, head, klen, d_k)

        q_with_bias_u = (
            query + self.pos_bias_u.view(1, 1, self.num_heads, self.head_dim)
        ).transpose(1, 2)
        # (batch, head, qlen, d_k)
        q_with_bias_v = (
            query + self.pos_bias_v.view(1, 1, self.num_heads, self.head_dim)
        ).transpose(1, 2)

        # (batch, head, qlen, klen)
        matrix_ac = torch.matmul(q_with_bias_u, key.permute(0, 2, 3, 1))
        # (batch, num_heads, klen, 2*klen-1)
        matrix_bd = torch.matmul(q_with_bias_v, p_k.permute(0, 2, 3, 1))
        matrix_bd = self.rel_shift(matrix_bd)  # shifting trick

        # if klen != qlen:
        #   import ipdb
        #  ipdb.set_trace(

        attn_score = (matrix_ac + matrix_bd) * self.scale

        # compute attention probability
        if attn_mask is not None:
            if attn_mask.ndim == 2:
                attn_mask = attn_mask.view(1, 1, qlen, klen)
            else:
                attn_mask = attn_mask.view(-1, self.num_heads, qlen, klen)

            if attn_mask.dtype == torch.bool:
                attn_score = attn_score.masked_fill(
                    attn_mask, self.attn_fill_value
                )
            else:
                attn_score += attn_mask

        if key_padding_mask is not None:
            attn_score = attn_score.masked_fill(
                key_padding_mask.view(bsz, 1, 1, klen), self.attn_fill_value,
            )

        attn_score = F.softmax(attn_score, dim=-1)
        attn_score = self.dropout_att(attn_score)
        x = torch.matmul(
            attn_score, value.transpose(1, 2)
        )  # (batch, head, time1, d_k)
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(bsz, -1, self.vhead_dim * self.num_heads)
        )  # (batch, time1, d_model)

        out = self.out_proj(x)
        if return_attn_weights:
            return out, attn_score
        return out

class MultiheadAttention(nn.Module):
    """ The class is a wrapper of MultiHead Attention for torch.nn.MultiHeadAttention.

    Reference: https://pytorch.org/docs/stable/nn.html

    Arguments
    ----------
    num_heads : int
        parallel attention heads.
    dropout : float
        a Dropout layer on attn_output_weights (default: 0.0).
    bias : bool
        add bias as module parameter (default: True).
    add_bias_kv : bool
        add bias to the key and value sequences at dim=0.
    add_zero_attn : bool
        add a new batch of zeros to the key and value sequences at dim=1.
    kdim : int
        total number of features in key (default: None).
    vdim : int
        total number of features in value (default: None).

    Example
    -------
    >>> inputs = torch.rand([8, 60, 512])
    >>> net = MultiheadAttention(nhead=8, d_model=inputs.shape[-1])
    >>> outputs, attn = net(inputs, inputs, inputs)
    >>> outputs.shape
    torch.Size([8, 60, 512])
    """

    def __init__(
        self,
        nhead,
        d_model,
        dropout=0.0,
        bias=True,
        add_bias_kv=False,
        add_zero_attn=False,
        kdim=None,
        vdim=None,
    ):
        super().__init__()

        self.att = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=nhead,
            dropout=dropout,
            bias=bias,
            add_bias_kv=add_bias_kv,
            add_zero_attn=add_zero_attn,
            kdim=kdim,
            vdim=vdim,
        )

    def forward(
        self,
        query,
        key,
        value,
        attn_mask: Optional[torch.Tensor] = None,
        key_padding_mask: Optional[torch.Tensor] = None,
        return_attn_weights: Optional[torch.Tensor] = True,
        pos_embs: Optional[torch.Tensor] = None,
    ):
        """
        Arguments
        ----------
        query : torch.Tensor
            (B, L, E) where L is the target sequence length,
            B is the batch size, E is the embedding dimension.
        key : torch.Tensor
            (B, S, E) where S is the source sequence length,
            B is the batch size, E is the embedding dimension.
        value : torch.Tensor
            (B, S, E) where S is the source sequence length,
            B is the batch size, E is the embedding dimension.
        key_padding_mask : torch.Tensor, optional
            (B, S) where B is the batch size, S is the source sequence
            length. If a ByteTensor is provided, the non-zero positions will
            be ignored while the position with the zero positions will be
            unchanged. If a BoolTensor is provided, the positions with the
            value of True will be ignored while the position with the value
            of False will be unchanged.
        attn_mask : torch.Tensor, optional
            2D mask (L, S) where L is the target sequence length, S is
            the source sequence length.
            3D mask (N*num_heads, L, S) where N is the batch
            size, L is the target sequence length, S is the source sequence
            length. attn_mask ensure that position i is allowed to attend the
            unmasked positions. If a ByteTensor is provided, the non-zero
            positions are not allowed to attend while the zero positions will
            be unchanged. If a BoolTensor is provided, positions with True is
            not allowed to attend while False values will be unchanged. If a
            FloatTensor is provided, it will be added to the attention weight.
        pos_embs: torch.Tensor, optional
            Positional embeddings added to the attention map of shape (L, S, E) or (L, S, 1).

        Outputs
        -------
        attn_output : torch.Tensor
            (B, L, E) where L is the target sequence length, B is the
            batch size, E is the embedding dimension.
        attn_output_weights : torch.Tensor
            (B, L, S) where B is the batch size, L is the target
            sequence length, S is the source sequence length.
        """
        # give tensors of shape (time, batch, fea)
        query = query.permute(1, 0, 2)
        key = key.permute(1, 0, 2)
        value = value.permute(1, 0, 2)

        # this will be legit because of https://github.com/pytorch/pytorch/blob/5288d05cfdda85c46c4df84617fa7f37c21b10b3/torch/nn/functional.py#L4946
        # we can inject relative learnable pos embeddings directly in MHA via the attn_mask
        if pos_embs is not None:
            if attn_mask is not None:
                attn_mask += pos_embs
            else:
                attn_mask = pos_embs

        output = self.att(
            query,
            key,
            value,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=return_attn_weights,
        )

        if return_attn_weights:
            output, attention_weights = output
            # reshape the output back to (batch, time, fea)
            output = output.permute(1, 0, 2)
            return output, attention_weights
        else:
            output = output.permute(1, 0, 2)
            return output

class PositionalwiseFeedForward(nn.Module):
    """The class implements the positional-wise feed forward module in
    “Attention Is All You Need”.

    Arguments
    ----------
    d_ffn: int
        Hidden layer size.
    input_shape : tuple, optional
        Expected shape of the input. Alternatively use ``input_size``.
    input_size : int, optional
        Expected size of the input. Alternatively use ``input_shape``.
    dropout: float, optional
        Dropout rate.
    activation: torch.nn.Module, optional
        activation functions to be applied (Recommendation: ReLU, GELU).

    Example
    -------
    >>> inputs = torch.rand([8, 60, 512])
    >>> net = PositionalwiseFeedForward(256, input_size=inputs.shape[-1])
    >>> outputs = net(inputs)
    >>> outputs.shape
    torch.Size([8, 60, 512])
    """

    def __init__(
        self,
        d_ffn,
        input_shape=None,
        input_size=None,
        dropout=0.0,
        activation=nn.ReLU,
    ):
        super().__init__()

        if input_shape is None and input_size is None:
            raise ValueError("Expected one of input_shape or input_size")

        if input_size is None:
            input_size = input_shape[-1]

        self.ffn = nn.Sequential(
            nn.Linear(input_size, d_ffn),
            activation(),
            nn.Dropout(dropout),
            nn.Linear(d_ffn, input_size),
        )

    def forward(self, x):
        # give a tensor of shap (time, batch, fea)
        x = x.permute(1, 0, 2)
        x = self.ffn(x)

        # reshape the output back to (batch, time, fea)
        x = x.permute(1, 0, 2)

        return x


"""Learning rate scheduler."""

import numpy as np
import torch
#from dotmap import DotMap
#from fairseq.optim.fairseq_optimizer import FairseqOptimizer
#from fairseq.optim.lr_scheduler import cosine_lr_scheduler
from torch import nn
from torch.optim.lr_scheduler import LambdaLR
from functools import wraps
import warnings
import math
from torch.optim.optimizer import Optimizer
import weakref

EPOCH_DEPRECATION_WARNING = (
    "The epoch parameter in `scheduler.step()` was not necessary and is being "
    "deprecated where possible. Please use `scheduler.step()` to step the "
    "scheduler. During the deprecation, if epoch is different from None, the "
    "closed form is used instead of the new chainable form, where available. "
    "Please open an issue if you are unable to replicate your use case: "
    "https://github.com/pytorch/pytorch/issues/new/choose."
)
# from transformers import (
#     get_constant_schedule,
#     get_linear_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup)
class _LRScheduler(object):

    def __init__(self, optimizer, last_epoch=-1, verbose=False):

        # Attach optimizer
        if not isinstance(optimizer, Optimizer):
            raise TypeError('{} is not an Optimizer'.format(
                type(optimizer).__name__))
        self.optimizer = optimizer

        # Initialize epoch and base learning rates
        if last_epoch == -1:
            for group in optimizer.param_groups:
                group.setdefault('initial_lr', group['lr'])
        else:
            for i, group in enumerate(optimizer.param_groups):
                if 'initial_lr' not in group:
                    raise KeyError("param 'initial_lr' is not specified "
                                   "in param_groups[{}] when resuming an optimizer".format(i))
        self.base_lrs = [group['initial_lr'] for group in optimizer.param_groups]
        self.last_epoch = last_epoch

        # Following https://github.com/pytorch/pytorch/issues/20124
        # We would like to ensure that `lr_scheduler.step()` is called after
        # `optimizer.step()`
        def with_counter(method):
            if getattr(method, '_with_counter', False):
                # `optimizer.step()` has already been replaced, return.
                return method

            # Keep a weak reference to the optimizer instance to prevent
            # cyclic references.
            instance_ref = weakref.ref(method.__self__)
            # Get the unbound method for the same purpose.
            func = method.__func__
            cls = instance_ref().__class__
            del method

            @wraps(func)
            def wrapper(*args, **kwargs):
                instance = instance_ref()
                instance._step_count += 1
                wrapped = func.__get__(instance, cls)
                return wrapped(*args, **kwargs)

            # Note that the returned function here is no longer a bound method,
            # so attributes like `__func__` and `__self__` no longer exist.
            wrapper._with_counter = True
            return wrapper

        self.optimizer.step = with_counter(self.optimizer.step)
        self.optimizer._step_count = 0
        self._step_count = 0
        self.verbose = verbose

        self.step()

    def state_dict(self):
        """Returns the state of the scheduler as a :class:`dict`.

        It contains an entry for every variable in self.__dict__ which
        is not the optimizer.
        """
        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}

    def load_state_dict(self, state_dict):
        """Loads the schedulers state.

        Args:
            state_dict (dict): scheduler state. Should be an object returned
                from a call to :meth:`state_dict`.
        """
        self.__dict__.update(state_dict)

    def get_last_lr(self):
        """ Return last computed learning rate by current scheduler.
        """
        return self._last_lr

    def get_lr(self):
        # Compute learning rate using chainable form of the scheduler
        raise NotImplementedError

    def print_lr(self, is_verbose, group, lr, epoch=None):
        """Display the current learning rate.
        """
        if is_verbose:
            if epoch is None:
                print('Adjusting learning rate'
                      ' of group {} to {:.4e}.'.format(group, lr))
            else:
                print('Epoch {:5d}: adjusting learning rate'
                      ' of group {} to {:.4e}.'.format(epoch, group, lr))

    def step(self, epoch=None):
        # Raise a warning if old pattern is detected
        # https://github.com/pytorch/pytorch/issues/20124
        if self._step_count == 1:
            if not hasattr(self.optimizer.step, "_with_counter"):
                warnings.warn("Seems like `optimizer.step()` has been overridden after learning rate scheduler "
                              "initialization. Please, make sure to call `optimizer.step()` before "
                              "`lr_scheduler.step()`. See more details at "
                              "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning)

            # Just check if there were two first lr_scheduler.step() calls before optimizer.step()
            elif self.optimizer._step_count < 1:
                warnings.warn("Detected call of `lr_scheduler.step()` before `optimizer.step()`. "
                              "In PyTorch 1.1.0 and later, you should call them in the opposite order: "
                              "`optimizer.step()` before `lr_scheduler.step()`.  Failure to do this "
                              "will result in PyTorch skipping the first value of the learning rate schedule. "
                              "See more details at "
                              "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning)
        self._step_count += 1

        class _enable_get_lr_call:

            def __init__(self, o):
                self.o = o

            def __enter__(self):
                self.o._get_lr_called_within_step = True
                return self

            def __exit__(self, type, value, traceback):
                self.o._get_lr_called_within_step = False

        with _enable_get_lr_call(self):
            if epoch is None:
                self.last_epoch += 1
                values = self.get_lr()
            else:
                warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
                self.last_epoch = epoch
                if hasattr(self, "_get_closed_form_lr"):
                    values = self._get_closed_form_lr()
                else:
                    values = self.get_lr()

        for i, data in enumerate(zip(self.optimizer.param_groups, values)):
            param_group, lr = data
            param_group['lr'] = lr
            self.print_lr(self.verbose, i, lr, epoch)

        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]

class CosineAnnealingLR(_LRScheduler):
    r"""Set the learning rate of each parameter group using a cosine annealing
    schedule, where :math:`\eta_{max}` is set to the initial lr and
    :math:`T_{cur}` is the number of epochs since the last restart in SGDR:

    .. math::
        \begin{aligned}
            \eta_t & = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1
            + \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right),
            & T_{cur} \neq (2k+1)T_{max}; \\
            \eta_{t+1} & = \eta_{t} + \frac{1}{2}(\eta_{max} - \eta_{min})
            \left(1 - \cos\left(\frac{1}{T_{max}}\pi\right)\right),
            & T_{cur} = (2k+1)T_{max}.
        \end{aligned}

    When last_epoch=-1, sets initial lr as lr. Notice that because the schedule
    is defined recursively, the learning rate can be simultaneously modified
    outside this scheduler by other operators. If the learning rate is set
    solely by this scheduler, the learning rate at each step becomes:

    .. math::
        \eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
        \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right)

    It has been proposed in
    `SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this only
    implements the cosine annealing part of SGDR, and not the restarts.

    Args:
        optimizer (Optimizer): Wrapped optimizer.
        T_max (int): Maximum number of iterations.
        eta_min (float): Minimum learning rate. Default: 0.
        last_epoch (int): The index of last epoch. Default: -1.
        verbose (bool): If ``True``, prints a message to stdout for
            each update. Default: ``False``.

    .. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
        https://arxiv.org/abs/1608.03983
    """

    def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1, verbose=False):
        self.T_max = T_max
        self.eta_min = eta_min
        super(CosineAnnealingLR, self).__init__(optimizer, last_epoch, verbose)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", UserWarning)

        if self.last_epoch == 0:
            return self.base_lrs
        elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0:
            return [group['lr'] + (base_lr - self.eta_min) *
                    (1 - math.cos(math.pi / self.T_max)) / 2
                    for base_lr, group in
                    zip(self.base_lrs, self.optimizer.param_groups)]
        return [(1 + math.cos(math.pi * self.last_epoch / self.T_max)) /
                (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max)) *
                (group['lr'] - self.eta_min) + self.eta_min
                for group in self.optimizer.param_groups]

    def _get_closed_form_lr(self):
        return [self.eta_min + (base_lr - self.eta_min) *
                (1 + math.cos(math.pi * self.last_epoch / self.T_max)) / 2
                for base_lr in self.base_lrs]

def clip_gradient(model, clip: float):
    nn.utils.clip_grad_norm_(model.parameters(), clip)

class ConcatLR(torch.optim.lr_scheduler._LRScheduler):
    """
    From Over9000
    https://github.com/mgrankin/over9000/blob/master/train.py
    """
    def __init__(self, optimizer, scheduler1, scheduler2, total_steps,
                 pct_start=0.5, last_epoch=-1):
        self.scheduler1 = scheduler1
        self.scheduler2 = scheduler2
        self.step_start = float(pct_start * total_steps) - 1
        self.curr_epoch = 0
        super(ConcatLR, self).__init__(optimizer, last_epoch)

    def step(self):
        if self.curr_epoch <= self.step_start:
            self.scheduler1.step()
        else:
            self.scheduler2.step()
        self.curr_epoch += 1
        super().step()

    def get_lr(self):
        if self.curr_epoch <= self.step_start:
            return self.scheduler1.get_last_lr()
        else:
            return self.scheduler2.get_last_lr()

class TradeoffAnnealer:
    def __init__(self, c, num_steps=None):
        """
        Anneal the tradeoff between label and augmentation loss according
            to some schedule.

        :param c: config
        :param num_steps: int, provide when loading from checkpoint to fast-
            forward to that tradeoff value.
        """
        self.c = c
        self.name = self.c.exp_tradeoff_annealing

        self.num_steps = 0
        self.init_tradeoff = self.c.exp_tradeoff
        self.curr_tradeoff = self.c.exp_tradeoff
        self.max_steps = self.get_max_steps()
        self.step_map = {
            'constant': self.constant_step,
            'cosine': self.cosine_step,
            'linear_decline': self.linear_decline_step}

        if self.name not in self.step_map.keys():
            raise NotImplementedError

        self.step = self.step_map[self.name]

        if num_steps > 0:
            # If we are loading a model from checkpoint,
            # should update the annealer to that number of steps.
            for _ in range(num_steps):
                self.step()

            print(f'Fast-forwarded tradeoff annealer to step {num_steps}.')

        print(
            f'Initialized "{self.name}" augmentation/label tradeoff annealer. '
            f'Annealing to minimum value in {self.max_steps} steps.')

    def get_max_steps(self):
        # If annealing proportion is set to -1,
        if self.c.exp_tradeoff_annealing_proportion == -1:
            # and the optimizer proportion is set, we use the optimizer
            # proportion to determine how long it takes for the tradeoff to
            # anneal to 0.
            if self.c.exp_optimizer_warmup_proportion != -1:
                return int(np.ceil(self.c.exp_optimizer_warmup_proportion
                                   * self.c.exp_num_total_steps))
            # and the optimizer proportion is not set,
            # we take all steps to anneal.
            else:
                return self.c.exp_num_total_steps

        if (self.c.exp_tradeoff_annealing_proportion < 0
                or self.c.exp_tradeoff_annealing_proportion > 1):
            raise Exception('Invalid tradeoff annealing proportion.')

        # Otherwise, we use the tradeoff annealing proportion to determine
        # for how long we anneal.
        return int(np.ceil(self.c.exp_tradeoff_annealing_proportion
                           * self.c.exp_num_total_steps))

    def constant_step(self):
        self.num_steps += 1
        return self.curr_tradeoff

    def linear_decline_step(self):
        curr = self.num_steps
        max_val = self.init_tradeoff

        if self.num_steps <= self.max_steps:
            self.curr_tradeoff = max_val - (curr / self.max_steps) * max_val
        else:
            self.curr_tradeoff = 0

        self.num_steps += 1

        return self.curr_tradeoff

    def cosine_step(self):
        if self.num_steps <= self.max_steps:
            self.curr_tradeoff = self.init_tradeoff * (1 / 2) * (
                np.cos(np.pi * (self.num_steps / self.max_steps)) + 1)
        else:
            self.curr_tradeoff = 0

        self.num_steps += 1

        return self.curr_tradeoff

class LRScheduler:
    def __init__(self, c, name, optimizer):
        self.c = c
        self.name = name
        self.optimizer = optimizer
        self.num_steps = 0

        self.construct_auto_scheduler()

        print(f'Initialized "{name}" learning rate scheduler.')

    def construct_auto_scheduler(self):
        total_steps = self.c.exp_num_total_steps

        if self.c.exp_optimizer_warmup_proportion >= 0:
            num_warmup_steps = (
                    total_steps * self.c.exp_optimizer_warmup_proportion)
        else:
            num_warmup_steps = self.c.exp_optimizer_warmup_fixed_n_steps

        print(f'Warming up for {num_warmup_steps}/{total_steps} steps.')

        if self.name == 'constant':
            self.scheduler = get_constant_schedule(optimizer=self.optimizer)
        elif self.name == 'linear_warmup':
            self.scheduler = get_linear_schedule_with_warmup(
                optimizer=self.optimizer,
                num_warmup_steps=num_warmup_steps,
                num_training_steps=total_steps)
        # elif self.name == 'cosine_cyclic':
        #     args = dict(
        #         warmup_updates=num_warmup_steps,
        #         warmup_init_lr=1e-7,
        #         max_lr=self.c.exp_lr,
        #         lr=[1e-7],
        #         t_mult=2.,
        #         lr_period_updates=num_warmup_steps * 2,
        #         lr_shrink=0.5)
        #     optim = FairseqOptimizer(None)
        #     optim._optimizer = optim.optimizer = self.optimizer
        #     self.scheduler = cosine_lr_scheduler.CosineSchedule(
        #         optimizer=optim, args=DotMap(args))
        elif self.name == 'polynomial_decay_warmup':
            # Based on the fairseq implementation, which is based on BERT
            self.scheduler = get_polynomial_decay_schedule_with_warmup(
                optimizer=self.optimizer,
                num_warmup_steps=num_warmup_steps,
                num_training_steps=total_steps,
                lr_end=1e-7,
                power=1.0)
        elif self.name == 'flat_and_anneal':
            def d(x):
                return 1

            assert self.c.exp_optimizer_warmup_proportion >= 0

            # We use exp_optimizer_warmup_proportion to denote the
            # flat LR regime, prior to annealing
            dummy = LambdaLR(self.optimizer, d)
            cosine = CosineAnnealingLR(
                self.optimizer, int(total_steps * (
                    1 - self.c.exp_optimizer_warmup_proportion)))
            self.scheduler = ConcatLR(
                self.optimizer, dummy, cosine, total_steps,
                self.c.exp_optimizer_warmup_proportion)
        else:
            raise NotImplementedError

    def step(self):
        self.num_steps += 1
        c_lr = self.c.exp_lr
        num = self.num_steps
        tot = self.c.exp_num_total_steps

        if self.name == 'cosine_cyclic':
            self.scheduler.step_update(num_updates=num)
        else:
            self.scheduler.step()