"""
Migrate management command.
"""

from __future__ import print_function

from copy import deepcopy
import os
from itertools import chain
from optparse import make_option
import pprint
import subprocess
from StringIO import StringIO

from two2three.pgen2 import driver
from two2three import pygram, pytree

from django.core.management.base import BaseCommand
from django.utils.importlib import import_module

from south.migration import Migrations

from south_rebase.diff import ModelsDiff

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
        make_option('--git', action='store_const', dest='plan_generator',
                    const='git', default='interactive',
                    help='Rebase all migrations from last commit'),
    )

    def handle(self, app_name=None, **options):
        self.options = options

        if app_name is None:
            print('No app_label given')
            return

        migrations = Migrations(app_name)
        rebase_plan = self.split_migrations(migrations)
        if not rebase_plan.migrations:
            print('No migrations to rebase')
            return

        diffs = []
        for plan_elem in rebase_plan.migrations:
            """Prepare diffs"""
            diffs.append(self.diff_frozen_models(plan_elem['previous'],
                                                 plan_elem['migration']))

        prev_migration = rebase_plan.onto
        self.prepare_file_cache()
        for diff, migration in zip(
                diffs, (m['migration'] for m in rebase_plan.migrations)):
            """Rebase - rename migrations and apply diffs"""
            # Prepare new name for migration
            last_migration_no, _ = split_migration_name(prev_migration.name())
            _, migration_base_name = split_migration_name(migration.name())

            migration_path = self.find_migration_path(migration)
            migration_dir = os.path.dirname(migration_path)

            new_migration_filename = '%04i_%s.py' % (last_migration_no + 1,
                                                     migration_base_name)
            new_migration_path = os.path.join(
                migration_dir, new_migration_filename)

            # patch migration tree
            new_migration_tree = self.get_updated_migration(
                migration, prev_migration, diff=diff)

            self.write_file(new_migration_path,
                            unicode(new_migration_tree).encode('utf-8'))
            self.remove_file(migration_path)

            prev_migration = migration

        self.flush_file_cache()

    def write_file(self, path, content):
        self.files_to_write[path] = content

    def remove_file(self, path):
        self.files_to_remove.append(path)

    def prepare_file_cache(self):
        self.files_to_remove = []
        self.files_to_write = {}  # (filename, content)

    def flush_file_cache(self):
        if all(map(os.path.exists, self.files_to_remove)) and \
                not any(map(os.path.exists, self.files_to_write)):
            for path in self.files_to_remove:
                os.remove(path)
            for path, data in self.files_to_write.items():
                if not isinstance(data, str):
                    data = data.encode('utf-8')
                with open(path, 'w') as f:
                    f.write(data)

    def find_migration_path(self, migration):
        """Find path to file this migration comes from"""
        path = import_module(migration.migration_class().__module__).__file__
        if path.endswith('.pyc'):
            path = path[:-1]
        return path

    def split_migrations(self, all_migrations):
        """Pick splitting algorithm and return rebase_plan"""
        generators = dict(interactive=self.interactive_split_migrations,
                          git=self.git_split_migrations)
        return generators[self.options['plan_generator']](all_migrations)

    def interactive_split_migrations(self, all_migrations):
        """Prepare rebase_plan by asking user about conflicting migrations"""
        rebase_plan = RebasePlan()

        known_migrations_by_number = {}
        for m in all_migrations:
            num, base_name = split_migration_name(m.name())
            known_migrations_by_number.setdefault(num, [])
            known_migrations_by_number[num].append(m)

        rebase_plan.migrations = []
        for num, num_migrations in known_migrations_by_number.items():
            """For every migration number, ask about conflicts"""
            if len(num_migrations) > 1:
                selected_migration = interactive_select(
                    [(m.name(), m) for m in num_migrations],
                    'Colliding migration names:',
                    'pick the one to rebase')

                last_num_migrations = set(known_migrations_by_number[num - 1])
                """last_num_migrations are migrations for previous number,
                minus the one that is already selected for rebasing"""
                if rebase_plan.migrations:
                    last_num_migrations.discard(
                        rebase_plan.migrations[-1]['migration'])

                prev = only_element(last_num_migrations)
                rebase_plan.migrations.append({
                    'migration': selected_migration,
                    'previous': prev,
                })

        last_num_migrations = set(num_migrations)
        if rebase_plan.migrations:
            last_num_migrations.discard(
                rebase_plan.migrations[-1]['migration'])
        rebase_plan.onto = only_element(last_num_migrations)

        return rebase_plan

    def git_split_migrations(self, all_migrations):
        """return plan for rebasing all migrations from last commit"""
        repo_path = os.path.abspath(git_cdup())
        new_migration_paths = set()

        normalize = lambda p: os.path.realpath(p)

        for added, deleted, path in git_diff_tree():
            if added and '/migrations/' in path:
                if not deleted:
                    new_migration_paths.add(normalize(
                        os.path.join(repo_path, path)))
                else:
                    print(u'WARNING: Not using migration "%s" because '
                          u'it is modified and not new' % path)

        # Pick which migrations to rebase
        to_rebase = []
        other_migrations = []
        for migration in all_migrations:
            normalized = normalize(self.find_migration_path(migration))
            if normalized in new_migration_paths:
                to_rebase.append(migration)
            else:
                other_migrations.append(migration)

        first_num = min(migration_number(m) for m in to_rebase)

        others_last_num = max(migration_number(m)
                              for m in other_migrations)

        prev_num = first_num - 1
        last_common_migration = only_element(
            m for m in other_migrations if migration_number(m) == prev_num)

        rebase_plan = RebasePlan()
        for prev, migration in zip(chain([last_common_migration], to_rebase),
                                   to_rebase):
            rebase_plan.migrations.append({
                'previous': prev,
                'migration': migration,
            })
        rebase_plan.onto = only_element(
            m for m in other_migrations
            if migration_number(m) == others_last_num)

        return rebase_plan

    def diff_frozen_models(self, prev_migration, migration):
        """Return differences between frozen models"""
        app_label = migration.app_label()

        prev_models = prev_migration.migration_class().models
        models = migration.migration_class().models

        models_for_diff = {
            m for m in chain(prev_models.keys(), models.keys())
            if m.startswith('%s.' % app_label)}

        filter_for_diff = lambda ms: {k: v for k, v in ms.items()
                                      if k in models_for_diff}

        diff = ModelsDiff.diff(filter_for_diff(prev_models),
                               filter_for_diff(models))
        return diff

    def get_updated_migration(self, migration, prev_migration, diff):
        """Return tree of migration with diff applied"""
        self.driver = driver.Driver(pygram.python_grammar,
                                    convert=pytree.convert)
        migration_src = open(self.find_migration_path(migration)).read() + '\n'
        tree = self.driver.parse_string(migration_src)
        patched_tree = deepcopy(tree)
        self.patch_migration_tree(
            patched_tree,
            diff.apply(prev_migration.migration_class().models))
        return patched_tree

    def patch_migration_tree(self, tree, new_models):
        model_dict_assignment = find_assignments(
            find_classes(tree, name='Migration').next(),
            name='models').next()
        new_models_assignment_tree = self.driver.parse_string(
            'models = %s\n' % pprint_str(new_models))
        new_models_assignment = find_assignments(
            new_models_assignment_tree).next()
        model_dict_assignment.children[2] = new_models_assignment.children[2]

class RebasePlan(object):
    """Object representing a plan for rebasing."""
    def __init__(self):
        self.onto = None  # What is the new base migration?
        self.migrations = []
        """List of dicts of migrations to rebase (in target order). Used keys:
            - migration - the migration that needs rebasing
            - previous - a migration this one was based on.
        """

def split_migration_name(name):
    num, base = name.split('_', 1)
    return int(num), base

def migration_number(migration):
    return split_migration_name(migration.name())[0]

def iterate_tree(tree):
    yield tree
    for child in tree.children:
        for node in iterate_tree(child):
            yield node

def find_classes(tree, name=None):
    """Find all classes in tree"""
    for node in iterate_tree(tree):
        if node.type == pygram.python_symbols.classdef:
            if name is None or node.children[1].value == name:
                yield node

def find_assignments(tree, name=None):
    for node in iterate_tree(tree):
        if node.type == pygram.python_symbols.expr_stmt and \
                (name is None or name == node.children[0].value) and \
                node.children[1].value == '=':
            yield node

def interactive_select(options, pre_message=None, post_message=None):
    """Select one of options interactively"""
    if len(options) == 0:
        return
    if not isinstance(options[0], tuple):
        options = [(unicode(o), o) for o in options]
    if len(options) == 1:
        return options[0][1]

    def print_options():
        if pre_message is not None:
            print(pre_message)

        for i, (opt_name, opt_value) in enumerate(options):
            print('% 4i: %s' % (i, opt_name))

        if post_message is not None:
            print(post_message)

    option_index = -1

    while not 0 <= option_index < len(options):
        try:
            print_options()
            option_index = int(raw_input())
        except ValueError:
            pass

    return options[option_index][1]

def only_element(collection):
    """Return the only element of collection, or throw if it has more"""
    coll_iter = iter(collection)
    try:
        elem = coll_iter.next()
    except StopIteration:
        raise ValueError('Empty collection')
    try:
        coll_iter.next()
        raise ValueError('Collection has more elements')
    except StopIteration:
        return elem

def pprint_str(obj, **kwargs):
    assert 'stream' not in kwargs
    stream = StringIO()
    pprint.pprint(obj, stream=stream, **kwargs)
    return stream.getvalue()

def git_diff_tree(treeish='HEAD', other_treeish=None):
    """Does git diff-tree and yields tuples of (added, deleted, file_path)"""
    args = 'git diff-tree --numstat -r -z'.split() + [treeish]
    if other_treeish is not None:
        args.append(other_treeish)
    raw_output = subprocess.check_output(args)
    # Drop the first line - it's the commit hash
    lines = raw_output.split(chr(0))[1:]
    for line in filter(None, lines):
        added, deleted, path = tuple(line.split('\t'))
        yield (int(added), int(deleted), path.decode('utf-8'))

def git_cdup():
    """Return path to root of git project"""
    return subprocess.check_output('git rev-parse --show-cdup'.split()).strip()


from datetime import date, datetime, timedelta, timezone

