import os
import numpy as np
from pdb2sql.interface import interface

try:
    import freesasa

except ImportError:
    print('Freesasa not found')

class BSA(object):

    def __init__(self, pdb_data, sqldb=None, chainA='A', chainB='B'):
        '''Compute the burried surface area feature

        Freesasa is required for this feature.

        https://freesasa.github.io

        >>> wget http://freesasa.github.io/freesasa-2.0.2.tar.gz
        >>> tar -xvf freesasa-2.0.3.tar.gz
        >>> cd freesasa
        >>> ./configure CFLAGS=-fPIC (--prefix /home/<user>/)
        >>> make
        >>> make install

        Since release 2.0.3 the python bindings are separate module
        >>> pip install freesasa

        Args :
            pdb_data (list(byte) or str): pdb data or filename of the pdb
            sqldb (pdb2sql.interface instance or None, optional) if the sqldb is None the sqldb will be created
            chainA (str, optional): name of the first chain
            chainB (str, optional): name of the second chain

        Example :

        >>> bsa = BSA('1AK4.pdb')
        >>> bsa.get_structure()
        >>> bsa.get_contact_residue_sasa()
        >>> bsa.sql.close()

        '''

        self.pdb_data = pdb_data
        if sqldb is None:
            self.sql = interface(pdb_data)
        else:
            self.sql = sqldb
        self.chains_label = [chainA, chainB]

        freesasa.setVerbosity(freesasa.nowarnings)

    def get_structure(self):
        """Get the pdb structure of the molecule."""

        # we can have a str or a list of bytes as input
        if isinstance(self.pdb_data, str):
            self.complex = freesasa.Structure(self.pdb_data)
        else:
            self.complex = freesasa.Structure()
            atomdata = self.sql.get(
                'name,resName,resSeq,chainID,x,y,z')
            for atomName, residueName, residueNumber, chainLabel, x, y, z in atomdata:
                atomName = '{:>2}'.format(atomName[0])
                self.complex.addAtom(
                    atomName, residueName, residueNumber, chainLabel, x, y, z)
        self.result_complex = freesasa.calc(self.complex)

        self.chains = {}
        self.result_chains = {}
        for label in self.chains_label:
            self.chains[label] = freesasa.Structure()
            atomdata = self.sql.get(
                'name,resName,resSeq,chainID,x,y,z', chainID=label)
            for atomName, residueName, residueNumber, chainLabel, x, y, z in atomdata:
                atomName = '{:>2}'.format(atomName[0])
                self.chains[label].addAtom(
                    atomName, residueName, residueNumber, chainLabel, x, y, z)
            self.result_chains[label] = freesasa.calc(
                self.chains[label])

    def get_contact_residue_sasa(self, cutoff=8.5):
        """Compute the feature value."""

        self.bsa_data = {}
        self.bsa_data_xyz = {}

        res = self.sql.get_contact_residues(cutoff=cutoff)
        keys = list(res.keys())
        res = res[keys[0]]+res[keys[1]]

        for r in res:

            # define the selection string and the bsa for the complex
            select_str = ('res, (resi %d) and (chain %s)' %
                          (r[1], r[0]),)
            asa_complex = freesasa.selectArea(
                select_str, self.complex, self.result_complex)['res']

            # define the selection string and the bsa for the isolated
            select_str = ('res, resi %d' % r[1],)
            asa_unbound = freesasa.selectArea(
                select_str, self.chains[r[0]], self.result_chains[r[0]])['res']

            # define the bsa
            bsa = asa_unbound-asa_complex

            # define the xyz key : (chain,x,y,z)
            chain = {'A': 0, 'B': 1}[r[0]]
            xyz = np.mean(self.sql.get(
                'x,y,z', resSeq=r[1], chainID=r[0]), 0)
            xyzkey = tuple([chain]+xyz.tolist())

            # put the data in dict
            self.bsa_data[r] = [bsa]


"""Some helper functions."""
import json
import numpy as np
import random
import vocabulary as tweet_vocal

def RandomizeAnTweetFromFile():
  """Used in PostTweet. Randomize a sentence from a list."""
  tweet_lst = tweet_vocal.tweet_lst
  samples = 1
  if random.random() > 0.5:
    samples = 2
  sentence = ''
  for i in range(samples):
    sentence += np.random.choice(tweet_lst, 1)[0]
  return sentence

# key_dict = {entities:{user_mentions:[id]}, user:id}
# same_level_obj = [retweeted_status]
def RetriveUserIds(json_obj):
  """Get all user ids in a list of Tweet objects."""
  rst_set = set()

  def _helper(current_obj):
    if 'entities' in current_obj:
      if 'user_mentions' in current_obj['entities']:
        if current_obj['entities']['user_mentions']:
          for user_obj in current_obj['entities']['user_mentions']:
            rst_set.add(user_obj['id'])

    if 'user' in current_obj:
      rst_set.add(current_obj['user']['id'])

    if 'retweeted_status' in current_obj:
      _helper(current_obj['retweeted_status'])

  for tweet_obj in json_obj:
    _helper(tweet_obj)

  return list(rst_set)

def randomizeSource(prob_dict):
  """ Weighted randomized an item in prob_dict according to their weight.

  Args:
    prob_dict: item (str) -> probability (float)

  Returns:
    item (str)
  """
  keys = list(prob_dict)
  probs = [prob_dict[k] for k in keys]
  draw = np.random.choice(keys, 1, p=probs)[0]
  return draw

def randomizeObjs(
  obj_list, num_returns=1, return_key=None,
  whether_weighted=False, weighted_key=None,
  weighted_order=False):
  """ Randomize objects from a list of objects.

  All objects should have same structure.
  Args:
    obj_list: a list of objects.
    return_key: if it's none, the function returns the selected object.
                Otherwise, return the value of the return_key in this selected
                object. We don't support nested keys
    whether_weighted: if it's False, the function would randomize one uniformly.
                      Else, the function will take weighted_key as the key to
                      randomize one item. To enable the weighted randomization,
                      weighted_key or weighted_order must be defined.
    weighted_key: A key in object. We don't support nested keys, and the values
                  associating with the key must be numbers.
    weighted_order: weighted probability based on the order(index) of the objects
                    in this list. The smaller index (such 0, 1) has a higher
                    probability

  Returns:
    An object or an attribute in object

  """
  if not obj_list:
    return None
  if not whether_weighted:
    selected_obj = np.random.choice(obj_list, num_returns)
  else:
    if weighted_key:
      weight_lst = [float(o[weighted_key]) for o in obj_lst]
    elif weighted_order:
      lst_len = len(obj_list)
      weight_lst = [lst_len - idx for idx in range(lst_len)]
    total_weight = sum(weight_lst)
    weight_lst = [w/total_weight for w in weight_lst]
    selected_obj = np.random.choice(obj_list, num_returns, weight_lst)

  if num_returns == 1:
    if return_key:
      selected_obj = selected_obj[0][return_key]
    else:
      selected_obj = selected_obj[0]
  else:
    if return_key:
      tmp_obj = []
      for obj in selected_obj:
        tmp_obj.append(obj[return_key])
      selected_obj = tmp_obj
  return selected_obj


# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from abc import ABC, abstractmethod
from .crf import CRF
from .attention import BilinearMatrixAttention

import pyximport; pyximport.install(setup_args={'include_dirs': np.get_include()})
from .calgorithm import parse_proj, parse_mwe, parse_mwe_ud
from .crf_algorithm import crf_marginal

class ParserModule(ABC):

    @property
    @classmethod
    @abstractmethod
    def name(cls):
        pass

    @staticmethod
    @abstractmethod
    def load_data(parser, graph):
        pass

    @staticmethod
    @abstractmethod
    def batch_label(batch):
        pass

    @abstractmethod
    def evaluate(self, results, pred, gold, mask):
        pass

    @abstractmethod
    def metrics(self, results):
        pass

class SequenceLabeler(nn.Module, ParserModule):

    def __init__(self, parser, layer_size, hidden_size, label_size, dropout=0., crf=False):
        super(SequenceLabeler, self).__init__()
        print("build sequence labeling network...", self.__class__.name, "crf:", crf)

        self.use_crf = crf
        ## add two more label for downlayer lstm, use original label size for CRF
        self.label_size = label_size + 2

        lst = []
        for i in range(layer_size):
            if i == 0:
                lst.append(nn.Linear(parser._bilstm_dims, hidden_size))
            else:
                lst.append(nn.Linear(hidden_size, hidden_size))

            lst.append(nn.PReLU())
            lst.append(nn.Dropout(dropout))

        if layer_size > 0:
            lst.append(nn.Linear(hidden_size, self.label_size))
        else:
            lst.append(nn.Linear(parser._bilstm_dims, self.label_size))

        self.transform = nn.Sequential(*lst)

        if self.use_crf:
            self.crf = CRF(label_size)
        else:
            self.loss = nn.NLLLoss(ignore_index=0, reduction='sum')

    def calculate_loss(self, lstm_features, batch):
        batch_label = self.batch_label(batch)
        mask = batch["mask"]
        mask_h = batch["mask_h"]
        outs = self.transform(lstm_features)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)
        if self.use_crf:
            total_loss = self.crf.neg_log_likelihood_loss(outs, mask_h, batch_label)
            total_loss = total_loss / float(batch_size)
            scores, tag_seq = self.crf._viterbi_decode(outs, mask_h)
        else:
            outs = outs.view(batch_size * seq_len, -1)
            score = F.log_softmax(outs, 1)
            total_loss = self.loss(score, batch_label.view(batch_size * seq_len))
            total_loss = total_loss / float(batch_size)
            _, tag_seq  = torch.max(score, 1)
            tag_seq = tag_seq.view(batch_size, seq_len)

        return total_loss, tag_seq

    def forward(self, parser, lstm_features, batch):
        batch_label = self.batch_label(batch)
        mask = batch["mask"]
        mask_h = batch["mask_h"]
        outs = self.transform(lstm_features)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)
        if self.use_crf:
            scores, tag_seq = self.crf._viterbi_decode(outs, mask_h)
        else:
            outs = outs.view(batch_size * seq_len, -1)
            _, tag_seq  = torch.max(outs, 1)
            tag_seq = tag_seq.view(batch_size, seq_len)
            ## filter padded position with zero
            tag_seq = mask.long() * tag_seq

        batch["pred_" + self.name] = tag_seq

        return tag_seq

    def scores(self, lstm_features, batch):
        mask = batch["mask"]
        mask_h = batch["mask_h"]
        batch_size = mask.size(0)
        seq_len = mask.size(1)
        if self.use_crf:
            outs = self.transform(lstm_features).view(batch_size, seq_len, -1)
            scores = outs.cpu().data.numpy().astype('float64')
            transitions = self.crf.transitions.cpu().data.numpy().astype('float64')
            lengths = torch.sum(mask_h.long(), dim=1).view(batch_size,).long().cpu().data.numpy()
            ret = torch.zeros_like(outs)

            for i in range(batch_size):
                length = lengths[i]
                ret[i, :length, :] = torch.Tensor(crf_marginal(scores[i, :length, :], transitions))
            return ret
        else:
            outs = self.transform(lstm_features).view(batch_size * seq_len, -1)
            return F.log_softmax(outs, dim=1).view(batch_size, seq_len, -1)

    def evaluate(self, results, pred, gold, mask):
        overlaped = (pred == gold).byte()
        correct = float((overlaped * mask).sum())
        total = float(mask.sum())
        results["{}-c".format(self.__class__.name)] += correct
        results["{}-t".format(self.__class__.name)] += total

    def metrics(self, results):
        correct = results["{}-c".format(self.__class__.name)]
        total = results["{}-t".format(self.__class__.name)]
        results["metrics/{}-acc".format(self.__class__.name)] = correct / (total + 1e-10) * 100.
        del results["{}-c".format(self.__class__.name)]
        del results["{}-t".format(self.__class__.name)]

