#!/usr/bin/env python
"""
just write from scratch ...

so, we need:
- an image encoder model. we'll use something like a lenet-5 cnn
- sender network, that takes the encoding as input (?)
- receiver network
    - ... and we will dot product the output of the receiver with the various encoded images
"""
import argparse, time, os, math, json, sys, contextlib, random

import torch
from torch import autograd, nn, optim
import torch.nn.functional as F
import numpy as np

from ulfs.utils import die
from ulfs.params import Params
from ulfs import name_utils, rl_common, tensor_utils,  utils, nn_modules, metrics
from ulfs.runner_base_v1 import RunnerBase
from ulfs.stats import Stats

from data_code.clevr.three_shapes_dataset import Dataset
from ilm.cnn_models import *

@contextlib.contextmanager
def torch_random_state():
    rnd_tch = torch.get_rng_state()
    rnd_cuda = torch.cuda.get_rng_state()
    np_state = np.random.get_state()
    rand_state = random.getstate()
    yield
    torch.set_rng_state(rnd_tch)
    torch.cuda.set_rng_state(rnd_cuda)
    np.random.set_state(np_state)
    random.setstate(rand_state)

class LangSenderModel(nn.Module):
    """
    generator model
    we'll use a differentiable teacher-forcing model, with fixed utterance length
    """
    def __init__(
            self, opt_name, embedding_size, vocab_size, utt_len, rnn_type, num_layers, input_size,
            dropout
        ):
        self.embedding_size = embedding_size
        self.utt_len = utt_len
        self.vocab_size = vocab_size
        self.rnn_type = rnn_type
        self.num_layers = num_layers

        super().__init__()
        self.h_in = nn.Linear(input_size, embedding_size)
        if rnn_type == 'SRU':
            from sru import SRU
            RNN = SRU
        else:
            RNN = getattr(nn, f'{rnn_type}')
        rnn_params = {
            'input_size': embedding_size,
            'hidden_size': embedding_size,
            'num_layers': num_layers,
            'dropout': dropout
        }
        if rnn_type == 'SRU':
            rnn_params['rescale'] = False
            rnn_params['use_tanh'] = True
        self.rnn = RNN(**rnn_params)
        self.h_out = nn.Linear(embedding_size, vocab_size)
        self.drop = nn.Dropout(dropout)

        Opt = getattr(optim, opt_name)
        self.opt = Opt(lr=0.001, params=self.parameters())

    def forward(self, thoughts):
        """
        thoughts ae [N][input_size]
        """
        N, K = thoughts.size()
        embs = self.h_in(thoughts)
        embs = self.drop(embs)
        device = embs.device

        if self.rnn_type in ['SRU']:
            h = torch.zeros(self.num_layers, N, self.embedding_size, dtype=torch.float32, device=device)
            h[0] = embs
        elif self.rnn_type in ['GRU']:
            h = torch.zeros(self.num_layers, N, self.embedding_size, dtype=torch.float32, device=device)
            h[0] = embs
        else:
            raise Exception(f'unrecognized rnn type {self.rnn_type}')

        fake_input = torch.zeros(self.utt_len, N, self.embedding_size, dtype=torch.float32, device=device)
        if self.rnn_type == 'SRU':
            output, state = self.rnn(fake_input, h)
        elif self.rnn_type in ['GRU']:
            output, h = self.rnn(fake_input, h)
        else:
            raise Exception(f'rnn type {self.rnn_type} not recognized')
        utts = self.h_out(output)
        return utts

class LangReceiverModel(nn.Module):
    def __init__(
            self, opt_name, embedding_size, vocab_size, utt_len,
            rnn_type, dropout, num_layers, output_size
        ):
        self.rnn_type = rnn_type
        self.embedding_size = embedding_size

        super().__init__()
        self.embedding = nn_modules.EmbeddingAdapter(vocab_size, embedding_size)
        if rnn_type == 'SRU':
            from sru import SRU
            RNN = SRU
        else:
            RNN = getattr(nn, f'{rnn_type}')
        self.rnn = RNN(
            input_size=embedding_size,
            hidden_size=embedding_size,
            num_layers=num_layers
        )
        self.h_out = nn.Linear(embedding_size, output_size)
        Opt = getattr(optim, opt_name)
        self.opt = Opt(lr=0.001, params=self.parameters())

    def forward(self, utts, do_predict_correct=False):
        embs = self.embedding(utts)
        seq_len, batch_size, embedding_size = embs.size()
        output, state = self.rnn(embs)
        state = state[-1]
        x = self.h_out(state)
        return x

class SenderPathway(nn.Module):
    def __init__(self, opt_name, cnn, lang_sender, clip_grad):
        super().__init__()
        self.cnn = cnn
        self.lang_sender = lang_sender
        self.clip_grad = clip_grad

        Opt = getattr(optim, opt_name)
        self.params = self.parameters()
        self.opt = Opt(lr=0.001, params=self.params)

    def forward(self, images_t):
        """
        assumes that images_t just contains a batch of single images, not multipel images per example
        """
        image_enc = self.cnn(images_t)
        utt_logits = self.lang_sender(image_enc)
        return utt_logits

    def sup_train_batch(self, images, utts):
        utts_logits = self(images)

        # logits = self.model(meanings)
        _, utts_pred = utts_logits.max(dim=-1)
        correct = utts_pred == utts
        acc = correct.float().mean().item()

        crit = nn_modules.GeneralCrossEntropyLoss()
        loss = crit(utts_logits, utts)
        self.opt.zero_grad()
        loss.backward()
        if self.clip_grad is not None and self.clip_grad > 0:
            torch.nn.utils.clip_grad_norm_(self.params, self.clip_grad)
        self.opt.step()
        return loss.item(), acc

class ReceiverPathway(nn.Module):
    def __init__(self, opt_name, cnn, lang_receiver, clip_grad):
        super().__init__()
        self.cnn = cnn
        self.lang_receiver = lang_receiver
        self.clip_grad = clip_grad

        Opt = getattr(optim, opt_name)
        self.params = self.parameters()
        self.opt = Opt(lr=0.001, params=self.params)

    def forward(self, utts, images):
        """
        we assume images are all the receiver images, ie the goal image, and some distractors
        the sender image is not included in images
        """
        lang_enc = self.lang_receiver(utts)  # [N][E]

        d = images.size()
        images_flat_t = tensor_utils.merge_dims(images, 0, 1)  # [M * N][C][H][W]
        images_enc_flat = self.cnn(images_flat_t)   # [M * N][E]
        images_enc = tensor_utils.split_dim(images_enc_flat, 0, d[0], d[1])  # [M][N][E]
        lang_enc_flat_exp = lang_enc.unsqueeze(0).expand_as(images_enc)  # [M][N][E]
        lang_enc_flat_exp = tensor_utils.merge_dims(lang_enc_flat_exp.contiguous(), 0, 1)   # [M * N][E]

        dp_left = lang_enc_flat_exp.unsqueeze(-2)  # [M * N][1][E]
        dp_right = images_enc_flat.unsqueeze(-1)   # [M * N][E][1]
        dp = torch.bmm(dp_left, dp_right)     # [M * N][1][1]
        dp = dp.view(d[0], d[1])   # [M][N]
        dp = dp.transpose(0, 1)   # [N][M]
        return dp

    def sup_train_batch(self, images, utts):
        """
        we assume a batch of single images, no distractors added
        """
        d = images.size()
        with autograd.no_grad():
            images_enc = self.cnn(images)   # [N][E]

        lang_enc = self.lang_receiver(utts)  # [N][E]
        crit = nn.MSELoss()
        loss = crit(lang_enc, images_enc)
        self.opt.zero_grad()
        loss.backward()
        if self.clip_grad is not None and self.clip_grad > 0:
            torch.nn.utils.clip_grad_norm_(self.params, self.clip_grad)
        self.opt.step()
        acc = 0  # placeholder
        return loss.item(), acc

class Agent(nn.Module):
    def __init__(self, p, image_size, img_embedding_size):
        super().__init__()
        self.p = p
        self.img_embedding_size = img_embedding_size

        CNN = globals()[p.conv_class]
        # share cnn between sender and receiver for now (can try unsharing later)
        self.cnn = CNN(dropout=p.dropout, num_layers=p.num_conv_layers, image_size=image_size)
        self.lang_sender = LangSenderModel(
            opt_name=p.opt,
            embedding_size=p.embedding_size,
            vocab_size=p.vocab_size,
            utt_len=p.utt_len,
            rnn_type=p.rnn_type,
            num_layers=p.num_layers,
            input_size=self.img_embedding_size,
            dropout=p.dropout
        )
        self.lang_receiver = LangReceiverModel(
            opt_name=p.opt,
            embedding_size=p.embedding_size,
            vocab_size=p.vocab_size,
            utt_len=p.utt_len,
            rnn_type=p.rnn_type,
            dropout=p.dropout,
            num_layers=p.num_layers,
            output_size=self.img_embedding_size
        )
        self.sender_pathway = SenderPathway(opt_name=p.opt, cnn=self.cnn, lang_sender=self.lang_sender, clip_grad=p.clip_grad)
        self.receiver_pathway = ReceiverPathway(opt_name=p.opt, cnn=self.cnn, lang_receiver=self.lang_receiver, clip_grad=p.clip_grad)
        if p.enable_cuda:
            self.lang_sender = self.lang_sender.cuda()
            self.lang_receiver = self.lang_receiver.cuda()
            self.cnn = self.cnn.cuda()
        Opt = getattr(optim, p.opt)
        self.opt_both = Opt(
            lr=0.001,
            params=list(self.lang_sender.parameters()) + list(self.lang_receiver.parameters()) + list(self.cnn.parameters())
        )

    def state_dict(self):
        return {
            'lang_sender_state': self.lang_sender.state_dict(),
            'lang_receiver_state': self.lang_receiver.state_dict(),
            'cnn_state': self.cnn.state_dict(),
            'opt_both_state': self.opt_both.state_dict()
        }

    def load_state_dict(self, statedict):
        self.lang_sender.load_state_dict(statedict['lang_sender_state'])
        self.lang_receiver.load_state_dict(statedict['lang_receiver_state'])
        self.opt_both.load_state_dict(statedict['opt_both_state'])

class SoftmaxLink(object):
    def __init__(self, p):
        pass

    def sample_utterances(self, utt_probs):
        return utt_probs

    def sample_image_choice(self, dp):
        return dp

    def calc_loss(self, image_choice):
        dp = image_choice
        batch_size, _ = dp.size()
        crit = nn.CrossEntropyLoss()
        loss = crit(dp, torch.zeros(batch_size, dtype=torch.int64, device=dp.device))
        loss_v = loss.item()
        return loss, loss_v

class RLLink(object):
    def __init__(self, p):
        self.s_ent = p.s_ent
        self.r_ent = p.r_ent

    def sample_utterances(self, utt_probs, training):
        # print('rllink training', self.training)
        self.s_sender = rl_common.draw_categorical_sample(
            action_probs=utt_probs,
            batch_idxes=None,
            training=training
        )
        utts = self.s_sender.actions.detach()
        return utts

    def sample_image_choice(self, training, dp):
        self.s_recv = rl_common.draw_categorical_sample(
            action_probs=dp,
            batch_idxes=None,
            training=training
        )
        return (self.s_sender.greedy_matches, self.s_recv.greedy_matches), self.s_recv.actions

    def calc_loss(self, image_choice, training):
        # s_recv = image_choice
        s_recv = self.s_recv
        rewards = (s_recv.actions == 0).float()

        # for reporting purposes:
        loss_v = - rewards.mean().item()

        # lets baseline the reward first
        rewards_mean = rewards.mean().item()
        rewards_std = rewards.std().item()
        rewards = rewards - rewards_mean
        if rewards_std > 1e-1:
            rewards = rewards / rewards_std

        if training:
            rl_loss = self.s_sender.calc_loss(rewards) + s_recv.calc_loss(rewards)
            loss_all = rl_loss
            ent_loss = 0
            if self.s_ent is not None and self.s_ent > 0:
                ent_loss -= self.s_sender.entropy * self.s_ent
            if self.r_ent is not None and self.r_ent > 0:
                ent_loss -= s_recv.entropy * self.r_ent
            loss_all += ent_loss
        else:
            # rl_loss = 0
            loss_all = 0

        loss = loss_all
        return loss, loss_v

