import logging
import multiprocessing

import gunicorn.app.base

from app import configure, create_app, start, load_static_assets, \
    load_static_lib
from app.config import parser, settings

gunicorn.SERVER_SOFTWARE = 'gunicorn'  # hide gunicorn version

LOCAL_ENVS = ["LOCAL", "DEV"]

logger = logging.getLogger(__name__)

class Application(gunicorn.app.base.BaseApplication):
    def __init__(self, app, options=None):
        self.options = options or {}
        self.application = app
        super(Application, self).__init__()

    def load_config(self):
        config = dict([(key, value) for key, value in self.options.items()
                       if key in self.cfg.settings and value is not None])
        for key, value in config.items():
            self.cfg.set(key.lower(), value)

    def load(self):
        return self.application

def init_app():
    args = vars(parser.parse_args())
    # Get the configuration values from runtime/env/parameters
    # Essentially puts them in the settings variable
    configure(**args)
    return create_app()

def _post_fork(server=None, w=None):
    _config_logging()
    start()

def _config_logging():
    for logger in 'gunicorn.access', 'gunicorn.error':
        logging.getLogger(logger).propagate = True
        logging.getLogger(logger).handlers = []

if __name__ == '__main__':
    app = init_app()

    env_name = settings.get("ENV_NAME")

    # Make templates auto reload automatically when
    # Running locally or dev mode
    if env_name in LOCAL_ENVS:
        app.config['TEMPLATES_AUTO_RELOAD'] = True

    default_workers = (multiprocessing.cpu_count() * 2) + 1
    opts = {
        "accesslog": settings.get("ACCESS_LOG"),
        "access_log_format": settings.get("ACCESS_LOG_FORMAT"),
        "bind": settings.get("BIND"),
        "errorlog": settings.get("ERROR_LOG"),
        "keepalive": settings.get("KEEP_ALIVE"),
        "post_fork": _post_fork,
        "proc_name": settings.get("APP_NAME"),
        "max_requests": settings.get("MAX_REQUESTS"),
        "max_requests_jitter": settings.get("MAX_REQUESTS_JITTER"),
        "worker_class": settings.get("WORKER_CLASS"),
        "workers": settings.get("WORKERS") or (1 if env_name in LOCAL_ENVS
                                               else default_workers)
    }

    # This allows running the web application in local dev mode
    # This does not require login/auth from flask
    # This does not require a web server like nginx
    if env_name in LOCAL_ENVS:
        app.add_url_rule("/assets/<path:path>",
                        methods=['GET'], view_func=load_static_assets)
        app.add_url_rule("/lib/<path:path>",
                        methods=['GET'], view_func=load_static_lib)

        opts["reload"] = settings.get("RELOAD") or True
        opts['timeout'] = 86400
        logger.debug(opts)

    Application(app, opts).run()


from linotp.lib.crypto.utils import compare

"""
Declare the SecretObject to encapsulate security aspects
"""

import base64
import binascii
import ctypes
import hmac
import json
import logging
import os
import stat
import struct
import sys
from crypt import crypt as libcrypt

from Cryptodome.Cipher import AES
from pysodium import crypto_scalarmult_curve25519 as calc_dh

import linotp
from linotp.lib.crypto import utils

log = logging.getLogger(__name__)

class SecretObj(object):
    """
    High level interface to security operations

    This provides high level security operations without
    needing access to the secure data

    This is to be used by token implementations and
    classes that need encrypted data such as the
    database fields

    The encryption operations themselves are realised
    using a SecurityModule (such as HSM, PKCS11)

    The class implementation ensures that secret keys
    are not left around in memory after an operation
    has been carried out.

    It is possible to use this in two modes: With HSM,
    where operations are passed to the HSM, and without
    where a potentially degraded implementation is used.
    This is to provide the functionality during startup,
    before the HSM is ready.
    """

    def __init__(self, val, iv, preserve=True, hsm=None):
        self.val = val
        self.iv = iv
        self.bkey = None
        self.preserve = preserve
        self.hsm = hsm

    def getKey(self):
        log.debug("Warning: Requesting secret key as plaintext.")
        return utils.decrypt(self.val, self.iv, hsm=self.hsm)

    def calc_dh(self, partition, data):
        """
        encapsulate the Diffi Helmann calculation

        as the server secret key is a sensitive data, we try to encapsulate
        it and care for the cleanup

        :param partition: the id of the server secret key
        :param :
        """
        server_secret_key = utils.get_dh_secret_key(partition)
        hmac_secret = calc_dh(server_secret_key, data)

        utils.zerome(server_secret_key)

        return hmac_secret

    def compare(self, key):
        bhOtpKey = binascii.unhexlify(key)
        enc_otp_key = utils.encrypt(bhOtpKey, self.iv, hsm=self.hsm)
        otpKeyEnc = binascii.hexlify(enc_otp_key)

        return otpKeyEnc == self.val

    def compare_password(self, password):
        """
        compare the password of the password token

        the password token contains the unix hashed (hmac256) password format
        and is using the standard libcryp password hash compare. the iv is used
        as indicator for the new password format, which is :1:

        - legacy -
        the seed for some tokens contains the encrypted password
        insetead of decrypting the password and running the comparison,
        the new otp will be encrypted as well.

        :param password: the password - for the password token this is the
                         to be compared password

        :return: boolean
        """

        if self.iv == b":1:":
            return utils.compare_password(password, self.val.decode("utf-8"))

        # the legacy comparison: compare the ecrypted password

        enc_otp_key = utils.encrypt(password, self.iv, hsm=self.hsm)

        return compare(
            binascii.hexlify(enc_otp_key), binascii.hexlify(self.val)
        )

    def hmac_digest(self, data_input, hash_algo=None, bkey=None):
        b_key = bkey

        if not bkey:
            b_key = self._setupKey_()

        data = data_input

        if not hash_algo:
            hash_algo = utils.get_hashalgo_from_description("sha1")

        h_digest = utils.hmac_digest(
            bkey=b_key, data_input=data, hsm=self.hsm, hash_algo=hash_algo
        )

        if not bkey:
            self._clearKey_(preserve=self.preserve)

        return h_digest

    def aes_decrypt(self, data_input):
        """
        support inplace aes decryption for the yubikey

        :param data_input: data, that should be decrypted
        :return: the decrypted data
        """
        self._setupKey_()
        aes = AES.new(self.bkey, AES.MODE_ECB)
        msg_bin = aes.decrypt(data_input)
        self._clearKey_(preserve=self.preserve)
        return msg_bin

    @staticmethod
    def encrypt(seed: str, iv=None, hsm=None):
        if not iv:
            iv = utils.geturandom(16)
        enc_seed = utils.encrypt(seed, iv, hsm=hsm)
        return iv, enc_seed

    @staticmethod
    def decrypt(enc_seed, iv=None, hsm=None):
        dec_seed = utils.decrypt(enc_seed, iv=iv, hsm=hsm)
        return dec_seed

    @staticmethod
    def hash_pin(pin):
        """
        hash a given pin

        :param pin:
        :return: a concatenated 'iv:hashed_pin'
        """

        iv = utils.geturandom(16)
        hashed_pin = utils.hash_digest(pin.encode("utf-8"), iv)
        return iv, hashed_pin

    @staticmethod
    def check_hashed_pin(pin: str, hashed_pin: bytes, iv: bytes) -> bool:
        """
        check a hashed against a given pin

        :param hashed_pin: hex binary
        :param iv: hex binary iv from former decryption step
        :param pin: string
        :return: boolean
        """

        hash_pin = utils.hash_digest(pin.encode("utf-8"), iv)

        # TODO: position independend compare
        if hashed_pin == hash_pin:
            return True

        return False

    @staticmethod
    def encrypt_pin(pin: str):
        """
        encrypt a given pin

        :param pin:
        :return: a concatenated 'iv:crypt'
        """

        iv = utils.geturandom(16)
        enc_pin = utils.encryptPin(pin.encode("utf-8"), iv=iv)

        return enc_pin

    @staticmethod
    def check_encrypted_pin(pin: str, encrypted_pin: bytes, iv: bytes) -> bool:
        """
        check an encrypted against a given pin

        :param encrypted_pin: hex binary
        :param iv: hex binary iv from former decryption step
        :param pin: string
        :return: boolean
        """

        crypted_pin = utils.encryptPin(pin.encode("utf-8"), iv)

        # TODO: position independend compare
        if encrypted_pin == crypted_pin.encode("utf-8"):
            return True

        return False

    @staticmethod
    def decrypt_pin(pin, hsm=None):
        dec_pin = utils.decryptPin(pin, hsm=hsm)
        return dec_pin

    def encryptPin(self):
        self._setupKey_()
        res = utils.encryptPin(self.bkey)
        self._clearKey_(preserve=self.preserve)
        return res

    def _setupKey_(self):
        if not hasattr(self, "bkey"):
            self.bkey = None

        if self.bkey is None:
            self.bkey = binascii.unhexlify(
                self.decrypt(self.val, self.iv, hsm=self.hsm)
            )

        return self.bkey

    def _clearKey_(self, preserve=False):
        if preserve is False:
            if not hasattr(self, "bkey"):
                self.bkey = None

            if self.bkey is not None:
                utils.zerome(self.bkey)
                del self.bkey

    def __del__(self):
        self._clearKey_()

    def __enter__(self):
        self._clearKey_()

    def __exit__(self, type, value, traceback):
        self._clearKey_()