class XPOSTagger(SequenceLabeler):

    name = "XPOS"

    @staticmethod
    def load_data(parser, graph):
        labels = [0] + [parser._xpos.get(n.xpos, 1) for n in graph.nodes[1:]]
        return {"xpos": labels}

    @staticmethod
    def batch_label(batch):
        return batch["xpos"]

class UPOSTagger(SequenceLabeler):

    name = "UPOS"

    @staticmethod
    def load_data(parser, graph):
        labels = [0] + [parser._upos.get(n.upos, 1) for n in graph.nodes[1:]]
        return {"upos": labels}

    @staticmethod
    def batch_label(batch):
        return batch["upos"]

class NERBIOTagger(SequenceLabeler):

    name = "NERBIO"

    @staticmethod
    def load_data(parser, graph):
        ner_dic = {}
        labels = [3 for i in range(len(graph.nodes))]
        for i in range(1, len(graph.nodes)):
            if graph.nodes[i].feats == "B":
                labels[i] = 1
            elif graph.nodes[i].feats == "I":
                labels[i] = 2
        return {"nerbio": labels}

    @staticmethod
    def batch_label(batch):
        return batch["nerbio"]

    def evaluate(self, results, pred, gold, mask):
        batch_size, length = mask.shape
        pred = pred.cpu().data.numpy()
        gold = gold.cpu().data.numpy()
        mask = mask.cpu().data.numpy()
        pred_set = set()
        gold_set = set()
        for i in range(batch_size):
            curstart = None
            for j in range(1, length):
                if mask[i, j] == 0:
                    if curstart is not None:
                        pred_set.add((i, curstart, j))
                    break
                if pred[i, j] == 1:
                    if curstart is not None:
                        pred_set.add((i, curstart, j))
                    curstart = j
                elif pred[i, j] == 2:
                    if curstart is None:
                        curstart = j
                elif pred[i, j] == 3:
                    if curstart is not None:
                        pred_set.add((i, curstart, j))
                    curstart = None

        for i in range(batch_size):
            curstart = None
            for j in range(1, length):
                if mask[i, j] == 0:
                    if curstart is not None:
                        gold_set.add((i, curstart, j))
                    break
                if gold[i, j] == 1:
                    if curstart is not None:
                        gold_set.add((i, curstart, j))
                    curstart = j
                elif gold[i, j] == 2:
                    if curstart is None:
                        curstart = j
                elif gold[i, j] == 3:
                    if curstart is not None:
                        gold_set.add((i, curstart, j))
                    curstart = None

        results["{}-p".format(self.__class__.name)] += len(pred_set)
        results["{}-r".format(self.__class__.name)] += len(gold_set)
        results["{}-c".format(self.__class__.name)] += len(pred_set.intersection(gold_set))

    def metrics(self, results):
        p = results["{}-p".format(self.__class__.name)]
        r = results["{}-r".format(self.__class__.name)]
        c = results["{}-c".format(self.__class__.name)]
        precision = c / (p + 1e-6)
        recall = c / (r + 1e-6)
        results["metrics/{}-p".format(self.__class__.name)] = precision * 100.
        results["metrics/{}-r".format(self.__class__.name)] = recall * 100.
        results["metrics/{}-f1".format(self.__class__.name)] = 2. / (1./(precision + 1e-6) + 1./(recall+1e-6)) * 100.
        del results["{}-p".format(self.__class__.name)]
        del results["{}-r".format(self.__class__.name)]
        del results["{}-c".format(self.__class__.name)]

class PointerSelector(nn.Module, ParserModule):

    def __init__(self, parser, hidden_size, dropout=0.):
        super(PointerSelector, self).__init__()
        print("build pointer selector ...", self.__class__.name)
        self.head_mlp = nn.Sequential(
            nn.Linear(parser._bilstm_dims, hidden_size),
            nn.LeakyReLU(0.1),
            nn.Dropout(dropout),
        )

        self.dep_mlp = nn.Sequential(
            nn.Linear(parser._bilstm_dims, hidden_size),
            nn.LeakyReLU(0.1),
            nn.Dropout(dropout),
        )

        self.attention = BilinearMatrixAttention(hidden_size, hidden_size, True)
        self.loss = nn.NLLLoss(ignore_index=-1, reduction='sum')

    def calculate_loss(self, lstm_features, batch):
        batch_label = self.batch_label(batch)
        mask = batch["mask"]
        mask_h = self.batch_cand_mask(batch)

        heads = self.head_mlp(lstm_features)
        deps = self.dep_mlp(lstm_features)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)

        mask_att = mask.unsqueeze(2) * mask_h.unsqueeze(1)
        scores = self.attention(deps, heads).masked_fill((1-mask_att).bool(), float("-inf")).view(batch_size * seq_len, -1)
        scores = F.log_softmax(scores, dim=1)

        _, tag_seq  = torch.max(scores, 1)
        tag_seq = tag_seq.view(batch_size, seq_len) + 1

        total_loss = self.loss(scores, (batch_label - 1).view(batch_size * seq_len))
        total_loss = total_loss / float(batch_size)

        return total_loss, tag_seq

    def forward(self, parser, lstm_features, batch):
        batch_label = self.batch_label(batch)
        mask = batch["mask"]

        mask_h = self.batch_cand_mask(batch)
        heads = self.head_mlp(lstm_features)
        deps = self.dep_mlp(lstm_features)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)

        mask_att = mask.unsqueeze(2) * mask_h.unsqueeze(1)
        scores = self.attention(deps, heads).masked_fill((1-mask_att).bool(), float("-inf")).view(batch_size * seq_len, -1)
        scores = F.log_softmax(scores, dim=1)

        _, tag_seq  = torch.max(scores, 1)
        tag_seq = tag_seq.view(batch_size, seq_len) + 1

        scores = scores.view(batch_size, seq_len, -1).cpu().data.numpy().astype('float64')
        word_length = batch["word_length"].cpu().data.numpy()
        for i in range(batch_size):
            l = int(word_length[i])
            s = scores[i, :l, :l].T
            heads = parse_proj(s)
            tag_seq[i, :l] = torch.Tensor(heads + 1)

        ## filter padded position with zero
        tag_seq = mask.long() * tag_seq

        batch["pred_head"] = tag_seq

        return tag_seq

    def evaluate(self, results, pred, gold, mask):
        overlaped = (pred == gold).byte()
        # correct = np.sum(overlaped * mask)
        correct = float((overlaped * mask).sum())
        total = float(mask.sum())
        results["{}-c".format(self.__class__.name)] += correct
        results["{}-t".format(self.__class__.name)] += total

    def metrics(self, results):
        correct = results["{}-c".format(self.__class__.name)]
        total = results["{}-t".format(self.__class__.name)]
        results["metrics/{}-acc".format(self.__class__.name)] = correct / (total + 1e-10) * 100.
        del results["{}-c".format(self.__class__.name)]
        del results["{}-t".format(self.__class__.name)]

    @staticmethod
    @abstractmethod
    def batch_cand_mask(batch):
        pass

class HSelParser(PointerSelector):

    name = "HSel"

    @staticmethod
    def load_data(parser, graph):
        return {"head": graph.heads + 1}

    @staticmethod
    def batch_label(batch):
        return batch["head"]

    @staticmethod
    def batch_cand_mask(batch):
        return batch["mask_h"]

class RelLabeler(nn.Module, ParserModule):

    name = "Rel"

    def __init__(self, parser, hidden_size, dropout=0.):
        super(RelLabeler, self).__init__()
        print("build rel labeler...", self.__class__.name)
        self.head_mlp = nn.Sequential(
            nn.Linear(parser._bilstm_dims, hidden_size),
            nn.LeakyReLU(0.1),
            nn.Dropout(dropout),
        )

        self.dep_mlp = nn.Sequential(
            nn.Linear(parser._bilstm_dims, hidden_size),
            nn.LeakyReLU(0.1),
            nn.Dropout(dropout),
        )

        self.attention = nn.Bilinear(hidden_size, hidden_size, len(parser._rels) + 1, True)
        self.bias_x = nn.Linear(hidden_size, len(parser._rels) + 1, False)
        self.bias_y = nn.Linear(hidden_size, len(parser._rels) + 1, False)
        self.loss = nn.NLLLoss(ignore_index=-1, reduction='sum')

    def calculate_loss(self, lstm_features, batch):
        batch_label = self.batch_label(batch)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)

        mask = batch["mask"]
        head = torch.abs(batch["head"] - 1)

        ran = torch.arange(batch_size, device=head.get_device()).unsqueeze(1) * seq_len
        idx = (head + ran).view(batch_size * seq_len)

        heads = self.head_mlp(lstm_features).view(batch_size * seq_len, -1)[idx]
        deps = self.dep_mlp(lstm_features).view(batch_size * seq_len, -1)

        scores = self.attention(deps, heads) + self.bias_x(heads) + self.bias_y(deps)
        scores = F.log_softmax(scores, dim=1)

        _, tag_seq  = torch.max(scores, 1)
        tag_seq = tag_seq.view(batch_size, seq_len)

        total_loss = self.loss(scores, batch_label.view(batch_size * seq_len))
        total_loss = total_loss / float(batch_size)

        return total_loss, tag_seq

    def forward(self, parser, lstm_features, batch):
        batch_label = self.batch_label(batch)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)

        mask = batch["mask"]
        if "pred_head" in batch:
            head = torch.abs(batch["pred_head"] - 1)
        else:
            head = torch.abs(batch["head"] - 1)

        ran = torch.arange(batch_size, device=head.get_device()).unsqueeze(1) * seq_len
        idx = (head + ran).view(batch_size * seq_len)

        heads = self.head_mlp(lstm_features).view(batch_size * seq_len, -1)[idx]
        deps = self.dep_mlp(lstm_features).view(batch_size * seq_len, -1)

        scores = self.attention(deps, heads) + self.bias_x(heads) + self.bias_y(deps)
        scores = F.log_softmax(scores, dim=1)

        _, tag_seq  = torch.max(scores, 1)
        tag_seq = tag_seq.view(batch_size, seq_len)
        ## filter padded position with zero
        tag_seq = mask.long() * tag_seq

        batch["pred_rel"] = tag_seq

        return tag_seq

    def scores(self, lstm_features, batch):
        batch_label = self.batch_label(batch)
        batch_size = batch_label.size(0)
        seq_len = batch_label.size(1)

        heads = self.head_mlp(lstm_features)
        deps = self.dep_mlp(lstm_features)

        ret = torch.matmul(deps.unsqueeze(1), self.attention.weight) @ torch.transpose(heads.unsqueeze(1), -1, -2)
        ret = ret + self.bias_x(heads).transpose(-1, -2).unsqueeze(-2)
        ret = ret + self.bias_y(deps).transpose(-1, -2).unsqueeze(-1)
        ret = ret + self.attention.bias.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)

        return F.log_softmax(ret, dim=1).transpose(-1, -2)

    def evaluate(self, results, pred, gold, mask):
        overlaped = (pred == gold).byte()
        correct = float((overlaped * mask).sum())
        total = float(mask.sum())
        results["{}-c".format(self.__class__.name)] += correct
        results["{}-t".format(self.__class__.name)] += total

    def metrics(self, results):
        correct = results["{}-c".format(self.__class__.name)]
        total = results["{}-t".format(self.__class__.name)]
        results["metrics/{}-acc".format(self.__class__.name)] = correct / (total + 1e-10) * 100.
        del results["{}-c".format(self.__class__.name)]
        del results["{}-t".format(self.__class__.name)]

    @staticmethod
    def load_data(parser, graph):
        labels = [0] + [parser._rels.get(r, 1) for r in graph.rels[1:]]
        return {"rel": labels}

    @staticmethod
    def batch_label(batch):
        return batch["rel"]