class RefTaskGame(object):
    def __init__(self, sender_pathway, receiver_pathway, link):
        self.sender_pathway = sender_pathway
        self.receiver_pathway = receiver_pathway
        self.link = link

    def forward(self, images_t, training):
        utt_logits = self.sender_pathway(images_t[0])
        utt_probs = F.softmax(utt_logits, dim=-1)

        utts = self.link.sample_utterances(utt_probs, training=training)

        self.dp = self.receiver_pathway(utts=utts, images=images_t[1:])
        self.dp = F.softmax(self.dp, dim=-1)

        (sender_greedy, receiver_greedy), pred = self.link.sample_image_choice(training=training, dp=self.dp)

        # _, pred = self.dp.max(dim=-1)
        acc = (pred == 0).float().mean().item()
        return (sender_greedy, receiver_greedy), utts.detach(), acc

    def calc_loss(self, training):
        loss, loss_v = self.link.calc_loss(self.dp, training=training)
        # return (sender_greedy, receiver_greedy), loss, loss_v
        return loss, loss_v

def batched_run_nograd(params, model, inputs, batch_size, input_batch_dim, output_batch_dim):
    """
    assumes N is multiple of batch_size
    """
    p = params
    N = inputs.size(input_batch_dim)
    num_batches = (N + batch_size - 1) // batch_size
    outputs = None
    count = 0
    for b in range(num_batches):
        b_start = b * batch_size
        b_end = min(b_start + batch_size, N)
        count += (b_end - b_start)
        input_batch = inputs.narrow(dim=input_batch_dim, start=b_start, length=b_end - b_start)
        if p.enable_cuda:
            input_batch = input_batch.cuda()
        with torch.no_grad():
            output_batch = model(input_batch).detach().cpu()
        if outputs is None:
            out_size_full = list(output_batch.size())
            out_size_full[output_batch_dim] = N
            outputs = torch.zeros(*out_size_full, dtype=output_batch.dtype, device='cpu')
        out_narrow = outputs.narrow(dim=output_batch_dim, start=b_start, length=b_end - b_start)
        out_narrow[:] = output_batch
    assert count == N
    return outputs

class Runner(RunnerBase):
    def __init__(self):
        super().__init__(
            save_as_statedict_keys=['teacher'],
            additional_save_keys=[],
            step_key='training_step'
        )

    def setup(self, p):
        if p.seed is not None:
            torch.manual_seed(p.seed)
            np.random.seed(p.seed)
            random.seed(p.seed)
            torch.backends.cudnn.deterministic = True
            print('seeding torch and numpy using ', p.seed)

        self.dataset = Dataset(data_dir=p.data_dir)

        # determine image size from dataset
        images_t, _ = self.dataset.sample_batch(batch_size=2)
        self.image_size = images_t.size(-1)
        print('self.image_size', self.image_size)

        CNN = globals()[p.conv_class]

        # determine output size from cnn:
        _cnn = CNN(dropout=p.dropout, num_layers=p.num_conv_layers, image_size=self.image_size)  # share cnn across everything
        print(_cnn.convnet)
        print('convnet output size', _cnn.output_size)
        if p.enable_cuda:
            _cnn = _cnn.cuda()
        if p.enable_cuda:
            images_t = images_t.cuda()
        with autograd.no_grad():
            enc_images_t = _cnn(images_t[0])
        _, self.img_embedding_size = enc_images_t.size()
        print('self.img_embedding_size', self.img_embedding_size)

        self.teacher = Agent(p=p, img_embedding_size=self.img_embedding_size, image_size=self.image_size)

        Link = globals()[f'{p.link}Link']
        self.link = Link(p=p)

        self.sup_train_N = int(self.dataset.N_train * p.sup_train_frac)
        print('sup_train_N', self.sup_train_N)

    def step(self, p):
        training_step = self.training_step
        step = self.training_step
        render = self.should_render()
        link = self.link

        sup_images = self.dataset.sample_images(self.sup_train_N)
        if p.enable_cuda:
            sup_images = sup_images.cuda()

        print('generating teacher utterances...', end='', flush=True)
        _gen_start = time.time()
        self.teacher.eval()
        # with autograd.no_grad():
        #     sup_images_enc = self.teacher.cnn(sup_images).detach()
        sup_utts_logits = batched_run_nograd(
            params=p,
            model=self.teacher.sender_pathway,
            inputs=sup_images,
            batch_size=p.batch_size,
            input_batch_dim=0,
            output_batch_dim=1
        )
        # self.teacher.train()
        # print('sup_utts_logits.size()', sup_utts_logits.size())
        _, sup_utts = sup_utts_logits.max(dim=-1)
        # print('sup_utts.size()', sup_utts.size())
        print(' done in %.0f seconds' % (time.time() - _gen_start))

        student = Agent(p=p, img_embedding_size=self.img_embedding_size, image_size=self.image_size)

        # and then train each half supervised on this data
        # sender first...
        sup_sender_epochs = 0
        sup_receiver_epochs = 0
        sup_sender_acc = 0
        sup_receiver_acc = 0
        sup_sender_time = 0
        sup_receiver_time = 0
        if (
                (p.sup_acc is not None and p.sup_acc > 0)
                or (p.sup_ksteps is not None and p.sup_ksteps > 0)
            ) and (step > 0 or not p.train_e2e):
            for (agent_str, pathway) in [('send', student.sender_pathway), ('recv', student.receiver_pathway)]:
                # print('sup training on' + agent_str)
                _epoch = 0
                _sup_start = time.time()
                _last_print = time.time()
                sup_stats = Stats([
                    'loss_sum',
                    'acc_sum',
                    'episodes_count',
                ])
                while True:
                    b_idxes = torch.from_numpy(np.random.choice(self.sup_train_N, p.batch_size, replace=False))
                    b_utts = sup_utts[:, b_idxes]
                    b_images = sup_images[b_idxes]
                    if p.enable_cuda:
                        b_utts = b_utts.cuda()
                        b_images = b_images.cuda()
                    b_loss, b_acc = pathway.sup_train_batch(images=b_images, utts=b_utts)
                    sup_stats.episodes_count += 1
                    sup_stats.loss_sum += b_loss
                    sup_stats.acc_sum += b_acc

                    _epoch += 1
                    _done_training = False
                    if p.sup_acc is not None and epoch_acc >= p.sup_acc:
                        # print('done sup training (reason: acc)')
                        _done_training = True
                    if p.sup_ksteps is not None and _epoch >= p.sup_ksteps * 1000:
                        # print('done sup training (reason: steps)')
                        _done_training = True
                    if _done_training or time.time() - _last_print >= 30.0:
                        _elapsed_time = time.time() - _sup_start
                        _loss = sup_stats.loss_sum / sup_stats.episodes_count
                        _acc = sup_stats.acc_sum / sup_stats.episodes_count
                        log_dict = {
                            'record_type': f'sup_{agent_str}',
                            'agent': agent_str,
                            'ilm_epoch': step,
                            'epoch': _epoch,
                            'sps': int(_epoch / _elapsed_time),
                            'sup_time': int(_elapsed_time),
                            'loss': _loss,
                            'acc': _acc,
                        }
                        formatstr = (
                            '{record_type} g={ilm_epoch} e={epoch} '
                            't={sup_time:.0f} '
                            'sps={sps:.0f} '
                            'loss={loss:.3f} '
                            'acc={acc:.3f} '
                        )
                        self.print_and_log(log_dict, formatstr=formatstr)
                        sup_stats.reset()
                        _last_print = time.time()
                    if _done_training:
                        # print('done training for pathway', pathway.__class__.__name__)
                        break
                if pathway == student.sender_pathway:
                    sup_sender_epochs = _epoch
                    sup_sender_acc = _acc
                    sup_sender_time = time.time() - _sup_start
                elif pathway == student.receiver_pathway:
                    sup_receiver_epochs = _epoch
                    sup_receiver_acc = _acc
                    sup_receiver_time = time.time() - _sup_start
                else:
                    raise Exception('invalid pathway value')
            # print('done supervised training')

        e2e_time = 0
        if p.train_e2e:
            # then train end to end for a bit, as decoder-encoder, looking at reconstruction accuracy
            # we'll do this on the same meanings as we got from the teacher? or different ones? or
            # just rnadomly sampled from everything except heldout?
            # maybe train on everything except holdout?
            last_print = time.time()
            _e2e_start = time.time()
            e2e_stats = Stats([
                'episodes_count',
                'e2e_loss_sum',
                'e2e_acc_sum',
                'sender_greedy_sum',
                'receiver_greedy_sum',
            ])
            epoch = 0
            student.train()
            ref_task_game = RefTaskGame(
                sender_pathway=student.sender_pathway, receiver_pathway=student.receiver_pathway, link=link)
            student_params = student.parameters()
            while True:
                images_t, labels = self.dataset.sample_batch(batch_size=p.batch_size)
                if p.enable_cuda:
                    images_t, labels = images_t.cuda(), labels.cuda()

                (sender_greedy, receiver_greedy), _, acc = ref_task_game.forward(images_t, training=True)
                loss, loss_v = ref_task_game.calc_loss(training=True)

                e2e_stats.e2e_loss_sum += loss_v
                e2e_stats.e2e_acc_sum += acc
                e2e_stats.sender_greedy_sum += sender_greedy
                e2e_stats.receiver_greedy_sum += receiver_greedy
                e2e_stats.episodes_count += 1

                student.opt_both.zero_grad()
                loss.backward()
                if p.clip_grad is not None and p.clip_grad > 0:
                    torch.nn.utils.clip_grad_norm_(student_params, p.clip_grad)
                student.opt_both.step()

                _done_training = False
                if p.e2e_acc is not None and acc >= p.e2e_acc:
                    # print('reached target e2e acc %.3f' % acc, ' => breaking')
                    _done_training = True
                if p.e2e_ksteps is not None and epoch >= p.e2e_ksteps * 1000:
                    # print('reached target e2e step', epoch, ' => breaking')
                    _done_training = True
                save_e2e = p.save_e2e_everyk is not None and p.save_e2e_everyk > 0 and (epoch % (p.save_e2e_everyk * 1000)) == 0
                if time.time() - last_print >= self.render_every_seconds or _done_training or save_e2e:
                    holdout_acc_sum = 0
                    holdout_ep_count = 0
                    holdout_rho_sum = 0
                    holdout_send_greed_sum = 0
                    holdout_recv_greed_sum = 0
                    # self.teacher.lang_sender.eval()
                    student.eval()
                    ho_utts_l = []
                    ho_labels_l = []
                    with torch_random_state():
                        for i, (batch_image, batch_labels) in enumerate(self.dataset.iter_holdout(batch_size=p.batch_size)):
                            if p.enable_cuda:
                                batch_image = batch_image.cuda()
                                batch_labels = batch_labels.cuda()
                            with autograd.no_grad():
                                _, utts, acc = ref_task_game.forward(batch_image, training=False)
                                if utts.dtype == torch.float32:
                                    _, utts = utts.max(dim=-1)
                                holdout_acc_sum += acc
                                holdout_ep_count += 1
                                utts = utts.transpose(0, 1)
                                holdout_rho_sum += metrics.topographic_similarity(utts, batch_labels)
                            ho_utts_l.append(utts)
                            ho_labels_l.append(batch_labels)
                    ho_utts = torch.cat(ho_utts_l)
                    ho_labels = torch.cat(ho_labels_l)
                    if save_e2e:
                        samples_filename = p.utt_samples.format(epoch=epoch)
                        with open(samples_filename, 'wb') as f:
                            torch.save({'samples': {'utts': ho_utts, 'labels': ho_labels}, 'meta': p.__dict__}, f)
                        print('saved samples to ' + samples_filename)
                        model_save = p.model_save.format(epoch=epoch)
                        self.save_to(model_save)
                        print('saved model to ' + model_save)
                    student.train()
                    # self.teacher.lang_sender.train()
                    rho = holdout_rho_sum / holdout_ep_count
                    holdout_acc = holdout_acc_sum / holdout_ep_count
                    acc = e2e_stats.e2e_acc_sum / e2e_stats.episodes_count
                    loss = e2e_stats.e2e_loss_sum / e2e_stats.episodes_count
                    sender_greedy = e2e_stats.sender_greedy_sum / e2e_stats.episodes_count
                    receiver_greedy = e2e_stats.receiver_greedy_sum / e2e_stats.episodes_count

                    _elapsed_time = time.time() - _e2e_start
                    log_dict = {
                        'record_type': 'e2e',
                        'ilm_epoch': step,
                        'epoch': epoch,
                        'sps': int(epoch / _elapsed_time),
                        'e2e_time': int(_elapsed_time),
                        'acc': acc,
                        'holdout_acc': holdout_acc,
                        'rho': rho,
                        'loss': loss,
                        'send_greed': sender_greedy,
                        'recv_greed': receiver_greedy,
                    }
                    formatstr = (
                        '{record_type} '
                        'g={ilm_epoch} '
                        'e={epoch} '
                        't={e2e_time:.0f} '
                        'sps={sps:.0f} '
                        'acc={acc:.3f} '
                        'loss={loss:.3f} '
                        'ho_acc={holdout_acc:.3f} '
                        'rho={rho:.3f} '
                        's_g={send_greed:.3f} '
                        'r_g={recv_greed:.3f} '
                    )
                    self.print_and_log(log_dict, formatstr=formatstr)

                    e2e_stats.reset()
                    last_print = time.time()
                if _done_training:
                    break
                epoch += 1
            e2e_time = time.time() - _e2e_start
            e2e_acc = acc
            e2e_holdout_acc = holdout_acc
            e2e_rho = rho
            e2e_send_greedy = sender_greedy
            e2e_recv_greedy = receiver_greedy

            self.teacher = student

        if True:
            log_dict = {
                'type': 'ilm',
                'sps': int(step / (time.time() - self.start_time)),
                'elapsed_time': time.time() - self.start_time,
                'e2e_time': e2e_time,
                'e2e_acc': e2e_acc,
                'e2e_holdout_acc': e2e_holdout_acc,
                'e2e_rho': e2e_rho,
                'e2e_send_greedy': e2e_send_greedy,
                'e2e_recv_greedy': e2e_recv_greedy,
                'sup_sender_epochs': sup_sender_epochs,
                'sup_receiver_epochs': sup_receiver_epochs,
                'sup_sender_acc': sup_sender_acc,
                'sup_receiver_acc': sup_receiver_acc,
                'sup_sender_time': sup_sender_time,
                'sup_receiver_time': sup_receiver_time,
            }

            formatstr = (
                '{type} '
                'g={training_step} '
                't={elapsed_time:.0f} '
                'sps={sps:.0f}\n'
                '    sup_snd[e={sup_sender_epochs} acc={sup_sender_acc:.3f} t={sup_sender_time:.0f}]\n'
                '    sup_rcv[e={sup_receiver_epochs} acc={sup_receiver_acc:.3f} t={sup_receiver_time:.0f}]\n'
                '    e2e[acc={e2e_acc:.3f} t={e2e_time:.0f} sg={e2e_send_greedy:.3f} rg={e2e_recv_greedy:.3f}]\n'
                'ho_acc={e2e_holdout_acc:.3f} '
                'rho={e2e_rho:.3f} '
            )
            self.print_and_log(log_dict, formatstr=formatstr)
        if p.max_gen is not None and step + 1 >= p.max_gen:
            print('reached max generations', p.max_gen, '=> terminating')
            self.finish = True

