"""
This class is responsible for the whole experimentals
"""

# TODO: implement unfinished
import argparse
import colorsys
import datetime
import os
from PIL import Image, ImageFont, ImageDraw
import pandas as pd
import matplotlib.pyplot as plt

import numpy as np
import keras.backend as K
from keras.layers import Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
from keras.utils import multi_gpu_model

from utils.io_utils import write_csv
from yolo3.model import preprocess_true_boxes, yolo_body, yolo_loss, yolo_eval
from yolo3.utils import get_random_data, letterbox_image

from utils.coco_tools import ExportGroundtruthToCOCO, ExportDetectionsToCOCO, COCOWrapper, COCOEvalWrapper

def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
    '''data generator for fit_generator'''
    n = len(annotation_lines)
    i = 0
    while True:
        image_data = []
        box_data = []
        for b in range(batch_size):
            if i == 0:
                np.random.shuffle(annotation_lines)
            image, box = get_random_data(annotation_lines[i], input_shape, random=True)
            image_data.append(image)
            box_data.append(box)
            i = (i + 1) % n
        image_data = np.array(image_data)
        box_data = np.array(box_data)
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)

def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
    n = len(annotation_lines)
    if n == 0 or batch_size <= 0: return None
    return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)

class Trainer(object):
    _defaults = {
        "score": 0.3,
        "iou": 0.45,
        "gpu_num": 1,
    }

    def __init__(self,
                 train_annotation_path='train.txt',
                 valid_annotation_path='valid.txt',
                 log_dir='logs/000/',
                 output_dir='outputs/000',
                 classes_path='model_data/coco_classes.txt',
                 anchors_path='model_data/yolo_anchors.txt',
                 input_shape=(416, 416),
                 gpu_num=1
                 ):
        self.train_annotation_path = train_annotation_path
        self.valid_annotation_path = valid_annotation_path
        self.log_dir = log_dir
        self.output_dir = output_dir
        self.classes_path = classes_path
        self.anchors_path = anchors_path
        self.class_names = self.get_classes()
        self.num_classes = len(self.class_names)
        self.anchors = self.get_anchors()

        if not os.path.exists(log_dir):
            os.makedirs(log_dir)

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        self.training_log_file = open(os.path.join(self.log_dir, "progress.csv"), "w")

        self.categories_dict = [{'id': i, 'name': name} for i, name in enumerate(self.class_names)] # for COCO API

        # Generate colors for drawing bounding boxes.
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        np.random.seed(None)  # Reset seed to default.

        self.sess = K.get_session()

        self.input_shape = input_shape
        self.gpu_num = gpu_num
        self.score = 0.3
        self.iou = 0.45
        self.ready_for_model = False
        self.ready_for_data = False
        self.is_trained = False

    def print_categories_dict(self): # for debugging
        print(self.categories_dict)

    def get_classes(self):
        '''loads the classes'''
        with open(self.classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def get_anchors(self):
        '''loads the anchors from a file'''
        with open(self.anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def create_model(self, load_pretrained=True, freeze_body=2, weights_path='model_data/yolo_weights.h5'):
        '''create the training model'''
        # K.clear_session()  # get a new session
        image_input = Input(shape=(None, None, 3))
        h, w = self.input_shape
        num_anchors = len(self.anchors)

        y_true = [Input(shape=(h // {0: 32, 1: 16, 2: 8}[l], w // {0: 32, 1: 16, 2: 8}[l],
                               num_anchors // 3, self.num_classes + 5)) for l in range(3)]

        self.model_body = yolo_body(image_input, num_anchors // 3, self.num_classes)
        print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, self.num_classes))

        if self.gpu_num>=2:
            self.model_body = multi_gpu_model(self.model_body, gpus=self.gpu_num)

        # Generate YOLOv3 core model for learning
        if load_pretrained:
            self.model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
            print('Load weights {}.'.format(weights_path))

            if freeze_body in [1, 2]:
                # Freeze darknet53 body or freeze all but 3 output layers.
                num = (185, len(self.model_body.layers) - 3)[freeze_body - 1]
                for i in range(num): self.model_body.layers[i].trainable = False
                print('Freeze the first {} layers of total {} layers.'.format(num, len(self.model_body.layers)))

        model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
                            arguments={'anchors': self.anchors, 'num_classes': self.num_classes, 'ignore_thresh': 0.5})(
            [*self.model_body.output, *y_true])
        model = Model([self.model_body.input, *y_true], model_loss)

        # Generate output tensor targets for filtered bounding boxes
        self.input_image_shape = K.placeholder(shape=(2,))

        boxes, scores, classes = yolo_eval(self.model_body.output, self.anchors,
                len(self.class_names), self.input_image_shape,
                score_threshold=self.score, iou_threshold=self.iou)

        self.ready_for_model = True
        return model, boxes, scores, classes

    def make_detection_list(self, annotation_lines):
        assert self.ready_for_data, "Ready for validation data first"
        image_ids = []
        detection_boxes = []
        detection_scores = []
        detection_classes = []
        for line in annotation_lines:
            line = line.split()
            image_id = line[0]
            image = Image.open(line[0])
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
            image_data = np.array(boxed_image, dtype='float32')
            # preprocessing image
            image_data /= 255.
            image_data = np.expand_dims(image_data, 0)  # Add batch dimension.
            # inference
            out_boxes, out_scores, out_classes = self.sess.run(
                [self.boxes, self.scores, self.classes],
                feed_dict={
                    self.model_body.input: image_data,
                    self.input_image_shape: [image.size[1], image.size[0]],
                    K.learning_phase(): 0
                })
            image_ids.append(image_id)
            detection_scores.append(np.array(out_scores))
            detection_classes.append(np.array(out_classes))
            if len(out_boxes) == 0:
                detection_boxes.append(np.zeros((0, 4)))
            else:
                detection_boxes.append(
                    np.array([np.array([left, top, right, bottom]) for top, left, bottom, right in out_boxes])
                )
            del image

        return image_ids, detection_boxes, detection_scores, detection_classes

    def make_groundtruth_list(self, annotation_lines):
        assert self.ready_for_data, "Ready for validation data first"
        image_ids = []
        detection_boxes = []
        detection_classes = []
        for line in annotation_lines:
            line = line.split()
            image_id = line[0]
            line_boxes_ids = [l.split(",") for l in line[1:]]
            # print(line[1:])
            if len(line_boxes_ids) == 0:
                boxes = np.zeros((0, 4))
            else:
                boxes = np.array([np.array([float(xmin), float(ymin), float(xmax), float(ymax)])
                                  for xmin, ymin, xmax, ymax, _ in line_boxes_ids])
            classes = np.array([int(class_id) for _, _, _, _, class_id in line_boxes_ids])
            image_ids.append(image_id)
            detection_boxes.append(boxes)
            detection_classes.append(classes)
        return image_ids, detection_boxes, detection_classes

    def train(self, freeze_body=2, weights_path='model_data/yolo_weights.h5', batch_size=32, epochs=1000):

        self.model, self.boxes, self.scores, self.classes = self.create_model(
            load_pretrained=True,
            freeze_body=freeze_body,
            weights_path=weights_path
        )

        # logging = TensorBoard(log_dir=self.log_dir)
        # checkpoint = ModelCheckpoint(self.log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
        #                              monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
        # reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
        # early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)

        # load annotation data as list
        with open(self.train_annotation_path) as f:
            self.train_lines = f.readlines()
        with open(self.valid_annotation_path) as f:
            self.valid_lines = f.readlines()

        self.ready_for_data = True

        # shuffle annotation data for train
        np.random.seed(10101)
        np.random.shuffle(self.train_lines)
        np.random.seed(None)

        num_train = len(self.train_lines)
        num_valid = len(self.valid_lines)

        self.model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        self.train_data_generator = data_generator_wrapper(self.train_lines,
                                                      batch_size,
                                                      self.input_shape,
                                                      self.anchors,
                                                      self.num_classes)
        self.valid_data_generator = data_generator_wrapper(self.valid_lines,
                                                      batch_size,
                                                      self.input_shape,
                                                      self.anchors,
                                                      self.num_classes)

        first_valid_result = self.model.evaluate_generator(self.valid_data_generator, max(1, num_valid // batch_size))
        print("[Epoch 0] (at the pretrained stage) [validation loss : %.4f]" % first_valid_result)
        self.evaluate(0)
        self.sample_detection(0)

        start_time = datetime.datetime.now()
        steps_per_epoch = max(1, num_train // batch_size)
        for epoch in range(epochs):
            epoch += 1
            #K.set_learning_phase(1) # for training
            for batch_i in range(steps_per_epoch):
                batch_i += 1
                # train manually
                X, y = next(self.train_data_generator)
                train_loss = self.model.train_on_batch(X, y)

                elapsed_time = datetime.datetime.now() - start_time
                print("[Epoch %d] [Batch %d/%d] [YOLO loss: %f] time: %s"
                      % (epoch, batch_i, steps_per_epoch, train_loss, elapsed_time))

            valid_loss = self.model.evaluate_generator(self.valid_data_generator, max(1, num_valid // batch_size))
            print("[Epoch %d] [validation loss: %f]" % (epoch, valid_loss))
            self.evaluate(epoch)
            #self.sample_detection(epoch)

            if not os.path.exists("{}/models".format(self.output_dir)):
                os.makedirs("{}/models".format(self.output_dir), exist_ok=True)

            self.model.save_weights(os.path.join(self.output_dir, "models/ep%03d-val_loss%.4f.weights" % (epoch, valid_loss)))

    def train_automatically(self, freeze_body=2, weights_path='model_data/yolo_weights.h5', batch_size=32):

        self.model, self.boxes, self.scores, self.classes = self.create_model(
            load_pretrained=True,
            freeze_body=freeze_body,
            weights_path=weights_path
        )

        logging = TensorBoard(log_dir=self.log_dir)
        checkpoint = ModelCheckpoint(self.log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
                                     monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
        reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
        early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)

        # load annotation data as list
        with open(self.train_annotation_path) as f:
            self.train_lines = f.readlines()
        with open(self.valid_annotation_path) as f:
            self.valid_lines = f.readlines()

        self.ready_for_data = True

        # shuffle annotation data for train
        np.random.seed(10101)
        np.random.shuffle(self.train_lines)
        np.random.seed(None)

        num_train = len(self.train_lines)
        num_valid = len(self.valid_lines)

        self.model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        self.train_data_generator = data_generator_wrapper(self.train_lines,
                                                           batch_size,
                                                           self.input_shape,
                                                           self.anchors,
                                                           self.num_classes)
        self.valid_data_generator = data_generator_wrapper(self.valid_lines,
                                                           batch_size,
                                                           self.input_shape,
                                                           self.anchors,
                                                           self.num_classes)

        first_valid_result = self.model.evaluate_generator(self.valid_data_generator, max(1, num_valid // batch_size))
        print("Epoch 0 (at the pretrained stage) [validation loss : %.4f]" % first_valid_result)
        self.evaluate(0)
        self.sample_detection(0)

        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_valid, batch_size))
        self.model.fit_generator(self.train_data_generator,
                            steps_per_epoch=max(1, num_train // batch_size),
                            validation_data=self.valid_data_generator,
                            validation_steps=max(1, num_valid // batch_size),
                            epochs=100,
                            initial_epoch=0,
                            callbacks=[logging, checkpoint])

        self.model.save_weights(self.log_dir + 'trained_weights.h5')
        print("training finished correctly")

    def evaluate(self, epoch, out_path_base="evaluation"):
        # for train data
        image_ids_gt, groundtruth_boxes_list, groundtruth_classes_list = self.make_groundtruth_list(self.train_lines)
        image_ids_dt, detection_boxes_list, detection_scores_list, detection_classes_list = self.make_detection_list(self.train_lines)

        groundtruth_list = ExportGroundtruthToCOCO(
            image_ids_gt, groundtruth_boxes_list, groundtruth_classes_list, self.categories_dict
        )
        detections_list = ExportDetectionsToCOCO(
            image_ids_dt, detection_boxes_list, detection_scores_list, detection_classes_list, self.categories_dict
        )

        groundtruth = COCOWrapper(groundtruth_list)
        detections = groundtruth.LoadAnnotations(detections_list)
        evaluator = COCOEvalWrapper(groundtruth, detections)
        summary_metrics_train, per_cat_ap_train = evaluator.ComputeMetrics()

        # for validation data
        image_ids_gt, groundtruth_boxes_list, groundtruth_classes_list = self.make_groundtruth_list(self.valid_lines)
        image_ids_dt, detection_boxes_list, detection_scores_list, detection_classes_list = self.make_detection_list(self.valid_lines)

        groundtruth_list = ExportGroundtruthToCOCO(
            image_ids_gt, groundtruth_boxes_list, groundtruth_classes_list, self.categories_dict
        )
        detections_list = ExportDetectionsToCOCO(
            image_ids_dt, detection_boxes_list, detection_scores_list, detection_classes_list, self.categories_dict
        )

        groundtruth = COCOWrapper(groundtruth_list)
        detections = groundtruth.LoadAnnotations(detections_list)
        evaluator = COCOEvalWrapper(groundtruth, detections)
        summary_metrics_valid, per_cat_ap_valid = evaluator.ComputeMetrics()

        # # print evaluated metrics to std_out
        # print("Evaluation on TRAIN DATA -----")
        # for metric_name, metric_value in summary_metrics_train:
        #     print(metric_name + " : " + metric_value)
        # print("Evaluation on VALIDATION DATA -----")
        # for metric_name, metric_value in summary_metrics_valid:
        #     print(metric_name + " : " + metric_value)

        # save evaluated metrics to single csv file
        outlines = [epoch] + list(summary_metrics_train.values()) + list(summary_metrics_valid.values())
        self.training_log_file.write(",".join([str(x) for x in outlines])+"\n")
        self.training_log_file.flush()

        # save evaluated metrics to csv file by epoch
        # summary_filename = out_path_base + "_summary-ep%3d.csv" % epoch
        # summary_path = os.path.join(self.log_dir, summary_filename)

        # per_cat_filename = out_path_base + "_per_category-ep%3d.csv" % epoch
        # per_cat_path = os.path.join(self.log_dir, per_cat_filename)

        # write_csv(summary_path, summary_metrics_valid)
        # write_csv(per_cat_path, per_cat_ap_valid)

        # plot training process
        names = ["epoch"] \
                + [n + "_train" for n in list(summary_metrics_train.keys())] \
                + [n + "_valid" for n in list(summary_metrics_valid.keys())]
        d = pd.read_csv(os.path.join(self.log_dir, "progress.csv"), header=None, names=names)
        if d.shape[0]==1: return
        d = d.interpolate()
        p = d.plot(x="epoch",
                   y=['Precision/mAP_train', 'Precision/mAP_valid',
                      'Precision/mAP@.50IOU_train', 'Precision/mAP@.50IOU_valid',
                      'Recall/AR@1_train', 'Recall/AR@1_valid'])
        fig = p.get_figure()
        fig.savefig(os.path.join(self.log_dir, "graph.png"))
        plt.close()

    def sample_detection(self, epoch):
        os.makedirs("%s/sample_images" % self.output_dir, exist_ok=True)
        # r, c = 1, 2
        samples = self.valid_lines[:3]
        for idx, line in enumerate(samples):
            line = line.split()
            image_gt = Image.open(line[0])
            image_dt = Image.open(line[0])
            line_boxes_ids = [l.split(",") for l in line[1:]]
            gt_boxes = np.array([np.array([float(xmin), float(ymin), float(xmax), float(ymax)])
                              for xmin, ymin, xmax, ymax, _ in line_boxes_ids])
            gt_classes = np.array([int(class_id) for _, _, _, _, class_id in line_boxes_ids])

            new_image_size = (image_gt.width - (image_gt.width % 32),
                              image_gt.height - (image_gt.height % 32))
            boxed_image = letterbox_image(image_gt, new_image_size)
            image_data = np.array(boxed_image, dtype='float32')

            image_data /= 255.
            image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

            out_boxes, out_scores, out_classes = self.sess.run(
                [self.boxes, self.scores, self.classes],
                feed_dict={
                    self.model_body.input: image_data,
                    self.input_image_shape: [image_gt.size[1], image_gt.size[0]],
                    K.learning_phase(): 0
                })

            font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                                      size=np.floor(3e-2 * image_gt.size[1] + 0.5).astype('int32'))
            thickness = (image_gt.size[0] + image_gt.size[1]) // 300

            # ground truth result
            for i, c in reversed(list(enumerate(gt_classes))):
                groundtruth_class = self.class_names[c]
                box = gt_boxes[i]
                label = '{}'.format(groundtruth_class)
                draw = ImageDraw.Draw(image_gt)
                label_size = draw.textsize(label, font)
                left, top, right, bottom = box
                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(image_gt.size[1], np.floor(bottom + 0.5).astype('int32'))
                right = min(image_gt.size[0], np.floor(right + 0.5).astype('int32'))

                if top - label_size[1] >= 0:
                    text_origin = np.array([left, top - label_size[1]])
                else:
                    text_origin = np.array([left, top + 1])

                for i in range(thickness):
                    draw.rectangle(
                        [left + i, top + i, right - i, bottom - i],
                        outline=self.colors[c])
                draw.rectangle(
                    [tuple(text_origin), tuple(text_origin + label_size)],
                    fill=self.colors[c])
                draw.text(text_origin, label, fill=(0, 0, 0), font=font)
                del draw

            # for detection result
            for i, c in reversed(list(enumerate(out_classes))):
                predicted_class = self.class_names[c]
                box = out_boxes[i]
                score = out_scores[i]
                label = '{} {:.2f}'.format(predicted_class, score)
                draw = ImageDraw.Draw(image_dt)
                label_size = draw.textsize(label, font)
                left, top, right, bottom = box
                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(image_gt.size[1], np.floor(bottom + 0.5).astype('int32'))
                right = min(image_gt.size[0], np.floor(right + 0.5).astype('int32'))

                if top - label_size[1] >= 0:
                    text_origin = np.array([left, top - label_size[1]])
                else:
                    text_origin = np.array([left, top + 1])

                for i in range(thickness):
                    draw.rectangle(
                        [left + i, top + i, right - i, bottom - i],
                        outline=self.colors[c])
                draw.rectangle(
                    [tuple(text_origin), tuple(text_origin + label_size)],
                    fill=self.colors[c])
                draw.text(text_origin, label, fill=(0, 0, 0), font=font)
                del draw

            # Rescale images 0 - 1
            imgs = [np.array(image_gt, dtype='float32') / 255., np.array(image_dt, dtype='float32') / 255.]

            titles = ["Ground Truth", "Detected by YOLOv3"]

            fig, axs = plt.subplots(1, 2)
            for i in range(2):
                axs[i].imshow(imgs[i])
                axs[i].set_title(titles[i])
                axs[i].axis('off')
            fig.savefig(os.path.join(self.output_dir, "sample_images/sample_ep%04d-%d.png" % (epoch, idx)))
            plt.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        prog='retraining.py',
        description='Retraining yolov3 on your own dataset.',
        add_help=True,
    )

    parser.add_argument('--train_path', type=str, help='path to train annotation txt path', required=True)
    parser.add_argument('--valid_path', type=str, help='path to validation annotation txt path', required=True)
    parser.add_argument('-b', '--batch_size', type=int, default=64, help='batch size for training')
    parser.add_argument('--log_dir', type=str, default='logs/000', help='path to log directory')
    parser.add_argument('-o', '--output_dir', type=str, default='outputs/000', help='path to output directory')
    parser.add_argument('--classes_path', type=str, default='model_data/coco_classes.txt')
    parser.add_argument('--anchors_path', type=str, default='model_data/yolo_anchors.txt')
    parser.add_argument('-s', '--input_size', type=str, default='416x416', help='HEIGHT x WIDTH')
    parser.add_argument('--pretrained_weights', type=str, default='model_data/yolo_weights.h5', help='path to pretrained yolov3 weights')
    parser.add_argument('-n', '--gpu_num', type=int, default=1, help='number of GPU')

    args = parser.parse_args()
    train_annotation_path = args.train_path
    valid_annotation_path = args.valid_path
    batch_size = args.batch_size
    log_dir = args.log_dir
    output_dir = args.output_dir
    classes_path = args.classes_path
    anchors_path = args.anchors_path
    height, width = args.input_size.split("x")
    input_shape = (int(height), int(width))
    weights_path = args.pretrained_weights
    gpu_num = args.gpu_num

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    trainer = Trainer(
        train_annotation_path=train_annotation_path,
        valid_annotation_path=valid_annotation_path,
        log_dir=log_dir,
        output_dir = output_dir,
        classes_path=classes_path,
        anchors_path=anchors_path,
        input_shape=input_shape,
        gpu_num=gpu_num
    )

    trainer.train(
        weights_path=weights_path,
        batch_size=batch_size,
        freeze_body=0
    )


"""Downstream expert for Query-by-Example Spoken Term Detection on QUESST 2014."""

from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np
import pyximport
import torch
import torch.nn as nn
from lxml import etree
from torch.utils.data import DataLoader
from tqdm import tqdm

pyximport.install(setup_args={"include_dirs": np.get_include()})

from .dataset import QUESST14Dataset
from .segmental_dtw import segmental_dtw

class DownstreamExpert(nn.Module):
    """
    Used to handle downstream-specific operations
    eg. downstream forward, metric computation, contents to log
    """

    def __init__(
        self, upstream_dim: int, downstream_expert: dict, expdir: str, **kwargs
    ):
        super(DownstreamExpert, self).__init__()
        self.upstream_dim = upstream_dim
        self.datarc = downstream_expert["datarc"]
        self.expdir = Path(expdir)
        self.test_dataset = QUESST14Dataset(**self.datarc)

    def _get_dataloader(self, dataset):
        return DataLoader(
            dataset,
            shuffle=False,
            batch_size=self.datarc["batch_size"],
            drop_last=False,
            num_workers=self.datarc["num_workers"],
            collate_fn=dataset.collate_fn,
        )

    # Interface
    def get_train_dataloader(self):
        return None

    # Interface
    def get_dev_dataloader(self):
        return None

    # Interface
    def get_test_dataloader(self):
        return self._get_dataloader(self.test_dataset)

    # Interface
    def forward(
        self,
        features,
        audio_names,
        records,
        **kwargs,
    ):
        for feature, audio_name in zip(features, audio_names):
            feature = feature.detach().cpu().numpy().astype(np.double)
            # feature = feature[feature.argmax(1) != 0]
            records["features"].append(feature)
            records["audio_names"].append(audio_name)

        return torch.zeros(1)

    # interface
    def log_records(self, records, **kwargs):
        """Perform DTW and save results."""
        queries = records["features"][: self.test_dataset.n_queries]
        docs = records["features"][self.test_dataset.n_queries :]
        query_names = records["audio_names"][: self.test_dataset.n_queries]
        doc_names = records["audio_names"][self.test_dataset.n_queries :]
        results = defaultdict(list)
        scores = []

        # Calculate matching scores
        with ProcessPoolExecutor(self.datarc["num_workers"]) as executor:
            futures = []

            for query, query_name in zip(queries, query_names):
                query_name = query_name.replace(".wav", "")
                for doc, doc_name in zip(docs, doc_names):
                    doc_name = doc_name.replace(".wav", "")
                    futures.append(
                        executor.submit(match, query, doc, query_name, doc_name)
                    )

            for future in tqdm(
                as_completed(futures), total=len(futures), ncols=0, desc="DTW"
            ):
                query_name, doc_name, score = future.result()
                results[query_name].append((doc_name, score))
                scores.append(score)

        # Determine score threshold
        scores = sorted(scores)
        score_thresh = scores[int(0.99 * len(scores))]
        score_min = scores[0]

        # Build XML tree
        root = etree.Element(
            "stdlist",
            termlist_filename="benchmark.stdlist.xml",
            indexing_time="1.00",
            language="english",
            index_size="1",
            system_id="benchmark",
        )
        for query_name, doc_scores in results.items():
            term_list = etree.SubElement(
                root,
                "detected_termlist",
                termid=query_name,
                term_search_time="1.0",
                oov_term_count="1",
            )
            for doc_name, score in doc_scores:
                etree.SubElement(
                    term_list,
                    "term",
                    file=doc_name,
                    channel="1",
                    tbeg="0.000",
                    dur="0.00",
                    score=f"{score - score_min:.4f}",
                    decision="YES" if score > score_thresh else "NO",
                )

        # Output XML
        tree = etree.ElementTree(root)
        tree.write(
            str(self.expdir / "benchmark.stdlist.xml"),
            encoding="UTF-8",
            pretty_print=True,
        )

def match(query, doc, query_name, doc_name):
    cost = segmental_dtw(query, doc)
    return query_name, doc_name, -1 * cost


#!/usr/bin/env python
import time
import argparse
from pathlib import Path
from astropy.time import Time
from scipy.optimize import leastsq
import fitsio
import numpy as np
import textwrap
from . import sdss_paths

__version__ = '3.2.2'

class APOGEERaw:
    """A class to parse raw data from APOGEE. The purpose of this class is to
    read raw image files from /data/apogee/archive, regardless of any future
    changes in SDSS. This
    will hopefully help SDSS-V logging.

    Methods:
        compute_offets: Returns the offsets of an arc/object to compute the
        relative dither

        ap_test: Returns a list of faint fibers and missing fibers from a flat
        """

    def __init__(self, fil, args, ext=1, ):
        self.file = Path(fil)
        if not self.file.exists():
            raise FileNotFoundError(f"Could not find file {self.file.absolute()}")
        self.ext = ext
        self.args = args
        try:
            header = fitsio.read_header(fil, ext=ext)
        except OSError:
            header = fitsio.read_header(fil, ext=ext)
        if (header['DITHPIX'] - 12.994) < 0.05:
            self.dither = 'A'
        elif (header['DITHPIX'] - 13.494) < 0.05:
            self.dither = 'B'
        else:
            self.dither = '{:.1f}'.format(header['DITHPIX'])
        self.exp_time = header['EXPTIME']
        self.isot = Time(header['DATE-OBS'])  # Local
        self.field_id = int(header["FIELDID"])
        self.plate_id = header['PLATEID']
        if "CONFIGID" in header.keys():
            if isinstance(header["CONFIGID"], int):
                self.design_id = header["DESIGNID"]
                self.config_id = header["CONFIGID"]
            else:
                self.design_id = 0
                self.config_id = 0
        else:
            self.design_id = 0
            self.config_id = 0

        if "CARTID" in header.keys():
            if isinstance(header["CARTID"], int):
                self.cart_id = f"{header['CARTID']:.0f}"
            else:
                self.cart_id = header["CARTID"]
        else:
            self.cart_id = 0

        self.exp_id = int(str(fil).split('-')[-1].split('.')[0])
        if header['EXPTYPE'].capitalize() == 'Arclamp':
            if header['LAMPUNE']:
                self.exp_type = 'UNe Arc'
            elif header['LAMPTHAR']:
                self.exp_type = 'ThAr Arc'
            elif "FPI" in header["OBSCMNT"]:
                self.exp_type = "FPI Lamp"
            else:
                self.exp_type = "Arc"
                if args.verbose:
                    print('Could not process exposure type of {}'.format(self.file))
        else:
            self.exp_type = header['EXPTYPE'].capitalize()
        self.n_read = header['NREAD']
        self.lead = header['PLATETYP']
        self.img_type = header['IMAGETYP'].capitalize()
        if header['EXPTYPE'] == 'OBJECT':
            self.seeing = header['SEEING']
        else:
            self.seeing = 0.0
        self.mjd = self.file.absolute().parent.name
        self.quickred_data = np.array([[]])
        self.quickred_file = sdss_paths.ap_qr / f"{self.mjd}/apq-{self.exp_id}.fits"
        self.utr_file = ''
        self.utr_data = np.array([[]])

    # noinspection PyTupleAssignmentBalance,PyTypeChecker
    def compute_offset(self, fibers=(60, 70), w0: int=1105, dw:int=40, sigma:float=1.2745):
        """This is based off of apogeeThar.OneFileFitting written by Elena. It
        is supposed to generate a float for the pixel offsets of an APOGEE
        ThAr cal. Here is how it works:
        It opens a quickred file, which is of shape n_fiber*n_dispersion_pixels,
        and then it averages the fibers inside the fibers tuple. It then based
        off of the provided w0 (mean) and sigma, it creates a gaussian function
        and then creates a function called err_func that compares the gaussian
        with a slice of the data, from w0-dw/2 to w0+dw/2. It then uses scipy's
        least squared equation solver to find the difference between w0 given as
        an input and the actual w0 of the spectral line. This only works if you
        pick a prominent line to go off of. The default parameters are given for
        ThAr lines, but UNe lines could also be used, with the following inputs:
        fibers: (30, 35)
        w0: 1761
        dw: 20
        sigma: 3
        """
        w0 = int(w0)
        dw = int(dw)
        if self.quickred_data.size == 0:
            try:
                self.quickred_file.exists()
            except OSError:
                pass
            if self.quickred_file.exists():
                self.quickred_data = fitsio.read(self.quickred_file, 3)[0][0]
            else:
                self.quickred_file = (sdss_paths.ap_qr
                                      / 'quickred/{}/ap1D-a-{}.fits.fz'
                                        ''.format(self.mjd, self.exp_id))
                if not self.quickred_file.exists():
                    print(f"Offsets for {self.file.name} could not be read")
                    return np.nan
                self.quickred_data = fitsio.read(self.quickred_file, 1)
        lower = w0 - dw // 2
        upper = w0 + dw // 2
        line_inds = np.arange(self.quickred_data.shape[1])[lower:upper]
        try:
            line = np.average(self.quickred_data[fibers[0]:fibers[1], lower:upper],
                          axis=0)
        except ZeroDivisionError:
            print(f"Couldn't find dither offsets of image {self.file} with"
                  f" shape {self.quickred_data.shape}")
            return np.nan

        def fit_func(w, x):
            return np.exp(-0.5 * ((x - w) / sigma) ** 2)

        def err_func(w, x, y):
            return fit_func(w, x) - y

        w_model, success = leastsq(err_func, w0, args=(line_inds, line))

        diff = w_model[0] - w0
        if np.abs(diff) > 10:
            print('A large dither was reported for exposure {}: {:.3f}'
                  '\n  This may mean the zero point needs to be'
                  ' adjusted, currently it is {}'
                  ''.format(self.exp_id, diff, w0))
        return diff

    def ap_test(self, ws=(550, 910), master_col=None, plot=False,
                print_it=False):
        if master_col is None:
            raise ValueError("APTest didn't receive a valid master_col: {}"
                             "".format(master_col))
        if self.quickred_data.size == 0:
            if self.quickred_file.exists():
                self.quickred_data = fitsio.read(self.quickred_file, 3)[0][0]
            else:
                self.quickred_file = (sdss_paths.ap_qr
                                      / 'quickred/{}/ap1D-a-{}.fits.fz'
                                        ''.format(self.mjd, self.exp_id))
                if not self.quickred_file.exists():
                    print(f"Offsets for {self.file.name} could not be read")
                    return [], [], np.nan
                self.quickred_data = fitsio.read(self.quickred_file, 1)
        slc = np.median(self.quickred_data[:, ws[0]:ws[1]], axis=1)
        flux_ratio = np.flipud(slc / master_col)
        bad_data = ((flux_ratio == np.inf)
                    | (flux_ratio == -np.inf)
                    | np.isnan(flux_ratio))
        flux_ratio[bad_data] = np.nan
        missing = flux_ratio < 0.2
        faint = (flux_ratio < 0.7) & (0.2 <= flux_ratio)
        bright = ~missing & ~faint
        i_missing = np.where(missing)[0].astype(int)
        i_faint = np.where(faint)[0].astype(int)
        i_bright = np.where(bright)[0]
        missing_bundles = self.create_bundles(i_missing)
        faint_bundles = self.create_bundles(i_faint)
        if print_it:
            print(textwrap.fill('Missing Fibers: {}'.format(missing_bundles),
                                80))
            print(textwrap.fill('Faint Fibers: {}'.format(faint_bundles), 80))
            print()

        if plot:
            import matplotlib.pyplot as plt
            fig = plt.figure(figsize=(9, 4))
            ax = fig.gca()
            x = np.arange(len(flux_ratio)) + 1
            ax.plot(x[i_bright], flux_ratio[i_bright], 'o', c=(0, 0.6, 0.533))
            ax.plot(x[i_faint], flux_ratio[i_faint], 'o', c=(0.933, 0.466, 0.2))
            ax.plot(x[i_missing], flux_ratio[i_missing], 'o',
                    c=(0.8, 0.2, 0.066))
            ax.set_xlabel('Fiber ID')
            ax.set_ylabel('Throughput Efficiency')
            ax.axis([1, 300, -0.2, 1.35])
            ax.grid(True)
            ax.axhline(0.7, c=(0, 0.6, 0.533))
            ax.axhline(0.2, c=(0.933, 0.466, 0.2))
            ax.set_title('APOGEE Fiber Relative Intensity {}'.format(
                self.exp_id), size=15)
            plt.show()

        return missing_bundles, faint_bundles, flux_ratio

    @staticmethod
    def create_bundles(subset):
        """This method converts an array of ints into a list of strings that
        describe a large series of fibers and ints for lone fibers.

        Ex: [1, 2, 3, 5] -> ['1 - 3', 5]

        """
        if len(subset) == 0:
            return []
        bundles = [subset[0]]
        b = 0
        for fib in subset[1:]:
            if isinstance(bundles[b], np.int64):
                if bundles[b] + 1 == fib:
                    if fib % 30 == 0:
                        bundles.append(fib)
                        b += 1
                    else:
                        # All strings are created here
                        bundles[b] = '{} - {}'.format(bundles[b], fib)
                else:
                    bundles.append(fib)
                    b += 1
            elif isinstance(bundles[b], str):
                if int(bundles[b].split()[-1]) + 1 == fib:
                    if fib % 30 == 1:
                        bundles.append(fib)
                        b += 1
                    else:
                        bundles[b] = '{} - {}'.format(
                            bundles[b].split()[0], fib)
                else:
                    bundles.append(fib)
                    b += 1
        # for i, bundle in enumerate(bundles):
        #     if isinstance(bundle, str):
        #         low, high = np.array(bundle.split(' - ')).astype(int)
        #         if ((low - 1) // 30) == ((high - 1) // 30):
        #             bundles[i] = '{} bundle'.format(low)
        return bundles

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--today', action='store_true',
                        help="Whether or not you want to search for today's"
                             " data, whether or not the night is complete."
                             " Note: must be run after 00:00Z")
    parser.add_argument('-m', '--mjd',
                        help='If not today (-t), the self.mjd to search')
    parser.add_argument('-v', '--verbose', action='count', default=1,
                        help='Show details, can be stacked')
    args = parser.parse_args()
    if args.today:
        mjd_today = int(Time.now().sjd)
        data_dir = sdss_paths.ap_archive / f"{mjd_today}/"
    elif args.self.mjd:
        data_dir = sdss_paths.ap_archive / f"{args.self.mjd}"
    else:
        raise Exception('No date specified')
    # print(data_dir)
    for path in data_dir.rglob('apq*.fits'):
        print(path)

if __name__ == '__main__':
    main()


import tr_helpers
import tensorflow as tf
import numpy as np
import collections
import time
from collections import deque, OrderedDict
from tensorboardX import SummaryWriter
import gym
import vecenv
from tf_moving_mean_std import MovingMeanStd

def swap_and_flatten01(arr):
    """
    swap and then flatten axes 0 and 1
    """
    s = arr.shape
    return arr.swapaxes(0, 1).reshape(s[0] * s[1], *s[2:])

class A2CAgent:
    def __init__(self, sess, name, observation_space, is_discrete, action_space, config):
        assert is_discrete
        observation_shape = observation_space.shape
        self.name = name
        self.config = config
        self.env_name = config['ENV_NAME']
        self.ppo = config['PPO']
        self.is_adaptive_lr = config['LR_SCHEDULE'] == 'ADAPTIVE'
        self.is_polynom_decay_lr = config['LR_SCHEDULE'] == 'POLYNOM_DECAY'
        self.is_exp_decay_lr = config['LR_SCHEDULE'] == 'EXP_DECAY'
        self.lr_multiplier = tf.constant(1, shape=(), dtype=tf.float32)
        if self.is_adaptive_lr:
            self.lr_threshold = config['LR_THRESHOLD']
        self.e_clip = config['E_CLIP']
        self.clip_value = config['CLIP_VALUE']
        self.network = config['NETWORK']
        self.rewards_shaper = config['REWARD_SHAPER']
        self.num_actors = config['NUM_ACTORS']
        self.vec_env = vecenv.create_vec_env(self.env_name, self.num_actors)
        self.steps_num = config['STEPS_NUM']
        self.seq_len = self.config['SEQ_LEN']
        self.normalize_advantage = config['NORMALIZE_ADVANTAGE']
        self.normalize_input = self.config['NORMALIZE_INPUT']

        self.state_shape = observation_shape#(observation_shape[0], observation_shape[1], observation_shape[2] // 3 + 2)
        self.critic_coef = config['CRITIC_COEF']
        self.writer = SummaryWriter()
        self.sess = sess
        self.grad_norm = config['GRAD_NORM']
        self.gamma = self.config['GAMMA']
        self.lam = self.config['LAMBDA']
        self.dones = np.asarray([False]*self.num_actors, dtype=np.bool)
        self.current_rewards = np.asarray([0]*self.num_actors, dtype=np.float32)
        self.game_rewards = deque([], maxlen=1000)
        self.obs_ph = tf.placeholder('uint8', (None, ) + self.state_shape, name = 'obs')
        self.target_obs_ph = tf.placeholder('uint8', (None, ) + self.state_shape, name = 'target_obs')
        self.actions_num = action_space.n
        self.actions_ph = tf.placeholder('int32', (None,), name = 'actions')

        self.old_logp_actions_ph = tf.placeholder('float32', (None, ), name = 'old_logpactions')
        self.rewards_ph = tf.placeholder('float32', (None,), name = 'rewards')
        self.old_values_ph = tf.placeholder('float32', (None,), name = 'old_values')
        self.advantages_ph = tf.placeholder('float32', (None,), name = 'advantages')
        self.learning_rate_ph = tf.placeholder('float32', (), name = 'lr_ph')

        self.epoch_num = tf.Variable( tf.constant(0, shape=(), dtype=tf.float32), trainable=False)#, name = 'epochs')
        self.update_epoch_op = self.epoch_num.assign(self.epoch_num + 1)
        self.current_lr = self.learning_rate_ph

        if self.is_adaptive_lr:
            self.lr_threshold = config['LR_THRESHOLD']
        if self.is_polynom_decay_lr:
            self.lr_multiplier = tf.train.polynomial_decay(1.0, global_step=self.epoch_num, decay_steps=config['MAX_EPOCHS'], end_learning_rate=0.001, power=tr_helpers.get_or_default(config, 'DECAY_POWER', 1.0))
        if self.is_exp_decay_lr:
            self.lr_multiplier = tf.train.exponential_decay(1.0, global_step=self.epoch_num, decay_steps=config['MAX_EPOCHS'],  decay_rate = config['DECAY_RATE'])

        if self.normalize_input:
            self.moving_mean_std = MovingMeanStd(shape = observation_space.shape, epsilon = 1e-5, decay = 0.99)
            self.input_obs = self.moving_mean_std.normalize(self.obs_ph, train=True)
            self.input_target_obs = self.moving_mean_std.normalize(self.target_obs_ph, train=False)
        else:
            self.input_obs = self.obs_ph
            self.input_target_obs = self.target_obs_ph

        self.input_obs = tf.to_float(self.input_obs) / 255.0
        self.input_target_obs = tf.to_float(self.input_target_obs) / 255.0

        games_num = self.config['MINIBATCH_SIZE'] // self.seq_len # it is used only for current rnn implementation

        self.num_vel_inputs = 2 * 4
        self.vels_ph = tf.placeholder(tf.float32, [self.config['MINIBATCH_SIZE'], self.num_vel_inputs])
        self.target_vels_ph = tf.placeholder(tf.float32, [self.num_actors, self.num_vel_inputs])

        self.train_dict = {
            'name' : 'agent',
            'inputs' : self.input_obs,
            'batch_num' : self.config['MINIBATCH_SIZE'],
            'games_num' : games_num,
            'actions_num' : self.actions_num,
            'vels_ph' : self.vels_ph,
            'prev_actions_ph' : self.actions_ph
        }

        self.run_dict = {
            'name' : 'agent',
            'inputs' : self.input_target_obs,
            'batch_num' : self.num_actors,
            'games_num' : self.num_actors,
            'actions_num' : self.actions_num,
            'vels_ph' : self.target_vels_ph,
            'prev_actions_ph' : None
        }

        self.states = None
        if self.network.is_rnn():
            self.neglogp_actions ,self.state_values, self.action, self.entropy, self.states_ph, self.masks_ph, self.lstm_state, self.initial_state = self.network(self.train_dict, reuse=False)
            self.target_neglogp, self.target_state_values, self.target_action, _,  self.target_states_ph, self.target_masks_ph, self.target_lstm_state, self.target_initial_state = self.network(self.run_dict, reuse=True)
            self.states = self.target_initial_state
        else:
            self.neglogp_actions ,self.state_values, self.action, self.entropy = self.network(self.train_dict, reuse=False)
            self.target_neglogp, self.target_state_values, self.target_action, _ = self.network(self.run_dict, reuse=True)

        curr_e_clip = self.e_clip * self.lr_multiplier
        if (self.ppo):
            self.prob_ratio = tf.exp(self.old_logp_actions_ph - self.neglogp_actions)
            self.pg_loss_unclipped = -tf.multiply(self.advantages_ph, self.prob_ratio)
            self.pg_loss_clipped = -tf.multiply(self.advantages_ph, tf.clip_by_value(self.prob_ratio, 1.- curr_e_clip, 1.+ curr_e_clip))
            self.actor_loss = tf.reduce_mean(tf.maximum(self.pg_loss_unclipped, self.pg_loss_clipped))
        else:
            self.actor_loss = tf.reduce_mean(self.neglogp_actions * self.advantages_ph)

        self.c_loss = (tf.squeeze(self.state_values) - self.rewards_ph)**2
        if self.clip_value:
            self.cliped_values = self.old_values_ph + tf.clip_by_value(tf.squeeze(self.state_values) - self.old_values_ph, - curr_e_clip, curr_e_clip)
            self.c_loss_clipped = tf.square(self.cliped_values - self.rewards_ph)
            self.critic_loss = tf.reduce_mean(tf.maximum(self.c_loss, self.c_loss_clipped))
        else:
            self.critic_loss = tf.reduce_mean(self.c_loss)

        self.kl_approx = 0.5 * tf.stop_gradient(tf.reduce_mean((self.old_logp_actions_ph - self.neglogp_actions)**2))

        entropy_coef = self.config['ENTROPY_COEF'] * self.lr_multiplier
        self.loss = self.actor_loss + 0.5 * self.critic_coef * self.critic_loss -  entropy_coef * self.entropy

        self.train_step = tf.train.AdamOptimizer(self.current_lr * self.lr_multiplier)
        self.weights = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='agent')

        grads = tf.gradients(self.loss, self.weights)
        if self.config['TRUNCATE_GRADS']:
            grads, _ = tf.clip_by_global_norm(grads, self.grad_norm)
        grads = list(zip(grads, self.weights))
        self.train_op = self.train_step.apply_gradients(grads)
        self.saver = tf.train.Saver()
        self.sess.run(tf.global_variables_initializer())

    def update_epoch(self):
        return self.sess.run([self.update_epoch_op])[0]

    def get_action_values(self, obs):
        run_ops = [self.target_action, self.target_state_values, self.target_neglogp]
        if self.network.is_rnn():
            run_ops.append(self.target_lstm_state)
            return self.sess.run(run_ops, {self.target_obs_ph : obs[0], self.target_vels_ph : obs[1], self.target_states_ph : self.states, self.target_masks_ph : self.dones})
        else:
            return (*self.sess.run(run_ops, {self.target_obs_ph : obs[0], self.target_vels_ph : obs[1]}), None)

    def get_values(self, obs):
        if self.network.is_rnn():
            return self.sess.run([self.target_state_values], {self.target_obs_ph : obs[0], self.target_vels_ph : obs[1], self.target_states_ph : self.states, self.target_masks_ph : self.dones})
        else:
            return self.sess.run([self.target_state_values], {self.target_obs_ph : obs[0], self.target_vels_ph : obs[1]})

    def play_steps(self):
        # Here, we init the lists that will contain the mb of experiences
        mb_obs0, mb_obs1, mb_rewards, mb_actions, mb_values, mb_dones, mb_neglogpacs = [],[],[],[],[],[],[]
        mb_states = []
        epinfos = []
        # For n in range number of steps
        for _ in range(self.steps_num):
            if self.network.is_rnn():
                mb_states.append(self.states)

            actions, values, neglogpacs, self.states = self.get_action_values(self.obs)
            actions = np.squeeze(actions)
            values = np.squeeze(values)
            neglogpacs = np.squeeze(neglogpacs)
            mb_obs0.append(self.obs[0])
            mb_obs1.append(self.obs[1])
            mb_actions.append(actions)
            mb_values.append(values)
            mb_neglogpacs.append(neglogpacs)
            mb_dones.append(self.dones)

            self.obs, rewards, self.dones, infos = self.vec_env.step(actions)
            self.current_rewards += rewards

            for reward, done in zip(self.current_rewards, self.dones):
                if done:
                    self.game_rewards.append(reward)

            self.current_rewards = self.current_rewards * (1.0 -self.dones)

            shaped_rewards = self.rewards_shaper(rewards)
            epinfos.append(infos)
            mb_rewards.append(shaped_rewards)

        #using openai baseline approach
        mb_obs0 = np.asarray(mb_obs0, dtype=np.uint8)
        mb_obs1 = np.asarray(mb_obs1, dtype=np.float32)
        mb_rewards = np.asarray(mb_rewards, dtype=np.float32)
        mb_actions = np.asarray(mb_actions, dtype=np.float32)
        mb_values = np.asarray(mb_values, dtype=np.float32)
        mb_neglogpacs = np.asarray(mb_neglogpacs, dtype=np.float32)
        mb_dones = np.asarray(mb_dones, dtype=np.bool)
        mb_states = np.asarray(mb_states, dtype=np.float32)
        last_values = self.get_values(self.obs)
        last_values = np.squeeze(last_values)

        mb_returns = np.zeros_like(mb_rewards)
        mb_advs = np.zeros_like(mb_rewards)
        lastgaelam = 0

        for t in reversed(range(self.steps_num)):
            if t == self.steps_num - 1:
                nextnonterminal = 1.0 - self.dones
                nextvalues = last_values
            else:
                nextnonterminal = 1.0 - mb_dones[t+1]
                nextvalues = mb_values[t+1]

            delta = mb_rewards[t] + self.gamma * nextvalues * nextnonterminal  - mb_values[t]
            mb_advs[t] = lastgaelam = delta + self.gamma * self.lam * nextnonterminal * lastgaelam

        mb_returns = mb_advs + mb_values
        if self.network.is_rnn():
            result = (*map(swap_and_flatten01, (mb_obs0, mb_obs1, mb_returns, mb_dones, mb_actions, mb_values, mb_neglogpacs, mb_states )), epinfos)
        else:
            result = (*map(swap_and_flatten01, (mb_obs0, mb_obs1, mb_returns, mb_dones, mb_actions, mb_values, mb_neglogpacs)), None, epinfos)

        return result

    def save(self, fn):
        self.saver.save(self.sess, fn)

    def restore(self, fn):
        self.saver.restore(self.sess, fn)

    def train(self):
        self.obs = self.vec_env.reset()
        batch_size = self.steps_num * self.num_actors
        minibatch_size = self.config['MINIBATCH_SIZE']
        mini_epochs_num = self.config['MINI_EPOCHS']
        num_minibatches = batch_size // minibatch_size
        last_lr = self.config['LEARNING_RATE']
        frame = 0
        update_time = 0
        last_mean_rewards = -100500
        play_time = 0
        epoch_num = 0
        max_epochs = tr_helpers.get_or_default(self.config, 'MAX_EPOCHS', 1e6)
        print('max epochs', max_epochs)
        while True:
            play_time_start = time.time()
            epoch_num = self.update_epoch()
            frame += batch_size
            obses, vels, returns, dones, actions, values, neglogpacs, lstm_states, _ = self.play_steps()
            advantages = returns - values
            if self.normalize_advantage:
                advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

            a_losses = []
            c_losses = []
            entropies = []
            kls = []
            play_time_end = time.time()
            play_time = play_time_end - play_time_start
            update_time_start = time.time()
            if self.network.is_rnn():
                total_games = batch_size // self.seq_len
                num_games_batch = minibatch_size // self.seq_len
                game_indexes = np.arange(total_games)
                flat_indexes = np.arange(total_games * self.seq_len).reshape(total_games, self.seq_len)
                lstm_states = lstm_states[::self.seq_len]
                for _ in range(0, mini_epochs_num):
                    np.random.shuffle(game_indexes)

                    for i in range(0, num_minibatches):
                        batch = range(i * num_games_batch, (i + 1) * num_games_batch)
                        mb_indexes = game_indexes[batch]
                        mbatch = flat_indexes[mb_indexes].ravel()

                        dict = {}
                        dict[self.old_values_ph] = values[mbatch]
                        dict[self.old_logp_actions_ph] = neglogpacs[mbatch]
                        dict[self.advantages_ph] = advantages[mbatch]
                        dict[self.rewards_ph] = returns[mbatch]
                        dict[self.actions_ph] = actions[mbatch]
                        dict[self.obs_ph] = obses[mbatch]
                        dict[self.vels_ph] = vels[mbatch]
                        dict[self.masks_ph] = dones[mbatch]
                        dict[self.states_ph] = lstm_states[batch]

                        dict[self.learning_rate_ph] = last_lr
                        run_ops = [self.actor_loss, self.critic_loss, self.entropy, self.kl_approx, self.current_lr, self.lr_multiplier, self.train_op]
                        run_ops.append(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
                        a_loss, c_loss, entropy, kl, last_lr, lr_mul, _, _ = self.sess.run(run_ops, dict)
                        a_losses.append(a_loss)
                        c_losses.append(c_loss)
                        kls.append(kl)
                        entropies.append(entropy)
            else:
                for _ in range(0, mini_epochs_num):
                    permutation = np.random.permutation(batch_size)
                    obses = obses[permutation]
                    returns = returns[permutation]

                    actions = actions[permutation]
                    values = values[permutation]
                    neglogpacs = neglogpacs[permutation]
                    advantages = advantages[permutation]

                    for i in range(0, num_minibatches):
                        batch = range(i * minibatch_size, (i + 1) * minibatch_size)
                        dict = {self.obs_ph: obses[batch], self.actions_ph : actions[batch], self.rewards_ph : returns[batch],
                                self.advantages_ph : advantages[batch], self.old_logp_actions_ph : neglogpacs[batch], self.old_values_ph : values[batch]}
                        dict[self.vels_ph] = vels[batch]
                        dict[self.learning_rate_ph] = last_lr
                        run_ops = [self.actor_loss, self.critic_loss, self.entropy, self.kl_approx, self.current_lr, self.lr_multiplier, self.train_op]

                        run_ops.append(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
                        a_loss, c_loss, entropy, kl, last_lr, lr_mul, _, _ = self.sess.run(run_ops, dict)
                        a_losses.append(a_loss)
                        c_losses.append(c_loss)
                        kls.append(kl)
                        entropies.append(entropy)
            update_time_end = time.time()
            update_time = update_time_end - update_time_start
            sum_time = update_time + play_time
            if True:
                print('Frames per seconds: ', batch_size / sum_time)
                self.writer.add_scalar('performance/fps', batch_size / sum_time, frame)
                self.writer.add_scalar('performance/upd_time', update_time, frame)
                self.writer.add_scalar('performance/play_time', play_time, frame)
                self.writer.add_scalar('losses/a_loss', np.mean(a_losses), frame)
                self.writer.add_scalar('losses/c_loss', np.mean(c_losses), frame)
                self.writer.add_scalar('losses/entropy', np.mean(entropies), frame)
                self.writer.add_scalar('info/last_lr', last_lr * lr_mul, frame)
                self.writer.add_scalar('info/lr_mul', lr_mul, frame)
                self.writer.add_scalar('info/e_clip', self.e_clip * lr_mul, frame)
                self.writer.add_scalar('info/kl', np.mean(kls), frame)
                self.writer.add_scalar('epochs', epoch_num, frame)
                if len(self.game_rewards) > 0:
                    mean_rewards = np.mean(self.game_rewards)
                    self.writer.add_scalar('mean_rewards', mean_rewards, frame)
                    if mean_rewards > last_mean_rewards:
                        print('saving next best rewards: ', mean_rewards)
                        last_mean_rewards = mean_rewards
                        self.save("./nn/" + self.name + self.env_name)
                        if last_mean_rewards > self.config['SCORE_TO_WIN']:
                            print('Network won!')
                            return
                if epoch_num > max_epochs:
                    print('MAX EPOCHS NUM!')
                    return
                update_time = 0


import os
import sys
import json
import time
import base64
import logging
import requests
import datetime
import numpy as np
from PIL import Image
from io import BytesIO
from pathlib import Path
import onnxruntime as rt
from torchvision import transforms
from inference_schema.schema_decorators import input_schema, output_schema
from inference_schema.parameter_types.standard_py_parameter_type import StandardPythonParameterType

session, transform, classes, input_name, model_stamp = None, None, None, None, None
logger = logging.getLogger()

def init():
    global session, transform, classes, input_name, logger, model_stamp
    logger.info('Attempting to load model artifacts')

    if 'AZUREML_MODEL_DIR' in os.environ:
        logger.info('using AZUREML_MODEL_DIR')
        root_dir = Path(os.environ['AZUREML_MODEL_DIR']).resolve() / 'model'
    else:
        logger.info('using local')
        root_dir = Path('outputs/model').resolve()

    logger.info(f'using model path {root_dir}')
    meta_file = root_dir / 'meta.json'
    model_file = root_dir / 'model.onnx'
    logger.info(f'metadata path: {meta_file}')
    logger.info(f'model path: {model_file}')

    logger.info('loading metadata')
    with open(meta_file, 'r') as f:
        model_meta = json.load(f)
    logger.info(f'metadata load complete: {model_meta}')

    classes = model_meta['classes']
    model_stamp = model_meta['timestamp']

    logger.info('loading model')
    session = rt.InferenceSession(str(model_file))
    input_name = session.get_inputs()[0].name
    logger.info(f'model load complete (entry: {input_name})')
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    logger.info(f'transforms initialized')
    logger.info(f'init complete!')

@input_schema('image', StandardPythonParameterType("https://aiadvocate.z5.web.core.windows.net/scissors.png"))
@output_schema(StandardPythonParameterType({
  'time': StandardPythonParameterType(0.060392),
  'prediction': StandardPythonParameterType("paper"),
  'scores': StandardPythonParameterType({
    'none': StandardPythonParameterType(0.20599432),
    'paper': StandardPythonParameterType(0.31392053),
    'rock': StandardPythonParameterType(0.2621823),
    'scissors': StandardPythonParameterType(0.21790285)
  }),
  'timestamp': StandardPythonParameterType(datetime.datetime.now().isoformat()),
  'model_update': StandardPythonParameterType(datetime.datetime.now().isoformat()),
  'message': StandardPythonParameterType("Success!")
}))
def run(image):
    global session, transform, classes, input_name, logger, model_stamp

    print('starting inference clock')
    prev_time = time.time()

    # process data
    try:
        if image.startswith('http'):
            print(f'loading web image {image}')
            response = requests.get(image)
            img = Image.open(BytesIO(response.content))
        elif image.startswith('data:image'):
            data = image.split(",")
            print(f'loading base64 image {data[0]}')
            b64image = base64.b64decode(data[1])
            img = Image.open(BytesIO(b64image))
        else:
            print(f'loading base64 image')
            b64image = base64.b64decode(image)
            img = Image.open(BytesIO(b64image))

        v = transform(img.convert('RGB'))

        # predict with model
        print('pre-prediction')
        pred_onnx = session.run(None, {input_name: v.unsqueeze(0).numpy()})[0][0]
        print('prediction complete')

        predictions = {}
        for i in range(len(classes)):
            predictions[classes[i]] = float(pred_onnx[i])

        print('preparing payload')
        payload = {
            'time': float(0),
            'prediction': classes[int(np.argmax(pred_onnx))],
            'scores': predictions,
            'timestamp': datetime.datetime.now().isoformat(),
            'model_update': model_stamp,
            'message': 'Success!'
        }

    except Exception as e:
        predictions = {}
        for i in range(len(classes)):
            predictions[classes[i]] = float(0)

        print('preparing payload')
        payload = {
            'time': float(0),
            'prediction': "none",
            'scores': predictions,
            'timestamp': datetime.datetime.now().isoformat(),
            'model_update': model_stamp,
            'message': f'{e}'
        }

    current_time = time.time()
    print('stopping clock')
    inference_time = datetime.timedelta(seconds=current_time - prev_time)
    payload['time'] = float(inference_time.total_seconds())

    print(f'payload: {json.dumps(payload)}')
    print('inference complete')

    return payload

if __name__ == '__main__':
    init()
    def inf(uri, truth):
        print(f'---->Inference with {truth}:')
        o = run(uri)
        print(json.dumps(o, indent=4))
        print(f'---->End Inference with [{truth}] => [{o["prediction"]}]')

    inf('https://aiadvocate.z5.web.core.windows.net/rock.png', 'rock')
    inf('https://aiadvocate.z5.web.core.windows.net/paper.png', 'paper')
    inf('https://aiadvocate.z5.web.core.windows.net/scissors.png', 'scissors')
    inf('bad_uri', 'Bad Uri')

    with open('testimage.txt', 'r') as f:
        img = f.read()
        inf(img, 'rock')

    with open('fullimage.txt', 'r') as f:
        img = f.read()
        inf(img, 'none')


import torch
import argparse
import pathlib
import numpy as np
import logging
import sys
import os
from tqdm import tqdm

from vision.datasets.voc_dataset import VOCDataset
from vision.utils.misc import str2bool, Timer
from utils.utils import create_logger
from vision.ssd.ssd_mobile import create_ssd_lite, create_ssd_lite_predictor
from vision.utils.multadds_params_count import comp_multadds, count_parameters_in_MB
from eval_metrics import compute_average_precision_per_class, group_annotation_by_class

parser = argparse.ArgumentParser(description="SSDLite Evaluation.")
# It can be mobilenetv2, shufflenetv2, mnasnet, resnet18, darts, csonas, nasnet
parser.add_argument('--net', type=str, default="darts")
parser.add_argument("--eval_epoch", type=str, default="199")
parser.add_argument('--gpu', type=int, default='0')
parser.add_argument('--dataset_path', type=str, default='/raid/huangsh/datasets/PASCAL_VOC/VOCdevkit/VOC2007')
parser.add_argument("--use_cuda", type=str2bool, default=True)
parser.add_argument("--use_2007_metric", type=str2bool, default=True)
parser.add_argument("--nms_method", type=str, default="hard")
parser.add_argument("--iou_threshold", type=float, default=0.5, help="The threshold of Intersection over Union.")
args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = '{}'.format(args.gpu)
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")

if __name__ == '__main__':
    timer = Timer()
    #args.eval_epoch =
    dataset = VOCDataset(args.dataset_path, is_test=True)
    checkpoint_dir = 'outputs/voc-320-{}-cosine-e200-lr0.010'.format(args.net)
    args.eval_dir = "{}/eval_results".format(checkpoint_dir)
    eval_path = pathlib.Path(args.eval_dir)
    eval_path.mkdir(exist_ok=True)
    args.model = '{}/{}-Epoch-{}.pth'.format(checkpoint_dir, args.net, args.eval_epoch)
    args.label_file = './voc-model-labels.txt'
    class_names = [name.strip() for name in open(args.label_file).readlines()]
    if not os.path.isdir(args.eval_dir):
        os.makedirs(args.eval_dir)
    logger = create_logger(args.eval_dir, is_train=False)

    timer.start("Load Model")
    net = create_ssd_lite(len(class_names), arch=args.net, is_test=True)
    net.eval()
    params = count_parameters_in_MB(net)
    logger.info("Params = %.2fMB" % params)
    mult_adds = comp_multadds(net, input_size=(3, 320, 320))
    logger.info("Mult-Adds = %.2fMB" % mult_adds)
    net.load(args.model)
    net = net.to(DEVICE)
    predictor = create_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
    logger.info(f'It took {timer.end("Load Model")} seconds to load the model.')

    results = []
    timer.start("Detection")
    for i in tqdm(range(len(dataset))):
        timer.start("Load Image")
        image = dataset.get_image(i)
        load_t = timer.end("Load Image")
        timer.start("Predict")
        boxes, labels, probs = predictor.predict(image)
        predict_t = timer.end("Predict")
        indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i
        boxes = boxes + 1.0     # matlab's indexes start from 1 for pascal voc
        result = torch.cat([indexes.reshape(-1, 1), labels.reshape(-1, 1).float(), probs.reshape(-1, 1), boxes], dim=1)
        if result.size(1) == 7:
            results.append(result)

    logger.info(args.model)
    logger.info("     ### Detecting {} imgs with using {:.3f} ###     ".format(len(dataset), timer.end("Detection")))
    results = torch.cat(results)
    results = results
    for class_index, class_name in enumerate(class_names):
        if class_index == 0:
            continue  # ignore background
        prediction_path = eval_path / f"det_test_{class_name}.txt"
        with open(prediction_path, "w") as f:
            sub = results[results[:, 1] == class_index, :]      # all
            for i in range(sub.size(0)):
                prob_box = sub[i, 2:].numpy()
                img_id = dataset.ids[int(sub[i, 0])]
                print(str(img_id) + " " + " ".join([str(v) for v in prob_box]), file=f)

    aps = []
    logger.info("\n\nAverage Precision Per-class:")
    true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(dataset)
    for class_index, class_name in enumerate(class_names):
        if class_index == 0:
            continue
        prediction_path = eval_path / f"det_test_{class_name}.txt"
        ap = compute_average_precision_per_class(true_case_stat[class_index], all_gb_boxes[class_index], all_difficult_cases[class_index],
                                                 prediction_path, args.iou_threshold, args.use_2007_metric
                                                 )
        aps.append(ap)
        logger.info(f"{class_name}: {ap}")
    logger.info(f"\nAverage Precision Across All Classes:{sum(aps)/len(aps)}")