import numpy as np
import random

from basic.read_data import DataSet

def get_feed_dict(model, batch, is_train, supervised=True):
  assert isinstance(batch, DataSet)
  config = model.config

  N, M, JX, JQ, VW, VC, d, W = \
    config.batch_size, config.max_num_sents, config.max_sent_size, \
    config.max_ques_size, config.word_vocab_size, config.char_vocab_size, config.hidden_size, config.max_word_size
  if config.split_supports:
    M = 1

  feed_dict = {}

  if config.split_supports:
    x, cx, x_mask = [], [], []
  else:
    x = np.zeros([N, M, JX], dtype='int32')
    cx = np.zeros([N, M, JX, W], dtype='int32')
    x_mask = np.zeros([N, M, JX], dtype='bool')

  q = np.zeros([N, JQ], dtype='int32')
  cq = np.zeros([N, JQ, W], dtype='int32')
  q_mask = np.zeros([N, JQ], dtype='bool')
  x_group = np.zeros([N], dtype='int32')

  q_sub = np.zeros([N, config.max_ques_sub_size], dtype='int32')
  cq_sub = np.zeros([N, config.max_ques_sub_size, W], dtype='int32')
  qsub_mask = np.zeros([N, config.max_ques_sub_size], dtype='bool')

  feed_dict[model.x] = x
  feed_dict[model.x_mask] = x_mask
  feed_dict[model.cx] = cx
  feed_dict[model.q] = q
  feed_dict[model.cq] = cq
  feed_dict[model.q_sub] = q_sub
  feed_dict[model.cq_sub] = cq_sub
  feed_dict[model.q_mask] = q_mask
  feed_dict[model.q_sub_mask] = qsub_mask
  feed_dict[model.is_train] = is_train

  feed_dict[model.x_group] = x_group

  if config.use_glove_for_unk:
    feed_dict[model.new_emb_mat] = batch.shared['new_emb_mat']

  X = batch.data['x']
  CX = batch.data['cx']

  if supervised:
    if config.split_supports is True:
      y = np.zeros([N, M, config.max_para_size], dtype='bool')
      y2 = np.zeros([N, M, config.max_para_size], dtype='bool')
      wy = np.zeros([N, M, config.max_para_size], dtype='bool')
    else:
      y = np.zeros([N, M, JX], dtype='bool')
      y2 = np.zeros([N, M, JX], dtype='bool')
      wy = np.zeros([N, M, JX], dtype='bool')
    na = np.zeros([N], dtype='bool')
    feed_dict[model.y] = y
    feed_dict[model.y2] = y2
    feed_dict[model.wy] = wy
    feed_dict[model.na] = na

    if config.supervise_final_doc or (config.oracle is not None):
      answer_doc_ids = - np.ones([N, 10], dtype='int32')
      answer_word_ids = - np.ones([N, 10], dtype='int32')
      feed_dict[model.answer_doc_ids] = answer_doc_ids
      feed_dict[model.answer_word_ids] = answer_word_ids
      for i, (doc_id, word_id, nai) in enumerate(zip(batch.data['answer_doc_ids'], batch.data['answer_ids_in_doc'], batch.data['na'])):
        if nai:
          if config.supervise_final_doc or config.oracle is not None:
            answer_doc_ids[i, 0] = 0
            answer_word_ids[i, 0] = 0
          continue
        assert len(doc_id) == len(word_id), (len(doc_id), len(word_id))
        if config.shuffle_answer_doc_ids:
          joint_id = list(zip(doc_id, word_id))
          np.random.shuffle(joint_id)  # shuffle doc_ids and word_ids jointly to randomize the supervision label.
          doc_id, word_id = zip(*joint_id)

        for j in range(min(10, len(doc_id))):
          if config.select_top_n_doc > 0:
            if doc_id[j] < config.select_top_n_doc:
              answer_doc_ids[i, j] = doc_id[j]
              answer_word_ids[i, j] = word_id[j]
            else:
              answer_doc_ids[i, 0] = np.random.randint(0, config.select_top_n_doc)
              answer_word_ids[i, 0] = 0
              break
          else:
            answer_doc_ids[i, j] = doc_id[j]
            answer_word_ids[i, j] = word_id[j]

      for i,adis in enumerate(answer_doc_ids):
        if config.supervise_final_doc and adis[0] == -1:
          answer_doc_ids[i, 0] = 0
          answer_word_ids[i, 0] = 0

    if config.supervise_first_doc:
      first_doc_ids = np.zeros([N], dtype='int32')
      feed_dict[model.first_doc_ids] = first_doc_ids
      if config.select_top_n_doc == 0:
        for i, doc_id in enumerate(batch.data['first_doc_ids']):
          first_doc_ids[i] = doc_id

    if config.mac_prediction == 'candidates':
      cand_spanss = batch.data['cand_span']
      cand_span_y = batch.data['cand_span_y']
      cand_wordss = batch.data['cand_word']
      cand_word_y = batch.data['cand_word_y']

      max_cand_size = 0
      for ci in cand_spanss:
        max_cand_size = max(max_cand_size, len(ci))

      candidate_spans = np.zeros([N, M, max_cand_size, 2], dtype='int32')
      candidate_span_y = np.zeros([N, M], dtype='int32')
      num_exceed_cand = np.zeros([N, M, max_cand_size], dtype='int32')
      feed_dict[model.candidate_spans] = candidate_spans
      feed_dict[model.candidate_span_y] = candidate_span_y
      feed_dict[model.num_exceed_cand] = num_exceed_cand

      for i, (cand_spans, nai) in enumerate(zip(cand_spanss, batch.data['na'])):
        num_exceed_candi = 0
        num_exceed_cand_bfr_y = 0
        for j, cand_span in enumerate(cand_spans):
          assert cand_span[0][0] == 0, (cand_span)
          assert cand_span[1][0] == 0, (cand_span)

          # Ignore candidate spans that exceed max_para_size.
          if cand_span[0][1] > config.max_para_size or cand_span[1][1] > config.max_para_size:
            num_exceed_candi += 1
            if not nai:
              assert j != cand_span_y[i][0]
              if j < cand_span_y[i][0]:
                num_exceed_cand_bfr_y += 1
            continue
          num_exceed_cand[i, 0, j-num_exceed_candi] = num_exceed_candi
          candidate_spans[i, 0, j-num_exceed_candi, 0] = cand_span[0][1]
          candidate_spans[i, 0, j-num_exceed_candi, 1] = cand_span[1][1]

        if nai:
          #na[i] = nai
          continue

        assert cand_span_y[i][0] < len(cand_spans), (cand_span_y[i][0], len(cand_spans))
        candidate_span_y[i, 0] = cand_span_y[i][0] - num_exceed_cand_bfr_y

    for i, (xi, cxi, yi, nai) in enumerate(zip(X, CX, batch.data['y'], batch.data['na'])):
      if nai:
        na[i] = nai
        continue
      start_idx, stop_idx = random.choice(yi)
      j, k = start_idx
      j2, k2 = stop_idx

      y[i, j, k] = True
      y2[i, j2, k2-1] = True
      if j == j2:
        wy[i, j, k:k2] = True
      else:
        wy[i, j, k:len(batch.data['x'][i][j])] = True
        wy[i, j2, :k2] = True

  def _get_word(word):
    d = batch.shared['word2idx']
    for each in (word, word.lower(), word.capitalize(), word.upper()):
      if each in d:
        return d[each]
    if config.use_glove_for_unk:
      d2 = batch.shared['new_word2idx']
      for each in (word, word.lower(), word.capitalize(), word.upper()):
        if each in d2:
          return d2[each] + len(d)
    return 1

  def _get_char(char):
    d = batch.shared['char2idx']
    if char in d:
      return d[char]
    return 1

  if model.period_id is None:
    model.period_id = _get_word('.')
    print("period_id: %d" %model.period_id)

  max_sent_size = 0
  for xi in X:
    for xij in xi:
      max_sent_size = max(max_sent_size, len(xij))

  if config.split_supports is True:
    assert max_sent_size <= config.max_sent_size, (max_sent_size)

  for i, xi in enumerate(X):
    word_count = 0

    for j, xij in enumerate(xi):
      if j == config.max_num_sents:
        raise Exception("Exceed max_num_sents.")
        break
      if config.split_supports and word_count >= config.max_para_size:
        if config.supervise_first_doc and j <= first_doc_ids[i]:
          first_doc_ids[i] = 0
        if config.supervise_final_doc:
          for count,id in enumerate(answer_doc_ids[i]):
            if j <= id:
              if config.supervise_final_doc and count == 0:  # Must have a valid doc id as the label of the final attention loss.
                answer_doc_ids[i] = [0] + [-1]*9
                answer_word_ids[i] = [0] + [-1]*9
              else:
                answer_doc_ids[i][count:] = - np.ones(10 - count)
                answer_word_ids[i][count:] = - np.ones(10 - count)
              break
        break

      if config.split_supports:  # Add every sentence as a separate batch.
        _xij = np.zeros([1, max_sent_size], dtype='int32')
        _xij_mask = np.zeros([1, max_sent_size], dtype='bool')
        x.append(_xij)
        x_mask.append(_xij_mask)

      for k, xijk in enumerate(xij):
        if config.split_supports and word_count >= config.max_para_size:
        # same as word_count > config.max_sent_size if para_size_th and sent_size_th is set to same.
          break
        if k == config.max_sent_size:
          break
        each = _get_word(xijk)
        assert isinstance(each, int), each
        if config.split_supports:
          x[-1][0, k] = each
          x_mask[-1][0, k] = True
        else:
          x[i, j, k] = each
          x_mask[i, j, k] = True

        word_count += 1
      x_group[i] += 1

    if config.split_supports:
      assert word_count <= config.max_para_size, (word_count)

  for i, cxi in enumerate(CX):
    # Create a batch
    word_count = 0

    for j, cxij in enumerate(cxi):
      if j == config.max_num_sents:
        raise Exception("Exceed max_num_sents.")
        break
      if config.split_supports and word_count >= config.max_para_size:
        break

      if config.split_supports:  # Add every sentence as a separate batch.
        _cxij = np.zeros([1, max_sent_size, W], dtype='int32')
        cx.append(_cxij)

      for k, cxijk in enumerate(cxij):
        if config.split_supports and word_count >= config.max_para_size:
          break
        if k == config.max_sent_size:
          break

        for l, cxijkl in enumerate(cxijk):
          if l == config.max_word_size:
            break
          if config.split_supports:
            cx[-1][0, k, l] = _get_char(cxijkl)
          else:
            cx[i, j, k, l] = _get_char(cxijkl)

        word_count += 1 #

    if config.split_supports:
      assert word_count <= config.max_para_size, (word_count)

  for i, qi in enumerate(batch.data['q']):
    for j, qij in enumerate(qi):
      q[i, j] = _get_word(qij)
      q_mask[i, j] = True

  for i, cqi in enumerate(batch.data['cq']):
    for j, cqij in enumerate(cqi):
      for k, cqijk in enumerate(cqij):
        cq[i, j, k] = _get_char(cqijk)
        if k + 1 == config.max_word_size:
          break

  if config.get_query_subject:
    for i, qi in enumerate(batch.data['q2']):
      for j, qij in enumerate(qi):
        q_sub[i, j] = _get_word(qij)
        qsub_mask[i, j] = True

    for i, cqi in enumerate(batch.data['cq2']):
      for j, cqij in enumerate(cqi):
        for k, cqijk in enumerate(cqij):
          cq_sub[i, j, k] = _get_char(cqijk)
          if k + 1 == config.max_word_size:
            break

  all_period_loc_plus_1 = []
  x_len = np.sum(np.asarray(x_mask).astype(int), axis=2)
  for i, xi in enumerate(x):
    period_found = [j + 1 for j, period in enumerate(list(xi[0])) if period == model.period_id]
    if len(period_found) > 10:
      all_period_loc_plus_1.append([period_found[:10]])
    else:
      to_append = period_found + [x_len[i][0]]*(10-len(period_found))
      all_period_loc_plus_1.append([to_append])
  x_sents_len = all_period_loc_plus_1
  feed_dict[model.x_sents_len] = np.stack(x_sents_len)
  if supervised:
    if config.split_supports is False:
      assert np.sum(~(x_mask | ~wy)) == 0  # if x_mask == 0, then wy must be 0

  num_examples = len(batch.data['x'])
  if num_examples < config.batch_size and (config.mac_reasoning_unit == 'attention-lstm'): # Last batch every epoch
    cutoff_1 = np.sum(x_group[:(config.batch_size-num_examples)])
    feed_dict[model.x] = np.stack(x+x[:cutoff_1])
    feed_dict[model.cx] = np.stack(cx+cx[:cutoff_1])
    feed_dict[model.x_mask] = np.stack(x_mask+x_mask[:cutoff_1])
    feed_dict[model.x_sents_len] = np.stack(x_sents_len+x_sents_len[:cutoff_1])

    cutoff_2 = config.batch_size - num_examples
    feed_dict[model.x_group] = np.concatenate([x_group[:num_examples], x_group[:cutoff_2]])
    feed_dict[model.q] = np.concatenate([q[:num_examples], q[:cutoff_2]])
    feed_dict[model.cq] = np.concatenate([cq[:num_examples], cq[:cutoff_2]])
    feed_dict[model.q_sub] = np.concatenate([q_sub[:num_examples], q_sub[:cutoff_2]])
    feed_dict[model.cq_sub] = np.concatenate([cq_sub[:num_examples], cq_sub[:cutoff_2]])
    feed_dict[model.q_mask] = np.concatenate([q_mask[:num_examples], q_mask[:cutoff_2]])
    feed_dict[model.q_sub_mask] = np.concatenate([qsub_mask[:num_examples], qsub_mask[:cutoff_2]])
  else:
    feed_dict[model.x] = np.stack(x)
    feed_dict[model.cx] = np.stack(cx)
    feed_dict[model.x_mask] = np.stack(x_mask)

  assert len(x) == len(x_mask)
  assert len(x) == sum(x_group), (len(x), sum(x_group))
  return feed_dict


