import numpy as np
import cv2
import imageio
import tensorflow as tf
import json

import os
import sys
import argparse

import object_detection.object_detector as obj
import action_detection.action_detector as act

import time
DISPLAY = False
SHOW_CAMS = False
def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-v', '--video_path', type=str, required=False, default="")
    parser.add_argument('-d', '--display', type=str, required=False, default="True")

    args = parser.parse_args()
    display = (args.display == "True" or args.display == "true")

    #actor_to_display = 6 # for cams

    video_path = args.video_path
    basename = os.path.basename(video_path).split('.')[0]
    out_vid_path = "./output_videos/%s_output.mp4" % (basename if not SHOW_CAMS else basename+'_cams_actor_%.2d' % actor_to_display)
    #out_vid_path = './output_videos/testing.mp4'

    # video_path = "./tests/chase1Person1View3Point0.mp4"
    # out_vid_path = 'output.mp4'

    main_folder = './'

    # NAS

    obj_detection_model =  'ssd_mobilenet_v2_coco_2018_03_29'
    obj_detection_graph = os.path.join("object_detection", "weights", obj_detection_model, "frozen_inference_graph.pb")

    print("Loading object detection model at %s" % obj_detection_graph)

    obj_detector = obj.Object_Detector(obj_detection_graph)
    tracker = obj.Tracker()

    print("Reading video file %s" % video_path)
    reader = imageio.get_reader(video_path, 'ffmpeg')
    action_freq = 8
    # fps_divider = 1
    print('Running actions every %i frame' % action_freq)
    fps = reader.get_meta_data()['fps'] #// fps_divider
    W, H = reader.get_meta_data()['size']
    T = tracker.timesteps
    if not display:
        writer = imageio.get_writer(out_vid_path, fps=fps)
        print("Writing output to %s" % out_vid_path)

    # act_detector = act.Action_Detector('i3d_tail')
    # ckpt_name = 'model_ckpt_RGB_i3d_pooled_tail-4'
    act_detector = act.Action_Detector('soft_attn')
    #ckpt_name = 'model_ckpt_RGB_soft_attn-16'
    #ckpt_name = 'model_ckpt_soft_attn_ava-23'
    ckpt_name = 'model_ckpt_soft_attn_pooled_cosine_drop_ava-130'

    #input_frames, temporal_rois, temporal_roi_batch_indices, cropped_frames = act_detector.crop_tubes_in_tf([T,H,W,3])
    memory_size = act_detector.timesteps - action_freq
    updated_frames, temporal_rois, temporal_roi_batch_indices, cropped_frames = act_detector.crop_tubes_in_tf_with_memory([T,H,W,3], memory_size)

    rois, roi_batch_indices, pred_probs = act_detector.define_inference_with_placeholders_noinput(cropped_frames)

    ckpt_path = os.path.join(main_folder, 'action_detection', 'weights', ckpt_name)
    act_detector.restore_model(ckpt_path)

    prob_dict = {}
    frame_cnt = 0
    for cur_img in reader:
        frame_cnt += 1
        #tracker.add_frame(cur_img)
        print("frame_cnt: %i" %frame_cnt)
        # Object Detection
        expanded_img = np.expand_dims(cur_img, axis=0)
        #expanded_img = np.tile(expanded_img, [10,1,1,1]) # test the speed
        t1 = time.time()
        detection_list = obj_detector.detect_objects_in_np(expanded_img)
        detection_info = [info[0] for info in detection_list]
        t2 = time.time(); print('obj det %.2f seconds' % (t2-t1))
        tracker.update_tracker(detection_info, cur_img)
        t3 = time.time(); print('tracker %.2f seconds' % (t3-t2))
        no_actors = len(tracker.active_actors)

        if tracker.active_actors and frame_cnt % action_freq == 0:
            probs = []

            cur_input_sequence = np.expand_dims(np.stack(tracker.frame_history[-action_freq:], axis=0), axis=0)

            rois_np, temporal_rois_np = tracker.generate_all_rois()
            if no_actors > 14:
                no_actors = 14
                rois_np = rois_np[:14]
                temporal_rois_np = temporal_rois_np[:14]

            #feed_dict = {input_frames:cur_input_sequence,
            feed_dict = {updated_frames:cur_input_sequence, # only update last #action_freq frames
                         temporal_rois: temporal_rois_np,
                         temporal_roi_batch_indices: np.zeros(no_actors),
                         rois:rois_np,
                         roi_batch_indices:np.arange(no_actors)}
            run_dict = {'pred_probs': pred_probs}
            if SHOW_CAMS:
                run_dict['cropped_frames'] = cropped_frames
                #import pdb;pdb.set_trace()
                run_dict['final_i3d_feats'] =  act_detector.act_graph.get_collection('final_i3d_feats')[0]
                #run_dict['cls_weights'] = [var for var in tf.global_variables() if var.name == "CLS_Logits/kernel:0"][0]
                run_dict['cls_weights'] = act_detector.act_graph.get_collection('variables')[-2] # this is the kernel
            #import pdb;pdb.set_trace()
            out_dict = act_detector.session.run(run_dict, feed_dict=feed_dict)
            probs = out_dict['pred_probs']
            # associate probs with actor ids
            print_top_k = 5
            for bb in range(no_actors):
                act_probs = probs[bb]
                order = np.argsort(act_probs)[::-1]
                cur_actor_id = tracker.active_actors[bb]['actor_id']
                print("Person %i" % cur_actor_id)
                cur_results = []
                for pp in range(print_top_k):
                    print('\t %s: %.3f' % (act.ACTION_STRINGS[order[pp]], act_probs[order[pp]]))
                    cur_results.append((act.ACTION_STRINGS[order[pp]], act_probs[order[pp]]))
                prob_dict[cur_actor_id] = cur_results

            t5 = time.time(); print('action %.2f seconds' % (t5-t3))
        # # Action detection
        # no_actors = len(tracker.active_actors)
        # #batch_np = np.zeros([no_actors, act_detector.timesteps] + act_detector.input_size + [3], np.uint8)
        # batch_list = []
        # rois_np = np.zeros([no_actors, 4])
        # batch_indices_np = np.array(range(no_actors))
        # for bb, actor_info in enumerate(tracker.active_actors):
        #     actor_no = actor_info['actor_id']
        #     tube, roi = tracker.crop_person_tube(actor_no)
        #     #batch_np[bb, :] = tube
        #     batch_list.append(tube)
        #     rois_np[bb]= roi
        #t4 = time.time(); print('cropping %.2f seconds' % (t4-t3))

        # if tracker.active_actors:
        #     batch_np = np.stack(batch_list, axis=0)
        #     max_batch_size = 10
        #     prob_list = []
        #     cur_index = 0
        #     while cur_index < no_actors:
        #         cur_batch = batch_np[cur_index:cur_index+max_batch_size]
        #         cur_roi = rois_np[cur_index:cur_index+max_batch_size]
        #         cur_indices = batch_indices_np[cur_index:cur_index+max_batch_size] - cur_index
        #         feed_dict = {input_seq:cur_batch, rois:cur_roi, roi_batch_indices:cur_indices}
        #         #t51 = time.time(); print('action before run %.2f seconds' % (t51-t4))
        #         cur_probs = act_detector.session.run(pred_probs, feed_dict=feed_dict)
        #         #t52 = time.time(); print('action after run %.2f seconds' % (t52-t51))
        #         prob_list.append(cur_probs)
        #         cur_index += max_batch_size
        #     probs = np.concatenate(prob_list, axis=0)

        #t5 = time.time(); print('action %.2f seconds' % (t5-t4))
        # Print top_k probs
        #out_img = visualize_detection_results(cur_img, tracker.active_actors, prob_dict)
        if frame_cnt > 16:
            out_img = visualize_detection_results(tracker.frame_history[-16], tracker.active_actors, prob_dict)
            if SHOW_CAMS:
                if tracker.active_actors:
                    actor_indices = [ii for ii in range(no_actors) if tracker.active_actors[ii]['actor_id'] == actor_to_display]
                    if actor_indices:
                        out_img = visualize_cams(out_img, cur_input_sequence, out_dict, actor_indices[0])
                    else:
                        continue
                else:
                    continue
            if display:
                cv2.imshow('results', out_img[:,:,::-1])
                cv2.waitKey(10)
            else:
                writer.append_data(out_img)

    if not display:
        writer.close()