from dateutil.relativedelta import relativedelta

from ..errors import DatabaseConnectionError
from ..services.database import ConnectionPool
from ..services.press_release_link import PressReleaseLinkService
from ..views.view import View

class PressReleaseView(View):
    """旭川市新型コロナウイルス感染症報道発表日データ

    旭川市新型コロナウイルス感染症報道発表日データをFlaskへ渡すデータにする

    Attributes:
        latest_date (date): 最新の報道発表日

    """

    def __init__(self, pool: ConnectionPool):
        """
        Args:
            table_name (str): テーブル名
            pool (:obj:`ConnectionPool`): SimpleConnectionPoolを要素に持つオブジェクト

        """
        self.__latest_date = self._get_latest_date(pool)

    @property
    def latest_date(self):
        return self.__latest_date

    def _get_latest_date(self, pool: ConnectionPool) -> date:
        """最新の報道発表日の日付を返す

        Args:
            pool (:obj:`ConnectionPool`): SimpleConnectionPoolを要素に持つオブジェクト

        Returns:
            latest_date (date): 最新の報道発表日の日付データ

        """
        now = datetime.now(timezone(timedelta(hours=+9), "JST"))
        try:
            press_release_link_service = PressReleaseLinkService(pool)
            latest_date = press_release_link_service.get_latest_publication_date()
        except DatabaseConnectionError:
            # エラーが起きた場合現在日付を基準とする。
            # このとき市の発表が16時になることが多いので16時より前なら前日を基準とする。
            latest_date = now.date()
            if now.hour < 16:
                latest_date = latest_date - relativedelta(days=1)

        return latest_date


import math
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import os

from collections import OrderedDict

from fairseq import utils, distributed_utils
from fairseq.models import (
    FairseqEncoderDecoderModel,
    BaseFairseqModel,
    register_model,
    register_model_architecture,
)
from fairseq.models.transformer import (
    base_architecture,
    Embedding,
    TransformerModel,
    FairseqEncoder,
    FairseqIncrementalDecoder,
    EncoderOut,
)

from fairseq.modules import (
    AdaptiveSoftmax,
    LayerNorm,
    PositionalEmbedding,
    SinusoidalPositionalEmbedding,
    TransformerDecoderLayer,
    TransformerEncoderLayer,
)

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    try:
        tensors_gather = [torch.zeros_like(tensor)
            for _ in range(torch.distributed.get_world_size())]
        # rank = torch.distributed.get_rank()
        # tensors_gather[rank] = tensor
        # output = torch.cat(tensors_gather, dim=0)
        # torch.distributed.all_reduce(output)
        # print(rank, output[0], force=True)
        torch.distributed.all_gather(tensors_gather, tensor.detach(), async_op=False)
        output = torch.cat(tensors_gather, dim=0)
        return output
    except AssertionError as e:
        print('|WARNING :', e)
        return tensor

@register_model('backsum_transformer')
class BacksumTransformerModel(BaseFairseqModel):
    """Train Transformer models for multiple language pairs simultaneously.

    Requires `--task multilingual_translation`.

    We inherit all arguments from TransformerModel and assume that all language
    pairs use a single Transformer architecture. In addition, we provide several
    options that are specific to the multilingual setting.

    Args:
        --share-encoder-embeddings: share encoder embeddings across all source languages
        --share-decoder-embeddings: share decoder embeddings across all target languages
        --share-encoders: share all encoder params (incl. embeddings) across all source languages
        --share-decoders: share all decoder params (incl. embeddings) across all target languages
    """

    def __init__(self, args, encoders, decoders, task):
        super().__init__()
        assert encoders.keys() == decoders.keys()
        self.keys = list(encoders.keys())
        for key in self.keys:
            assert isinstance(encoders[key], FairseqEncoder)
            assert isinstance(decoders[key], FairseqIncrementalDecoder)

        self.models = nn.ModuleDict({
            key: FairseqEncoderDecoderModel(encoders[key], decoders[key])
            for key in self.keys
        })
        assert args.share_encoders
        if (task.momentum_contrast_loss_ratio > 0.0 or task.moco_steps is not None
        or task.byol_ratio > 0.0 or task.byol_steps is not None
        or task.parallel_byol_ratio > 0.0 or task.parallel_byol_steps is not None
        or task.gold_gen_byol_ratio > 0.0 or task.gold_gen_steps is not None
        or task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None):
            self.encoder_extra_fc = nn.Sequential(nn.Linear(args.encoder_embed_dim, args.encoder_embed_dim * 4), nn.ReLU(), nn.Linear(args.encoder_embed_dim * 4, args.encoder_embed_dim))
            if not args.cross_byol:
                self.encoder_extra_fc2 = nn.Sequential(LayerNorm(args.encoder_embed_dim), nn.Linear(args.encoder_embed_dim, args.encoder_embed_dim * 4), nn.ReLU(), nn.Linear(args.encoder_embed_dim * 4, args.encoder_embed_dim))

            self.momentum_encoder = copy.deepcopy(encoders[self.keys[0]])
            self.momentum_encoder.requires_grad_(False)
            if task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None:
                self.momentum_decoder = copy.deepcopy(decoders[self.keys[0]])
                self.momentum_decoder.requires_grad_(False)
            self.momentum_encoder_extra_fc = copy.deepcopy(self.encoder_extra_fc)
            self.momentum_encoder_extra_fc.requires_grad_(False)
        if task.momentum_contrast_loss_ratio > 0.0 or task.moco_steps is not None:
            momentum_bank = torch.nn.functional.normalize( torch.randn([args.encoder_embed_dim, args.momentum_contrast_capcity], device='cuda'), dim=0 )
            if args.fp16:
                momentum_bank = momentum_bank.half()

            self.register_buffer('momentum_bank', momentum_bank)
            self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

        if (task.byol_ratio > 0.0 or task.byol_steps is not None
        or task.parallel_byol_ratio or task.parallel_byol_steps is not None
        or task.gold_gen_byol_ratio or task.gold_gen_steps is not None
        or task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None):
            from fairseq.modules.multihead_attention import MultiheadAttention
            self.cross_attention = MultiheadAttention(
                args.encoder_embed_dim,
                args.encoder_attention_heads,
                kdim=args.encoder_embed_dim,
                vdim=args.encoder_embed_dim,
                dropout=args.attention_dropout,
                encoder_decoder_attention=True,
            )

        # lang embedding for cnndm and inshorts, 0 for cnndm, 1 for inshorts .
        if args.share_decoders and args.share_encoders:
            self.lang_embed_token = Embedding(len(args.langs), args.encoder_embed_dim, padding_idx=None)
            for k in self.models.keys():
                self.models[k].encoder.lang_embed_token = self.lang_embed_token
                self.models[k].decoder.lang_embed_token = self.lang_embed_token

    @torch.no_grad()
    def dequeue_and_enqueue(self, keys):
        keys = concat_all_gather(keys)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        momentum_contrast_capcity = self.momentum_bank.shape[1]
        # if ptr + batch_size > momentum_contrast_capcity:
        #     tail = momentum_contrast_capcity - ptr
        #     head = ptr + batch_size - momentum_contrast_capcity
        #     self.momentum_bank[:, ptr: ] = keys[:tail].T
        #     self.momentum_bank[:, :head] = keys[-head:].T
        # else:
        # # replace the keys at ptr (dequeue and enqueue)
        #     self.momentum_bank[:, ptr:ptr + batch_size] = keys.T
        self.momentum_bank.copy_(torch.cat([self.momentum_bank[:, batch_size:], keys.T], dim=1))
        ptr = (ptr + batch_size) % momentum_contrast_capcity # move pointer

        self.queue_ptr[0] = ptr

    def byol_cross_attention(self, q, k):
        x, attn = self.cross_attention(
            query=k,
            key=q,
            value=q
        )
        return x

    def forward(self, src_tokens, src_lengths, prev_output_tokens, src_lang, tgt_lang, src_lang_tokens=None, tgt_lang_tokens=None, **kwargs):
        key = '{}-{}'.format(src_lang, tgt_lang)
        encoder_out = self.models[key].encoder(src_tokens, src_lengths, lang_tokens=src_lang_tokens, **kwargs)
        decoder_out = self.models[key].decoder(
                prev_output_tokens, encoder_out, lang_tokens=tgt_lang_tokens, **kwargs,
            )
        decoder_out[1]['model_key'] = key
        return decoder_out

    def get_normalized_probs(self, net_output, log_probs, sample=None):
        key = net_output[1]['model_key']
        return self.models[key].get_normalized_probs(net_output, log_probs, sample)

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        TransformerModel.add_args(parser)
        parser.add_argument('--share-encoder-embeddings', action='store_true',
                            help='share encoder embeddings across languages')
        parser.add_argument('--share-decoder-embeddings', action='store_true',
                            help='share decoder embeddings across languages')
        parser.add_argument('--share-encoders', action='store_true',
                            help='share encoders across languages')
        parser.add_argument('--share-decoders', action='store_true',
                            help='share decoders across languages')
        parser.add_argument('--load-decoders', action='store_true')

    @classmethod
    def build_model(cls, args, task):
        """Build a new model instance."""

        # make sure all arguments are present in older models
        base_multilingual_architecture(args)

        if not hasattr(args, 'max_source_positions'):
            args.max_source_positions = 1024
        if not hasattr(args, 'max_target_positions'):
            args.max_target_positions = 1024

        # src_langs = [lang_pair.split('-')[0] for lang_pair in task.model_lang_pairs]
        # tgt_langs = [lang_pair.split('-')[1] for lang_pair in task.model_lang_pairs]
        src_langs = task.langs
        tgt_langs = task.langs

        if args.share_encoders:
            args.share_encoder_embeddings = True
        if args.share_decoders:
            args.share_decoder_embeddings = True

        def build_embedding(dictionary, embed_dim, path=None):
            num_embeddings = len(dictionary)
            padding_idx = dictionary.pad()
            emb = Embedding(num_embeddings, embed_dim, padding_idx)
            # if provided, load from preloaded dictionaries
            if path:
                embed_dict = utils.parse_embedding(path)
                utils.load_embedding(embed_dict, dictionary, emb)
            return emb

        # build shared embeddings (if applicable)
        shared_encoder_embed_tokens, shared_decoder_embed_tokens = None, None
        if args.share_all_embeddings:
            if args.encoder_embed_dim != args.decoder_embed_dim:
                raise ValueError(
                    '--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim')
            if args.decoder_embed_path and (
                    args.decoder_embed_path != args.encoder_embed_path):
                raise ValueError('--share-all-embeddings not compatible with --decoder-embed-path')
            shared_encoder_embed_tokens = build_embedding(
                task.dict,
                embed_dim=args.encoder_embed_dim,
                path=args.encoder_embed_path,
            )
            shared_decoder_embed_tokens = shared_encoder_embed_tokens
            args.share_decoder_input_output_embed = True
        else:
            if args.share_encoder_embeddings:
                shared_encoder_embed_tokens = (
                    build_embedding(
                        task.dict,
                        embed_dim=args.encoder_embed_dim,
                        path=args.encoder_embed_path,
                    )
                )
            if args.share_decoder_embeddings:
                shared_decoder_embed_tokens = (
                    build_embedding(
                        task.dict,
                        embed_dim=args.decoder_embed_dim,
                        path=args.decoder_embed_path,
                    )
                )

        # encoders/decoders for each language
        lang_encoders, lang_decoders = {}, {}

        def get_encoder(lang):
            if lang not in lang_encoders:
                if shared_encoder_embed_tokens is not None:
                    encoder_embed_tokens = shared_encoder_embed_tokens
                else:
                    encoder_embed_tokens = build_embedding(
                        task.dict, args.encoder_embed_dim, args.encoder_embed_path
                    )
                lang_encoders[lang] = TransformerEncoder(args, task.dict, encoder_embed_tokens)
            return lang_encoders[lang]

        def get_decoder(lang):
            if lang not in lang_decoders:
                if shared_decoder_embed_tokens is not None:
                    decoder_embed_tokens = shared_decoder_embed_tokens
                else:
                    decoder_embed_tokens = build_embedding(
                        task.dict, args.decoder_embed_dim, args.decoder_embed_path
                    )
                lang_decoders[lang] = TransformerDecoder(args, task.dict, decoder_embed_tokens)
            return lang_decoders[lang]

        # shared encoders/decoders (if applicable)
        shared_encoder, shared_decoder = None, None
        if args.share_encoders:
            shared_encoder = get_encoder(src_langs[0])
        if args.share_decoders:
            shared_decoder = get_decoder(tgt_langs[0])

        encoders, decoders = OrderedDict(), OrderedDict()
        for lang_pair in task.model_lang_pairs:
            src, tgt = lang_pair.split('-')
            encoders[lang_pair] = shared_encoder if shared_encoder is not None else get_encoder(src)
            decoders[lang_pair] = shared_decoder if shared_decoder is not None else get_decoder(tgt)

        return cls(args, encoders, decoders, task)

    def load_state_dict(self, state_dict, strict=True, args=None):
        models_key = [k for k in state_dict.keys() if k.startswith('models')]
        if len(models_key) != 0:
            state_dict_models = state_dict.copy()
            for k, _ in state_dict.items():
                if not k.startswith('models.'):
                    assert k.startswith('lang_embed_token.') or k.startswith('momentum') or k.startswith('encoder_extra_fc') or k == "queue_ptr" or k.startswith("cross_attention")
                    continue
                assert k.startswith('models.')
                lang_pair = k.split('.')[1]
                if lang_pair not in self.models:
                    del state_dict_models[k]
            # for k, _ in state_dict.items():
            #     for lang_pair in self.models:
            #         if k.startswith('models.{}.encoder'.format(lang_pair)):
            #             new_k = k.replace('models.{}.encoder'.format(lang_pair), 'momentum_encoders.{}'.format(lang_pair))
            #             if new_k not in state_dict:
            #                 state_dict_models[new_k] = state_dict[k]
            if "encoder_extra_fc2.3.weight" not in state_dict_models:
                for k in ["encoder_extra_fc2.{}.weight".format(i) for i in range(3)]:
                    if k in state_dict_models:
                        del state_dict_models[k]
                        del state_dict_models[k.replace('weight', "bias")]
            missing_keys, unexpected_keys = super().load_state_dict(state_dict_models, strict=False)
            # assert unexpected_keys == []
            print("----------------------- unexpected_keys --------------")
            print(unexpected_keys)
            print("----------------------missing_keys-------------------")
            print(missing_keys)
            for key in unexpected_keys:
                assert key.startswith('encoder_extra_fc2')
            for key in missing_keys:
                assert key.startswith('momentum') or key.startswith('encoder_extra_fc') or key == "queue_ptr" or key.startswith("cross_attention")
            if missing_keys != []:
                self.momentum_encoder.load_state_dict(self.models[self.keys[0]].encoder.state_dict())
        else:
            state_dict_subset = state_dict.copy()
            print('Load {} parameters'.format(len(state_dict_subset)))
            if not args.share_encoders:
                encoder_key = [k for k in state_dict_subset.keys() if k.startswith('encoder')]
                for key in encoder_key:
                    del state_dict_subset[key]
            if not args.share_decoders and not args.load_decoders:
                decoder_key = [k for k in state_dict_subset.keys() if k.startswith('decoder')]
                for key in decoder_key:
                    del state_dict_subset[key]
            print('Reserve {} parameters'.format(len(state_dict_subset)))
            if not "encoder.lang_embed_token.weight" in state_dict_subset and hasattr(self, 'lang_embed_token'):
                state_dict_subset["encoder.lang_embed_token.weight"] = self.lang_embed_token.weight
            if not "decoder.lang_embed_token.weight" in state_dict_subset and hasattr(self, 'lang_embed_token'):
                state_dict_subset["decoder.lang_embed_token.weight"] = self.lang_embed_token.weight
            for k, m in self.models.items():
                missing_keys, unexpected_keys = m.load_state_dict(state_dict_subset, False)
                if args.share_encoders and args.share_decoders:
                    assert missing_keys + unexpected_keys == []
                    break

            if hasattr(self, "momentum_encoder"):
                encoder_state = self.models[self.keys[0]].encoder.state_dict()
                encoder_state.pop('lang_embed_token.weight', None)
                self.momentum_encoder.load_state_dict(encoder_state)