#############################################################################
#
# VFRAME
# MIT License
# Copyright (c) 2020 Adam Harvey and VFRAME
# https://vframe.io
#
#############################################################################

from os.path import join
import time
from pathlib import Path

import cv2 as cv
import numpy as np

from vframe.models.cvmodels import ClassifyResult, ClassifyResults
from vframe.models.cvmodels import DetectResult, DetectResults
from vframe.models.cvmodels import BenchmarkResult
from vframe.utils import file_utils, im_utils
from vframe.settings import app_cfg

class NetProc:
  """Basic Net Processor for OpenCV DNN models
  """

  def __init__(self, dnn_cfg):
    """Instantiate an DNN network model
    """
    self.log = app_cfg.LOG
    self.dnn_cfg = dnn_cfg

    # build dnn net
    if self.dnn_cfg.fp_config is None:
      self.net = cv.dnn.readNet(self.dnn_cfg.fp_model)
    else:
      self.net = cv.dnn.readNet(self.dnn_cfg.fp_model, self.dnn_cfg.fp_config)
    self.net.setPreferableBackend(self.dnn_cfg.dnn_backend)
    self.net.setPreferableTarget(self.dnn_cfg.dnn_target)

    # load class labels
    if self.dnn_cfg.labels_exist:
      self.labels = file_utils.load_txt(dnn_cfg.fp_labels)  # line-delimited list of class labels
    else:
      self.log.debug(f'labels file does not exist: {dnn_cfg.fp_labels}. Using empty label set.')
      self.labels = []

    self.__post_init__()

  def __post_init__(self):
    """Override me
    """
    pass

  def _pre_process(self, im):
    """Pre-process image
    """
    cfg = self.dnn_cfg

    if cfg.width == cfg.height and not cfg.width == cfg.height:
      cfg.width = min(cfg.width, cfg.height)
      cfg.height = cfg.width
      self.log.warning(f'Width and height must be equal. Forcing to lowest size: {cfg.width}')

    self.frame_dim_orig = im.shape[:2][::-1]
    im = im_utils.resize(im, width=cfg.width, height=cfg.height, force_fit=cfg.fit)
    self.frame_dim_resized = im.shape[:2][::-1]
    dim = self.frame_dim_resized if cfg.fit else cfg.size
    blob = cv.dnn.blobFromImage(im, cfg.scale, dim, cfg.mean, crop=cfg.crop, swapRB=cfg.rgb)
    self.net.setInput(blob)

  def _post_process(self, outs):
    """Post process net output and return ProcessorResult
    """
    self.log.error('Override this')
    return ProcessorResult(0, 0.0)

  def features(self):
    """Probably only used for classification networks. Consider removing.
    """
    self.log.error('Override this')
    return np.array()

  def fps(self, im=None, n_iters=10):
    """Benchmark model FPS on image
    :param im: (np.ndarray) image
    :pram n_iters: (int) iterations
    """
    if im is None:
      im = im_utils.create_blank_im(640, 480)
    _ = self.infer(im)  # warmup
    start_time = time.time()
    for i in range(n_iters):
      _ = self.infer(im)
    fps = n_iters / (time.time() - start_time)
    return fps

  def infer(self, im):
    """Runs pre-processor, inference, and post-processor
    """
    self._pre_process(im)
    if self.dnn_cfg.layers:
      outs = self.net.forward(self.dnn_cfg.layers)
    else:
      outs = self.net.forward()
    results = self._post_process(outs)
    return results

  def _perf_ms(self):
    """Returns network forward pass performance time in milliseconds
    """
    t, _ = self.net.getPerfProfile()
    return t * 1000.0 / cv.getTickFrequency()

# -----------------------------------------------------------------------------
#
# Classification networks
#
# -----------------------------------------------------------------------------

class ClassificationProc(NetProc):
  """Classification inference processor
  """
  limit = 1

  def _post_process(self, outs):
    preds = outs.flatten()
    classify_results = []
    if self.limit > 1:
      idxs = np.argsort(preds)[::-1][:self.limit]  # top N indices
    else:
      idxs = [np.argmax(preds)]
    for idx in idxs:
      if preds[idx] > self.dnn_cfg.threshold:
        classify_result = ClassifyResult(idx, preds[idx], self.labels[idx])
        classify_results.append(classify_result)

    return ClassifyResults(classify_results)

  def features(self, im):
    cfg = self.dnn_cfg
    self._pre_process(im)
    feat_vec = np.squeeze(self.net.forward(cfg.features).flatten())
    return feat_vec / np.linalg.norm(feat_vec)

# -----------------------------------------------------------------------------
#
# Detection base
#
# -----------------------------------------------------------------------------

class DetectionProc(NetProc):

  def _nms(self, detect_results):
    """Apply non-maximum suppression and filter detection results
    :param detect_results: List[DetectResult]
    :returns List[DetectResult]
    """
    confidences = [float(d.confidence) for d in detect_results]
    boxes = [d.bbox.xywh_int for d in detect_results]
    idxs = cv.dnn.NMSBoxes(boxes, confidences, self.dnn_cfg.threshold, self.dnn_cfg.nms_threshold)
    detect_results_nms = [detect_results[i[0]] for i in idxs]
    return detect_results_nms


import numpy as np
import pandas as pd
import torch
import pickle
from tqdm import tqdm