class MWEJointDecoder(nn.Module, ParserModule):

    name = "JointMWE"

    def __init__(self, hselparser, rellabeler, nerbiotagger):
        super(MWEJointDecoder, self).__init__()
        print("build MWE joint decoder ...")

        class Object(object):
            pass

        self.c = Object()

        self.c.hselparser = hselparser
        self.c.rellabeler = rellabeler
        self.c.nerbiotagger = nerbiotagger

    @staticmethod
    def load_data(parser, graph):
        return {}

    @staticmethod
    def batch_label(batch):
        return batch["head"]

    def evaluate(self, results, pred, gold, mask):
        overlaped = (pred == gold).byte()
        correct = float((overlaped * mask).sum())
        total = float(mask.sum())
        results["{}-c".format(self.__class__.name)] += correct
        results["{}-t".format(self.__class__.name)] += total

    def metrics(self, results):
        correct = results["{}-c".format(self.__class__.name)]
        total = results["{}-t".format(self.__class__.name)]
        results["metrics/{}-acc".format(self.__class__.name)] = correct / (total + 1e-10) * 100.
        del results["{}-c".format(self.__class__.name)]
        del results["{}-t".format(self.__class__.name)]

    def calculate_loss(self, lstm_features, batch):
        return None, None

    def forward(self, parser, lstm_features, batch):
        mask = batch["mask"]
        batch_size = mask.size(0)
        seq_len = mask.size(1)

        mask_h = self.c.hselparser.batch_cand_mask(batch)
        heads = self.c.hselparser.head_mlp(lstm_features)
        deps = self.c.hselparser.dep_mlp(lstm_features)

        mask_att = mask.unsqueeze(2) * mask_h.unsqueeze(1)
        scores = self.c.hselparser.attention(deps, heads).masked_fill((1-mask_att).bool(), float("-inf")).view(batch_size * seq_len, -1)
        scores = F.log_softmax(scores, dim=1)

        bio_scores = self.c.nerbiotagger.scores(lstm_features, batch).cpu().data.numpy().astype('float64')

        _, tag_seq  = torch.max(scores, 1)
        tag_seq = tag_seq.view(batch_size, seq_len) + 1
        _, rel_seq  = torch.max(scores, 1)
        rel_seq = rel_seq.view(batch_size, seq_len) + 1

        rel_scores = self.c.rellabeler.scores(lstm_features, batch).cpu().data.numpy().astype('float64')

        scores = scores.view(batch_size, seq_len, -1).cpu().data.numpy().astype('float64')
        word_length = batch["word_length"].cpu().data.numpy()
        for i in range(batch_size):
            l = int(word_length[i])
            s = scores[i, :l, :l].T
            t = bio_scores[i, :l]

            if "mwe_NNP" in parser._rels:
                heads, rels = parse_mwe(s, rel_scores[i, :, :l, :l], t * 1., parser._rels["mwe_NNP"])
            else:
                heads, rels = parse_mwe_ud(s, rel_scores[i, :, :l, :l], t * 1., parser._rels["flat"], parser._rels["punct"])
            tag_seq[i, :l] = torch.Tensor(heads + 1)
            rel_seq[i, :l] = torch.Tensor(rels)

        ## filter padded position with zero
        tag_seq = mask.long() * tag_seq
        rel_seq = mask.long() * rel_seq

        batch["pred_head"] = tag_seq
        batch["pred_rel"] = rel_seq

        return tag_seq


# 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.
"""External function interface to CuDNN v7 library."""
# pylint: disable-msg=C0103
import ctypes
import numpy as np
import tvm

import tvm._ffi
from tvm import te

# algos can be read from cudnn.h
_FWD_ALGOS = [
    "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM",
    "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM",
    "CUDNN_CONVOLUTION_FWD_ALGO_GEMM",
    "CUDNN_CONVOLUTION_FWD_ALGO_DIRECT",
    "CUDNN_CONVOLUTION_FWD_ALGO_FFT",
    "CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING",
    "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD",
    "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED",
    "CUDNN_CONVOLUTION_FWD_ALGO_COUNT",
]

_BWD_FILTER_ALGOS = [
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0",
    # non-deterministic
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1",
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT",
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_3",
    # non-deterministic, algo0 with workspaceS
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD",
    # not implemented
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED",
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING",
    "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT",
]

_BWD_DATA_ALGOS = [
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_0",
    # non-deterministic
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_1",
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT",
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING",
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD",
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED",
    "CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT",
]

_ALGO_TYPE = ["fwd", "bwd_filter", "bwd_data"]

def exists():
    """
    Checks whether the local machine can use CuDNN.

    Returns
    -------
        exists: bool

            True if CuDNN support is enabled and a CuDNN-capable GPU
            exists.  Otherwise, False.
    """
    func = tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
    if func is None:
        return False

    return bool(func())

def algo_to_index(algo_type, algo_name):
    """Return a index represents the algorithm, which can be used in
    calling CuDNN function

    Parameters
    ----------
        algo_type : str
            ["fwd", "bwd_filter", "bwd_data]

        algo_name : str
            algorithm name in cudnn definition
            fwd = [
                "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM",
                "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM",
                "CUDNN_CONVOLUTION_FWD_ALGO_GEMM",
                "CUDNN_CONVOLUTION_FWD_ALGO_DIRECT",
                "CUDNN_CONVOLUTION_FWD_ALGO_FFT",
                "CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING",
                "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD",
                "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED",
                "CUDNN_CONVOLUTION_FWD_ALGO_COUNT",
            ]
            bwd_filter = [
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0",
                # non-deterministic
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1",
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT",
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_3",
                # non-deterministic, algo0 with workspaceS
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD",
                # not implemented
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED",
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING",
                "CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT",
            ]
            bwd_data = [
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_0",
                # non-deterministic
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_1",
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT",
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING",
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD",
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED",
                "CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT",
            ]

    Returns
    -------
        algo: int
            algorithm index

    """
    idx = -1
    if algo_type == "fwd":
        idx = _FWD_ALGOS.index(algo_name)
    elif algo_type == "bwd_filter":
        idx = _BWD_FILTER_ALGOS.index(algo_name)
    elif algo_type == "bwd_data":
        idx = _BWD_DATA_ALGOS.index(algo_name)
    assert idx >= 0
    return idx

def _get_np_int32_array_handle(arr):
    """Return a void_p handle for a numpy array

    Parameters
    ----------
    arr: numpy.NDArray
        source numpy array

    Returns
    -------
    ptr:  ctypes.c_void_p
        pointer to the data
    """
    assert arr.dtype == np.int32
    ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
    return ctypes.cast(ptr, ctypes.c_void_p)

def _prepare_global_func_params(dims, pad, stride, dilation, x_shape=None, w_shape=None):
    full_dims = dims + 2
    if x_shape:
        assert isinstance(x_shape, list)
        assert len(x_shape) == full_dims
    if w_shape:
        assert isinstance(w_shape, list)
        assert len(w_shape) == full_dims

    pad = (
        np.full(dims, pad, dtype=np.int32)
        if isinstance(pad, int)
        else np.array(pad, dtype=np.int32)
    )
    stride = (
        np.full(dims, stride, dtype=np.int32)
        if isinstance(stride, int)
        else np.array(stride, dtype=np.int32)
    )
    dilation = (
        np.full(dims, dilation, dtype=np.int32)
        if isinstance(dilation, int)
        else np.array(dilation, dtype=np.int32)
    )

    xshape = np.array(x_shape, dtype=np.int32) if x_shape else None
    wshape = np.array(w_shape, dtype=np.int32) if x_shape else None

    return pad, stride, dilation, xshape, wshape