np.random.seed(10)
COLORS = np.random.randint(0, 255, [1000, 3])
def visualize_detection_results(img_np, active_actors, prob_dict):
    score_th = 0.30
    action_th = 0.20

    # copy the original image first
    disp_img = np.copy(img_np)
    H, W, C = img_np.shape
    #for ii in range(len(active_actors)):
    for ii in range(len(active_actors)):
        cur_actor = active_actors[ii]
        actor_id = cur_actor['actor_id']
        cur_act_results = prob_dict[actor_id] if actor_id in prob_dict else []
        try:
            cur_box, cur_score, cur_class = cur_actor['all_boxes'][-16], cur_actor['all_scores'][0], 1
        except IndexError:
            continue

        if cur_score < score_th:
            continue

        top, left, bottom, right = cur_box

        left = int(W * left)
        right = int(W * right)

        top = int(H * top)
        bottom = int(H * bottom)

        conf = cur_score
        #label = bbox['class_str']
        # label = 'Class_%i' % cur_class
        label = obj.OBJECT_STRINGS[cur_class]['name']
        message = '%s_%i: %% %.2f' % (label, actor_id,conf)
        action_message_list = ["%s:%.3f" % (actres[0][0:7], actres[1]) for actres in cur_act_results if actres[1]>action_th]
        # action_message = " ".join(action_message_list)

        raw_colors = COLORS[actor_id]
        rect_color = tuple(int(raw_color) for raw_color in raw_colors)
        text_color = tuple(255-color_value for color_value in rect_color)

        cv2.rectangle(disp_img, (left,top), (right,bottom), rect_color, 3)

        font_size =  max(0.5,(right - left)/50.0/float(len(message)))
        cv2.rectangle(disp_img, (left, top-int(font_size*40)), (right,top), rect_color, -1)
        cv2.putText(disp_img, message, (left, top-12), 0, font_size, text_color, 1)

        #action message writing
        cv2.rectangle(disp_img, (left, top), (right,top+10*len(action_message_list)), rect_color, -1)
        for aa, action_message in enumerate(action_message_list):
            offset = aa*10
            cv2.putText(disp_img, action_message, (left, top+5+offset), 0, 0.5, text_color, 1)

    return disp_img

def visualize_cams(image, input_frames, out_dict, actor_idx):
    #classes = ["walk", "bend", "carry"]
    #classes = ["sit", "ride"]
    classes = ["talk to", "watch (a", "listen to"]
    action_classes = [cc for cc in range(60) if any([cname in act.ACTION_STRINGS[cc] for cname in classes])]

    feature_activations = out_dict['final_i3d_feats']
    cls_weights = out_dict['cls_weights']
    input_frames = out_dict['cropped_frames'].astype(np.uint8)
    probs = out_dict["pred_probs"]

    class_maps = np.matmul(feature_activations, cls_weights)
    min_val = np.min(class_maps[:,:, :, :, :])
    max_val = np.max(class_maps[:,:, :, :, :]) - min_val

    normalized_cmaps = np.uint8((class_maps-min_val)/max_val * 255.)

    t_feats = feature_activations.shape[1]
    t_input = input_frames.shape[1]
    index_diff = (t_input) // (t_feats+1)

    img_new_height = 400
    img_new_width = int(image.shape[1] / float(image.shape[0]) * img_new_height)
    img_to_show = cv2.resize(image.copy(), (img_new_width,img_new_height))[:,:,::-1]
    #img_to_concat = np.zeros((400, 800, 3), np.uint8)
    img_to_concat = np.zeros((400, 400, 3), np.uint8)
    for cc in range(len(action_classes)):
        cur_cls_idx = action_classes[cc]
        act_str = act.ACTION_STRINGS[action_classes[cc]]
        message = "%s:%%%.2d" % (act_str[:20], 100*probs[actor_idx, cur_cls_idx])
        for tt in range(t_feats):
            cur_cam = normalized_cmaps[actor_idx, tt,:,:, cur_cls_idx]
            cur_frame = input_frames[actor_idx, (tt+1) * index_diff, :,:,::-1]

            resized_cam = cv2.resize(cur_cam, (100,100))
            colored_cam = cv2.applyColorMap(resized_cam, cv2.COLORMAP_JET)

            overlay = cv2.resize(cur_frame.copy(), (100,100))
            overlay = cv2.addWeighted(overlay, 0.5, colored_cam, 0.5, 0)

            img_to_concat[cc*125:cc*125+100, tt*100:(tt+1)*100, :] = overlay
        cv2.putText(img_to_concat, message, (20, 13+100+125*cc), 0, 0.5, (255,255,255), 1)

    final_image = np.concatenate([img_to_show, img_to_concat], axis=1)
    return final_image[:,:,::-1]

if __name__ == '__main__':
    main()


# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed 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.
# ==============================================================================
"""Prioritized Experience Replay (proportional variant) DQN agent class."""

# pylint: disable=g-bad-import-order

from typing import Any, Callable, Mapping, Text

from absl import logging
import dm_env
import jax
import jax.numpy as jnp
import numpy as np
import optax
import rlax

from dqn_zoo import parts
from dqn_zoo import processors
from dqn_zoo import replay as replay_lib

# Batch variant of double_q_learning.
_batch_double_q_learning = jax.vmap(rlax.double_q_learning)

