from __future__ import annotations

import dataclasses
from typing import Dict, List, Any, Tuple, Optional
from enum import Enum
import io
import numpy as np

@dataclasses.dataclass
class Source:
    filepath: str
    name: str
    num_pages: int
    pages: List[Section]
    images: List[List[Any]]
    page_images: List[Any]
    authors: str
    url: str

    def __str__(self):
        return "<Source file={}, num_pages={}, num_images={}>".format(self.name, len(self.pages), len(self.images))

    def serialise(self) -> Tuple[bytes, Any]:
        '''Serialise the Source data into a tuple of raw byte data and structured json'''
        meta = {
			"source": self.name,
			"authors": self.authors,
			"url": self.url,
			"filepath": self.filepath,
			"num_pages": self.num_pages,
			"pages": [s.to_tuple() for s in self.pages],
		}

		### Write images to a compressed bytestream
        ims = {}
        if self.page_images:
            for i,im in enumerate(self.page_images):
                ims["page_{}".format(i)] = im

        if self.images:
            for i,images in enumerate(self.images):
                for j,im in enumerate(images):
                    ims["page_{}_image_{}".format(i, j)] = im

        stream = io.BytesIO()
        np.savez_compressed(stream, **ims)
        stream.seek(0)
        return [stream.read(), meta]

    @staticmethod
    def deserialise(images: bytes, data: Any) -> Source:
        '''Convert serialised data back into a source'''

        ### Load images and make sure we order them correctly
        stream = io.BytesIO(images)
        image_dict = np.load(stream)

        page_images = [None for i in range(int(data["num_pages"]))]
        loaded_images = [[] for i in range(int(data["num_pages"]))]

        for k in image_dict.keys():
            if "image" in k:
                parts = k.split("_")
                page = int(parts[1])
                image_num = int(parts[3])

                while len(loaded_images[page]) <= image_num:
                    loaded_images[page].append(None)

                loaded_images[page][image_num] = image_dict[k]
            else:
                page_images[int(k[5:])] = image_dict[k]

        ### Load
        s = Source(
            name=data["source"],
            authors=data["authors"],
            url=data["url"],
            filepath=data["filepath"],
            num_pages = data["num_pages"],
            pages=[Section.from_tuple(s) for s in data["pages"]],
            page_images = page_images,
            images=loaded_images
        )
        return s

@dataclasses.dataclass
class Bound:
    left: float
    top: float
    width: float
    height: float

    def __str__(self):
        return f"Bound<{self.left}, {self.top}, {self.right()}, {self.bottom()}>"

    def right(self) -> float:
        '''Returns right edge of bounding box'''
        return self.left + self.width

    def bottom(self) -> float:
        '''Returns top of bounding box'''
        return self.top + self.height

    @staticmethod
    def from_dict(data: Dict[str, float]) -> Bound:
        return Bound(
            left=data["left"],
            top=data["top"],
            height=data["height"],
            width=data["width"]
        )

    def to_dict(self) -> Dict[str, float]:
        return {
            "left":self.left,
            "top":self.top,
            "height":self.height,
            "width":self.width
        }

    @staticmethod
    def merge(bounds: List[Bound]) -> Bound:
        '''Merge a set of bounds into a single larger bound'''
        left = 10000000
        right = 0
        top = 10000000
        bottom = 0
        for b in bounds:
            if b.left < left:
                left = b.left
            if b.right() > right:
                right = b.right()
            if b.top < top:
                top = b.top
            if b.bottom() > bottom:
                bottom = b.bottom()

        return Bound(left=left, top=top, width=right-left, height=bottom-top)

@dataclasses.dataclass
class Line:
    '''Container class storing one or more lines of text alongside their attribues and bounding box'''
    id: str
    text: str
    bound: Bound
    page: int
    attributes: List[str]

    @staticmethod
    def merge(lines: List[Line], join_char=" ") -> Line:
        text = join_char.join(l.text for l in lines)
        bound = Bound.merge(l.bound for l in lines)
        attrib = []
        # Sort lines
        lines.sort(key=lambda l: l.bound.left)

        for l in lines:
            attrib += l.attributes
        return Line(id = lines[0].id, text=text, bound=bound, page=lines[0].page, attributes=attrib)

    @staticmethod
    def from_tuple(tuple: List[Any]) -> Line:
        return Line(id=str(tuple[0]), text=tuple[1], bound=Bound.from_dict(tuple[2]), page=tuple[3], attributes=tuple[4])

    def to_tuple(self) -> List[Any]:
        return [self.id, self.text, self.bound.to_dict(), self.page, self.attributes]

class Section:
    '''Container holding multiple lines with a single bounding box'''

    class SortOrder(Enum):
        NoSort = 0
        Vertical = 1
        Horizontal = 2

    def __init__(self, lines: List[Line] = None, attributes: List[str] = None,
            sort_order: Section.SortOrder=SortOrder.Vertical,
            bound: Optional[Bound]=None, ids: Optional[List[str]]=None, page=-1):
        self.lines = lines if lines else []
        self.ids = ids if ids else {l.id: l for l in self.lines}
        self.bound = bound if bound else Bound.merge(l.bound for l in self.lines)
        self.page = page
        self.attributes = attributes if attributes else []
        self.sort_order = sort_order

        if self.page == -1 and len(self.lines) > 0:
            self.page = min([l.page for l in self.lines])

    def is_empty(self) -> bool:
        '''Returns true if this section contains no lines'''
        return len(self.lines) == 0

    def add_line(self, line: Line, sort: bool=True, sort_order: Section.SortOrder=None) -> None:
        '''Add a line to the section. Will dynamically update the bound'''
        #Ignore if this line is already present
        if line.id in self.ids:
            return

        self.lines.append(line)
        self.ids[line.id] = line
        self.bound = Bound.merge([self.bound, line.bound])

        if sort:
            self.sort(sort_order)

    def add_section(self, section: Section, sort: bool=True, sort_order: Section.SortOrder=None) -> None:
        '''Add all lines from another section to this one'''
        for line in section.lines:
            self.add_line(line, sort=False)
        self.attributes += section.attributes
        if sort:
            self.sort(sort_order=sort_order)

        if self.page == -1:
            self.page = section.page
        elif section.page != -1:
            self.page = min(self.page, section.page)

    def remove_line(self, line: Line) -> None:
        '''Delete a line from this section. Expensive!'''
        keep_lines = []
        for l in self.lines:
            if l.id != line.id:
                keep_lines.append(l)
        self.ids.pop(line.id)

        self.lines = keep_lines
        self.bound = Bound.merge([l.bound for l in self.lines])

    def sort(self, sort_order: Section.SortOrder=None) -> None:
        '''Sort lines within section'''
        if sort_order == None:
            sort_order = self.sort_order

        if sort_order == Section.SortOrder.Vertical:
            self.lines.sort(key=lambda x: x.bound.top + 100*self.page)
        elif sort_order == Section.SortOrder.Horizontal:
            self.lines.sort(key=lambda x: x.bound.left + 100*self.page)
        elif sort_order == Section.SortOrder.NoSort:
            pass

    def get_line_by_id(self, line_id: str) -> Optional[Line]:
        '''Returns a line by the line id'''
        if line_id in self.ids:
            return self.ids[line_id]
        else:
            return None

    def get_line_attributes(self) -> List[str]:
        '''Returns a list of all attributes attached to these lines'''
        attribs = []
        for l in self.lines:
            attribs += l.attributes
        return attribs

    def get_section_text(self, join_char="\n") -> str:
        '''Returns the total section text'''
        text = ""
        for i in range(len(self.lines)):
            if len(self.lines[i].text) > 0 and self.lines[i].text[-1] == "-":
                text += self.lines[i].text[:-1]
            else:
                text += self.lines[i].text + join_char
        return text

    def __contains__(self, line: Line) -> bool:
        return line.id in self.ids

    def __len__(self) -> int:
        return len(self.lines)

    def to_tuple(self) -> List[Any]:
        lines = [l.to_tuple() for l in self.lines]
        bound = self.bound.to_dict()
        return [lines, bound, self.attributes, self.sort_order.value, self.page]

    @staticmethod
    def from_tuple(data: Any) -> Section:
        lines = [Line.from_tuple(l) for l in data[0]]
        ids = [l.id for l in lines]
        return Section(
            lines=lines,
            ids=ids,
            bound=Bound.from_dict(data[1]),
            attributes=data[2],
            sort_order=data[3],
            page=data[4] if len(data) > 4 else -1
        )


# Copyright 2019 Xilinx Inc.
#
# 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.

import os
import math
import time
import threading
import itertools
import collections
import numpy as np
from input_fn import calib_input
from tqdm import tqdm
import tensorflow as tf
from tensorflow.contrib import vitis_vai
from tensorflow.python.client import timeline
from yolov3_predictor import yolo_predictor
from config import Yolov3VOCConfig as Config

tf.app.flags.DEFINE_string('input_graph', '../float_model/yolov3_voc.pb',
                           'TensorFlow \'GraphDef\' file to load.')
tf.app.flags.DEFINE_string('eval_image_path', '',
                           'The directory where put the eval images')
tf.app.flags.DEFINE_string('eval_image_list', '',
                           'file has validation images list')
tf.app.flags.DEFINE_string('input_node', 'input_1', 'input node of pb model')
tf.app.flags.DEFINE_string(
    'output_node', 'conv2d_59/BiasAdd,conv2d_67/BiasAdd,conv2d_75/BiasAdd',
    'ouput node of pb model')
tf.app.flags.DEFINE_integer('input_height', 416, 'input height of pb model')
tf.app.flags.DEFINE_integer('input_width', 416, 'input width of pb model')
tf.app.flags.DEFINE_string('result_file', 'voc2007_pred_results.txt',
                           'The directory of output results')
tf.app.flags.DEFINE_boolean('use_quantize', False, 'quantize or not')
tf.app.flags.DEFINE_integer('nthreads', '12',
                            'The thread number for running evaluation.')
tf.app.flags.DEFINE_integer('batch', '8',
                            'The batch size for running evaluation.')
tf.app.flags.DEFINE_string('mode', '', 'accuracy or perf mode')

FLAGS = tf.app.flags.FLAGS

def write_items_to_file(image_id, items, fw):
    for item in items:
        fw.write(image_id + " " +
                 " ".join([str(comp) for comp in item]) + "\n")

def write_detection_result(detection_result, fw):
    for img_id, items in detection_result.items():
        write_items_to_file(img_id, items, fw)

def detection_thread(images, image_ids, input_image_shapes, class_names, t_id, detection_result):
    batch, threads = FLAGS.batch, FLAGS.nthreads
    begin = t_id * batch
    step = threads * batch
    image_num = len(images)
    for j in range(begin, image_num, step):
        run_size = batch
        if((image_num - j) < batch):
            run_size = image_num - j
        begin, end = j, j + run_size
        input_images, img_ids, input_shapes = images[begin:
                                                     end], image_ids[begin: end], input_image_shapes[begin: end]

        nms_boxes, nms_scores, nms_classes, valid_detections = sess.run(
            [pred_boxes, pred_scores, pred_classes, pred_detections],
            feed_dict={input_x: input_images, input_image_shape: input_shapes})

        if FLAGS.mode == "perf":
            continue

        for b, img_id, input_shape in zip(range(0, run_size), img_ids, input_shapes):
            valid_num = valid_detections[b]  # 1
            out_boxes = nms_boxes[b]  # [max_detections, 4]
            out_boxes = out_boxes[:valid_num]  # [valid_num, 4]
            out_class = nms_classes[b]  # [max_detections]
            out_class = out_class[:valid_num]  # [valid_num]
            out_scores = nms_scores[b]
            out_scores = out_scores[:valid_num]

            items = []
            for i, c in reversed(list(enumerate(out_class))):
                predicted_class = class_names[int(c)]
                box = out_boxes[i]
                score = out_scores[i]

                top, left, bottom, right = box
                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(input_shape[0], np.floor(
                    bottom + 0.5).astype('int32'))
                right = min(input_shape[1], np.floor(
                    right + 0.5).astype('int32'))
                item = [predicted_class, score, left, top, right, bottom]
                items.append(item)
            detection_result[img_id] = items