def correct_predictions(pred_prob_list, series_len_list):
    eps = 0.000001
    pred_prob_list_corrected = np.zeros(pred_prob_list.shape, dtype=np.float32)
    start = 0
    for i in tqdm(range(len(series_len_list))):
        end = series_len_list[i]
        negative_exam_for_pe = pred_prob_list[start+0]
        indeterminate = pred_prob_list[start+1]
        chronic_pe = pred_prob_list[start+2]
        acute_and_chronic_pe = pred_prob_list[start+3]
        central_pe = pred_prob_list[start+4]
        leftsided_pe = pred_prob_list[start+5]
        rightsided_pe = pred_prob_list[start+6]
        rv_lv_ratio_gte_1 = pred_prob_list[start+7]
        rv_lv_ratio_lt_1 = pred_prob_list[start+8]
        image_pe = pred_prob_list[start+9:end]

        loss_weight_list = np.zeros(pred_prob_list[start:end].shape, dtype=np.float32)
        loss_weight_list[0] = 0.0736196319
        loss_weight_list[1] = 0.09202453988
        loss_weight_list[2] = 0.1042944785
        loss_weight_list[3] = 0.1042944785
        loss_weight_list[4] = 0.1877300613
        loss_weight_list[5] = 0.06257668712
        loss_weight_list[6] = 0.06257668712
        loss_weight_list[7] = 0.2346625767
        loss_weight_list[8] = 0.0782208589
        loss_weight_list[9:] = 0.07361963*0.005

        if (np.amax(image_pe)<=0.5) and (int(negative_exam_for_pe>0.5)+int(indeterminate>0.5)==1) and (int(chronic_pe>0.5)+int(acute_and_chronic_pe>0.5)==0) and (int(central_pe>0.5)+int(leftsided_pe>0.5)+int(rightsided_pe>0.5)==0) and (int(rv_lv_ratio_gte_1>0.5)+int(rv_lv_ratio_lt_1>0.5)==0):
            pred_prob_list_corrected[start:end] = pred_prob_list[start:end]
        elif (np.amax(image_pe)>0.5) and (int(negative_exam_for_pe>0.5)+int(indeterminate>0.5)==0) and (int(chronic_pe>0.5)+int(acute_and_chronic_pe>0.5)<2) and (int(central_pe>0.5)+int(leftsided_pe>0.5)+int(rightsided_pe>0.5)>0) and (int(rv_lv_ratio_gte_1>0.5)+int(rv_lv_ratio_lt_1>0.5)==1):
            pred_prob_list_corrected[start:end] = pred_prob_list[start:end]
        else:
            to_neg = pred_prob_list[start:end].copy()
            for n in range(len(image_pe)):
                if image_pe[n]>0.5:
                    to_neg[9+n] = 0.5
            if negative_exam_for_pe>0.5 and indeterminate>0.5:
                if negative_exam_for_pe>indeterminate:
                    to_neg[1] = 0.5
                else:
                    to_neg[0] = 0.5
            elif negative_exam_for_pe<=0.5 and indeterminate<=0.5:
                if negative_exam_for_pe>indeterminate:
                    to_neg[0] = 0.5+eps
                else:
                    to_neg[1] = 0.5+eps
            if chronic_pe>0.5:
                to_neg[2] = 0.5
            if acute_and_chronic_pe>0.5:
                to_neg[3] = 0.5
            if central_pe>0.5:
                to_neg[4] = 0.5
            if leftsided_pe>0.5:
                to_neg[5] = 0.5
            if rightsided_pe>0.5:
                to_neg[6] = 0.5
            if rv_lv_ratio_gte_1>0.5:
                to_neg[7] = 0.5
            if rv_lv_ratio_lt_1>0.5:
                to_neg[8] = 0.5

            to_pos = pred_prob_list[start:end].copy()
            if np.amax(image_pe)<=0.5:
                max_idx = np.argmax(image_pe)
                to_pos[9+max_idx] = 0.5+eps
            if negative_exam_for_pe>0.5:
                to_pos[0] = 0.5
            if indeterminate>0.5:
                to_pos[1] = 0.5
            if chronic_pe>0.5 and acute_and_chronic_pe>0.5:
                if chronic_pe>acute_and_chronic_pe:
                    to_pos[3] = 0.5
                else:
                    to_pos[2] = 0.5
            if central_pe<=0.5 and leftsided_pe<=0.5 and rightsided_pe<=0.5:
                if central_pe>leftsided_pe and central_pe>rightsided_pe:
                    to_pos[4] = 0.5+eps
                if leftsided_pe>central_pe and leftsided_pe>rightsided_pe:
                    to_pos[5] = 0.5+eps
                if rightsided_pe>central_pe and rightsided_pe>leftsided_pe:
                    to_pos[6] = 0.5+eps
            if rv_lv_ratio_gte_1>0.5 and rv_lv_ratio_lt_1>0.5:
                if rv_lv_ratio_gte_1>rv_lv_ratio_lt_1:
                    to_pos[8] = 0.5
                else:
                    to_pos[7] = 0.5
            elif rv_lv_ratio_gte_1<=0.5 and rv_lv_ratio_lt_1<=0.5:
                if rv_lv_ratio_gte_1>rv_lv_ratio_lt_1:
                    to_pos[7] = 0.5+eps
                else:
                    to_pos[8] = 0.5+eps

            loss_weight_list1 = torch.tensor(loss_weight_list, dtype=torch.float32)
            pred_prob_list1 = torch.tensor(pred_prob_list[start:end], dtype=torch.float32)
            pred_prob_list_neg = torch.tensor(to_neg, dtype=torch.float32)
            pred_prob_list_pos = torch.tensor(to_pos, dtype=torch.float32)
            #print(loss_weight_list1.shape, pred_prob_list1.shape, pred_prob_list_neg.shape, pred_prob_list_pos.shape)
            to_neg_loss = ((torch.nn.BCELoss(reduction='none')(pred_prob_list1, pred_prob_list_neg)*loss_weight_list1).sum() / loss_weight_list1.sum()).numpy()
            to_pos_loss = ((torch.nn.BCELoss(reduction='none')(pred_prob_list1, pred_prob_list_pos)*loss_weight_list1).sum() / loss_weight_list1.sum()).numpy()

            if to_neg_loss>to_pos_loss:
                pred_prob_list_corrected[start:end] = to_pos
            else:
                pred_prob_list_corrected[start:end] = to_neg

        start = series_len_list[i]
    return pred_prob_list_corrected

def check_consistency(sub, test):

    '''
    Checks label consistency and returns the errors

    Args:
    sub   = submission dataframe (pandas)
    test  = test.csv dataframe (pandas)
    '''

    # EXAM LEVEL
    for i in test['StudyInstanceUID'].unique():
        df_tmp = sub.loc[sub.id.str.contains(i, regex = False)].reset_index(drop = True)
        df_tmp['StudyInstanceUID'] = df_tmp['id'].str.split('_').str[0]
        df_tmp['label_type']       = df_tmp['id'].str.split('_').str[1:].apply(lambda x: '_'.join(x))
        del df_tmp['id']
        if i == test['StudyInstanceUID'].unique()[0]:
            df = df_tmp.copy()
        else:
            df = pd.concat([df, df_tmp], axis = 0)
    df_exam = df.pivot(index = 'StudyInstanceUID', columns = 'label_type', values = 'label')

    # IMAGE LEVEL
    df_image = sub.loc[sub.id.isin(test.SOPInstanceUID)].reset_index(drop = True)
    df_image = df_image.merge(test, how = 'left', left_on = 'id', right_on = 'SOPInstanceUID')
    df_image.rename(columns = {"label": "pe_present_on_image"}, inplace = True)
    del df_image['id']

    # MERGER
    df = df_exam.merge(df_image, how = 'left', on = 'StudyInstanceUID')
    ids    = ['StudyInstanceUID', 'SeriesInstanceUID', 'SOPInstanceUID']
    labels = [c for c in df.columns if c not in ids]
    df = df[ids + labels]

    # SPLIT NEGATIVE AND POSITIVE EXAMS
    df['positive_images_in_exam'] = df['StudyInstanceUID'].map(df.groupby(['StudyInstanceUID']).pe_present_on_image.max())
    df_pos = df.loc[df.positive_images_in_exam >  0.5]
    df_neg = df.loc[df.positive_images_in_exam <= 0.5]

    # CHECKING CONSISTENCY OF POSITIVE EXAM LABELS
    rule1a = df_pos.loc[((df_pos.rv_lv_ratio_lt_1  >  0.5)  &
                         (df_pos.rv_lv_ratio_gte_1 >  0.5)) |
                        ((df_pos.rv_lv_ratio_lt_1  <= 0.5)  &
                         (df_pos.rv_lv_ratio_gte_1 <= 0.5))].reset_index(drop = True)
    rule1a['broken_rule'] = '1a'
    rule1b = df_pos.loc[(df_pos.central_pe    <= 0.5) &
                        (df_pos.rightsided_pe <= 0.5) &
                        (df_pos.leftsided_pe  <= 0.5)].reset_index(drop = True)
    rule1b['broken_rule'] = '1b'
    rule1c = df_pos.loc[(df_pos.acute_and_chronic_pe > 0.5) &
                        (df_pos.chronic_pe           > 0.5)].reset_index(drop = True)
    rule1c['broken_rule'] = '1c'
    rule1d = df_pos.loc[(df_pos.indeterminate        > 0.5) |
                        (df_pos.negative_exam_for_pe > 0.5)].reset_index(drop = True)
    rule1d['broken_rule'] = '1d'

    # CHECKING CONSISTENCY OF NEGATIVE EXAM LABELS
    rule2a = df_neg.loc[((df_neg.indeterminate        >  0.5)  &
                         (df_neg.negative_exam_for_pe >  0.5)) |
                        ((df_neg.indeterminate        <= 0.5)  &
                         (df_neg.negative_exam_for_pe <= 0.5))].reset_index(drop = True)
    rule2a['broken_rule'] = '2a'
    rule2b = df_neg.loc[(df_neg.rv_lv_ratio_lt_1     > 0.5) |
                        (df_neg.rv_lv_ratio_gte_1    > 0.5) |
                        (df_neg.central_pe           > 0.5) |
                        (df_neg.rightsided_pe        > 0.5) |
                        (df_neg.leftsided_pe         > 0.5) |
                        (df_neg.acute_and_chronic_pe > 0.5) |
                        (df_neg.chronic_pe           > 0.5)].reset_index(drop = True)
    rule2b['broken_rule'] = '2b'

    # MERGING INCONSISTENT PREDICTIONS
    errors = pd.concat([rule1a, rule1b, rule1c, rule1d, rule2a, rule2b], axis = 0)

    # OUTPUT
    print('Found', len(errors), 'inconsistent predictions')
    return errors

gt_list_seresnext50_128 = np.load('predictions/gt_list_seresnext50_128.npy')
loss_weight_list_seresnext50_128 = np.load('predictions/loss_weight_list_seresnext50_128.npy')
pred_prob_list_seresnext50_128 = np.load('predictions/pred_prob_list_seresnext50_128.npy')
gt_list_seresnext101_128 = np.load('predictions/gt_list_seresnext101_128.npy')
loss_weight_list_seresnext101_128 = np.load('predictions/loss_weight_list_seresnext101_128.npy')
pred_prob_list_seresnext101_128 = np.load('predictions/pred_prob_list_seresnext101_128.npy')

gt_list_seresnext50_192 = np.load('predictions/gt_list_seresnext50_192.npy')
loss_weight_list_seresnext50_192 = np.load('predictions/loss_weight_list_seresnext50_192.npy')
pred_prob_list_seresnext50_192 = np.load('predictions/pred_prob_list_seresnext50_192.npy')
gt_list_seresnext101_192 = np.load('predictions/gt_list_seresnext101_192.npy')
loss_weight_list_seresnext101_192 = np.load('predictions/loss_weight_list_seresnext101_192.npy')
pred_prob_list_seresnext101_192 = np.load('predictions/pred_prob_list_seresnext101_192.npy')

gt_list = gt_list_seresnext50_128
loss_weight_list = loss_weight_list_seresnext50_128
pred_prob_list = (pred_prob_list_seresnext50_128 + pred_prob_list_seresnext101_128) / 2.0

gt_list1 = torch.tensor(gt_list, dtype=torch.float32)
loss_weight_list1 = torch.tensor(loss_weight_list, dtype=torch.float32)
pred_prob_list1 = torch.tensor(pred_prob_list, dtype=torch.float32)
kaggle_loss = torch.nn.BCELoss(reduction='none')(pred_prob_list1, gt_list1)
kaggle_loss = (kaggle_loss*loss_weight_list1).sum() / loss_weight_list1.sum()
print(kaggle_loss.numpy())

series_len_list = np.load('predictions/series_len_list.npy')
pred_prob_list = correct_predictions(pred_prob_list, series_len_list)

gt_list1 = torch.tensor(gt_list, dtype=torch.float32)
loss_weight_list1 = torch.tensor(loss_weight_list, dtype=torch.float32)
pred_prob_list1 = torch.tensor(pred_prob_list, dtype=torch.float32)
kaggle_loss = torch.nn.BCELoss(reduction='none')(pred_prob_list1, gt_list1)
kaggle_loss = (kaggle_loss*loss_weight_list1).sum() / loss_weight_list1.sum()
print(kaggle_loss.numpy())