class PrioritizedDqn:
  """Prioritized Experience Replay (proportional variant) DQN agent."""

  def __init__(
      self,
      preprocessor: processors.Processor,
      sample_network_input: jnp.ndarray,
      network: parts.Network,
      optimizer: optax.GradientTransformation,
      transition_accumulator: replay_lib.TransitionAccumulator,
      replay: replay_lib.PrioritizedTransitionReplay,
      batch_size: int,
      exploration_epsilon: Callable[[int], float],
      min_replay_capacity_fraction: float,
      learn_period: int,
      target_network_update_period: int,
      grad_error_bound: float,
      rng_key: parts.PRNGKey,
  ):
    self._preprocessor = preprocessor
    self._replay = replay
    self._transition_accumulator = transition_accumulator
    self._batch_size = batch_size
    self._exploration_epsilon = exploration_epsilon
    self._min_replay_capacity = min_replay_capacity_fraction * replay.capacity
    self._learn_period = learn_period
    self._target_network_update_period = target_network_update_period

    # Initialize network parameters and optimizer.
    self._rng_key, network_rng_key = jax.random.split(rng_key)
    self._online_params = network.init(network_rng_key,
                                       sample_network_input[None, ...])
    self._target_params = self._online_params
    self._opt_state = optimizer.init(self._online_params)

    # Other agent state: last action, frame count, etc.
    self._action = None
    self._frame_t = -1  # Current frame index.
    self._max_seen_priority = 1.

    # Define jitted loss, update, and policy functions here instead of as
    # class methods, to emphasize that these are meant to be pure functions
    # and should not access the agent object's state via `self`.

    def loss_fn(online_params, target_params, transitions, weights, rng_key):
      """Calculates loss given network parameters and transitions."""
      _, *apply_keys = jax.random.split(rng_key, 4)
      q_tm1 = network.apply(online_params, apply_keys[0],
                            transitions.s_tm1).q_values
      q_t = network.apply(online_params, apply_keys[1],
                          transitions.s_t).q_values
      q_target_t = network.apply(target_params, apply_keys[2],
                                 transitions.s_t).q_values
      td_errors = _batch_double_q_learning(
          q_tm1,
          transitions.a_tm1,
          transitions.r_t,
          transitions.discount_t,
          q_target_t,
          q_t,
      )
      td_errors = rlax.clip_gradient(td_errors, -grad_error_bound,
                                     grad_error_bound)
      losses = rlax.l2_loss(td_errors)
      assert losses.shape == (self._batch_size,) == weights.shape
      # This is not the same as using a huber loss and multiplying by weights.
      loss = jnp.mean(losses * weights)
      return loss, td_errors

    def update(rng_key, opt_state, online_params, target_params, transitions,
               weights):
      """Computes learning update from batch of replay transitions."""
      rng_key, update_key = jax.random.split(rng_key)
      d_loss_d_params, td_errors = jax.grad(
          loss_fn, has_aux=True)(online_params, target_params, transitions,
                                 weights, update_key)
      updates, new_opt_state = optimizer.update(d_loss_d_params, opt_state)
      new_online_params = optax.apply_updates(online_params, updates)
      return rng_key, new_opt_state, new_online_params, td_errors

    self._update = jax.jit(update)

    def select_action(rng_key, network_params, s_t, exploration_epsilon):
      """Samples action from eps-greedy policy wrt Q-values at given state."""
      rng_key, apply_key, policy_key = jax.random.split(rng_key, 3)
      q_t = network.apply(network_params, apply_key, s_t[None, ...]).q_values[0]
      a_t = rlax.epsilon_greedy().sample(policy_key, q_t, exploration_epsilon)
      return rng_key, a_t

    self._select_action = jax.jit(select_action)

  def step(self, timestep: dm_env.TimeStep) -> parts.Action:
    """Selects action given timestep and potentially learns."""
    self._frame_t += 1

    timestep = self._preprocessor(timestep)

    if timestep is None:  # Repeat action.
      action = self._action
    else:
      action = self._action = self._act(timestep)

      for transition in self._transition_accumulator.step(timestep, action):
        self._replay.add(transition, priority=self._max_seen_priority)

    if self._replay.size < self._min_replay_capacity:
      return action

    if self._frame_t % self._learn_period == 0:
      self._learn()

    if self._frame_t % self._target_network_update_period == 0:
      self._target_params = self._online_params

    return action

  def reset(self) -> None:
    """Resets the agent's episodic state such as frame stack and action repeat.

    This method should be called at the beginning of every episode.
    """
    self._transition_accumulator.reset()
    processors.reset(self._preprocessor)
    self._action = None

  def _act(self, timestep) -> parts.Action:
    """Selects action given timestep, according to epsilon-greedy policy."""
    s_t = timestep.observation
    self._rng_key, a_t = self._select_action(self._rng_key, self._online_params,
                                             s_t, self.exploration_epsilon)
    return parts.Action(jax.device_get(a_t))

  def _learn(self) -> None:
    """Samples a batch of transitions from replay and learns from it."""
    logging.log_first_n(logging.INFO, 'Begin learning', 1)
    transitions, indices, weights = self._replay.sample(self._batch_size)
    self._rng_key, self._opt_state, self._online_params, td_errors = (
        self._update(
            self._rng_key,
            self._opt_state,
            self._online_params,
            self._target_params,
            transitions,
            weights,
        ))
    assert weights.shape == td_errors.shape
    priorities = jnp.abs(td_errors)
    priorities = jax.device_get(priorities)
    max_priority = priorities.max()
    self._max_seen_priority = np.max([self._max_seen_priority, max_priority])
    self._replay.update_priorities(indices, priorities)

  @property
  def online_params(self) -> parts.NetworkParams:
    """Returns current parameters of Q-network."""
    return self._online_params

  @property
  def exploration_epsilon(self) -> float:
    """Returns epsilon value currently used by (eps-greedy) behavior policy."""
    return self._exploration_epsilon(self._frame_t)

  @property
  def importance_sampling_exponent(self) -> float:
    """Returns current importance sampling exponent of prioritized replay."""
    return self._replay.importance_sampling_exponent

  @property
  def max_seen_priority(self) -> float:
    """Returns maximum seen replay priority up until this time."""
    return self._max_seen_priority

  def get_state(self) -> Mapping[Text, Any]:
    """Retrieves agent state as a dictionary (e.g. for serialization)."""
    state = {
        'rng_key': self._rng_key,
        'frame_t': self._frame_t,
        'opt_state': self._opt_state,
        'online_params': self._online_params,
        'target_params': self._target_params,
        'replay': self._replay.get_state(),
        'max_seen_priority': self._max_seen_priority,
    }
    return state

  def set_state(self, state: Mapping[Text, Any]) -> None:
    """Sets agent state from a (potentially de-serialized) dictionary."""
    self._rng_key = state['rng_key']
    self._frame_t = state['frame_t']
    self._opt_state = jax.device_put(state['opt_state'])
    self._online_params = jax.device_put(state['online_params'])
    self._target_params = jax.device_put(state['target_params'])
    self._replay.set_state(state['replay'])
    self._max_seen_priority = state['max_seen_priority']


#!/usr/bin/env python3
#
# Copyright (C) 2019-2020 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#
# This file is based on https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/python_demos/speech_recognition_demo
#

import wave
import timeit
import argparse
import yaml
import numpy as np
from tqdm import tqdm
from deepspeech_openvino.utils.context_timer import Timer
from deepspeech_openvino.utils.deep_speech_pipeline import DeepSpeechPipeline, PROFILES
from deepspeech_openvino.deepspeech_asr_base import DeepSpeechASR

def preprocess_sound(audio_rev, stt):
    sampling_rate = 16000
    audio = np.frombuffer(audio_rev, dtype=np.dtype("<h"))
    # normalize to -1 to 1, int 16 to float32
    audio = audio / np.float32(32768)
    audio = audio.reshape(-1, 1)
    audio_features = stt.extract_mfcc(audio, sampling_rate=sampling_rate)
    return audio_features

def get_profile(profile_name):
    if profile_name in PROFILES:
        return PROFILES[profile_name]
    with open(profile_name, "rt") as f:
        profile = yaml.safe_load(f)
    return profile

beam_width = 500
max_candidates = 1

class DeepSpeechv8(DeepSpeechASR):
    def __init__(
        self,
        model,
        device,
        sample_rate,
        alphabet_cfg="",
        lang_model="",
        SPEECH_LIB_PATH="",
    ):
        # TODO asserts on model= have .xml xtension, sample rate should be a number,
        self.lm = lang_model
        super().__init__(
            model=model,
            device=device,
            sample_rate=sample_rate,
            alphabet_cfg=alphabet_cfg,
            SPEECH_LIB_PATH=SPEECH_LIB_PATH,
        )

    def _load_model_to_device(self):
        self.profile = get_profile("mds08x_en")
        self.stt = DeepSpeechPipeline(
            model=self.model_xml,
            lm=self.lm,
            beam_width=beam_width,
            max_candidates=max_candidates,
            profile=self.profile,
            device=self.device,
        )

    def push_data(self, buffer):
        audio_features = preprocess_sound(buffer, self.stt)
        # character_probs = self.stt.extract_per_frame_probs(audio_features, wrap_iterator=tqdm)
        character_probs = self.stt.extract_per_frame_probs(audio_features)
        transcription = self.stt.decode_probs(character_probs)
        # print(transcription)
        for candidate in transcription:
            self.text = "{} ".format(candidate["text"])
            # print(self.text)
        self.output_text = self.text

    def _transcribe(self, features):
        pass

    def get_result(self, final, finish_processing):
        # TODO: do something with final
        result = self.output_text
        if finish_processing:
            self.exec_net = ""
            import time

            # self.stt.activate_model(self.stt.default_device)
        return result


from __future__ import division

import os
import numpy as np
from torch.autograd import Variable
from torch.utils import data as data_
import torch
import random
import ipdb

from src.head_detector_vgg16 import Head_Detector_VGG16
from trainer import Head_Detector_Trainer
from src.config import opt
import src.utils as utils
from data.dataset import Dataset, inverse_normalize
import src.array_tool as at
from src.vis_tool import visdom_bbox
from src.bbox_tools import bbox_iou

phases = ['train', 'val', 'test']

data_check_flag = False

def eval(dataloader, head_detector):
    """
    Given the dataloader of the test split compute the
    average corLoc of the dataset using the head detector
    model given as the argument to the function.
    """
    test_img_num = 0
    test_corrLoc = 0.0
    for _, (img, bbox_, scale) in enumerate(dataloader):
        scale = at.scalar(scale)
        img, bbox = img.cuda().float(), bbox_.cuda()
        img, bbox = Variable(img), Variable(bbox)
        pred_bboxes_, _ = head_detector.predict(img, scale, mode='evaluate')
        gt_bboxs = at.tonumpy(bbox_)[0]
        pred_bboxes_ = at.tonumpy(pred_bboxes_)
        if pred_bboxes_.shape[0] == 0:
            test_img_num += 1
            continue
        else:
            ious = bbox_iou(pred_bboxes_, gt_bboxs)
            max_ious = ious.max(axis=1)
            corr_preds = np.where(max_ious >= 0.5)[0]
            num_boxs = gt_bboxs.shape[0]
            num_corr_preds = len(corr_preds)
            test_corrLoc += num_corr_preds / num_boxs
            test_img_num += 1
    return test_corrLoc / test_img_num