def run_detection_batch(images, img_ids, input_shapes, class_names):
    thread_list = []
    detection_result = {}
    for t_id in range(0, FLAGS.nthreads):
        t = threading.Thread(
            target=detection_thread,
            args=(images, img_ids, input_shapes, class_names, t_id, detection_result))
        thread_list.append(t)

    st = time.perf_counter()
    for t in thread_list:
        t.start()

    for t in thread_list:
        t.join()
    et = time.perf_counter()

    return detection_result, (et-st)

if __name__ == "__main__":

    config = Config()
    class_names = config.classes
    predictor = yolo_predictor(config)

    sess = tf.compat.v1.Session()
    with tf.io.gfile.GFile(FLAGS.input_graph, 'rb') as f:
        graph_def = tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())
        # Create vitis-ai wego graph for cloud DPU acceleration
        graph_def = vitis_vai.create_wego_graph(
            target="DPUCVDX8H_ISA1_F2W2_8PE",
            input_graph_def=graph_def,
            feed_dict={
                "input_1": [-1, FLAGS.input_height, FLAGS.input_width, 3]}
        )
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')
    sess.run(tf.compat.v1.global_variables_initializer())

    input_x = sess.graph.get_tensor_by_name(FLAGS.input_node + ':0')
    output_y = []
    for node in FLAGS.output_node.split(","):
        output_y.append(sess.graph.get_tensor_by_name(node + ":0"))

    input_image_shape = tf.compat.v1.placeholder(tf.int32, shape=(None, 2))
    pred_boxes, pred_scores, pred_classes, pred_detections = predictor.predict(
        output_y, input_image_shape)

    mode = FLAGS.mode
    if mode != "accuracy" and mode != "perf":
        raise ValueError(
            "Unsupported mode, support values: [ %s, %s]." %
            ("accuracy", "perf"))

    load_size = None if mode == "accuracy" else 960

    print("[INFO] loading images...")
    img_list, img_ids, img_input_shapes = calib_input(
        FLAGS.eval_image_path,
        FLAGS.eval_image_list,
        FLAGS.input_height,
        FLAGS.input_width,
        load_size)

    if mode == "perf":
        r_n = 20
        print("[INFO] start running perf...")
        print("[INFO] repeat running %d times with %d images." %
              (r_n, load_size))
        t = 0.0
        for i in tqdm(range(r_n)):
            _, t_ = run_detection_batch(
                img_list, img_ids, img_input_shapes, class_names)
            t += t_
        print("===================== Perf Result =====================")
        print("[Total Images] %d" % (r_n * load_size))
        print("[Total Time]   %0.6fs" % float(t))
        print("[FPS]          %0.2f" % (float(load_size) / (t / r_n)))
    else:
        fw = open(FLAGS.result_file, "w")
        print("[INFO] start running detection...")
        detection_result, d_time = run_detection_batch(
            img_list, img_ids, img_input_shapes, class_names)
        print("[INFO] writing detection result to file: ", FLAGS.result_file)
        write_detection_result(detection_result, fw)
        print("[INFO] evaluating detection results...")
        print("===================== Accuracy Result =====================")
        fw.close()

    sess.close()


import torch
import numpy as np
from detection.utils.general import non_max_suppression
from detection.utils.torch_utils import select_device
from detection.utils.datasets import letterbox

def detect(img, model):
    # assume img is raw image from cv2 with shape (height, widht, 3) and BGR
    # Then we will need to transpose it ot (channels, height, widht) and RGB
    origin_height, origin_width, _ = img.shape
    img = letterbox(img)
    # img = letterbox(img)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)

    # Initialize
    device = select_device()
    # half precision only supported on CUDA
    # device = 'gpu'
    half = device != 'cpu'

    # img = img -  [144.7748, 107.7354, 99.4750]
    # img[0] = img[0] - 114.7748
    img[0] = img[0] - 114.7748
    img[1] = img[1] - 107.7354
    img[2] = img[2] - 99.4750
    # invert RGB to tensor
    img = torch.from_numpy(img).to(device)
    # img = torch.from_numpy(img)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    # Inference
    pred = model(img, augment=True)[0]

    # Apply NMS
    pred = non_max_suppression(pred, 0.01, 0.5, classes=0, agnostic=True)
    # Process detections
    # only one img
    det = pred[0]
    if det is not None and len(det):
        det = det.cpu().numpy()
        det = det[det.argmax(axis=0)[-2]]
        # Write results
        x_center = (det[0] + det[2]) / 2
        y_center = (det[1] + det[3]) / 2
        # convert it to original size position
        x_center = round(x_center/640 * origin_width)
        y_center = round(y_center/384 * origin_height)
        return int(x_center), int(y_center)


import os
import time
import config
import numpy as np
from PIL import Image
import tensorflow as tf
from dataReader import Reader
from model.yolo3_model import yolo
from collections import defaultdict
from yolo_predict import yolo_predictor
from utils import draw_box, load_weights, letterbox_image, voc_ap

# 指定使用GPU的Index
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu_index

def train():
    """
    Introduction
    ------------
        训练模型
    """
    train_reader = Reader('train', config.data_dir, config.anchors_path, config.num_classes, input_shape = config.input_shape, max_boxes = config.max_boxes)
    train_data = train_reader.build_dataset(config.train_batch_size)
    is_training = tf.placeholder(tf.bool, shape = [])
    iterator = train_data.make_one_shot_iterator()
    images, bbox, bbox_true_13, bbox_true_26, bbox_true_52 = iterator.get_next()
    images.set_shape([None, config.input_shape, config.input_shape, 3])
    bbox.set_shape([None, config.max_boxes, 5])
    grid_shapes = [config.input_shape // 32, config.input_shape // 16, config.input_shape // 8]
    bbox_true_13.set_shape([None, grid_shapes[0], grid_shapes[0], 3, 5 + config.num_classes])
    bbox_true_26.set_shape([None, grid_shapes[1], grid_shapes[1], 3, 5 + config.num_classes])
    bbox_true_52.set_shape([None, grid_shapes[2], grid_shapes[2], 3, 5 + config.num_classes])
    draw_box(images, bbox)
    model = yolo(config.norm_epsilon, config.norm_decay, config.anchors_path, config.classes_path, config.pre_train)
    bbox_true = [bbox_true_13, bbox_true_26, bbox_true_52]
    output = model.yolo_inference(images, config.num_anchors / 3, config.num_classes, is_training)
    loss = model.yolo_loss(output, bbox_true, model.anchors, config.num_classes, config.ignore_thresh)
    l2_loss = tf.losses.get_regularization_loss()
    loss += l2_loss
    tf.summary.scalar('loss', loss)
    merged_summary = tf.summary.merge_all()
    global_step = tf.Variable(0, trainable = False)
    lr = tf.train.exponential_decay(config.learning_rate, global_step, decay_steps = 2000, decay_rate = 0.8)
    optimizer = tf.train.AdamOptimizer(learning_rate = lr)
    # 如果读取预训练权重，则冻结darknet53网络的变量
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        if config.pre_train:
            train_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='yolo')
            train_op = optimizer.minimize(loss = loss, global_step = global_step, var_list = train_var)
        else:
            train_op = optimizer.minimize(loss = loss, global_step = global_step)
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()
    with tf.Session(config = tf.ConfigProto(log_device_placement = False)) as sess:
        ckpt = tf.train.get_checkpoint_state(config.model_dir)
        if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):
            print('restore model', ckpt.model_checkpoint_path)
            saver.restore(sess, ckpt.model_checkpoint_path)
        else:
            sess.run(init)
        if config.pre_train is True:
            load_ops = load_weights(tf.global_variables(scope = 'darknet53'), config.darknet53_weights_path)
            sess.run(load_ops)
        summary_writer = tf.summary.FileWriter(config.log_dir, sess.graph)
        loss_value = 0
        for epoch in range(config.Epoch):
            for step in range(int(config.train_num / config.train_batch_size)):
                start_time = time.time()
                train_loss, summary, global_step_value, _ = sess.run([loss, merged_summary, global_step, train_op], {is_training : True})
                loss_value += train_loss
                duration = time.time() - start_time
                examples_per_sec = float(duration) / config.train_batch_size
                format_str = ('Epoch {} step {},  train loss = {} ( {} examples/sec; {} ''sec/batch)')
                print(format_str.format(epoch, step, loss_value / global_step_value, examples_per_sec, duration))
                summary_writer.add_summary(summary = tf.Summary(value = [tf.Summary.Value(tag = "train loss", simple_value = train_loss)]), global_step = step)
                summary_writer.add_summary(summary, step)
                summary_writer.flush()
            # 每3个epoch保存一次模型
            if epoch % 3 == 0:
                checkpoint_path = os.path.join(config.model_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step = global_step)