id_list = np.load('predictions/id_list.npy')
train_df = pd.read_csv('../../input/train.csv')
study_id_list = train_df['StudyInstanceUID'].values
series_id_list = train_df['SeriesInstanceUID'].values
image_id_list = train_df['SOPInstanceUID'].values
with open('../process_input/split2/series_list_valid.pickle', 'rb') as f:
    series_list_valid = pickle.load(f)
for i in range(len(series_list_valid)):
    series_list_valid[i] = series_list_valid[i].split('_')[0]
series_list_valid = set(series_list_valid)
valid_study_id_list = []
valid_series_id_list = []
valid_image_id_list = []
for i in range(len(study_id_list)):
    if study_id_list[i] in series_list_valid:
        valid_study_id_list.append(study_id_list[i])
        valid_series_id_list.append(series_id_list[i])
        valid_image_id_list.append(image_id_list[i])
valid_df = pd.DataFrame(data={'StudyInstanceUID': valid_study_id_list, 'SeriesInstanceUID': valid_series_id_list, 'SOPInstanceUID': valid_image_id_list})
sub_df = pd.DataFrame(data={'id': id_list, 'label': pred_prob_list})
print(len(id_list), len(series_len_list), len(sub_df), len(valid_df))

errors = check_consistency(sub_df, valid_df)


import numpy as np
import cv2
import matplotlib.pyplot as plt
import mtcnn
from mtcnn.mtcnn import MTCNN
import os
import tqdm

path = '...' #@@@ OVERRIDE: Path of a folder containing images

# Initialize Multi-task Cascaded Convolutional Neural Network
detector = MTCNN()
all_ids = sorted(os.listdir(path))

# Create an empty list to store the ids of images without a face
no_face_ids = []
for id in tqdm.tqdm(all_ids):
    img = cv2.imread(os.path.join(path, id))
    faces = detector.detect_faces(img)
    if not faces:
        no_face_ids.append(id)

# Save in a txt file the ids of images with no face
with open(os.path.join(path, 'ids_no_face.txt'), 'w') as f:
    f.write('\n'.join(no_face_ids))

# Visuaize face bbox
def visualize_face_detector(path, id):
    detector = MTCNN()
    img = cv2.imread(os.path.join(path, id))
    faces = detector.detect_faces(img)
    for result in faces:
        x, y, w, h = result['box']
        x1, y1 = x + w, y + h
        print(x, y, x1, y1)
        cv2.rectangle(img, (x, y), (x1, y1), (255, 0, 0), 2)
        plt.imshow(img);


import torch
from torch import nn, optim, utils
import numpy as np
import os
import time
import dill
import json
import random
import pathlib
import warnings
from tqdm import tqdm
import visualization
import evaluation
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from argument_parser import args
from model.trajectron import Trajectron
from model.model_registrar import ModelRegistrar
from model.model_utils import cyclical_lr
from model.dataset import EnvironmentDataset, collate
from tensorboardX import SummaryWriter
# torch.autograd.set_detect_anomaly(True)

from snce.contrastive import SocialNCE
from snce.model import ProjHead, EventEncoder

np.set_printoptions(precision=2, suppress=True)
torch.set_printoptions(precision=2, sci_mode=False)

if not torch.cuda.is_available() or args.device == 'cpu':
    args.device = torch.device('cpu')
else:
    if torch.cuda.device_count() == 1:
        # If you have CUDA_VISIBLE_DEVICES set, which you should,
        # then this will prevent leftover flag arguments from
        # messing with the device allocation.
        args.device = 'cuda:0'

    args.device = torch.device(args.device)

if args.eval_device is None:
    args.eval_device = torch.device('cpu')

# This is needed for memory pinning using a DataLoader (otherwise memory is pinned to cuda:0 by default)
torch.cuda.set_device(args.device)

if args.seed is not None:
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    os.environ['PYTHONHASHSEED'] = str(args.seed)
    # torch.set_deterministic(True)           # reproducibility for torch==1.7
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.enabled = False

def seed_worker(worker_id):                 # https://pytorch.org/docs/stable/notes/randomness.html
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