def train():
    # Get the dataset
    for phase in phases:
        if phase == 'train':
            train_data_list_path = os.path.join(opt.data_root_path,'brainwash_train.idl')
            train_data_list = utils.get_phase_data_list(train_data_list_path)
        elif phase == 'val':
            val_data_list_path = os.path.join(opt.data_root_path,'brainwash_val.idl')
            val_data_list = utils.get_phase_data_list(val_data_list_path)
        elif phase == 'test':
            test_data_list_path = os.path.join(opt.data_root_path,'brainwash_test.idl')
            test_data_list = utils.get_phase_data_list(test_data_list_path)

    print "Number of images for training: %s" %(len(train_data_list))
    print "Number of images for val: %s" %(len(val_data_list))
    print "Number of images for test: %s" %(len(test_data_list))

    # Just to visualize if the dataset is correctly loaded
    if data_check_flag:
        utils.check_loaded_data(train_data_list[random.randint(1,len(train_data_list))])
        utils.check_loaded_data(val_data_list[random.randint(1,len(val_data_list))])
        utils.check_loaded_data(test_data_list[random.randint(1,len(test_data_list))])

    # Load the train dataset
    dataset = Dataset(train_data_list)

    # Load the test dataset
    test_dataset = Dataset(val_data_list)
    print "Load data."

    dataloader = data_.DataLoader(dataset, batch_size=1,shuffle=True, num_workers=1)
    test_dataloader = data_.DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=1)

    # Initialize the head detector.
    head_detector_vgg16 = Head_Detector_VGG16(ratios=[1], anchor_scales=[8,16])
    print('model construct completed')

    trainer = Head_Detector_Trainer(head_detector_vgg16).cuda()
    best_map = 0.0
    lr_ = opt.lr

    for epoch in range(opt.epoch):
        trainer.reset_meters()
        # train_img_num = 0
        # train_corrLoc = 0.0
        for ii, (img, bbox_, scale) in enumerate(dataloader):
            scale = at.scalar(scale)
            img, bbox = img.cuda().float(), bbox_.cuda()
            img, bbox = Variable(img), Variable(bbox)
            _, _, _ = trainer.train_step(img, bbox, scale)

            if (ii+1) % opt.plot_every == 0:
                # if os.path.exists(opt.debug_file):
                #     ipdb.set_trace()
                trainer.vis.plot_many(trainer.get_meter_data())
                ori_img_ = inverse_normalize(at.tonumpy(img[0]))
                gt_img = visdom_bbox(ori_img_, at.tonumpy(bbox_[0]))
                trainer.vis.img('gt_img', gt_img)
                rois,_ = trainer.head_detector.predict(img, scale=scale, mode='visualize')
                pred_img = visdom_bbox(ori_img_, at.tonumpy(rois))
                trainer.vis.img('pred_img', pred_img)
                trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm')

        # avg_train_CorrLoc = train_corrLoc / train_img_num

        avg_test_CorrLoc = eval(test_dataloader, head_detector_vgg16)

        print("Epoch {} of {}.".format(epoch+1, opt.epoch))
        # print("  training average corrLoc accuracy:\t\t{:.3f}".format(avg_train_CorrLoc))
        print("  test average corrLoc accuracy:\t\t{:.3f}".format(avg_test_CorrLoc))

        model_save_path = trainer.save(best_map=avg_test_CorrLoc)

        # if avg_test_CorrLoc >= best_map:
        #     best_map = avg_test_CorrLoc
        #     best_path = trainer.save(best_map=best_map)
        if epoch == 8:
            trainer.load(model_save_path)
            trainer.head_detector.scale_lr(opt.lr_decay)
            lr_ = lr_ * opt.lr_decay

if __name__ == '__main__':

    train()


from __future__ import print_function
import pickle
from model.toolkits.interactions import close_contacts
from model.toolkits.interactions import hydrophobic_contacts
from model.toolkits.interactions import salt_bridges
import numpy as np
import pandas as pd
# from bitarray import bitarray
try:
    # Open Babel >= 3.0
    from openbabel import openbabel as ob
except ImportError:
    import openbabel as ob
import sys
import os
import argparse
# from time import time
from model.toolkits.parse_conf import parse_config_vina, parse_protein_vina, parse_ligand_vina
# from model.toolkits.parse_conf import parse_config
# from model.toolkits.parse_docking_conf import parse_vina_conf
from model.toolkits.PARAMETERS import HYDROPHOBIC, AROMATIC, HBOND, ELECTROSTATIC, HBOND_ANGLE,\
    AROMATIC_ANGLE_LOW, AROMATIC_ANGLE_HIGH
from model.obbl import Molecule
from model.toolkits.spatial import angle, distance
from model.toolkits.interactions import hbonds, pi_stacking, salt_bridges, \
    hydrophobic_contacts, close_contacts, halogenbonds
from model.toolkits.pocket import pocket_atoms
from model.IFP import cal_Interactions,  get_Molecules, cal_IFP
from pathos.multiprocessing import Pool
from functools import partial
from tqdm.auto import tqdm

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="config",
                        default='config.txt')
    parser.add_argument("--save", help="dataframe name of IFP",
                        default='IFP')
    parser.add_argument("--n_jobs", help="number of threads", type=int,
                        default=10)
    args = parser.parse_args()
    return args

def save_obj(obj, name):
    os.system('mkdir obj')
    with open('obj/' + name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name):
    with open('obj/' + name + '.pkl', 'rb') as f:
        return pickle.load(f)

def walk_folder(path, suffix):
    processed = []
    files = os.listdir(path)
    print(f"{len(files)} files have been detected!")
    outpdbqt = []
    for file in files:
        # print(file)
        if suffix in file:
            outpdbqt.append(file)
    #         base_name = os.path.basename(file)
    #         # print(base_name)
    #         simple_name = base_name.replace('_', '.').split('.')
    #         simple_name = simple_name[0]
    #         processed.append({'simple_name': simple_name,
    #                           'base_name': base_name, 'full_name': file})
    # # print(processed)
    return outpdbqt

def concat_df(dic_main, dic):
    # df_interaction = ''
    # for i in range(len(df_res)):
    #     if i == 0:
    #         df_interaction = df_res[0]
    #     else:
    #         for key in df_interaction.keys():
    #             df_interaction[key] = pd.concat(
    #                 [df_interaction[key], df_res[i][key]])
    # return df_interaction
    for key in dic_main.keys():
        if len(dic_main[key]) == 0:
            dic_main[key] = dic[key]
        else:
            dic_main[key] = pd.concat([dic_main[key], dic[key]])
    return dic_main

def cal_interacions_run(ligand, mol_p, config):
    ''' This funciton will calculate all the intercations between the protein and all the docking poses of the ligand.
    '''
    ligand_name = ligand['simple_name']
    ligand_folder = config['ligand_folder']
    ligand = parse_ligand_vina(os.path.join(
        ligand_folder, ligand['base_name']))
    ligand_poses = ligand['docked_ligands']
    interaction_poses = []
    for ipose in range(len(ligand_poses)):
        mol_l = Molecule(ligand_poses[ipose], protein=False)
        df_res = cal_Interactions(
            mol_p, mol_l, config)
        print(f"\n{ligand_name}-{ipose}\n")
        for key in df_res.keys():
            df_res[key]['Molecule'] = f'{ligand_name}-{ipose}'
        interaction_poses.append(df_res)
    return interaction_poses
    #     df_Interaction = concat_df(df_Interaction, df_res)

def IFP(ligand, config):
    '''The  interaction types of different ligand poses will be transformed into interaction fingerprint.
    '''
    protein = config['protein']
    protein = parse_protein_vina(protein)
    mol_p = Molecule(protein['protein'], protein=True)
    base_name = os.path.basename(ligand)

    # print(base_name)
    simple_name = base_name.replace('_out.pdbqt', '')
    processed = {'simple_name': simple_name,
                 'base_name': base_name, 'full_name': ligand}
    interaction_poses = cal_interacions_run(processed, mol_p, config)
    reference_atom = load_obj('refer_atoms_list')
    reference_res = load_obj('refer_res_list')
    IFP_poses = []
    for ipose in range(len(interaction_poses)):
        if ipose < 5:
            # the IFPs are python list.
            AAIFP, RESIFP = cal_IFP(
                interaction_poses[ipose], reference_atom, reference_res)
            AAIFP = [f'{simple_name}_{ipose}']+AAIFP
            RESIFP = [f'{simple_name}_{ipose}']+RESIFP
            IFP_poses.append([AAIFP, RESIFP])
    return IFP_poses