def eval(model_path, min_Iou = 0.5, yolo_weights = None):
    """
    Introduction
    ------------
        计算模型在coco验证集上的MAP, 用于评价模型
    """
    ground_truth = {}
    class_pred = defaultdict(list)
    gt_counter_per_class = defaultdict(int)
    input_image_shape = tf.placeholder(dtype = tf.int32, shape = (2,))
    input_image = tf.placeholder(shape = [None, 416, 416, 3], dtype = tf.float32)
    predictor = yolo_predictor(config.obj_threshold, config.nms_threshold, config.classes_path, config.anchors_path)
    boxes, scores, classes = predictor.predict(input_image, input_image_shape)
    val_Reader = Reader("val", config.data_dir, config.anchors_path, config.num_classes, input_shape = config.input_shape, max_boxes = config.max_boxes)
    image_files, bboxes_data = val_Reader.read_annotations()
    with tf.Session() as sess:
        if yolo_weights is not None:
            with tf.variable_scope('predict'):
                boxes, scores, classes = predictor.predict(input_image, input_image_shape)
            load_op = load_weights(tf.global_variables(scope = 'predict'), weights_file = yolo_weights)
            sess.run(load_op)
        else:
            saver = tf.train.Saver()
            saver.restore(sess, model_path)
        for index in range(len(image_files)):
            val_bboxes = []
            image_file = image_files[index]
            file_id = os.path.split(image_file)[-1].split('.')[0]
            for bbox in bboxes_data[index]:
                left, top, right, bottom, class_id = bbox[0], bbox[1], bbox[2], bbox[3], bbox[4]
                class_name = val_Reader.class_names[int(class_id)]
                bbox = [float(left), float(top), float(right), float(bottom)]
                val_bboxes.append({"class_name" : class_name, "bbox": bbox, "used": False})
                gt_counter_per_class[class_name] += 1
            ground_truth[file_id] = val_bboxes
            image = Image.open(image_file)
            resize_image = letterbox_image(image, (416, 416))
            image_data = np.array(resize_image, dtype = np.float32)
            image_data /= 255.
            image_data = np.expand_dims(image_data, axis = 0)

            out_boxes, out_scores, out_classes = sess.run(
                [boxes, scores, classes],
                feed_dict = {
                    input_image : image_data,
                    input_image_shape : [image.size[1], image.size[0]]
                })
            print("detect {}/{} found boxes: {}".format(index, len(image_files), len(out_boxes)))
            for o, c in enumerate(out_classes):
                predicted_class = val_Reader.class_names[c]
                box = out_boxes[o]
                score = out_scores[o]

                top, left, bottom, right = box
                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
                right = min(image.size[0], np.floor(right + 0.5).astype('int32'))

                bbox = [left, top, right, bottom]
                class_pred[predicted_class].append({"confidence": str(score), "file_id": file_id, "bbox": bbox})

    # 计算每个类别的AP
    sum_AP = 0.0
    count_true_positives = {}
    for class_index, class_name in enumerate(sorted(gt_counter_per_class.keys())):
        count_true_positives[class_name] = 0
        predictions_data = class_pred[class_name]
        # 该类别总共有多少个box
        nd = len(predictions_data)
        tp = [0] * nd  # true positive
        fp = [0] * nd  # false positive
        for idx, prediction in enumerate(predictions_data):
            file_id = prediction['file_id']
            ground_truth_data = ground_truth[file_id]
            bbox_pred = prediction['bbox']
            Iou_max = -1
            gt_match = None
            for obj in ground_truth_data:
                if obj['class_name'] == class_name:
                    bbox_gt = obj['bbox']
                    bbox_intersect = [max(bbox_pred[0], bbox_gt[0]), max(bbox_gt[1], bbox_pred[1]), min(bbox_gt[2], bbox_pred[2]), min(bbox_gt[3], bbox_pred[3])]
                    intersect_weight = bbox_intersect[2] - bbox_intersect[0] + 1
                    intersect_high = bbox_intersect[3] - bbox_intersect[1] + 1
                    if intersect_high > 0 and intersect_weight > 0:
                        union_area = (bbox_pred[2] - bbox_pred[0] + 1) * (bbox_pred[3] - bbox_pred[1] + 1) + (bbox_gt[2] - bbox_gt[0] + 1) * (bbox_gt[3] - bbox_gt[1] + 1) - intersect_weight * intersect_high
                        Iou = intersect_high * intersect_weight / union_area
                        if Iou > Iou_max:
                            Iou_max = Iou
                            gt_match = obj
            if Iou_max > min_Iou:
                if not gt_match['used'] and gt_match is not None:
                    tp[idx] = 1
                    gt_match['used'] = True
                else:
                    fp[idx] = 1
            else:
                fp[idx] = 1
        # 计算精度和召回率
        sum_class = 0
        for idx, val in enumerate(fp):
            fp[idx] += sum_class
            sum_class += val
        sum_class = 0
        for idx, val in enumerate(tp):
            tp[idx] += sum_class
            sum_class += val
        rec = tp[:]
        for idx, val in enumerate(tp):
            rec[idx] = tp[idx] / gt_counter_per_class[class_name]
        prec = tp[:]
        for idx, val in enumerate(tp):
            prec[idx] = tp[idx] / (fp[idx] + tp[idx])

        ap, mrec, mprec = voc_ap(rec, prec)
        sum_AP += ap
    MAP = sum_AP / len(gt_counter_per_class) * 100
    print("The Model Eval MAP: {}".format(MAP))

if __name__ == "__main__":
    train()
    # 计算模型的Map
    # eval(config.model_dir, yolo_weights = config.yolo3_weights_path)


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

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

from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.algo import gail
from a2c_ppo_acktr.arguments import get_args
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
from a2c_ppo_acktr.utils import get_vec_normalize
from evaluation import evaluate

def main():
    #Get arguments from command line
    args = get_args()

    #Set random seeds
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    #Helps with reproducibility?
    if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

    #Where to log
    if args.i != None:
        args.log_dir += str(args.i)
    args.log_dir += "-" + str(args.seed)
    log_dir = os.path.expanduser(args.log_dir)
    eval_log_dir = log_dir + "_eval"
    utils.cleanup_log_dir(log_dir)
    utils.cleanup_log_dir(eval_log_dir)

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

    # envs_gail = make_vec_envs(args.env_name, args.seed, 1,
                         # args.gamma, '/tmp/gym/', device, False)

    envs_seed = make_vec_envs(args.env_name, args.seed, args.num_processes,
                         args.gamma, args.log_dir, device, False)

    aenv_name = args.env_name

    genv_name = args.env_name
    # genv_name = args.env_name

    if args.sub:
        args.env_name = "Sub" + args.env_name
        genv_name = "Sub" + genv_name

    if args.train_gail:
        init_gail = 29
    else:
        init_gail = -1
    update_freq = 5

    #Create base policy
    actor_critic_g = Policy(
        envs_seed.observation_space.shape,
        envs_seed.action_space,
        base_kwargs={'recurrent': args.recurrent_policy})
    actor_critic_g.to(device)

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

    if not(args.train_gail):
        actor_critic_g, ob_rms = \
                torch.load(os.path.join(args.gail_agent_dir, genv_name + "gail.pt"))

        # vec_norm = get_vec_normalize(envs_gail)
        # vec_norm2 = get_vec_normalize(envs_seed)
        # if vec_norm2 is not None:

            # vec_norm2.eval()

            # vec_norm2.obs_rms = ob_rms

    #Set algorithm based on input
    gailfilt = lambda x: np.clip((x - ob_rms.mean) / np.sqrt(ob_rms.var + utils.get_vec_normalize(envs_seed).epsilon), -utils.get_vec_normalize(envs_seed).clip_obs, utils.get_vec_normalize(envs_seed).clip_obs)

    if args.algo == 'a2c':
        agent_g = algo.A2C_ACKTR(
            actor_critic_g,
            args.value_loss_coef,
            args.entropy_coef,
            lr=args.lr,
            eps=args.eps,
            alpha=args.alpha,
            max_grad_norm=args.max_grad_norm)
        agent_s = algo.A2C_ACKTR(
            actor_critic_s,
            args.value_loss_coef,
            args.entropy_coef,
            lr=args.lr,
            eps=args.eps,
            alpha=args.alpha,
            max_grad_norm=args.max_grad_norm)
    elif args.algo == 'ppo':
        agent_g = algo.PPO(
            actor_critic_g,
            args.clip_param,
            args.ppo_epoch,
            args.num_mini_batch,
            args.value_loss_coef,
            args.entropy_coef,
            lr=args.lr,
            eps=args.eps,
            max_grad_norm=args.max_grad_norm)
        agent_s = algo.SEED(
            actor_critic_s,
            actor_critic_g,
            args.clip_param,
            args.ppo_epoch,
            args.num_mini_batch,
            args.value_loss_coef,
            args.entropy_coef,
            args.kl_coef,
            lr=args.lr,
            eps=args.eps,
            steps=args.num_steps*args.num_processes,
            env_steps=args.num_env_steps,
            max_grad_norm=args.max_grad_norm)
    elif args.algo == 'acktr':
        agent_g = algo.A2C_ACKTR(
            actor_critic_g, args.value_loss_coef, args.entropy_coef, acktr=True)
        agent_s = algo.A2C_ACKTR(
            actor_critic_s, args.value_loss_coef, args.entropy_coef, acktr=True)

    if args.gail:
        assert len(envs_seed.observation_space.shape) == 1
        #Create discriminator
        acti = 0
        space = None
        if envs_seed.action_space.__class__.__name__ == "Discrete":
            acti = envs_seed.action_space.n
            space = acti
        else:
            acti = envs_seed.action_space.shape[0]

        print(acti)
        if args.train_gail:
            discr = gail.Discriminator(
                envs_seed.observation_space.shape[0] + acti, 100,
                device, space)
        else:
            discr = torch.load(os.path.join(args.gail_agent_dir, genv_name + "discr.pt"))[0]
            print(discr)
        #Get filename based on environment
        file_name = os.path.join(
            args.gail_experts_dir, "trajs_{}.pt".format(
                "mountaincarcontinuous"))
        #Store dataset
        expert_dataset = gail.ExpertDataset(
            file_name, num_trajectories=4, subsample_frequency=2)
        drop_last = len(expert_dataset) > args.gail_batch_size
        gail_train_loader = torch.utils.data.DataLoader(
            dataset=expert_dataset,
            batch_size=args.gail_batch_size,
            shuffle=True,
            drop_last=drop_last)

    # rollouts_g = RolloutStorage(args.num_steps, args.num_processes,
                              # envs_gail.observation_space.shape, envs_gail.action_space,
                              # actor_critic_g.recurrent_hidden_state_size)
    rollouts_s = RolloutStorage(args.num_steps, args.num_processes,
                                envs_seed.observation_space.shape, envs_seed.action_space,
                                actor_critic_s.recurrent_hidden_state_size)

    orig_obs = torch.zeros(args.num_steps + 1, args.num_processes, *envs_seed.observation_space.shape)

    # obs_gail = envs_gail.reset()
    # rollouts_g.obs[0].copy_(obs_gail)
    # rollouts_g.to(device)

    obs_seed = envs_seed.reset()
    orig_obs[0].copy_(torch.from_numpy(utils.get_vec_normalize(envs_seed).get_original_obs()))
    rollouts_s.obs[0].copy_(obs_seed)
    rollouts_s.to(device)

    episode_rewards_g = deque(maxlen=10)
    episode_rewards_s = deque(maxlen=10)

    episode_rewards_s.append(0)
    episode_rewards_g.append(0)

    start = time.time()
    num_updates = int(
        args.num_env_steps) // args.num_steps // args.num_processes
    for j in range(num_updates):
        # vec_norm_g = get_vec_normalize(envs_gail)
        # vec_norm_s = get_vec_normalize(envs_seed)
        # if vec_norm_s is not None:
            # vec_norm_s.eval()
            # vec_norm_s.obs_rms = ob_rms
        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            utils.update_linear_schedule(
                agent_g.optimizer, j, num_updates,
                agent_g.optimizer.lr if args.algo == "acktr" else args.lr)
            utils.update_linear_schedule(
                agent_s.optimizer, j, num_updates,
                agent_s.optimizer.lr if args.algo == "acktr" else args.lr)

        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                # if (j < init_gail or j%update_freq==0) and args.train_gail:
                    # value_g, action_g, action_log_prob_g, recurrent_hidden_states_g = actor_critic_g.act(
                        # rollouts_g.obs[step], rollouts_g.recurrent_hidden_states[step],
                        # rollouts_g.masks[step])
                if j > init_gail:
                    value_s, action_s, action_log_prob_s, recurrent_hidden_states_s = actor_critic_s.act(
                        rollouts_s.obs[step], rollouts_s.recurrent_hidden_states[step],
                        rollouts_s.masks[step])
            # Obser reward and next obs
            # if (j < init_gail or j%update_freq==0) and args.train_gail:
                # obs_g, reward_g, done_g, infos_g = envs_gail.step(torch.squeeze(action_g))
                # obs_g, reward_g, done_g, infos_g = envs_gail.step(action_g)

                # for info in infos_g:
                    # if 'episode' in info.keys():
                        # episode_rewards_g.append(info['episode']['r'])

                # If done then clean the history of observations.
                # masks_g = torch.FloatTensor(
                    # [[0.0] if done_ else [1.0] for done_ in done_g])
                # bad_masks_g = torch.FloatTensor(
                    # [[0.0] if 'bad_transition' in info.keys() else [1.0]
                     # for info in infos_g])

                # rollouts_g.insert(obs_g, recurrent_hidden_states_g, action_g,
                                # action_log_prob_g, value_g, reward_g, masks_g, bad_masks_g)

            if j > init_gail:
                # obs_s, reward_s, done_s, infos_s = envs_seed.step(torch.squeeze(action_s))
                obs_s, reward_s, done_s, infos_s = envs_seed.step(action_s)

                for info in infos_s:
                    if 'episode' in info.keys():
                        episode_rewards_s.append(info['episode']['r'])
                masks_s = torch.FloatTensor(
                    [[0.0] if done_ else [1.0] for done_ in done_s])
                bad_masks_s = torch.FloatTensor(
                    [[0.0] if 'bad_transition' in info.keys() else [1.0]
                     for info in infos_s])
                rollouts_s.insert(obs_s, recurrent_hidden_states_s, action_s,
                            action_log_prob_s, value_s, reward_s, masks_s, bad_masks_s)
                orig_obs[step+1].copy_(torch.from_numpy(utils.get_vec_normalize(envs_seed).get_original_obs()))

        with torch.no_grad():
            if (j < init_gail or j%update_freq==0) and args.train_gail:
                next_value_g = actor_critic_g.get_value(
                    rollouts_g.obs[-1], rollouts_g.recurrent_hidden_states[-1],
                    rollouts_g.masks[-1]).detach()
            if j > init_gail:
                next_value_s = actor_critic_s.get_value(
                    rollouts_s.obs[-1], rollouts_s.recurrent_hidden_states[-1],
                    rollouts_s.masks[-1]).detach()

        total = 0
        if args.gail:
            # if j >= 10:
                # envs_gail.venv.eval()
            # gail_epoch = args.gail_epoch
            # if j < 10:
                # gail_epoch = 100  # Warm up
            # if (j < init_gail or j%update_freq==0) and args.train_gail:
                # for _ in range(gail_epoch):
                    # discr.update(gail_train_loader, rollouts_g,
                                 # utils.get_vec_normalize(envs_gail)._obfilt)

            for step in range(args.num_steps):
                if (j < init_gail or j%update_freq==0) and args.train_gail:
                    rollouts_g.rewards[step], h = discr.predict_reward(
                        rollouts_g.obs[step], rollouts_g.actions[step], args.gamma,
                        rollouts_g.masks[step])
                    total += h
                if j > init_gail:
                    if j < num_updates/2:
                        r_s, h = discr.predict_reward_seed(
                            orig_obs[step], rollouts_s.actions[step], args.gamma,
                            rollouts_s.masks[step], gailfilt)
                        r_s = torch.clamp(r_s, -5, 5)

                        rollouts_s.rewards[step]= -args.demonstration_coef*r_s*2*(.5-j/num_updates)+rollouts_s.rewards[step]

                    else:
                        rollouts_s.rewards[step] = rollouts_s.rewards[step]
            total = total/args.num_steps/args.num_processes

        if (j < init_gail or j%update_freq==0) and args.train_gail:
            rollouts_g.compute_returns(next_value_g, args.use_gae, args.gamma,
                                     args.gae_lambda, args.use_proper_time_limits)
        if j > init_gail:
            rollouts_s.compute_returns(next_value_s, args.use_gae, args.gamma,
                                     args.gae_lambda, args.use_proper_time_limits)

        # if j % 2 == 0:
            # value_loss_g, action_loss_g, dist_entropy_g = agent_g.update(rollouts_g)
            # rollouts_g.after_update()
        # else:
            # value_loss_s, action_loss_s, dist_entropy_s = agent_s.update(rollouts_s)
            # rollouts_s.after_update()
        if (j < init_gail or j%update_freq==0) and args.train_gail:
            value_loss_g, action_loss_g, dist_entropy_g = agent_g.update(rollouts_g)
            rollouts_g.after_update()

        if j > init_gail:
            value_loss_s, action_loss_s, dist_entropy_s = agent_s.update(rollouts_s, orig_obs, gailfilt)
            rollouts_s.after_update()
            orig_obs[0].copy_(orig_obs[-1])

        # save for every interval-th episode or for the last epoch
        if (j % args.save_interval == 0
                or j == num_updates - 1) and args.save_dir != "":
            save_path = os.path.join(args.save_dir, args.algo)
            try:
                os.makedirs(save_path)
            except OSError:
                pass
            env_name = args.env_name.split(":")[-1]
            # torch.save([
                # actor_critic_g,
                # getattr(get_vec_normalize(envs_gail), 'obs_rms', None)
            # ], os.path.join(save_path, env_name + "gail.pt"))
            it = ""
            if args.i != None:
                it = args.i
            torch.save([
                actor_critic_s,
                getattr(get_vec_normalize(envs_seed), 'obs_rms', None)
            ], os.path.join(save_path, env_name + str(it) + "seed.pt"))

        if j % args.log_interval == 0 and len(episode_rewards_s) > 1:
            total_num_steps = (j + 1) * args.num_processes * args.num_steps
            end = time.time()
            # print("GAIL")
            # print(total)
            # print(
                # "Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.3f}/{:.3f}, min/max reward {:.3f}/{:.3f}\n"
                # .format(j, total_num_steps,
                        # int(total_num_steps / (end - start)),
                        # len(episode_rewards_g), np.mean(episode_rewards_g),
                        # np.median(episode_rewards_g), np.min(episode_rewards_g),
                        # np.max(episode_rewards_g)))
            if j > init_gail:
                print(episode_rewards_s)
                print(
                    "Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.3f}/{:.3f}, min/max reward {:.3f}/{:.3f}\n"
                    .format(j, total_num_steps,
                            int(total_num_steps / (end - start)),
                            len(episode_rewards_s), np.mean(episode_rewards_s),
                            np.median(episode_rewards_s), np.min(episode_rewards_s),
                            np.max(episode_rewards_s)))
        num = args.num_processes
        if (args.eval_interval is not None and len(episode_rewards_s) > 1
                and j % args.eval_interval == 0):
            obs_rms = utils.get_vec_normalize(envs_seed).obs_rms
            evaluate(actor_critic_s, obs_rms, aenv_name, args.seed,
                     num, eval_log_dir, device)