class TransformerEncoder(FairseqEncoder):
    """
    Transformer encoder consisting of *args.encoder_layers* layers. Each layer
    is a :class:`TransformerEncoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): encoding dictionary
        embed_tokens (torch.nn.Embedding): input embedding
    """

    def __init__(self, args, dictionary, embed_tokens):
        super().__init__(dictionary)
        self.register_buffer('version', torch.Tensor([3]))

        self.dropout = args.dropout
        self.encoder_layerdrop = args.encoder_layerdrop

        embed_dim = embed_tokens.embedding_dim
        self.padding_idx = embed_tokens.padding_idx
        self.max_source_positions = args.max_source_positions

        self.embed_tokens = embed_tokens

        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)

        self.embed_positions = PositionalEmbedding(
            args.max_source_positions, embed_dim, self.padding_idx,
            learned=args.encoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.layer_wise_attention = getattr(args, 'layer_wise_attention', False)

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerEncoderLayer(args)
            for i in range(args.encoder_layers)
        ])

        if args.encoder_normalize_before:
            self.layer_norm = LayerNorm(embed_dim)
        else:
            self.layer_norm = None
        if getattr(args, 'layernorm_embedding', False):
            self.layernorm_embedding = LayerNorm(embed_dim)
        else:
            self.layernorm_embedding = None

    def forward_embedding(self, src_tokens, lang_tokens):
        # embed tokens and positions
        x = embed = self.embed_scale * self.embed_tokens(src_tokens)
        if self.embed_positions is not None:
            x = embed + self.embed_positions(src_tokens)
        # add lang_embedding
        if lang_tokens is not None and hasattr(self, 'lang_embed_token'):
            x.add_(self.lang_embed_token(src_tokens.clone().fill_(lang_tokens)))
        if self.layernorm_embedding:
            x = self.layernorm_embedding(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        return x, embed

    def forward(self, src_tokens, src_lengths, cls_input=None, lang_tokens=None, return_all_hiddens=False, **unused):
        """
        Args:
            src_tokens (LongTensor): tokens in the source language of shape
                `(batch, src_len)`
            src_lengths (torch.LongTensor): lengths of each source sentence of
                shape `(batch)`
            return_all_hiddens (bool, optional): also return all of the
                intermediate hidden states (default: False).

        Returns:
            namedtuple:
                - **encoder_out** (Tensor): the last encoder layer's output of
                  shape `(src_len, batch, embed_dim)`
                - **encoder_padding_mask** (ByteTensor): the positions of
                  padding elements of shape `(batch, src_len)`
                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
                  of shape `(batch, src_len, embed_dim)`
                - **encoder_states** (List[Tensor]): all intermediate
                  hidden states of shape `(src_len, batch, embed_dim)`.
                  Only populated if *return_all_hiddens* is True.
        """
        if self.layer_wise_attention:
            return_all_hiddens = True

        x, encoder_embedding = self.forward_embedding(src_tokens, lang_tokens)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        # compute padding mask
        encoder_padding_mask = src_tokens.eq(self.padding_idx)
        if not encoder_padding_mask.any():
            encoder_padding_mask = None

        encoder_states = [] if return_all_hiddens else None

        # encoder layers
        for layer in self.layers:
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if not self.training or (dropout_probability > self.encoder_layerdrop):
                x = layer(x, encoder_padding_mask)
                if return_all_hiddens:
                    encoder_states.append(x)

        if self.layer_norm:
            x = self.layer_norm(x)
            if return_all_hiddens:
                encoder_states[-1] = x

        return EncoderOut(
            encoder_out=x,  # T x B x C
            encoder_padding_mask=encoder_padding_mask,  # B x T
            encoder_embedding=encoder_embedding,  # B x T x C
            encoder_states=encoder_states,  # List[T x B x C]
        )

    def reorder_encoder_out(self, encoder_out, new_order):
        """
        Reorder encoder output according to *new_order*.

        Args:
            encoder_out: output from the ``forward()`` method
            new_order (LongTensor): desired order

        Returns:
            *encoder_out* rearranged according to *new_order*
        """
        if encoder_out.encoder_out is not None:
            encoder_out = encoder_out._replace(
                encoder_out=encoder_out.encoder_out.index_select(1, new_order)
            )
        if encoder_out.encoder_padding_mask is not None:
            encoder_out = encoder_out._replace(
                encoder_padding_mask=encoder_out.encoder_padding_mask.index_select(0, new_order)
            )
        if encoder_out.encoder_embedding is not None:
            encoder_out = encoder_out._replace(
                encoder_embedding=encoder_out.encoder_embedding.index_select(0, new_order)
            )
        if encoder_out.encoder_states is not None:
            for idx, state in enumerate(encoder_out.encoder_states):
                encoder_out.encoder_states[idx] = state.index_select(1, new_order)
        return encoder_out

    def max_positions(self):
        """Maximum input length supported by the encoder."""
        if self.embed_positions is None:
            return self.max_source_positions
        return min(self.max_source_positions, self.embed_positions.max_positions())

    def buffered_future_mask(self, tensor):
        dim = tensor.size(0)
        if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device:
            self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
            if self._future_mask.size(0) < dim:
                self._future_mask = torch.triu(utils.fill_with_neg_inf(self._future_mask.resize_(dim, dim)), 1)
        return self._future_mask[:dim, :dim]

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                print('deleting {0}'.format(weights_key))
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)
        for i in range(len(self.layers)):
            # update layer norms
            self.layers[i].upgrade_state_dict_named(state_dict, "{}.layers.{}".format(name, i))

        version_key = '{}.version'.format(name)
        if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) < 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict[version_key] = torch.Tensor([1])
        return state_dict