if __name__ == '__main__':
    utils.clean_argv()
    runner = Runner()

    runner.add_param('--ds-ref', type=str)
    runner.add_param('-s', type=int, help='num shapes')
    runner.add_param('--opt', type=str, default='RMSprop')
    runner.add_param('--conv-class', type=str, default='CNNALPoolingAll')
    runner.add_param('--num-conv-layers', type=int, default=8)
    runner.add_param('--data-family', type=str, default='objects_gl')
    runner.add_param('--data-dir', type=str, default='~/data/{data_family}/{ds_ref}')
    runner.add_param('--utt-samples', type=str, default='tmp/{ref}_samples_{epoch}.pth')
    runner.add_param('--model-save', type=str, default='tmp/{ref}_model_{epoch}.pth')

    runner.add_param('--seed', type=int)
    runner.add_param('--batch-size', type=int, default=32)
    runner.add_param('--link', type=str, default='RL')
    runner.add_param('--clip-grad', type=float, default=0)
    runner.add_param('--sup-acc', type=float)
    runner.add_param('--e2e-acc', type=float)
    runner.add_param('--ilm', type=str)
    runner.add_param('--max-gen', type=int)
    runner.add_param('-f', type=float, default=0.4, help='supervised train fraction')
    runner.add_param('--no-train-e2e', action='store_true')
    runner.add_param('--save-e2e-everyk', type=int, default=100)

    runner.add_param('--embedding-size', type=int, default=50)
    runner.add_param('--vocab-size', type=int, default=100, help='excludes any terminator')
    runner.add_param('--model', type=str, default='RNN')
    runner.add_param('--rnn-type', type=str, default='GRU')
    runner.add_param('--num-layers', type=int, default=1)
    runner.add_param('--dropout', type=float, default=0.5)
    runner.add_param('--utt-len', type=int, default=6)
    runner.add_param('--nle', type=str, default='2,3', help='negative log10 entropy reg')

    runner.parse_args()
    args = runner.params
    if args.ilm is not None:
        args.e2e_ksteps, args.sup_ksteps = [float(v) for v in args.ilm.split(',')]
    else:
        args.e2e_ksteps, args.sup_ksteps = None, None
    del args.__dict__['ilm']
    args.utt_samples = args.utt_samples.format(ref=args.ref, epoch='{epoch}')
    args.model_save = args.model_save.format(ref=args.ref, epoch='{epoch}')
    args.s_ent, args.r_ent = [math.pow(10, -float(v)) for v in args.nle.split(',')]
    print(f'ent reg {args.s_ent:.1e} {args.r_ent:.1e}')
    del args.__dict__['nle']
    args.shapes = args.s
    del args.__dict__['s']
    args.sup_train_frac = args.f
    del args.__dict__['f']
    if args.shapes is not None:
        args.ds_ref = {
            1: 'dsd39_1sb',
            2: 'dsd37_2s_ho2',
            3: 'dsd38_3s_ho2'
            # 2: 'dsd32_twoshape_123',
            # 3: 'dsd33_threeshapes_123'
        }[args.shapes]
    assert args.ds_ref is not None
    utils.reverse_args(runner.params, 'no_train_e2e', 'train_e2e')
    runner.params.data_dir = runner.params.data_dir.format(**runner.params.__dict__)
    print('runner.params', runner.params)
    runner.setup_base()
    runner.run_base()


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import division
import numpy as np

from datetime import datetime
from lib.ReservoirNetworkController import ReservoirNetworkController
from lib.AtlasJointsInfo import AtlasJointsInfo