if __name__ == "__main__":
    main()


# -*- coding: UTF-8 -*-

"""
iDARTS - predict
Predict splicing given sequence features and expression profiles
"""

import os
import h5py
from tqdm import tqdm
import numpy as np
from pkg_resources import resource_filename
from . import config
from keras.models import model_from_json

import logging
logger = logging.getLogger('iDARTS.predict')

def get_file_num_lines(filename):
    fp = open(filename)
    count = 0
    for line in fp:
        count += 1
    fp.close()
    return count

def load_dl_model(event_type):
    model_list = []
    model_config_fopen = open(config.MODEL_DIR + '/model.{}.json'.format(event_type))
    model_config = model_config_fopen.read()
    model_config_fopen.close()
    for model_num in xrange(1, 6):
        model = model_from_json(model_config)
        model.load_weights(config.MODEL_DIR + '/model.{}.{}.hd5'.format(event_type, model_num))
        model_list.append(model)
    return model_list

def get_dnn_pred(X, model_list):
    y = np.mean(np.array([model.predict(X).flatten() for model in model_list]), axis = 0)
    return y

def get_rbp_list():
    rbp_list = []
    fp = open(config.RBP_LIST_FN)
    fp.readline()
    for line in fp:
        sp = line.strip().split('\t')
        rbp_list.append(sp[0])
    fp.close()
    return rbp_list

def get_user_exp_profiles(expr_fn, exp_max_value):
    rbp_list = get_rbp_list()
    rbp_dict = {rbp:True for rbp in rbp_list}
    fp = open(expr_fn)
    header = fp.readline().strip().split('\t')
    tissue_list = header[1].split(',')
    tissue_num = len(tissue_list)
    dict_rbp2exp = {}
    for line in fp:
        sp = line.strip().split('\t')
        gene_id = sp[0].split('.')[0]
        if gene_id in rbp_dict:
            dict_rbp2exp[gene_id] = [float(x) for x in sp[1].split(',')]
    fp.close()
    if len(dict_rbp2exp) != len(rbp_list):
        logger.info("RBP expression profile generated failed, please check RBP Gene lists")
        sys.stderr.write("RBP expression profile generated failed, please check RBP Gene lists")
        sys.exit(0)
    tissue2exp_matrix = []
    for rbp_id in rbp_list:
        tissue2exp_matrix.append(dict_rbp2exp[rbp_id])
    tissue2exp_matrix = np.array(tissue2exp_matrix)
    tissue2exp_matrix = tissue2exp_matrix.T / exp_max_value
    user_exp_dict = {}
    user_exp_dict['tissue2exp_matrix'] = tissue2exp_matrix
    user_exp_dict['tissue_list'] = tissue_list
    user_exp_dict['tissue_num'] = tissue_num
    return user_exp_dict

def make_prediction(args):
    event_type = args.event_type
    input_file = args.input
    out_file = args.output
    model_list = load_dl_model(event_type)
    event_config = h5py.File(config.MODEL_DIR + '/{}.config.hd5'.format(event_type), 'r')
    feature_max_value = event_config.get('feature_max_value')[:]
    if args.expr != None:
        exp_max_value = event_config.get('exp_max_value')[:]
        expr_fn = args.expr
        user_exp_dict = get_user_exp_profiles(expr_fn, exp_max_value)
        tissue_list = user_exp_dict['tissue_list']
        tissue_num = user_exp_dict['tissue_num']
        tissue2exp_matrix = user_exp_dict['tissue2exp_matrix']
    else:
        tissue_list = list(event_config.get('tissue_list')[:])
        tissue_num = len(tissue_list)
        tissue2exp_matrix = event_config.get('tissue2exp_matrix')[:]
    num_lines = get_file_num_lines(input_file) - 1
    fp = open(input_file)
    header = fp.readline().strip().split('\t')
    if event_type == 'A3SS' or event_type == 'A5SS' or event_type == 'RI':
        logger.info('A3SS, A5SS, and RI still under development')
        sys.exit(0)
        feature_start_idx = header.index('9G8.C_F')
    else:
        feature_start_idx = header.index('LogLen.C1')
    fw = open(out_file, 'w')
    fw.write(header[0] + '\t' + '\t'.join(tissue_list) + '\n')
    for line in tqdm(fp, total = num_lines):
        sp = line.strip().split('\t')
        feature = np.array([float(x) for x in sp[feature_start_idx:]]) / feature_max_value
        X = np.concatenate((np.repeat([feature], tissue_num, axis = 0), tissue2exp_matrix), axis = 1)
        y = get_dnn_pred(X, model_list)
        fw.write(sp[0] + '\t' + '\t'.join([str(x) for x in y]) + '\n')
    fp.close()

def parser(args):
    make_prediction(args)


import argparse
import csv
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import re
import numpy as bp
import scipy
import pysam
import multiprocessing
import gzip
#from vcftidy import *