def conv_output_shape(
    tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1
):
    """Get output shape of 2D or 3D convolution

    Paramters
    ---------
    tensor_format: int
        0: CUDNN_TENSOR_NCHW
        1: CUDNN_TENSOR_NHWC
        2: CUDNN_TENSOR_NCHW_VECT_C
    pad: int or list
        padding
    stride: int or list
        stride
    dilation: int or list
        dilation
    x_shape: list
        input shape
    w_shape: list
        weight shape
    data_dtype: str
        data type
    conv_dtype: str
        convolution type
    groups: int
        number of groups

    Returns
    -------
    oshape: list
        output shape
    """

    assert len(x_shape) == len(w_shape)
    assert len(x_shape) in (4, 5)

    if tensor_format == 0:
        n_output = x_shape[0]
        c_output = w_shape[0]
        x_chan = x_shape[1]
        w_chan_input = w_shape[1]
        x_shape = x_shape[2:]
        w_shape = w_shape[2:]

    elif tensor_format == 1:
        n_output = x_shape[0]
        c_output = w_shape[0]
        x_chan = x_shape[-1]
        w_chan_input = w_shape[-1]
        assert len(x_shape) == 4, "CuDNN layout NHWC is only well-defined for 4d tensors"
        x_shape = x_shape[1:-1]
        w_shape = w_shape[1:-1]

    elif tensor_format == 2:
        n_output = x_shape[0]
        c_output = w_shape[0]
        x_chan = x_shape[1]
        w_chan_input = w_shape[1]
        w_lanes = tvm.runtime.DataType(conv_dtype).lanes
        assert w_lanes == 1
        x_shape = x_shape[2:]
        w_shape = w_shape[2:]

    else:
        raise ValueError("Unknown CuDNN tensor format: '{}'".format(tensor_format))

    x_lanes = tvm.runtime.DataType(data_dtype).lanes
    assert x_chan * x_lanes == w_chan_input * groups, (
        "Mismatched dimensions, data has {} channels/group "
        "(dimension {} with {} lanes/value, {} groups), "
        "but weights require {} input channels/group"
    ).format(x_chan // groups, x_chan, x_lanes, groups, w_chan_input)

    output_dims = []
    for x_shape_i, w_shape_i, pad_i, stride_i, dilation_i in zip(
        x_shape, w_shape, pad, stride, dilation
    ):
        output_dim = 1 + (x_shape_i + 2 * pad_i - (((w_shape_i - 1) * dilation_i) + 1)) // stride_i
        output_dims.append(output_dim)

    if tensor_format in [0, 2]:
        output = [n_output, c_output, *output_dims]
    elif tensor_format == 1:
        output = [n_output, *output_dims, c_output]
    else:
        raise ValueError("Unknown CuDNN tensor format: '{}'".format(tensor_format))

    return output

def _conv_output_shape_from_cudnn(
    tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1
):
    """Get output shape of 2D or 3D convolution.  The output of this
    function should be identical to that of conv_output_shape, but
    requires a GPU with CuDNN to be present.  This is maintained for
    testing purposes to validate the output of conv_output_shape.

    Paramters
    ---------
    tensor_format: int
        0: CUDNN_TENSOR_NCHW
        1: CUDNN_TENSOR_NHWC
        2: CUDNN_TENSOR_NCHW_VECT_C
    pad: int or list
        padding
    stride: int or list
        stride
    dilation: int or list
        dilation
    x_shape: list
        input shape
    w_shape: list
        weight shape
    data_dtype: str
        data type
    conv_dtype: str
        convolution type
    groups: int
        number of groups

    Returns
    -------
    oshape: list
        output shape

    """
    dims = len(x_shape)
    assert dims in (4, 5)

    pad, stride, dilation, xshape, wshape = _prepare_global_func_params(
        dims - 2, pad, stride, dilation, x_shape, w_shape
    )
    oshape = np.zeros((dims), dtype=np.int32)

    func = tvm._ffi.get_global_func("tvm.contrib.cudnn.conv.output_shape_from_cudnn")
    func(
        tensor_format,
        dims - 2,
        _get_np_int32_array_handle(pad),
        _get_np_int32_array_handle(stride),
        _get_np_int32_array_handle(dilation),
        _get_np_int32_array_handle(xshape),
        _get_np_int32_array_handle(wshape),
        _get_np_int32_array_handle(oshape),
        data_dtype,
        conv_dtype,
        groups,
    )
    return list(oshape)

def conv_find_algo(
    tensor_format,
    pad,
    stride,
    dilation,
    x_shape,
    w_shape,
    y_shape,
    data_dtype,
    conv_dtype,
    groups=1,
):
    """Choose the best algo for the given input.

    Paramters
    ---------
    tensor_format: int
        0: CUDNN_TENSOR_NCHW
        1: CUDNN_TENSOR_NHWC
        2: CUDNN_TENSOR_NCHW_VECT_C
    pad: int or list
        padding
    stride: int or list
        stride
    dilation: int or list
        dilation
    x_shape: list
        input shape
    w_shape: list
        weight shape
    y_shape: list
        output shape
    data_dtype: str
        data type
    conv_dtype: str
        convolution type
    groups: int
        number of groups

    Returns
    -------
    algo: int
        algo chosen by CUDNN
    """
    dims = len(x_shape)
    assert dims in (4, 5)

    pad, stride, dilation, xshape, wshape = _prepare_global_func_params(
        dims - 2, pad, stride, dilation, x_shape, w_shape
    )
    yshape = np.array(y_shape, dtype=np.int32)
    func = tvm._ffi.get_global_func("tvm.contrib.cudnn.conv.find_algo")
    return func(
        tensor_format,
        dims - 2,
        _get_np_int32_array_handle(pad),
        _get_np_int32_array_handle(stride),
        _get_np_int32_array_handle(dilation),
        _get_np_int32_array_handle(xshape),
        _get_np_int32_array_handle(wshape),
        _get_np_int32_array_handle(yshape),
        data_dtype,
        conv_dtype,
        groups,
    )

def conv_forward(x, w, pad, stride, dilation, conv_mode, tensor_format, algo, conv_dtype, groups=1):
    """Create an extern op that compute 2D or 3D convolution with CuDNN

    Parameters
    ----------
    x: Tensor
        input feature map
    w: Tensor
        convolution weight
    pad: int or list
        padding
    stride: int or list
        stride
    dilation: int or list
        dilation
    conv_mode: int
        0: CUDNN_CONVOLUTION
        1: CUDNN_CROSS_CORRELATION
    tensor_format: int
        0: CUDNN_TENSOR_NCHW
        1: CUDNN_TENSOR_NHWC
        2: CUDNN_TENSOR_NCHW_VECT_C
    algo: int
        Forward algorithm, get index from ```algo_to_index``` function
        if algo == -1, the best algo will be chosen by CUDNN
    conv_dtype: str
        convolution type
    groups: int
        the number of groups

    Returns
    -------
    y: Tensor
        The result tensor
    """
    dims = len(x.shape)
    assert dims in (4, 5)

    conv_dtype = x.dtype if conv_dtype is None else conv_dtype
    pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation)

    x_shape = list(x.shape)

    if isinstance(x.shape[0], tvm.tir.expr.IntImm):
        oshape = conv_output_shape(
            tensor_format,
            pad,
            stride,
            dilation,
            x_shape,
            list(w.shape),
            x.dtype,
            conv_dtype,
            groups,
        )
        if algo == -1:
            # For now if we try to call `cudnnFindConvolutionForwardAlgorithm` when
            # using INT8 data type, CuDNN will crash down.
            # On the other hand, CuDNN only support IMPLICIT_PRECOMP_GEMM at NHWC format
            if tensor_format == 1 and conv_dtype == "int32":
                algo = 1
            else:
                algo = conv_find_algo(
                    tensor_format,
                    pad,
                    stride,
                    dilation,
                    list(x.shape),
                    list(w.shape),
                    oshape,
                    x.dtype,
                    conv_dtype,
                    groups,
                )
    else:
        # The dynamic batch size case, pretend this is a single batch
        x_shape[0] = 1
        oshape = conv_output_shape(
            tensor_format,
            pad,
            stride,
            dilation,
            x_shape,
            list(w.shape),
            x.dtype,
            conv_dtype,
            groups,
        )
        oshape[0] = x.shape[0]
        # This picks CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM
        # It seems this is the fastest among algorithms that are always applicable
        algo = 1

    if dims == 4:
        return te.extern(
            oshape,
            [x, w],
            lambda ins, outs: tvm.tir.call_packed(
                "tvm.contrib.cudnn.conv2d.forward",
                conv_mode,
                tensor_format,
                algo,
                pad[0],
                pad[1],
                stride[0],
                stride[1],
                dilation[0],
                dilation[1],
                ins[0],
                ins[1],
                outs[0],
                conv_dtype,
                groups,
            ),
            name="y",
        )

    return te.extern(
        oshape,
        [x, w],
        lambda ins, outs: tvm.tir.call_packed(
            "tvm.contrib.cudnn.conv3d.forward",
            conv_mode,
            tensor_format,
            algo,
            pad[0],
            pad[1],
            pad[2],
            stride[0],
            stride[1],
            stride[2],
            dilation[0],
            dilation[1],
            dilation[2],
            ins[0],
            ins[1],
            outs[0],
            conv_dtype,
            groups,
        ),
        name="y",
    )

def softmax(x, axis=-1):
    """Compute softmax using CuDNN

    Parameters
    ----------
    x : tvm.te.Tensor
        The input tensor

    axis : int
        The axis to compute the softmax

    Returns
    -------
    ret : tvm.te.Tensor
        The result tensor
    """
    return te.extern(
        x.shape,
        [x],
        lambda ins, outs: tvm.tir.call_packed(
            "tvm.contrib.cudnn.softmax.forward", ins[0], outs[0], axis
        ),
        name="y",
    )

def log_softmax(x, axis=-1):
    """Compute log_softmax using CuDNN

    Parameters
    ----------
    x : tvm.te.Tensor
        The input tensor

    axis : int
        The axis to compute log softmax over

    Returns
    -------
    ret : tvm.te.Tensor
        The result tensor
    """
    return te.extern(
        x.shape,
        [x],
        lambda ins, outs: tvm.tir.call_packed(
            "tvm.contrib.cudnn.log_softmax.forward", ins[0], outs[0], axis
        ),
        name="y",
    )


import logging
import os
import random
import tempfile
import time
import zipfile
from collections import deque
from contextlib import contextmanager
from typing import Generator

import cloudpickle
import numpy as np
import tensorflow as tf
from baselines.common import colorize, dataset, explained_variance, tf_util, zipsame
from baselines.common.cg import cg
from baselines.common.mpi_adam import MpiAdam
from cluster_work import ClusterWork
from mpi4py import MPI

from kb_learning.envs import MultiObjectTargetPoseEnv, NormalizeActionWrapper
from kb_learning.policy_networks.mlp_policy import MlpPolicyNetwork
from kb_learning.policy_networks.swarm_policy import SwarmPolicyNetwork
from kb_learning.tools.trpo_tools import ActWrapper, add_vtarg_and_adv, flatten_lists, traj_segment_generator

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["KMP_BLOCKTIME"] = "0"
os.environ["KMP_AFFINITY"] = "granularity=fine,compact,1,0"

logger = logging.getLogger('trpo')