def main(config, args):
    debug = 0  # if in debug model 0: False  1: True
    df_Interaction = {'df_hbond': '', 'df_halogen': '',
                      'df_elecpair': '', 'df_hydrophobic': '', 'df_pistack': ''}
    protein = config['protein']
    ligand_folder = config['ligand_folder']
    ligands = walk_folder(ligand_folder, '_out.pdbqt')
    protein = parse_protein_vina(protein)
    mol_p = Molecule(protein['protein'], protein=True)

    with Pool(args.n_jobs) as pool:
        IFP_p = partial(
            IFP,  config=config)
        res_list = [x for x in tqdm(
            pool.imap(IFP_p, list(ligands)),
            total=len(ligands),
            miniters=50
        )
            if x is not None]
    AAIFP_full = []
    ResIFP_full = []
    for IFP_poses in res_list:
        for iPose in IFP_poses:
            AAIFP_full.append(iPose[0])
            ResIFP_full.append(iPose[1])

    # df_Interaction = concat_df(df_res)
    reference_atom = load_obj('refer_atoms_list')
    reference_res = load_obj('refer_res_list')
    colname = ['Molecule']
    for iatm in reference_atom:
        for iifp in ['hbd', 'halg', 'elec', 'hrdr', 'pipi']:
            colname.append(f'{iatm}-{iifp}')

    AAIFP_full = pd.DataFrame(
        AAIFP_full, columns=colname)
    colname = ['Molecule']
    for ires in reference_res:
        for iifp in ['hbd', 'halg', 'elec', 'hrdr', 'pipi']:
            colname.append(f'{ires}-{iifp}')
    ResIFP_full = pd.DataFrame(
        ResIFP_full, columns=colname)
    AAIFP_full.to_csv(f'{args.save}_AAIFP.csv', index=None)
    ResIFP_full.to_csv(f'{args.save}_ResIFP.csv', index=None)

if __name__ == "__main__":
    args = get_parser()
    config = parse_config_vina(args.config)
    main(config, args)


import numpy as np
import pickle
import copy
import sys
import argparse
import logging
import os
from logger import get_logger
from tqdm import tqdm
from collections import deque

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.tensorboard import SummaryWriter

from mlc import step_hmlc_K
from mlc_utils import clone_parameters, tocuda, DummyScheduler

from models import *
from meta_models import *

parser = argparse.ArgumentParser(description='MLC Training Framework')
parser.add_argument('--dataset', type=str, choices=['cifar10', 'cifar100', 'clothing1m'], default='cifar10')
parser.add_argument('--method', default='hmlc_K_mix', type=str, choices=['hmlc_K_mix', 'hmlc_K'])
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--data_seed', type=int, default=1)
parser.add_argument('--epochs', '-e', type=int, default=75, help='Number of epochs to train.')
parser.add_argument('--num_iterations', default=100000, type=int)
parser.add_argument('--every', default=100, type=int, help='Eval interval (default: 100 iters)')
parser.add_argument('--bs', default=32, type=int, help='batch size')
parser.add_argument('--test_bs', default=100, type=int, help='batch size')
parser.add_argument('--gold_bs', type=int, default=32)
parser.add_argument('--cls_dim', type=int, default=64, help='Label embedding dim (Default: 64)')
parser.add_argument('--grad_clip', default=0.0, type=float, help='max grad norm (default: 0, no clip)')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum for optimizer')
parser.add_argument('--main_lr', default=3e-4, type=float, help='lr for main net')
parser.add_argument('--meta_lr', default=3e-5, type=float, help='lr for meta net')
parser.add_argument('--optimizer', default='adam', type=str, choices=['adam', 'sgd', 'adadelta'])
parser.add_argument('--opt_eps', default=1e-8, type=float, help='eps for optimizers')
#parser.add_argument('--tau', default=1, type=float, help='tau')
parser.add_argument('--wdecay', default=5e-4, type=float, help='weight decay (default: 5e-4)')

# noise parameters
parser.add_argument('--corruption_type', default='unif', type=str, choices=['unif', 'flip'])
parser.add_argument('--corruption_level', default='-1', type=float, help='Corruption level')
parser.add_argument('--gold_fraction', default='-1', type=float, help='Gold fraction')

parser.add_argument('--sparsemax', default=False, action='store_true', help='Use softmax instead of softmax for meta model (default: False)')
parser.add_argument('--tie', default=False, action='store_true', help='Tie label embedding to the output classifier output embedding of metanet (default: False)')

parser.add_argument('--runid', default='exp', type=str)
parser.add_argument('--queue_size', default=1, type=int, help='Number of iterations before to compute mean loss_g')

############## LOOK-AHEAD GRADIENT STEPS FOR MLC ##################
parser.add_argument('--gradient_steps', default=1, type=int, help='Number of look-ahead gradient steps for meta-gradient (default: 1)')

# CIFAR
# Positional arguments
parser.add_argument('--data_path', default='data', type=str, help='Root for the datasets.')
# Optimization options
parser.add_argument('--nosgdr', default=False, action='store_true', help='Turn off SGDR.')

# Acceleration
parser.add_argument('--prefetch', type=int, default=2, help='Pre-fetching threads.')
# i/o
parser.add_argument('--logdir', type=str, default='runs', help='Log folder.')
parser.add_argument('--local_rank', type=int, default=-1, help='local rank (-1 for local training)')

args = parser.parse_args()

# //////////////////////////// set logging ///////////////////////
filename = '_'.join([args.dataset, args.method, args.corruption_type, args.runid, str(args.epochs), str(args.seed), str(args.data_seed)])
logfile = 'logs/' + filename + '.log'
logger = get_logger(logfile, args.local_rank)
# //////////////////////////////////////////////////////////////////

logger.info(args)
logger.info('CUDA available:' + str(torch.cuda.is_available()))

# cuda set up
torch.cuda.set_device(0) # local GPU

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

# loss function for hard label and soft labels
hard_loss_f = F.cross_entropy
from mlc_utils import soft_cross_entropy as soft_loss_f

# //////////////////////// defining model ////////////////////////

def get_data(dataset, gold_fraction, corruption_prob, get_C):
    if dataset == 'cifar10' or dataset == 'cifar100':
        sys.path.append('CIFAR')

        from data_helper_cifar import prepare_data
        args.use_mwnet_loader = True # use exactly the same loader as in the mwnet paper
        logger.info('================= Use the same dataloader as in MW-Net =========================')
        return prepare_data(gold_fraction, corruption_prob, get_C, args)
    elif dataset == 'clothing1m':
        sys.path.append('CLOTHING1M')

        from data_helper_clothing1m import prepare_data
        return prepare_data(args)

def build_models(dataset, num_classes):
    cls_dim = args.cls_dim # input label embedding dimension

    if dataset in ['cifar10', 'cifar100']:
        from CIFAR.resnet import resnet32

        # main net
        model = resnet32(num_classes)
        main_net = model

        # meta net
        hx_dim = 64 #0 if isinstance(model, WideResNet) else 64 # 64 for resnet-32
        meta_net = MetaNet(hx_dim, cls_dim, 128, num_classes, args)

    elif dataset == 'clothing1m': # use pretrained ResNet-50 model
        model = ResNet50(num_classes)
        main_net = model

        hx_dim = 2048 # from resnet50
        meta_net = MetaNet(2048, cls_dim, 128, num_classes, args)

    main_net = main_net.cuda()
    meta_net = meta_net.cuda()

    logger.info('========== Main model ==========')
    logger.info(model)
    logger.info('========== Meta model ==========')
    logger.info(meta_net)

    return main_net, meta_net

def setup_training(main_net, meta_net, exp_id=None):

    # ============== setting up from scratch ===================
    # set up optimizers and schedulers
    # meta net optimizer
    optimizer = torch.optim.Adam(meta_net.parameters(), lr=args.meta_lr,
                                 weight_decay=0, #args.wdecay, # meta should have wdecay or not??
                                 amsgrad=True, eps=args.opt_eps)
    scheduler = DummyScheduler(optimizer)

    # main net optimizer
    main_params = main_net.parameters()

    if args.optimizer == 'adam':
        main_opt = torch.optim.Adam(main_params, lr=args.main_lr, weight_decay=args.wdecay, amsgrad=True, eps=args.opt_eps)
    elif args.optimizer == 'sgd':
        main_opt = torch.optim.SGD(main_params, lr=args.main_lr, weight_decay=args.wdecay, momentum=args.momentum)

    if args.dataset in ['cifar10', 'cifar100']:
        # follow MW-Net setting
        main_schdlr = torch.optim.lr_scheduler.MultiStepLR(main_opt, milestones=[80,100], gamma=0.1)
    elif args.dataset in ['clothing1m']:
        main_schdlr = torch.optim.lr_scheduler.MultiStepLR(main_opt, milestones=[5], gamma=0.1)
    else:
        main_schdlr = DummyScheduler(main_opt)

    last_epoch = -1

    return main_net, meta_net, main_opt, optimizer, main_schdlr, scheduler, last_epoch