from pyramid.security import Allow
from pyramid.security import Deny
from pyramid.security import Everyone
from pyramid.traversal import find_interface
from pyramid.traversal import resource_path

from karl.models.interfaces import IIntranets
from karl.models.interfaces import IUserAddedGroup
from karl.models.interfaces import IUserRemovedGroup

from karl.security.policy import ADMINISTRATOR_PERMS
from karl.security.policy import GUEST_PERMS
from karl.security.policy import MODERATOR_PERMS
from karl.security.policy import MEMBER_PERMS
from karl.security.policy import CREATE
from karl.security.policy import NO_INHERIT
from karl.security.workflow import postorder
from karl.security.workflow import acl_diff
from karl.security.workflow import reset_security_workflow

from karl.utils import find_catalog
from karl.utils import find_community
from karl.utils import find_peopledirectory_catalog
from karl.utils import find_users
from karl.utils import find_profiles
from karl.views.communities import get_community_groups

#------------------------------------------------------------------------------
#   Helper methods
#------------------------------------------------------------------------------

_marker = object()

def _reindex(ob, texts=False):
    catalog = find_catalog(ob)
    if catalog is None:
        return # Will be true for a mailin test trace

    # XXX reindexing the 'path' index can be removed once we've
    # removed the last ACLChecker spelled in catalog queries from the
    # code; this is the "old" way of doing security filtering.
    path_index = catalog['path']
    path_index.reindex_doc(ob.docid, ob)

    # In some cases changing the workflow state of an object can change its
    # ranking in text search.
    if texts:
        text_index = catalog['texts']
        text_index.reindex_doc(ob.docid, ob)

    # if the object is folderish, we need to reindex it plus all its
    # subobjects' 'allowed' index entries recursively; each object's
    # allowed value depends on its parents in the lineage
    allowed_index = catalog.get('allowed')
    if allowed_index is not None:
        for node in postorder(ob):
            if hasattr(node, 'docid'):
                allowed_index.reindex_doc(node.docid, node)

def _reindex_peopledir(profile):
    catalog = find_peopledirectory_catalog(profile)

    # The docid stored as an attribute on the profile is the docid for the main
    # site catalog.  The peopledir catalog has a different set of docids, so it
    # is necessary to consult the document map.
    path = resource_path(profile)
    docid = catalog.document_map.docid_for_address(path)
    catalog.reindex_doc(docid, profile)

def ts(*args):
    return '\t'.join([unicode(x) for x in args])

#------------------------------------------------------------------------------
#   Electors
#------------------------------------------------------------------------------

def not_intranets_containment(context):
    return not intranets_containment(context)

def intranets_containment(context):
    if find_interface(context, IIntranets) is not None:
        return True
    return False

def private_community_containment(context):
    if not_intranets_containment(context):
        community = find_community(context)
        if community is None:
            return False
        return getattr(community, 'security_state', None) == 'private'
    return False

def public_community_containment(context):
    if not_intranets_containment(context):
        community = find_community(context)
        if community is None:
            return False
        return getattr(community, 'security_state', None) == 'public'
    return False

#------------------------------------------------------------------------------
#   Workflow for communities
#------------------------------------------------------------------------------
def community_to_private(ob, info):
    community = find_community(ob)
    acl = []
    moderators_group_name = community.moderators_group_name
    members_group_name = community.members_group_name
    acl.append((Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS))
    acl.append((Allow, moderators_group_name, MODERATOR_PERMS))
    acl.append((Allow, members_group_name, MEMBER_PERMS))
    acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(community, acl)
    if added or removed:
        community.__acl__ = acl
        msg = ts('community-private', resource_path(community), added, removed)
    _reindex(community)
    return msg

