# Copyright (c) 2018, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause

"""BERT finetuning runner."""

from __future__ import absolute_import, division, print_function

import argparse
import csv
import logging
import os
import random
import sys
import codecs
import numpy as np
import torch
import torch.nn as nn
from collections import defaultdict
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
from scipy.stats import beta
from torch.nn import CrossEntropyLoss, MSELoss
from scipy.special import softmax
from sklearn.metrics import f1_score

from transformers.models.roberta.tokenization_roberta import RobertaTokenizer
from transformers.optimization import AdamW
from transformers.models.roberta.modeling_roberta import RobertaModel#RobertaForSequenceClassification

p = os.path.abspath('../')
if p not in sys.path:
    sys.path.append(p)
from load_data import load_DocNLI

logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt = '%m/%d/%Y %H:%M:%S',
                    level = logging.INFO)
logger = logging.getLogger(__name__)

bert_hidden_dim = 1024
pretrain_model_dir = 'roberta-large' #'roberta-large' , 'roberta-large-mnli', 'bert-large-uncased'

def store_transformers_models(model, tokenizer, output_dir, flag_str):
    '''
    store the model
    '''
    output_dir+='/'+flag_str
    # if not os.path.exists(output_dir):
    #     os.makedirs(output_dir)
    print('starting model storing....')
    # model.save_pretrained(output_dir)
    torch.save(model.state_dict(), output_dir)
    # tokenizer.save_pretrained(output_dir)
    print('store succeed')

class RobertaForSequenceClassification(nn.Module):
    def __init__(self, tagset_size):
        super(RobertaForSequenceClassification, self).__init__()
        self.tagset_size = tagset_size

        self.roberta_single= RobertaModel.from_pretrained(pretrain_model_dir)
        self.single_hidden2tag = RobertaClassificationHead(bert_hidden_dim, tagset_size)

    def forward(self, input_ids, input_mask):
        outputs_single = self.roberta_single(input_ids, input_mask, None)
        hidden_states_single = outputs_single[1]#torch.tanh(self.hidden_layer_2(torch.tanh(self.hidden_layer_1(outputs_single[1])))) #(batch, hidden)

        score_single = self.single_hidden2tag(hidden_states_single) #(batch, tag_set)
        return score_single

class RobertaClassificationHead(nn.Module):
    """wenpeng overwrite it so to accept matrix as input"""

    def __init__(self, bert_hidden_dim, num_labels):
        super(RobertaClassificationHead, self).__init__()
        self.dense = nn.Linear(bert_hidden_dim, bert_hidden_dim)
        self.dropout = nn.Dropout(0.1)
        self.out_proj = nn.Linear(bert_hidden_dim, num_labels)

    def forward(self, features):
        x = features#[:, 0, :]  # take <s> token (equiv. to [CLS])
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x