class TransformerDecoder(FairseqIncrementalDecoder):
    """
    Transformer decoder consisting of *args.decoder_layers* layers. Each layer
    is a :class:`TransformerDecoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): decoding dictionary
        embed_tokens (torch.nn.Embedding): output embedding
        no_encoder_attn (bool, optional): whether to attend to encoder outputs
            (default: False).
    """

    def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
        super().__init__(dictionary)
        self.register_buffer('version', torch.Tensor([3]))

        self.dropout = args.dropout
        self.decoder_layerdrop = args.decoder_layerdrop
        self.share_input_output_embed = args.share_decoder_input_output_embed

        input_embed_dim = embed_tokens.embedding_dim
        embed_dim = args.decoder_embed_dim
        self.output_embed_dim = args.decoder_output_dim

        self.padding_idx = embed_tokens.padding_idx
        self.max_target_positions = args.max_target_positions

        self.embed_tokens = embed_tokens

        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)

        self.project_in_dim = Linear(input_embed_dim, embed_dim, bias=False) if embed_dim != input_embed_dim else None

        self.embed_positions = PositionalEmbedding(
            args.max_target_positions, embed_dim, self.padding_idx,
            learned=args.decoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.cross_self_attention = getattr(args, 'cross_self_attention', False)
        self.layer_wise_attention = getattr(args, 'layer_wise_attention', False)

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerDecoderLayer(args, no_encoder_attn)
            for _ in range(args.decoder_layers)
        ])

        self.adaptive_softmax = None

        self.project_out_dim = Linear(embed_dim, self.output_embed_dim, bias=False) \
            if embed_dim != self.output_embed_dim and not args.tie_adaptive_weights else None

        if args.adaptive_softmax_cutoff is not None:
            self.adaptive_softmax = AdaptiveSoftmax(
                len(dictionary),
                self.output_embed_dim,
                options.eval_str_list(args.adaptive_softmax_cutoff, type=int),
                dropout=args.adaptive_softmax_dropout,
                adaptive_inputs=embed_tokens if args.tie_adaptive_weights else None,
                factor=args.adaptive_softmax_factor,
                tie_proj=args.tie_adaptive_proj,
            )
        elif not self.share_input_output_embed:
            self.embed_out = nn.Parameter(torch.Tensor(len(dictionary), self.output_embed_dim))
            nn.init.normal_(self.embed_out, mean=0, std=self.output_embed_dim ** -0.5)

        if args.decoder_normalize_before and not getattr(args, 'no_decoder_final_norm', False):
            self.layer_norm = LayerNorm(embed_dim)
        else:
            self.layer_norm = None
        if getattr(args, 'layernorm_embedding', False):
            self.layernorm_embedding = LayerNorm(embed_dim)
        else:
            self.layernorm_embedding = None

    def forward(
        self,
        prev_output_tokens,
        encoder_out=None,
        incremental_state=None,
        features_only=False,
        lang_tokens=None,
        **extra_args
    ):
        """
        Args:
            prev_output_tokens (LongTensor): previous decoder outputs of shape
                `(batch, tgt_len)`, for teacher forcing
            encoder_out (optional): output from the encoder, used for
                encoder-side attention
            incremental_state (dict): dictionary used for storing state during
                :ref:`Incremental decoding`
            features_only (bool, optional): only return features without
                applying output layer (default: False).

        Returns:
            tuple:
                - the decoder's output of shape `(batch, tgt_len, vocab)`
                - a dictionary with any model-specific outputs
        """
        x, extra = self.extract_features(
            prev_output_tokens,
            lang_tokens=lang_tokens,
            encoder_out=encoder_out,
            incremental_state=incremental_state,
            **extra_args
        )
        if not features_only:
            x = self.output_layer(x)
        return x, extra

    def extract_features(
        self,
        prev_output_tokens,
        encoder_out=None,
        incremental_state=None,
        full_context_alignment=False,
        alignment_layer=None,
        alignment_heads=None,
        lang_tokens=None,
        **unused,
    ):
        """
        Similar to *forward* but only return features.

        Includes several features from "Jointly Learning to Align and
        Translate with Transformer Models" (Garg et al., EMNLP 2019).

        Args:
            full_context_alignment (bool, optional): don't apply
                auto-regressive mask to self-attention (default: False).
            alignment_layer (int, optional): return mean alignment over
                heads at this layer (default: last layer).
            alignment_heads (int, optional): only average alignment over
                this many heads (default: all heads).

        Returns:
            tuple:
                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
                - a dictionary with any model-specific outputs
        """
        if alignment_layer is None:
            alignment_layer = len(self.layers) - 1

        # embed positions
        positions = self.embed_positions(
            prev_output_tokens,
            incremental_state=incremental_state,
        ) if self.embed_positions is not None else None

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:].clone()
            if positions is not None:
                positions = positions[:, -1:].clone()

        # embed tokens and positions
        x = self.embed_scale * self.embed_tokens(prev_output_tokens)

        if self.project_in_dim is not None:
            x = self.project_in_dim(x)

        if positions is not None:
            x += positions

        # add lang embedding
        if lang_tokens is not None and hasattr(self, 'lang_embed_token'):
            x += self.lang_embed_token(prev_output_tokens.clone().fill_(lang_tokens))

        if self.layernorm_embedding:
            x = self.layernorm_embedding(x)

        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        self_attn_padding_mask = None
        if self.cross_self_attention or prev_output_tokens.eq(self.padding_idx).any():
            self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)

        # decoder layers
        attn = None
        inner_states = [x]
        for idx, layer in enumerate(self.layers):
            encoder_state = None
            if encoder_out is not None:
                if self.layer_wise_attention:
                    encoder_state = encoder_out.encoder_states[idx]
                else:
                    encoder_state = encoder_out.encoder_out

            if incremental_state is None and not full_context_alignment:
                self_attn_mask = self.buffered_future_mask(x)
            else:
                self_attn_mask = None

            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if not self.training or (dropout_probability > self.decoder_layerdrop):
                x, layer_attn = layer(
                    x,
                    encoder_state,
                    encoder_out.encoder_padding_mask if encoder_out is not None else None,
                    incremental_state,
                    self_attn_mask=self_attn_mask,
                    self_attn_padding_mask=self_attn_padding_mask,
                    need_attn=(idx == alignment_layer),
                    need_head_weights=(idx == alignment_layer),
                )
                inner_states.append(x)
                if layer_attn is not None and idx == alignment_layer:
                    attn = layer_attn.float()

        if attn is not None:
            if alignment_heads is not None:
                attn = attn[:alignment_heads]

            # average probabilities over heads
            attn = attn.mean(dim=0)

        if self.layer_norm:
            x = self.layer_norm(x)

        # T x B x C -> B x T x C
        x = x.transpose(0, 1)

        if self.project_out_dim is not None:
            x = self.project_out_dim(x)

        return x, {'attn': attn, 'inner_states': inner_states}

    def output_layer(self, features, **kwargs):
        """Project features to the vocabulary size."""
        if self.adaptive_softmax is None:
            # project back to size of vocabulary
            if self.share_input_output_embed:
                return F.linear(features, self.embed_tokens.weight)
            else:
                return F.linear(features, self.embed_out)
        else:
            return features

    def max_positions(self):
        """Maximum output length supported by the decoder."""
        if self.embed_positions is None:
            return self.max_target_positions
        return min(self.max_target_positions, self.embed_positions.max_positions())

    def buffered_future_mask(self, tensor):
        dim = tensor.size(0)
        if (
            not hasattr(self, '_future_mask')
            or self._future_mask is None
            or self._future_mask.device != tensor.device
            or self._future_mask.size(0) < dim
        ):
            self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
        return self._future_mask[:dim, :dim]

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)

        for i in range(len(self.layers)):
            # update layer norms
            layer_norm_map = {
                '0': 'self_attn_layer_norm',
                '1': 'encoder_attn_layer_norm',
                '2': 'final_layer_norm'
            }
            for old, new in layer_norm_map.items():
                for m in ('weight', 'bias'):
                    k = '{}.layers.{}.layer_norms.{}.{}'.format(name, i, old, m)
                    if k in state_dict:
                        state_dict['{}.layers.{}.{}.{}'.format(name, i, new, m)] = state_dict[k]
                        del state_dict[k]

        version_key = '{}.version'.format(name)
        if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) <= 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict[version_key] = torch.Tensor([1])

        return state_dict

@register_model_architecture('backsum_transformer', 'backsum_transformer')
def base_multilingual_architecture(args):
    base_architecture(args)
    args.share_encoder_embeddings = getattr(args, 'share_encoder_embeddings', False)
    args.share_decoder_embeddings = getattr(args, 'share_decoder_embeddings', False)
    args.share_encoders = getattr(args, 'share_encoders', False)
    args.share_decoders = getattr(args, 'share_decoders', False)