class RecessiveModel:
	def __init__(self, af = 1e-2, SBPV_cutoff=1e-5, DP_cutoff=8, AB_cutoff1=0.25, AB_cutoff2=0.8):
		self.SBPV_cutoff = SBPV_cutoff
		self.DP_cutoff = DP_cutoff
		self.AB_cutoff1 = AB_cutoff1
		self.AB_cutoff2 = AB_cutoff2
		self.LGD = set(["splice_acceptor_variant", "splice_donor_variant", "stop_gained",
			"stop_lost", "start_lost", "frameshift_variant"])
		self.CADD_cutoff = 25
		self.REVEL_cutoff = 0.5
		self.AF_cutoff = af
		#self.C = ["syn","lgd","dmis","lgd/dmis"]
		self.C = ["syn","lgd","mis","cadd15","cadd20","cadd25","revel.5","mvp2.85","mpc1","lgd_cadd25"]
		return

	# compute NHet, NHom, AC and AF of each site, within a population, Write to a VCF file.
	def ComputeSiteAF(self, InpVCF, Indvs, prefix, OutVCF):
		fin = gzip.open(InpVCF, 'rt')
		fout = open(OutVCF, 'w')
		for l in fin:
			if l.startswith("##"):
				fout.write(l)
				continue
			elif l.startswith("#"):
				llist = l.strip().split("\t")
				#Head = llist[]
				indvs = llist[9:]
				index = self.GetIndex(Indvs, indvs)
				fout.write("##ComputeAF={}\n".format(prefix))
				fout.write("\t".join(llist[:9])+"\n")
			else:
				llist = l.strip().split("\t")
				length = (len(llist[4].split(","))+1)
				AFs = [0] * length
				ACs = [0] * length
				AC_Het = [0] * length
				AC_Hom = [0] * length
				AN = 0
				GTs = llist[9:]
				for idx in index:
					GT = self.GenotypeQC(llist[8], GTs[idx])
					if GT:
						A1, A2 = GT[0], GT[1]
						ACs[A1] += 1
						ACs[A2] += 1
						if A1 != A2:
							AC_Het[A2] += 1
						else:
							AC_Hom[A2] += 1
						AN += 2
				for i in range(length):
					try:
						AFs[i] = str(float(ACs[i])/AN)
					except ZeroDivisionError:
						AFs[i] = "0"
					ACs[i] = str(ACs[i])
					AC_Het[i] = str(AC_Het[i])
					AC_Hom[i] = str(AC_Hom[i])
				New = ";{}_AN={};{}_AC={};{}_AF={};{}_AC_Het={};{}_AC_Hom={}".format(prefix, AN, prefix, ",".join(ACs[1:]), prefix, ",".join(AFs[1:]), prefix, ",".join(AC_Het[1:]), prefix, ",".join(AC_Hom[1:]))
				llist[7] = llist[7] + New
				fout.write("\t".join(llist[:8])+"\n")
		return
	def GetIndex(self, Indvs, indvs):
		res = []
		for indv in Indvs:
			try:
				res.append(indvs.index(indv))
			except:
				continue
		return res

	# Return genotype as [A1, A2] if pass QC , return False otherwise
	def GenotypeQC(self, fmt, gt_dat):
		tmp = {}
		for k,v in zip(fmt.split(":"), gt_dat.split(":")):
			tmp[k] = v
		GT = tmp["GT"].split("/")
		if tmp["GQ"] == ".":
			#return False
			return [0,0]
		elif float(tmp["GQ"]) < 60:
			return False
		if GT[0] == "." or GT[1] == ".":
			return False
		if tmp["GT"] != "0/0":
			#print(fmt, gt_dat)
			#if "." in tmp["SBPV"]:
			#	pass
			#elif float(tmp["SBPV"]) < self.SBPV_cutoff:
			#	return False
			if float(tmp["DP"]) < self.DP_cutoff:
				return False
			if float(tmp["AD"].split(",")[1])/float(tmp["DP"]) < self.AB_cutoff1: # or float(tmp["AD"].split(",")[1])/float(tmp["DP"]) > self.AB_cutoff2:
				return False
		return [int(i) for i in GT]

	def LoadPedigree(self, PedFil, Samples):
		PedFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/recessive/EUR_Fams.ped"
		Fams = []
		reader = csv.reader(open(PedFil, 'rt'), delimiter="\t")
		PreFamID, tmp = None, None
		for row in reader:
			row.append(Samples.index(row[1])) # Add sample index in VCF header, to locate genotype
			FamID = row[0]
			if FamID != PreFamID:
				if tmp != None:
					Fams.append(tmp)
				tmp = Family(FamID)
				PreFamID = FamID
				tmp.Proband = Sample(row)
			else:
				if row[1] == tmp.Proband.Father:
					tmp.Father = Sample(row)
				elif row[1] == tmp.Proband.Mother:
					tmp.Mother = Sample(row)
				else:
					tmp.Siblings.append(Sample(row))
		return Fams

	def getINFO(self, info_string):
		infolist = info_string.split(';')
		infodict = {}
		for kv in infolist:
			kv = kv.split('=')
			if len(kv) == 2:
				k, v = kv
				infodict[k] = v
		return infodict

	def match_allele_csq(self, Ref, Alts, csq_head, csq_string):
		# Trim Leading Base
		Alts = Alts.split(",")
		if len(list(set([x[0] for x in Alts])))==1 and Ref[0] == list(set([x[0] for x in Alts]))[0]:
			_Ref = Ref[1:] if len(Ref[1:]) >0 else "-"
			_Alts = [Alt[1:] if len(Alt[1:]) >0 else "-" for Alt in Alts]
			#print (Ref, Alts, ";", _Ref, _Alts)
		else:
			_Alts = Alts
		res = {}
		csqs = csq_string.split(",")
		csqs = [dict(zip(csq_head, vep.split("|"))) for vep in csqs]
		for i, Alt in enumerate(Alts):
			res[Alt] = []
			for j, csq in enumerate(csqs):
				if csq["Allele"] == _Alts[i]:
					csq["Consequence"] = csq["Consequence"] .split("&")
					res[Alt].append(csq)
		return res

	def search_severe_consequence(self, var_k, Allele_CSQ_dict, Alt):
		severe_consequence = None
		severe_trans = None
		for i in range(len(Allele_CSQ_dict[Alt])):
			consequence = Allele_CSQ_dict[Alt][i]["Consequence"]
			Transcript = Allele_CSQ_dict[Alt][i]["Feature"]
			#if var_k == "11:7994466:T:G":
			#   print(consequence, Transcript)
			#if set(consequence).intersection(self.LGD) >= 1:
			if len(set(consequence).intersection(self.LGD))>= 1:
				return i, consequence, Transcript
			elif consequence[0] == "missense_variant":
				return i, consequence, Transcript
			elif consequence[0] == "synonymous_variant":
				#print(consequence)
				severe_consequence = consequence
				severe_trans = Transcript
		if severe_consequence == None:
			return 0, "non-coding", None
		else:
			return 0, severe_consequence, severe_trans

	def Recessive(self, Chr, GenotypeFil, VEPFil, AFFil, GenecodeFil):
		#GenotypeFil = pysam.TabixFile(GenotypeFil)
		#VEPFil = pysam.TabixFile(VEPFil)
		#AFFil = pysam.TabixFile(AFFil)
		#Genes, Transtripts = LoadGeneCode(GenecodeFil)
		#CSQ_header = [X.strip().split("Format: ")[1].rstrip('>\"').split("|") for X in VEPFil.header if X.startswith("##INFO=<ID=CSQ")][0]
		GenotypeFil = gzip.open(GenotypeFil, 'rb')
		VEPFil = gzip.open(VEPFil, 'rb')
		DIR = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/Filt/"
		OutFil1 = open(DIR+"SPARK30K.Anno.Filt.Chr{}.vcf".format(Chr), 'wb')
		OutFil2 = open(DIR+"SPARK30K.Genotypes.Filt.Chr{}.vcf".format(Chr), 'wb')
		# Skip Header Genotype
		while 1:
			l = GenotypeFil.readline()#.decode("utf-8")
			OutFil2.write(l)
			if l.decode("utf-8").startswith("#C"):
				break
		# Skip Header VEP
		while 1:
			l = VEPFil.readline()
			OutFil1.write(l)
			if l.decode("utf-8").startswith("##INFO=<ID=CSQ"):
				CSQ_header = l.decode("utf-8").split("Format: ")[1].rstrip('>\"').split("|")
			if l.decode("utf-8").startswith("#C"):
				break

		while 1: # iterate through genes
			VEPLine = VEPFil.readline()
			GenotypesLine = GenotypeFil.readline()
			if not VEPLine:
				break
			llist = VEPLine.decode("utf-8").split("\t")
			#llist2 = GenotypesLine.decode("utf-8").split("\t")
			Chr, Pos, Ref, Alts = llist[0], llist[1], llist[3], llist[4]
			infodict = self.getINFO(llist[7])
			Allele_CSQ_dict = self.match_allele_csq(Ref, Alts, CSQ_header, infodict["CSQ"])
			KEEP = False
			for i, Alt in enumerate(Alts.split(",")):
				var_k = "{}:{}:{}:{}".format(Chr, Pos, Ref, Alt)
				Allele_CSQ_dict[Alt][0]["gnomADg_AF_NFE"] = Allele_CSQ_dict[Alt][0]["gnomADg_AF_NFE"].split("&")[0]
				Allele_CSQ_dict[Alt][0]["gnomADe_AF_NFE"] = Allele_CSQ_dict[Alt][0]["gnomADe_AF_NFE"].split("&")[0]
				vep = Allele_CSQ_dict[Alt][0]
				gnomADg_af = 0 if (vep["gnomADg_AF_NFE"] == "" or vep["gnomADg_AF_NFE"] == ".")\
						else float(vep["gnomADg_AF_NFE"])
				gnomADe_af = 0 if (vep["gnomADe_AF_NFE"] == "" or vep["gnomADe_AF_NFE"] == ".")\
						else float(vep["gnomADe_AF_NFE"])
				if gnomADg_af > self.AF_cutoff:
					continue
				idx_anno, cons, trans = self.search_severe_consequence(var_k, Allele_CSQ_dict, Alt)
				#if cons[0] == "synonymous_variant":
				#	print(cons)
				if len(set(cons).intersection(self.LGD))>= 1 or "synonymous_variant" in set(cons) or "missense_variant" in set(cons):
					KEEP = True
			if KEEP:
				OutFil1.write(VEPLine)
				OutFil2.write(GenotypesLine)
		return

def GetOptions():
	parser = argparse.ArgumentParser()
	#parser.add_argument("-v", "--vcf", type=str, required=True, help="<Required> VCF file")
	#parser.add_argument("-l", "--list", type=str, required=True, help="<Required> Indv List file")
	#parser.add_argument("-o", "--out", type=str, help="<Required> Output VCF file")
	parser.add_argument("--chr", type=str, help="<Required> Output VCF file")
	parser.add_argument("--af", type=float, help="<Required> Output VCF file")
	args = parser.parse_args()
	#if args.out == None:
	#	args.out = "test.out.vcf"
	return args