class SpikingNeuralNetwork:
    INPUT_LAYER_SIZE = 28
    OUTPUT_LAYER_SIZE = 28
    SIMULATION_TIME_INTERVAL = 15 # msec
    RESERVOIR_NETWORK_SIZE = 500
    MAX_INPUT_RATE = 1000
    MIN_INPUT_RATE = 50

    def __init__(self):
        self._output_layer_activations = np.zeros(self.OUTPUT_LAYER_SIZE)
        self._hidden_layer_weights = np.zeros([self.RESERVOIR_NETWORK_SIZE,
                                               self.OUTPUT_LAYER_SIZE])
        self._hidden_layer_neuron_mapping =  \
            self._initialize_hidden_layer_neuron_mapping()
        for i in xrange(len(self._hidden_layer_neuron_mapping)):
            one_neuron_mapping = self._hidden_layer_neuron_mapping[i]
            for c in xrange(len(one_neuron_mapping)):
                hidden_neuron_index = one_neuron_mapping[c]
                self._hidden_layer_weights[hidden_neuron_index][i] = \
                    np.random.rand(1)[0]
        self._hidden_layer_biases = np.random.rand(self.OUTPUT_LAYER_SIZE)
        self._joints_info_provider = AtlasJointsInfo()
        self._input_layer = np.zeros(self.INPUT_LAYER_SIZE)
        self._hidden_layer = self._initialize_reservoir_network()
        self._output_layer = np.zeros(self.OUTPUT_LAYER_SIZE)

    def _initialize_reservoir_network(self):
        snn_controller = ReservoirNetworkController(
                self.RESERVOIR_NETWORK_SIZE, self.INPUT_LAYER_SIZE)
        return snn_controller

    def process_input(self, state):
        # TODO so far we are using only position, eventually should use more
        # parameters
        start_time = datetime.now()
        self._set_position_values(self._normalize_position_input(state.position))
        stop_time = datetime.now()
        delta = stop_time - start_time
        print "converting input in: " + str(delta.seconds) + "." + str(
                delta.microseconds/1000)
        start_time = stop_time
        self._apply_poisson_group_input()
        stop_time = datetime.now()
        delta = stop_time - start_time
        print "applying Poisson group in: " + str(delta.seconds) + "." + str(
                delta.microseconds/1000)
        start_time = stop_time
        self._hidden_layer.run_simulation(self.SIMULATION_TIME_INTERVAL)
        stop_time = datetime.now()
        delta = stop_time - start_time
        print "simulation running in: " + str(delta.seconds) + "." + str(
                delta.microseconds/1000)
        start_time = stop_time
        self._decode_snn_output()
        stop_time = datetime.now()
        delta = stop_time - start_time
        print "decoding output in: " + str(delta.seconds) + "." + str(
                delta.microseconds/1000)

    def _set_position_values(self, position):
        for i in xrange(len(position)):
            self._input_layer[i] = position[i]

    def _normalize_position_input(self, position_input):
        result = []
        for i in xrange(len(position_input)):
            result.append((position_input[i] -
                           self._joints_info_provider.get_min_value_for_joint(
                               i)) / (
                              self._joints_info_provider.get_max_value_for_joint(
                                  i) -
                              self._joints_info_provider.get_min_value_for_joint(
                                  i)))
        # output values are between 0 and 1
        return result

    def get_input_layer_values(self):
        return self._input_layer

    def _apply_poisson_group_input(self):
        firing_rates = np.zeros(len(self._input_layer))
        for i in xrange(len(self._input_layer)):
            firing_rates[i] = self._convert_to_rate(self._input_layer[i])
        self._hidden_layer.set_poisson_group_rates(firing_rates)

    def _convert_to_rate(self, input):
        # input values are between 0 and 1 (after normalization)
        rate = input*(self.MAX_INPUT_RATE - self.MIN_INPUT_RATE) + \
               self.MIN_INPUT_RATE
        return rate

    def _decode_snn_output(self):
        firing_rates = self._hidden_layer.get_reservoir_firing_rates_output()
        # print "SpikingNeuralNetwork._decode_snn_output(): " \
        #       "np.amax(firing_rates) = " + str(
        #         np.amax(firing_rates))
        # print "SpikingNeuralNetwork._decode_snn_output(): " \
        #       "np.amin(firing_rates) = " + str(
        #         np.amin(firing_rates))
        self._compute_activations_from_reservoir(firing_rates)
        for i in xrange(self.OUTPUT_LAYER_SIZE):
            self._output_layer[i] = self._output_layer_activations[i]

    def _normalize_firing_rates_output(self, input_value):
        max_value = self.MAX_INPUT_RATE
        min_value = self.MIN_INPUT_RATE
        result = []
        for i in xrange(len(input_value)):
            result.append((input_value[i] - min_value) / (max_value - min_value))
        return result

    def get_output_layer_values(self):
        return self._output_layer

    def _compute_activations_from_reservoir(self, firing_rates):
        for i in xrange(len(self._hidden_layer_neuron_mapping)):
            one_neuron_mapping = self._hidden_layer_neuron_mapping[i]
            z = 0
            for c in xrange(len(one_neuron_mapping)):
                hidden_neuron_index = one_neuron_mapping[c]
                rate = firing_rates[hidden_neuron_index]
                # TODO think of a more elegant way to deal with it
                if rate == 0:
                    rate = 0.0001
                # TODO: computing activation from reservoir as
                # rate/self.MAX_INPUT_RATE is
                # quite questionable!!!
                z = z + self._hidden_layer_weights[hidden_neuron_index][i]*(
                    rate/self.MAX_INPUT_RATE)
            z = z + self._hidden_layer_biases[i]
            activation = self._sigmoid(z)
            self._output_layer_activations[i] = activation

    def _initialize_hidden_layer_neuron_mapping(self):
        mapping = []
        for i in xrange(self.OUTPUT_LAYER_SIZE):
            # TODO should be np.random.randint(10)
            connected_neurons_num = np.random.randint(1,10,1)[0]
            one_neuron_mapping = np.random.randint(1,
                                                self.RESERVOIR_NETWORK_SIZE,
                                                   connected_neurons_num)
            mapping.append(one_neuron_mapping)
        return mapping

    def _sigmoid(self, z):
        return 1.0/(1.0+np.exp(-z))

    def get_hidden_layer_weights_for_output_neuron(self, neuron_idx):
        # print "entering get_hidden_layer_weights_for_output_neuron"
        hidden_layer_weights_T = self._hidden_layer_weights.T
        # print "len(hidden_layer_weights_T[neuron_idx]) = " + str(len(
        #         hidden_layer_weights_T[neuron_idx]))
        # print "hidden_layer_weights_T[neuron_idx] : " + str(
        #         hidden_layer_weights_T[neuron_idx])
        # print "leaving get_hidden_layer_weights_for_output_neuron"
        return hidden_layer_weights_T[neuron_idx]

    def get_hidden_layer_firing_rates(self):
        return self._hidden_layer.get_reservoir_firing_rates_output()

    def get_mapping_for_output_neuron(self, neuron_idx):
        return self._hidden_layer_neuron_mapping[neuron_idx]

    def get_hidden_layer_biases(self):
        return self._hidden_layer_biases

    def set_hidden_layer_weights_for_neuron(self, neuron_idx,
                                            hidden_layer_weights):
        # print "entering set_hidden_layer_weights_for_neuron"
        # print "len(self._hidden_layer_weights.T[neuron_idx]) = " + str(len(
        #         self._hidden_layer_weights.T[neuron_idx]))
        # print "len(hidden_layer_weights) = " + str(len(hidden_layer_weights))
        self._hidden_layer_weights.T[neuron_idx] = hidden_layer_weights
        # print "leaving set_hidden_layer_weights_for_neuron"

    def set_hidden_layer_biases(self, hidden_layer_biases):
        # print "entering SpikingNeuralNetwork.set_hidden_layer_biases"
        self._hidden_layer_biases = hidden_layer_biases
        # print "leaving SpikingNeuralNetwork.set_hidden_layer_biases"

    def recalculate_output_layer(self):
        self._decode_snn_output()


#!/usr/bin/env python
# coding: utf-8

import numpy as np
import math
import tensorflow.keras.backend as K
import tensorflow as tf

def xywh_to_x1y1x2y2(boxes):
    return tf.concat([boxes[..., :2] - boxes[..., 2:] * 0.5, boxes[..., :2] + boxes[..., 2:] * 0.5], axis=-1)

# x,y,w,h
def bbox_iou(boxes1, boxes2):
    boxes1_area = boxes1[..., 2] * boxes1[..., 3]  # w * h
    boxes2_area = boxes2[..., 2] * boxes2[..., 3]

    # (x, y, w, h) -> (x0, y0, x1, y1)
    boxes1 = xywh_to_x1y1x2y2(boxes1)
    boxes2 = xywh_to_x1y1x2y2(boxes2)

    # coordinates of intersection
    top_left = tf.maximum(boxes1[..., :2], boxes2[..., :2])
    bottom_right = tf.minimum(boxes1[..., 2:], boxes2[..., 2:])
    intersection_xy = tf.maximum(bottom_right - top_left, 0.0)

    intersection_area = intersection_xy[..., 0] * intersection_xy[..., 1]
    union_area = boxes1_area + boxes2_area - intersection_area

    return 1.0 * intersection_area / (union_area + tf.keras.backend.epsilon())

def bbox_giou(boxes1, boxes2):
    boxes1_area = boxes1[..., 2] * boxes1[..., 3]  # w*h
    boxes2_area = boxes2[..., 2] * boxes2[..., 3]

    # (x, y, w, h) -> (x0, y0, x1, y1)
    boxes1 = xywh_to_x1y1x2y2(boxes1)
    boxes2 = xywh_to_x1y1x2y2(boxes2)

    top_left = tf.maximum(boxes1[..., :2], boxes2[..., :2])
    bottom_right = tf.minimum(boxes1[..., 2:], boxes2[..., 2:])

    intersection_xy = tf.maximum(bottom_right - top_left, 0.0)
    intersection_area = intersection_xy[..., 0] * intersection_xy[..., 1]

    union_area = boxes1_area + boxes2_area - intersection_area

    iou = 1.0 * intersection_area / (union_area + tf.keras.backend.epsilon())

    enclose_top_left = tf.minimum(boxes1[..., :2], boxes2[..., :2])
    enclose_bottom_right = tf.maximum(boxes1[..., 2:], boxes2[..., 2:])

    enclose_xy = enclose_bottom_right - enclose_top_left
    enclose_area = enclose_xy[..., 0] * enclose_xy[..., 1]

    giou = iou - tf.math.divide_no_nan(enclose_area - union_area, enclose_area)

    return giou