class InputExample(object):
    """A single training/test example for simple sequence classification."""

    def __init__(self, guid, text_a, text_b=None, label=None):
        """Constructs a InputExample.

        Args:
            guid: Unique id for the example.
            text_a: string. The untokenized text of the first sequence. For single
            sequence tasks, only this sequence must be specified.
            text_b: (Optional) string. The untokenized text of the second sequence.
            Only must be specified for sequence pair tasks.
            label: (Optional) string. The label of the example. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text_a = text_a
        self.text_b = text_b
        self.label = label

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, label_id):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_id = label_id

class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""

    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    @classmethod
    def _read_tsv(cls, input_file, quotechar=None):
        """Reads a tab separated value file."""
        with open(input_file, "r") as f:
            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
            lines = []
            for line in reader:
                if sys.version_info[0] == 2:
                    line = list(unicode(cell, 'utf-8') for cell in line)
                lines.append(line)
            return lines

class RteProcessor(DataProcessor):
    """Processor for the RTE data set (GLUE version)."""

    def get_MNLI_train_and_dev(self, train_filename, dev_filename_list):
        '''
        classes: ["entailment", "neutral", "contradiction"]
        '''
        examples_per_file = []
        for filename in [train_filename]+dev_filename_list:
            examples=[]
            readfile = codecs.open(filename, 'r', 'utf-8')
            line_co=0
            for row in readfile:
                if line_co>0:
                    line=row.strip().split('\t')
                    guid = "train-"+str(line_co-1)
                    # text_a = 'MNLI. '+line[8].strip()
                    text_a = line[8].strip()
                    text_b = line[9].strip()
                    label = line[-1].strip() #["entailment", "neutral", "contradiction"]
                    examples.append(
                        InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
                line_co+=1
            readfile.close()
            print('loaded  MNLI size:', len(examples))
            examples_per_file.append(examples)
        dev_examples = []
        for listt in examples_per_file[1:]:
            dev_examples+=listt
        return examples_per_file[0], dev_examples #train, dev

    def get_labels(self):
        'here we keep the three-way in MNLI training '
        return ["entailment", "not_entailment"]
        # return ["entailment", "neutral", "contradiction"]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            if i == 0:
                continue
            guid = "%s-%s" % (set_type, line[0])
            text_a = line[1]
            text_b = line[2]
            label = line[-1]
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
        return examples

def convert_examples_to_features(examples, label_list, max_seq_length,
                                 tokenizer, output_mode,
                                 cls_token_at_end=False,
                                 cls_token='[CLS]',
                                 cls_token_segment_id=1,
                                 sep_token='[SEP]',
                                 sep_token_extra=False,
                                 pad_on_left=False,
                                 pad_token=0,
                                 pad_token_segment_id=0,
                                 sequence_a_segment_id=0,
                                 sequence_b_segment_id=1,
                                 mask_padding_with_zero=True):
    """ Loads a data file into a list of `InputBatch`s
        `cls_token_at_end` define the location of the CLS token:
            - False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]
            - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
        `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
    """

    label_map = {label : i for i, label in enumerate(label_list)}

    features = []
    for (ex_index, example) in enumerate(examples):
        if ex_index % 10000 == 0:
            logger.info("Writing example %d of %d" % (ex_index, len(examples)))

        tokens_a = tokenizer.tokenize(example.text_a)

        tokens_b = None
        if example.text_b:
            tokens_b = tokenizer.tokenize(example.text_b)
            # Modifies `tokens_a` and `tokens_b` in place so that the total
            # length is less than the specified length.
            # Account for [CLS], [SEP], [SEP] with "- 3". " -4" for RoBERTa.
            special_tokens_count = 4 if sep_token_extra else 3
            _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - special_tokens_count)
        else:
            # Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
            special_tokens_count = 3 if sep_token_extra else 2
            if len(tokens_a) > max_seq_length - special_tokens_count:
                tokens_a = tokens_a[:(max_seq_length - special_tokens_count)]

        # The convention in BERT is:
        # (a) For sequence pairs:
        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
        #  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1
        # (b) For single sequences:
        #  tokens:   [CLS] the dog is hairy . [SEP]
        #  type_ids:   0   0   0   0  0     0   0
        #
        # Where "type_ids" are used to indicate whether this is the first
        # sequence or the second sequence. The embedding vectors for `type=0` and
        # `type=1` were learned during pre-training and are added to the wordpiece
        # embedding vector (and position vector). This is not *strictly* necessary
        # since the [SEP] token unambiguously separates the sequences, but it makes
        # it easier for the model to learn the concept of sequences.
        #
        # For classification tasks, the first vector (corresponding to [CLS]) is
        # used as as the "sentence vector". Note that this only makes sense because
        # the entire model is fine-tuned.
        tokens = tokens_a + [sep_token]
        if sep_token_extra:
            # roberta uses an extra separator b/w pairs of sentences
            tokens += [sep_token]
        segment_ids = [sequence_a_segment_id] * len(tokens)

        if tokens_b:
            tokens += tokens_b + [sep_token]
            segment_ids += [sequence_b_segment_id] * (len(tokens_b) + 1)

        if cls_token_at_end:
            tokens = tokens + [cls_token]
            segment_ids = segment_ids + [cls_token_segment_id]
        else:
            tokens = [cls_token] + tokens
            segment_ids = [cls_token_segment_id] + segment_ids

        input_ids = tokenizer.convert_tokens_to_ids(tokens)

        # The mask has 1 for real tokens and 0 for padding tokens. Only real
        # tokens are attended to.
        input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)

        # Zero-pad up to the sequence length.
        padding_length = max_seq_length - len(input_ids)
        if pad_on_left:
            input_ids = ([pad_token] * padding_length) + input_ids
            input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
            segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids
        else:
            input_ids = input_ids + ([pad_token] * padding_length)
            input_mask = input_mask + ([0 if mask_padding_with_zero else 1] * padding_length)
            segment_ids = segment_ids + ([pad_token_segment_id] * padding_length)

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length

        if output_mode == "classification":
            label_id = label_map[example.label]
        elif output_mode == "regression":
            label_id = float(example.label)
        else:
            raise KeyError(output_mode)

        # if ex_index < 5:
        #     logger.info("*** Example ***")
        #     logger.info("guid: %s" % (example.guid))
        #     logger.info("tokens: %s" % " ".join(
        #             [str(x) for x in tokens]))
        #     logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
        #     logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
        #     logger.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
        #     logger.info("label: %s (id = %d)" % (example.label, label_id))

        features.append(
                InputFeatures(input_ids=input_ids,
                              input_mask=input_mask,
                              segment_ids=segment_ids,
                              label_id=label_id))
    return features

def _truncate_seq_pair(tokens_a, tokens_b, max_length):
    """Truncates a sequence pair in place to the maximum length."""

    # This is a simple heuristic which will always truncate the longer sequence
    # one token at a time. This makes more sense than truncating an equal percent
    # of tokens from each, since if one sequence is very short then each token
    # that's truncated likely contains more information than a longer sequence.
    while True:
        total_length = len(tokens_a) + len(tokens_b)
        if total_length <= max_length:
            break
        if len(tokens_a) > len(tokens_b):
            tokens_a.pop()
        else:
            tokens_b.pop()

def main():
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument("--task_name",
                        default=None,
                        type=str,
                        required=True,
                        help="The name of the task to train.")
    ## Other parameters
    parser.add_argument("--cache_dir",
                        default="",
                        type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument("--data_label",
                        default="",
                        type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument("--max_seq_length",
                        default=128,
                        type=int,
                        help="The maximum total input sequence length after WordPiece tokenization. \n"
                             "Sequences longer than this will be truncated, and sequences shorter \n"
                             "than this will be padded.")
    parser.add_argument("--do_train",
                        action='store_true',
                        help="Whether to run training.")
    parser.add_argument("--do_eval",
                        action='store_true',
                        help="Whether to run eval on the dev set.")
    parser.add_argument("--do_lower_case",
                        action='store_true',
                        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--train_batch_size",
                        default=16,
                        type=int,
                        help="Total batch size for training.")
    parser.add_argument("--eval_batch_size",
                        default=64,
                        type=int,
                        help="Total batch size for eval.")
    parser.add_argument("--learning_rate",
                        default=1e-5,
                        type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--num_train_epochs",
                        default=3.0,
                        type=float,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--warmup_proportion",
                        default=0.1,
                        type=float,
                        help="Proportion of training to perform linear learning rate warmup for. "
                             "E.g., 0.1 = 10%% of training.")
    parser.add_argument("--no_cuda",
                        action='store_true',
                        help="Whether not to use CUDA when available")
    parser.add_argument("--local_rank",
                        type=int,
                        default=-1,
                        help="local_rank for distributed training on gpus")
    parser.add_argument('--seed',
                        type=int,
                        default=42,
                        help="random seed for initialization")
    parser.add_argument('--gradient_accumulation_steps',
                        type=int,
                        default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument('--fp16',
                        action='store_true',
                        help="Whether to use 16-bit float precision instead of 32-bit")
    parser.add_argument('--loss_scale',
                        type=float, default=0,
                        help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
                             "0 (default value): dynamic loss scaling.\n"
                             "Positive power of 2: static loss scaling value.\n")
    parser.add_argument('--server_ip', type=str, default='', help="Can be used for distant debugging.")
    parser.add_argument('--server_port', type=str, default='', help="Can be used for distant debugging.")

    args = parser.parse_args()

    processors = {
        "rte": RteProcessor
    }

    output_modes = {
        "rte": "classification"
    }

    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        n_gpu = torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        n_gpu = 1
        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.distributed.init_process_group(backend='nccl')
    logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
        device, n_gpu, bool(args.local_rank != -1), args.fp16))

    if args.gradient_accumulation_steps < 1:
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
                            args.gradient_accumulation_steps))

    args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

    if not args.do_train and not args.do_eval:
        raise ValueError("At least one of `do_train` or `do_eval` must be True.")

    task_name = args.task_name.lower()

    if task_name not in processors:
        raise ValueError("Task not found: %s" % (task_name))

    processor = processors[task_name]()
    output_mode = output_modes[task_name]

    test_examples = load_DocNLI('test', hypo_only=False)

    label_list = ["entailment", "not_entailment"]#, "contradiction"]
    num_labels = len(label_list)
    print('num_labels:', num_labels,  ' test size:', len(test_examples))

    # device = torch.device('cpu')
    model = RobertaForSequenceClassification(num_labels)
    tokenizer = RobertaTokenizer.from_pretrained(pretrain_model_dir, do_lower_case=args.do_lower_case)
    model.load_state_dict(torch.load('DocNLI.pretrained.RoBERTA.model.pt', map_location=device))

    model.to(device)

    '''load test set'''
    test_features = convert_examples_to_features(
        test_examples, label_list, args.max_seq_length, tokenizer, output_mode,
        cls_token_at_end=False,#bool(args.model_type in ['xlnet']),            # xlnet has a cls token at the end
        cls_token=tokenizer.cls_token,
        cls_token_segment_id=0,#2 if args.model_type in ['xlnet'] else 0,
        sep_token=tokenizer.sep_token,
        sep_token_extra=True,#bool(args.model_type in ['roberta']),           # roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805
        pad_on_left=False,#bool(args.model_type in ['xlnet']),                 # pad on the left for xlnet
        pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
        pad_token_segment_id=0)#4 if args.model_type in ['xlnet'] else 0,)

    test_all_input_ids = torch.tensor([f.input_ids for f in test_features], dtype=torch.long)
    test_all_input_mask = torch.tensor([f.input_mask for f in test_features], dtype=torch.long)
    test_all_segment_ids = torch.tensor([f.segment_ids for f in test_features], dtype=torch.long)
    test_all_label_ids = torch.tensor([f.label_id for f in test_features], dtype=torch.long)

    test_data = TensorDataset(test_all_input_ids, test_all_input_mask, test_all_segment_ids, test_all_label_ids)
    test_sampler = SequentialSampler(test_data)
    test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=args.eval_batch_size)

    '''
    start evaluate on dev set after this epoch
    '''
    model.eval()
    final_test_performance = evaluation(test_dataloader, device, model)
    print('final_test_performance:', final_test_performance)

def evaluation(dev_dataloader, device, model):
    eval_loss = 0
    nb_eval_steps = 0
    preds = []
    gold_label_ids = []
    # print('Evaluating...')
    for input_ids, input_mask, segment_ids, label_ids in dev_dataloader:

        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        segment_ids = segment_ids.to(device)
        label_ids = label_ids.to(device)
        gold_label_ids+=list(label_ids.detach().cpu().numpy())

        with torch.no_grad():
            logits = model(input_ids, input_mask)
        if len(preds) == 0:
            preds.append(logits.detach().cpu().numpy())
        else:
            preds[0] = np.append(preds[0], logits.detach().cpu().numpy(), axis=0)

        nb_eval_steps+=1
        print('eval_steps:', nb_eval_steps, '/', len(dev_dataloader))

    preds = preds[0]

    pred_probs = softmax(preds,axis=1)
    pred_label_ids = list(np.argmax(pred_probs, axis=1))

    gold_label_ids = gold_label_ids
    assert len(pred_label_ids) == len(gold_label_ids)
    # print('gold_label_ids:', gold_label_ids)
    # print('pred_label_ids:', pred_label_ids)
    f1 = f1_score(gold_label_ids, pred_label_ids, pos_label= 0, average='binary')
    return f1

if __name__ == "__main__":
    main()

'''

CUDA_VISIBLE_DEVICES=3 python -u test_on_docNLI_RoBERTa.py --task_name rte --do_train --do_lower_case --data_label DUC --num_train_epochs 20 --train_batch_size 32 --eval_batch_size 128 --learning_rate 1e-6 --max_seq_length 512 --seed 42

'''


from asyncio.windows_events import NULL
from synthesizer.inference import Synthesizer
from pydantic import BaseModel, Field
from encoder import inference as speacker_encoder
import torch
import os
from pathlib import Path
from enum import Enum
import ppg_extractor as Extractor
import ppg2mel as Convertor
import librosa
from scipy.io.wavfile import write
import re
import numpy as np
from mkgui.base.components.types import FileContent
from vocoder.hifigan import inference as gan_vocoder
from typing import Any
import matplotlib.pyplot as plt

# Constants
AUDIO_SAMPLES_DIR = 'samples\\'
EXT_MODELS_DIRT = "ppg_extractor\\saved_models"
CONV_MODELS_DIRT = "ppg2mel\\saved_models"
VOC_MODELS_DIRT = "vocoder\\saved_models"
TEMP_SOURCE_AUDIO = "wavs/temp_source.wav"
TEMP_TARGET_AUDIO = "wavs/temp_target.wav"
TEMP_RESULT_AUDIO = "wavs/temp_result.wav"

# Load local sample audio as options TODO: load dataset
if os.path.isdir(AUDIO_SAMPLES_DIR):
    audio_input_selection = Enum('samples', list((file.name, file) for file in Path(AUDIO_SAMPLES_DIR).glob("*.wav")))
# Pre-Load models
if os.path.isdir(EXT_MODELS_DIRT):
    extractors =  Enum('extractors', list((file.name, file) for file in Path(EXT_MODELS_DIRT).glob("**/*.pt")))
    print("Loaded extractor models: " + str(len(extractors)))
else:
    raise Exception(f"Model folder {EXT_MODELS_DIRT} doesn't exist.")

if os.path.isdir(CONV_MODELS_DIRT):
    convertors =  Enum('convertors', list((file.name, file) for file in Path(CONV_MODELS_DIRT).glob("**/*.pth")))
    print("Loaded convertor models: " + str(len(convertors)))
else:
    raise Exception(f"Model folder {CONV_MODELS_DIRT} doesn't exist.")

if os.path.isdir(VOC_MODELS_DIRT):
    vocoders =  Enum('vocoders', list((file.name, file) for file in Path(VOC_MODELS_DIRT).glob("**/*gan*.pt")))
    print("Loaded vocoders models: " + str(len(vocoders)))
else:
    raise Exception(f"Model folder {VOC_MODELS_DIRT} doesn't exist.")

class Input(BaseModel):
    local_audio_file: audio_input_selection = Field(
        ..., alias="输入语音（本地wav）",
        description="选择本地语音文件."
    )
    upload_audio_file: FileContent = Field(default=None, alias="或上传语音",
        description="拖拽或点击上传.", mime_type="audio/wav")
    local_audio_file_target: audio_input_selection = Field(
        ..., alias="目标语音（本地wav）",
        description="选择本地语音文件."
    )
    upload_audio_file_target: FileContent = Field(default=None, alias="或上传目标语音",
        description="拖拽或点击上传.", mime_type="audio/wav")
    extractor: extractors = Field(
        ..., alias="编码模型",
        description="选择语音编码模型文件."
    )
    convertor: convertors = Field(
        ..., alias="转换模型",
        description="选择语音转换模型文件."
    )
    vocoder: vocoders = Field(
        ..., alias="语音解码模型",
        description="选择语音解码模型文件(目前只支持HifiGan类型)."
    )

class AudioEntity(BaseModel):
    content: bytes
    mel: Any

class Output(BaseModel):
    __root__: tuple[AudioEntity, AudioEntity, AudioEntity]

    def render_output_ui(self, streamlit_app, input) -> None:  # type: ignore
        """Custom output UI.
        If this method is implmeneted, it will be used instead of the default Output UI renderer.
        """
        src, target, result = self.__root__

        streamlit_app.subheader("Synthesized Audio")
        streamlit_app.audio(result.content, format="audio/wav")

        fig, ax = plt.subplots()
        ax.imshow(src.mel, aspect="equal", interpolation="none")
        ax.set_title("mel spectrogram(Source Audio)")
        streamlit_app.pyplot(fig)
        fig, ax = plt.subplots()
        ax.imshow(target.mel, aspect="equal", interpolation="none")
        ax.set_title("mel spectrogram(Target Audio)")
        streamlit_app.pyplot(fig)
        fig, ax = plt.subplots()
        ax.imshow(result.mel, aspect="equal", interpolation="none")
        ax.set_title("mel spectrogram(Result Audio)")
        streamlit_app.pyplot(fig)

def convert(input: Input) -> Output:
    """convert(转换)"""
    # load models
    extractor = Extractor.load_model(Path(input.extractor.value))
    convertor = Convertor.load_model(Path(input.convertor.value))
    # current_synt = Synthesizer(Path(input.synthesizer.value))
    gan_vocoder.load_model(Path(input.vocoder.value))

    # load file
    if input.upload_audio_file != None:
        with open(TEMP_SOURCE_AUDIO, "w+b") as f:
            f.write(input.upload_audio_file.as_bytes())
            f.seek(0)
        src_wav, sample_rate = librosa.load(TEMP_SOURCE_AUDIO)
    else:
        src_wav, sample_rate  = librosa.load(input.local_audio_file.value)
        write(TEMP_SOURCE_AUDIO, sample_rate, src_wav) #Make sure we get the correct wav

    if input.upload_audio_file_target != None:
        with open(TEMP_TARGET_AUDIO, "w+b") as f:
            f.write(input.upload_audio_file_target.as_bytes())
            f.seek(0)
        ref_wav, _ = librosa.load(TEMP_TARGET_AUDIO)
    else:
        ref_wav, _  = librosa.load(input.local_audio_file_target.value)
        write(TEMP_TARGET_AUDIO, sample_rate, ref_wav) #Make sure we get the correct wav

    ppg = extractor.extract_from_wav(src_wav)
    # Import necessary dependency of Voice Conversion
    from utils.f0_utils import compute_f0, f02lf0, compute_mean_std, get_converted_lf0uv
    ref_lf0_mean, ref_lf0_std = compute_mean_std(f02lf0(compute_f0(ref_wav)))
    speacker_encoder.load_model(Path("encoder/saved_models/pretrained_bak_5805000.pt"))
    embed = speacker_encoder.embed_utterance(ref_wav)
    lf0_uv = get_converted_lf0uv(src_wav, ref_lf0_mean, ref_lf0_std, convert=True)
    min_len = min(ppg.shape[1], len(lf0_uv))
    ppg = ppg[:, :min_len]
    lf0_uv = lf0_uv[:min_len]
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    _, mel_pred, att_ws = convertor.inference(
        ppg,
        logf0_uv=torch.from_numpy(lf0_uv).unsqueeze(0).float().to(device),
        spembs=torch.from_numpy(embed).unsqueeze(0).to(device),
    )
    mel_pred= mel_pred.transpose(0, 1)
    breaks = [mel_pred.shape[1]]
    mel_pred= mel_pred.detach().cpu().numpy()

    # synthesize and vocode
    wav, sample_rate = gan_vocoder.infer_waveform(mel_pred)

    # write and output
    write(TEMP_RESULT_AUDIO, sample_rate, wav) #Make sure we get the correct wav
    with open(TEMP_SOURCE_AUDIO, "rb") as f:
        source_file = f.read()
    with open(TEMP_TARGET_AUDIO, "rb") as f:
        target_file = f.read()
    with open(TEMP_RESULT_AUDIO, "rb") as f:
        result_file = f.read()

    return Output(__root__=(AudioEntity(content=source_file, mel=Synthesizer.make_spectrogram(src_wav)), AudioEntity(content=target_file, mel=Synthesizer.make_spectrogram(ref_wav)), AudioEntity(content=result_file, mel=Synthesizer.make_spectrogram(wav))))


#!/usr/bin/env python

'''
Creates initial configuration for a system containing unbranched linear chains.

'''
import sys
import math
import numpy as np
from configuration import Configuration
from config_io import *

#-------------------------------------------------------------------------------

config = Configuration()

#Add simulation box and boundary condition. The box size can be updated later if
#necessary.
config.add_simbox(10.0, 10.0, 10.0, 1)

#Molecule details
num_mols = 50
config.add_molecule_type('CHN-UB', num_mols)

#Total number of atoms in the molecule
natm = 2 #int(sys.argv[1])

#Total number of bonds
nbnd = natm - 1

#Total number of angles. If no angles set num_angles to zero.
nang = natm - 2 if natm > 2 else 0

#Atom types
# natm atoms as point particles with mass 1.0
atm_t_bb = config.add_atom_type('M', 1, 1.0)

#Vdw interaction
eps = 1.0; sigma = 2.0; rcut = 2.0**(1.0/6)*sigma
config.add_ia_vdw('M', 'M', 1, np.array([eps, sigma, rcut]))

#A = 2.0; rcut = 1.0
#config.add_ia_vdw('M', 'M', 3, np.array([A, rcut]))

#Bond types
eps = 1.0; sigma = 2.0
bnd_t_bb = config.add_bond_type(3, np.array([30*eps/sigma**2, 1.5*sigma, eps, sigma]))

#Angle types
ang_t_bb = config.add_angle_type(1, np.array([2.0]))

#Tether types (No tethers)
#teth_t = config.add_tether_type(1, np.array([2.0, 1e-6]))

#Add atoms, bonds, and angles
len_bond = 1.94 #0.97*sigma #Equilibrium bond length
theta = (25.0/180.0)*math.pi
sep = 1.9 # Must be <= len_bond

#Loop over all molecules (here only a single molecule type)
imol = 1; iatm = 1; ibnd = 1; iang = 1 #Index of one past the last atom, bond, etc. added
for imol in range(1, num_mols+1):
    iatm_beg = iatm
    ibnd_beg = ibnd
    iang_beg = iang

    config.append_atom_unbonded( atm_t_bb, 0.0, sep=sep )
    iatm += 1
    config.append_atom_bonded(1, 0.0, len_bond, 'efjc', iatm-1, sep=sep)
    iatm += 1
    for i in range(3,natm+1):
        config.append_atom_bonded(atm_t_bb, 0.0, len_bond, 'efjc',
                iatm-1, im2=iatm-2, theta=theta, sep=sep)
        iatm += 1

    for i in range(1, nbnd+1):
        j = iatm_beg + i - 1
        config.add_bond(bnd_t_bb, j, j+1)
        ibnd += 1

    for i in range(1, nang+1):
        j = iatm_beg + i - 1
        config.add_angle(ang_t_bb, j, j+1, j+2)
        iang += 1

    config.add_molecule(1, natm, iatm_beg, nbnd, ibnd_beg, nang, iang_beg, 0, 0)

#config.add_tether(teth_t, 1, np.zeros((3,)) )

#config.update_simbox()
#config.to_center(1)
config.apply_pbc()
#config.add_flow_field( 0, np.array([]) )

write_ldf(config, 'tfr.txt', title='test_frame')
write_cfg(config, 'sus-%d.cfg'%natm, title='test_cfg')


import os
import math
import shutil
import unittest
import numpy as np

from signaturizer import Signaturizer

class TestSignaturizer(unittest.TestCase):

    def setUp(self):
        # path for test data
        test_dir = os.path.dirname(os.path.realpath(__file__))
        self.data_dir = os.path.join(test_dir, 'data')
        self.tmp_dir = os.path.join(test_dir, 'tmp')
        if os.path.exists(self.tmp_dir):
            shutil.rmtree(self.tmp_dir)
        os.mkdir(self.tmp_dir)
        self.test_smiles = [
            # Erlotinib
            'COCCOC1=C(C=C2C(=C1)C(=NC=N2)NC3=CC=CC(=C3)C#C)OCCOC',
            # Diphenhydramine
            'CN(C)CCOC(C1=CC=CC=C1)C2=CC=CC=C2'
        ]
        self.invalid_smiles = ['C', 'C&', 'C']
        self.tautomer_smiles = ['CC(O)=Nc1ccc(O)cc1', 'CC(=O)NC1=CC=C(C=C1)O']
        self.inchi = [
            'InChI=1S/C8H9NO2/c1-6(10)9-7-2-4-8(11)5-3-7/h2-5,11H,1H3,(H,9,10)'
        ]

    def tearDown(self):
        if os.path.exists(self.tmp_dir):
            shutil.rmtree(self.tmp_dir, ignore_errors=True)
            pass

    def test_predict(self):
        # load reference predictions
        ref_file = os.path.join(self.data_dir, 'pred.npy')
        pred_ref = np.load(open(ref_file, 'rb'))
        # load module and predict
        module_dir = os.path.join(self.data_dir, 'B1')
        module = Signaturizer(module_dir, local=True)
        res = module.predict(self.test_smiles)
        np.testing.assert_almost_equal(pred_ref, res.signature[:])
        # test saving to file
        destination = os.path.join(self.tmp_dir, 'pred.h5')
        res = module.predict(self.test_smiles, destination)
        self.assertTrue(os.path.isfile(destination))
        np.testing.assert_almost_equal(pred_ref, res.signature[:])
        self.assertEqual(len(res.applicability[:]), 2)
        self.assertFalse(np.isnan(res.applicability[0]))
        # test prediction of invalid SMILES
        res = module.predict(self.invalid_smiles)
        for comp in res.signature[0]:
            self.assertFalse(math.isnan(comp))
        for comp in res.signature[1]:
            self.assertTrue(math.isnan(comp))
        for comp in res.signature[2]:
            self.assertFalse(math.isnan(comp))

    def test_predict_multi(self):
        # load multiple model and check that results stacked correctly
        module_dirs = list()
        A1_path = os.path.join(self.data_dir, 'A1')
        B1_path = os.path.join(self.data_dir, 'B1')
        module_dirs.append(A1_path)
        module_dirs.append(B1_path)
        module_A1B1 = Signaturizer(module_dirs, local=True)
        res_A1B1 = module_A1B1.predict(self.test_smiles)
        self.assertEqual(res_A1B1.signature.shape[0], 2)
        self.assertEqual(res_A1B1.signature.shape[1], 128 * 2)

        module_A1 = Signaturizer(A1_path, local=True)
        res_A1 = module_A1.predict(self.test_smiles)
        np.testing.assert_almost_equal(res_A1B1.signature[:, :128],
                                       res_A1.signature)

        module_B1 = Signaturizer(B1_path, local=True)
        res_B1 = module_B1.predict(self.test_smiles)
        np.testing.assert_almost_equal(res_A1B1.signature[:, 128:],
                                       res_B1.signature)

        res = module_A1B1.predict(self.invalid_smiles)
        for comp in res.signature[0]:
            self.assertFalse(math.isnan(comp))
        for comp in res.signature[1]:
            self.assertTrue(math.isnan(comp))
        for comp in res.signature[2]:
            self.assertFalse(math.isnan(comp))
        self.assertTrue(all(res.failed == [False, True, False]))

    def test_predict_global_remote(self):
        module = Signaturizer(['GLOBAL'])
        res = module.predict(self.test_smiles)
        self.assertEqual(res.signature.shape[0], 2)
        self.assertEqual(res.signature.shape[1], 128 * 25)

    def test_overwrite(self):
        module_dir = os.path.join(self.data_dir, 'B1')
        module = Signaturizer(module_dir, local=True)
        destination = os.path.join(self.tmp_dir, 'pred.h5')
        module.predict(self.test_smiles, destination)
        # repeating writing will result in an exception
        with self.assertRaises(Exception):
            module.predict(self.test_smiles, destination)

    def test_tautomers(self):
        module = Signaturizer('A1')
        res = module.predict(self.tautomer_smiles)
        self.assertTrue(all(res.signature[0] == res.signature[1]))

    def test_inchi(self):
        module = Signaturizer('A1')
        res_inchi = module.predict(self.inchi, keytype='InChI')
        res_smiles = module.predict([self.tautomer_smiles[0]])
        self.assertTrue(all(res_inchi.signature[0] == res_smiles.signature[0]))

    def test_all_single(self):
        module = Signaturizer('A1')
        res_all = module.predict(self.test_smiles)
        for idx, smile in enumerate(self.test_smiles):
            res_single = module.predict([smile])
            np.testing.assert_almost_equal(res_all.signature[idx],
                                           res_single.signature[0])


from absl import app, flags, logging
from absl.flags import FLAGS
import os
import shutil
import tensorflow as tf
from core.yolov4 import YOLO, decode, compute_loss, decode_train
from core.dataset import Dataset
from core.config import cfg
import numpy as np
from core import utils
from core.utils import freeze_all, unfreeze_all

### - Adding lines to remove the error - tensorflow.python.framework.errors_impl.NotFoundError: No algorithm worked! [Op:Conv2D]
#'''
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

config = ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
#config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
#'''
###

flags.DEFINE_string('model', 'yolov4', 'yolov4, yolov3')
flags.DEFINE_string('weights', './data/yolov4.weights', 'pretrained weights')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')

def main(_argv):
    #physical_devices = tf.config.experimental.list_physical_devices('GPU')
    #if len(physical_devices) > 0:
    #    tf.config.experimental.set_memory_growth(physical_devices[0], True) # as an error was throwing https://github.com/hunglc007/tensorflow-yolov4-tflite/issues/171

    '''
    config = ConfigProto()
    #config.gpu_options.per_process_gpu_memory_fraction = 0.2
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    '''
    #test code
    '''
    config = tf.compat.v1.ConfigProto(gpu_options=
                                      tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.8)
                                      # device_count = {'GPU': 1}
                                      )
    config.gpu_options.allow_growth = True
    session = tf.compat.v1.Session(config=config)
    tf.compat.v1.keras.backend.set_session(session)
    '''
    #test code ends here

    '''
    gpus = tf.config.list_physical_devices('GPU')
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])
    '''

    trainset = Dataset(FLAGS, is_training=True)
    testset = Dataset(FLAGS, is_training=False)
    logdir = "./data/log"
    isfreeze = False
    #steps_per_epoch = len(trainset)
    steps_per_epoch = 200  # experimental to train the networks fast
    first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
    second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
    global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
    warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
    total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
    # train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period

    input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
    IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH

    freeze_layers = utils.load_freeze_layer(FLAGS.model, FLAGS.tiny)

    feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
    if FLAGS.tiny:
        bbox_tensors = []
        for i, fm in enumerate(feature_maps):
            if i == 0:
                bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
            else:
                bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
            bbox_tensors.append(fm)
            bbox_tensors.append(bbox_tensor)
    else:
        bbox_tensors = []
        for i, fm in enumerate(feature_maps):
            if i == 0:
                bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 8, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
            elif i == 1:
                bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 16, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
            else:
                bbox_tensor = decode_train(fm, cfg.TRAIN.INPUT_SIZE // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
            bbox_tensors.append(fm)
            bbox_tensors.append(bbox_tensor)

    model = tf.keras.Model(input_layer, bbox_tensors)
    model.summary()

    if FLAGS.weights == None:
        print("Training from scratch")
    else:
        if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
            utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
        else:
            model.load_weights(FLAGS.weights)
        print('Restoring weights from: %s ... ' % FLAGS.weights)

    optimizer = tf.keras.optimizers.Adam()
    if os.path.exists(logdir): shutil.rmtree(logdir)
    writer = tf.summary.create_file_writer(logdir)

    # define training step function
    # @tf.function
    def train_step(image_data, target):
        with tf.GradientTape() as tape:
            pred_result = model(image_data, training=True)
            giou_loss = conf_loss = prob_loss = 0

            # optimizing process
            for i in range(len(freeze_layers)):
                conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
                loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
                giou_loss += loss_items[0]
                conf_loss += loss_items[1]
                prob_loss += loss_items[2]

            total_loss = giou_loss + conf_loss + prob_loss

            gradients = tape.gradient(total_loss, model.trainable_variables)
            optimizer.apply_gradients(zip(gradients, model.trainable_variables))
            tf.print("=> STEP %4d/%4d   lr: %.6f   giou_loss: %4.2f   conf_loss: %4.2f   "
                     "prob_loss: %4.2f   total_loss: %4.2f" % (global_steps, total_steps, optimizer.lr.numpy(),
                                                               giou_loss, conf_loss,
                                                               prob_loss, total_loss))
            # update learning rate
            global_steps.assign_add(1)
            if global_steps < warmup_steps:
                lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
            else:
                lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
                    (1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
                )
            optimizer.lr.assign(lr.numpy())

            # writing summary data
            with writer.as_default():
                tf.summary.scalar("lr", optimizer.lr, step=global_steps)
                tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
                tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
                tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
                tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
            writer.flush()
    def test_step(image_data, target):
        with tf.GradientTape() as tape:
            pred_result = model(image_data, training=True)
            giou_loss = conf_loss = prob_loss = 0

            # optimizing process
            for i in range(len(freeze_layers)):
                conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
                loss_items = compute_loss(pred, conv, target[i][0], target[i][1], STRIDES=STRIDES, NUM_CLASS=NUM_CLASS, IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
                giou_loss += loss_items[0]
                conf_loss += loss_items[1]
                prob_loss += loss_items[2]

            total_loss = giou_loss + conf_loss + prob_loss

            tf.print("=> TEST STEP %4d   giou_loss: %4.2f   conf_loss: %4.2f   "
                     "prob_loss: %4.2f   total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
                                                               prob_loss, total_loss))

    for epoch in range(first_stage_epochs + second_stage_epochs):
        if epoch < first_stage_epochs:
            if not isfreeze:
                isfreeze = True
                for name in freeze_layers:
                    freeze = model.get_layer(name)
                    freeze_all(freeze)
        elif epoch >= first_stage_epochs:
            if isfreeze:
                isfreeze = False
                for name in freeze_layers:
                    freeze = model.get_layer(name)
                    unfreeze_all(freeze)
        for image_data, target in trainset:
            train_step(image_data, target)
        for image_data, target in testset:
            test_step(image_data, target)
        model.save_weights("./checkpoints/yolov4_1")
        #model.save("./saved_keras_model")
        model.save("pruned_model.h5")

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


import copy
import glob
import os
import time
import pdb
from collections import deque

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from a2c_ppo_acktr import algo
from a2c_ppo_acktr.arguments import get_args
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
from a2c_ppo_acktr.utils import get_vec_normalize, update_linear_schedule
from a2c_ppo_acktr.visualize import visdom_plot
from mymodels import TemporalDifferenceModule, CollectSamples, Logger

args = get_args()
assert args.algo in ['a2c', 'ppo', 'acktr']
if args.recurrent_policy:
    assert args.algo in ['a2c', 'ppo'], \
        'Recurrent policy is not implemented for ACKTR'

args.num_processes=8
num_updates = int(args.num_env_steps) // args.num_steps // args.num_processes
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)

if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

try:
    os.makedirs(args.log_dir, exist_ok=True)
except OSError:
    files = glob.glob(os.path.join(args.log_dir, '*.monitor.csv'))
    for f in files:
        os.remove(f)

eval_log_dir = args.log_dir + "_eval"

try:
    os.makedirs(eval_log_dir)
except OSError:
    files = glob.glob(os.path.join(eval_log_dir, '*.monitor.csv'))
    for f in files:
        os.remove(f)

def main():
    torch.set_num_threads(1)
    device = torch.device("cuda:0" if args.cuda else "cpu")

    if args.vis:
        from visdom import Visdom
        viz = Visdom(port=args.port)
        win = None
    envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
                        args.gamma, args.log_dir, args.add_timestep, device, True)

    frame_skip = 4          # frame skip
    if args.tb_dir[-1] != '/':
        args.tb_dir = args.tb_dir + '/'
    logger = Logger(args.tb_dir)
    logger.write_settings(args)
    if args.use_tdm:

        # beta scheduler
        if args.beta_schedule == 'const':
            beta_func = lambda x : float(args.beta_int)
        elif args.beta_schedule == 'sqrt':
            beta_func = lambda x : 1./np.sqrt(x+2)
        elif args.beta_schedule == 'log':
            beta_func = lambda x : 1./np.log(x+2)
        elif args.beta_schedule == 'linear':
            beta_func = lambda x : 1./(x+2)

        # bonus function variations
        if args.bonus_func == 'linear':
            bonus_func = lambda x : x+1
        elif args.bonus_func == 'square':
            bonus_func = lambda x : (x+1)**2
        elif args.bonus_func == 'sqrt':
            bonus_func = lambda x : (x+1)**(1/2)
        elif args.bonus_func == 'log':
            bonus_func = lambda x : np.log(x+1)

        # temporal difference module
        tdm = TemporalDifferenceModule(inputSize= 2*int(envs.observation_space.shape[0]),
                                      outputSize=args.time_intervals,
                                      num_fc_layers=int(args.num_layers),
                                      depth_fc_layers=int(args.fc_width),
                                      lr=float(args.opt_lr),
                                      buffer_max_length = args.buffer_max_length,
                                      buffer_RL_ratio=args.buffer_RL_ratio,
                                      frame_skip=frame_skip,
                                      tdm_epoch=args.tdm_epoch,
                                      tdm_batchsize=args.tdm_batchsize,
                                      logger=logger,
                                      bonus_func=bonus_func).to(device)

        #collect random trajectories
        sample_collector = CollectSamples(envs, args.num_processes, initial=True)
        tdm.buffer_rand = sample_collector.collect_trajectories(args.num_rollouts, args.steps_per_rollout)

        # initial training
        tdm.update()

    actor_critic = Policy(envs.observation_space.shape, envs.action_space,
        base_kwargs={'recurrent': args.recurrent_policy})
    actor_critic.to(device)

    if args.algo == 'a2c':
        agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef,
                               args.entropy_coef, lr=args.lr,
                               eps=args.eps, alpha=args.alpha,
                               max_grad_norm=args.max_grad_norm)
    elif args.algo == 'ppo':
        agent = algo.PPO(actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch,
                         args.value_loss_coef, args.entropy_coef, lr=args.lr,
                               eps=args.eps,
                               max_grad_norm=args.max_grad_norm)
    elif args.algo == 'acktr':
        agent = algo.A2C_ACKTR(actor_critic, args.value_loss_coef,
                               args.entropy_coef, acktr=True)

    rollouts = RolloutStorage(args.num_steps, args.num_processes,
                        envs.observation_space.shape, envs.action_space,
                        actor_critic.recurrent_hidden_state_size)
    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)
    episode_rewards = deque(maxlen=10)
    start = time.time()
    for j in range(num_updates):

        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            if args.algo == "acktr":
                # use optimizer's learning rate since it's hard-coded in kfac.py
                update_linear_schedule(agent.optimizer, j, num_updates, agent.optimizer.lr)
            else:
                update_linear_schedule(agent.optimizer, j, num_updates, args.lr)

        if args.algo == 'ppo' and args.use_linear_clip_decay:
            agent.clip_param = args.clip_param  * (1 - j / float(num_updates))

        # acting
        for step in range(args.num_steps):

            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                        rollouts.obs[step],
                        rollouts.recurrent_hidden_states[step],
                        rollouts.masks[step])

            # Obser reward and next obs
            # envs.render()

            obs_old = obs.clone()
            obs, reward, done, infos = envs.step(action)

            for info in infos:
                if 'episode' in info.keys():
                    episode_rewards.append(info['episode']['r'])

            # If done then clean the history of observations.
            masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])

            #compute intrinsic bonus
            if args.use_tdm:
                tdm.symm_eval = True if step == args.num_steps-1 else False
                reward_int = tdm.compute_bonus(obs_old, obs).float()
                reward += beta_func(step + j*args.num_steps) * reward_int.cpu().unsqueeze(1)

                if (j % args.log_interval == 0) and (step == args.num_steps-1):
                    logger.add_reward_intrinsic(reward_int, (j+1)*args.num_steps*args.num_processes)

            #saving to buffer.
            rollouts.insert(obs, recurrent_hidden_states, action, action_log_prob, value, reward, masks)

        # saving to buffer and periodic updating parameters
        if (args.use_tdm):
            tdm.buffer_RL_temp.append((rollouts.obs, rollouts.masks))
            if (j%args.num_steps==0 and j>0):
                tdm.update()

        with torch.no_grad():
            next_value = actor_critic.get_value(rollouts.obs[-1],
                                                rollouts.recurrent_hidden_states[-1],
                                                rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)

        value_loss, action_loss, dist_entropy = agent.update(rollouts)

        rollouts.after_update()

        # save for every interval-th episode or for the last epoch
        # no
        # save every 1-million steps
        if (((j+1)*args.num_steps*args.num_processes)%1e6 == 0 or j == num_updates - 1) and args.save_dir != "":
            save_path = os.path.join(args.save_dir, args.algo)
            try:
                os.makedirs(save_path)
            except OSError:
                pass

            # A really ugly way to save a model to CPU
            save_model = actor_critic
            if args.cuda:
                save_model = copy.deepcopy(actor_critic).cpu()

            save_model = [save_model,
                          getattr(get_vec_normalize(envs), 'ob_rms', None)]

            if j == num_updates - 1:
                save_here = os.path.join(save_path, args.env_name + "_step_{}M.pt".format((j+1)*args.num_steps*args.num_processes//1e6))
            else:
                save_here = os.path.join(save_path, args.env_name + "_final.pt")
            torch.save(save_model, save_here) # saved policy.

        total_num_steps = (j + 1) * args.num_processes * args.num_steps

        # printing outputs
        if j % args.log_interval == 0 and len(episode_rewards) > 1:
            end = time.time()
            print("Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n".
                format(j,
                       total_num_steps,
                       int(total_num_steps / (end - start)),
                       len(episode_rewards),
                       np.mean(episode_rewards),
                       np.median(episode_rewards),
                       np.min(episode_rewards),
                       np.max(episode_rewards), dist_entropy,
                       value_loss, action_loss))
            logger.add_reward(episode_rewards, (j+1)*args.num_steps*args.num_processes)

        #
        # if j % args.tb_interval == 0:
        #     # mean/std or median/1stqt?
        #     logger.add_tdm_loss(loss, self.epoch_count*i)

        # evaluation process
        # if (args.eval_interval is not None
        #         and len(episode_rewards) > 1
        #         and j % args.eval_interval == 0):
        #     eval_envs = make_vec_envs(
        #         args.env_name, args.seed + args.num_processes, args.num_processes,
        #         args.gamma, eval_log_dir, args.add_timestep, device, True)
        #
        #     vec_norm = get_vec_normalize(eval_envs)
        #     if vec_norm is not None:
        #         vec_norm.eval()
        #         vec_norm.ob_rms = get_vec_normalize(envs).ob_rms
        #
        #     eval_episode_rewards = []
        #
        #     obs = eval_envs.reset()
        #     eval_recurrent_hidden_states = torch.zeros(args.num_processes,
        #                     actor_critic.recurrent_hidden_state_size, device=device)
        #     eval_masks = torch.zeros(args.num_processes, 1, device=device)
        #
        #     while len(eval_episode_rewards) < 10:
        #         with torch.no_grad():
        #             _, action, _, eval_recurrent_hidden_states = actor_critic.act(
        #                 obs, eval_recurrent_hidden_states, eval_masks, deterministic=True)
        #
        #         # Obser reward and next obs
        #         # envs.render()
        #         obs, reward, done, infos = eval_envs.step(action)
        #
        #         eval_masks = torch.FloatTensor([[0.0] if done_ else [1.0]
        #                                         for done_ in done])
        #         for info in infos:
        #             if 'episode' in info.keys():
        #                 eval_episode_rewards.append(info['episode']['r'])
        #
        #     eval_envs.close()
        #
        #     print(" Evaluation using {} episodes: mean reward {:.5f}\n".
        #         format(len(eval_episode_rewards),
        #                np.mean(eval_episode_rewards)))

        # # plotting
        # if args.vis and j % args.vis_interval == 0:
        #     try:
        #         # Sometimes monitor doesn't properly flush the outputs
        #         win = visdom_plot(viz, win, args.log_dir, args.env_name,
        #                           args.algo, args.num_env_steps)
        #     except IOError:
        #         pass
    #if done save:::::::::::
    logger.save()

if __name__ == "__main__":
    main()


'''Train CIFAR10 with PyTorch.'''
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import pickle
import torchvision
import torchvision.transforms as transforms
from build_vocab import Vocabulary
import numpy as np
import os
import random
import warnings
import argparse
from model import *
from Decoder import *
from data_loader import CocoDataset, collate_fn
from torch.utils.tensorboard import SummaryWriter

from utils import Config
from torch.optim.lr_scheduler import MultiStepLR
#os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from torchvision.models import ResNet
import torch.distributed as dist
import torch.utils.data
import torch.utils.data.distributed

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc = 0  # best test accuracy
start_epoch = 0  # start from epoch 0 or last checkpoint epoch

def parse_args():
    parser = argparse.ArgumentParser(description='PyTorch CoCo Training')
    parser.add_argument('config', type=str, default='configs/baseline.py',
                        help='train config file path')
    parser.add_argument('--vocab_path', type=str, default='data/vocab.pkl',
                        help='path for vocabulary wrapper')
    parser.add_argument('--image_dir', type=str, default='data/train2014',
                        help='directory for resized images')
    parser.add_argument('--caption_path', type=str, default='data/annotations/captions_train2014.json',
                        help='path for train annotation json file')
    parser.add_argument('--seed', default=None, type=int,
                        help='seed for initializing training. ')
    parser.add_argument('--gpu', default=None, type=int,
                        help='GPU id to use.')

    args = parser.parse_args()
    return args

# Training
def train_ic(epoch,cfg,net,decoder,linear,bn,optimizer_ic,criterion,data_loader,writer):
    caption_ite = 0
    net.train()
    #build the frontend
    decoder.train()
    linear.train()
    bn.train()
    total_step = len(data_loader)
    print('Train Image Caption.')
    for i, (images, captions, lengths) in enumerate(data_loader):
        # Set mini-batch dataset
        images = images.to(device)
        captions = captions.to(device)
        targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
        targets = targets.to(device)
        decoder.zero_grad()
        net.zero_grad()
        features = F.adaptive_avg_pool2d(net(images), 1)
        features = features.view(features.size(0), -1)
        features = bn(linear(features))
        outputs = decoder(features, captions, lengths)

        loss_ic = criterion(outputs, targets)
        loss_ic.backward()
        optimizer_ic.step()

        writer.add_scalar('IC_Loss/train', loss_ic, epoch * total_step + i)
        writer.add_scalar('Perplexity/train', np.exp(loss_ic.item()), epoch * total_step + i)

        # Print log info
        if i % cfg.log_step == 0:
            print('IC-Epoch [{}], Step [{}/{}], Loss: {:.4f}, Perplexity: {:5.4f}'
                  .format(epoch, len(data_loader), i, loss_ic.item(), np.exp(loss_ic.item())))

        # Save the model checkpoints
        if (i + 1) % cfg.save_step == 0:
            if not os.path.exists(cfg.model_path):
                os.makedirs(cfg.model_path)
            torch.save(decoder.state_dict(), os.path.join(
                cfg.model_path, 'new_decoder-{}-{}.ckpt'.format(epoch + 1, i + 1)), _use_new_zipfile_serialization=False)
            torch.save(net.state_dict(), os.path.join(
                cfg.model_path, 'new_encoder-{}-{}.ckpt'.format(epoch + 1, i + 1)), _use_new_zipfile_serialization=False)
    return net

def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    cfg.config = args.config
    print(cfg)
    if cfg.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if cfg.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely '
                      'disable data parallelism.')

    if cfg.dist_url == "env://" and cfg.world_size == -1:
        cfg.world_size = int(os.environ["WORLD_SIZE"])

    # cfg.distributed = cfg.world_size > 1

    ngpus_per_node = torch.cuda.device_count()
    main_worker(cfg.gpu, ngpus_per_node, cfg)

def main_worker(gpu, ngpus_per_node, cfg):
    if cfg.gpu is not None:
        print("Use GPU: {} for training".format(cfg.gpu))

    if cfg.distributed:
        print('init distributing process')
        if cfg.dist_url == "env://" and cfg.rank == -1:
            cfg.rank = int(os.environ["RANK"])
        dist.init_process_group(backend=cfg.dist_backend, init_method=cfg.dist_url,
                                world_size=cfg.world_size, rank=cfg.rank)

    # Data
    print('==> Preparing data..')
    # Load vocabulary wrapper for image caption
    with open(cfg.vocab_path, 'rb') as f:
        vocab = pickle.load(f)
    # Image preprocessing, normalization for the pretrained resnet
    # cifar cls, use resized 36x36 image
    if cfg.task == 'cifar_cls':
        transform = transforms.Compose([
            transforms.RandomCrop(cfg.crop_size, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))])

    # imagenet cls, 224x224
    # same as MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
    if cfg.task == 'imagenet_cls':
        transform = transforms.Compose([
            transforms.RandomGrayscale(p=0.2),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))])

    # coco det, 1333x800
    # same as MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
    if cfg.task == 'coco_det':
        transform = transforms.Compose([
            transforms.RandomGrayscale(p=0.2),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406),
                                 (0.229, 0.224, 0.225))])

    # COCO caption dataset
    coco = CocoDataset(root=cfg.image_dir,
                       json=cfg.caption_path,
                       vocab=vocab,
                       transform=transform)
    #Build data loader for image caption training
    if cfg.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(coco)
    else:
        train_sampler = None

    # Data loader for COCO dataset
    # This will return (images, captions, lengths) for each iteration.
    # images: a tensor of shape (batch_size, 3, 224, 224).
    # captions: a tensor of shape (batch_size, padded_length).
    # lengths: a list indicating valid length for each caption. length is (batch_size).
    data_loader = torch.utils.data.DataLoader(dataset=coco,
                                              batch_size=cfg.batch_size,
                                              shuffle=(train_sampler is None),
                                              num_workers=cfg.num_workers,
                                              collate_fn=collate_fn,
                                              pin_memory=True,
                                              sampler=train_sampler)

    # Build the Decoder models
    decoder = DecoderRNN(cfg.model['embed_size'], cfg.model['hidden_size'], len(vocab), cfg.model['num_layers'])

    if cfg.model['net'] == 'densenet121':
        linear_ic = nn.Linear(1024, 256)
        bn_ic = nn.BatchNorm1d(256, momentum=0.01)
        net = DenseNet121()

    if cfg.model['net'] == 'densenet169':
        linear_ic = nn.Linear(4096, 256)
        bn_ic = nn.BatchNorm1d(256, momentum=0.01)
        net = DenseNet169()

    if cfg.model['net'] == 'resnet34':
        linear_ic = nn.Linear(512, 256)
        bn_ic = nn.BatchNorm1d(256, momentum=0.01)
        net = ResNet34()

    if cfg.model['net'] == 'resnet50':
        linear_ic = nn.Linear(2048, 256)
        bn_ic = nn.BatchNorm1d(256, momentum=0.01)
        net = ResNet50()

    if cfg.model['net'] == 'resnet101':
        linear_ic = nn.Linear(2048, 256)
        bn_ic = nn.BatchNorm1d(256, momentum=0.01)
        net = ResNet101()

    print('cfg.distributed:', cfg.distributed)
    if cfg.distributed:
        linear_ic.cuda()
        bn_ic.cuda()
        net.cuda()
        decoder.cuda()
        # DistributedDataParallel will divide and allocate batch_size to all
        # available GPUs if device_ids are not set
        linear_ic = torch.nn.parallel.DistributedDataParallel(linear_ic)
        bn_ic = torch.nn.parallel.DistributedDataParallel(bn_ic)
        net = torch.nn.parallel.DistributedDataParallel(net)
        decoder = torch.nn.parallel.DistributedDataParallel(decoder)
    else:
        torch.cuda.set_device(device)
        linear_ic.cuda(cfg.gpu)
        bn_ic.cuda(cfg.gpu)
        net.cuda(cfg.gpu)
        decoder.cuda(cfg.gpu)

    criterion = nn.CrossEntropyLoss()
    # Optimizer for image classificaation
    # optimizer = optim.Adam(list(net.parameters()), lr=cfg.lr)

    optimizer_ic = optim.Adam(list(net.parameters()) + list(linear_ic.parameters())+ list(decoder.parameters())+list(bn_ic.parameters()), lr=cfg.lr) #0.0001
    scheduler = MultiStepLR(optimizer_ic, milestones=[60,120,160], gamma=0.1)

    if cfg.loading:
        # Load checkpoint.
        print('==> Resuming from checkpoint..')
        # assert os.path.isdir(cfg.checkpoint), 'Error: no checkpoint directory found!'
        checkpoint = torch.load(cfg.checkpoint)
        net.load_state_dict(checkpoint)
        # best_acc = checkpoint['acc']
        start_epoch = int(cfg.checkpoint.split('/')[-1].split('-')[1])
    else:
        start_epoch = 0

    #scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_ic, T_max=200)
    log_dir = 'log/' + cfg.config.split('/')[1][:-3]
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    writer = SummaryWriter(log_dir=log_dir)
    #start training
    for epoch in range(start_epoch,cfg.num_epochs):
        if cfg.distributed:
            train_sampler.set_epoch(epoch)
        net=train_ic(epoch, cfg, net=net, decoder=decoder,linear=linear_ic,bn=bn_ic, optimizer_ic=optimizer_ic,criterion=criterion,data_loader=data_loader,writer=writer)
        scheduler.step()

if __name__ == '__main__':
    main()


from enum import Enum
import operator
import itertools
import collections
import sys
import numpy as np
import re

from .alignment import AlnStats

class Operation(Enum):
    AlnMatch = 0
    Insertion = 1
    Deletion = 2
    Skip = 3
    Soft = 4
    Hard = 5
    Padding = 6
    SeqMatch = 7
    SeqMismatch = 8

    def consumes_read(self):
        return _consumes_read[self.value]

    def consumes_ref(self):
        return _consumes_ref[self.value]

    def consumes_both(self):
        return _consumes_both[self.value]

    def __str__(self):
        return _letters[self.value]

_letters = 'MIDNSHP=X'
_operations = dict((letter, Operation(num)) for num, letter in enumerate(_letters))
_consumes_read = [True, True, False, False, True, False, False, True, True]
_consumes_ref = [True, False, True, True, False, False, False, True, True]
_consumes_both = list(itertools.starmap(operator.and_, zip(_consumes_read, _consumes_ref)))

AlignedRegion = collections.namedtuple('AlignedRegion', 'start1 end1 start2 end2')

class Cigar:
    def __init__(self, line):
        self._tuples = []
        current = 0
        for char in line:
            if char.isdigit():
                current = current * 10 + int(char)
            else:
                self._tuples.append((current, _operations[char]))
                current = 0
        self._tuples = tuple(self._tuples)
        self.__init_lengths()
        self._index = None

    @classmethod
    def from_tuples(cls, tuples, lengths_from=None):
        self = cls.__new__(cls)
        self._tuples = tuple(tuples)
        if lengths_from:
            self._aligned_len = lengths_from._aligned_len
            self._ref_len = lengths_from._ref_len
            self._read_len = lengths_from._read_len
        else:
            self.__init_lengths()
        self._index = None
        return self

    @classmethod
    def from_pysam_tuples(cls, tuples):
        self = cls.__new__(cls)
        self._tuples = tuple((length, Operation(op)) for op, length in tuples)
        self.__init_lengths()
        self._index = None
        return self

    def init_index(self):
        self._index = CigarIndex(self)

    def init_proxy_index(self):
        """
        Initializes fake proxy index, that does not do anything, but allows to run functions such as aligned_region().
        """
        self._index = ProxyIndex()

    def __init_lengths(self):
        self._read_len = 0
        self._ref_len = 0
        self._aligned_len = 0
        for len, op in self._tuples:
            if op.consumes_read():
                self._read_len += len
            if op.consumes_ref():
                self._ref_len += len
            if op.consumes_both():
                self._aligned_len += len

    def reversed(self):
        """
        Returns reversed cigar.
        """
        return Cigar.from_tuples(self._tuples[::-1], self)

    def __iter__(self):
        return iter(self._tuples)

    def iter_from(self, start_ix):
        return itertools.islice(self._tuples, start_ix, len(self._tuples))

    @property
    def read_len(self):
        return self._read_len

    @property
    def ref_len(self):
        return self._ref_len

    @property
    def aligned_len(self):
        return self._aligned_len

    @property
    def index(self):
        return self._index

    """
    Returns i-th pair (length: int, operation: Operation).
    """
    def __getitem__(self, i):
        return self._tuples[i]

    """
    Finds aligned region in two sequences.

    Parameters:
    - ref_start: 0-based start of the region in the first sequence (should be >= 0),
    - ref_end: 0-based exclusive end of the region in the first sequence (should be <= ref_len),
    - alt_size_diff: difference between reference and alternative variants. Positive if alternative is longer.

    Returns pair `read_start, read_end`, output interval may be empty.

    TODO:
    Process case when the variant overlaps homopolymer: for example
        chr1:160696714-160696714        catatGtcatg
        chr22:18944546-18944546         catatGttcat
    can be stored as GT-, GTT or as G-T, GTT. We need to account for both cases.
    """
    def aligned_region(self, ref_start, ref_end, alt_size_diff=0):
        cigar_start, ref_pos, read_pos = self._index.find_by_ref(ref_start)
        read_start = None

        for cigar_ix in range(cigar_start, len(self._tuples)):
            length, op = self._tuples[cigar_ix]
            if op.consumes_ref():
                if read_start is None and ref_pos <= ref_start < ref_pos + length:
                    # Found start in the first sequence.
                    if op.consumes_read():
                        read_start = read_pos + ref_start - ref_pos
                    else:
                        # this position is deleted in the second sequence.
                        read_start = read_pos

                if ref_pos < ref_end <= ref_pos + length:
                    # Found end in the first sequence.
                    if op.consumes_read():
                        read_end = read_pos + ref_end - ref_pos
                    else:
                        read_end = read_pos

                    if alt_size_diff > 0 and ref_end - ref_pos == length and cigar_ix + 1 < len(self._tuples):
                        # Include case when the alternative variant is in the insertion in the second sequence.
                        next_len, next_op = self._tuples[cigar_ix + 1]
                        if not next_op.consumes_ref():
                            read_end += min(next_len, alt_size_diff)

                    assert read_start is not None
                    return (read_start, read_end)
                ref_pos += length

            if op.consumes_read():
                read_pos += length
        assert False

    """
    Should be extended CIGAR: with X/= instead of M.
    Returns iterator, each element: `AlignedRegion`.
    """
    def find_differences(self):
        pos1 = 0
        pos2 = 0
        curr = None
        for length, op in self._tuples:
            if op.consumes_both():
                if op == Operation.SeqMismatch:
                    if curr:
                        yield curr
                    curr = AlignedRegion(start1=pos1, end1=pos1 + length, start2=pos2, end2=pos2 + length)
                elif op == Operation.AlnMatch:
                    raise ValueError('Cannot call find_differences() on non-extended CIGAR')

                pos1 += length
                pos2 += length

            # Both positions should be non-zero
            elif op.consumes_ref():
                assert pos1 and pos2
                if curr:
                    if curr.start1 < pos1 - 1:
                        yield curr
                yield AlignedRegion(start1=pos1 - 1, end1=pos1 + length, start2=pos2 - 1, end2=pos2)
                curr = None
                pos1 += length

            elif op.consumes_read():
                assert pos1 and pos2
                if curr:
                    if curr.start1 < pos1 - 1:
                        yield curr
                yield AlignedRegion(start1=pos1 - 1, end1=pos1, start2=pos2 - 1, end2=pos2 + length)
                curr = None
                pos2 += length
        if curr:
            yield curr

    def remove_clipping(self):
        """
        Returns tuple `(left_clipping, right_clipping, new_cigar)`.
        """
        left = 0
        right = 0
        new_tuples = []

        for length, op in self._tuples:
            if op == Operation.Soft or op == Operation.Hard:
                if new_tuples:
                    right += length
                else:
                    left += length
            else:
                new_tuples.append((length, op))
        return left, right, Cigar.from_tuples(new_tuples)

    def get_clipping(self):
        left_len, left_op = self._tuples[0]
        right_len, right_op = self._tuples[-1]
        # Do not expect to see hard clipping.
        assert left_op.consumes_read() and right_op.consumes_read()

        left = 0 if left_op.consumes_ref() else left_len
        right = 0 if right_op.consumes_ref() else right_len
        return left, right

    def __str__(self):
        return ''.join('%d%s' % t for t in self._tuples)

    def to_str(self, delim=' '):
        return delim.join('%d%s' % t for t in self._tuples)

    def __len__(self):
        return len(self._tuples)

    def to_full(self, ref_seq, read_seq):
        new_tuples = []
        pos1 = 0
        pos2 = 0

        for length, op in self._tuples:
            if op.consumes_both():
                curr_len = 0
                curr_op = Operation.SeqMatch
                for i in range(length):
                    new_op = Operation.SeqMatch if read_seq[pos1 + i] == ref_seq[pos2 + i] else Operation.SeqMismatch
                    if curr_op == new_op:
                        curr_len += 1
                    else:
                        if curr_len:
                            new_tuples.append((curr_len, curr_op))
                        curr_op = new_op
                        curr_len = 1
                new_tuples.append((curr_len, curr_op))
                pos1 += length
                pos2 += length

            elif op.consumes_read():
                new_tuples.append((length, op))
                pos1 += length
            elif op.consumes_ref():
                new_tuples.append((length, op))
                pos2 += length

        return Cigar.from_tuples(new_tuples, lengths_from=self)

    def to_extended_with_md(self, read_seq, md_tag):
        """
        Returns pair (extended CIGAR, reference_sequence).
        """
        new_tuples = []
        read_pos = 0
        ref_pos = 0
        ref_seq = ''
        # Use [A-Z^] instead of [0-9] because tag starts and ends with a number.
        # Parentheses needed to capture [A-Z^] as well as numbers.
        md_tag = re.split(r'([A-Z^]+)', md_tag)
        for i in range(0, len(md_tag), 2):
            md_tag[i] = int(md_tag[i])
        # Skip first, if it is == 0.
        md_ix = int(md_tag[0] == 0)
        md_shift = 0

        for length, op in self._tuples:
            if op.consumes_both():
                while length > 0:
                    # Sequence
                    if md_ix % 2:
                        subseq = md_tag[md_ix]
                        entry_len = len(subseq) - md_shift
                        pos_inc = min(length, entry_len)
                        new_op = Operation.SeqMismatch

                        ref_seq += subseq[md_shift : md_shift + pos_inc]
                        if entry_len > length:
                            md_shift += length
                        else:
                            md_ix += 1 + int(md_tag[md_ix + 1] == 0)
                            md_shift = 0

                    # Number
                    else:
                        new_op = Operation.SeqMatch
                        entry_len = md_tag[md_ix] - md_shift
                        pos_inc = min(length, entry_len)
                        ref_seq += read_seq[read_pos : read_pos + pos_inc]

                        if entry_len > length:
                            md_shift += length
                        else:
                            md_ix += 1
                            md_shift = 0

                    ref_pos += pos_inc
                    read_pos += pos_inc
                    if new_tuples and new_tuples[-1][1] == new_op:
                        new_tuples[-1] = (pos_inc + new_tuples[-1][0], new_op)
                    else:
                        new_tuples.append((pos_inc, new_op))
                    length -= pos_inc

            elif op.consumes_read():
                read_pos += length
                new_tuples.append((length, op))

            elif op.consumes_ref():
                assert md_tag[md_ix][0] == '^'
                ref_seq += md_tag[md_ix][1:]
                # Add 2 if next entry is 0. We should not get out of the list because it should end with a number.
                md_ix += 1 + int(md_tag[md_ix + 1] == 0)

                ref_pos += length
                new_tuples.append((length, op))

        assert md_ix == len(md_tag) and read_pos == self._read_len \
            and ref_pos == self._ref_len and len(ref_seq) == ref_pos
        return Cigar.from_tuples(new_tuples, lengths_from=self), ref_seq

    def calculate_stats(self):
        """
        Returns `AlnStats`. The Cigar should be full.
        """
        stats = AlnStats()
        for length, op in self._tuples:
            if op.consumes_both():
                if op == Operation.SeqMismatch:
                    stats.add_mismatches(length)
                else:
                    assert op == Operation.SeqMatch
            elif op.consumes_read() and op != Operation.Soft:
                stats.add_insertions(length)
            elif op.consumes_ref():
                stats.add_deletions(length)
        stats.update_from_cigar(self)
        return stats

    def calculate_stats_with_seqs(self, read_seq, ref_seq):
        """
        Returns `AlnStats`. The Cigar can contain M operations as well as X and =.
        """
        stats = AlnStats()
        pos1 = 0
        pos2 = 0
        for length, op in self._tuples:
            if op.consumes_both():
                stats.add_mismatches(sum(read_seq[pos1 + i] != ref_seq[pos2 + i] for i in range(length)))
                pos1 += length
                pos2 += length
            elif op.consumes_read():
                if op != Operation.Soft:
                    stats.add_insertions(length)
                pos1 += length
            elif op.consumes_ref():
                stats.add_deletions(length)
                pos2 += length
        stats.update_from_cigar(self)
        return stats

    def no_gaps(self):
        return all(op.consumes_both() for _, op in self._tuples[1:-1])

    def to_short(self):
        new_tuples = []
        curr_match = 0
        for length, op in self._tuples:
            if op.consumes_both():
                curr_match += length
            else:
                if curr_match:
                    new_tuples.append((curr_match, Operation.AlnMatch))
                    curr_match = 0
                new_tuples.append((length, op))
        if curr_match:
            new_tuples.append((curr_match, Operation.AlnMatch))
        return Cigar.from_tuples(new_tuples, lengths_from=self)

    def to_pysam_tuples(self):
        return [(op.value, len) for len, op in self._tuples]

    def aligned_seqs(self, read_seq, ref_seq):
        read_res = ''
        ref_res = ''

        pos1 = 0
        pos2 = 0
        for length, op in self._tuples:
            if op.consumes_both():
                subs1 = read_seq[pos1 : pos1 + length]
                subs2 = ref_seq[pos2 : pos2 + length]
                if op == Operation.SeqMismatch:
                    read_res += subs1.lower()
                    ref_res += subs2.lower()
                    if subs1 == subs2:
                        sys.stderr.write('Error: operation %dX at positions %d,%d corresponds to sequence match %s\n'
                            % (length, pos1, pos2, subs1))
                else:
                    read_res += subs1
                    ref_res += subs2
                    if op == Operation.SeqMatch and subs1 != subs2:
                        sys.stderr.write('Error: operation %d= at positions %d,%d corresponds to sequence '
                            'mismatch %s != %s\n' % (length, pos1, pos2, subs1, subs2))
                pos1 += length
                pos2 += length

            elif op.consumes_read():
                ref_res += '-' * length
                read_res += read_seq[pos1 : pos1 + length]
                pos1 += length
            elif op.consumes_ref():
                ref_res += ref_seq[pos2 : pos2 + length]
                read_res += '-' * length
                pos2 += length

        if pos1 != len(read_seq) or pos2 != len(ref_seq):
            sys.stderr.write('Error: CIGAR length does not match sequences length: '
                'read_len = %d, read_seq = %d, ref_len = %d, ref_seq = %d\n'
                % (pos1, len(read_seq), pos2, len(ref_seq)))
        return read_res, ref_res

    def aligned_pairs(self, ref_start=0, read_start=0):
        """
        Returns iterator over pairs of indices, aligned to each other.
        """
        assert not (ref_start and read_start)
        if ref_start:
            cigar_start, ref_pos, read_pos = self._index.find_by_ref(ref_start)
        elif read_start:
            cigar_start, ref_pos, read_pos = self._index.find_by_read(read_start)
        else:
            cigar_start = ref_pos = read_pos = 0

        for length, op in itertools.islice(self._tuples, cigar_start, len(self._tuples)):
            if op.consumes_both():
                for j in range(max(ref_start - ref_pos, 0), length):
                    yield (read_pos + j, ref_pos + j)
                read_pos += length
                ref_pos += length
            elif op.consumes_read():
                read_pos += length
            elif op.consumes_ref():
                ref_pos += length

    def read_region(self, ref_start, ref_end):
        cigar_start, ref_pos, read_pos = self._index.find_by_ref(ref_start)
        read_start = None

        for length, op in itertools.islice(self._tuples, cigar_start, len(self._tuples)):
            if ref_pos == ref_end:
                assert read_start is not None
                return read_start, read_pos

            cons_read = op.consumes_read()
            if op.consumes_ref():
                if ref_pos <= ref_start < ref_pos + length:
                    assert read_start is None
                    read_start = read_pos + (ref_start - ref_pos if cons_read else 1)
                if ref_pos <= ref_end < ref_pos + length:
                    assert read_start is not None
                    read_end = read_pos + (ref_end - ref_pos if cons_read else 1)
                    return read_start, read_end
                ref_pos += length

            if cons_read:
                read_pos += length

        assert ref_end == ref_pos
        assert read_start is not None
        return read_start, read_pos

    def subcigar(self, ref_start, ref_end):
        """
        Returns tuple (ref_start, ref_end, read_start, read_end, subcigar).
        Note, that if the first or last operation within the region of interest is not M,
        returned ref_start or ref_end may be different from input ref_start or ref_end.

        If the resulting alignment will be empty, returns None.
        """
        assert ref_start >= 0 and ref_end <= self._ref_len
        cigar_start, ref_pos, read_pos = self._index.find_by_ref(ref_start)
        read_start = None
        new_tuples = []

        for length, op in itertools.islice(self._tuples, cigar_start, len(self._tuples)):
            if ref_pos >= ref_end:
                assert ref_pos == ref_end
                break

            cons_ref = op.consumes_ref()
            cons_read = op.consumes_read()
            if cons_ref and cons_read:
                if ref_pos <= ref_start < ref_pos + length:
                    assert read_start is None
                    read_start = read_pos + ref_start - ref_pos
                intersection = min(ref_end, ref_pos + length) - max(ref_start, ref_pos)
                if intersection > 0:
                    new_tuples.append((intersection, op))
                if ref_pos < ref_end <= ref_pos + length:
                    read_pos += ref_end - ref_pos
                    break

                ref_pos += length
                read_pos += length

            elif cons_ref:
                if ref_pos <= ref_start < ref_pos + length:
                    # Alignment starts with Deletion.
                    ref_start = ref_pos + length
                    if ref_start >= ref_end:
                        return None
                elif ref_pos < ref_end <= ref_pos + length:
                    # Alignment ends with Deletion.
                    ref_end = ref_pos
                    break
                elif ref_pos >= ref_start:
                    new_tuples.append((length, op))
                ref_pos += length

            elif cons_read:
                if read_start is not None:
                    new_tuples.append((length, op))
                read_pos += length

        read_end = read_pos
        new_cigar = Cigar.from_tuples(new_tuples)
        assert new_cigar.ref_len == ref_end - ref_start
        assert new_cigar.read_len == read_end - read_start
        return ref_start, ref_end, read_start, read_end, new_cigar

    def revert(self, strand):
        new_tuples = []
        for length, op in self._tuples:
            if op.consumes_both():
                new_tuples.append((length, op))
            elif op == Operation.Insertion:
                new_tuples.append((length, Operation.Deletion))
            elif op == Operation.Deletion:
                new_tuples.append((length, Operation.Insertion))
            else:
                raise RuntimeError('Encountered operation %s while trying to revert cigar "%s"' % (op, self.to_str()))
        res = Cigar.from_tuples(new_tuples if strand else reversed(new_tuples), self)
        res._read_len = self._ref_len
        res._ref_len = self._read_len
        return res

    def __eq__(self, other):
        return self._tuples == other._tuples

    def has_true_clipping(self, base_qualities, max_len=2, low_bq=10):
        """
        Suppose a read has C clipped bases and there are B bases with quality <= low_bq (on the same side of the read).
        Returns True if C > B + max_len (for left OR for right side).
        """
        left, right = self.get_clipping()
        if base_qualities is None:
            return left > max_len or right > max_len

        if left:
            low_bq_left = sum(bq <= low_bq for bq in base_qualities[:left])
            if left > low_bq_left + max_len:
                return True
        if right:
            low_bq_right = sum(bq <= low_bq for bq in base_qualities[-right:])
            if right > low_bq_right + max_len:
                return True
        return False

class CigarIndex:
    def __init__(self, cigar, step_size=100):
        """
        Saves indices of CIGAR tuples to allow faster search within duplication.
        """
        cigar_indices = []
        ref_positions = []
        read_positions = []

        prev_ref_pos = ref_pos = 0
        prev_read_pos = read_pos = 0

        for cigar_ix, (length, op) in enumerate(cigar):
            if ref_pos > prev_ref_pos + step_size or read_pos > prev_read_pos + step_size:
                cigar_indices.append(cigar_ix)
                ref_positions.append(ref_pos)
                read_positions.append(read_pos)
                prev_ref_pos = ref_pos
                prev_read_pos = read_pos

            if op.consumes_ref():
                ref_pos += length
            if op.consumes_read():
                read_pos += length
        self._cigar_indices = tuple(cigar_indices)
        self._ref_positions = tuple(ref_positions)
        self._read_positions = tuple(read_positions)

    def find_by_ref(self, ref_pos):
        """
        Returns tuple (cigar_ix, ref_pos', read_pos') with the biggest cigar_ix such that ref_pos' <= ref_pos.
        """
        i = np.searchsorted(self._ref_positions, ref_pos, side='right') - 1
        if i >= 0:
            return (self._cigar_indices[i], self._ref_positions[i], self._read_positions[i])
        return (0, 0, 0)

    def find_by_read(self, read_pos):
        """
        Returns tuple (cigar_ix, ref_pos', read_pos') with the biggest cigar_ix such that read_pos' <= read_pos.
        """
        i = np.searchsorted(self._read_positions, read_pos, side='right') - 1
        if i >= 0:
            return (self._cigar_indices[i], self._ref_positions[i], self._read_positions[i])
        return (0, 0, 0)

class ProxyIndex:
    def __init__(self):
        pass

    def find_by_ref(self, _ref_pos):
        return (0, 0, 0)

    def find_by_read(self, read_pos):
        return (0, 0, 0)


#! /usr/bin/env python
import numpy as np
import rospy
from DeltaGripper import DeltaGripper
from sensor_msgs.msg import Joy

# This ROS Node converts Joystick inputs from the joy node
# into commands for turtlesim or any other robot

# Receives joystick messages (subscribed to Joy topic)
# then converts the joysick inputs into Twist commands
# axis 1 aka left stick vertical controls linear speed
# axis 0 aka left stick horizonal controls angular speed
def callback(data):
    left_joy_stick_x = -data.axes[0]
    left_joy_stick_y = data.axes[1]
    left_up_trigger_z = data.buttons[4]
    left_down_trigger_z = -data.axes[3]
    right_joy_stick_x = data.axes[2]
    right_joy_stick_y = data.axes[5]
    right_up_trigger_z = data.buttons[5]
    right_down_trigger_z = -data.axes[4]

    bad_motor_up_button = data.buttons[3]
    bad_motor_down_button = data.buttons[1]

    delta_velocity = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).reshape((1,6))

    if(left_down_trigger_z > 0.8):
        delta_velocity[0][3] = 0.005
        delta_velocity[0][4] = 0.005
        delta_velocity[0][5] = 0.005
    elif(left_up_trigger_z > 0.8):
        delta_velocity[0][3] = -0.005
        delta_velocity[0][4] = -0.005
        delta_velocity[0][5] = -0.005
    if(right_down_trigger_z > 0.8):
        delta_velocity[0][0] = 0.005
        delta_velocity[0][1] = 0.005
        delta_velocity[0][2] = 0.005
    elif(right_up_trigger_z > 0.8):
        delta_velocity[0][0] = -0.005
        delta_velocity[0][1] = -0.005
        delta_velocity[0][2] = -0.005

    if(left_joy_stick_x < -0.8):
        delta_velocity[0][3] = 0.005
        delta_velocity[0][4] = 0.005
        delta_velocity[0][5] = -0.005
    elif(left_joy_stick_x > 0.8):
        delta_velocity[0][3] = -0.005
        delta_velocity[0][4] = -0.005
        delta_velocity[0][5] = 0.005
    if(right_joy_stick_x < -0.8):
        delta_velocity[0][0] = -0.005
        delta_velocity[0][1] = 0.005
        delta_velocity[0][2] = 0.005
    elif(right_joy_stick_x > 0.8):
        delta_velocity[0][0] = 0.005
        delta_velocity[0][1] = -0.005
        delta_velocity[0][2] = -0.005

    if(left_joy_stick_y < -0.8):
        delta_velocity[0][3] = -0.005
        delta_velocity[0][4] = 0.005
    elif(left_joy_stick_y > 0.8):
        delta_velocity[0][3] = 0.005
        delta_velocity[0][4] = -0.005
    if(right_joy_stick_y < -0.8):
        delta_velocity[0][1] = 0.005
        delta_velocity[0][2] = -0.005
    elif(right_joy_stick_y > 0.8):
        delta_velocity[0][1] = -0.005
        delta_velocity[0][2] = 0.005

    if(bad_motor_up_button > 0.8):
        delta_velocity[0][3] = -0.005
    elif(bad_motor_down_button > 0.8):
        delta_velocity[0][3] = 0.005

    for i in range(6):
        if delta_velocity[0][i] != 0:
            delta_gripper.move_delta_velocity(delta_velocity, [0.2])
            break

# Intializes everything
def start():
    # publishing to "turtle1/cmd_vel" to control turtle1

    #rospy.init_node('DeltaTeleop')
    global delta_gripper
    delta_gripper = DeltaGripper()
    # subscribed to joystick inputs on topic "joy"
    rospy.Subscriber("/joy", Joy, callback, queue_size=1)
    # starts the node
    rospy.spin()

if __name__ == '__main__':
    start()