def main():
	args = GetOptions()
	ins = RecessiveModel(args.af)
	#List = [l.strip() for l in open(args.list)]
	#ins.ComputeSiteAF(args.vcf, List, "EUR", args.out)
	Chr = args.chr
	GenotypeFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/GenotypesSplitbyChr/GATK4_20190729.chr{}.vcf.gz".format(Chr)
	VEPFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/VCF/SitesSplitbyChr/annotated/GATK4_20190729.chr{}.mappability.vcf.gz".format(Chr)
	AFFil = "/home/local/users/jw/Genetics_Projects/SPARK/30K_07/recessive/AF/GATK4_20190729.chr{}.eurAF.vcf.gz".format(Chr)
	genecode = "/home/local/users/jw/vep_data/homo_sapiens/GeneCodeV29/CHRs/genecodev29.{}.gtf".format(Chr)
	ins.Recessive(Chr, GenotypeFil ,VEPFil, AFFil, genecode)

if __name__=='__main__':
	main()


import numpy
import simtk.unit
import simtk.openmm as mm

print('OpenMM version: ', mm.version.full_version)
from openmmtools.constants import kB

class CustomizableGHMC(mm.CustomIntegrator):
    """Customizable GHMC: proposal can contain mix of deterministic and stochastic steps.
    If the provided `splitting` is deterministic, then we add the stochastic velocity update to the beginning of the splitting.

    The class of allowable splittings is the same as for the LangevinSplittingIntegrator
    """

    def __init__(self,
                 splitting="V R V",
                 temperature=298.0 * simtk.unit.kelvin,
                 collision_rate=91.0 / simtk.unit.picoseconds,
                 timestep=1.0 * simtk.unit.femtoseconds,
                 constraint_tolerance=1e-4,
                 ):
        """
        Parameters
        ----------
        splitting : string
            Sequence of R, V substeps to be executed each timestep.
            Forces are only used in V-step. Handle multiple force groups by appending the force group index
            to V-steps, e.g. "V0" will only use forces from force group 0. "V" will perform a step using all forces.

        temperature : numpy.unit.Quantity compatible with kelvin, default: 298.0*simtk.unit.kelvin
           Temperature of heat bath

        collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91.0/simtk.unit.picoseconds
           Collision rate

        timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds
           Integration timestep

        constraint_tolerance : float
            Numerical tolerance for solving constraints
        """

        # Compute constants
        kT = kB * temperature
        gamma = collision_rate
        kinetic_energy = "0.5 * m * v * v"

        # Convert splitting string into a list of all-caps strings
        splitting = splitting.upper().split()

        # Count how many times each step appears, so we know how big each R/V/O substep will be
        n_R = sum([letter == "R" for letter in splitting])
        n_V = sum([letter == "V" for letter in splitting])
        n_O = sum([letter == "O" for letter in splitting])

        # If the splitting is deterministic, add a velocity-randomization step to the beginning
        # of the procedure
        if n_O == 0:
            splitting = ["O"] + splitting

        # Check if the splitting string asks for multi-time-stepping.
        # If so, each force group should be integrated for a total length equal to dt
        if len(set([step for step in splitting if step[0] == "V"])) > 1:
            mts = True
            fgs = set([step[1:] for step in splitting if step[0] == "V"])
            n_Vs = dict()
            for fg in fgs:
                n_Vs[fg] = sum([step[1:] == fg for step in splitting])
        else:
            mts = False

        # Do a couple sanity checks on the splitting string
        # Make sure we contain at least one of R, V, O steps
        assert ("R" in splitting)
        assert ("V" in [s[0] for s in splitting])
        assert ("O" in splitting)

        # Make sure it contains no invalid characters
        assert (set(splitting).issubset(set("RVO").union(set(["V{}".format(i) for i in range(32)]))))

        # If the splitting string contains both "V" and a force-group-specific V0,V1,etc.,
        # then raise an error
        if mts and (n_V > 0):
            raise (ValueError("Splitting string includes an evaluation of all forces and "
                              "evaluation of subsets of forces."))

        # Define substep functions
        def R_step():
            # update positions (and velocities, if there are constraints)
            self.addComputePerDof("x", "x + ((dt / {}) * v)".format(n_R))
            self.addComputePerDof("x1", "x")  # save pre-constraint positions in x1
            self.addConstrainPositions()  # x is now constrained
            self.addComputePerDof("v", "v + ((x - x1) / (dt / {}))".format(n_R))
            self.addConstrainVelocities()

        def V_step(fg):
            """Deterministic velocity update, using only forces from force-group fg.
            Parameters
            ----------
            fg : string
                Force group to use in this substep.
                "" means all forces, "0" means force-group 0, etc.
            """
            # update velocities
            if mts:
                self.addComputePerDof("v", "v + ((dt / {}) * f{} / m)".format(n_Vs[fg], fg))
            else:
                self.addComputePerDof("v", "v + (dt / {}) * f / m".format(n_V))
            self.addConstrainVelocities()

        def O_step():
            # measure heat
            self.addComputeSum("old_ke", kinetic_energy)

            # update velocities
            self.addComputePerDof("v", "(a * v) + (b * sqrt(kT / m) * gaussian)")
            self.addConstrainVelocities()

            # measure heat
            self.addComputeSum("new_ke", kinetic_energy)
            self.addComputeGlobal("heat", "heat + (new_ke - old_ke)")

        def substep_function(step_string):
            if step_string == "R":
                R_step()
            elif step_string[0] == "V":
                V_step(step_string[1:])

        def compute_total_energy(name="e_old"):
            self.addComputeSum("ke", kinetic_energy)
            self.addComputeGlobal(name, "ke + energy")

        # Create a new CustomIntegrator
        super(CustomizableGHMC, self).__init__(timestep)

        # Initialize
        self.addGlobalVariable("kT", kT)

        # Velocity mixing parameter: current velocity component
        h = timestep
        self.addGlobalVariable("a", numpy.exp(-gamma * h))

        # Velocity mixing parameter: random velocity component
        self.addGlobalVariable("b", numpy.sqrt(1 - numpy.exp(- 2 * gamma * h)))

        # Positions before application of position constraints
        self.addPerDofVariable("x1", 0)

        # Set constraint tolerance
        self.setConstraintTolerance(constraint_tolerance)

        # Add bookkeeping variables
        self.addGlobalVariable("ke_old", 0)
        self.addGlobalVariable("ke_new", 0)
        self.addGlobalVariable("heat", 0)
        self.addGlobalVariable("shadow_work", 0)
        self.addGlobalVariable("e_old", 0)
        self.addGlobalVariable("e_new", 0)
        self.addGlobalVariable("acc_ratio", 0)
        self.addGlobalVariable("accept", 0)
        self.addGlobalVariable("naccept", 0)
        self.addGlobalVariable("ntrials", 0)
        self.addPerDofVariable("xold", 0)
        self.addPerDofVariable("vold", 0)

        # Compute energy of current state
        self.addUpdateContextState()
        self.addComputePerDof("xold", "x")
        self.addComputePerDof("vold", "v")
        compute_total_energy("e_old")

        # Reset "heat" to zero
        self.addComputeGlobal("heat", "0")

        # Integrate / generate proposal
        for i, step in enumerate(splitting):
            substep_function(step)

        # Compute M-H ratio in terms of energy change during the deterministic steps
        compute_total_energy("e_new")
        self.addComputeGlobal("shadow_work", "(e_new - e_old) - heat")
        self.addComputeGlobal("acc_ratio", "exp(- shadow_work / kT)")

        # Accept / reject : flip momenta upon rejection
        self.addComputeGlobal("accept", "step(acc_ratio - uniform)")
        self.beginIfBlock("accept != 1")
        self.addComputePerDof("x", "xold")
        self.addComputePerDof("v", "-vold")
        self.endBlock()

        # Update acceptance statistics
        self.addComputeGlobal("naccept", "naccept + accept")
        self.addComputeGlobal("ntrials", "ntrials + 1")