def community_to_public(ob, info):
    community = find_community(ob)
    acl = []
    moderators_group_name = community.moderators_group_name
    members_group_name = community.members_group_name
    acl.append((Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS))
    acl.append((Allow, moderators_group_name, MODERATOR_PERMS))
    acl.append((Allow, members_group_name, MEMBER_PERMS))
    acl.append((Allow, 'group.KarlStaff', GUEST_PERMS))
    acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(community, acl)
    if added or removed:
        community.__acl__ = acl
        msg = ts('community-public', resource_path(community), added, removed)
    _reindex(community)
    return msg

def community_to_intranet(ob, info):
    community = find_community(ob)
    acl = []
    moderators_group_name = community.moderators_group_name
    members_group_name = community.members_group_name
    acl.append((Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS))
    acl.append((Allow, moderators_group_name, MODERATOR_PERMS))
    acl.append((Allow, members_group_name, MEMBER_PERMS))
    # inherit from /offices
    #acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(community, acl)
    if added or removed:
        community.__acl__ = acl
        msg = ts('community-intranet', resource_path(community), added, removed)
    _reindex(community)
    return msg

#------------------------------------------------------------------------------
#   Special workflow for blog entries.
#------------------------------------------------------------------------------
"""
o Only the creator and the admin is allowed to edit or delete an entry.
"""
def blogentry_to_inherits(ob, info):
    acl = [
        (Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS),
        (Allow, ob.creator, MEMBER_PERMS),
        (Deny, Everyone, ('edit', 'delete')),
        # Note:  don't add NO_INHERIT, ergo fall back to __parent__
    ]
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('blogentry-inherits', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

def blogentry_to_private(ob, info):
    community = find_community(ob)
    acl = [(Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS)]
    acl.append((Allow, ob.creator, MEMBER_PERMS))
    moderators_group_name = community.moderators_group_name
    members_group_name = community.members_group_name
    acl.append((Allow, moderators_group_name, MODERATOR_PERMS))
    acl.append((Allow, members_group_name, GUEST_PERMS))
    acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('blogentry-private', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

#------------------------------------------------------------------------------
#   Special workflow for comments.
#------------------------------------------------------------------------------
"""
o Only the creator and the admin is allowed to edit or delete a comment.
"""
def comment_to_inherits(ob, info):
    acl = [
        (Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS),
        (Allow, ob.creator, MEMBER_PERMS),
        (Deny, Everyone, ('edit', 'delete')),
        # Note:  don't add NO_INHERIT, ergo fall back to __parent__
    ]
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('comment-inherits', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

#------------------------------------------------------------------------------
#   Workflow for forums and forum topics
#------------------------------------------------------------------------------

def forum_to_inherits(ob, info):
    acl = [(Allow, 'group.KarlStaff', (CREATE,))]
    # Note:  don't add NO_INHERIT, ergo fall back to __parent__
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('forum-inherited', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

def forum_topic_to_inherits(ob, info):
    acl = [
        (Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS),
        (Allow, ob.creator, MEMBER_PERMS),
        (Deny, Everyone, ('edit', 'delete')),
        # Note:  don't add NO_INHERIT, ergo fall back to __parent__
        ]
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('forum-topic-inherited', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

#------------------------------------------------------------------------------
#   Special workflow for profiles
#------------------------------------------------------------------------------

def to_profile_active(ob, info):
    acl  = [
        (Allow, ob.creator, MEMBER_PERMS + ('view_only',)),
    ]
    acl.append((Allow, 'group.KarlUserAdmin',
                ADMINISTRATOR_PERMS + ('view_only',)))
    acl.append((Allow, 'group.KarlAdmin',
                ADMINISTRATOR_PERMS + ('view_only',)))
    acl.append((Allow, 'group.KarlStaff',
                GUEST_PERMS + ('view_only',)))
    users = find_users(ob)
    user = users.get_by_id(ob.creator)
    if user is not None:
        groups = user['groups']
        for group, role in get_community_groups(groups):
            c_group = 'group.community:%s:%s' % (group, role)
            acl.append((Allow, c_group, GUEST_PERMS + ('view_only',)))
    acl.append((Allow, 'system.Authenticated', ('view_only',)))
    acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('to-active', resource_path(ob), added, removed)
    _reindex(ob, texts=True)
    _reindex_peopledir(ob)
    return msg

def to_profile_inactive(ob, info):
    acl  = [
        (Allow, 'system.Authenticated', ('view_only',)),
        (Allow, 'group.KarlUserAdmin', ADMINISTRATOR_PERMS + ('view_only',)),
        (Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS + ('view_only',)),
        NO_INHERIT,
    ]
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('to-inactive', resource_path(ob), added, removed)
    _reindex(ob, texts=True)
    _reindex_peopledir(ob)
    return msg

#------------------------------------------------------------------------------
#   Workflow for content within a community
#------------------------------------------------------------------------------

def content_to_inherits(ob, info):
    msg = None
    added, removed = acl_diff(ob, {})
    if hasattr(ob, '__acl__'):
        del ob.__acl__
        msg =  ts('content-inherited', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

def content_to_private(ob, info):
    community = find_community(ob)
    acl = []
    moderators_group_name = community.moderators_group_name
    members_group_name = community.members_group_name
    acl.append((Allow, 'group.KarlAdmin', ADMINISTRATOR_PERMS))
    acl.append((Allow, moderators_group_name, MODERATOR_PERMS))
    acl.append((Allow, members_group_name, MEMBER_PERMS))
    acl.append(NO_INHERIT)
    msg = None
    added, removed = acl_diff(ob, acl)
    if added or removed:
        ob.__acl__ = acl
        msg = ts('content-private', resource_path(ob), added, removed)
    _reindex(ob)
    return msg

#------------------------------------------------------------------------------
#   Workflow for intranet content
#
# This workflow applies to communities and community content in the context
# of intranet sites.
#------------------------------------------------------------------------------

intranet_content_to_inherits = content_to_inherits

# subscribers

def reset_profile(event):
    """ Subscriber for IUserAddedGroup and IUserRemovedGroup events.

    o Find the profile corresponding to the user and update its ACL
      by resettings its security workflow
    """
    if (not IUserAddedGroup.providedBy(event) and
        not IUserRemovedGroup.providedBy(event)):
        return

    profiles = find_profiles(event.site)
    profile = profiles.get(event.id)
    if profile is not None:
        reset_security_workflow(profile)


from logging import Logger

import pandas as pd
import torch
import torchaudio
from torch.utils.data import Dataset
from torchaudio.transforms import Spectrogram
from torchvision.transforms import Resize

class WavDataset(Dataset):
    """
    Dataset for comfortable getting audio files as tensors and their labels
    """
    def __init__(self, df: pd.DataFrame, n_fft: int, size: tuple[int, int], logger: Logger):
        """
        :param df: dataframe with paths to audio files and labels
        :param n_fft: using for transformation audio into spectrogram
        :param size: size of image for spectrogram
        :param logger: logger object
        """
        self.df = df
        self.spectrogram = Spectrogram(n_fft)
        self.resize = Resize(size)

        test_example, _ = self[0]
        logger.info(f"Dataset creation succeed. Output shape: {test_example.shape}")

    def __len__(self) -> int:
        return len(self.df)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, int]:
        """
        Transform an audio and return its spectrogram

        :param index: index of sample in dataset

        :return: spectrogram, label
        """
        row = self.df.iloc[index]
        path = row['path']
        label = row['label']
        wave, sample_rate = torchaudio.load(path)
        wave = wave[0]
        wave = self.spectrogram(wave)
        wave = 10 * torch.log10(wave)
        wave = wave[None]
        # wave = wave.repeat(3, 1, 1)
        wave = self.resize(wave)
        return wave, label


from pyramid.config import Configurator
from vctweb.resources import Root
from vctweb import views

def main(global_config, **settings):
    """ This function returns a Pyramid WSGI application.
    """
    config = Configurator(root_factory=Root, settings=settings)

    config.add_route('home', '/', view=views.home, renderer="templates/home.pt")
    config.add_route('patient_add', '/patients/add', view=views.patient_add, renderer="templates/patient_add.pt")
    config.add_route('patient_view', '/patients/{id}/', view=views.patient_view, renderer="templates/patient_view.pt")
    config.add_route('patient_items', '/patients/{id}/items', view=views.patient_items, renderer="templates/patient_items.pt")
    config.add_route('patient_edit', '/patients/{id}/edit', view=views.patient_edit, renderer="templates/patient_edit.pt")
    config.add_route('patient_item_add', '/patients/{id}/add', view=views.patient_item_add, renderer="templates/patient_item_add.pt")
    config.add_route('patient_item_view', '/items/{id}/', view=views.patient_item_view, renderer="templates/patient_item_view.pt")
    config.add_route('patient_item_edit', '/items/{id}/edit', view=views.patient_item_edit, renderer="templates/patient_item_edit.pt")
    config.add_route('patients', '/patients/', view=views.patients, renderer="templates/patients.pt")

    config.add_static_view('static', 'vctweb:static')
    return config.make_wsgi_app()


from __future__ import print_function, division, absolute_import
from fontTools import ttLib
from collections import OrderedDict
import argparse
import re

__version__ = '0.1'

uniname_RE = re.compile(r"""((uni)|(u))              # match "uni" or "u"
                            ((?(2)([0-9A-F]{4}))     # if "uni" match 4 digits
                             (?(3)([0-9A-F]{4,6})))  # if "u" match 4-6 digits
                            (?!.)""", re.X)          # don't match anythg else

def codepoint_from_uniname(name):
    m = re.match(uniname_RE, name)
    if m:
        codepoint_hex = m.group(4)
        return int(codepoint_hex, 16)
    else:
        raise ValueError("invalid uniname: '%s'" % name)

def parse_options(args):
    parser = argparse.ArgumentParser(
        description="Use GlyphOrderAndAliasDB to rename glyphs in TTF/OTF")
    parser.add_argument('goadb', metavar='GOADB.txt',
                        type=argparse.FileType('r'), help='input GOADB file.')
    parser.add_argument('infile', metavar='INPUT',
                        help='input font file.')
    parser.add_argument('outfile', metavar='OUTPUT', nargs='?',
                        help='output font file (default: overwrite input '
                        'file).')
    parser.add_argument('--version', action='version', version=__version__)
    options = parser.parse_args(args)
    if not options.outfile:
        # careful: this is going to overwrite the input file!
        options.outfile = options.infile
    return options

class GoadbManager(object):

    def __init__(self, goadbpath, fontpath):
        self.names, self.mapping = self.parse_goadb(goadbpath)
        self.font = ttLib.TTFont(fontpath, recalcBBoxes=False,
                                 recalcTimestamp=False)

    @staticmethod
    def parse_goadb(goadbfile):
        try:
            goadb = goadbfile.read()
        except AttributeError:
            with open(goadbfile, 'r') as f:
                goadb = f.read()
        names = OrderedDict()
        mapping = {}
        for i, line in enumerate(goadb.splitlines()):
            line = re.sub(r"#.*", "", line).strip()
            if not line:
                continue
            items = line.split()
            n = len(items)
            assert len(items) in (2, 3)
            if n == 2:
                final, friendly = items
            elif n == 3:
                final, friendly, uniname = items
                code = codepoint_from_uniname(uniname)
                if code not in mapping:
                    mapping[code] = final
                else:
                    raise AssertionError('uniname redefined: %s' % uniname)
            else:
                raise AssertionError(
                    "GOADB must have either 2 or 3 items per entry: "
                    "line %d: %s" % line)
            names[friendly] = final
        return names, mapping

    def rename_glyphs(self):
        # Do the renaming in newnames
        newnames = []
        for gname in self.font.getGlyphOrder():
            if gname in self.names:
                newnames.append(self.names[gname])
            else:
                newnames.append(gname)
        # must reload TTFont again
        font = ttLib.TTFont(self.font.reader.file.name)
        # set glyph order as newnames then load post table
        font.setGlyphOrder(newnames)
        font['post']
        self.font = font

    def save(self, outpath):
        self.font.save(outpath, reorderTables=False)

    def close(self):
        self.font.close()

def main(args=None):
    options = parse_options(args)
    manager = GoadbManager(options.goadb, options.infile)
    manager.rename_glyphs()
    manager.save(options.outfile)
    manager.close()

if __name__ == '__main__':
    main()


import unicodedata

from django.http import HttpRequest

from events.utils import (
    get_business_id,
    get_organization_name,
    PROVIDER_BUSINESS_ID_FIELD,
    PROVIDER_NAME_FIELD,
)

EVENT_BASE_DATA = {
    "offers": [
        {
            "is_free": True,
            "price": {"fi": "", "sv": "", "en": ""},
            "description": {"fi": "", "sv": "", "en": ""},
            "info_url": None,
        }
    ],
}

CREATE_EVENT_BASE_DATA = {
    "event_status": "EventScheduled",
    "type_id": "General",
    "data_source": "tet",
    "in_language": [{"@id": "http://localhost:8080/v1/language/fi/"}],
}

def _new_from(obj, keys):
    new = {}
    for key in keys:
        new[key] = obj[key]
    return new

def _shorten_description(descobj):
    shortened_obj = {}
    for lang in descobj.keys():
        desc = descobj[lang]
        if len(desc) <= 125:
            shortened_obj[lang] = desc
        else:
            shortened_obj[
                lang
            ] = f"{desc[:125]}{unicodedata.lookup('Horizontal ellipsis')}"

    return shortened_obj

def _set_event_owner(event, request):
    user = request.user

    if user.is_staff:
        if user.email:
            event["custom_data"]["editor_email"] = user.email
        # In TET settings we specify `USERNAME_CLAIM = "oid"`, which sets the user's object identifier as
        # Django username. This is added to provide "event ownership info" in case user.email is empty.
        event["custom_data"]["editor_oid"] = user.username

        event["provider"] = {
            PROVIDER_NAME_FIELD: get_organization_name(request),
        }
    else:  # user is logged in via suomi.fi
        event["provider"] = {
            PROVIDER_NAME_FIELD: get_organization_name(request),
            PROVIDER_BUSINESS_ID_FIELD: get_business_id(request),
        }

def enrich_create_event(event, publisher, request: HttpRequest):
    event.update(EVENT_BASE_DATA)
    event.update(CREATE_EVENT_BASE_DATA)

    if "publication_status" not in event or event["publication_status"] is None:
        event["publication_status"] = "draft"

    event["publisher"] = publisher
    event["short_description"] = _shorten_description(event["description"])

    _set_event_owner(event, request)

    return event

def enrich_update_event(event, request):
    event.update(EVENT_BASE_DATA)
    event["short_description"] = _shorten_description(event["description"])

    _set_event_owner(event, request)

    # Not sure why it resets publication status from draft to public without the following line
    # This is not a problem because we keep the two in sync
    # publication_status is what decides whether the event is shown in Youth UI
    event["publication_status"] = "public" if event["date_published"] else "draft"
    return event

def reduce_get_event(event):
    tetevent = _new_from(
        event,
        (
            "id",
            "name",
            "description",
            "location",
            "keywords",
            "date_published",
            "start_time",
            "end_time",
            "custom_data",
            "publication_status",
            "in_language",
            "images",
            "provider",
        ),
    )

    return tetevent


from preprocessing import *
from tqdm import tqdm
import argparse
import matplotlib.pyplot as plt
import bar_chart_race as bcr

def from_millis_to_minutes(millis: float) -> float:
    return round(millis / (1000*60), 2)

def from_minutes_to_millis(minutes: float) -> float:
    return (minutes * (1000*60))

def gather_durations_data(df, args):
    data = {}
    if args.decades:
        decades = df.groupby(pd.Grouper(key='date', freq='10YS'))
        for dt, decade in tqdm(decades):
            data[dt.year] = {}
            for duration in durations_millis:
                data[dt.year][f'{from_millis_to_minutes(duration)}'] = (decade.duration > duration).sum()
    else:
        for duration in durations_millis:
            data[from_millis_to_minutes(duration)] = (df.duration > duration).sum()
    return data

def gather_genres_data(df, args):
    df = df[df.genre.notna()]
    df['genre'] = df['genre'].apply(get_genre)
    df = df[df.genre != 'wtf']
    df.genre = df.genre.map(lambda i: genres_acc[i])
    data = {}
    selected_genres = set(args.genres)
    if args.decades:
        decades = df.groupby(pd.Grouper(key='date', freq='10YS'))
        for dt, decade in tqdm(decades):
            data[dt.year] = {}
            if len(args.genres) == 1 and "top" in args.genres[0]:
                top_k = int(args.genres[0].split(' ')[1])
                for genre_name, count in decade.groupby('genre').genre.count().nlargest(top_k).iteritems():
                    genre_name = 'prog. rock' if genre_name == 'progressive rock' else genre_name
                    data[dt.year][genre_name.capitalize()] = count
            else:
                for genre_name, genre in decade.groupby('genre'):
                    if genre_name in selected_genres or 'all' in selected_genres:
                        genre_name = 'prog. rock' if genre_name == 'progressive rock' else genre_name
                        data[dt.year][genre_name.capitalize()] = genre.genre.count()
    else:
        if len(args.genres) == 1 and "top" in args.genres[0]:
            top_k = int(args.genres[0].split(' ')[1])
            for genre_name, count in df.groupby('genre').genre.count().nlargest(top_k).iteritems():
                genre_name = 'prog. rock' if genre_name == 'progressive rock' else genre_name
                data[genre_name.capitalize()] = count
        else:
            for genre_name, genre in df.groupby('genre'):
                if genre_name in selected_genres or 'all' in selected_genres:
                    genre_name = 'prog. rock' if genre_name == 'progressive rock' else genre_name
                    data[genre_name.capitalize()] = genre.genre.count()
    return data

def chart_race_and_exit(args):
    df = pd.read_csv('data/results_genre.csv')
    df = df[df['artist'] != 'Glee Cast']
    df = df[df.genre.notna()]
    df['genre'] = df['genre'].apply(get_genre)
    df = df[df.genre != 'wtf']
    genres_acc_cap = list(map(lambda g: g.capitalize(), genres_acc))
    df.genre = df.genre.map(lambda i: genres_acc_cap[i])
    df.date = pd.to_datetime(df.date, format='%Y-%m-%d')
    df = df[df.date < datetime(year=2015, month=1, day=1)]
    groups = df.groupby(pd.Grouper(key='date', freq='1YS'))
    data = {}
    for date, group in groups:
        counts = group.genre.value_counts()
        counts = counts.to_dict()
        for g in (set(genres_acc_cap) - set(counts.keys())):
            counts[g] = 0.

        data[date] = counts
    data = pd.DataFrame(data).T
    print("Creating chart race and quitting")
    bcr.bar_chart_race(
        df=data,
        n_bars=8,
        filename='./chart_race_genres.gif',
        sort='desc',
        fixed_order=False,
        steps_per_period=20,
        period_length=800,
        period_label={'x': .98, 'y': .3, 'ha': 'right', 'va': 'center'},
        perpendicular_bar_func=None,
        period_fmt="%Y",
        cmap='dark12',
        bar_size=.70,
        scale='log',
        title=args.title
    )
    exit(0)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot random stuff duration-related')
    parser.add_argument('-d', '--durations', dest='durations', nargs='+', help='compute number of songs greater than these durations. Format: minutes.seconds')
    parser.add_argument('-g', '--genres', dest='genres', nargs='+', help='plot frequency of specified genres. Use "all" for all genres, use "top k" to select top k genres only')
    parser.add_argument('--chart_race', '-r', dest='chart_race', action='store_true', help='specify to output a bar chart race (only valid when -g is specified)')
    parser.add_argument('--decades', dest='decades', action='store_true', help='specify to group by decades')
    parser.add_argument('--normalize', dest='normalize', action='store_true', help='specify to normalize data')
    parser.add_argument('--decades_on_x', dest='decades_on_x', action='store_true', help='plots decades on x (plots durations on x otherwise)')
    parser.add_argument('--title', dest='title', help='title of the plot figure', default='')
    parser.add_argument('--x_label', dest='x_label', help='x label name', default='')
    parser.add_argument('--y_label', dest='y_label', help='y label name', default='')
    parser.add_argument('--figsize', dest='figsize', help='size of plot figure, eg. 12,10', default='12,10')
    parser.add_argument('--output_path', dest='output_path', help='output path to save plot, eg. decades.png')
    args = parser.parse_args()
    if args.chart_race and not args.genres:
        raise UserWarning("--chart_race specified without --genres. This will be ignored")
    elif args.chart_race and args.genres:
        chart_race_and_exit(args)
    durations_millis = list(map(lambda d: from_minutes_to_millis(float(d)), args.durations))
    df = pd.read_csv('data/results_genre.csv')
    df = df[df['artist'] != 'Glee Cast']
    df = aggregate_sixties_and_2020s(df)
    _, ax = plt.subplots(figsize=tuple(map(int, args.figsize.split(','))))

    if args.genres:
        data = gather_genres_data(df, args)
    else:
        data = gather_durations_data(df, args)

    data = pd.DataFrame(data) if args.decades else pd.DataFrame(data, index=['Frequency'])
    if args.normalize:
        for ind in data.index:
            data.loc[ind] = data.loc[ind] / data.loc[ind].max()
    if args.decades_on_x:
        data.T.plot.bar(title=args.title, xlabel=args.x_label, ylabel=args.y_label, ax=ax, legend=args.decades)
    else:
        data.plot.bar(title=args.title, xlabel=args.x_label, ylabel=args.y_label, ax=ax, legend=args.decades)

    if args.output_path:
        plt.savefig(args.output_path)
    else:
        plt.show()


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

df_adm = pd.read_csv('/Users/KexinHuang/Downloads/ADMISSIONS.csv')
df_adm.ADMITTIME = pd.to_datetime(df_adm.ADMITTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')
df_adm.DISCHTIME = pd.to_datetime(df_adm.DISCHTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')
df_adm.DEATHTIME = pd.to_datetime(df_adm.DEATHTIME, format = '%Y-%m-%d %H:%M:%S', errors = 'coerce')

df_adm = df_adm.sort_values(['SUBJECT_ID','ADMITTIME'])
df_adm = df_adm.reset_index(drop = True)
df_adm['NEXT_ADMITTIME'] = df_adm.groupby('SUBJECT_ID').ADMITTIME.shift(-1)
df_adm['NEXT_ADMISSION_TYPE'] = df_adm.groupby('SUBJECT_ID').ADMISSION_TYPE.shift(-1)

rows = df_adm.NEXT_ADMISSION_TYPE == 'ELECTIVE'
df_adm.loc[rows,'NEXT_ADMITTIME'] = pd.NaT
df_adm.loc[rows,'NEXT_ADMISSION_TYPE'] = np.NaN

df_adm = df_adm.sort_values(['SUBJECT_ID','ADMITTIME'])

#When we filter out the "ELECTIVE", we need to correct the next admit time for these admissions since there might be 'emergency' next admit after "ELECTIVE"
df_adm[['NEXT_ADMITTIME','NEXT_ADMISSION_TYPE']] = df_adm.groupby(['SUBJECT_ID'])[['NEXT_ADMITTIME','NEXT_ADMISSION_TYPE']].fillna(method = 'bfill')
df_adm['DAYS_NEXT_ADMIT']=  (df_adm.NEXT_ADMITTIME - df_adm.DISCHTIME).dt.total_seconds()/(24*60*60)
df_adm['OUTPUT_LABEL'] = (df_adm.DAYS_NEXT_ADMIT < 30).astype('int')
### filter out newborn and death
df_adm = df_adm[df_adm['ADMISSION_TYPE']!='NEWBORN']
df_adm = df_adm[df_adm.DEATHTIME.isnull()]
df_adm['DURATION'] = (df_adm['DISCHTIME']-df_adm['ADMITTIME']).dt.total_seconds()/(24*60*60)

df_notes = pd.read_csv('/Users/KexinHuang/Downloads/NOTEEVENTS.csv')
df_notes = df_notes.sort_values(by=['SUBJECT_ID','HADM_ID','CHARTDATE'])
df_adm_notes = pd.merge(df_adm[['SUBJECT_ID','HADM_ID','ADMITTIME','DISCHTIME','DAYS_NEXT_ADMIT','NEXT_ADMITTIME','ADMISSION_TYPE','DEATHTIME','OUTPUT_LABEL','DURATION']],
                        df_notes[['SUBJECT_ID','HADM_ID','CHARTDATE','TEXT','CATEGORY']],
                        on = ['SUBJECT_ID','HADM_ID'],
                        how = 'left')

df_adm_notes.ADMITTIME_C = df_adm_notes.ADMITTIME.apply(lambda x: str(x).split(' ')[0])
df_adm_notes['ADMITTIME_C'] = pd.to_datetime(df_adm_notes.ADMITTIME_C, format = '%Y-%m-%d', errors = 'coerce')
df_adm_notes['CHARTDATE'] = pd.to_datetime(df_adm_notes.CHARTDATE, format = '%Y-%m-%d', errors = 'coerce')

### If Discharge Summary
df_discharge = df_adm_notes[df_adm_notes['CATEGORY'] == 'Discharge summary']
# multiple discharge summary for one admission -> after examination -> replicated summary -> replace with the last one
df_discharge = (df_discharge.groupby(['SUBJECT_ID','HADM_ID']).nth(-1)).reset_index()
df_discharge=df_discharge[df_discharge['TEXT'].notnull()]

### If Less than n days on admission notes (Early notes)
def less_n_days_data (df_adm_notes, n):

    df_less_n = df_adm_notes[((df_adm_notes['CHARTDATE']-df_adm_notes['ADMITTIME_C']).dt.total_seconds()/(24*60*60))<n]
    df_less_n=df_less_n[df_less_n['TEXT'].notnull()]
    #concatenate first
    df_concat = pd.DataFrame(df_less_n.groupby('HADM_ID')['TEXT'].apply(lambda x: "%s" % ' '.join(x))).reset_index()
    df_concat['OUTPUT_LABEL'] = df_concat['HADM_ID'].apply(lambda x: df_less_n[df_less_n['HADM_ID']==x].OUTPUT_LABEL.values[0])

    return df_concat

df_less_2 = less_n_days_data(df_adm_notes, 2)
df_less_3 = less_n_days_data(df_adm_notes, 3)

import re
def preprocess1(x):
    y=re.sub('\\[(.*?)\\]','',x) #remove de-identified brackets
    y=re.sub('[0-9]+\.','',y) #remove 1.2. since the segmenter segments based on this
    y=re.sub('dr\.','doctor',y)
    y=re.sub('m\.d\.','md',y)
    y=re.sub('admission date:','',y)
    y=re.sub('discharge date:','',y)
    y=re.sub('--|__|==','',y)
    return y

def preprocessing(df_less_n):
    df_less_n['TEXT']=df_less_n['TEXT'].fillna(' ')
    df_less_n['TEXT']=df_less_n['TEXT'].str.replace('\n',' ')
    df_less_n['TEXT']=df_less_n['TEXT'].str.replace('\r',' ')
    df_less_n['TEXT']=df_less_n['TEXT'].apply(str.strip)
    df_less_n['TEXT']=df_less_n['TEXT'].str.lower()

    df_less_n['TEXT']=df_less_n['TEXT'].apply(lambda x: preprocess1(x))

    #to get 318 words chunks for readmission tasks
    from tqdm import tqdm
    df_len = len(df_less_n)
    want=pd.DataFrame({'ID':[],'TEXT':[],'Label':[]})
    for i in tqdm(range(df_len)):
        x=df_less_n.TEXT.iloc[i].split()
        n=int(len(x)/318)
        for j in range(n):
            want=want.append({'TEXT':' '.join(x[j*318:(j+1)*318]),'Label':df_less_n.OUTPUT_LABEL.iloc[i],'ID':df_less_n.HADM_ID.iloc[i]},ignore_index=True)
        if len(x)%318>10:
            want=want.append({'TEXT':' '.join(x[-(len(x)%318):]),'Label':df_less_n.OUTPUT_LABEL.iloc[i],'ID':df_less_n.HADM_ID.iloc[i]},ignore_index=True)

    return want

df_discharge = preprocessing(df_discharge)
df_less_2 = preprocessing(df_less_2)
df_less_3 = preprocessing(df_less_3)

### An example to get the train/test/split with random state:
### note that we divide on patient admission level and share among experiments, instead of notes level.
### This way, since our methods run on the same set of admissions, we can see the
### progression of readmission scores.

readmit_ID = df_adm[df_adm.OUTPUT_LABEL == 1].HADM_ID
not_readmit_ID = df_adm[df_adm.OUTPUT_LABEL == 0].HADM_ID
#subsampling to get the balanced pos/neg numbers of patients for each dataset
not_readmit_ID_use = not_readmit_ID.sample(n=len(readmit_ID), random_state=1)
id_val_test_t=readmit_ID.sample(frac=0.2,random_state=1)
id_val_test_f=not_readmit_ID_use.sample(frac=0.2,random_state=1)

id_train_t = readmit_ID.drop(id_val_test_t.index)
id_train_f = not_readmit_ID_use.drop(id_val_test_f.index)

id_val_t=id_val_test_t.sample(frac=0.5,random_state=1)
id_test_t=id_val_test_t.drop(id_val_t.index)

id_val_f=id_val_test_f.sample(frac=0.5,random_state=1)
id_test_f=id_val_test_f.drop(id_val_f.index)

# test if there is overlap between train and test, should return "array([], dtype=int64)"
(pd.Index(id_test_t).intersection(pd.Index(id_train_t))).values

id_test = pd.concat([id_test_t, id_test_f])
test_id_label = pd.DataFrame(data = list(zip(id_test, [1]*len(id_test_t)+[0]*len(id_test_f))), columns = ['id','label'])

id_val = pd.concat([id_val_t, id_val_f])
val_id_label = pd.DataFrame(data = list(zip(id_val, [1]*len(id_val_t)+[0]*len(id_val_f))), columns = ['id','label'])

id_train = pd.concat([id_train_t, id_train_f])
train_id_label = pd.DataFrame(data = list(zip(id_train, [1]*len(id_train_t)+[0]*len(id_train_f))), columns = ['id','label'])

#get discharge train/val/test

discharge_train = df_discharge[df_discharge.ID.isin(train_id_label.id)]
discharge_val = df_discharge[df_discharge.ID.isin(val_id_label.id)]
discharge_test = df_discharge[df_discharge.ID.isin(test_id_label.id)]

# subsampling for training....since we obtain training on patient admission level so now we have same number of pos/neg readmission
# but each admission is associated with different length of notes and we train on each chunks of notes, not on the admission, we need
# to balance the pos/neg chunks on training set. (val and test set are fine) Usually, positive admissions have longer notes, so we need
# find some negative chunks of notes from not_readmit_ID that we haven't used yet

df = pd.concat([not_readmit_ID_use, not_readmit_ID])
df = df.drop_duplicates(keep=False)
#check to see if there are overlaps
(pd.Index(df).intersection(pd.Index(not_readmit_ID_use))).values

# for this set of split with random_state=1, we find we need 400 more negative training samples
not_readmit_ID_more = df.sample(n=400, random_state=1)
discharge_train_snippets = pd.concat([df_discharge[df_discharge.ID.isin(not_readmit_ID_more)], discharge_train])

#shuffle
discharge_train_snippets = discharge_train_snippets.sample(frac=1, random_state=1).reset_index(drop=True)

#check if balanced
discharge_train_snippets.Label.value_counts()

discharge_train_snippets.to_csv('./discharge/train.csv')
discharge_val.to_csv('./discharge/val.csv')
discharge_test.to_csv('./discharge/test.csv')

### for Early notes experiment: we only need to find training set for 3 days, then we can test
### both 3 days and 2 days. Since we split the data on patient level and experiments share admissions
### in order to see the progression, the 2 days training dataset is a subset of 3 days training set.
### So we only train 3 days and we can test/val on both 2 & 3days or any time smaller than 3 days. This means
### if we train on a dataset with all the notes in n days, we can predict readmissions smaller than n days.

#for 3 days note, similar to discharge

early_train = df_less_3[df_less_3.ID.isin(train_id_label.id)]
not_readmit_ID_more = df.sample(n=500, random_state=1)
early_train_snippets = pd.concat([df_less_3[df_less_3.ID.isin(not_readmit_ID_more)], early_train])
#shuffle
early_train_snippets = early_train_snippets.sample(frac=1, random_state=1).reset_index(drop=True)
early_train_snippets.to_csv('./3days/train.csv')

early_val = df_less_3[df_less_3.ID.isin(val_id_label.id)]
early_val.to_csv('./3days/val.csv')

# we want to test on admissions that are not discharged already. So for less than 3 days of notes experiment,
# we filter out admissions discharged within 3 days
actionable_ID_3days = df_adm[df_adm['DURATION'] >= 3].HADM_ID
test_actionable_id_label = test_id_label[test_id_label.id.isin(actionable_ID_3days)]
early_test = df_less_3[df_less_3.ID.isin(test_actionable_id_label.id)]

early_test.to_csv('./3days/test.csv')

#for 2 days notes, we only obtain test set. Since the model parameters are tuned on the val set of 3 days

actionable_ID_2days = df_adm[df_adm['DURATION'] >= 2].HADM_ID

test_actionable_id_label_2days = test_id_label[test_id_label.id.isin(actionable_ID_2days)]

early_test_2days = df_less_2[df_less_2.ID.isin(test_actionable_id_label_2days.id)]

early_test_2days.to_csv('./2days/test.csv')


from abc import ABC

import loralib as lora
import torch
from chatgpt.dataset import RewardDataset
from chatgpt.models.loss import PairWiseLoss
from torch.optim import Adam, Optimizer
from torch.utils.data import DataLoader
from tqdm import tqdm

from .strategies import Strategy
from .utils import is_rank_0

class RewardModelTrainer(ABC):
    """
        Trainer to use while training reward model.

    Args:
        model (torch.nn.Module): the model to train
        strategy (Strategy): the strategy to use for training
        optim(Optimizer): the optimizer to use for training
        train_dataset (RewardDataset): the dataset to use for training
        eval_dataset (RewardDataset): the dataset to use for evaluation
        batch_size (int, defaults to 1): the batch size while training
        max_epochs (int, defaults to 2): the number of epochs to train
        optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer
    """

    def __init__(
        self,
        model,
        strategy: Strategy,
        optim: Optimizer,
        train_dataset: RewardDataset,
        eval_dataset: RewardDataset,
        batch_size: int = 1,
        max_epochs: int = 2,
    ) -> None:
        super().__init__()
        self.strategy = strategy
        self.epochs = max_epochs
        self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size)
        self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)

        self.model = strategy.setup_model(model)
        if "DDP" in str(self.strategy):
            self.model = self.model.module
        self.loss_fn = PairWiseLoss()
        self.optimizer = strategy.setup_optimizer(optim, self.model)

    def fit(self, use_lora):
        epoch_bar = tqdm(range(self.epochs), desc='Train epoch', disable=not is_rank_0())
        for epoch in range(self.epochs):
            step_bar = tqdm(range(self.train_dataloader.__len__()),
                            desc='Train step of epoch %d' % epoch,
                            disable=not is_rank_0())
            # train
            self.model.train()
            for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader:
                chosen_ids = chosen_ids.squeeze(1).cuda()
                c_mask = c_mask.squeeze(1).cuda()
                reject_ids = reject_ids.squeeze(1).cuda()
                r_mask = r_mask.squeeze(1).cuda()
                chosen_reward = self.model(chosen_ids, attention_mask=c_mask)
                reject_reward = self.model(reject_ids, attention_mask=r_mask)
                loss = self.loss_fn(chosen_reward, reject_reward)
                self.strategy.backward(loss, self.model, self.optimizer)
                self.strategy.optimizer_step(self.optimizer)
                self.optimizer.zero_grad()
                step_bar.update()
                step_bar.set_postfix({'loss': loss.item()})

            # eval
            self.model.eval()
            with torch.no_grad():
                dist = 0
                loss_sum = 0
                for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader:
                    chosen_ids = chosen_ids.squeeze(1).cuda()
                    c_mask = c_mask.squeeze(1).cuda()
                    reject_ids = reject_ids.squeeze(1).cuda()
                    r_mask = r_mask.squeeze(1).cuda()
                    chosen_reward = self.model(chosen_ids, attention_mask=c_mask)
                    reject_reward = self.model(reject_ids, attention_mask=r_mask)
                    dist += (chosen_reward - reject_reward).mean().item()
                    loss = self.loss_fn(chosen_reward, reject_reward)
                    loss_sum += loss.item()
                dist_mean = dist / self.eval_dataloader.__len__()
                loss_mean = loss_sum / self.eval_dataloader.__len__()
            epoch_bar.update()
            step_bar.set_postfix({'loss': loss_mean, 'dist_mean': dist_mean})
            step_bar.close()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Factory that returns all the algorithms fron the project that are finall.
"""
from __future__ import print_function

import importlib

# pylint-disable: wrong-import-order
from ag_frame.algorithms.representations.mutations import base
from ag_frame import exceptions

PREFIX = "ag_frame.algorithms.representations.mutations"
MUTATIONS = [
    "bit_mutate",
]

def mutation_factory():
    """Get all mutations as a list."""
    all_mutations = []
    for mutation_module in MUTATIONS:
        module_name = "{}.{}".format(PREFIX, mutation_module)
        module = importlib.import_module(module_name)
        for item in dir(module):
            item = getattr(module, item)
            try:
                if issubclass(item, base.BaseMutation):
                    item.is_implemented()
                else:
                    continue
            except (exceptions.MutationNotImplemented, TypeError):
                continue

            all_mutations.append(item)

    return all_mutations