def uniform_mix_C(num_classes, mixing_ratio):
    '''
    returns a linear interpolation of a uniform matrix and an identity matrix
    '''
    return mixing_ratio * np.full((num_classes, num_classes), 1 / num_classes) + \
        (1 - mixing_ratio) * np.eye(num_classes)

def flip_labels_C(num_classes, corruption_prob):
    '''
    returns a matrix with (1 - corruption_prob) on the diagonals, and corruption_prob
    concentrated in only one other entry for each row
    '''
    np.random.seed(args.seed)

    C = np.eye(num_classes) * (1 - corruption_prob)
    row_indices = np.arange(num_classes)
    for i in range(num_classes):
        C[i][np.random.choice(row_indices[row_indices != i])] = corruption_prob
    return C

# //////////////////////// run experiments ////////////////////////
def run():
    corruption_fnctn = uniform_mix_C if args.corruption_type == 'unif' else flip_labels_C
    filename = '_'.join([args.dataset, args.method, args.corruption_type, args.runid, str(args.epochs), str(args.seed), str(args.data_seed)])

    results = {}

    gold_fractions = [0.02]

    if args.gold_fraction != -1:
        assert args.gold_fraction >=0 and args.gold_fraction <=1, 'Wrong gold fraction!'
        gold_fractions = [args.gold_fraction]

    corruption_levels = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]

    if args.corruption_level != -1: # specied one corruption_level
        assert args.corruption_level >= 0 and args.corruption_level <=1, 'Wrong noise level!'
        corruption_levels = [args.corruption_level]

    for gold_fraction in gold_fractions:
        results[gold_fraction] = {}
        for corruption_level in corruption_levels:
            # //////////////////////// load data //////////////////////////////
            # use data_seed her
            gold_loader, silver_loader, valid_loader, test_loader, num_classes = get_data(args.dataset, gold_fraction, corruption_level, corruption_fnctn)

            # //////////////////////// build main_net and meta_net/////////////
            main_net, meta_net = build_models(args.dataset, num_classes)

            # //////////////////////// train and eval model ///////////////////
            exp_id = '_'.join([filename, str(gold_fraction), str(corruption_level)])
            test_acc, baseline_acc = train_and_test(main_net, meta_net, gold_loader, silver_loader, valid_loader, test_loader, exp_id)

            results[gold_fraction][corruption_level] = {}
            results[gold_fraction][corruption_level]['method'] = test_acc
            results[gold_fraction][corruption_level]['baseline'] = baseline_acc
            logger.info(' '.join(['Gold fraction:', str(gold_fraction), '| Corruption level:', str(corruption_level),
                  '| Method acc:', str(results[gold_fraction][corruption_level]['method']),
                                  '| Baseline acc:', str(results[gold_fraction][corruption_level]['baseline'])]))
            logger.info('')

    with open('out/' + filename, 'wb') as file:
        pickle.dump(results, file)
    logger.info("Dumped results_ours in file: " + filename)

def test(main_net, test_loader): # this could be eval or test
    # //////////////////////// evaluate method ////////////////////////
    correct = torch.zeros(1).cuda()
    nsamples = torch.zeros(1).cuda()

    # forward
    main_net.eval()

    for idx, (*data, target) in enumerate(test_loader):
        data, target = tocuda(data), tocuda(target)

        # forward
        with torch.no_grad():
            output = main_net(data)

        # accuracy
        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).sum().item()
        nsamples += len(target)

    test_acc = (correct / nsamples).item()

    # set back to train
    main_net.train()

    return test_acc

####################################################################################################
###  training code
####################################################################################################
def train_and_test(main_net, meta_net, gold_loader, silver_loader, valid_loader, test_loader, exp_id=None):
    writer = SummaryWriter(args.logdir + '/' + exp_id)

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

    main_net, meta_net, main_opt, optimizer, main_schdlr, scheduler, last_epoch = setup_training(main_net, meta_net, exp_id)

    # //////////////////////// switching on training mode ////////////////////////
    meta_net.train()
    main_net.train()

    # set up statistics
    best_params = None
    best_main_opt_sd = None
    best_main_schdlr_sd = None

    best_meta_params = None
    best_meta_opt_sd = None
    best_meta_schdlr_sd = None
    best_val_metric = float('inf')

    val_metric_queue = deque()
    # set done

    args.dw_prev = [0 for param in meta_net.parameters()] # 0 for previous iteration
    args.steps = 0

    silver_train_data = silver_loader.dataset.train_data
    silver_train_labels = silver_loader.dataset.train_labels
    gold_train_data = gold_loader.loader.dataset.train_data
    gold_train_labels = gold_loader.loader.dataset.train_labels
    work_dir = './savedir'

    np.save(os.path.join(work_dir, 'train_data.npy'), silver_train_data)
    np.save(os.path.join(work_dir, 'train_labels.npy'), silver_train_labels)

    np.save(os.path.join(work_dir,'meta_train_data.npy'), gold_train_data)
    np.save(os.path.join(work_dir,'meta_train_labels.npy'), gold_train_labels)

    for epoch in tqdm(range(last_epoch+1, args.epochs)):# change to epoch iteration
        logger.info('Epoch %d:' % epoch)

        for i, (*data_s, target_s) in enumerate(silver_loader):
            *data_g, target_g = next(gold_loader)#.next()

            data_g, target_g = tocuda(data_g), tocuda(target_g)
            data_s, target_s_ = tocuda(data_s), tocuda(target_s)

            # bi-level optimization stage
            eta = main_schdlr.get_lr()[0]
            if args.method == 'hmlc_K':
                loss_g, loss_s = step_hmlc_K(main_net, main_opt, hard_loss_f,
                                             meta_net, optimizer, soft_loss_f,
                                             data_s, target_s_, data_g, target_g,
                                             None, None,
                                             eta, args)
            elif args.method == 'hmlc_K_mix':
                # split the clean set to two, one for training and the other for meta-evaluation
                gbs = int(target_g.size(0) / 2)
                if type(data_g) is list:
                    data_c = [x[gbs:] for x in data_g]
                    data_g = [x[:gbs] for x in data_g]
                else:
                    data_c = data_g[gbs:]
                    data_g = data_g[:gbs]

                target_c = target_g[gbs:]
                target_g = target_g[:gbs]
                loss_g, loss_s = step_hmlc_K(main_net, main_opt, hard_loss_f,
                                             meta_net, optimizer, soft_loss_f,
                                             data_s, target_s_, data_g, target_g,
                                             data_c, target_c,
                                             eta, args)

            args.steps += 1
            if i % args.every == 0:
                writer.add_scalar('train/loss_g', loss_g.item(), args.steps)
                writer.add_scalar('train/loss_s', loss_s.item(), args.steps)

                ''' get entropy of predictions from meta-net '''
                logit_s, x_s_h = main_net(data_s, return_h=True)
                pseudo_target_s = meta_net(x_s_h.detach(), target_s_).detach()
                entropy = -(pseudo_target_s * torch.log(pseudo_target_s+1e-10)).sum(-1).mean()

                writer.add_scalar('train/meta_entropy', entropy.item(), args.steps)

                main_lr = main_schdlr.get_lr()[0]
                meta_lr = scheduler.get_lr()[0]
                writer.add_scalar('train/main_lr', main_lr, args.steps)
                writer.add_scalar('train/meta_lr', meta_lr, args.steps)
                writer.add_scalar('train/gradient_steps', args.gradient_steps, args.steps)

                logger.info('Iteration %d loss_s: %.4f\tloss_g: %.4f\tMeta entropy: %.3f\tMain LR: %.8f\tMeta LR: %.8f' %( i, loss_s.item(), loss_g.item(), entropy.item(), main_lr, meta_lr))

        # PER EPOCH PROCESSING

        # lr scheduler
        main_schdlr.step()
        #scheduler.step()

        # evaluation on validation set
        val_acc = test(main_net, valid_loader)
        test_acc = test(main_net, test_loader)

        logger.info('Val acc: %.4f\tTest acc: %.4f' % (val_acc, test_acc))
        if args.local_rank <=0: # single GPU or GPU 0
            writer.add_scalar('train/val_acc', val_acc, epoch)
            writer.add_scalar('test/test_acc', test_acc, epoch)

        if len(val_metric_queue) == args.queue_size:  # keep at most this number of records
            # remove the oldest record
            val_metric_queue.popleft()

        val_metric_queue.append(-val_acc)

        avg_val_metric = sum(list(val_metric_queue)) / len(val_metric_queue)
        if avg_val_metric < best_val_metric:
            best_val_metric = avg_val_metric

            best_params = copy.deepcopy(main_net.state_dict())

            best_main_opt_sd = copy.deepcopy(main_opt.state_dict())
            best_main_schdlr_sd = copy.deepcopy(main_schdlr.state_dict())

            best_meta_params = copy.deepcopy(meta_net.state_dict())
            best_meta_opt_sd = copy.deepcopy(optimizer.state_dict())
            best_meta_schdlr_sd = copy.deepcopy(scheduler.state_dict())

            # dump best to file also
            ####################### save best models so far ###################

            logger.info('Saving best models...')
            torch.save({
                'epoch': epoch,
                'val_metric': best_val_metric,
                'main_net': best_params,
                'main_opt': best_main_opt_sd,
                'main_schdlr': best_main_schdlr_sd,
                'meta_net': best_meta_params,
                'meta_opt': best_meta_opt_sd,
                'meta_schdlr': best_meta_schdlr_sd
            }, 'models/%s_best.pth' % exp_id)

        writer.add_scalar('train/val_acc_best', -best_val_metric, epoch) # write current best val_acc to tensorboard

    # //////////////////////// evaluating method ////////////////////////
    main_net.load_state_dict(best_params)
    test_acc = test(main_net, test_loader) # evaluate best params picked from validation

    writer.add_scalar('test/acc', test_acc, args.steps) # this test_acc should be roughly the best as it's taken from the best iteration
    logger.info('Test acc: %.4f' % test_acc)

    return test_acc, 0