"""
CNN-DQN based RL agent for Pong

Requirements (model training):
- pytorch
- gym
- matplotlib
- numpy
- pillow
- opencv-python
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np
import math
import random
import os
import sys
from collections import namedtuple

class Agent:

    NAME = 'AgentPongai'
    MODEL_PATH = 'models/'
    TRAINED_MODEL = 'model.mdl'
    N_ACTIONS = 3
    FRAME_WIDTH = 200
    FRAME_HEIGHT = 200
    EPS_START = 0.99
    EPS_END = 0.05
    EPS_DECAY = 150000
    LR = 1e-4

    # Common interface for both training and testing
    # supports 0 (frames are summed to each other rather than stacked), 2 or 4 frame stacking

    def __init__(   self, training=False, batch_size=64, gamma=0.99, memory_size=15000, \
                    stack_frames=4, downscale=True, priority_memory=True):
        self.training = training
        self.batch_size = batch_size
        self.gamma = gamma
        self.stack_frames = stack_frames if stack_frames in (0, 2, 4) else 0
        self.downscale = downscale
        self.priority_memory = priority_memory
        self.w = Agent.FRAME_WIDTH // 2 if downscale else Agent.FRAME_WIDTH
        self.h = Agent.FRAME_HEIGHT // 2 if downscale else Agent.FRAME_HEIGHT
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        in_channels = stack_frames if stack_frames in (2, 4) else 1
        self.policy_net = DQN(Agent.N_ACTIONS, in_channels, self.h, self.w).to(self.device)
        self.target_net = DQN(Agent.N_ACTIONS, in_channels, self.h, self.w).to(self.device)
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        self.memory = PriorityMemory(memory_size) if priority_memory else ReplayMemory(memory_size)
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=Agent.LR)
        self.epsilon = 1.0
        self.prev_state = [None] if self.stack_frames in (0, 2) else [None, None, None]
        self.steps_done = 0

    """
    Load the trained model for evaluation
    """
    def load_model(self):
        self.policy_net.load_state_dict(torch.load(Agent.TRAINED_MODEL, map_location=self.device))
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        self.policy_net.eval()

    """
    This needs to correspond to save
    """
    def load_trained_model(self):
        dir = os.path.join(os.getcwd(), Agent.MODEL_PATH, self.get_name())
        if not os.path.isdir(dir):
            sys.exit(f'Model loading failed: {dir} does not exist')
        model_files = os.listdir(dir)
        model_files.sort()
        model_path = os.path.join(dir, model_files[-1]) # Select the latest model
        self.policy_net.load_state_dict(torch.load(model_path, map_location=self.device))
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=Agent.LR)

    """
    Load a trained model from a specific path
    """
    def load_model_file(self, model_path):
        self.policy_net.load_state_dict(torch.load(model_path, map_location=self.device))
        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()
        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=Agent.LR)

    """
    Reset Agent when an episode has finished
    """
    def reset(self):
        self.prev_state = [None] if self.stack_frames in (0, 2) else [None, None, None]

    """
    frame: numpy array (200, 200, 3) -> convert to pytorch tensor
    """
    def get_action(self, frame=None):
        self.epsilon = Agent.EPS_END + (Agent.EPS_START - Agent.EPS_END) * math.exp(-1. * self.steps_done / Agent.EPS_DECAY)
        self.steps_done += 1
        action = torch.randint(0, Agent.N_ACTIONS, (1,)).item()
        if frame is None:
            return action
        state = self.__preprocess_frame(frame)
        if self.prev_state[0] != None and (not self.training or np.random.uniform(size=1)[0] > self.epsilon):
            # Select the 'best' action
            with torch.no_grad():
                state_diff = self.__process_states(state, self.prev_state)
                action = torch.argmax(self.policy_net(state_diff)).item()
        if len(self.prev_state) == 3:
            # 4 frames stacking enabled
            self.prev_state[0] = self.prev_state[1]
            self.prev_state[1] = self.prev_state[2]
        self.prev_state[-1] = state
        return action

    """
    Get Agent name, used by the environment
    """
    def get_name(self):
        return Agent.NAME

    """
    Save a Pytorch model to MODEL_PATH
    """
    def save(self):
        dir = os.path.join(os.getcwd(), Agent.MODEL_PATH, self.get_name())
        if os.path.isdir(dir):
            # This does not work for more than 100 models
            content = os.listdir(dir)
            content.sort()
            model_id = content[-1][-6:-4] if len(content) > 0 else -1 # Get the latest model number
            model_id = '%02d' % (int(model_id) + 1, ) # Increase it by one
        else:
            os.makedirs(dir)
            model_id = '00'
        model_path = os.path.join(dir, self.get_name() + model_id + '.pth')
        torch.save(self.policy_net.state_dict(), model_path)

    """
    Convert a rbg image frame given as a numpy array to a grayscale tensor
    frame: np.array with shape(200, 200, 3) containing int values 0 - 255
    """
    def __preprocess_frame(self, frame):
        frame = frame.astype(float)
        frame = np.true_divide(frame, 255.0)
        frame = 1.0 - frame # invert colors so that the background is zeros
        grayscale = 0.299 * frame[:,:,0] + 0.587 * frame[:,:,1] + 0.114 * frame[:,:,2]
        if self.downscale:
            grayscale = grayscale.reshape((1, Agent.FRAME_HEIGHT//2, 2, Agent.FRAME_WIDTH//2, 2)).max(4).max(2) # downscale by factor of 2
        return torch.from_numpy(grayscale).to(torch.float32).to(self.device)

    """
    Combine two consequtive states to create a combined state that shows movement.
    This is the state information that is feed on the DQN and stored.
    state: torch float32 tensor with values in the range of 0.0 and 1.0
    prev_state: torch float32 tensor with values in the range of 0.0 and 1.0
    """
    def __state_diff(self, state, prev_state):
        combined = torch.add(state, 0.5 * prev_state).view(1, 1, self.h, self.w) # multiply prev_state to distinguish differences
        combined -= torch.min(combined)
        combined /= torch.max(combined)
        return combined

    """
    Stack two consequtive states to create a combined state that shows movement.
    This is the state information that is feed on the DQN and stored.
    states: torch float32 tensor with values in the range of 0.0 and 1.0 as a list
    """
    def __stack_states(self, states):
        stacked = torch.stack(states).view(1, self.stack_frames, self.h, self.w)
        return stacked

    """
    This is just a wrapper for processing state and the previous state
    """
    def __process_states(self, state, prev_states):
        return self.__stack_states([state] + prev_states) if self.stack_frames > 0 else self.__state_diff(state, prev_states[0])

    # Training interface

    """
    Store actions, states, rewards for replay
    Note: state and next_state need to be image frame differences as tensors
    observations: 3 latest observations (frames)
    """
    def update_memory(self, observations, action, reward, done):
        action = torch.tensor([[action]], dtype=torch.long, device=self.device)
        reward = torch.tensor([reward], dtype=torch.float32, device=self.device)
        states = [self.__preprocess_frame(observation) for observation in observations]
        state = self.__process_states(states[-2], states[:-2])
        next_state = self.__process_states(states[-1], states[1:-1])
        done = torch.tensor([done], dtype=torch.bool, device=self.device)
        self.memory.push(state, action, next_state, reward, done)

    """
    Do a single batch update on the DQN network
    """
    def update_dqn(self):
        if len(self.memory) < self.batch_size:
            return
        if self.priority_memory:
            samples, indices, weights = self.memory.sample(self.batch_size)
            weights = torch.tensor(weights, dtype=torch.float32, device=self.device)
        else:
            samples = self.memory.sample(self.batch_size)
        batch = Transition(*zip(*samples))
        state_batch = torch.cat(batch.state)
        next_state_batch = torch.cat(batch.next_state)
        action_batch = torch.cat(batch.action)
        reward_batch = torch.cat(batch.reward)
        done_batch = torch.cat(batch.done)
        non_final_mask = [done_batch == False]
        non_final_next_states = next_state_batch[non_final_mask]

        # Compute Q(s, a) by selecting Q values matching selected actions
        state_action_values = self.policy_net(state_batch).gather(1, action_batch).squeeze(1)
        # Compute V(s_t+1)
        next_state_values = torch.zeros(self.batch_size, device=self.device)
        next_state_values[non_final_mask] = self.target_net(non_final_next_states).max(1)[0].detach()
        # Compute target Q values
        target_state_action_values = reward_batch + self.gamma * next_state_values

        # Compute loss and optimize the network
        loss = F.smooth_l1_loss(state_action_values, target_state_action_values, reduction='none')
        if self.priority_memory:
            priorities = torch.abs(loss).detach() + 1e-5
            loss = loss * weights
            self.memory.update_priorities(indices, priorities.cpu().numpy())

        self.optimizer.zero_grad()
        loss = loss.mean()
        loss.backward()
        for param in self.policy_net.parameters():
            param.grad.data.clamp_(-10, 10) # to make stable updates
        self.optimizer.step()

    """
    Double DQN: https://arxiv.org/pdf/1509.06461.pdf
    """
    def update_ddqn(self):
        if len(self.memory) < self.batch_size:
            return
        if self.priority_memory:
            samples, indices, weights = self.memory.sample(self.batch_size)
            weights = torch.tensor(weights, dtype=torch.float32, device=self.device)
        else:
            samples = self.memory.sample(self.batch_size)
        batch = Transition(*zip(*samples))
        state_batch = torch.cat(batch.state)
        next_state_batch = torch.cat(batch.next_state)
        action_batch = torch.cat(batch.action)
        reward_batch = torch.cat(batch.reward)
        done_batch = torch.cat(batch.done)
        non_final_mask = [done_batch == False]
        non_final_next_states = next_state_batch[non_final_mask]

        # Compute Q(s, a) by selecting Q values matching selected actions
        state_action_values = self.policy_net(state_batch).gather(1, action_batch).squeeze(1)
        # Compute V(s_t+1)
        next_state_values = torch.zeros(self.batch_size, device=self.device)
        greedy_actions = self.policy_net(non_final_next_states).max(1)[1].unsqueeze(1).detach()
        next_state_values[non_final_mask] = self.target_net(non_final_next_states).gather(1, greedy_actions).squeeze(1).detach()

        # Compute target Q values
        target_state_action_values = reward_batch + self.gamma * next_state_values

        # Compute loss and optimize the network
        loss = F.smooth_l1_loss(state_action_values, target_state_action_values, reduction='none')
        if self.priority_memory:
            priorities = torch.abs(loss).detach() + 1e-5
            loss = loss * weights
            self.memory.update_priorities(indices, priorities.cpu().numpy())

        self.optimizer.zero_grad()
        loss = loss.mean()
        loss.backward()
        for param in self.policy_net.parameters():
            param.grad.data.clamp_(-10, 10) # to make stable updates
        self.optimizer.step()

    """
    Update the target network to match the policy network
    Call this periodically from the training loop
    """
    def update_target_network(self):
        self.target_net.load_state_dict(self.policy_net.state_dict())

    """
    Set agent learning rate
    """
    def set_lr(self, lr):
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

"""
CNN-based neural network which is updated according to DQN update rules
"""
class DQN(nn.Module):

    def __init__(self, n_actions, in_channels, h=Agent.FRAME_HEIGHT, w=Agent.FRAME_WIDTH):
        super(DQN, self).__init__()
        # Conv2d: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d
        padding, dilation = 0, 1
        strides = [4, 2, 1]
        kernels = [(8, 8), (4, 4), (3, 3)]
        channels = [in_channels, 32, 64, 64]
        def out_size(in_size, kernel_size, stride):
            return math.floor((in_size + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1)

        self.conv1 = nn.Conv2d(channels[0], channels[1], kernels[0], padding=padding, dilation=dilation, stride=strides[0])
        self.conv2 = nn.Conv2d(channels[1], channels[2], kernels[1], padding=padding, dilation=dilation, stride=strides[1])
        self.conv3 = nn.Conv2d(channels[2], channels[3], kernels[2], padding=padding, dilation=dilation, stride=strides[2])
        h_out = out_size(out_size(out_size(h, kernels[0][0], strides[0]), kernels[1][0], strides[1]), kernels[2][0], strides[2])
        w_out = out_size(out_size(out_size(w, kernels[0][1], strides[0]), kernels[1][1], strides[1]), kernels[2][1], strides[2])
        flattened_size = h_out * w_out * channels[-1]
        self.fc1 = nn.Linear(flattened_size, 512)
        self.fc2 = nn.Linear(512, n_actions)

    """
    x : (batch_size, in_channels, image_width, image_height) = a batch of  in_channels images as a tensor
    """
    def forward(self, x):
        y = F.relu(self.conv1(x))
        y = F.relu(self.conv2(y))
        y = F.relu(self.conv3(y))
        y = torch.flatten(y, start_dim=1)
        y = F.relu(self.fc1(y))
        y = self.fc2(y)
        return y

"""
Directly from Pytorch example implementation for replay memory
https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html
"""
Transition = namedtuple('Transition',
                    ('state', 'action', 'next_state', 'reward', 'done'))
class ReplayMemory():
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def update_beta(self):
        pass

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

"""
https://arxiv.org/pdf/1511.05952.pdf
"""
class PriorityMemory():
    def __init__(self, capacity, alpha=0.6, beta0=0.4):
        self.capacity = capacity
        self.memory = []
        self.priorities = np.zeros((capacity,), dtype=np.float32)
        self.position = 0
        self.alpha = alpha
        self.beta0 = beta0
        self.beta = beta0

    def push(self, *args):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        # store transition with maximal priority
        self.memory[self.position] = Transition(*args)
        self.priorities[self.position] = 1.0 if len(self.memory) == 1 else np.max(self.priorities)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        N = len(self.memory)
        priorities = self.priorities[:N]
        P = np.power(priorities, self.alpha)
        P /= np.sum(P)

        indices = np.random.choice(N, size=batch_size, p=P, replace=False)
        samples = [self.memory[i] for i in indices]

        w = np.power(N * P[indices], -self.beta)
        w /= np.max(w)
        return samples, indices, w

    def update_priorities(self, indices, priorities):
        for idx, priority in zip(indices, priorities):
            self.priorities[idx] = priority

    def update_beta(self, interval=8000):
        self.beta = min(1.0, self.beta + (1.0 - self.beta0) / interval)

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

## Testing ##

def test_DQN():
    n_actions = 3
    dqn = DQN(n_actions, 1, h=200, w=200)
    test = torch.ones((1, 1, 200, 200))
    assert dqn(test).size() == torch.Size([1, n_actions])
    test2 = torch.zeros((13, 1, 200, 200))
    assert dqn(test2).size() == torch.Size([13, n_actions])

def test_Agent():
    agent = Agent()
    frame1 = np.ones((200, 200, 3))
    frame2 = np.ones((200, 200, 3))
    assert type(agent.get_action(frame1)) == int
    assert agent.get_action(frame2) in [i for i in range(Agent.N_ACTIONS)]

if __name__ == '__main__':
    test_DQN()
    test_Agent()


#!/usr/bin/env python

'''
    Make cross section.
'''

import numpy as np
import matplotlib.pyplot as plt

# our own functions.
from makeinit import makeinit
from freeman import rk4int
from makesos import sos

if __name__ == '__main__':

    # set model parameters.
    A, B, W = 1., 2., 0.25 # slow bar,
    ic_radius = 10. # maximal radius for init condition.
    N_pts = 8
    EJ_t = 3.

    # generate initial condition.
    ic_pts = makeinit(A, B, W, ic_radius, N_pts, EJ_t)

    # integrate to t=100, 8001 steps.
    orb = rk4int(ic_pts, A * A, B * B, W, 0.05, 64001)

    for orb_i in orb:
        k = sos(orb_i, axis=0, direction='+')
        plt.scatter(k[:, 1], k[:, 3], s=2)
    plt.show()


#!/usr/env/bin python

r"""
Compute integration metrics for paired samples
"""

import argparse
import functools
import pathlib

import anndata
import numpy as np
import pandas as pd
import yaml

import scglue
import scglue.metrics

def parse_args() -> argparse.Namespace:
    r"""
    Parse command line arguments
    """
    parser = argparse.ArgumentParser(
        description="Compute integration metrics for paired samples"
    )
    parser.add_argument(
        "-d", "--datasets", dest="datasets", type=pathlib.Path, required=True,
        nargs="+", help="Path to datasets (.h5ad)"
    )
    parser.add_argument(
        "-l", "--latents", dest="latents", type=pathlib.Path, required=True,
        nargs="+", help="Path to latent embeddings (.csv)"
    )
    parser.add_argument(
        "--cell-type", dest="cell_type", type=str, default="cell_type",
        help="Column name in obs specifying cell types"
    )
    parser.add_argument(
        "--domain", dest="domain", type=str, default="domain",
        help="Column name in obs specifying domain"
    )
    parser.add_argument(
        "-p", "--paired", dest="paired", default=False, action="store_true",
        help="Whether the latent embeddings are paired"
    )
    parser.add_argument(
        "-r", "--run-info", dest="run_info", type=pathlib.Path, required=True,
        help="Path to run info file (.yaml)"
    )
    parser.add_argument(
        "-o", "--output", dest="output", type=pathlib.Path, required=True,
        help="Path to output file (.yaml)"
    )
    return parser.parse_args()

def main(args: argparse.Namespace) -> None:
    r"""
    Main function
    """
    if len(args.datasets) != len(args.latents):
        raise RuntimeError("Datasets and latents should have the same number of entries!")

    print("[1/3] Reading data...")
    datasets = [anndata.read_h5ad(item, backed="r") for item in args.datasets]
    cell_types = [dataset.obs[args.cell_type].to_numpy() for dataset in datasets]
    domains = [dataset.obs[args.domain].to_numpy() for dataset in datasets]
    latents = [pd.read_csv(item, header=None, index_col=0).to_numpy() for item in args.latents]
    with args.run_info.open("r") as f:
        run_info = yaml.load(f, Loader=yaml.Loader)

    print("[2/3] Computing metrics...")
    masks = [np.apply_along_axis(lambda x: ~np.any(np.isnan(x)), 1, latent) for latent in latents]
    for i, mask in enumerate(masks):
        rm_pct = 100 * (1 - mask.sum() / mask.size)
        if rm_pct:
            print(f"Ignoring {rm_pct:.1f}% cells in dataset {i} due to missing values!")
    combined_cell_type = np.concatenate([cell_type[mask] for cell_type, mask in zip(cell_types, masks)])
    combined_domain = np.concatenate([domain[mask] for domain, mask in zip(domains, masks)])
    combined_latent = np.concatenate([latent[mask] for latent, mask in zip(latents, masks)])
    metrics = {
        "time": run_info["time"],
        "n_cells": run_info["n_cells"],
        "mean_average_precision": scglue.metrics.mean_average_precision(
            combined_latent, combined_cell_type
        ),
        "seurat_alignment_score": scglue.metrics.seurat_alignment_score(
            combined_latent, combined_domain, random_state=0
        )
    }
    if args.paired:
        if len(datasets) != 2:
            raise RuntimeError("Expect exactly two datasets in paired mode!")
        mask = functools.reduce(np.logical_and, masks)
        rm_pct = 100 * (1 - mask.sum() / mask.size)
        if rm_pct:
            print(f"Ignoring {rm_pct:.1f}% cells in all datasets due to missing values!")
        metrics["foscttm"] = np.concatenate(
            scglue.metrics.foscttm(*[latent[mask] for latent in latents])
        ).mean().item()
    else:
        metrics["foscttm"] = None

    print("[3/3] Saving results...")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w") as f:
        yaml.dump(metrics, f)

if __name__ == "__main__":
    main(parse_args())


import argparse
import os
import time
import onnx
import torch
import requests
from io import BytesIO
from PIL import Image, ImageOps
import cv2
from torchvision import transforms
from torchvision.transforms.functional import to_tensor, to_pil_image
import onnxruntime as ort
import numpy as np

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', type=str, default='')
    parser.add_argument('--filename', type=str, default=f"{time.time()}.jpg")
    args = parser.parse_args()

    project_path = "/home/alotaima/Projects/side/onlysudo/src/ai/anime"
    args.outdir = '/src/server/public/ai/anime'

    root_path = '/'.join(os.path.abspath(os.getcwd()).split('/')[:-1])
    output_path = f"{root_path}{args.outdir}/{args.filename}"

    if args.url == '':
        print('Need url or streamer name!')
        exit()

    args.size = 1024
    side_by_side = False
    device = "cuda"

    print('start')
    # # Load the ONNX model
    # model = onnx.load("animegan2_generator.onnx")
    # # Check that the model is well formed
    # onnx.checker.check_model(model)
    ort_session = ort.InferenceSession("animegan2_generator.onnx")

    response = requests.get(args.url, stream = True)
    img = Image.open(BytesIO(response.content)).convert("RGB")
    img = ImageOps.fit(img, (args.size, args.size), centering=(0.5, 0.5))

    w, h = img.size
    s = min(w, h)
    img = img.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
    img = img.resize((args.size, args.size), Image.LANCZOS)

    input = to_tensor(img).unsqueeze(0) * 2 - 1

    output = ort_session.run(
        None,
        {"actual_input_1": input.cpu().detach().numpy()},
    )[0]
    output = torch.tensor(output)[0]

    if side_by_side:
        output = torch.cat([input[0], output], dim=2)

    output = (output * 0.5 + 0.5).clip(0, 1)

    out = to_pil_image(output)

    out = np.concatenate([img, out], axis=1)

    cv2.imwrite(output_path, cv2.cvtColor(out, cv2.COLOR_BGR2RGB))

    print(f"Done! {output_path}")


#!/usr/bin/env python
from __future__ import print_function

import roslib
# roslib.load_manifest('my_package')

import rospy
from sensor_msgs.msg import CameraInfo
from sensor_msgs.msg import Image as msg_Image
from cv_bridge import CvBridge, CvBridgeError
import sys
import os

from numpy import median

from darknet_ros_msgs.msg import BoundingBox, BoundingBoxes, Artifact

import yaml

import time

class LocateArtifact:
    def __init__(self, camera_info_topic, depth_image_topic, color_image_topic, pose_topic, pub_color_topic, bounding_boxes_topic, artifact_topic):
        """
        Basic logic:
        """

        self.camera_info = None
        self.depth = None
        self.depth_keep = None
        self.pose_keep = None
        self.artifact = Artifact()
        self.bridge = CvBridge()
        self.can_publish_color = True

        rospy.Subscriber(camera_info_topic, CameraInfo, callback=self.callback_camera)

        rospy.Subscriber(depth_image_topic, msg_Image, callback=self.callback_depth)

        rospy.Subscriber(color_image_topic, msg_Image, callback=self.callback_color)

        #rospy.Subscriber(pose_topic, odometry_msgs, callback=self.callback_pose)

        rospy.Subscriber(bounding_boxes_topic, BoundingBoxes, callback=self.callback_boundingbox)

        self.pub_color_image = rospy.Publisher(pub_color_topic, msg_Image, queue_size=1)

        self.pub_artifact = rospy.Publisher(artifact_topic, Artifact, queue_size=1)

        # to ensure the color message can be resent after a time threshold
        self.start = time.time()
        self.time_threshold = 1 # 5 seconds

    def callback_depth(self, data):
        print ('recieved depth message')
        self.depth = data

    def callback_camera(self, data):
        print ('recieved camera message')
        self.camera_info = data

    #def callback_pose(self, data):
        #print ('recieved pose message')
        #self.pose = data

    def callback_color(self, data):
        print ('recieved color message')
        # record the depth, pose, color data with the 'same' timestamp
        if self.depth is None:
            print ('depth is None')
            return

        #print (self.can_publish_color)

        if not self.can_publish_color:
            end = time.time()
            if (end - self.start) > self.time_threshold:
                self.can_publish_color = True

        if self.can_publish_color:
            self.pub_color_image.publish(data)
            self.can_publish_color = False
            self.depth_keep = self.depth
            #self.pose_keep = self.pose

            self.start = time.time()
            print ('published color message')

    def callback_boundingbox(self, data):
        if self.depth_keep is None or self.camera_info is None:
            return

        cv_image = self.bridge.imgmsg_to_cv2(self.depth_keep, self.depth_keep.encoding)

        fx = self.camera_info.K[0]
        fy = self.camera_info.K[4]
        cx = self.camera_info.K[2]
        cy = self.camera_info.K[5]

        # iterate through the bounding boxes
        for box in data.bounding_boxes:
            # get the center of the box
            object_type = box.Class
            if object_type == 'none':
                continue

            prob = box.probability
            u = int( (box.xmin + box.xmax)/2 )
            v = int( (box.ymin + box.ymax)/2 )

            # find the depth of the box
            # depth = median of a grid of depth values around the center
            # depth = self.depth[x, y]
            z_list = []
            for i in range( max(0, u-20), min(u+20, self.depth_keep.width) ):
                for j in range( max(0, v-20), min(v+20, self.depth_keep.height) ):
                    z_list.append(cv_image[j, i])

            z = median(z_list)/1e3 # mm --> meter

            # calculate the (x, y, z) of the box
            x = (u-cx)*z/fx
            y = (v-cy)*z/fy

            print (object_type, prob, u, v)

            print ('Object position', x, y, z)

            # sometimes there is invalid estimate to be (0, 0, 0)
            if (x*x + y*y + z*z) < 1e-4:
                continue

            # publish an Artifact message with object type, probability, and (x, y, z)
            self.artifact.obj_class = object_type
            self.artifact.obj_prob = prob
            self.artifact.position.x = x
            self.artifact.position.y = y
            self.artifact.position.z = z
            self.pub_artifact.publish(self.artifact)

        self.can_publish_color = True

if __name__ == '__main__':
    if len(sys.argv) < 8:
        print ('Usage: \n\t locate_artifacts.py image_topic depth_topic camera_info_topic object_pos_topic')
    try:
        rospy.init_node('locate_artifacts')
        print ('arguments:', sys.argv)

        la = LocateArtifact(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[6], sys.argv[7])

        rospy.spin()

    except KeyboardInterrupt:
        print("Shutting down")