def bbox_ciou(boxes1, boxes2):
    '''
    ciou = iou - p2/c2 - av
    :param boxes1: (8, 13, 13, 3, 4)   pred_xywh
    :param boxes2: (8, 13, 13, 3, 4)   label_xywh
    :return:
    '''
    boxes1_x0y0x1y1 = tf.concat([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
                                 boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)
    boxes2_x0y0x1y1 = tf.concat([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
                                 boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)
    boxes1_x0y0x1y1 = tf.concat([tf.minimum(boxes1_x0y0x1y1[..., :2], boxes1_x0y0x1y1[..., 2:]),
                                 tf.maximum(boxes1_x0y0x1y1[..., :2], boxes1_x0y0x1y1[..., 2:])], axis=-1)
    boxes2_x0y0x1y1 = tf.concat([tf.minimum(boxes2_x0y0x1y1[..., :2], boxes2_x0y0x1y1[..., 2:]),
                                 tf.maximum(boxes2_x0y0x1y1[..., :2], boxes2_x0y0x1y1[..., 2:])], axis=-1)

    # area
    boxes1_area = (boxes1_x0y0x1y1[..., 2] - boxes1_x0y0x1y1[..., 0]) * (
                boxes1_x0y0x1y1[..., 3] - boxes1_x0y0x1y1[..., 1])
    boxes2_area = (boxes2_x0y0x1y1[..., 2] - boxes2_x0y0x1y1[..., 0]) * (
                boxes2_x0y0x1y1[..., 3] - boxes2_x0y0x1y1[..., 1])

    # top-left and bottom-right coord, shape: (8, 13, 13, 3, 2)
    left_up = tf.maximum(boxes1_x0y0x1y1[..., :2], boxes2_x0y0x1y1[..., :2])
    right_down = tf.minimum(boxes1_x0y0x1y1[..., 2:], boxes2_x0y0x1y1[..., 2:])

    # intersection area and iou
    inter_section = tf.maximum(right_down - left_up, 0.0)
    inter_area = inter_section[..., 0] * inter_section[..., 1]
    union_area = boxes1_area + boxes2_area - inter_area
    iou = inter_area / (union_area + 1e-9)

    # top-left and bottom-right coord of the enclosing rectangle, shape: (8, 13, 13, 3, 2)
    enclose_left_up = tf.minimum(boxes1_x0y0x1y1[..., :2], boxes2_x0y0x1y1[..., :2])
    enclose_right_down = tf.maximum(boxes1_x0y0x1y1[..., 2:], boxes2_x0y0x1y1[..., 2:])

    # diagnal ** 2
    enclose_wh = enclose_right_down - enclose_left_up
    enclose_c2 = K.pow(enclose_wh[..., 0], 2) + K.pow(enclose_wh[..., 1], 2)

    # center distances between two rectangles
    p2 = K.pow(boxes1[..., 0] - boxes2[..., 0], 2) + K.pow(boxes1[..., 1] - boxes2[..., 1], 2)

    # add av
    atan1 = tf.atan(boxes1[..., 2] / (boxes1[..., 3] + 1e-9))
    atan2 = tf.atan(boxes2[..., 2] / (boxes2[..., 3] + 1e-9))
    v = 4.0 * K.pow(atan1 - atan2, 2) / (math.pi ** 2)
    a = v / (1 - iou + v)

    ciou = iou - 1.0 * p2 / enclose_c2 - 1.0 * a * v
    return ciou

def yolo_loss(args, num_classes, iou_loss_thresh, anchors):
    conv_lbbox = args[2]   # (?, ?, ?, 3*(num_classes+5))
    conv_mbbox = args[1]   # (?, ?, ?, 3*(num_classes+5))
    conv_sbbox = args[0]   # (?, ?, ?, 3*(num_classes+5))
    label_sbbox = args[3]   # (?, ?, ?, 3, num_classes+5)
    label_mbbox = args[4]   # (?, ?, ?, 3, num_classes+5)
    label_lbbox = args[5]   # (?, ?, ?, 3, num_classes+5)
    true_bboxes = args[6]   # (?, 50, 4)
    pred_sbbox = decode(conv_sbbox, anchors[0], 8, num_classes)
    pred_mbbox = decode(conv_mbbox, anchors[1], 16, num_classes)
    pred_lbbox = decode(conv_lbbox, anchors[2], 32, num_classes)
    sbbox_ciou_loss, sbbox_conf_loss, sbbox_prob_loss = loss_layer(conv_sbbox, pred_sbbox, label_sbbox, true_bboxes, 8, num_classes, iou_loss_thresh)
    mbbox_ciou_loss, mbbox_conf_loss, mbbox_prob_loss = loss_layer(conv_mbbox, pred_mbbox, label_mbbox, true_bboxes, 16, num_classes, iou_loss_thresh)
    lbbox_ciou_loss, lbbox_conf_loss, lbbox_prob_loss = loss_layer(conv_lbbox, pred_lbbox, label_lbbox, true_bboxes, 32, num_classes, iou_loss_thresh)

    ciou_loss = (lbbox_ciou_loss + sbbox_ciou_loss + mbbox_ciou_loss) * 3.54
    conf_loss = (lbbox_conf_loss + sbbox_conf_loss + mbbox_conf_loss) * 64.3
    prob_loss = (lbbox_prob_loss + sbbox_prob_loss + mbbox_prob_loss) * 1

    return ciou_loss+conf_loss+prob_loss

def loss_layer(conv, pred, label, bboxes, stride, num_class, iou_loss_thresh):
    conv_shape = tf.shape(conv)
    batch_size = conv_shape[0]
    output_size = conv_shape[1]
    input_size = stride * output_size
    conv = tf.reshape(conv, (batch_size, output_size, output_size,
                             3, 5 + num_class))
    conv_raw_prob = conv[:, :, :, :, 5:]
    conv_raw_conf = conv[:, :, :, :, 4:5]

    pred_xywh = pred[:, :, :, :, 0:4]
    pred_conf = pred[:, :, :, :, 4:5]

    label_xywh = label[:, :, :, :, 0:4]
    respond_bbox = label[:, :, :, :, 4:5]
    label_prob = label[:, :, :, :, 5:]

    # Coordinate loss
    ciou = tf.expand_dims(bbox_giou(pred_xywh, label_xywh), axis=-1)  # (8, 13, 13, 3, 1)
    # ciou = tf.expand_dims(bbox_ciou(pred_xywh, label_xywh), axis=-1)  # (8, 13, 13, 3, 1)
    input_size = tf.cast(input_size, tf.float32)

    # loss weight of the gt bbox: 2-(gt area/img area)
    bbox_loss_scale = 2.0 - 1.0 * label_xywh[:, :, :, :, 2:3] * label_xywh[:, :, :, :, 3:4] / (input_size ** 2)
    ciou_loss = respond_bbox * bbox_loss_scale * (1 - ciou)  # iou loss for respond bbox

    # Classification loss for respond bbox
    prob_loss = respond_bbox * tf.nn.sigmoid_cross_entropy_with_logits(labels=label_prob, logits=conv_raw_prob)

    expand_pred_xywh = pred_xywh[:, :, :, :, np.newaxis, :]  # (?, grid_h, grid_w, 3, 1, 4)
    expand_bboxes = bboxes[:, np.newaxis, np.newaxis, np.newaxis, :, :]  # (?, 1, 1, 1, 70, 4)
    iou = bbox_iou(expand_pred_xywh, expand_bboxes)  # IoU between all pred bbox and all gt (?, grid_h, grid_w, 3, 70)
    max_iou = tf.expand_dims(tf.reduce_max(iou, axis=-1), axis=-1)  # max iou: (?, grid_h, grid_w, 3, 1)

    # ignore the bbox which is not respond bbox and max iou < threshold
    respond_bgd = (1.0 - respond_bbox) * tf.cast(max_iou < iou_loss_thresh, tf.float32)

    # Confidence loss
    conf_focal = tf.pow(respond_bbox - pred_conf, 2)

    conf_loss = conf_focal * (
            respond_bbox * tf.nn.sigmoid_cross_entropy_with_logits(labels=respond_bbox, logits=conv_raw_conf)
            +
            respond_bgd * tf.nn.sigmoid_cross_entropy_with_logits(labels=respond_bbox, logits=conv_raw_conf)
    )

    ciou_loss = tf.reduce_mean(tf.reduce_sum(ciou_loss, axis=[1, 2, 3, 4]))
    conf_loss = tf.reduce_mean(tf.reduce_sum(conf_loss, axis=[1, 2, 3, 4]))
    prob_loss = tf.reduce_mean(tf.reduce_sum(prob_loss, axis=[1, 2, 3, 4]))

    return ciou_loss, conf_loss, prob_loss

def decode(conv_output, anchors, stride, num_class):
    conv_shape = tf.shape(conv_output)
    batch_size = conv_shape[0]
    output_size = conv_shape[1]
    anchor_per_scale = len(anchors)
    conv_output = tf.reshape(conv_output, (batch_size, output_size, output_size, anchor_per_scale, 5 + num_class))
    conv_raw_dxdy = conv_output[:, :, :, :, 0:2]
    conv_raw_dwdh = conv_output[:, :, :, :, 2:4]
    conv_raw_conf = conv_output[:, :, :, :, 4:5]
    conv_raw_prob = conv_output[:, :, :, :, 5:]
    y = tf.tile(tf.range(output_size, dtype=tf.int32)[:, tf.newaxis], [1, output_size])
    x = tf.tile(tf.range(output_size, dtype=tf.int32)[tf.newaxis, :], [output_size, 1])
    xy_grid = tf.concat([x[:, :, tf.newaxis], y[:, :, tf.newaxis]], axis=-1)
    xy_grid = tf.tile(xy_grid[tf.newaxis, :, :, tf.newaxis, :], [batch_size, 1, 1, anchor_per_scale, 1])
    xy_grid = tf.cast(xy_grid, tf.float32)
    pred_xy = (tf.sigmoid(conv_raw_dxdy) + xy_grid) * stride
    pred_wh = (tf.exp(conv_raw_dwdh) * anchors)
    pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
    pred_conf = tf.sigmoid(conv_raw_conf)
    pred_prob = tf.sigmoid(conv_raw_prob)
    return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)


""" Additional helper functions dealing with transient-CW F(t0,tau) maps """

import numpy as np
import os
import logging
from time import time

# optional imports
import importlib as imp

def _optional_import(modulename, shorthand=None):
    """
    Import a module/submodule only if it's available.

    using importlib instead of __import__
    because the latter doesn't handle sub.modules

    Also including a special check to fail more gracefully
    when CUDA_DEVICE is set to too high a number.
    """

    if shorthand is None:
        shorthand = modulename
        shorthandbit = ""
    else:
        shorthandbit = " as " + shorthand

    try:
        globals()[shorthand] = imp.import_module(modulename)
        logging.debug("Successfully imported module %s%s." % (modulename, shorthandbit))
        success = True
    except ImportError as e:
        logging.debug("Failed to import module {:s}.".format(modulename))
        success = False

    return success

class pyTransientFstatMap(object):
    """
    simplified object class for a F(t0,tau) F-stat map (not 2F!)
    based on LALSuite's transientFstatMap_t type
    replacing the gsl matrix with a numpy array

    F_mn:   2D array of 2F values
    maxF:   maximum of F (not 2F!)
    t0_ML:  maximum likelihood transient start time t0 estimate
    tau_ML: maximum likelihood transient duration tau estimate
    """

    def __init__(self, N_t0Range, N_tauRange):
        self.F_mn = np.zeros((N_t0Range, N_tauRange), dtype=np.float32)
        # Initializing maxF to a negative value ensures
        # that we always update at least once and hence return
        # sane t0_d_ML, tau_d_ML
        # even if there is only a single bin where F=0 happens.
        self.maxF = float(-1.0)
        self.t0_ML = float(0.0)
        self.tau_ML = float(0.0)

# dictionary of the actual callable F-stat map functions we support,
# if the corresponding modules are available.
fstatmap_versions = {
    "lal": lambda multiFstatAtoms, windowRange: getattr(
        lalpulsar, "ComputeTransientFstatMap"
    )(multiFstatAtoms, windowRange, False),
    "pycuda": lambda multiFstatAtoms, windowRange: pycuda_compute_transient_fstat_map(
        multiFstatAtoms, windowRange
    ),
}

def init_transient_fstat_map_features(wantCuda=False, cudaDeviceName=None):
    """
    Initialization of available modules (or "features") for F-stat maps.

    Returns a dictionary of method names, to match fstatmap_versions
    each key's value set to True only if
    all required modules are importable on this system.
    """

    features = {}

    have_lal = _optional_import("lal")
    have_lalpulsar = _optional_import("lalpulsar")
    features["lal"] = have_lal and have_lalpulsar

    # import GPU features
    have_pycuda = _optional_import("pycuda")
    have_pycuda_drv = _optional_import("pycuda.driver", "drv")
    have_pycuda_gpuarray = _optional_import("pycuda.gpuarray", "gpuarray")
    have_pycuda_tools = _optional_import("pycuda.tools", "cudatools")
    have_pycuda_compiler = _optional_import("pycuda.compiler", "cudacomp")
    features["pycuda"] = (
        have_pycuda_drv
        and have_pycuda_gpuarray
        and have_pycuda_tools
        and have_pycuda_compiler
    )

    logging.debug("Got the following features for transient F-stat maps:")
    logging.debug(features)

    if wantCuda and features["pycuda"]:
        logging.debug("CUDA version: " + ".".join(map(str, drv.get_version())))

        drv.init()
        logging.debug(
            "Starting with default pyCUDA context,"
            " then checking all available devices..."
        )
        try:
            context0 = pycuda.tools.make_default_context()
        except pycuda._driver.LogicError as e:
            if e.message == "cuDeviceGet failed: invalid device ordinal":
                devn = int(os.environ["CUDA_DEVICE"])
                raise RuntimeError(
                    "Requested CUDA device number {} exceeds"
                    " number of available devices!"
                    " Please change through environment"
                    " variable $CUDA_DEVICE.".format(devn)
                )
            else:
                raise pycuda._driver.LogicError(e.message)

        num_gpus = drv.Device.count()
        logging.debug("Found {} CUDA device(s).".format(num_gpus))

        devices = []
        devnames = np.empty(num_gpus, dtype="S32")
        for n in range(num_gpus):
            devn = drv.Device(n)
            devices.append(devn)
            devnames[n] = devn.name().replace(" ", "-").replace("_", "-")
            logging.debug(
                "device {}: model: {}, RAM: {}MB".format(
                    n, devnames[n], devn.total_memory() / (2.0 ** 20)
                )
            )

        if "CUDA_DEVICE" in os.environ:
            devnum0 = int(os.environ["CUDA_DEVICE"])
        else:
            devnum0 = 0

        matchbit = ""
        if cudaDeviceName:
            # allow partial matches in device names
            devmatches = [
                devidx
                for devidx, devname in enumerate(devnames)
                if cudaDeviceName in devname
            ]
            if len(devmatches) == 0:
                context0.detach()
                raise RuntimeError(
                    'Requested CUDA device "{}" not found.'
                    " Available devices: [{}]".format(
                        cudaDeviceName, ",".join(devnames)
                    )
                )
            else:
                devnum = devmatches[0]
                if len(devmatches) > 1:
                    logging.warning(
                        'Found {} CUDA devices matching name "{}".'
                        " Choosing first one with index {}.".format(
                            len(devmatches), cudaDeviceName, devnum
                        )
                    )
            os.environ["CUDA_DEVICE"] = str(devnum)
            matchbit = '(matched to user request "{}")'.format(cudaDeviceName)
        elif "CUDA_DEVICE" in os.environ:
            devnum = int(os.environ["CUDA_DEVICE"])
        else:
            devnum = 0
        devn = devices[devnum]
        logging.info(
            "Choosing CUDA device {},"
            " of {} devices present: {}{}...".format(
                devnum, num_gpus, devn.name(), matchbit
            )
        )
        if devnum == devnum0:
            gpu_context = context0
        else:
            context0.pop()
            gpu_context = pycuda.tools.make_default_context()
            gpu_context.push()

        _print_GPU_memory_MB("Available")
    else:
        gpu_context = None

    return features, gpu_context

def call_compute_transient_fstat_map(
    version, features, multiFstatAtoms=None, windowRange=None
):
    """Choose which version of the ComputeTransientFstatMap function to call."""

    if version in fstatmap_versions:
        if features[version]:
            time0 = time()
            FstatMap = fstatmap_versions[version](multiFstatAtoms, windowRange)
            timingFstatMap = time() - time0
        else:
            raise Exception(
                "Required module(s) for transient F-stat map"
                ' method "{}" not available!'.format(version)
            )
    else:
        raise Exception(
            'Transient F-stat map method "{}"' " not implemented!".format(version)
        )
    return FstatMap, timingFstatMap

def reshape_FstatAtomsVector(atomsVector):
    """
    Make a dictionary of ndarrays out of a atoms "vector" structure.

    The input is a "vector"-like structure with times as the higher hierarchical
    level and a set of "atoms" quantities defined at each timestamp.
    The output is a dictionary with an entry for each quantity,
    which is a 1D ndarray over timestamps for that one quantity.
    """

    numAtoms = atomsVector.length
    atomsDict = {}
    atom_fieldnames = [
        "timestamp",
        "Fa_alpha",
        "Fb_alpha",
        "a2_alpha",
        "ab_alpha",
        "b2_alpha",
    ]
    atom_dtypes = [np.uint32, complex, complex, np.float32, np.float32, np.float32]
    for f, field in enumerate(atom_fieldnames):
        atomsDict[field] = np.ndarray(numAtoms, dtype=atom_dtypes[f])

    for n, atom in enumerate(atomsVector.data):
        for field in atom_fieldnames:
            atomsDict[field][n] = atom.__getattribute__(field)

    atomsDict["Fa_alpha_re"] = np.float32(atomsDict["Fa_alpha"].real)
    atomsDict["Fa_alpha_im"] = np.float32(atomsDict["Fa_alpha"].imag)
    atomsDict["Fb_alpha_re"] = np.float32(atomsDict["Fb_alpha"].real)
    atomsDict["Fb_alpha_im"] = np.float32(atomsDict["Fb_alpha"].imag)

    return atomsDict

def _get_absolute_kernel_path(kernel):
    pyfstatdir = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))
    kernelfile = kernel + ".cu"
    return os.path.join(pyfstatdir, "pyCUDAkernels", kernelfile)