@register_model_architecture('backsum_transformer', 'backsum_transformer_bart_large')
def multilingual_transformer_bart_large(args):
    args.encoder_embed_path = getattr(args, 'encoder_embed_path', None)
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4*1024)
    args.encoder_layers = getattr(args, 'encoder_layers', 12)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
    args.encoder_normalize_before = getattr(args, 'encoder_normalize_before', False)
    args.encoder_learned_pos = getattr(args, 'encoder_learned_pos', True)
    args.decoder_embed_path = getattr(args, 'decoder_embed_path', None)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', args.encoder_embed_dim)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', args.encoder_ffn_embed_dim)
    args.decoder_layers = getattr(args, 'decoder_layers', 12)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
    args.decoder_normalize_before = getattr(args, 'decoder_normalize_before', False)
    args.decoder_learned_pos = getattr(args, 'decoder_learned_pos', True)
    args.attention_dropout = getattr(args, 'attention_dropout', 0.)
    args.relu_dropout = getattr(args, 'relu_dropout', 0.)
    args.dropout = getattr(args, 'dropout', 0.1)
    args.max_target_positions = getattr(args, 'max_target_positions', 1024)
    args.max_source_positions = getattr(args, 'max_source_positions', 1024)
    args.adaptive_softmax_cutoff = getattr(args, 'adaptive_softmax_cutoff', None)
    args.adaptive_softmax_dropout = getattr(args, 'adaptive_softmax_dropout', 0)
    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
    args.share_all_embeddings = getattr(args, 'share_all_embeddings', True)

    args.decoder_output_dim = getattr(args, 'decoder_output_dim', args.decoder_embed_dim)
    args.decoder_input_dim = getattr(args, 'decoder_input_dim', args.decoder_embed_dim)

    args.no_scale_embedding = getattr(args, 'no_scale_embedding', True)
    args.layernorm_embedding = getattr(args, 'layernorm_embedding', True)

    args.activation_fn = getattr(args, 'activation_fn', 'gelu')
    args.pooler_activation_fn = getattr(args, 'pooler_activation_fn', 'tanh')
    args.pooler_dropout = getattr(args, 'pooler_dropout', 0.0)
    base_multilingual_architecture(args)


import torch
import cfg

class UpsampleLayer(torch.nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')

class ConvolutionalLayer(torch.nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):
        super().__init__()

        self.sub_module = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.LeakyReLU(0.1)
        )

    def forward(self, x):
        return self.sub_module(x)