if __name__ == '__main__':
    run()


import numpy as np
import pandas as pd
from Bio import AlignIO, SeqIO, Seq
from pathlib import Path
from dataclasses import dataclass
from typing import List
cwd = Path.cwd()  # Edit paths as appropriate

# Load reference
ref = SeqIO.read(cwd / 'Wuhan-Hu-1_NC_045512.2.fasta', format='fasta')

# Load genome annotations
meta = pd.read_csv(cwd / 'wuhan-hu-1_genome_annotations_V2.csv')
meta = meta.drop('Unnamed: 8', axis=1)

def get_protein(nuc_start, nuc_end):
    seq = ref.seq[nuc_start - 1:nuc_end]   # nuc_start, nuc_end are one-based indices
    protein_sequence = seq.translate()
    return protein_sequence

# Convert gene annotation positions to per position dataframe
final_df = pd.DataFrame()

# For parsing ORF1ab codons
orf1ab = meta.loc[np.logical_and(meta.region == 'ORF1ab', meta.region_type == 'coding'), 'protein_name']
orf1b = ['NSP12a (part 2)', 'NSP13', 'NSP14', 'NSP15a', 'NSP16']
orf1a = list(set(orf1ab) - set(orf1b))
orf1ab_length = 0
orf1a_length = 0
orf1b_length = 4401

for i in range(len(meta)):
    row = meta.iloc[i:i + 1, :]
    genome_start = row.genome_start.iloc[0]
    genome_end = row.genome_end.iloc[0]
    region_type = row.region_type.iloc[0]
    region = row.region.iloc[0]
    nucleotide_length = row.nucleotide_length.iloc[0]
    protein_name = row.protein_name.iloc[0]
    exploded_element = pd.concat([row] * row.nucleotide_length.iloc[0], axis=0, ignore_index=True)

    # Add nucleotide positions
    exploded_element.insert(loc=3,
                            column='nucleotide_pos',
                            value=np.arange(genome_start, genome_end + 1))

    # Add reference nucleotide
    exploded_element.insert(loc=3,
                            column='ref_nuc',
                            value=ref[genome_start - 1:genome_end])  # convert to zero-based to index sequence

    if region_type == 'coding':
        protein_length = int(row.protein_length.iloc[0])

        # Add position of nucleotide in codon
        exploded_element.insert(loc=3,
                                column='pos_in_codon',
                                value=[1, 2, 3] * protein_length)

        # Add amino acid
        protein_sequence = get_protein(genome_start, genome_end)

        exploded_element.insert(loc=3,
                                column='ref_AA',
                                value=[AA for AA in list(protein_sequence) for _ in range(3)])

    elif region_type == 'non-coding':
        for col in ['pos_in_codon', 'codon_number', 'codon_from_gene_start', 'ref_AA']:
            exploded_element.insert(loc=3,
                                    column=col,
                                    value=-1)
    else:
        raise TypeError('Did you add a new region_type?')

    # Non-continuous codon number from start of each protein
    if region_type == 'coding':
        if protein_name == 'NSP12a (part 2)' and region == 'ORF1ab':
            exploded_element.insert(loc=3,
                                    column='codon_number',
                                    value=[pos for pos in np.arange(protein_length) + 1 + 9 for _ in range(3)])
        else:
            exploded_element.insert(loc=3,
                                    column='codon_number',
                                    value=[pos for pos in np.arange(protein_length) + 1 for _ in range(3)])
    else:
        print(region_type, region, protein_name)

    # Add codon position with ORF1ab in continuous codon sequence accounting for -1 frameshift
    codon_start = 0

    if region_type == 'coding' and region == 'ORF1ab':
        # ORF1a
        if protein_name in orf1a:
            exploded_element.insert(loc=3,
                                    column='codon_from_gene_start',
                                    value=[pos for pos in np.arange(protein_length) + 1 + orf1a_length
                                           for _ in range(3)])

            if protein_name != "NSP11":
                orf1a_length += protein_length

        # Account for -1 frameshift
        elif protein_name in orf1b:
            exploded_element.insert(loc=3,
                                    column='codon_from_gene_start',
                                    value=[pos for pos in np.arange(protein_length) + 1 + orf1b_length
                                           for _ in range(3)])

            orf1b_length += protein_length

        else:
            raise RuntimeError('error in ORF1ab continuous codon number!')

    elif region_type == 'coding' and region != 'ORF1ab':
        exploded_element.insert(loc=3,
                                column='codon_from_gene_start',
                                value=[pos for pos in np.arange(protein_length) + 1
                                       for _ in range(3)])

    # Make dataframes for each possible nuc substitution
    df_list = [exploded_element.copy() for _ in range(4)]
    nuc_list = ['A', 'T', 'G', 'C']

    for df, nuc in zip(df_list, nuc_list):
        df.insert(loc=8, column='var_nuc', value=nuc)
        df.insert(loc=5, column='var_AA', value=-1)

        # Iterate through each codon to get variant AA
        if region_type == 'coding':
            assert len(df) % 3 == 0
            num_codons = len(df) // 3
            idx = 0

            for _ in range(num_codons):
                ori_seq = df.loc[idx:idx + 2, 'ref_nuc'].copy()
                for j in range(3):
                    var_seq = ori_seq.copy().reset_index(drop=True)
                    var_seq.loc[j] = df.loc[idx + j, 'var_nuc']
                    var_codon = Seq.Seq(var_seq.str.cat())
                    var_AA = var_codon.translate()
                    df.loc[idx + j, 'var_AA'] = var_AA

                idx += 3

        exploded_2_element = pd.concat(df_list, axis=0)

    # Append
    final_df = final_df.append(exploded_2_element)

# Filter off entries where ref and var nucleotides are the same
# final_df = final_df.loc[final_df.ref_nuc != final_df.var_nuc, :]

# Add mutation type
final_df.loc[:, 'mutation_type'] = -1
final_df.loc[final_df.loc[:, 'ref_AA'] == final_df.loc[:, 'var_AA'], 'mutation_type'] = 'S'
final_df.loc[final_df.loc[:, 'ref_AA'] != final_df.loc[:, 'var_AA'], 'mutation_type'] = 'NS'

# Save
final_df.to_csv(cwd / 'SARS-CoV-2_hookup_table_V3.csv', index=False)


import os
import numpy as np

from express.parsers import BaseParser
from express.parsers.apps.espresso import settings
from express.parsers.mixins.ionic import IonicDataMixin
from express.parsers.mixins.reciprocal import ReciprocalDataMixin
from express.parsers.mixins.electronic import ElectronicDataMixin
from express.parsers.utils import find_file, lattice_basis_to_poscar
from express.parsers.apps.espresso.settings import NEB_PATH_FILE_SUFFIX
from express.parsers.apps.espresso.formats.txt import EspressoTXTParser
from express.parsers.apps.espresso.formats.xml import EspressoXMLParser