def _print_GPU_memory_MB(key):
    mem_used_MB = drv.mem_get_info()[0] / (2.0 ** 20)
    mem_total_MB = drv.mem_get_info()[1] / (2.0 ** 20)
    logging.debug(
        "{} GPU memory: {:.4f} / {:.4f} MB free".format(key, mem_used_MB, mem_total_MB)
    )

def pycuda_compute_transient_fstat_map(multiFstatAtoms, windowRange):
    """
    GPU version of the function to compute transient-window "F-statistic map"
    over start-time and timescale {t0, tau}.
    Based on XLALComputeTransientFstatMap from LALSuite,
    (C) 2009 Reinhard Prix, licensed under GPL

    Returns a 2D matrix F_mn,
    with m = index over start-times t0,
    and  n = index over timescales tau,
    in steps of dt0  in [t0,  t0+t0Band],
    and         dtau in [tau, tau+tauBand]
    as defined in windowRange input.
    """

    if windowRange.type >= lalpulsar.TRANSIENT_LAST:
        raise ValueError(
            "Unknown window-type ({}) passed as input."
            " Allowed are [0,{}].".format(
                windowRange.type, lalpulsar.TRANSIENT_LAST - 1
            )
        )

    # internal dict for search/setup parameters
    tCWparams = {}

    # first combine all multi-atoms
    # into a single atoms-vector with *unique* timestamps
    tCWparams["TAtom"] = multiFstatAtoms.data[0].TAtom
    TAtomHalf = int(tCWparams["TAtom"] / 2)  # integer division
    atoms = lalpulsar.mergeMultiFstatAtomsBinned(multiFstatAtoms, tCWparams["TAtom"])

    # make a combined input matrix of all atoms vectors, for transfer to GPU
    tCWparams["numAtoms"] = atoms.length
    atomsDict = reshape_FstatAtomsVector(atoms)
    atomsInputMatrix = np.column_stack(
        (
            atomsDict["a2_alpha"],
            atomsDict["b2_alpha"],
            atomsDict["ab_alpha"],
            atomsDict["Fa_alpha_re"],
            atomsDict["Fa_alpha_im"],
            atomsDict["Fb_alpha_re"],
            atomsDict["Fb_alpha_im"],
        )
    )

    # actual data spans [t0_data, t0_data + tCWparams['numAtoms'] * TAtom]
    # in steps of TAtom
    tCWparams["t0_data"] = int(atoms.data[0].timestamp)
    tCWparams["t1_data"] = int(
        atoms.data[tCWparams["numAtoms"] - 1].timestamp + tCWparams["TAtom"]
    )

    logging.debug(
        "Transient F-stat map:"
        " t0_data={:d}, t1_data={:d}".format(tCWparams["t0_data"], tCWparams["t1_data"])
    )
    logging.debug(
        "Transient F-stat map:"
        " numAtoms={:d}, TAtom={:d},"
        " TAtomHalf={:d}".format(tCWparams["numAtoms"], tCWparams["TAtom"], TAtomHalf)
    )

    # special treatment of window_type = none
    # ==> replace by rectangular window spanning all the data
    if windowRange.type == lalpulsar.TRANSIENT_NONE:
        windowRange.type = lalpulsar.TRANSIENT_RECTANGULAR
        windowRange.t0 = tCWparams["t0_data"]
        windowRange.t0Band = 0
        windowRange.dt0 = tCWparams["TAtom"]  # irrelevant
        windowRange.tau = tCWparams["numAtoms"] * tCWparams["TAtom"]
        windowRange.tauBand = 0
        windowRange.dtau = tCWparams["TAtom"]  # irrelevant

    """ NOTE: indices {i,j} enumerate *actual* atoms and their timestamps t_i,
    * while the indices {m,n} enumerate the full grid of values
    * in [t0_min, t0_max]x[Tcoh_min, Tcoh_max] in steps of deltaT.
    * This allows us to deal with gaps in the data in a transparent way.
    *
    * NOTE2: we operate on the 'binned' atoms returned
    * from XLALmergeMultiFstatAtomsBinned(),
    * which means we can safely assume all atoms to be lined up
    * perfectly on a 'deltaT' binned grid.
    *
    * The mapping used will therefore be {i,j} -> {m,n}:
    *   m = offs_i  / deltaT
    *   start-time offset from t0_min measured in deltaT
    *   n = Tcoh_ij / deltaT
    *   duration Tcoh_ij measured in deltaT,
    *
    * where
    *   offs_i  = t_i - t0_min
    *   Tcoh_ij = t_j - t_i + deltaT
    *
    """

    # We allocate a matrix  {m x n} = t0Range * TcohRange elements
    # covering the full transient window-range [t0,t0+t0Band]x[tau,tau+tauBand]
    tCWparams["N_t0Range"] = int(
        np.floor(1.0 * windowRange.t0Band / windowRange.dt0) + 1
    )
    tCWparams["N_tauRange"] = int(
        np.floor(1.0 * windowRange.tauBand / windowRange.dtau) + 1
    )
    FstatMap = pyTransientFstatMap(tCWparams["N_t0Range"], tCWparams["N_tauRange"])

    logging.debug(
        "Transient F-stat map:"
        " N_t0Range={:d}, N_tauRange={:d},"
        " total grid points: {:d}".format(
            tCWparams["N_t0Range"],
            tCWparams["N_tauRange"],
            tCWparams["N_t0Range"] * tCWparams["N_tauRange"],
        )
    )

    if windowRange.type == lalpulsar.TRANSIENT_RECTANGULAR:
        FstatMap.F_mn = pycuda_compute_transient_fstat_map_rect(
            atomsInputMatrix, windowRange, tCWparams
        )
    elif windowRange.type == lalpulsar.TRANSIENT_EXPONENTIAL:
        FstatMap.F_mn = pycuda_compute_transient_fstat_map_exp(
            atomsInputMatrix, windowRange, tCWparams
        )
    else:
        raise ValueError(
            "Invalid transient window type {}"
            " not in [{}, {}].".format(
                windowRange.type, lalpulsar.TRANSIENT_NONE, lalpulsar.TRANSIENT_LAST - 1
            )
        )

    # out of loop: get max2F and ML estimates over the m x n matrix
    FstatMap.maxF = FstatMap.F_mn.max()
    maxidx = np.unravel_index(
        FstatMap.F_mn.argmax(), (tCWparams["N_t0Range"], tCWparams["N_tauRange"])
    )
    FstatMap.t0_ML = windowRange.t0 + maxidx[0] * windowRange.dt0
    FstatMap.tau_ML = windowRange.tau + maxidx[1] * windowRange.dtau

    logging.debug(
        "Done computing transient F-stat map."
        " maxF={:.4f}, t0_ML={}, tau_ML={}".format(
            FstatMap.maxF, FstatMap.t0_ML, FstatMap.tau_ML
        )
    )

    return FstatMap

def pycuda_compute_transient_fstat_map_rect(atomsInputMatrix, windowRange, tCWparams):
    """
    only GPU-parallizing outer loop,
    keeping partial sums with memory in kernel
    """

    # gpu data setup and transfer
    _print_GPU_memory_MB("Initial")
    input_gpu = gpuarray.to_gpu(atomsInputMatrix)
    Fmn_gpu = gpuarray.GPUArray(
        (tCWparams["N_t0Range"], tCWparams["N_tauRange"]), dtype=np.float32
    )
    _print_GPU_memory_MB("After input+output allocation:")

    # GPU kernel
    kernel = "cudaTransientFstatRectWindow"
    kernelfile = _get_absolute_kernel_path(kernel)
    partial_Fstat_cuda_code = cudacomp.SourceModule(open(kernelfile, "r").read())
    partial_Fstat_cuda = partial_Fstat_cuda_code.get_function(kernel)
    partial_Fstat_cuda.prepare("PIIIIIIIIP")

    # GPU grid setup
    blockRows = min(1024, tCWparams["N_t0Range"])
    blockCols = 1
    gridRows = int(np.ceil(1.0 * tCWparams["N_t0Range"] / blockRows))
    gridCols = 1

    # running the kernel
    logging.debug(
        "Calling pyCUDA kernel with a grid of {}*{}={} blocks"
        " of {}*{}={} threads each: {} total threads...".format(
            gridRows,
            gridCols,
            gridRows * gridCols,
            blockRows,
            blockCols,
            blockRows * blockCols,
            gridRows * gridCols * blockRows * blockCols,
        )
    )
    partial_Fstat_cuda.prepared_call(
        (gridRows, gridCols),
        (blockRows, blockCols, 1),
        input_gpu.gpudata,
        tCWparams["numAtoms"],
        tCWparams["TAtom"],
        tCWparams["t0_data"],
        windowRange.t0,
        windowRange.dt0,
        windowRange.tau,
        windowRange.dtau,
        tCWparams["N_tauRange"],
        Fmn_gpu.gpudata,
    )

    # return results to host
    F_mn = Fmn_gpu.get()

    _print_GPU_memory_MB("Final")

    return F_mn