class ResidualLayer(torch.nn.Module):

    def __init__(self, in_channels):
        super().__init__()

        self.sub_module = torch.nn.Sequential(
            ConvolutionalLayer(in_channels, in_channels // 2, 1, 1, 0),
            ConvolutionalLayer(in_channels // 2, in_channels, 3, 1, 1),
        )

    def forward(self, x):
        return x + self.sub_module(x)

class DownsamplingLayer(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()

        self.sub_module = torch.nn.Sequential(
            ConvolutionalLayer(in_channels, out_channels, 3, 2, 1)
        )

    def forward(self, x):
        return self.sub_module(x)

class ConvolutionalSet(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()

        self.sub_module = torch.nn.Sequential(
            ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),
            ConvolutionalLayer(out_channels, in_channels, 3, 1, 1),

            ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),
            ConvolutionalLayer(out_channels, in_channels, 3, 1, 1),

            ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),
        )

    def forward(self, x):
        return self.sub_module(x)

class DarkNet(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.trunk_52 = torch.nn.Sequential(
            ConvolutionalLayer(3, 32, 3, 1, 1),
            ConvolutionalLayer(32, 64, 3, 2, 1),

            ResidualLayer(64),
            DownsamplingLayer(64, 128),

            ResidualLayer(128),
            ResidualLayer(128),
            DownsamplingLayer(128, 256),

            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
        )
        self.trunk_26 = torch.nn.Sequential(
            DownsamplingLayer(256, 512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
        )
        self.trunk_13 = torch.nn.Sequential(
            DownsamplingLayer(512, 1024),
            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024)
        )

    def forward(self, x):
        h_52 = self.trunk_52(x)
        h_26 = self.trunk_26(h_52)
        h_13 = self.trunk_13(h_26)
        return h_52,h_26,h_13

class MyNet(torch.nn.Module):
    def __init__(self,cls_num):
        super().__init__()
        self.trunk_52 = torch.nn.Sequential(
            ConvolutionalLayer(3, 32, 3, 1, 1),
            ConvolutionalLayer(32, 64, 3, 2, 1),

            ResidualLayer(64),
            DownsamplingLayer(64, 128),

            ResidualLayer(128),
            ResidualLayer(128),
            DownsamplingLayer(128, 256),

            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
            ResidualLayer(256),
        )
        self.trunk_26 = torch.nn.Sequential(
            DownsamplingLayer(256, 512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
            ResidualLayer(512),
        )
        self.trunk_13 = torch.nn.Sequential(
            DownsamplingLayer(512, 1024),
            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024),
            ResidualLayer(1024)
        )
        self.convset_13 = torch.nn.Sequential(
            ConvolutionalSet(1024, 512)
        )
        self.detetion_13 = torch.nn.Sequential(
            ConvolutionalLayer(512, 1024, 3, 1, 1),
            torch.nn.Conv2d(1024, 3*(cls_num+5), 1, 1, 0)
        )
        self.up_26 = torch.nn.Sequential(
            ConvolutionalLayer(512, 256, 1, 1, 0),
            UpsampleLayer()
        )
        self.convset_26 = torch.nn.Sequential(
            ConvolutionalSet(768, 256)
        )
        self.detetion_26 = torch.nn.Sequential(
            ConvolutionalLayer(256, 512, 3, 1, 1),
            torch.nn.Conv2d(512, 3*(cls_num+5), 1, 1, 0)
        )
        self.up_52 = torch.nn.Sequential(
            ConvolutionalLayer(256, 128, 1, 1, 0),
            UpsampleLayer()
        )
        self.convset_52 = torch.nn.Sequential(
            ConvolutionalSet(384, 128)
        )
        self.detetion_52 = torch.nn.Sequential(
            ConvolutionalLayer(128, 256, 3, 1, 1),
            torch.nn.Conv2d(256, 3*(cls_num+5), 1, 1, 0)
        )

    def forward(self,x):
        h_52 = self.trunk_52(x)
        h_26 = self.trunk_26(h_52)
        h_13 = self.trunk_13(h_26)
        convset_out_13 = self.convset_13(h_13)
        detetion_out_13 = self.detetion_13(convset_out_13)

        up_out_26 = self.up_26(convset_out_13)
        route_out_26 = torch.cat((up_out_26, h_26), dim=1)
        convset_out_26 = self.convset_26(route_out_26)
        detetion_out_26 = self.detetion_26(convset_out_26)

        up_out_52 = self.up_52(convset_out_26)
        # print("up_out_52:",up_out_52.shape)
        # print("h_52:",h_52.shape)
        route_out_52 = torch.cat((up_out_52, h_52), dim=1)
        convset_out_52 = self.convset_52(route_out_52)
        detetion_out_52 = self.detetion_52(convset_out_52)
        return detetion_out_13, detetion_out_26, detetion_out_52

if __name__ == '__main__':
    trunk = MyNet()

    x = torch.Tensor(2, 3, 416, 416)

    y_13, y_26, y_52 = trunk(x)
    print(y_13.shape)
    print(y_26.shape)
    print(y_52.shape)
    # print(y_13.view(-1, 3, 5, 13, 13).shape)


import logging
from datetime import datetime

# 第三方 import
from croniter import croniter
from django.conf import settings
from django.contrib.auth.models import User
from django.db.models import Q
from django.utils import timezone

# 项目内 import
from apps.codeproj import core, models
from apps.codeproj.core import ScanSchemeManager
from apps.job.models import Job
from codedog.celery import celery_app
from util import errcode
from util.operationrecord import OperationRecordHandler
from util.webclients import AnalyseClient

logger = logging.getLogger(__name__)

def _now():
    now = timezone.make_aware(datetime.now(), timezone.get_current_timezone())
    return now

@celery_app.task
def sync_pkg_rule_map():
    """启动AnalysisServer规则包与规则关系同步任务
    """
    logger.info("Start to drop expired table")
    result = AnalyseClient().api("sync_pkg_rule_map", data=None)
    logger.info("drop expired table result: %s" % result)

@celery_app.task
def check_project_active():
    """检查项目是否运行正常

    1. 检查项目近一个月是否有启动扫描，如果没有则标记并发送邮件通知给分支项目创建同学
    2. 检查定时项目近10次扫描是否存在SCM相关异常，如果有，则关闭定时任务
    """
    logger.info("开始检查项目是否运行正常")
    logger.info("1. 开始检查项目近一个月是否有启动扫描")
    projects = models.Project.objects.filter(status=models.Project.StatusEnum.ACTIVE)
    codedog_user, created = User.objects.get_or_create(username=settings.DEFAULT_USERNAME)
    current_time = _now()
    for project in projects:
        # 长期未成功执行扫描
        last_success_job = Job.objects.filter(project=project).order_by("-create_time").first()
        last_active_time = last_success_job.create_time if last_success_job else project.created_time
        if current_time - last_active_time > settings.PROJECT_DISACTIVE_TIMEOUT:
            logger.info("[Project: %d][Creator: %s] 项目失活，代码库地址: %s，上一次活跃时间为: %s" % (
                project.id, project.creator, project.scm_url, last_active_time))
            models.Project.objects.filter(id=project.id).update(
                status=models.Project.StatusEnum.DISACTIVE, modifier=codedog_user, modified_time=current_time)
            OperationRecordHandler.add_project_operation_record(
                project, "项目失活", codedog_user, message="项目于%s自动标记为失活" % current_time
            )
    logger.info("2. 开始检查开启定时扫描的项目近10次扫描是否存在SCM相关异常")
    project_schedules = models.ScanSchedule.objects.filter(
        enabled=True,
        project__status=models.Project.StatusEnum.ACTIVE,
        project__scan_scheme__status=models.ScanScheme.StatusEnum.ACTIVE
    ).exclude(scan_sched=None)

    for project_sched in project_schedules:
        project = project_sched.project
        result_codes = Job.objects.filter(project=project).order_by(
            "-create_time").values_list("result_code", flat=True)[:10]
        result_codes = list(set(result_codes))
        if len(result_codes) == 1 and errcode.is_scm_error(result_codes[0]):
            msg = "项目近10次扫描错误码均为%s，关闭定时扫描任务" % result_codes[0]
            logger.warning("[Project: %s] %s" % (project.id, msg))
            OperationRecordHandler.add_project_operation_record(
                project, "项目更新配置", codedog_user, message=msg
            )
            project_sched.enabled = False
            project_sched.save()

@celery_app.task
def push_scanscheme_template(scheme_ids, ref_scheme_id, username, **kwargs):
    """推送扫描方案模板到相关子方案
    """
    user, _ = User.objects.get_or_create(username=username)
    ref_scheme = models.ScanScheme.objects.get(id=ref_scheme_id)
    op_message = ScanSchemeManager.get_sync_message(ref_scheme, **kwargs)
    scan_scheme_message = []
    for scheme_id in scheme_ids:
        scan_scheme = models.ScanScheme.objects.get(id=scheme_id)
        ScanSchemeManager.sync_with_ref_scheme(scan_scheme=scan_scheme,
                                               ref_scheme=ref_scheme,
                                               user=user,
                                               **kwargs)
        OperationRecordHandler.add_scanscheme_operation_record(scan_scheme, "方案模板同步", username, op_message)
        scan_scheme_message.append("%s" % scan_scheme)
    ref_scheme_message = "%s，同步扫描方案：%s" % (op_message, "、".join(scan_scheme_message))
    OperationRecordHandler.add_scanscheme_operation_record(ref_scheme, "方案模板同步", username, ref_scheme_message)

@celery_app.task
def handle_scheduled_projects():
    """根据project的配置定时创建扫描
    筛选范围：
    1. 配置了定时任务
    2. 项目为活跃状态
    3. 扫描方案为活跃状态
    """
    # Todo


from flask import Flask
from flask_login import LoginManager
from flask_cors import CORS

login_manager = LoginManager()

def create_app():
    application = Flask(__name__)
    application.secret_key = b'secret_to_be_changed'

    CORS(application,
         origins=['http://localhost:3000',
                  'http://weathrio-dev.us-east-1.elasticbeanstalk.com/'],
         supports_credentials=True)
    login_manager.init_app(application)

    return application


#@title utils for image cropping based on text coordinates
import copy
import cv2
import numpy as np
import imutils
from imutils.perspective import four_point_transform

def get_rotate_crop_image(img, points):
    '''
    img_height, img_width = img.shape[0:2]
    left = int(np.min(points[:, 0]))
    right = int(np.max(points[:, 0]))
    top = int(np.min(points[:, 1]))
    bottom = int(np.max(points[:, 1]))
    img_crop = img[top:bottom, left:right, :].copy()
    points[:, 0] = points[:, 0] - left
    points[:, 1] = points[:, 1] - top
    '''
    img_crop_width = int(
        max(
            np.linalg.norm(points[0] - points[1]),
            np.linalg.norm(points[2] - points[3])))
    img_crop_height = int(
        max(
            np.linalg.norm(points[0] - points[3]),
            np.linalg.norm(points[1] - points[2])))
    pts_std = np.float32([[0, 0], [img_crop_width, 0],
                          [img_crop_width, img_crop_height],
                          [0, img_crop_height]])
    M = cv2.getPerspectiveTransform(points, pts_std)
    dst_img = cv2.warpPerspective(
        img,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE,
        flags=cv2.INTER_CUBIC)
    # dst_img_height, dst_img_width = dst_img.shape[0:2]
    # if dst_img_height * 1.0 / dst_img_width >= 1.5:
    #     dst_img = np.rot90(dst_img)
    return dst_img

def check_vertical_image(img_crop):
    """
    If vertical image rotate 90 degree clockwise
    """
    height, width, channel = img_crop.shape

    ## (width * 1.5) instead of just width to avoid rotation on single character
    if height > (width * 1.3):
        try:
            # print("line image height greater than width: Rotating word -- clockwise 90 degree")
            img_crop = imutils.rotate_bound(img_crop, angle=90) # rotate word if vertical
            # height, width = width, height # swap values
        except Exception as e:
            pass
    return img_crop

def crop_img_boxes(img, dt_boxes, rotate_vertical=True, use_angle_cls=True, use_four_point_transform=False, obj_paddleocr_text_classifier=None):

    img_crop_list = list()
    angle_list = []
    # print(type(img_crop_list))

    for bno in range(len(dt_boxes)):
        tmp_box = copy.deepcopy(dt_boxes[bno])
        # print(tmp_box)
        # print(np.array(tmp_box).shape)
        # print(len(np.array(tmp_box).shape))

        if len(np.array(tmp_box).shape) == 1:
            ## rectangle
            x1, x2, y1, y2 = tmp_box
            img_crop = img[y1:y2, x1:x2] # .copy()
            # print(img_crop.shape)

            if rotate_vertical:
                img_crop = check_vertical_image(img_crop)

            img_crop_list.append(img_crop)
            # print(type(img_crop_list[0]))
        else:
            if use_four_point_transform:
                img_crop = four_point_transform(img, tmp_box)
            else:
                img_crop = get_rotate_crop_image(img, tmp_box)

            if rotate_vertical:
                ## check is vertical image -- if vertical image rotate 90 degree
                img_crop = check_vertical_image(img_crop)

            img_crop_list.append(img_crop)

    if use_angle_cls:
        # img_crop_list, angle_list, elapse = paddleocr_engine.text_classifier(
        #     img_crop_list
        # )

        if obj_paddleocr_text_classifier:
            img_crop_list, angle_list, elapse = obj_paddleocr_text_classifier(
                img_crop_list
            )
            print("cls num  : {}, elapse : {}".format(
                len(img_crop_list), elapse))
        else:
            print(f"Warning... obj_paddleocr_text_classifier not set.. skipping angle classification")

    return img_crop_list, angle_list


#!/usr/bin/env python3
import codecs
import argparse
import json

def togeo(file):
    out = ".".join([file.rsplit(".", 1)[0], "geo"])
    with codecs.open(file, "r", "ascii") as fin:
        d = json.load(fin)[0]

    with codecs.open(out, "w", "ascii") as fou:
        fou.write(
            "$MeshFormat\n4.1 0 {0}\n$EndMeshFormat\n$Nodes\n 1 {1} {2} {3}\n2 1 0 {1}\n".format(
                8,
                len(d["vertices"]),
                d["vertices"][0]["id"] + 1,
                d["vertices"][-1]["id"] + 1,
            )
        )
        for i in d["vertices"]:
            fou.write("{}\n".format(i["id"] + 1))
        for i in d["vertices"]:
            fou.write("{:.5f} {:.5f} {:.5f}\n".format(*i["value"]))
        fou.write(
            "$EndNodes\n$Elements\n1 {0} 1 {0}\n2 1 3 {0}\n".format(len(d["elements"]))
        )
        for i, j in enumerate(d["elements"]):
            fou.write("{} {} {} {} {}\n".format(i + 1, *[k + 1 for k in j]))
        fou.write("$EndElements\n")

def main():
    p = argparse.ArgumentParser(prog="togeo")
    p.add_argument("file", nargs="+")
    a = p.parse_args()
    for i in a.file:
        togeo(i)

if __name__ == "__main__":
    main()


#!/usr/bin/env python

'''
./evolutionModel_V2.py $family $workfolder $logfolder
./evolutionModel_V2.py MIPF0000632 Test LOG
'''
import os
import os.path
import sys
import subprocess
import random
import re

from Bio import SeqIO
from Bio import AlignIO
from Bio.Align.Applications import ClustalOmegaCommandline
from deap import algorithms
from deap import base
from deap import creator
from deap import tools
from ete3 import NCBITaxa

# Load databases
ncbi = NCBITaxa()
if os.path.isfile("./taxdump.tar.gz"):
    print("Database loaded")
else:
    ncbi.update_taxonomy_database()

# In this case I'd like to maximise the score (fitness function)
creator.create("FitnessMax", base.Fitness, weights=(1.0,))
creator.create("Individual", list, fitness=creator.FitnessMax)
toolbox = base.Toolbox()
# Variable codification. Meaning and limits:  <10-11-21, cavelandiah> #
MODE_MIN, MODE_MAX = 0, 1  # Mode: normal = 0, high = 1
CLADE_MIN, CLADE_MAX = 0, 3  # Clade = Metazoa =0, Vertebrata=1, Mammalia=2, Primates=3
CUTOFF_MIN, CUTOFF_MAX = 0, 4  # 101=0,100=1,90=2,80=3,70=4
IND_SIZE = 1  # Number of repetitions
# Register variables = alias, function_to_alias
toolbox.register("attr_int1", random.randint, MODE_MIN, MODE_MAX)
toolbox.register("attr_int2", random.randint, CLADE_MIN, CLADE_MAX)
toolbox.register("attr_int3", random.randint, CUTOFF_MIN, CUTOFF_MAX)
# Build the individual as = [ int, int, int]
toolbox.register("individual", tools.initCycle, creator.Individual, (toolbox.attr_int1, toolbox.attr_int2, toolbox.attr_int3), n=IND_SIZE)
# Population
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

## Def
def index_taxonomy(file_taxonomy):
    tax_dict = {}
    for line in open(file_taxonomy, "r"):
        line = line.strip()
        if line.startswith("#"):
            continue
        key = line.split()[-1]  # 6085
        span = line.split()[-2]  # Metazoa;Bilateria;Deuterostoma;Hemichordata;
        tax_dict[key] = span
    return tax_dict

def index_quality(file_quality):
    quality_dict = {}
    for line in open(file_quality, "r"):
        line = line.strip()
        quality_dict[line] = 1
    return quality_dict

def index_families(file_fam):
    fam_dict = dict()
    for line in open(file_fam, "r"):
        line = line.strip()
        fields = line.split()
        fam_dict.setdefault(fields[0], []).append(fields[2])
    return fam_dict

def translate(individual):
    # Translate the meaning of the individual into parameters of evaluation program
    translation = []
    for i in range(len(individual)):
        if i == 0:  # mode
            if individual[i] == 0:
                translation.append("normal")
            elif individual[i] == 1:
                translation.append("high")
        elif i == 1:  #clade
            if individual[i] == 0:
                translation.append("Metazoa")
            elif individual[i] == 1:
                translation.append("Vertebrata")
            elif individual[i] == 2:
                translation.append("Mammalia")
            elif individual[i] == 3:
                translation.append("Primates")
        elif i == 2:  #cutoff
            if individual[i] == 0:
                translation.append("60")
            elif individual[i] == 1:
                translation.append("50")
            elif individual[i] == 2:
                translation.append("90")
            elif individual[i] == 3:
                translation.append("80")
            elif individual[i] == 4:
                translation.append("70")
    return translation

def evaluate_if_no_structure(str):
    # Evaluate if SS_consensus has structure
    str_len = len(str)
    count_gaps = 0
    for i in str:
        match_gap = re.match(r'\.', i)
        if match_gap:
            count_gaps = count_gaps + 1
    if str_len == count_gaps:
        return 1
    return 0

def file_len(fname):
    # Account number of species + number of sequences
    species = dict()
    lines_number = 0
    with open(fname, 'r') as infile:
        for line in infile:
            line = line.strip()
            if re.search(r'^#', line):
                continue
            else:
                seq_header = line.split()[0]
                spe = seq_header.split('-')[0]
                species[spe] = 1
                lines_number = lines_number + 1
    number_species = len(species.keys())
    score = number_species + lines_number
    return score

def evaluate_energy(stofile):
    # Evaluate folding energy of the sto file
    first = ["RNAalifold", "--noPS", "-q", "-r", stofile]  # -r = RIBOSUM scoring
    f = subprocess.check_output(first)
    if not isinstance(f, str):
        f = f.decode("utf-8")
    energy = f.split(' (')[-1]
    energy = energy.split('=')[0]
    energy = float(energy)*-1  # Positive to valid energies
    return energy

def evaluate_structure(structure):
    pattern = "\)\.*\("
    parts_str = len(re.findall(pattern, structure))
    score = int(parts_str) * -10
    return score

def evaluate_clade(clade):
    # Reward some clades in order to have more diversity:  <10-11-21, cavelandiah> #
    score = None
    if clade == "Metazoa":
        score = 10
    elif clade == "Vertebrata":
        score = 5
    elif clade == "Mammalia":
        score = 2.5
    elif clade == "Primates":
        score = 1.25
    else:
        score = 0
    return score

def stockholm_evaluation(stoFile, clade):
    if not os.path.isfile(stoFile):
        return -100
    number_of_align_sequences = file_len(stoFile)
    if number_of_align_sequences == 0:
        return -100
    align = AlignIO.read(stoFile, "stockholm")
    structure = align.column_annotations['secondary_structure']  # Obtain SStr
    eval_empty = evaluate_if_no_structure(structure)
    if eval_empty == 1:
        return -100
    else:
        clade_contr = evaluate_clade(clade) # weight based on clade
        folding_energy = evaluate_energy(stoFile)
        mirna_folding = evaluate_structure(structure) # Negative value if found additional stems in str
        score = float(clade_contr) + (float(number_of_align_sequences)*float(folding_energy)) + float(mirna_folding)
        return score

def generate_matrix(family, fasta, folder):
    outmatrix = folder+"/"+family+".distmat"
    clustalomega_cline = ClustalOmegaCommandline(infile=fasta, percentid=True, distmat_full=True, distmat_out=outmatrix, verbose=False, auto=False, force=True)
    clustalomega_cline()
    with open(outmatrix, 'r') as fin:
        data = fin.read().splitlines(True)
    with open(outmatrix, 'w') as fout:
        fout.writelines(data[1:])
    return outmatrix

def check(list1, val):
    count_no = 0
    n = len(list1) - 1 # n-1 max
    for x in list1:
        if ((float(x) < float(val)) or (float(x) == float(100))) and (float(x) != float(101)):
            count_no = count_no + 1
    count_no_prop = count_no / n
    if float(count_no_prop) > float(0.7):
        return 0
    else:
        return 1

def generate_valid_identity(matrix_file, cutoff):
    list = dict()
    with open(matrix_file, 'r') as infile:
        count = 0
        for line in infile:
            line = line.strip()
            ids = line.split()[0]
            values = line.split()[1:]
            values[count] = float(200) # Change diagonal to 200.0
            values = [float(i) for i in values] #Convert int -> float
            if (check(values, float(cutoff)) == 1):
                list[ids] = 1
            else:
                continue
            count = count + 1
    return list

def build_fasta_file(variables, family, sequences, taxonomy, quality, fam_relation, folder):
    mode = variables[0]
    clade = variables[1]
    cutoff = variables[2]
    seq_dict = {}
    fasta_seq = SeqIO.parse(open(sequences), 'fasta')
    ids_family = fam_relation[family]
    out_subset = open(folder+"/"+family+"_subset.fa","w")
    out_subset_name = folder+"/"+family+"_subset.fa"
    for fasta in fasta_seq:
        name, seq = fasta.description, fasta.seq
        header = name.split()
        id_mirbase = header[1]
        if id_mirbase in ids_family:
            out_subset.write(">"+str(name)+"\n"+str(seq)+"\n")
    out_subset.close()
    matrix_id = generate_matrix(family, out_subset_name, folder)
    valid_seqs_id = generate_valid_identity(matrix_id, cutoff)
    fasta_seq2 = SeqIO.parse(open(sequences), 'fasta')
    for fasta in fasta_seq2:
        vector_selection = [0, 0, 0]
        name, seq = fasta.description, fasta.seq
        header = name.split()
        infer_specie = header[2]+" "+header[3]
        id_mirbase = header[1]
        if id_mirbase in ids_family:
            tax_id = ncbi.get_name_translator([infer_specie])
            tax_id = tax_id.get(infer_specie)[0]
            if str(tax_id) in taxonomy.keys():
                lineage = taxonomy.get(str(tax_id))
                if clade.lower() in lineage.lower():
                    vector_selection[0] = 1
        if id_mirbase in quality:
            vector_selection[1] = 1
        id_name = header[0]
        if id_name in valid_seqs_id:
            vector_selection[2] = 1
        vector_string = ''.join(str(e) for e in vector_selection)
        seq_dict.setdefault(vector_string, []).append(id_mirbase)
    # Print selected [1, 1/0, 1 ]
    mode_numb = 0
    if mode == "high":
        mode_numb = 1
    subset = [1, mode_numb, 1]
    string_subset = "".join(str(f) for f in subset)
    out_name = "selected.fa"
    out_selected = open(out_name,'w')
    if string_subset in seq_dict:
        selected = seq_dict[string_subset]
    else:
        return out_name
    fasta_seq2 = SeqIO.parse(open(sequences), 'fasta')
    for fasta in fasta_seq2:
        name, seq = fasta.description, fasta.seq
        header = name.split()
        id_mirbase = header[1]
        if id_mirbase in selected:
            out_selected.write(">"+str(name)+"\n"+str(seq)+"\n")
    out_selected.close()
    return out_name

def doalifold(alnfile,outdir,short):
    try:
        if os.path.isfile(outdir+"/"+short+".stk"):
            output = outdir+"/"+short+".stk"
            return output
        else:
            first = ["RNAalifold", "--noPS", "-q", "-r", "--aln-stk="+short, alnfile]  # -r = RIBOSUM scoring
            f = subprocess.check_output(first)
            output = outdir+"/"+short+".stk"
            return output
    except Exception:
       sys.exit()

def evaluate(family, output_folder, output_folder_complete, logfolder, taxonomy, quality, mapping_file, individual):
    result = None
    current = os.getcwd()
    values = translate(individual)
    out_file = str(family)+"_"+str(values[0])+"_"+str(values[1])+"_"+str(values[2])+".sto"
    short = str(family)+"_"+str(values[0])+"_"+str(values[1])+"_"+str(values[2])
    fasta_file_subset = build_fasta_file(values, family, hairpin_fasta, taxonomy, quality, mapping_file, current)
    if os.path.getsize(fasta_file_subset) > 0:
        clustalomega_cline = ClustalOmegaCommandline(infile=fasta_file_subset, outfile=out_file, outfmt='st', verbose=False, auto=False, force=True)
        clustalomega_cline()
    else:
        score = -1000
        return score,
    if os.path.isfile(out_file):
        result = doalifold(out_file, current,short)
        score = stockholm_evaluation(result, values[1])
    else:
        score = -1000
    return score,

def mutVector(individual, indp, MODE_MIN, MODE_MAX, CLADE_MIN, CLADE_MAX, CUTOFF_MIN, CUTOFF_MAX):
    for i in range(len(individual)):
        # The mutation threshold is the same as defined
        if random.random() < indp:
            if i == 0:  # Mode
                individual[i] = random.randint(MODE_MIN, MODE_MAX)
            elif i == 1:  # Clade
                individual[i] = random.randint(CLADE_MIN, CLADE_MAX)
            elif i == 2:  # Cutoff
                individual[i] = random.randint(CUTOFF_MIN, CUTOFF_MAX)
    return individual,

def crossVector(child1, child2):
    new_child1 = [child1[0], child2[1], child1[2]]
    new_child2 = [child2[0], child1[1], child2[2]]
    return new_child1, new_child2

def analyse_winners(vector):
    limit = 10
    if len(vector) < limit:
        return 0
    elif len(vector) == limit:
        lista = dict()
        for i in range(len(vector)):
            concatenate = ",".join(vector[i])
            lista[concatenate] = 1
        uniq_set = len(lista.keys())
        if (uniq_set > 1):
            return 0
        elif (uniq_set == 1):
            return 2
        elif (uniq_set < 1):
            return 0
    elif len(vector) > limit:
        vector = vector[-limit:] # restrict to last 5 elements
        lista = dict()
        for i in range(len(vector)):
            concatenate = ",".join(vector[i])
            lista[concatenate] = 1
        uniq_set = len(lista.keys())
        if (uniq_set > 1):
            return 0
        elif (uniq_set == 1):
            return 2
        elif (uniq_set < 1):
            return 0

## Init
family = sys.argv[1]
output_folder_complete = sys.argv[2]
logfolder = sys.argv[3]
output_folder = output_folder_complete.split('/')[-1] #Last part of output folder
taxonomy = index_taxonomy("organisms.txt")
quality = index_quality("miRNA_high_v22.list")
mapping_file = index_families("mapping_file_original_cleaned_mirbase_v22_families.txt")
hairpin_fasta = "hairpin-metazoa.fa"

toolbox.register("mate", crossVector)
toolbox.register("mutate", mutVector)  # , (toolbox.attr_int1, toolbox.attr_int2, toolbox.attr_int3))
toolbox.register("select", tools.selRoulette)
toolbox.register("evaluate", evaluate, family, output_folder, output_folder_complete, logfolder, taxonomy, quality, mapping_file)

# Population and start fitness
pop = toolbox.population(n=5)
fitnesses = list(map(toolbox.evaluate, pop))
for ind, fit in zip(pop, fitnesses):
    ind.fitness.values = fit

# Step switches
g=0
switch=0
# Cross child:  <10-11-21, cavelandiah> #
CXBP=0.9
# Mutation in individual:  <10-11-21, cavelandiah> #
MUTPB=0.2

# All fitness for the population:  <10-11-21, cavelandiah> #
fits = [ind.fitness.values[0] for ind in pop]

logs = open(logfolder+"/"+family+"_log.txt", 'a')
selected_winner = []
#while g < 100: #max(fits) < 8:
# Iteration steps, until the same winning is conserved over 20 rounds:  <10-11-21, cavelandiah> #
while switch < 1 and g < 20:
    g = g + 1
    # print("-- Generation %i --" % g)
    # Select the next generation individuals
    offspring = toolbox.select(pop, len(pop))
    # Clone the selected individuals
    offspring = list(map(toolbox.clone, offspring))

#    #Apply crossover on the offspring
    for child1, child2 in zip(offspring[::2], offspring[1::2]):
        if random.random() < CXBP:
            # P(cross both childs < 0.7)
            toolbox.mate(child1, child2)
            del child1.fitness.values
            #del child2.fitness.values

    # Apply mutation on the offsprint
    for mutant in offspring:
        if random.random() < MUTPB:
            # The P(mutate_individual < 0.1) and each site 0.5
            toolbox.mutate(mutant, 0.5, MODE_MIN, MODE_MAX, CLADE_MIN, CLADE_MAX, CUTOFF_MIN, CUTOFF_MAX)
            del mutant.fitness.values

    # Evaluate the individuals with an invalid fitness
    invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
    fitnesses = toolbox.map(toolbox.evaluate, invalid_ind)
    for ind, fit in zip(invalid_ind, fitnesses):
        ind.fitness.values = fit

    # The population is entirely replaced by the offspring
    pop[:] = offspring

    # Gather all the fitnesses in one list and print the stats
    fits = [ind.fitness.values[0] for ind in pop]
    # STATS of generations:  <10-11-21, cavelandiah> #
    maximum = max(fits)
    winner = fits.index(maximum)
    winnerTr = translate(pop[winner])
    newline = ",".join(winnerTr)
    print(str(g)+" "+str(max(fits))+" "+str(newline)+" "+str(fits))
    logs.write(str(g)+" "+str(newline)+" "+str(maximum)+"\n")
    selected_winner.append(winnerTr)
    switch = analyse_winners(selected_winner)
    if maximum < 0:
        logs.write("No viable alignment")


import torch
import torch.nn as nn
import torch.nn.functional as F

class ResNeXtBottleneck(nn.Module):
    def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32, dilate=1):
        super(ResNeXtBottleneck, self).__init__()
        D = out_channels // 2
        self.conv_reduce = nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False)
        self.conv_conv = nn.Conv2d(D, D, kernel_size=3, stride=stride, padding=dilate, dilation=dilate,
                                   groups=cardinality,
                                   bias=False)
        self.conv_expand = nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False)

    def forward(self, x):
        if isinstance(x, tuple):
            bottleneck = self.conv_reduce.forward(torch.cat([x[0], x[1]], 1))
            bottleneck = F.relu(bottleneck, inplace=True)
            bottleneck = self.conv_conv.forward(bottleneck)
            bottleneck = F.relu(bottleneck, inplace=True)
            bottleneck = self.conv_expand.forward(bottleneck)
            return x[0] + bottleneck, x[1]
        else:
            bottleneck = self.conv_reduce.forward(x)

            bottleneck = F.relu(bottleneck, inplace=True)
            bottleneck = self.conv_conv.forward(bottleneck)
            bottleneck = F.relu(bottleneck, inplace=True)
            bottleneck = self.conv_expand.forward(bottleneck)
            return x + bottleneck