def main():
    # Load hyperparameters from json
    if args.load_dir and os.path.exists(args.load_dir):
        print(f"Loaded hyperparams of the pretrained model from {args.load_dir}")
        with open(os.path.join(args.load_dir, 'config.json'), 'r') as config_json:
            hyperparams = json.load(config_json)
    else:
        if not os.path.exists(args.conf):
            raise RuntimeError('Config json not found!')

        print(f"Load hyperparams config from {args.conf}")
        with open(args.conf, 'r', encoding='utf-8') as conf_json:
            hyperparams = json.load(conf_json)

        # Add hyperparams from arguments
        hyperparams['dynamic_edges'] = args.dynamic_edges
        hyperparams['edge_state_combine_method'] = args.edge_state_combine_method
        hyperparams['edge_influence_combine_method'] = args.edge_influence_combine_method
        hyperparams['edge_addition_filter'] = args.edge_addition_filter
        hyperparams['edge_removal_filter'] = args.edge_removal_filter
        hyperparams['batch_size'] = args.batch_size
        hyperparams['k_eval'] = args.k_eval
        hyperparams['offline_scene_graph'] = args.offline_scene_graph
        hyperparams['incl_robot_node'] = args.incl_robot_node
        hyperparams['node_freq_mult_train'] = args.node_freq_mult_train
        hyperparams['node_freq_mult_eval'] = args.node_freq_mult_eval
        hyperparams['scene_freq_mult_train'] = args.scene_freq_mult_train
        hyperparams['scene_freq_mult_eval'] = args.scene_freq_mult_eval
        hyperparams['scene_freq_mult_viz'] = args.scene_freq_mult_viz
        hyperparams['edge_encoding'] = not args.no_edge_encoding
        hyperparams['use_map_encoding'] = args.map_encoding
        hyperparams['augment'] = args.augment
        hyperparams['override_attention_radius'] = args.override_attention_radius

    print('-----------------------')
    print('| TRAINING PARAMETERS |')
    print('-----------------------')
    print('| batch_size: %d' % args.batch_size)
    print('| device: %s' % args.device)
    print('| eval_device: %s' % args.eval_device)
    print('| Offline Scene Graph Calculation: %s' % args.offline_scene_graph)
    print('| EE state_combine_method: %s' % args.edge_state_combine_method)
    print('| EIE scheme: %s' % args.edge_influence_combine_method)
    print('| dynamic_edges: %s' % args.dynamic_edges)
    print('| robot node: %s' % args.incl_robot_node)
    print('| edge_addition_filter: %s' % args.edge_addition_filter)
    print('| edge_removal_filter: %s' % args.edge_removal_filter)
    print('| MHL: %s' % hyperparams['minimum_history_length'])
    print('| PH: %s' % hyperparams['prediction_horizon'])
    print('| Weight: %s' % args.contrastive_weight)
    # print('| Weight: %s' % args.contrastive_temperature)
    # print('| Weight: %s' % args.contrastive_dim)
    print('-----------------------')

    log_writer = None
    model_dir = None
    if not args.debug:
        if args.load_dir and os.path.exists(args.load_dir):
            model_dir = args.load_dir
        else:
            # Create the log and model directiory if they're not present.
            model_dir = os.path.join(args.log_dir,
                                     'snce' + args.log_tag)
            pathlib.Path(model_dir).mkdir(parents=True, exist_ok=True)

            # Save config to model directory
            with open(os.path.join(model_dir, 'config.json'), 'w') as conf_json:
                json.dump(hyperparams, conf_json)

        log_writer = SummaryWriter(log_dir=model_dir)

    # Load training and evaluation environments and scenes
    train_scenes = []
    train_data_path = os.path.join(args.data_dir, args.train_data_dict)
    with open(train_data_path, 'rb') as f:
        train_env = dill.load(f, encoding='latin1')

    for attention_radius_override in args.override_attention_radius:
        node_type1, node_type2, attention_radius = attention_radius_override.split(' ')
        train_env.attention_radius[(node_type1, node_type2)] = float(attention_radius)

    if train_env.robot_type is None and hyperparams['incl_robot_node']:
        train_env.robot_type = train_env.NodeType[0]  # TODO: Make more general, allow the user to specify?
        for scene in train_env.scenes:
            scene.add_robot_from_nodes(train_env.robot_type)

    train_scenes = train_env.scenes
    train_scenes_sample_probs = train_env.scenes_freq_mult_prop if args.scene_freq_mult_train else None

    hyperparams['contrastive_weight'] = args.contrastive_weight

    train_dataset = EnvironmentDataset(train_env,
                                       hyperparams['state'],
                                       hyperparams['pred_state'],
                                       scene_freq_mult=hyperparams['scene_freq_mult_train'],
                                       node_freq_mult=hyperparams['node_freq_mult_train'],
                                       hyperparams=hyperparams,
                                       min_history_timesteps=hyperparams['minimum_history_length'],
                                       min_future_timesteps=hyperparams['prediction_horizon'],
                                       return_robot=not args.incl_robot_node)
    train_data_loader = dict()
    for node_type_data_set in train_dataset:
        if len(node_type_data_set) == 0:
            continue

        node_type_dataloader = utils.data.DataLoader(node_type_data_set,
                                                     collate_fn=collate,
                                                     pin_memory=False if args.device == 'cpu' else True,
                                                     batch_size=args.batch_size,
                                                     shuffle=True,
                                                     num_workers=args.preprocess_workers,
                                                     worker_init_fn=seed_worker)
        train_data_loader[node_type_data_set.node_type] = node_type_dataloader

    print(f"Loaded training data from {train_data_path}")

    eval_scenes = []
    eval_scenes_sample_probs = None
    if args.eval_every is not None:
        eval_data_path = os.path.join(args.data_dir, args.eval_data_dict)
        with open(eval_data_path, 'rb') as f:
            eval_env = dill.load(f, encoding='latin1')

        for attention_radius_override in args.override_attention_radius:
            node_type1, node_type2, attention_radius = attention_radius_override.split(' ')
            eval_env.attention_radius[(node_type1, node_type2)] = float(attention_radius)

        if eval_env.robot_type is None and hyperparams['incl_robot_node']:
            eval_env.robot_type = eval_env.NodeType[0]  # TODO: Make more general, allow the user to specify?
            for scene in eval_env.scenes:
                scene.add_robot_from_nodes(eval_env.robot_type)

        eval_scenes = eval_env.scenes
        eval_scenes_sample_probs = eval_env.scenes_freq_mult_prop if args.scene_freq_mult_eval else None

        eval_dataset = EnvironmentDataset(eval_env,
                                          hyperparams['state'],
                                          hyperparams['pred_state'],
                                          scene_freq_mult=hyperparams['scene_freq_mult_eval'],
                                          node_freq_mult=hyperparams['node_freq_mult_eval'],
                                          hyperparams=hyperparams,
                                          min_history_timesteps=hyperparams['minimum_history_length'],
                                          min_future_timesteps=hyperparams['prediction_horizon'],
                                          return_robot=not args.incl_robot_node)
        eval_data_loader = dict()
        for node_type_data_set in eval_dataset:
            if len(node_type_data_set) == 0:
                continue

            node_type_dataloader = utils.data.DataLoader(node_type_data_set,
                                                         collate_fn=collate,
                                                         pin_memory=False if args.eval_device == 'cpu' else True,
                                                         batch_size=args.eval_batch_size,
                                                         shuffle=True,
                                                         num_workers=args.preprocess_workers)
            eval_data_loader[node_type_data_set.node_type] = node_type_dataloader

        print(f"Loaded evaluation data from {eval_data_path}")

    # Offline Calculate Scene Graph
    if hyperparams['offline_scene_graph'] == 'yes':
        print(f"Offline calculating scene graphs")
        for i, scene in enumerate(train_scenes):
            scene.calculate_scene_graph(train_env.attention_radius,
                                        hyperparams['edge_addition_filter'],
                                        hyperparams['edge_removal_filter'])
            print(f"Created Scene Graph for Training Scene {i}")

        for i, scene in enumerate(eval_scenes):
            scene.calculate_scene_graph(eval_env.attention_radius,
                                        hyperparams['edge_addition_filter'],
                                        hyperparams['edge_removal_filter'])
            print(f"Created Scene Graph for Evaluation Scene {i}")

    model_registrar = ModelRegistrar(model_dir, args.device)

    #################################
    #               NCE             #
    #################################

    print('contrastive_sampling:', args.contrastive_sampling)
    head_projection = ProjHead(feat_dim=128, hidden_dim=32, head_dim=8).to(args.device)
    encoder_sample = EventEncoder(hidden_dim=8, head_dim=8).to(args.device)
    snce = SocialNCE(head_projection=head_projection, encoder_sample=encoder_sample, sampling=args.contrastive_sampling)

    #################################

    if args.load_dir and os.path.exists(args.load_dir):
        print(f"Loaded the pretrained model from {args.load_dir}")
        model_registrar.load_models()
        if not args.pretrain:
            head_file = os.path.join(model_dir, 'head_projection.pt')
            if os.path.exists(head_file):
                info = head_projection.load_state_dict(torch.load(head_file, map_location=args.device), strict=True)
                print('Loaded head_projection:\n', info)
            else:
                print('Skipped head_projection')
            encoder_file = os.path.join(model_dir, 'encoder_sample.pt')
            if os.path.exists(encoder_file):
                info = encoder_sample.load_state_dict(torch.load(encoder_file, map_location=args.device), strict=True)
                print('Loaded encoder_sample:\n', info)
            else:
                print('Skipped encoder_sample')
        print('===================\n')

    #################################

    trajectron = Trajectron(model_registrar,
                            hyperparams,
                            log_writer,
                            args.device)

    trajectron.set_environment(train_env, snce)
    trajectron.set_annealing_params()
    print('Created Training Model.')

    eval_trajectron = None
    if args.eval_every is not None or args.vis_every is not None:
        eval_trajectron = Trajectron(model_registrar,
                                     hyperparams,
                                     log_writer,
                                     args.eval_device)
        eval_trajectron.set_environment(eval_env)
        eval_trajectron.set_annealing_params()
    print('Created Evaluation Model.')

    optimizer = dict()
    lr_scheduler = dict()
    for node_type in train_env.NodeType:
        if node_type not in hyperparams['pred_state']:
            continue

        if args.pretrain:
            print('>>> Pretrain SSL')
            param_optim = [{'params': head_projection.parameters()},
                            {'params': encoder_sample.parameters()}]
        else:
            print('>>> Joint train')
            if args.contrastive_weight > 0:
                param_optim = [{'params': model_registrar.get_all_but_name_match('map_encoder').parameters()},
                        {'params': head_projection.parameters()},
                        {'params': encoder_sample.parameters()},
                        {'params': model_registrar.get_name_match('map_encoder').parameters(), 'lr':0.0008}]
            else:
                param_optim = [{'params': model_registrar.get_all_but_name_match('map_encoder').parameters()},
                        {'params': model_registrar.get_name_match('map_encoder').parameters(), 'lr':0.0008}]
        optimizer[node_type] = optim.Adam(param_optim, lr=hyperparams['learning_rate'])

        # Set Learning Rate
        if hyperparams['learning_rate_style'] == 'const':
            lr_scheduler[node_type] = optim.lr_scheduler.ExponentialLR(optimizer[node_type], gamma=1.0)
        elif hyperparams['learning_rate_style'] == 'exp':
            lr_scheduler[node_type] = optim.lr_scheduler.ExponentialLR(optimizer[node_type],
                                                                       gamma=hyperparams['learning_decay_rate'])

    #################################
    #           TRAINING            #
    #################################
    curr_iter_node_type = {node_type: 0 for node_type in train_data_loader.keys()}
    for epoch in range(1, args.train_epochs + 1):
        model_registrar.to(args.device)
        train_dataset.augment = args.augment
        for node_type, data_loader in train_data_loader.items():
            curr_iter = curr_iter_node_type[node_type]
            pbar = tqdm(data_loader, ncols=80)
            for batch in pbar:
                trajectron.set_curr_iter(curr_iter)
                trajectron.step_annealers(node_type)
                optimizer[node_type].zero_grad()
                train_loss, loss_task, loss_nce = trajectron.train_loss(batch, node_type)
                # pbar.set_description(f"Epoch {epoch}, {node_type} L: {train_loss.item():.2f}")
                pbar.set_description(f"Epoch {epoch}, A: {train_loss.item():.4f}, T: {loss_task.item():.4f}, S: {loss_nce.item():.4f}")
                train_loss.backward()
                # Clipping gradients.
                if hyperparams['grad_clip'] is not None:
                    nn.utils.clip_grad_value_(model_registrar.parameters(), hyperparams['grad_clip'])
                optimizer[node_type].step()

                # Stepping forward the learning rate scheduler and annealers.
                lr_scheduler[node_type].step()

                if not args.debug:
                    log_writer.add_scalar(f"{node_type}/train/learning_rate",
                                          lr_scheduler[node_type].get_last_lr()[0],
                                          curr_iter)
                    log_writer.add_scalar(f"{node_type}/train/loss", train_loss, curr_iter)

                curr_iter += 1
            curr_iter_node_type[node_type] = curr_iter
        train_dataset.augment = False
        if args.eval_every is not None or args.vis_every is not None:
            eval_trajectron.set_curr_iter(epoch)

        #################################
        #        VISUALIZATION          #
        #################################
        if args.vis_every is not None and not args.debug and epoch % args.vis_every == 0 and epoch > 0:
            max_hl = hyperparams['maximum_history_length']
            ph = hyperparams['prediction_horizon']
            with torch.no_grad():
                # Predict random timestep to plot for train data set
                if args.scene_freq_mult_viz:
                    scene = np.random.choice(train_scenes, p=train_scenes_sample_probs)
                else:
                    scene = np.random.choice(train_scenes)
                timestep = scene.sample_timesteps(1, min_future_timesteps=ph)
                predictions = trajectron.predict(scene,
                                                 timestep,
                                                 ph,
                                                 min_future_timesteps=ph,
                                                 z_mode=True,
                                                 gmm_mode=True,
                                                 all_z_sep=False,
                                                 full_dist=False)

                # Plot predicted timestep for random scene
                fig, ax = plt.subplots(figsize=(10, 10))
                visualization.visualize_prediction(ax,
                                                   predictions,
                                                   scene.dt,
                                                   max_hl=max_hl,
                                                   ph=ph,
                                                   map=scene.map['VISUALIZATION'] if scene.map is not None else None)
                ax.set_title(f"{scene.name}-t: {timestep}")
                log_writer.add_figure('train/prediction', fig, epoch)

                model_registrar.to(args.eval_device)
                # Predict random timestep to plot for eval data set
                if args.scene_freq_mult_viz:
                    scene = np.random.choice(eval_scenes, p=eval_scenes_sample_probs)
                else:
                    scene = np.random.choice(eval_scenes)
                timestep = scene.sample_timesteps(1, min_future_timesteps=ph)
                predictions = eval_trajectron.predict(scene,
                                                      timestep,
                                                      ph,
                                                      num_samples=20,
                                                      min_future_timesteps=ph,
                                                      z_mode=False,
                                                      full_dist=False)

                # Plot predicted timestep for random scene
                fig, ax = plt.subplots(figsize=(10, 10))
                visualization.visualize_prediction(ax,
                                                   predictions,
                                                   scene.dt,
                                                   max_hl=max_hl,
                                                   ph=ph,
                                                   map=scene.map['VISUALIZATION'] if scene.map is not None else None)
                ax.set_title(f"{scene.name}-t: {timestep}")
                log_writer.add_figure('eval/prediction', fig, epoch)

                # Predict random timestep to plot for eval data set
                predictions = eval_trajectron.predict(scene,
                                                      timestep,
                                                      ph,
                                                      min_future_timesteps=ph,
                                                      z_mode=True,
                                                      gmm_mode=True,
                                                      all_z_sep=True,
                                                      full_dist=False)

                # Plot predicted timestep for random scene
                fig, ax = plt.subplots(figsize=(10, 10))
                visualization.visualize_prediction(ax,
                                                   predictions,
                                                   scene.dt,
                                                   max_hl=max_hl,
                                                   ph=ph,
                                                   map=scene.map['VISUALIZATION'] if scene.map is not None else None)
                ax.set_title(f"{scene.name}-t: {timestep}")
                log_writer.add_figure('eval/prediction_all_z', fig, epoch)

        #################################
        #           EVALUATION          #
        #################################
        if args.eval_every is not None and not args.debug and epoch % args.eval_every == 0 and epoch > 0:
            max_hl = hyperparams['maximum_history_length']
            ph = hyperparams['prediction_horizon']
            model_registrar.to(args.eval_device)
            with torch.no_grad():
                # Calculate evaluation loss
                for node_type, data_loader in eval_data_loader.items():
                    eval_loss = []
                    print(f"Starting Evaluation @ epoch {epoch} for node type: {node_type}")
                    pbar = tqdm(data_loader, ncols=80)
                    for batch in pbar:
                        eval_loss_node_type = eval_trajectron.eval_loss(batch, node_type)
                        pbar.set_description(f"Epoch {epoch}, {node_type} L: {eval_loss_node_type.item():.2f}")
                        eval_loss.append({node_type: {'nll': [eval_loss_node_type]}})
                        del batch

                    evaluation.log_batch_errors(eval_loss,
                                                log_writer,
                                                f"{node_type}/eval_loss",
                                                epoch)

                # Predict batch timesteps for evaluation dataset evaluation
                eval_batch_errors = []
                for scene in tqdm(eval_scenes, desc='Sample Evaluation', ncols=80):
                    timesteps = scene.sample_timesteps(args.eval_batch_size)

                    predictions = eval_trajectron.predict(scene,
                                                          timesteps,
                                                          ph,
                                                          num_samples=50,
                                                          min_future_timesteps=ph,
                                                          full_dist=False)

                    eval_batch_errors.append(evaluation.compute_batch_statistics(predictions,
                                                                                scene.dt,
                                                                                max_hl=max_hl,
                                                                                ph=ph,
                                                                                node_type_enum=eval_env.NodeType,
                                                                                map=scene.map))

                evaluation.log_batch_errors(eval_batch_errors,
                                            log_writer,
                                            'eval',
                                            epoch,
                                            bar_plot=['kde'],
                                            box_plot=['ade', 'fde'])

                # Predict maximum likelihood batch timesteps for evaluation dataset evaluation
                eval_batch_errors_ml = []
                for scene in tqdm(eval_scenes, desc='MM Evaluation', ncols=80):
                    timesteps = scene.sample_timesteps(scene.timesteps)

                    predictions = eval_trajectron.predict(scene,
                                                          timesteps,
                                                          ph,
                                                          num_samples=1,
                                                          min_future_timesteps=ph,
                                                          z_mode=True,
                                                          gmm_mode=True,
                                                          full_dist=False)

                    eval_batch_errors_ml.append(evaluation.compute_batch_statistics(predictions,
                                                                                scene.dt,
                                                                                max_hl=max_hl,
                                                                                ph=ph,
                                                                                map=scene.map,
                                                                                node_type_enum=eval_env.NodeType,
                                                                                kde=False))

                evaluation.log_batch_errors(eval_batch_errors_ml,
                                            log_writer,
                                            'eval/ml',
                                            epoch)

        if args.save_every is not None and args.debug is False and epoch % args.save_every == 0:
            if not args.pretrain:
            #     torch.save(head_projection.state_dict(), os.path.join(model_dir, 'head_projection-%d.pt' % epoch))
            #     torch.save(encoder_sample.state_dict(), os.path.join(model_dir, 'encoder_sample-%d.pt' % epoch))
            # else:
                model_registrar.save_models(epoch, 'w-{:.4f}-s-{:d}'.format(args.contrastive_weight, args.seed))

    if args.pretrain:
        torch.save(head_projection.state_dict(), os.path.join(model_dir, 'head_projection.pt'))
        torch.save(encoder_sample.state_dict(), os.path.join(model_dir, 'encoder_sample.pt'))