def pycuda_compute_transient_fstat_map_exp(atomsInputMatrix, windowRange, tCWparams):
    """exponential window, inner and outer loop GPU-parallelized"""

    # gpu data setup and transfer
    _print_GPU_memory_MB("Initial")
    input_gpu = gpuarray.to_gpu(atomsInputMatrix)
    Fmn_gpu = gpuarray.GPUArray(
        (tCWparams["N_t0Range"], tCWparams["N_tauRange"]), dtype=np.float32
    )
    _print_GPU_memory_MB("After input+output allocation:")

    # GPU kernel
    kernel = "cudaTransientFstatExpWindow"
    kernelfile = _get_absolute_kernel_path(kernel)
    partial_Fstat_cuda_code = cudacomp.SourceModule(open(kernelfile, "r").read())
    partial_Fstat_cuda = partial_Fstat_cuda_code.get_function(kernel)
    partial_Fstat_cuda.prepare("PIIIIIIIIIP")

    # GPU grid setup
    blockRows = min(32, tCWparams["N_t0Range"])
    blockCols = min(32, tCWparams["N_tauRange"])
    gridRows = int(np.ceil(1.0 * tCWparams["N_t0Range"] / blockRows))
    gridCols = int(np.ceil(1.0 * tCWparams["N_tauRange"] / blockCols))

    # running the kernel
    logging.debug(
        "Calling kernel with a grid of {}*{}={} blocks"
        " of {}*{}={} threads each: {} total threads...".format(
            gridRows,
            gridCols,
            gridRows * gridCols,
            blockRows,
            blockCols,
            blockRows * blockCols,
            gridRows * gridCols * blockRows * blockCols,
        )
    )
    partial_Fstat_cuda.prepared_call(
        (gridRows, gridCols),
        (blockRows, blockCols, 1),
        input_gpu.gpudata,
        tCWparams["numAtoms"],
        tCWparams["TAtom"],
        tCWparams["t0_data"],
        windowRange.t0,
        windowRange.dt0,
        windowRange.tau,
        windowRange.dtau,
        tCWparams["N_t0Range"],
        tCWparams["N_tauRange"],
        Fmn_gpu.gpudata,
    )

    # return results to host
    F_mn = Fmn_gpu.get()

    _print_GPU_memory_MB("Final")

    return F_mn


# GPL 2.0
#
# written by Gregg Rice, gmr@unc.edu, rice.gregg@gmail.com
# copywrite 2013-2015
# all rights reserved
#

##################################################################################
# render one or two CT secondary structures with helices shown as semicircular arcs
# This simple version of the script does not plot any additional data above the structure(s).
# functionsplotArcRibbons by Steve Busan, modified in 2014 with permission for python hooks Gregg Rice
# om function arcplot
# ----

import RNAtools as RNA
import numpy as np
import sys, argparse, os, traceback, copy, re, math

from matplotlib import rc
#rc('text',usetex=True)

def rgb_int2pct(rgbArr):
    """
    converts RGB 255 values to a 0 to 1 scale

    (255,255,255) --> (1,1,1)
    """
    out = []
    for rgb in rgbArr:
        out.append((rgb[0]/255.0, rgb[1]/255.0, rgb[2]/255.0))
    return out

def parseArgs():
    prs = argparse.ArgumentParser()
    prs.add_argument("dotplot", type=str, help='input dotplot file to get pairing probabilities')
    prs.add_argument("outputPDF",type=str, help="Name of the output PDF graphic")
    prs.add_argument("--referenceCT",type=str, help="reference ct file used for getting the RNA sequence")
    prs.add_argument("--pkDS", type=str, help="Double stranded pk file to optionally add in black arcs for pseudoknots. Two column space seperated text file with pairs in columns")
    prs.add_argument("--secondaryStructure", action="store_true", default=False, help="Plot the secondary structure on top")

    out = prs.parse_args()

    return out

def plotArcRibbons(pairedNuc, ax, color, alpha=1.0, flip=False):
    """
    pairedNuc is a vector length N (seqeunce length) with non-zero values setting the connections
    this is a non-redundant .ct column. Example below
    0 3 0 0 0
      |___|   <-- arc between pos 2 and 4

    ax is the plot axis object

    color is the color of the arc, and alpha is the transparency of the arc

    flip will plot the arcs as rainbows rather than smiles
    """

    from matplotlib.path import Path
    import matplotlib.patches as patches

    handleLengthFactor = 4*(math.sqrt(2)-1)/3
    #vert = 100.0

    i = 0
    while i < len(pairedNuc):
        if pairedNuc[i] > i+1:
            outerPair = [i+0.5,pairedNuc[i]+0.5]
            # find the right side of helix
            lastPairedNuc = pairedNuc[i]
            offset = 1
            while i+offset<len(pairedNuc) and abs(pairedNuc[i+offset]-lastPairedNuc) == 1:
                lastPairedNuc = pairedNuc[i+offset]
                offset += 1
            innerPair = [i+offset+0.5, pairedNuc[i+offset-1]-0.5]
            i += offset-1
            outerRadius = (outerPair[1]-outerPair[0])/2.0
            innerRadius = (innerPair[1]-innerPair[0])/2.0
            #print "innerPair %s, outerPair %s"%(str(innerPair),str(outerPair))
            verts = [
            (outerPair[0], 0), # outer left

            (outerPair[0], -handleLengthFactor*outerRadius), # outer left control 1
            (outerPair[0]+outerRadius-handleLengthFactor*outerRadius, -outerRadius), # outer left control 2
            (outerPair[0]+outerRadius, -outerRadius), # outer center
            (outerPair[0]+outerRadius+handleLengthFactor*outerRadius, -outerRadius), # outer right control 1
            (outerPair[1], -handleLengthFactor*outerRadius), # outer right control 2

            (outerPair[1], 0), # outer right
            (innerPair[1], 0), # inner right

            (innerPair[1], -handleLengthFactor*innerRadius), # inner right control 1
            (innerPair[0]+innerRadius+handleLengthFactor*innerRadius, -innerRadius), # inner right control 2
            (innerPair[0]+innerRadius, -innerRadius), # inner center
            (innerPair[0]+innerRadius-handleLengthFactor*innerRadius, -innerRadius), # inner right control 1
            (innerPair[0], -handleLengthFactor*innerRadius), # inner right control 2

            (innerPair[0], 0), # inner left
            (outerPair[0], 0) # outer left duplicate point
            ]

            if flip:
                flip_offset=2

                verts = [
                (outerPair[0], 0 + flip_offset), # outer left

                (outerPair[0], handleLengthFactor*outerRadius + flip_offset), # outer left control 1
                (outerPair[0]+outerRadius-handleLengthFactor*outerRadius, outerRadius + flip_offset), # outer i left control 2
                (outerPair[0]+outerRadius, outerRadius + flip_offset), # outer center
                (outerPair[0]+outerRadius+handleLengthFactor*outerRadius, outerRadius + flip_offset), # outer right control 1
                (outerPair[1], handleLengthFactor*outerRadius + flip_offset), # outer right control 2

                (outerPair[1], 0 + flip_offset), # outer right
                (innerPair[1], 0 + flip_offset), # inner right

                (innerPair[1], handleLengthFactor*innerRadius + flip_offset), # inner right control 1
                (innerPair[0]+innerRadius+handleLengthFactor*innerRadius, innerRadius + flip_offset), # inner right control 2
                (innerPair[0]+innerRadius, innerRadius + flip_offset), # inner center
                (innerPair[0]+innerRadius-handleLengthFactor*innerRadius, innerRadius + flip_offset), # inner right control 1
                (innerPair[0], handleLengthFactor*innerRadius + flip_offset), # inner right control 2

                (innerPair[0], 0 + flip_offset), # inner left
                (outerPair[0], 0 + flip_offset) # outer left duplicate point
                ]

            for n in range(len(verts)):
                verts[n] = [verts[n][0], verts[n][1]-1.2]#/vert-0.013]
            codes = [
            Path.MOVETO,

            Path.CURVE4,
            Path.CURVE4,
            Path.CURVE4,
            Path.CURVE4,

            Path.LINETO,
            Path.LINETO,
            Path.LINETO,

            Path.CURVE4,
            Path.CURVE4,
            Path.CURVE4,
            Path.CURVE4,

            Path.LINETO,
            Path.LINETO,
            Path.CLOSEPOLY,
            ]
            path = Path(verts, codes)
            patch = patches.PathPatch(path, facecolor=color, linewidth=0, edgecolor='none', alpha=alpha)
            ax.add_patch(patch)
            patch.set_clip_on(False)
        i += 1