class DResNeXtBottleneck(nn.Module):
    """
    RexNeXt bottleneck type C (https://github.com/facebookresearch/ResNeXt/blob/master/models/resnext.lua)
    """

    def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32):
        """ Constructor
        Args:
            in_channels: input channel dimensionality
            out_channels: output channel dimensionality
            stride: conv stride. Replaces pooling layer.
            cardinality: num of convolution groups.
        """
        super(DResNeXtBottleneck, self).__init__()
        D = out_channels // 2
        self.conv_reduce = nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False)
        self.conv_conv = nn.Conv2d(D, D, kernel_size=3, stride=stride, padding=1, groups=cardinality, bias=False)
        self.conv_expand = nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.shortcut = nn.Sequential()
        if in_channels != out_channels:
            self.shortcut.add_module('shortcut_conv',
                                     nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0,
                                               bias=False))

    def forward(self, x):
        bottleneck = self.conv_reduce.forward(x)
        bottleneck = F.leaky_relu(bottleneck, 0.2, True)
        bottleneck = self.conv_conv.forward(bottleneck)
        bottleneck = F.leaky_relu(bottleneck, 0.2, True)
        bottleneck = self.conv_expand.forward(bottleneck)
        residual = self.shortcut.forward(x)
        return residual + bottleneck