if __name__ == '__main__':
    main()


import time
import json
import os

from src.utils import *
from src.notify import send_message

import numpy as np
from PIL import Image
import nxbt

TURTWIG = 0
CHIMCHAR = 1
PIPLUP = 2

def detect_shiny_starter(img_fn, timeout=17, framerate=30, timing_threshold=11.55):
    hp_bar_time = time.time()
    timeout_start = time.time()
    while time.time() - timeout_start < timeout:
        t = time.time()

        img = img_fn()

        # Get the opponents health bar image region
        height, width, _ = img.shape
        xl = int(width * 0.90)
        xr = int(width * 0.96)
        yb = int(height * 0.13)
        yt = int(height * 0.115)
        hp_bar = img[yt:yb, xl:xr, :]

        # Get the mean of the hp bar
        red_mean = hp_bar[:,:,0:1].mean()
        green_mean = hp_bar[:,:,1:2].mean()
        blue_mean = hp_bar[:,:,2:3].mean()

        # Get the diff of the hp_bar's channels
        # from the expected hp_bar colour
        red_diff = abs(red_mean - 100)
        green_diff = abs(green_mean - 200)
        blue_diff = abs(blue_mean - 100)

        # Check diffs
        if red_diff < 35 and green_diff < 35 and blue_diff < 35:
            appearance_time = time.time() - hp_bar_time
            print("Appearance time", appearance_time)
            if appearance_time > timing_threshold:
                print("Shiny detected")
                return True
            break

        # print(red_mean, green_mean, blue_mean, red_diff, green_diff, blue_diff)

        t = time.time() - t
        sleep_time = 1/framerate - t
        if sleep_time > 0:
            time.sleep(sleep_time)

    return False

def reset_hunt():
    stats = None
    if os.path.isfile("stats.json"):
        with open("stats.json", "r") as f:
            stats = json.load(f)
    else:
        stats = {
            "reset_count": 0,
            "issues": 0,
            "log": []
        }

    pokemon = TURTWIG

    starter_select_check = np.array(Image.open("./check-imgs/starter-select-check.png"))
    yes_select_check = np.array(Image.open("./check-imgs/yes-select-check.png"))
    hand_turtwig = np.array(Image.open("./check-imgs/hand-turtwig.png"))
    hand_chimchar = np.array(Image.open("./check-imgs/hand-chimchar.png"))
    hand_piplup = np.array(Image.open("./check-imgs/hand-piplup.png"))

    # Initialize an emulated controller and connect
    # to an available Switch.
    nx = nxbt.Nxbt()
    controller_index = nx.create_controller(
        nxbt.PRO_CONTROLLER,
        reconnect_address=nx.get_switch_addresses())
    nx.wait_for_connection(controller_index)

    add_to_stat_log(stats, "Controller Connected")

    # Activate controller and wait for Switch to recognize it
    nx.macro(controller_index, "B 0.1s\n 0.1s")
    nx.macro(controller_index, "B 0.1s\n 0.1s")
    nx.macro(controller_index, "B 0.1s\n 0.1s")

    while True:
        add_to_stat_log(stats, f"Reset Count {stats['reset_count']}")

        # Assume we are on the Home screen with BDSP selected
        press_button(nx, controller_index, "A")
        time.sleep(1.3)
        press_button(nx, controller_index, "A")
        time.sleep(0.5)
        press_button(nx, controller_index, "A")
        time.sleep(0.5)
        press_button(nx, controller_index, "A")
        time.sleep(2)

        # Wait for the game to boot
        game_loading = False
        game_loaded = False
        t = time.time()
        while game_loaded is not True:
            img = get_image()

            loading_mean = np.mean(relative_crop(img, 0.0, 0.025, 1.0, 0.975))

            if time.time() - t > 80:
                add_to_stat_log(stats, "Something went wrong. Game load > 80s.")
                exit_and_reset(nx, controller_index, stats)

            # Check if the game crashed
            # Get independent means of r, g, and b
            rgb_means = np.mean(relative_crop(img, 0.0, 0.025, 1.0, 0.975), (0,1))
            if abs(rgb_means[0] - 55) < 2 and abs(rgb_means[1] - 55) < 2 and abs(rgb_means[2] - 55) < 2:
                add_to_stat_log(stats, "Game crashed...")
                game_loading = False
                game_loaded = False
                press_button(nx, controller_index, "A")
                time.sleep(2)

            # Check if the game is loading
            if game_loading is False and loading_mean < 18:
                add_to_stat_log(stats, "Game is loading...")
                game_loading = True

            # Check if we've exited the loader
            if game_loading and loading_mean >= 18:
                add_to_stat_log(stats, "Game is loaded")
                game_loaded = True
            # If we haven't, keep pressing A
            else:
                press_button(nx, controller_index, "A")

        # Wait a bit to make sure the game has fully loaded in
        time.sleep(2)

        # Move up once on the dpad
        add_to_stat_log(stats, "Moving up...")
        tilt_stick(nx, controller_index, "L_STICK", 0, 100)

        # Wait for rival to talk
        time.sleep(1)

        # Skip dialogue and enter
        press_button(nx, controller_index, "A")

        # Wait for scene to load in
        time.sleep(3)

        # Mash A until we see a black screen
        # which signals that we're choosing a pokemon
        add_to_stat_log(stats, "Skipping dialogue...")
        t = time.time()
        select_loading = False
        select_loaded = False
        while select_loaded is not True:
            img = get_image()
            select_mean = np.mean(img)

            if time.time() - t > 70:
                add_to_stat_log(stats, "Something went wrong. Dialogue > 70s.")
                exit_and_reset(nx, controller_index, stats)

            # Check if the game is loading
            if select_loading is False and select_mean < 18:
                add_to_stat_log(stats, "Entering selection screen...")
                select_loading = True

            # Check if we've exited the loader
            if select_loading and select_mean >= 18:
                add_to_stat_log(stats, "Entered selection screen")
                select_loaded = True
            # If we haven't, keep pressing A
            else:
                press_button(nx, controller_index, "A")

        # Check that we are on the starter selection screen
        time.sleep(5)
        img = relative_crop(get_image(), 0.7, 0.75, 0.95, 0.92)
        img_mse = mse(starter_select_check, img)
        add_to_stat_log(stats, f"Starter MSE: {img_mse}")
        if img_mse > 15:
            add_to_stat_log(stats, "Something went wrong. Not on starter selection.")
            exit_and_reset(nx, controller_index, stats)
            continue
        add_to_stat_log(stats, "Passed selection screen check")

        # Pressing A until we're on the yes/no starter select prompt
        t = time.time()
        add_to_stat_log(stats, "Pressing A until starter select hand appears....")
        yes_select_present = False
        while True:
            mses = []
            for i in range(50):
                img_area = relative_crop(get_image(), 0.325, 0.35, 0.35, 0.325)
                mses.append(mse(hand_turtwig, img_area))

            img_mse = np.min(mses)
            add_to_stat_log(stats, f"Hand Select MSE: {img_mse}")

            if img_mse < 30:
                yes_select_present = True
                break

            press_button(nx, controller_index, "A", duration=2.0)

            if time.time() - t > 30:
                break

        if yes_select_present is False:
            add_to_stat_log(stats, "Something went wrong. Hand not seen in time.")
            exit_and_reset(nx, controller_index, stats)
            continue

        # Pressing A until we're on the yes/no starter select prompt
        t = time.time()
        add_to_stat_log(stats, "Moving right until hand selects the chosen starter...")
        yes_select_present = False
        while True:
            mses = []
            for i in range(50):
                img = get_image()
                height, width, _ = img.shape
                if pokemon == CHIMCHAR:
                    xl = 0.46
                    xr = 0.48
                    yb = 0.5
                    yt = 0.46
                elif pokemon == PIPLUP:
                    xl = 0.65
                    xr = 0.67
                    yb = 0.35
                    yt = 0.325
                else:
                    xl = 0.325
                    xr = 0.35
                    yb = 0.35
                    yt = 0.325
                img_area = relative_crop(img, xl, xr, yb, yt)

                if pokemon == CHIMCHAR:
                    mses.append(mse(hand_chimchar, img_area))
                elif pokemon == PIPLUP:
                    mses.append(mse(hand_piplup, img_area))
                else:
                    mses.append(mse(hand_turtwig, img_area))

            img_mse = np.min(mses)
            add_to_stat_log(stats, f"Hand Select MSE: {img_mse}")

            if img_mse < 30:
                yes_select_present = True
                break

            press_button(nx, controller_index, "DPAD_RIGHT")
            time.sleep(0.75)

            if time.time() - t > 30:
                break

        if yes_select_present is False:
            add_to_stat_log(stats, "Something went wrong. Hand not selected in time.")
            exit_and_reset(nx, controller_index, stats)
            continue

        # Pressing A until we're on the yes/no starter select prompt
        t = time.time()
        add_to_stat_log(stats, "Pressing A until starter yes/no selection...")
        yes_select_present = False
        while True:
            img = relative_crop(get_image(), 0.8, 0.825, 0.75, 0.675)
            img_mse = mse(yes_select_check, img)
            add_to_stat_log(stats, f"Yes Select MSE: {img_mse}")

            if img_mse < 15:
                yes_select_present = True
                break

            press_button(nx, controller_index, "A", duration=2.0)

            if time.time() - t > 30:
                break

        if yes_select_present is False:
            add_to_stat_log(stats, "Something went wrong. Starter not selected in time.")
            exit_and_reset(nx, controller_index, stats)
            continue

        # Select yes and check the selection
        add_to_stat_log(stats, "Starter present. Selecting...")
        for i in range(3):
            tilt_stick(nx, controller_index, "L_STICK", 0, 100, duration=0.2)
        time.sleep(1)
        press_button(nx, controller_index, "A")
        press_button(nx, controller_index, "A")
        add_to_stat_log(stats, "Starter selected")

        wait_for_battle(get_image)
        add_to_stat_log(stats, "Battle detected")

        t = time.time()
        result = detect_shiny_starter(get_image)
        add_to_stat_log(stats, f"HP Bar Time: {time.time() - t}")

        if result:
            add_to_stat_log(stats, "Shiny detected")
            return True
        else:
            add_to_stat_log(stats, "Shiny not detected")

        # Shiny not detected. Resetting.
        press_button(nx, controller_index, "HOME")
        time.sleep(2.5)

        # Checking if we're on the home menu
        hm_check = relative_crop(get_image(), 0.0, 0.05, 0.8, 0.7)
        hm_mean = hm_check.transpose(2,0,1).mean(axis=(1,2))

        add_to_stat_log(stats, "Checking if we're on the home menu...")
        if abs(hm_mean[0] - 55) < 2 and abs(hm_mean[1] - 55) < 2 and abs(hm_mean[2] - 55) < 2:
            add_to_stat_log(stats, "On home menu")
        else:
            add_to_stat_log(stats, "Not on home menu. Attempting to exit to home menu...")
            on_home_menu = False
            while on_home_menu is False:
                press_button(nx, controller_index, "B")
                time.sleep(0.5)
                press_button(nx, controller_index, "B")
                time.sleep(0.5)
                press_button(nx, controller_index, "HOME")
                time.sleep(1)

                # Checking if we're on the home menu
                img = get_image()
                height, width, _ = img.shape
                xl = int(width * 0.0)
                xr = int(width * 0.05)
                yb = int(height * 0.8)
                yt = int(height * 0.7)
                hm_check = img[yt:yb, xl:xr, :]
                hm_mean = hm_check.transpose(2,0,1).mean(axis=(1,2))

                if abs(hm_mean[0] - 55) < 2 and abs(hm_mean[1] - 55) < 2 and abs(hm_mean[2] - 55) < 2:
                    add_to_stat_log(stats, "On home menu")
                    on_home_menu = True

        press_button(nx, controller_index, "X")
        time.sleep(0.75)
        press_button(nx, controller_index, "A")
        time.sleep(3)

        stats['reset_count'] += 1
        with open("stats.json", "w") as f:
            json.dump(stats, f)