def arcplot(outPath="arcs.pdf",title="",seq=["A"],pairedNucArr=[], arcColors = [],alpha=[1.0], maxDistance=None, ct=None, secondary_structure=False):
    """
    draws arcs for many sets of nucleotides, pairedNucArr is a 2D array containing all sets of plotting
    elements. arcColor is the same length as pairedNucArr but contains color strings

    pairedNucArr is M x N where M is the number of plotted elements and N is the length of the RNA to plot
    arcColors is a M x 3 array of RGB tuples for each element
    alpha is a M x 1 array of alpha values to plot

    secondary structure is a booleon that sets whether to plot the structure in the ct file above the arcs representation
    seq is length N

    function duplicated by gmr for plotting several structures on the same plot
    """

    def findMaxDistance(pairingArrays):
        """
        finds the maximum pairing distance from a 2D array of all arc elements
        """
        maxDistance = 0

        for pairedNucA in pairingArrays:
            for i in range(len(pairedNucA)):
                fromNuc = i+1
                toNuc = pairedNucA[i]
                if toNuc == 0:
                    toNuc = fromNuc
                dist = toNuc-fromNuc
                if dist > maxDistance:
                    maxDistance = dist

        return maxDistance

    import matplotlib as mp
    mp.use('Agg')
    mp.rcParams['xtick.major.size'] = 8
    mp.rcParams['xtick.major.width'] = 2.5
    mp.rcParams['xtick.direction'] = 'out'
    mp.rcParams['xtick.minor.size'] = 4
    mp.rcParams['xtick.minor.width'] = 1

    import matplotlib.pyplot as plot
    import matplotlib.patches as patches
    import matplotlib.gridspec as gridspec

    num = range(1,len(seq)+1)

    # adjust the scale factor to fit on a page --gmr
    if len(num)>10000:
        scaleFactor = 0.0005
    if len(num)>500:
        scaleFactor = 0.005
    else:
        scaleFactor = 0.05

    # find longest base-pair and scale height of plot to fit this arc
    if not maxDistance:
        maxDistance = findMaxDistance(pairedNucArr)

    # set up figure dimensions
    figWidth = len(seq)*scaleFactor
    figHeight = maxDistance/2.0*scaleFactor

    # double the figure height if plotting the secondary structure above the plot
    if secondary_structure:
        figHeight = maxDistance * scaleFactor

    fig = plot.figure(figsize=(figWidth, figHeight)) # 500*scaleFactor

    if not secondary_structure:
        ax2 = plot.subplot(111)
    else:
        ax2 = plot.subplot(212)

    plot.xlim(0,len(seq))
    #ax2 = plot.gca()

    # ticks on top
    ax2.get_xaxis().tick_top()

    # determine tick locations based on sequece length
    from matplotlib.ticker import MultipleLocator, FormatStrFormatter
    if len(seq) <= 500:
        majorLocator = MultipleLocator(100)
        minorLocator = MultipleLocator(10)
        interval = 10
    elif len(seq) <= 10000:
        majorLocator = MultipleLocator(500)
        minorLocator = MultipleLocator(100)
        interval = 5
    else:
        majorLocator = MultipleLocator(2500)
        minorLocator = MultipleLocator(500)
        interval = 5

    majorFormatter = FormatStrFormatter('%i')
    minorFormatter = FormatStrFormatter('%i')

    ax2.xaxis.set_major_locator(majorLocator)
    ax2.xaxis.set_major_formatter(majorFormatter)
    ax2.xaxis.set_minor_locator(minorLocator)
    ax2.xaxis.set_minor_formatter(minorFormatter)

    plot.subplots_adjust(hspace=0)
    plot.ylim((-float(maxDistance)/2.0,0))
    #plot.ylim((-1.0,0))

    ax2.set_frame_on(False)
    ax2.axes.get_yaxis().set_visible(False)
    #ax2.axes.get_xaxis().tick_bottom()
    xlabels = ax2.axes.get_xaxis().get_majorticklabels()
    xlabels[0].set_visible(False)
    for label in xlabels:
        label.set_weight('bold')
        label.set_size(14)
        #label.set_rotation(30)
    xlabels = ax2.axes.get_xaxis().get_minorticklabels()
    xlabels[0].set_visible(False)
    labelCount = 0
    for label in xlabels:
        label.set_size(7)
        #label.set_rotation(30)
        # also need to hide minor tick labels that overlap major tick labels
        if labelCount%interval==1:
            label.set_visible(False)
        labelCount += 1

    xticks = ax2.axes.get_xaxis().get_major_ticks()
    xticks[0].set_visible(False)
    xticks = ax2.axes.get_xaxis().get_minor_ticks()
    xticks[0].set_visible(False)

    #bothColor = "green"
    #aColor = "red"
    #bColor = "purple"

    # plot the arcs
    for arc in range(len(pairedNucArr)):
        #if arc % 500 == 0:
        #    print arc, len(pairedNucArr)
        plotArcRibbons(pairedNucArr[arc], ax2, arcColors[arc], alpha=alpha[arc])

    #plotArcRibbons(aOnlyPaired, ax2, aColor, alpha=alpha)
    #plotArcRibbons(bOnlyPaired, ax2, bColor, alpha=alpha)
    #plotArcRibbons(bothPaired, ax2, bothColor, alpha=alpha)

    xmax, xmin, ymin, ymax = plot.axis()

    # put nuc sequence on axis
    if len(seq) <= 500:
        fontProp = mp.font_manager.FontProperties(family = "monospace",
                                              style="normal",
                                              weight="extra bold",
                                              size="4")
        for i in range(len(seq)):
            nuc = seq[i]
            if nuc == "T":
                nuc = "U"
            col = "black"
            plot.annotate(nuc, xy=(i+0.5,ymax),fontproperties=fontProp,color=col,annotation_clip=False,verticalalignment="top")

    if secondary_structure:
        ax1 = plot.subplot(211)
        ax1.set_frame_on(False)
        ax1.axes.get_yaxis().set_visible(False)
        ax2.axes.get_xaxis().set_visible(False)
        plot.xlim(0,len(seq))
        plot.ylim(0,maxDistance/2.0)

        ax1.spines["bottom"].set_position(("axes", -0.025))

        # get the helices from the ct file
        helix = ct.extractHelices(fillPairs=False)

        plot_set = []
        plot_color = []
        ss_color = (0.8,0.8,0.8)

        for k in helix:
            temp = np.zeros_like(ct.ct)
            for nt in helix[k]:
                temp[nt[0]-1] = nt[1]
                plot_set.append(temp)
                plot_color.append(ss_color)
            #print(helix[k])
        #print(helix)
        for hel, col in zip(plot_set, plot_color):
            plotArcRibbons(hel, ax1, col, alpha=1, flip=True)
        #plotArcRibbons(helix, ax1, (0.8,0.8,0.8), alpha=1)
        # xticks = ax1.axes.get_xaxis().get_major_ticks()
        # for k in xticks:
        #     k.set_visible(False)
        # xticks = ax1.axes.get_xaxis().get_minor_ticks()
        # for k in xticks:
        #     k.set_visible(False)

    plot.savefig(outPath,dpi=100,bbox_inches="tight")

def splitPlot(dpObj, ctObj, pk=None, outFile="arcs.pdf", secondary_structure=False):
    """
    splitplot takes a RNAtools dotplot object and a ct object to generate a the figure

    pk is a list pseudoknotted pairs to plot in the bottom panel

    outfile is the destination for the matplotlib output, usually a pdf file or a png

    secondary_structure is a booleon that sets whether to plot the secondary structure above the arc plot
    """
    x = dpObj
    y = ctObj

    # binning is in log10 scale
    #binning = [0.0,0.09691,0.5228,1.0,2.0]
    binning = [1.5228, 1.0, 0.5228, 0.09691, 0.0]
    # 3% 10% 30% 80% 100%

    alphaList = [0.7, 0.7, 0.7, 0.3]
    alphaList = [0.3, 0.7, 0.7, 0.7]
    alphaList = [1.0,1.0,1.0,1.0,]
    #alphaList = [1.0, 1.0, 1.0, 1.0]
    #colorList  = ["red", "orange", "yellow", "green","blue", "violet"]
    #colorList  = [(215, 25, 28), (253, 174, 97), (171, 221, 164), (43, 131, 186)]
    colorList  = [ (43, 131, 186), (171, 221, 164), (253, 174, 97), (215, 25, 28)]
    colorList  = [ (150,150,150), (255,204,0),  (72,143,205) ,(81, 184, 72)  ]
    colorList = rgb_int2pct(colorList)

    nucArr = []
    colors = []
    alpha  = []

    # bin the pairs by cutoff
    for i in range(0, len(binning)-1):

        probPairs = x.requireProb(binning[i],binning[i+1]).pairList()

        for pair in probPairs:

            temp = np.zeros_like(y.ct)
            temp[pair[0]-1] = pair[1]

            #tempCT = RNA.CT()
            #tempCT.pair2CT([pair],y.seq)

            #nucArr.append(tempCT.stripCT())
            nucArr.append(temp)

            #add a color from the choice list
            colors.append(colorList[i])
            alpha.append(alphaList[i])

    if pk:
        for pair in pk:
            temp = np.zeros_like(y.ct)
            temp[pair[0]-1] = pair[1]
            nucArr.append(temp)
            colors.append((0,0,0))
            alpha.append(0.8)
    #nucArr.append(y.stripCT())
    #colors.append("gray")

    arcplot(outPath=outFile, pairedNucArr=nucArr, arcColors=colors,seq=y.seq, alpha=alpha, maxDistance=None, secondary_structure=secondary_structure, ct=y)

def readPKfile(fIN):
    out = []
    for i in open(fIN).readlines():
        if i.lstrip()[0] == '#':
            continue
        pk = map(int, i.rstrip().split())
        out.append((pk[0],pk[1]))
    return out

def main():
    # parse the command line options
    arg = parseArgs()

    # read in the reference files
    x = RNA.dotPlot( arg.dotplot )
    if arg.referenceCT:
        y = RNA.CT( arg.referenceCT )
    else:
        # if we dont have a sequence create an empty CT file
        y = RNA.CT()
        # make the sequence all N as a placeholder
        seq = "N"*x.length
        y.pair2CT([], seq)

    #print(y)

    # add correction for slipped base pairs
    #x.averageSlippedBPs(y,predictedOnly=True)

    # read in the pseudoknots file if it exists
    if arg.pkDS:
        pkntsFile = arg.pkDS
        pknts = readPKfile(pkntsFile)
    else:
        pknts = None

    # main plotting loop
    splitPlot(x, y, pk=pknts, outFile=arg.outputPDF, secondary_structure=arg.secondaryStructure)

if __name__ == '__main__':
    main()


"""
Copyright (C) 2012 Alan J Lockett

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import numpy as np

from .basic import PopulationDistribution, GaussianProposal
from pyec.distribution.bayes.mutators import StructureMutator
from pyec.distribution.bayes.structure.proposal import StructureProposal
from pyec.distribution.ec.mutators import Gaussian, Bernoulli

from pyec.config import Config
from pyec.history import DoubleMarkovHistory

_ = Config

class SimulatedAnnealingAcceptance(PopulationDistribution):
   """A selector that implements the acceptance probability for simulated
   annealing. Computes the acceptance ratio and samples it. If the population
   size is greater than one, then this distribution maintains an array
   of accepted values, and can be used to run multiple concurrent
   markov chains.

   Config parameters

   * schedule -- The cooling schedule to use. May be any callable function, or
                 a string, one of ("log", "linear", "discount"). If it is a
                 callable, then it will be passed the ``updates`` value in
                 :class:`History` and should return a floating point value
                 that will divide the exponent in the Boltzmann distribution.
                 That is, ``schedule(n)`` should converge to zero as n goes to
                 infinity.
   * learningRate -- A divisor for the cooling schedule, used for built-in
                     schedules "log" and "linear". As a divisor, it divides the
                     temperature but multiplies the exponent.
   * temp0 -- An initial temperature for the temperature decay in "discount"
   * restart -- A probability of restarting, tested at each update
   * divisor -- A divisor that divides the ``updates`` property of
                :class:`History`, scaling the rate of decline in the temperature
   * discount -- The decay factor for the built-in discount schedule; ``temp0``
                 is multiplied by ``discount`` once each time the ``update``
                 method is called

   """
   config = Config(schedule="log",
                   learningRate = 1.0,
                   temp0 = 1.0,
                   restart = 0.0,
                   divisor = 100.0,
                   discount = .99,
                   populationSize = 1,
                   history = DoubleMarkovHistory)

   def compatible(self, history):
      return (hasattr(history, 'lastPopulation')
              and hasattr(history, 'penultimate')
              and hasattr(history, 'reportAcceptance'))

   def batch(self, popSize):
      temp = self.temperature()
      last = self.history.lastPopulation()
      penultimate = self.history.penultimate()
      if penultimate is None:
         return [x for x,s in last]
      scoreProposed = np.array([s for x,s in last])
      scoreAccepted = np.array([s for x,s in penultimate])
      exponent = (scoreProposed - scoreAccepted) / temp
      if self.config.minimize:
         exponent = -exponent
      probs = np.minimum(1.0, np.exp(exponent))
      selection = np.random.binomial(1, probs, np.shape(probs))
      accepted = 0.0
      result = []
      for i,sel in enumerate(selection):
         if sel > 0.5:
            result.append(last[i][0])
            accepted += 1.0
         else:
            result.append(penultimate[i][0])
      self.history.reportAcceptance(accepted / popSize)
      return result

   def temperature(self):
      n = 1 + float(self.history.updates) / self.config.divisor
      if hasattr(self.config.schedule, '__call__'):
         return self.config.schedule(n)
      elif self.config.schedule == "linear":
         return 1. / (n * self.config.learningRate)
      elif self.config.schedule == "log":
         return 1. / (np.log(n) * self.config.learningRate)
      elif self.config.schedule == "discount":
         return 1. / (self.config.temp0 * (self.config.discount ** n))

# Euclidean space
RealSimulatedAnnealing = (
   SimulatedAnnealingAcceptance << GaussianProposal[_(sd=.005)]
   #GaussianProposal[_(sd=.005)] << SimulatedAnnealingAcceptance
)

# fixed-length bit strings
BinarySimulatedAnnealing = (
   SimulatedAnnealingAcceptance << Bernoulli[_(p=.01)]
   #Bernoulli[_(p=.01)] << SimulatedAnnealingAcceptance
)

# Structure search in a Bayes net, use a
# pyec.distribution.bayes.space.BayesNetStructure space for searching.
BayesNetSimulatedAnnealing = (
   SimulatedAnnealingAcceptance[_(schedule="linear",
                                  divisor=100.)] <<
   StructureMutator[_(branchFactor=5)]
)[_(minimize=False)]