class TRPOLearner(ClusterWork):
    _restore_supported = True
    _default_params = {
        'sampling': {
            'timesteps_per_batch': 2048,
            'done_after_steps':    1024,
            'num_objects':         None,
            'reward_function':     None,
            'swarm_reward':        True,
            'agent_reward':        True,
            'schedule': 'linear'
        },
        'trpo':     {
            'gamma':         0.99,
            'lambda':        0.98,
            'max_kl':        0.01,
            'entcoeff':      0.0,
            'cg_iterations': 10,
            'cg_damping':    1e-2,
            'vf_step_size':  1e-3,
            'vf_iterations': 5,
        },
        'policy':   {
            'type':             'swarm',
            'swarm_net_size':   (64,),
            'swarm_net_type':   'mean',
            'objects_net_size': (64,),
            'objects_net_type': 'mean',
            'extra_net_size':   (64,),
            'concat_net_size':  (64,),
            'feature_net_size': (100, 50, 25),
            'pd_bias_init':     None
        },
        'buffer_len': 500
    }

    def __init__(self):
        super().__init__()
        self.num_workers = None
        self.rank = None

        self.session = None
        self.graph = None
        self.env = None

        self.policy_fn = None

        self.timesteps_per_batch = None
        self.seg_gen = None

        self.loss_names = None

        self.pi = None

        self.vfadam = None

        self.assign_old_eq_new = None
        self.compute_losses = None
        self.compute_lossandgrad = None
        self.compute_fvp = None
        self.compute_vflossandgrad = None
        self.vf_loss = None

        self.get_flat = None
        self.set_from_flat = None

        self.sampling_schedule = None
        self.gamma = None
        self.lam = None
        self.max_kl = None
        self.cg_iters = None
        self.cg_damping = None

        self.vf_iters = None
        self.vf_stepsize = None

        self.episodes_so_far = 0
        self.timesteps_so_far = 0
        self.iters_so_far = 0
        self.tstart = time.time()
        self.ep_length_buffer = None  # rolling buffer for episode lengths
        self.ep_return_buffer = None  # rolling buffer for episode rewards

    @contextmanager
    def timed(self, msg):
        if self.rank == 0:
            logger.info(colorize(msg, color='magenta'))
            tstart = time.time()
            yield
            logger.info(colorize("done in %.3f seconds" % (time.time() - tstart), color='magenta'))
        else:
            yield

    def allmean(self, x):
        assert isinstance(x, np.ndarray)
        out = np.empty_like(x)
        self._COMM.Allreduce(x, out, op=MPI.SUM)
        out /= self.num_workers
        return out

    def reset(self, config: dict, rep: int):
        os.environ["OMP_NUM_THREADS"] = '1'

        sess_conf = tf.ConfigProto(allow_soft_placement=True,
                                   inter_op_parallelism_threads=1,
                                   intra_op_parallelism_threads=1)
        self.graph = tf.Graph()
        self.session = tf_util.make_session(config=sess_conf, make_default=True, graph=self.graph)

        np.random.seed(self._seed)
        tf.set_random_seed(self._seed)
        random.seed(self._seed)

        if not self._COMM:
            logger.warning('setting self._COMM to MPI.COMM_WORLD')
            self._COMM = MPI.COMM_WORLD

        # create env
        self.env = MultiObjectTargetPoseEnv(configuration=config['env_config'],
                                            done_after_steps=self._params['sampling']['done_after_steps'])

        wrapped_env = NormalizeActionWrapper(self.env)

        # create policy function
        def policy_fn(name, **_policy_params):
            if self._params['policy']['type'] == 'swarm':
                return SwarmPolicyNetwork(name=name, **_policy_params)
            elif self._params['policy']['type'] == 'mlp':
                return MlpPolicyNetwork(name=name, **_policy_params)
            else:
                raise NotImplementedError

        _pd_bias_init = self._params['policy']['pd_bias_init'] or np.array([.0] * self.env.action_space.shape[0] * 2)
        policy_params = dict(ob_space=self.env.observation_space,
                             ac_space=self.env.action_space,  # only the type and shape of the action space is important
                             # we observe all other agents with (r, sin(a), cos(a), sin(th), cos(th), lin_vel, rot_vel)
                             agent_dims=self.env.kilobots_observation_space.shape[0] // self.env.num_kilobots,
                             num_agent_observations=self.env.num_kilobots,
                             # we observe all objects with (r, sin(a), cos(a), sin(th), cos(th), valid_indicator)
                             object_dims=self.env.object_observation_space.shape[0],
                             num_object_observations=len(config['env_config'].objects),
                             swarm_net_size=self._params['policy']['swarm_net_size'],
                             swarm_net_type=self._params['policy']['swarm_net_type'],
                             objects_net_size=self._params['policy']['objects_net_size'],
                             objects_net_type=self._params['policy']['objects_net_type'],
                             extra_net_size=self._params['policy']['extra_net_size'],
                             concat_net_size=self._params['policy']['concat_net_size'],
                             feature_net_size=self._params['policy']['feature_net_size'],
                             pd_bias_init=_pd_bias_init)

        self.policy_fn = policy_fn

        self.num_workers = self._COMM.Get_size()
        self.rank = self._COMM.Get_rank()
        logger.debug('initializing with rank {} of {}'.format(self.rank, self.num_workers))
        np.set_printoptions(precision=3)
        # Setup losses and stuff
        # ----------------------------------------
        pi = policy_fn("pi", **policy_params)
        oldpi = policy_fn("oldpi", **policy_params)
        atarg = tf.placeholder(dtype=tf.float32, shape=[None])  # Target advantage function (if applicable)
        ret = tf.placeholder(dtype=tf.float32, shape=[None])  # Empirical return

        ob = tf_util.get_placeholder_cached(name="ob")
        ac = pi.pdtype.sample_placeholder([None])

        kl_old_new = oldpi.pd.kl(pi.pd)
        entropy = pi.pd.entropy()
        mean_kl = tf.reduce_mean(kl_old_new)
        mean_entropy = tf.reduce_mean(entropy)
        entropy_bonus = self._params['trpo']['entcoeff'] * mean_entropy

        vf_error = tf.reduce_mean(tf.square(pi.vpred - ret))

        log_likelihood_ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac))  # advantage * pnew / pold
        surrogate_gain = tf.reduce_mean(log_likelihood_ratio * atarg)

        optimgain = surrogate_gain + entropy_bonus
        losses = [optimgain, mean_kl, entropy_bonus, surrogate_gain, mean_entropy]
        self.loss_names = ["optimizer_gain", "mean_kl", "entropy_loss", "surrogate_gain", "mean_entropy"]

        dist = mean_kl

        all_var_list = pi.get_trainable_variables()
        var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
        vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]

        self.vfadam = MpiAdam(vf_var_list, comm=self._COMM)

        self.get_flat = tf_util.GetFlat(var_list)
        self.set_from_flat = tf_util.SetFromFlat(var_list)
        kl_gradients = tf.gradients(dist, var_list)
        flat_tangent = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan")
        shapes = [var.get_shape().as_list() for var in var_list]
        start = 0
        tangents = []
        for shape in shapes:
            sz = tf_util.intprod(shape)
            tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
            start += sz
        gvp = tf.add_n([tf.reduce_sum(g * tangent) for (g, tangent) in zipsame(kl_gradients, tangents)])
        fvp = tf_util.flatgrad(gvp, var_list)

        self.assign_old_eq_new = tf_util.function([], [], updates=[tf.assign(oldv, newv)
                                                                   for (oldv, newv) in
                                                                   zipsame(oldpi.get_variables(), pi.get_variables())])

        self.compute_losses = tf_util.function([ob, ac, atarg], losses)
        self.compute_lossandgrad = tf_util.function([ob, ac, atarg], losses + [tf_util.flatgrad(optimgain, var_list)])
        self.compute_fvp = tf_util.function([flat_tangent, ob, ac, atarg], fvp)
        self.compute_vflossandgrad = tf_util.function([ob, ret], tf_util.flatgrad(vf_error, vf_var_list))
        self.vf_loss = tf_util.function([ob, ret], vf_error)

        self.pi = ActWrapper(pi, policy_params)

        tf_util.initialize()
        th_init = self.get_flat()
        self._COMM.Bcast(th_init, root=0)
        self.set_from_flat(th_init)
        self.vfadam.sync()
        logger.info("Init param sum {}".format(th_init.sum()))

        self.seg_gen = traj_segment_generator(pi, wrapped_env, self._params['sampling']['timesteps_per_batch'],
                                              stochastic=True)

        self.gamma = self._params['trpo']['gamma']
        self.lam = self._params['trpo']['lambda']
        self.max_kl = self._params['trpo']['max_kl']
        self.sampling_schedule = self._params['sampling']['schedule']
        self.vf_iters = self._params['trpo']['vf_iterations']
        self.vf_stepsize = self._params['trpo']['vf_step_size']

        self.cg_iters = self._params['trpo']['cg_iterations']
        self.cg_damping = self._params['trpo']['cg_damping']

        self.ep_length_buffer = deque(maxlen=self._params['buffer_len'])
        self.ep_return_buffer = deque(maxlen=self._params['buffer_len'])

    def _get_schedule_factor(self, schedule):
        if schedule == 'constant':
            factor = 1.0
        elif schedule == 'linear':
            factor = max(1.0 - float(self._it) / self._iterations, 0)
        elif schedule == 'sqrt':
            factor = np.sqrt(1.0 - float(self._it) / self._iterations)
        else:
            raise NotImplementedError
        return factor

    def iterate(self, config: dict, rep: int, n: int):
        # set random seed for repetition and iteration
        np.random.seed(self._seed)
        tf.set_random_seed(self._seed)
        random.seed(self._seed)

        time_update_start = time.time()

        theta_new = None
        mean_losses = None

        curr_sampling_factor = self._get_schedule_factor(self.sampling_schedule)
        self.env.progress_factor = curr_sampling_factor
        with self.timed("sampling"):
            seg = self.seg_gen.__next__()

        add_vtarg_and_adv(seg, self.gamma, self.lam)

        # ob, ac, a_target, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
        ob, ac, a_target, td_lambda_return = seg["ob"], seg["ac"], seg["adv"], seg["tdlamret"]
        v_prediction_before = seg["vpred"]  # predicted value function before udpate
        # standardized advantage function estimate
        a_target_mean = a_target.mean()
        a_target_mean = self._COMM.allreduce(a_target_mean, MPI.SUM) / self.num_workers
        a_target_std = (a_target - a_target_mean).std()
        a_target_std = self._COMM.allreduce(a_target_std, MPI.SUM) / self.num_workers
        a_target = (a_target - a_target_mean) / a_target_std

        # if hasattr(self.pi, "ret_rms"): self.pi.ret_rms.update(td_lambda_return)
        # if hasattr(self.pi, "ob_rms"): self.pi.ob_rms.update(ob)  # update running mean/std for policy

        # subsample data to take only 20% for computing the Fisher-vector products
        args = seg["ob"], seg["ac"], a_target
        fvp_arguments = [arr[::5] for arr in args]

        def fisher_vector_product(p):
            return self.allmean(self.compute_fvp(p, *fvp_arguments)) + self.cg_damping * p

        self.assign_old_eq_new()  # set old parameter values to new parameter values
        with self.timed("computegrad"):
            *loss_before, g = self.compute_lossandgrad(*args)
        loss_before = self.allmean(np.array(loss_before))
        g = self.allmean(g)
        if np.allclose(g, 0):
            logger.info("Got zero gradient. not updating")
        else:
            if g.dot(g) > 1e20:
                logger.warning("Gradient norm very large")
                g = g * 1e-4
            with self.timed("cg"):
                step_direction = cg(fisher_vector_product, g, cg_iters=self.cg_iters)
            assert np.isfinite(step_direction).all()
            shs = .5 * step_direction.dot(fisher_vector_product(step_direction))
            lm = np.sqrt(shs / self.max_kl)
            # logger.info("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
            full_step = step_direction / lm
            expected_improvement = g.dot(full_step)
            surrogate_loss_before = loss_before[0]
            step_size = 1.0
            theta_before = self.get_flat()
            # line search with step_direction from gradient to ensure loss improves and KL constraint is not violated
            for _ in range(10):
                theta_new = theta_before + full_step * step_size
                self.set_from_flat(theta_new)
                mean_losses = surrogate_loss, kl_loss, *_ = self.allmean(np.array(self.compute_losses(*args)))
                improve = surrogate_loss - surrogate_loss_before
                logger.info("Expected: {}".format(expected_improvement))
                logger.info("Actual: {}".format(improve))
                if not np.isfinite(mean_losses).all():
                    logger.info("Got non-finite value of losses -- bad!")
                elif kl_loss > self.max_kl * 1.5:
                    logger.info("violated KL constraint. shrinking step.")
                elif improve < 0:
                    logger.info("surrogate didn't improve. shrinking step.")
                else:
                    logger.info("Stepsize OK!")
                    break
                step_size *= .5
            else:
                logger.info("couldn't compute a good step")
                self.set_from_flat(theta_before)
            if self.num_workers > 1 and self.iters_so_far % 20 == 0:
                # list of tuples
                parameter_sums = self._COMM.allgather((theta_new.sum(), self.vfadam.getflat().sum()))
                assert all(np.allclose(ps, parameter_sums[0]) for ps in parameter_sums[1:])

        with self.timed("vf"):
            for _ in range(self.vf_iters):
                for (mbob, mbret) in dataset.iterbatches((seg["ob"], seg["tdlamret"]),
                                                         include_final_partial_batch=False, batch_size=64):
                    g = self.allmean(self.compute_vflossandgrad(mbob, mbret))
                    self.vfadam.update(g, self.vf_stepsize)

        vf_loss = self.vf_loss(ob, td_lambda_return)

        vf_loss_mean = np.mean(self._COMM.allgather(vf_loss))
        len_ret_local = (seg["ep_lens"], seg["ep_rets"])  # local values
        len_ret_pairs = self._COMM.allgather(len_ret_local)  # list of tuples
        ep_lengths, ep_returns = map(flatten_lists, zip(*len_ret_pairs))
        self.ep_length_buffer.extend(ep_lengths)
        self.ep_return_buffer.extend(ep_returns)

        self.episodes_so_far += len(ep_lengths)
        self.timesteps_so_far += sum(ep_lengths)
        self.iters_so_far += 1

        if self._COMM.rank == 0:
            time_elapsed = time.time() - time_update_start

            return_values = dict(zip(self.loss_names, mean_losses))

            return_values.update({
                "ev_tdlam_before":         explained_variance(v_prediction_before, td_lambda_return),
                "vf_loss":                 vf_loss_mean,
                "episode_return_mean":     np.mean(self.ep_return_buffer),
                "episode_return_min":      np.min(self.ep_return_buffer),
                "episode_return_max":      np.max(self.ep_return_buffer),
                "episode_length_mean":     np.mean(self.ep_length_buffer),
                "episodes_this_iteration": len(ep_lengths),
                "episodes":                self.episodes_so_far,
                "timesteps":               self.timesteps_so_far,
                "frames_per_second":       int(np.sum(ep_lengths) / time_elapsed)
            })

            important_values = ["episode_return_mean", "ev_tdlam_before"]

            for k, v in return_values.items():
                if k in important_values:
                    logger.info(colorize("{}: {}".format(k, v), color='green', bold=True))
                else:
                    logger.info("{}: {}".format(k, v))

            return return_values

    def finalize(self):
        if self.env:
            self.env.close()
        self.session.close()
        del self.graph

    def save_state(self, config: dict, rep: int, n: int):
        # save parameters at every iteration
        policy_path = os.path.join(self._log_path_it, 'policy.pkl')
        self.pi.save(policy_path)

    def restore_state(self, config: dict, rep: int, n: int):
        # resore parameters
        policy_path = os.path.join(self._log_path_it, 'policy.pkl')
        logger.info('Restoring parameters from {}'.format(policy_path))
        # self.pi = ActWrapper.load(policy_path, self.policy_fn)

        with open(policy_path, "rb") as f:
            model_data, act_params = cloudpickle.load(f)

        with tempfile.TemporaryDirectory() as td:
            arc_path = os.path.join(td, "packed.zip")
            with open(arc_path, "wb") as f:
                f.write(model_data)

            zipfile.ZipFile(arc_path, 'r', zipfile.ZIP_DEFLATED).extractall(td)
            self.pi.load_state(os.path.join(td, "model"))

        return True

    @classmethod
    def plot_results(cls, configs_results: Generator):
        pass