if __name__ == "__main__":
    reset_hunt()
    send_message("Found a shiny")


__description__ = \
"""
Wrap Gblocks for interfacing with phylopandas.
"""
__author__ = "Michael J. Harms (harmsm@gmail.com)"
__date__ = "2019-08-16"

import phylopandas as phy

import pandas as pd
import numpy as np

import math, sys, subprocess, string, random, re, os

def _qual_arg(user_value,
              python_arg_name,
              gblock_arg_name,
              allowable):
    """
    Construct and sanity check a qualitative argument to
    send to gblocks.

    user_value: value to try to send to gblocks
    python_arg_name: name of python argument (for error string)
    gblock_arg_name: name of argument in gblocks
    allowable: dictionary of allowable values mapping python to
               whatever should be jammed into gblocks
    """

    if user_value in allowable.keys():
        return "-{}={}".format(gblock_arg_name,allowable[user_value])
    else:
        err = "\n\n{} '{}' not recognized\n".format(python_arg_name,
                                                user_value)
        err += "must be one of:\n"

        allowed = list(allowable)
        allowed.sort()
        for a in allowed:
            err += "    {}\n".format(a)

        raise ValueError(err)

def _quant_arg(user_value,
               python_arg_name,
               gblock_arg_name,
               min_allowed,
               max_allowed):
    """
    Construct and sanity check a quantitative argument to
    send to gblocks.

    user_value: value to try to send to gblocks
    python_arg_name: name of python argument (for error string)
    gblock_arg_name: name of argument in gblocks
    min_allowed: minimum allowable value
    max_allowed: maximum allowable value
    """

    if user_value < min_allowed or user_value > max_allowed:
        err = "\n\nvalue of {} invalid ('{}').\n".format(python_arg_name,
                                                     user_value)
        err += "must be between {} and {}.\n".format(min_allowed,
                                                     max_allowed)
        raise ValueError(err)

    return "-{}={}".format(gblock_arg_name,int(round(user_value,0)))

def run(df,
        sequence_type="protein",
        min_conserved=None,
        min_flank=None,
        max_contig_noncon=8,
        min_block_length=10,
        allowed_gap="no",
        use_sim_matrix=True,
        min_initial_block=None,
        keep_tmp=False):
    """

    Run Gblocks on a phylopandas dataframe.

    Full docs for software are here:
    http://molevol.cmima.csic.es/castresana/Gblocks/Gblocks_documentation.html#Installation

    As of this writing (August 2019), Gblocks can be installed using
    bioconda (https://bioconda.github.io/).

    sequence_type: type of sequence
                   default: protein
                   allowed: ("protein","dna","codon")
                   gblock flag: -t
    min_conserved: minimum number of sequences for conserved position
                   default: num_seqs*0.5 + 1
                   allowed: (num_seqs*0.5 + 1,num_seqs)
                   gblock flag: -b1
    min_flank: minimum number of sequences for a flank position
                   default: 0.85*num_seqs
                   allowed: (min_conserved,num_seqs)
                   gblock flag: -b2
    max_contig_noncon: maximum number of contiguous non-conserved positions
                   default: 8
                   allowed: any integer > 0
                   gblock flag: -b3
    min_block_length: minimum length of a block
                   default: 10
                   allowed: any integer > 2
                   gblock flag: -b4
    allowed_gap: allowed gap positions
                   default: "no"
                   allowed: ("no","half","all")
                   gblock flag: -b5
    use_sim_matrix: use similarity matrix
                   default: True
                   allowed: True,False (True only allowed for proteins)
                   gblock flag: -b6
    min_initial_block: minimum length of an initial block
                   default: min_block_length
                   allowed: any integer > 2
                   gblock flag: -b0
    keep_tmp: keep temporary files and write out reports
                   default: False
                   allowed: False,True
    """

    tmp_file_root = "".join([random.choice(string.ascii_letters)
                             for i in range(5)])

    input_file = "{}.fasta".format(tmp_file_root)
    output_file = "{}.fasta-gb".format(tmp_file_root)
    output_summary = "{}.fasta-gb.htm".format(tmp_file_root)
    gblocks_stdout = "{}.stdout".format(tmp_file_root)
    gblocks_stderr = "{}.stderr".format(tmp_file_root)

    phy.seqio.write.to_fasta(df,id_col="uid",filename=input_file)

    cmd = ["Gblocks",input_file]

    num_seqs = len(df.uid)

    # Parse sequence type
    a = _qual_arg(sequence_type,
                  "sequence_type",
                  "t",
                  {"protein":"p","dna":"d","codon":"c"})
    cmd.append(a)

    # Parse min conserved
    min_allowed = int(math.ceil(0.5*num_seqs) + 1)
    max_allowed = num_seqs
    if min_conserved is None:
        min_conserved = min_allowed
    a = _quant_arg(min_conserved,"min_conserved","b1",
                   min_allowed,max_allowed)
    cmd.append(a)

    # Parse min flank
    min_allowed = min_conserved
    max_allowed = sys.maxsize
    if min_flank is None:
        min_flank = int(math.ceil(0.85*num_seqs))
    a = _quant_arg(min_flank,"min_flank","b2",
                   min_allowed,max_allowed)
    cmd.append(a)

    # Parse max_contig_noncon
    a = _quant_arg(max_contig_noncon,
                   "max_contig_noncon","b3",
                   1,sys.maxsize)
    cmd.append(a)

    # Parse min_block_length
    a = _quant_arg(min_block_length,
                   "min_block_length","b4",
                   2,sys.maxsize)
    cmd.append(a)

    # Parse allowed_gap
    a = _qual_arg(allowed_gap,
                  "allowed_gap",
                  "b5",
                  {"no":"n","half":"h","all":"a"})
    cmd.append(a)

    # parse use_sim_matrix
    use_sim_matrix = bool(use_sim_matrix)
    if use_sim_matrix and sequence_type != "protein":
        use_sim_matrix = False
    a = _qual_arg(use_sim_matrix,
                  "use_sim_matrix",
                  "b6",
                  {True:"y",False:"n"})
    cmd.append(a)

    # parse min_initial_block
    if min_initial_block is None:
        min_initial_block = min_block_length
    a = _quant_arg(min_initial_block,
                   "min_initial_block","b0",
                   2,sys.maxsize)
    cmd.append(a)

    if keep_tmp:
        cmd.append("-d=y")
    else:
        cmd.append("-d=n")

    # Run gblocks
    stdout = open(gblocks_stdout,"w")
    stderr = open(gblocks_stderr,"w")

    try:
        output = subprocess.run(cmd,
                                stdout=stdout,
                                stderr=stderr)
    except FileNotFoundError:
        err = "Gblocks does not appear to be in your path\n"
        raise RuntimeError(err)

    stdout.close()
    stderr.close()

    # As far as I can tell, Gblocks *always* returns 1.  At least on my
    # MacBook (Mojave 10.14.6, bioconda build.)
    if output.returncode != 1:

        err = "Gblocks failed for some reason\n"

        err += "\n\n------stdout-----\n\n"
        err += "".join(open(gblocks_stdout).read())
        err += "\n\n------stderr-----\n\n"
        err += "".join(open(gblocks_stderr).read())

        raise RuntimeError(err)

    # Parse output
    new_seqs = phy.read_fasta(output_file)
    new_seqs = pd.DataFrame({"uid":new_seqs.id,
                             "sequence":new_seqs.sequence})

    num_old_columns = len(df.loc[0].sequence)
    num_new_columns = len(new_seqs.loc[0].sequence)
    if num_new_columns == 0:
        err = "\n\nGblocks removed all columns from alignment.\n"
        err += "Trying building a better alignment, editing it manually,\n"
        err += "or altering one or more Gblocks settings to be less aggressive.\n"
        err += "See help string on this function for details.\n\n"

        raise RuntimeError(err)

    print("Gblocks reduced the number of columns from {} to {}.\n".format(num_old_columns,
                                                                          num_new_columns))

    # Drop the sequence information from the original frame
    to_merge = df.copy()
    to_merge = to_merge.drop(labels=["sequence"],axis=1)

    # Merge sequence information back in, this time aligned.
    output = pd.merge(to_merge, new_seqs, on=['uid'], how='left')

    # Make sure no na got introduced
    if np.sum(pd.isnull(output.sequence)) > 0:
        err = "an unknown error caused sequences to be lost during alignment.\n"
        err += "temporary files {} and {} saved.\n".format(input_file,output_file)
        raise RuntimeError(err)

    # Remove temporary files
    if not keep_tmp:
        os.remove(input_file)
        os.remove(output_file)
        os.remove(output_summary)
        os.remove(gblocks_stdout)
        os.remove(gblocks_stderr)

    return output