class EspressoParser(BaseParser, IonicDataMixin, ElectronicDataMixin, ReciprocalDataMixin):
    """
    Espresso parser class.
    """

    def __init__(self, *args, **kwargs):
        super(EspressoParser, self).__init__(*args, **kwargs)
        self.work_dir = self.kwargs["work_dir"]
        self.stdout_file = self.kwargs["stdout_file"]
        self.txt_parser = EspressoTXTParser(self.work_dir)
        self.xml_parser = EspressoXMLParser(self.find_xml_file())

    def find_xml_file(self):
        """
        Finds XML file.

        Note: Sternheimer GW XML file of the first process (gw0) is returned if this a Sternheimer GW calculation.

        Returns:
            str
        """
        is_sternheimer_gw = self._is_sternheimer_gw_calculation()
        for root, dirs, files in os.walk(self.work_dir, followlinks=True):
            for file_ in [f for f in files if settings.XML_DATA_FILE == f]:
                file_path = os.path.join(root, file_)
                if not is_sternheimer_gw or (is_sternheimer_gw and settings.STERNHEIMER_GW0_DIR_PATTERN in file_path):
                    return file_path

    def total_energy(self):
        """
        Returns total energy.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.total_energy
        """
        return self.txt_parser.total_energy(self._get_file_content(self.stdout_file))

    def fermi_energy(self):
        """
        Returns fermi energy.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.fermi_energy
        """
        return self.xml_parser.fermi_energy()

    def nspins(self):
        """
        Returns the number of spins.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.nspins
        """
        return self.xml_parser.nspins()

    def _is_sternheimer_gw_calculation(self):
        """
        Checks whether this is a Sternheimer GW calculation.

        The calculation is considered Sternheimer GW if "SternheimerGW" is written on top of the stdout file.

        NOTE: DO NOT READ THE WHOLE FILE INTO MEMORY AS IT IS BIG.

        Returns:
             bool
        """
        if os.path.exists(self.stdout_file):
            with open(self.stdout_file, "r") as f:
                for index, line in enumerate(f):
                    if index > 50: return False
                    if settings.STERNHEIMER_GW_TITLE in line:
                        return True

    def eigenvalues_at_kpoints(self):
        """
        Returns eigenvalues for all kpoints.

        NOTE: eigenvalues are extracted from Sternheimer GW stdout file if this is a Sternheimer GW calculation.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.eigenvalues_at_kpoints
        """
        if self._is_sternheimer_gw_calculation():
            text = self._get_file_content(self.stdout_file)
            inverse_reciprocal_lattice_vectors = self.xml_parser.get_inverse_reciprocal_lattice_vectors()
            return self.txt_parser.eigenvalues_at_kpoints_from_sternheimer_gw_stdout(text,
                                                                                inverse_reciprocal_lattice_vectors)
        else:
            return self.xml_parser.eigenvalues_at_kpoints()

    def ibz_k_points(self):
        """
        Returns ibz_k_points.

        Note:
            The function assumes that kpoints extracted from parsed source are inside the irreducible wedge of the
            Brillouin zone. Without checking whether it is the case or not.

        Reference:
            func: express.parsers.mixins.reciprocal.ReciprocalDataMixin.ibz_k_points
        """
        return np.array([eigenvalueData["kpoint"] for eigenvalueData in self.eigenvalues_at_kpoints()])

    def dos(self):
        """
        Returns density of states.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.dos
        """
        return self.txt_parser.dos()

    def initial_basis(self):
        """
        Returns initial basis.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.initial_basis
        """
        return self.txt_parser.initial_basis(self._get_file_content(self.stdout_file))

    def initial_lattice_vectors(self):
        """
        Returns initial lattice vectors.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.initial_lattice_vectors
        """
        return self.txt_parser.initial_lattice_vectors(self._get_file_content(self.stdout_file))

    def final_basis(self):
        """
        Returns final basis.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.basis
        """
        return self.xml_parser.final_basis()

    def final_lattice_vectors(self):
        """
        Returns final lattice.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.lattice_vectors
        """
        return self.xml_parser.final_lattice_vectors(reciprocal=False)

    def convergence_electronic(self):
        """
        Extracts convergence electronic.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.convergence_electronic
        """
        return self.txt_parser.convergence_electronic(self._get_file_content(self.stdout_file))

    def convergence_ionic(self):
        """
        Returns convergence ionic.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.convergence_ionic
        """
        return self.txt_parser.convergence_ionic(self._get_file_content(self.stdout_file))

    def stress_tensor(self):
        """
        Returns stress tensor.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.stress_tensor
        """
        return self.txt_parser.stress_tensor(self._get_file_content(self.stdout_file))

    def pressure(self):
        """
        Returns pressure.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.pressure
        """
        return self.txt_parser.pressure(self._get_file_content(self.stdout_file))

    def total_force(self):
        """
        Returns total force.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.total_force
        """
        return self.txt_parser.total_force(self._get_file_content(self.stdout_file))

    def atomic_forces(self):
        """
        Returns forces that is exerted on each atom by its surroundings.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.atomic_forces
        """
        return self.txt_parser.atomic_forces(self._get_file_content(self.stdout_file))

    def total_energy_contributions(self):
        """
        Extracts total energy contributions.

        Reference:
            func: express.parsers.mixins.electronic.ElectronicDataMixin.total_energy_contributions
        """
        return self.txt_parser.total_energy_contributions(self._get_file_content(self.stdout_file))

    def zero_point_energy(self):
        """
        Returns zero point energy.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.zero_point_energy
        """
        return self.txt_parser.zero_point_energy(self._get_file_content(self.stdout_file))

    def phonon_dos(self):
        """
        Returns phonon dos.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.phonon_dos
        """
        return self.txt_parser.phonon_dos()

    def phonon_dispersions(self):
        """
        Returns phonon dispersions.

        Reference:
            func: express.parsers.mixins.ionic.IonicDataMixin.phonon_dispersions
        """
        return self.txt_parser.phonon_dispersions()

    def _find_neb_dat_file(self):
        neb_path_file = find_file(NEB_PATH_FILE_SUFFIX, self.work_dir)
        if neb_path_file: return "{}.dat".format(neb_path_file[:neb_path_file.rfind(".")])

    def reaction_coordinates(self):
        """
        Returns reaction coordinates.

        Returns:
             list
        """
        neb_dat_file = self._find_neb_dat_file()
        return self.txt_parser.reaction_coordinates(self._get_file_content(neb_dat_file))

    def reaction_energies(self):
        """
        Returns reaction energies.

        Returns:
             list
        """
        neb_dat_file = self._find_neb_dat_file()
        return self.txt_parser.reaction_energies(self._get_file_content(neb_dat_file))

    def _get_esm_file(self):
        return find_file(".esm1", self.work_dir)

    def potential_profile(self):
        return self.txt_parser.potential_profile(self._get_file_content(self._get_esm_file()))

    def charge_density_profile(self):
        return self.txt_parser.charge_density_profile(self._get_file_content(self._get_esm_file()))

    def _is_pw_scf_output_file(self, path):
        """
        Checks whether the given file is PWSCF output file.

        The file is considered PWSCF output file if "Program PWSCF" is written on top of the file.

        NOTE: DO NOT READ THE WHOLE FILE INTO MEMORY AS IT COULD BE BIG.

        Returns:
             bool
        """
        if os.path.exists(path):
            with open(path, "r") as f:
                for index, line in enumerate(f):
                    if index > 50: return False
                    if settings.PWSCF_OUTPUT_FILE_REGEX in line:
                        return True

    def _find_pw_scf_output_files(self):
        pw_scf_output_files = []
        for root, dirs, files in os.walk(self.work_dir, followlinks=True):
            for file in files:
                path = os.path.join(root, file)
                if self._is_pw_scf_output_file(path):
                    pw_scf_output_files.append(path)
        return pw_scf_output_files

    def initial_structure_strings(self):
        structures = []
        for pw_scf_output_file in self._find_pw_scf_output_files():
            try:
                basis = self.txt_parser.initial_basis(self._get_file_content(pw_scf_output_file))
                lattice = self.txt_parser.initial_lattice_vectors(self._get_file_content(pw_scf_output_file))
                structures.append(lattice_basis_to_poscar(lattice, basis))
            except:
                raise
        return structures

    def final_structure_strings(self):
        structures = []
        for pw_scf_output_file in self._find_pw_scf_output_files():
            try:
                basis = self.txt_parser.final_basis(self._get_file_content(pw_scf_output_file))
                lattice = self.txt_parser.final_lattice_vectors(self._get_file_content(pw_scf_output_file))
                structures.append(lattice_basis_to_poscar(lattice, basis))
            except:
                pass
        return structures