import numpy as np
import configparser
import json
from datetime import datetime
from datetime import timedelta
import os
import logging
import pyvisa
from pydlcp import arduino_board, hotplate, errors, keithley, impedance_analyzer as ia, datastorage, bts
from apscheduler.schedulers.background import BackgroundScheduler
import platform
from typing import List

# Different data type definitions
ard_list = List[arduino_board.ArduinoBoard]
bts_list = List[bts.BTS]
hp_list = List[hotplate.Hotplate]
vcr_type = np.dtype([('V', 'd'), ('C', 'd'), ('R', 'd')])
dlcp_type = np.dtype([('osc_level', 'd'),
                      ('bias', 'd'),
                      ('nominal_bias', 'd'),
                      ('V', 'd'),
                      ('C', 'd'),
                      ('R', 'd')])

class BTSController:
    """
    This class provides methods to control a BTS and DLCP experiment interacting with different instruments and saving
    to a h5 data store.

    Attributes
    ----------
    _activeCVUnit: List[str]
        The unit that the impedance analyzer is currently locked to.
    _arduinos:
        A list with handles to different ArduinoBoard instances controlling the different test units.
    _availableResources: List[str]
        A list of available pyvisa resources. Used to verify that the resource given by the constructor is available.
    _btsAcquisitions: bts_list
        A list with the different BTS instances representing a different experimental BTS condition.
    _cvSweepParams: dict
        A dictionary with the CV sweep acquisition parameters
    _data_paths: List[str]
        A list containing the paths where the output data a logs will be saved
    debug: bool
        True if we want to execute in debugging mode
    _configMeasurementRequiredOptions: dict
        A dictionary containing rules for validating the acquisition parameters and settings
    _configSystemRequiredOptions: dict
        A dictionary containing rules for validating the system configuration if it is not in its default state
    _dlcpParams: dict
        A dictionary with the DLCP acquisition parameters
    _finishedUnits: List[int]
        A list with the number ids of the units that completed the BTS measurement
    _hotPlates: hp_list
        A list containing instances of the hotplates that control the temperature on each test unit
    _impedanceAnalyzer: ia.ImpedanceAnalyzer
        An instance to the impedance analyzer object
    _keithley: keithley.Keithley
        An instance to the keithley source-meter object
    _loggingLevels: dict
        A map of logging level strings to integers, as defined in the logging module
    _mainLogger: logging.Logger
        The logger to the class, used to handle any messages within the class.
    _measurementConfig: config.ConfigParser
        An instance of ConfigParser that contains all the acquisition parameters
    _physicalTestUnits: int
        The number of hardware units available for use.
    _resourceManager:pyvisa.highlevel.ResourceManager
        An instance of pyvisa's resource manager to instantiate the instruments from.
    _scheduler: BackgroundScheduler
        An instance of the BackgroundScheduler
    _schedulerRunning: bool
        True if the scheduler is running, false otherwise
    _testUnits: int
        The number of test units configured to use in the BTS measurement

    Methods
    -------
    ramp_up(self, test_unit: int):
        Set's the hotplate temperature to the stress temperature. And sets the BTS flag to 'heating_up'. If the unit's
        fan is currently on, turns it off.

    ramp_down(self, test_unit: int):
        Set the hotplate temperature to 25 °C and turns the unit's fan on if it is off. Sets the BTS flag to
        'cooling_down'. Disconnects all the pins in the test unit.

    start_temperature_log(self, test_unit: int):
        Starts the temperature log. Instructs the scheduler to call the method '_log_temperature' according to the
        configured setting. Starts the class scheduler if it is not already running. Sets the '_schedulerRunning' to
        True.

    start_bts(self, test_unit: int):
        Starts the bias-temperature stress. This method should be called once the test unit has reached the stress
        temperature. It turn's on the voltage (connects all pins on the unit to the voltage source and turns the source
        on if it is not on already). Set the BTS flag to 'running_stress'. Instructs the scheduler to call the method
        'stop_bts' at time 1 stress_interval unit time later than the current time.

    stop_bts(self, test_unit: int):
        Calls 'ramp_down' method.

    _log_temperature(self, test_unit: int):
        Measures the temperature and leakage current through all the devices connected in parallel to the keithley
        source meter and saves the log to each of the device's datastores.
    """
    _activeCVUnit = []
    _arduinos: ard_list = []
    _btsAcquisitions: bts_list = []
    _cvSweepParams: dict = None
    _data_paths: List[str] = []
    _dlcpParams: dict = None
    _finishedUnits: List[int] = []
    _hotPlates: hp_list = []
    _impedanceAnalyzer: ia.ImpedanceAnalyzer = None
    _keithley: keithley.Keithley = None
    _loggingLevels = {'NOTSET': logging.NOTSET,
                      'DEBUG': logging.DEBUG,
                      'INFO': logging.INFO,
                      'WARNING': logging.WARNING,
                      'ERROR': logging.ERROR,
                      'CRITICAL': logging.CRITICAL}
    _mainLogger: logging.Logger = None
    _measurementConfig: configparser.ConfigParser = None
    _schedulerRunning = False
    _testUnits: int = 0

    def __init__(self, config_file_url: str, **kwargs):
        if not isinstance(config_file_url, str):
            raise TypeError('The first argument should be an instance of str.')
        self.debug: bool = kwargs.get('debug', False)
        system_option_requirements_json = kwargs.get('system_option_requirements_json',
                                                     'system_config_required_options.json')
        measurement_option_requirements_json = kwargs.get('measurement_options_requirements_json',
                                                          'measurement_config_required_options.json')

        # Load validation rules for the system configuration file
        self._configSystemRequiredOptions = self._read_json_file(system_option_requirements_json)
        # Load validation rules for the measurement configuration file
        self._configMeasurementRequiredOptions = self._read_json_file(measurement_option_requirements_json)

        # Load the system configuration file
        config = configparser.ConfigParser()
        config.read(config_file_url)

        self._physicalTestUnits = config.getint(section='global', option='test_units')

        # If the system configuration file is valid, then store it in the object
        if self._validate_config(config, self._configSystemRequiredOptions):
            self._systemConfig = config

        self._resourceManager: pyvisa.highlevel.ResourceManager = pyvisa.highlevel.ResourceManager()
        self._availableResources = self._resourceManager.list_resources()
        self._scheduler: BackgroundScheduler = BackgroundScheduler()

    def ramp_up(self, test_unit: int):
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        hp: hotplate.Hotplate = self._hotPlates[test_unit]
        hp.set_temperature(int(bts_acquisition.temperature))
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        if a.fan_status:
            a.fan_off()
        a.disconnect_all_pins()
        bts_acquisition.status = 'heating_up'

    def ramp_down(self, test_unit: int):
        hp: hotplate.Hotplate = self._hotPlates[test_unit]
        hp.set_temperature(25)
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        if not a.fan_status:
            a.fan_on()
        a.disconnect_all_pins()
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        bts_acquisition.status = 'cooling_down'

    def start_temperature_log(self, test_unit: int):
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        self._scheduler.add_job(func=self._log_temperature, trigger='interval', args=[test_unit],
                                seconds=bts_acquisition.temperature_sampling_interval,
                                id='temperature_log_unit{0}'.format(test_unit))
        if not self._schedulerRunning:
            self._scheduler.start()
            self._schedulerRunning = True

    def start_bts(self, test_unit: int):
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        a.connect_all_pins()
        a.connect_keithley()
        # If the voltage source is off, turn it on
        if not self._keithley.source_on:
            self._keithley.turn_source_on()
        now: datetime = datetime.now()
        later = now + timedelta(seconds=bts_acquisition.stress_interval)
        self._scheduler.add_job(func=self.stop_bts, trigger='date', args=[test_unit],
                                date=later, id='bts_unit{0}'.format(test_unit))
        bts_acquisition.status = 'running_stress'
        bts_acquisition.start_bts_interval = now
        if not self._schedulerRunning:
            self._scheduler.start()
            self._schedulerRunning = True

    def stop_bts(self, test_unit: int):
        # Maybe we need to do something else here, else change all method calls to just, ramp_down.
        self.ramp_down(test_unit=test_unit)

    def _log_temperature(self, test_unit: int):
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        now = datetime.now()
        temperature = a.temperature
        current = self._keithley.current
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        clean_devices = bts_acquisition.clean_devices
        contaminated_devices = bts_acquisition.contaminated_devices
        dt = bts_acquisition.time_delta_total(current_datetime=now)
        for device in clean_devices:
            h5_storage: datastorage.BTSH5Store = bts_acquisition.get_device_storage(device=device, clean=True)
            h5_storage.log_temperature_current(time=dt,
                                               temperature=temperature,
                                               current=current)

        for device in contaminated_devices:
            h5_storage: datastorage.BTSH5Store = bts_acquisition.get_device_storage(device=device, clean=False)
            h5_storage.log_temperature_current(time=dt,
                                               temperature=temperature,
                                               current=current)

    def start_measurement(self):
        now = datetime.now()
        for test_unit in range(self._testUnits):
            self._log_cv_sweep(test_unit=test_unit)
            status_job_id = 'check_status{0}'.format(test_unit)
            self._scheduler.add_job(func=self._check_status, trigger='interval', args=[test_unit], id=status_job_id,
                                    seconds=10)
            self.start_temperature_log(test_unit=test_unit)

        if not self._schedulerRunning:
            self._scheduler.start()
            self._schedulerRunning = True

    def _check_status(self, test_unit: int):
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        now = datetime.now()
        temperature = a.temperature
        # If the device temperature equals the target temperature and the status is ramping up, start stress
        if np.isclose(a.temperature, bts_acquisition.temperature,
                      atol=2) and bts_acquisition.status == 'heating_up':
            self.start_bts(test_unit=test_unit)
        elif a.temperature <= 26 and bts_acquisition.status == 'cooling_down':
            dt = bts_acquisition.time_delta_bts(current_datetime=datetime.now())
            bts_acquisition.accumulate_interval(dt=dt)
            bts_acquisition.status = "idle"
            self._log_cv_sweep(test_unit=test_unit)
        elif bts_acquisition.status == 'finished':
            self._unit_finished(test_unit=test_unit)

    def _unit_finished(self, test_unit: int):
        a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
        if (a.temperature > 30) and (not a.fan_status):
            a.fan_on()
        elif a.fan_status:
            a.fan_off()

        # Stop logging temperature
        self._scheduler.remove_job(job_id='temperature_log_unit{0}'.format(test_unit))
        status_job_id = 'check_status{0}'.format(test_unit)
        # Stop checking if the unit is ready
        self._scheduler.remove_job(job_id=status_job_id)
        if test_unit not in self._finishedUnits:
            self._finishedUnits.append(test_unit)
        if len(self._finishedUnits) == self._testUnits:
            self._scheduler.shutdown()
            self._schedulerRunning = False

    def _log_cv_sweep(self, test_unit: int):
        # Make sure the board is idle
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        # If the current callback is in a scheduler queue, remove the job from the scheduler
        job_id = 'bts_cv_test_unit{0}'.format(test_unit)
        if job_id in self._scheduler.get_jobs():
            self._scheduler.remove_job(job_id=job_id)
        # Check if the measurement is not finished
        if bts_acquisition.status == "idle":
            # If the impedance analyzer is busy, schedule this measurement for later
            if len(self._activeCVUnit) > 0 or self._impedanceAnalyzer.status == "running":
                next_date = datetime.now() + timedelta(seconds=self._impedanceAnalyzer.wait_time)
                self._scheduler.add_job(self._log_cv_sweep(test_unit=test_unit),
                                        trigger='date', run_date=next_date,
                                        args=[test_unit], id=job_id)
            else:
                bts_acquisition.status = "running_cv"
                # Lock the Impedance Analyzer measurement to the test_unit
                self._activeCVUnit.append(test_unit)
                # Get the arduino for the selected board
                a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
                if a.keithley_connected:
                    a.disconnect_keithley()
                if a.fan_status:  # If fan is on
                    a.fan_off()
                # Make sure no other pins are connected to the impedance analyzer
                for i in range(self._testUnits):
                    ai: arduino_board.ArduinoBoard = self._arduinos[i]
                    b: bts.BTS = self._btsAcquisitions[i]
                    if b.status != 'running_stress':
                        ai.disconnect_all_pins()
                # loop over all clean pins
                for d, p in zip(bts_acquisition.clean_devices, bts_acquisition.clean_pins):
                    # turn the pin on
                    a.pin_on(pin_number=p)
                    # collect the data from the impedance analyzer
                    data: vcr_type = self._impedanceAnalyzer.cv_sweep(
                        voltage_start=float(self._cvSweepParams['voltage_start']),
                        voltage_step=float(self._cvSweepParams['voltage_step']),
                        voltage_stop=float(self._cvSweepParams['voltage_stop']),
                        frequency=float(self._cvSweepParams['frequency']),
                        integration_time=self._cvSweepParams['integration_time'],
                        noa=int(self._cvSweepParams['number_of_averages']),
                        osc_amplitude=float(self._cvSweepParams['osc_amplitude']),
                        sweep_direction=self._cvSweepParams['sweep_direction']
                    )
                    # turn the pin off
                    a.pin_off(pin_number=p)
                    # get the storage for the pin
                    h5_ds = bts_acquisition.get_device_storage(device=d, clean=True)
                    # append the data to the h5 storage
                    h5_ds.append_cv(time=bts_acquisition.accumulated_stress_time, cv_data=data)
                # loop over al contaminated pins
                for d, p in zip(bts_acquisition.contaminated_devices, bts_acquisition.contaminated_pins):
                    a.pin_on(pin_number=p)
                    # collect the data from the impedance analyzer
                    data: vcr_type = self._impedanceAnalyzer.cv_sweep(
                        voltage_start=float(self._cvSweepParams['voltage_start']),
                        voltage_step=float(self._cvSweepParams['voltage_step']),
                        voltage_stop=float(self._cvSweepParams['voltage_stop']),
                        frequency=float(self._cvSweepParams['frequency']),
                        integration_time=self._cvSweepParams['integration_time'],
                        noa=int(self._cvSweepParams['number_of_averages']),
                        osc_amplitude=float(self._cvSweepParams['osc_amplitude']),
                        sweep_direction=self._cvSweepParams['sweep_direction']
                    )
                    # turn the pin off
                    a.pin_off(pin_number=p)
                    # get the storage for the pin
                    h5_ds = bts_acquisition.get_device_storage(device=d, clean=False)
                    # append the data to the h5 storage
                    h5_ds.append_cv(time=bts_acquisition.accumulated_stress_time, cv_data=data)

                # Unlock the impedance analyzer
                self._activeCVUnit.remove(test_unit)
                # Run a DLCP
                self._log_dlcp(test_unit=test_unit)
                if bts_acquisition.status != 'finished':
                    self.ramp_up(test_unit=test_unit)
                else:
                    self._unit_finished(test_unit=test_unit)

    def _log_dlcp(self, test_unit: int):
        # Make sure the board is idle
        bts_acquisition: bts.BTS = self._btsAcquisitions[test_unit]
        # If exists remove the job from the scheduler
        scheduled_jobs = self._scheduler.get_jobs()
        job_id = 'bts_dlcp_test_unit{0}'.format(test_unit)
        if job_id in scheduled_jobs:
            self._scheduler.remove_job(job_id=job_id)
        # Check if the measurement is not finished
        if bts_acquisition.status == "idle":
            # Check if there is something else being measured
            if len(self._activeCVUnit) > 0 or self._impedanceAnalyzer.status == "running":
                next_date = datetime.now() + timedelta(seconds=self._impedanceAnalyzer.wait_time)
                self._scheduler.add_job(self._log_dlcp(test_unit=test_unit),
                                        trigger='date', run_date=next_date,
                                        args=[test_unit], id=job_id)
            else:
                bts_acquisition.status = "running_dlcp"
                # Lock the Impedance Analyzer measurement
                self._activeCVUnit.append(test_unit)
                # Get the arduino for the selected board
                a: arduino_board.ArduinoBoard = self._arduinos[test_unit]
                # Make sure no other pins are connected to the impedance analyzer
                for i in range(self._testUnits):
                    ai: arduino_board.ArduinoBoard = self._arduinos[i]
                    b: bts.BTS = self._btsAcquisitions[i]
                    if b.status != 'running_stress':
                        ai.disconnect_all_pins()
                # Lock the impedance analyzer to this unit
                self._activeCVUnit.remove(test_unit)
                # loop over all clean pins
                for d, p in zip(bts_acquisition.clean_devices, bts_acquisition.clean_pins):
                    # turn the pin on
                    a.pin_on(pin_number=p)
                    # collect the data from the impedance analyzer
                    data: dlcp_type = self._impedanceAnalyzer.dlcp_sweep(
                        nominal_bias=float(self._dlcpParams['nominal_bias']),
                        osc_start=float(self._dlcpParams['start_amplitude']),
                        osc_step=float(self._dlcpParams['step_amplitude']),
                        osc_stop=float(self._dlcpParams['stop_amplitude']),
                        frequency=float(self._cvSweepParams['frequency']),
                        integration_time=self._cvSweepParams['integration_time'],
                        noa=int(self._cvSweepParams['number_of_averages']),
                    )
                    # turn the pin off
                    a.pin_off(pin_number=p)
                    # get the storage for the pin
                    h5_ds = bts_acquisition.get_device_storage(device=d, clean=True)
                    # append the data to the h5 storage
                    h5_ds.append_dlcp(time=bts_acquisition.accumulated_stress_time, dlcp_data=data)
                # loop over al contaminated pins
                for d, p in zip(bts_acquisition.contaminated_devices, bts_acquisition.contaminated_pins):
                    a.pin_on(pin_number=p)
                    # collect the data from the impedance analyzer
                    data: dlcp_type = self._impedanceAnalyzer.dlcp_sweep(
                        nominal_bias=float(self._dlcpParams['nominal_bias']),
                        osc_start=float(self._dlcpParams['start_amplitude']),
                        osc_step=float(self._dlcpParams['step_amplitude']),
                        osc_stop=float(self._dlcpParams['stop_amplitude']),
                        frequency=float(self._cvSweepParams['frequency']),
                        integration_time=self._cvSweepParams['integration_time'],
                        noa=int(self._cvSweepParams['number_of_averages']),
                    )
                    # turn the pin off
                    a.pin_off(pin_number=p)
                    # get the storage for the pin
                    h5_ds = bts_acquisition.get_device_storage(device=d, clean=False)
                    # append the data to the h5 storage
                    h5_ds.append_dlcp(time=bts_acquisition.accumulated_stress_time, dlcp_data=data)
                bts_acquisition.status = "idle"

    def load_test_config(self, config: configparser.ConfigParser):
        if not isinstance(config, configparser.ConfigParser):
            raise TypeError('The argument should be an instance of configparser.ConfigParser.')
        if self.debug:
            self._print('Loading measurement configuration...')  # No logger yet...

        if self._validate_config(config, self._configMeasurementRequiredOptions):
            all_sections = config.sections()
            test_units = len([u for u in all_sections if 'bts' in u])
            if test_units > self._physicalTestUnits:
                msg = 'Measurement requesting {0:d} units. Available units: {1}'.format(test_units,
                                                                                self._physicalTestUnits)
                raise errors.ConfigurationError(message=msg)
            self._testUnits = test_units
            self._measurementConfig = config
            now = datetime.now()
            time_stamp = now.strftime('%Y%m%d')

            self._cvSweepParams = dict(config.items('cv_sweep'))
            self._dlcpParams = dict(config.items('dlcp'))

            # Create a directory structure by
            # +-- Temperature + Applied Electric Field
            #       +-- Device 1 - Clean
            #       +-- Device 2 - Clean
            #       | ...
            #       +-- Device 7 - Contaminated
            #       +-- Device 8 - Contaminated
            root_path: str = config.get(section='general', option='base_path')
            if platform.system() == 'Windows':
                root_path = '\\\\?\\' + root_path
            base_path = os.path.join(root_path, time_stamp)
            if self.debug:
                self._print('Creating base path at {0}'.format(base_path))  # No logger yet...
            self._create_path(base_path)
            # Create main logger
            self._mainLogger = self._create_logger(base_path, name='Main Logger', level='CRITICAL', console=True)

            for i in range(self._testUnits):
                section_name = 'bts{0}'.format(i)
                params = self._get_bts_metadata(section_name=section_name)
                # Create a BTS acquisition object
                bts_acquisition = bts.BTS(temperature=params['stress_temperature'], bias=params['stress_bias'],
                                          stress_interval=params['stress_interval'],
                                          temperature_sampling_interval=params['temperature_sampling_interval'],
                                          max_time=params['max_stress_time'])

                data_folder = '{0:.0f}C-{1:.1f}'.format(params['stress_temperature'], params['stress_bias'])
                relative_path = os.path.join(base_path, data_folder)
                self._data_paths.append(relative_path)
                self._create_path(relative_path)
                # Create separate h5 files for each pin
                clean_devices = config.get(section=section_name, option='clean_devices').split(',')
                clean_pins = config.get(section=section_name, option='clean_pins').split(',')
                contaminated_devices = config.get(section=section_name, option='contaminated_devices').split(',')
                contaminated_pins = config.get(section=section_name, option='contaminated_pins').split(',')
                for d, p in zip(clean_devices, clean_pins):
                    h5_name = os.path.join(relative_path, '{0}_{1}-clean.h5'.format(params['clean_wafer_id'], d))
                    ds = datastorage.BTSH5Store(h5_name)
                    metadata = params
                    metadata['device_id'] = d
                    ds.metadata(metadata=metadata, group="/bts")
                    ds.metadata(metadata=self._cvSweepParams, group='/bts')
                    bts_acquisition.append_device(device=d, pin=p, h5_storage=ds, clean=True)

                for d, p in zip(contaminated_devices, contaminated_pins):
                    h5_name = os.path.join(relative_path,
                                           '{0}_{1}-contaminated.h5'.format(params['contaminated_wafer_id'], d))
                    ds = datastorage.BTSH5Store(h5_name)
                    metadata = params
                    metadata['device_id'] = d
                    ds.metadata(metadata=metadata, group="/bts")
                    bts_acquisition.append_device(device=d, pin=p, h5_storage=ds, clean=False)

                self._btsAcquisitions.append(bts_acquisition)
            self._print('Loaded acquisition parameters successfully.', level='INFO')

    def _get_bts_metadata(self, section_name: str):
        if not isinstance(section_name, str):
            raise TypeError('The argument should be an instance of str.')
        config: configparser.ConfigParser = self._measurementConfig
        options = dict(config.items(section=section_name))
        options['stress_voltage'] = float(options['stress_voltage'])
        options['stress_temperature'] = float(options['stress_temperature'])
        options['max_stress_time'] = float(options['max_stress_time'])
        options['stress_interval'] = float(options['stress_interval'])
        options['temperature_sampling_interval'] = float(options['temperature_sampling_interval'])
        options['thickness'] = float(options['thickness'])

        return options

    def _validate_config(self, config: configparser.ConfigParser, required_options) -> bool:
        if not isinstance(config, configparser.ConfigParser):
            raise TypeError('The configuration argument must be an instance of configparser.ConfigParser.')
        for section in required_options:
            if section == 'arduino' or section == 'hotplate' or section == "bts":
                n = self._testUnits + 1
                section_units = ['{0}{1:d}'.format(section, i) for i in range(1, n)]
                for s in section_units:
                    if config.has_section(s):
                        for o in required_options[section]:
                            if not config.has_option(s, o):
                                msg = 'Config file must have option \'{0}\' for section \'{1}\''.format(o, s)
                                raise errors.BTSSystemConfigError(message=msg, test_units=self._testUnits)
                    else:
                        msg = "Config file must have section '{0}'.".format(s)
                        raise errors.BTSSystemConfigError(message=msg, test_units=self._testUnits)
            else:
                if config.has_section(section):
                    for o in required_options[section]:
                        if not config.has_option(section, o):
                            msg = 'Config file must have option \'{0}\' for section \'{1}\''.format(o, section)
                            raise errors.BTSSystemConfigError(message=msg, test_units=self._testUnits)
                else:
                    msg = "Config file must have section '{0}'.".format(section)
                    raise errors.BTSSystemConfigError(message=msg, test_units=self._testUnits)
        return True

    def resource_available(self, resource_address: str):
        if resource_address not in self._availableResources:
            return False
        return True

    def init_impedance_analyzer(self):
        address = self._systemConfig.get(section='impedance_analyzer', option='address')
        name = self._systemConfig.get(section='impedance_analyzer', option='name')
        if not self.resource_available(resource_address=address):
            msg = 'The impedance analyzer is not available at address \'{0}\''.format(address)
            raise errors.InstrumentError(address=address, resource_name=name, message=msg)

        self._impedanceAnalyzer = ia.ImpedanceAnalyzer(address=address, resource_manager=self._resourceManager)

    def init_keithley(self):
        address = self._systemConfig.get(section='keithley', option='address')
        name = self._systemConfig.get(section='keithley', option='name')
        if not self.resource_available(resource_address=address):
            msg = 'The Keithley source-meter is not available at address \'{0}\''.format(address)
            raise errors.InstrumentError(address=address, resource_name=name, message=msg)
        self._keithley = keithley.Keithley(address=address, resource_manager=self._resourceManager)

    def init_arduino_boards(self):
        for i in range(self._testUnits):
            section_name = 'arduino{0:d}'.format(i)
            address = self._systemConfig.get(section=section_name, option='address')
            name = self._systemConfig.get(section=section_name, option='name')
            all_options = dict(self._systemConfig.items(section=section_name))
            pin_mappings = {
                'keithley': all_options['pin_keithley'],
                'fan': all_options['pin_fan']
            }
            import re
            pattern = re.compile(r'pin(\d+)')
            for o in all_options:
                match = re.match(pattern=pattern, string=o)
                if match:
                    pin = int(match.groups()[0])
                    pin_mappings[pin] = all_options[o]
            if self.resource_available(resource_address=address):
                self._arduinos.append(arduino_board.ArduinoBoard(address=address, name=name,
                                                                 pin_mappings=pin_mappings))
            else:
                msg = 'The arduino board \'{0}\' was not found in address: \'{1}\''.format(name, address)
                raise errors.InstrumentError(address=address, resource_name=name, message=msg)

    def init_hotplates(self, calibration_file: str):
        calibration = configparser.ConfigParser()
        calibration.read(calibration_file)

        for i in range(self._testUnits):
            section_name = 'hotplate{0:d}'.format(i)
            address: str = self._systemConfig.get(section=section_name, option='address')
            name: str = self._systemConfig.get(section=section_name, option='name')
            if self.resource_available(resource_address=address):
                hp = hotplate.Hotplate(address=address, name=name)
                hp.load_calibration(calibration)
                self._hotPlates.append(hp)
            else:
                msg = "The hotplate '{0}' was not found in address: '{1}'".format(name, address)
                raise errors.InstrumentError(address=address, resource_name=name, message=msg)

    @staticmethod
    def _read_json_file(filename: str):
        file = open(filename, 'r')
        json_data = json.load(file)
        file.close()
        return json_data

    @staticmethod
    def _create_path(path: str):
        if not os.path.exists(path):
            os.makedirs(path)

    @staticmethod
    def _create_logger(path: str, name: str = 'experiment_logger',
                       level: [str, int] = 'CRITICAL', console: bool = False) -> logging.Logger:
        experiment_logger = logging.getLogger(name)
        experiment_logger.setLevel(level)
        filename = 'progress.log'
        log_file = os.path.join(path, filename)
        # create file handler which logs even critical messages
        fh = logging.FileHandler(log_file)
        fh.setLevel(level=level)
        experiment_logger.addHandler(fh)

        # create formatter and add it to the handlers
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)

        if console:
            ch = logging.StreamHandler()
            ch.setLevel(level=level)
            ch.setFormatter(formatter)
            experiment_logger.addHandler(ch)

        return experiment_logger

    def _print(self, msg: str, level="INFO"):
        level_no = self._loggingLevels[level]
        if self._mainLogger is None:
            print(msg)
        elif isinstance(self._mainLogger, logging.Logger):

            self._mainLogger.log(level_no, msg)