from tensorflow.keras.models import model_from_json
import numpy as np
import tensorflow as tf

class FacialExpressionModel(object):

    EMOTIONS_LIST = ["Angry", "Disgust", "Fear", "Happy", "Neutral", "Sad", "Surprised"]

    def __init__(self, model_json_file, model_weights_file):
        with open(model_json_file,"r") as json_file:
            loaded_model_json = json_file.read()
            self.loaded_model = model_from_json(loaded_model_json)

        self.loaded_model.load_weights(model_weights_file)
        self.loaded_model.make_predict_function()

    def predict_emotion(self, img):
        self.preds = self.loaded_model.predict(img)
        return FacialExpressionModel.EMOTIONS_LIST[np.argmax(self.preds)]


import math
import numpy as np
import rdkit
from rdkit import RDLogger

from MolEvol.common.parse_args import args
from MolEvol.common import update_args, get_properties
from MolEvol.common.eval import eval_all_v4
from MolEvol.graph_completion_model import *
from MolEvol.graph_completion_model.finetune import score_func

import torch

def evaluate_molecules(g_list, s_list=None):
    sg_pairs = []
    if s_list is None:
        for g in g_list:
            sg_pairs.append((None, g))
    else:
        for g, s in zip(g_list, s_list):
            sg_pairs.append((s, g))

    print('Get Properties...')
    if os.path.exists('{}/sgs_tuples_weighted'.format(args.iter_dir)):
        sgs_tuples = torch.load('{}/sgs_tuples_weighted'.format(args.iter_dir))
    else:
        sgs_tuples = get_properties(sg_pairs)
        sgs_tuples = sorted(sgs_tuples, key=lambda x: -score_func((x[2], x[3], x[4], x[5])))
        torch.save(sgs_tuples, '{}/sgs_tuples_weighted'.format(args.iter_dir))
    print('Get Statistics...')

    ans_str = '{} Our Metric {}\n'.format('#' * 10, '#' * 10)
    a, b, c, d, e = eval_all_v4(sgs_tuples, args.ref_path)
    ans_str += 'Success rate: %f\n' % a
    ans_str += 'Novelty: %f\n' % b
    ans_str += 'Diversity: %f\n' % c
    ans_str += 'GU: %f\n' % d
    ans_str += 'GNU: %f\n' % e

    for topk in [args.eval_num // 2]:
        tmp_res = sgs_tuples[:topk]
        _, _, x, y, qed, sa = zip(*tmp_res)
        scr = [score_func((x_i, y_i, qed_i, sa_i)) for x_i, y_i, qed_i, sa_i in zip(x, y, qed, sa)]

        ans_str += '{} Top {} Stats {}\n'.format('=' * 12, topk, '=' * 12)
        ans_str += 'Mean of (x, y, qed, sa, score): {:.4f}, {:.4f}, {:.4f}, {:.4f}, {:.4f}\n'.format(
            sum(x) / topk, sum(y) / topk, sum(qed) / topk, sum(sa) / topk, sum(scr) / topk
        )

    print(ans_str)
    return ans_str

def eval_molevol(rationale_all_dist):
    eval_dict = {k: (math.exp(v[0]), v[1]) for k, v in
                 sorted(rationale_all_dist.items(), key=lambda item: -item[1][0])[:300]}
    total_score = sum([v[0] for k, v in eval_dict.items()])
    eval_dict = {k: (round(v[0] / total_score * args.eval_num), v[1]) for k, v in eval_dict.items()}
    print(args.iter_dir)
    g_list = []
    s_list = []
    for k, v in eval_dict.items():
        decode_len = len(v[1])
        s_list += [k] * v[0]
        if decode_len > v[0]:
            g_list += random.sample(v[1], k=v[0])
        else:
            g_list += v[1] + random.choices(v[1], k=v[0] - decode_len)
    ans_str = evaluate_molecules(g_list, s_list)

    output = '{}/result_overall_weighted.txt'.format(args.iter_dir)
    with open(output, 'w') as f:
        f.write(ans_str)

def main(args):
    update_args(args, iter_num=0)
    for la in range(1, args.last_round + 1):
        update_args(args, iter_num=la)
        rationale_all_dist = torch.load('{}/rationale_all_dist'.format(args.iter_dir))
        eval_molevol(rationale_all_dist)

if __name__ == '__main__':
    lg = rdkit.RDLogger.logger()
    lg.setLevel(rdkit.RDLogger.CRITICAL)

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    args.ncpu = 20
    args.novelty_filter = True
    args.disc_ok = True
    args.prop4 = True
    args.local_search = True
    args.l2x_gcn = True
    args.exp_dir = 'experiment/[MolEvol]{}*{}'.format(args.rounds, args.epoch)

    assert args.disc_ok
    args.init_model = 'models/chembl-molgen/model.25'
    print(args)

    os.makedirs(args.exp_dir, exist_ok=True)
    cmd = 'cp -r MolEvol {}'.format(args.exp_dir)
    print("Moving main files to {}...".format(args.exp_dir))
    os.system(cmd)
    with open('{}/config'.format(args.exp_dir), 'w') as g:
        g.write(str(args))

    main(args)


#!/usr/bin/env python

import argparse
import numpy as np
from Bio import Seq
from Bio import SeqIO
from Bio import SeqRecord
ap = argparse.ArgumentParser(description="Compare PSVs derived from the reference to the psv graph")
ap.add_argument("--ref", help="Alignment of reference hits to the assembly")
ap.add_argument("--refFasta", help="Genome file.")
ap.add_argument("--vcf", help="Original vcf psv file.")
ap.add_argument("--psv", help="PSVs, one line per cluster.")
ap.add_argument("--nfti", help="read nucfreq.", default=None)
ap.add_argument("--nfto", help="Write nucfreq of matching positions.", default=None)
ap.add_argument("--writevcf", help="Output file.", default=None)

args = ap.parse_args()

ref = []
refFile = open(args.ref)
prevContig = None
posList = []
vcfFile = open(args.vcf)
vcf = []
vcfLines = []
refFastaFile = open(args.refFasta)
refSeq = SeqIO.read(refFastaFile, "fasta")

for line in vcfFile:
    if line[0] != "#":
        vcf.append(int(line.split()[1]))
        vcfLines.append(line.split())
        vcfLines[-1][7] = "AN=2;AO=NA;RO=NA;AL=NA"
        vcfLines[-1][8] = "GT:GQ"
        vcfLines[-1][9] = "0|1:100"

refCounts = {}
for line in refFile:
    vals = line.split()
    pos = int(vals[1])
    contig = vals[5]
    if pos + 1 not in refCounts:
        refCounts[pos+1] = 0
    refCounts[pos+1] +=1
    if contig != prevContig:
        if prevContig is not None:
            ref.append(posList)
        posList = []
    posList.append(pos+1)
    prevContig = contig
ref.append(posList)

multi = []
for rc in refCounts:
    if refCounts[rc] > 1:
        multi.append(rc)
multi.sort()
print len(multi)
for r in range(0,len(ref)):
    ref[r] = np.setdiff1d(ref[r], multi)

psvFile = open(args.psv)
psv = []

for line in psvFile:
    vals = line.split()
    psv.append([vcf[int(v)] for v in vals])

sims = []
for p in range(0,len(psv)):
    sim = [len(np.intersect1d(psv[p], ref[r])) for r in range(0,len(ref))]
    print str(p) + "\t" + " ".join(["{:3d}".format(d) for d in sim])
    sims.append(sim)

if args.nfti is not None or args.writevcf is not None:
    if args.nfti is not None:
        nfti = open(args.nfti)
        nfto = open(args.nfto,'w')
        nft = [ [int(v) for v in line.split()[1:]] for line in nfti ]

    for p in range(0,len(psv)):
        maxSim = max(sims[p])
        for i in range(0,len(sims[p])):
            if sims[p][i] == maxSim:
                break
        psvMatch = np.intersect1d(psv[p], ref[i])
        f = 0
        m = 0
        psvFilt = []
        nftFilt = []
        while (m < len(psvMatch) and f < len(vcfLines)):
            if psvMatch[m] == int(vcfLines[f][1]):
                psvFilt.append(vcfLines[f])
                if args.nfti is not None:
                    v = nftFilt[f] + [p]
                    nftFilt.append(v)

                m+=1
                f+=1
            elif psvMatch[m] < int(vcfLines[f][1]):
                m+=1
            else:
                f+=1
        if args.writevcf is not None:
            vcfOutName = args.writevcf + "." + str(p) + ".vcf"
            vcfOut = open(vcfOutName, 'w')
            vcfOut.write("""##fileformat=VCFv4.1
##INFO=<ID=AO,Number=A,Type=Integer,Description="Alternate allele observations">
##INFO=<ID=RO,Number=1,Type=Integer,Description="Reference allele observations">
##contig=<ID={},length={}>
#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	sample
""".format(refSeq.id, len(refSeq.seq)))

            vcfOut.write("\n".join(["\t".join(pf) for pf in psvFilt]) + "\n")
            vcfOut.close()
        if args.nfti is not None:
            nfto.write("\n".join(["\t".join(nf for nf in nftFilt)]) + "\n")

    if args.nfti is not None:
        nfto.close()