class Tunnel(nn.Module):
    def __init__(self, len=1, *args):
        super(Tunnel, self).__init__()

        tunnel = [DResNeXtBottleneck(*args) for _ in range(len)]
        self.tunnel = nn.Sequential(*tunnel)

    def forward(self, x):
        return self.tunnel(x)

class DilateTunnel(nn.Module):
    def __init__(self, depth=4):
        super(DilateTunnel, self).__init__()

        tunnel = [ResNeXtBottleneck(dilate=1) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(dilate=2) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(dilate=4) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(dilate=8) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(dilate=1) for _ in range(14)]

        self.tunnel = nn.Sequential(*tunnel)

    def forward(self, x):
        return self.tunnel(x)

class Pyramid(nn.Module):
    def __init__(self):
        super(Pyramid, self).__init__()
        self.entrance1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=7, stride=1, padding=3),
                                       nn.ReLU(inplace=True))
        self.entrance2 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
                                       nn.ReLU(inplace=True))
        self.entrance3 = nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1)

        self.up3 = nn.Sequential(nn.Conv2d(256, 32 * 4, 3, 1, 1, bias=False),
                                 nn.PixelShuffle(2),
                                 nn.ReLU(inplace=True))
        self.con2 = nn.Conv2d(160 + 32, 32, 5, 1, 2, bias=False)

        self.up2 = nn.Sequential(nn.Conv2d(32, 32 * 4, 3, 1, 1, bias=False),
                                 nn.PixelShuffle(2),
                                 nn.ReLU(inplace=True))
        self.con1 = nn.Conv2d(96 + 32, 32, 5, 1, 2, bias=False)

        tunnel = [DResNeXtBottleneck() for _ in range(30)]
        self.tunnel3 = nn.Sequential(*tunnel)

        depth = 3
        tunnel = [ResNeXtBottleneck(64, 32, cardinality=8, dilate=1) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=2) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=4) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=2),
                   ResNeXtBottleneck(64, 32, cardinality=8, dilate=1)]
        self.tunnel2 = nn.Sequential(*tunnel)

        tunnel = [ResNeXtBottleneck(64, 32, cardinality=8, dilate=1) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=2) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=4) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=8) for _ in range(depth)]
        tunnel += [ResNeXtBottleneck(64, 32, cardinality=8, dilate=2),
                   ResNeXtBottleneck(64, 32, cardinality=8, dilate=1)]
        self.tunnel1 = nn.Sequential(*tunnel)

        self.exit = nn.Conv2d(32, 3, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        lv1 = self.entrance1(x)
        lv2 = self.entrance2(lv1)
        lv3 = self.entrance3(lv2)  # need discussion

        lv3_up = self.up3(lv3)
        lv3 = self.up3(self.tunnel3(lv3))

        lv2_up = self.up2(lv3)
        lv2 = self.tunnel2((self.con2(torch.cat([lv2, lv3, lv3_up], 1)), lv3))[0] + lv3
        lv2 = self.up2(lv2)

        lv1 = self.tunnel1((self.con1(torch.cat([lv1, lv2, lv2_up], 1)), lv2))[0] + lv2
        return self.exit(lv1)

class PatchD(nn.Module):
    def __init__(self, ndf=64):
        super(PatchD, self).__init__()

        sequence = [
            nn.Conv2d(6, ndf, kernel_size=4, stride=2, padding=1, bias=False),  # 128
            nn.LeakyReLU(0.2, True),

            Tunnel(2, ndf, ndf),
            DResNeXtBottleneck(ndf, ndf * 2, 2),  # 64

            Tunnel(3, ndf * 2, ndf * 2),
            DResNeXtBottleneck(ndf * 2, ndf * 4, 2),  # 32

            Tunnel(4, ndf * 4, ndf * 4),
            DResNeXtBottleneck(ndf * 4, ndf * 8, 2),  # 16

            Tunnel(4, ndf * 8, ndf * 8),
            DResNeXtBottleneck(ndf * 8, ndf * 16, 2),  # 8

            Tunnel(2, ndf * 16, ndf * 16),
            DResNeXtBottleneck(ndf * 16, ndf * 32, 2),  # 4

            nn.Conv2d(ndf * 32, 1, kernel_size=4, stride=1, padding=0, bias=False)

        ]

        self.model = nn.Sequential(*sequence)

    def forward(self, input):
        return self.model(input)

        # TODO: fix relu bug