import tensorflow as tf
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)
from absl import app, flags, logging
from absl.flags import FLAGS
import core.utils as utils
from core.yolov4 import filter_boxes
from core.config import cfg
from core.utils import load_multiple_img
from tensorflow.python.saved_model import tag_constants
from PIL import Image
import cv2
import numpy as np
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
import os

"""
+multiple result save
+pretrained weight use
+early stopping
+negative sample training
+nms for other shape of yolo && for batch
+preprocess해서 model 결과 본 걸 다시 원래 이미지처럼 바꾸는 postprocess (resize, pad 등)
? classweight
"""

# 고정
flags.DEFINE_string('framework', 'tf', '(tf, tflite, trt)')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_float('iou', 0.5, 'iou threshold')
flags.DEFINE_float('score', 0.25, 'score threshold')
# flags.DEFINE_string('weights', r'D:\ckeckpoint\210413_custom','path to weights file')

# detection시 수정
converted_weight = r"D:\Public\JHS\YOLO\converted_weighs"   # yolo-custom-tiny 모델 pb 변환
# trained_np = r"D:\checkpoint\notf_ckeckpoint-epoch500_ok"
trained_tf = r"D:\yolov4-tflite-train_tf-epoch1000"
trained_np = r"D:\checkpoint\notf_ckpt-last"

flags.DEFINE_string('weights', trained_tf,'path to weights file')
flags.DEFINE_string('save_result_img', r"./data/detection/tfdata", 'folder to save result')
flags.DEFINE_boolean('use_trainres', True, "false for detect using converted weight(weights->pb) / True for training from scratch")
flags.DEFINE_integer('size', 640, 'resize images to')
flags.DEFINE_integer('num_detection_layer', 1, "3:yolov4 2:yolo-tiny 1:custom model")  # 디폴트 custom model
# flags.DEFINE_string('image', r"D:\Public\JHS\SIMPLE_DATA_TEST", 'path to input image (.bmp면 하나만 detect ,폴더명이면 해당 폴더 이미지 전부')
flags.DEFINE_string('image', r".\data\testtest", 'path to input image (.bmp면 하나만 detect ,폴더명이면 해당 폴더 이미지 전부')
# flags.DEFINE_string('image', r"D:\Public\JHS\Splash\OD\yolo\test", 'path to input image (.bmp면 하나만 detect ,폴더명이면 해당 폴더 이미지 전부')
flags.DEFINE_boolean('save_multiple', True, 'detection 결과 한꺼번에 여러개 보려면 True')  # convert weight의 경우 multiple 결과 이상함

def nms_convertedw(boxes, num_classes, iou_threshold=0.5, score_threshold=0.25,max_boxes_per_class=50):
    batch, num_bbox, xywhcp = tf.keras.backend.int_shape(boxes)

    boxes = np.array(boxes).reshape(1, -1, xywhcp)  # (batch, -1, xywhcp)
    boxes_conf = boxes[..., 4:]
    boxes_conf = np.max(boxes_conf, axis=-1, keepdims=True)
    conf_mask = boxes_conf >= score_threshold
    boxes_conf = boxes_conf[conf_mask]
    boxes_conf = np.reshape(boxes_conf, newshape=(1, boxes_conf.shape[0], 1))

    conf_mask = np.tile(conf_mask, (1, 1, 4 + num_classes))
    boxes = boxes[conf_mask]
    boxes = np.reshape(boxes, newshape=(1, -1, 4 + num_classes))

    boxes_classes = np.reshape(np.argmax(boxes[..., 4:], axis=-1), newshape=(1, -1, 1))
    boxes_coord = boxes[..., :4]

    bboxes_coord = []
    bboxes_scores = []
    bboxes_classes = []
    for class_ind in range(num_classes):
        mask_class = boxes_classes[..., 0:1] == class_ind
        boxes_class = boxes_classes[mask_class]
        boxes_conf_class = boxes_conf[mask_class]

        mask_class = np.tile(mask_class, (1, 1, 4))
        boxes_coord_class = boxes_coord[mask_class]
        boxes_coord_class = np.reshape(boxes_coord_class, (1, -1, 4))

        # conf 내림차순 정렬
        sorted_idx = np.argsort(boxes_conf_class)
        sorted_idx = sorted_idx[::-1]

        boxes_class = np.reshape(boxes_class, newshape=(len(sorted_idx), 1))
        boxes_class = boxes_class[sorted_idx]
        # boxes_class = np.expand_dims(boxes_class, axis=0)

        boxes_conf_class = np.reshape(boxes_conf_class, newshape=(len(sorted_idx), 1))
        boxes_conf_class = boxes_conf_class[sorted_idx]
        # boxes_conf_class = np.expand_dims(boxes_conf_class, axis=0)

        boxes_coord_class = np.reshape(boxes_coord_class, newshape=(len(sorted_idx), 4))
        boxes_coord_class = boxes_coord_class[sorted_idx]
        # boxes_coord_class = np.expand_dims(boxes_coord_class, axis=0)

        best_conf_ind = 0
        num_process = boxes_class.shape[0]
        while best_conf_ind + 1 < num_process:
            iou_scores = utils.bbox_iou(boxes_coord_class[best_conf_ind:best_conf_ind + 1, :],
                                  boxes_coord_class[best_conf_ind + 1:, :])
            iou_mask = iou_scores < iou_threshold
            iou_mask = np.reshape(iou_mask, newshape=(-1, 1))

            boxes_class = np.vstack([boxes_class[:best_conf_ind, :],
                                     np.expand_dims(boxes_class[best_conf_ind + 1:, :][iou_mask], axis=-1)])
            boxes_conf_class = np.vstack([boxes_conf_class[:best_conf_ind, :],
                                          np.expand_dims(boxes_conf_class[best_conf_ind + 1:, :][iou_mask], axis=-1)])

            iou_mask = np.tile(iou_mask, (1, 4))
            boxes_coord_class = np.vstack([boxes_coord_class[:best_conf_ind, :],
                                           np.reshape(boxes_coord_class[best_conf_ind + 1:, :][iou_mask],
                                                      newshape=(-1, 4))])

            best_conf_ind += 1
            num_process, _ = np.array(boxes_coord_class).shape

        bboxes_coord.append(boxes_coord_class)
        bboxes_scores.append(boxes_conf_class)
        bboxes_classes.append(boxes_class)

    bboxes_coord = np.vstack(bboxes_coord)
    bboxes_scores = np.vstack(bboxes_scores)
    bboxes_classes = np.vstack(bboxes_classes)
    return bboxes_coord, bboxes_scores, bboxes_classes

def nms(boxes, num_classes, iou_threshold=0.5, score_threshold=0.25,max_boxes_per_class=50):
    batch, feath, featw, num_anchors, xywhcp = tf.keras.backend.int_shape(boxes)

    boxes = np.array(boxes).reshape(1, -1, xywhcp)  # (batch, -1, xywhcp)
    boxes_conf = boxes[..., 4:5]
    boxes_classprob = boxes[..., 5:]

    boxes_conf = boxes_conf * boxes_classprob
    boxes_conf = np.max(boxes_conf, axis=-1, keepdims=True)
    conf_mask = boxes_conf >= score_threshold
    boxes_conf = boxes_conf[conf_mask]
    boxes_conf = np.reshape(boxes_conf, newshape=(1, boxes_conf.shape[0], 1))
    # print("boxes_conf:{}".format(boxes_conf))

    conf_mask = np.tile(conf_mask, (1, 1, 5 + num_classes))
    boxes = boxes[conf_mask]
    boxes = np.reshape(boxes, newshape=(1, len(boxes) // (5 + num_classes), 5 + num_classes))

    boxes_classes = np.reshape(np.argmax(boxes[..., 5:], axis=-1), newshape=(1, -1, 1))
    boxes_coord = boxes[..., :4]

    bboxes_coord = []
    bboxes_scores = []
    bboxes_classes = []
    for class_ind in range(num_classes):
        mask_class = boxes_classes[..., 0:1] == class_ind
        boxes_class = boxes_classes[mask_class]
        boxes_conf_class = boxes_conf[mask_class]

        mask_class = np.tile(mask_class, (1, 1, 4))
        boxes_coord_class = boxes_coord[mask_class]
        boxes_coord_class = np.reshape(boxes_coord_class, (1, -1, 4))

        # conf 내림차순 정렬
        sorted_idx = np.argsort(-boxes_conf_class)
        # sorted_idx = sorted_idx[::-1]

        boxes_class = np.reshape(boxes_class, newshape=(len(sorted_idx), 1))
        boxes_class = boxes_class[sorted_idx]
        # boxes_class = np.expand_dims(boxes_class, axis=0)

        boxes_conf_class = np.reshape(boxes_conf_class, newshape=(len(sorted_idx), 1))
        boxes_conf_class = boxes_conf_class[sorted_idx]
        # boxes_conf_class = np.expand_dims(boxes_conf_class, axis=0)

        boxes_coord_class = np.reshape(boxes_coord_class, newshape=(len(sorted_idx), 4))
        boxes_coord_class = boxes_coord_class[sorted_idx]
        # boxes_coord_class = np.expand_dims(boxes_coord_class, axis=0)

        best_conf_ind = 0
        num_process = boxes_class.shape[0]
        while best_conf_ind+1 < num_process:
            iou_scores = utils.bbox_iou(boxes_coord_class[best_conf_ind:best_conf_ind + 1, :],
                                  boxes_coord_class[best_conf_ind + 1:, :])
            iou_mask = iou_scores < iou_threshold
            iou_mask = np.reshape(iou_mask, newshape=(-1, 1))

            boxes_class = np.vstack([boxes_class[:best_conf_ind + 1, :],
                                     np.expand_dims(boxes_class[best_conf_ind + 1:, :][iou_mask], axis=-1)])
            boxes_conf_class = np.vstack([boxes_conf_class[:best_conf_ind + 1, :],
                                          np.expand_dims(boxes_conf_class[best_conf_ind + 1:, :][iou_mask], axis=-1)])

            iou_mask = np.tile(iou_mask, (1, 4))
            boxes_coord_class = np.vstack([boxes_coord_class[:best_conf_ind + 1, :],
                                           np.reshape(boxes_coord_class[best_conf_ind + 1:, :][iou_mask],
                                                      newshape=(-1, 4))])

            best_conf_ind += 1
            num_process, _ = np.array(boxes_coord_class).shape

        bboxes_coord.append(boxes_coord_class)
        bboxes_scores.append(boxes_conf_class)
        bboxes_classes.append(boxes_class)

    # max_bbox = max_boxes_per_class * num_classes
    bboxes_coord = np.vstack(bboxes_coord)
    bboxes_scores = np.vstack(bboxes_scores)
    bboxes_classes = np.vstack(bboxes_classes)
    return bboxes_coord, bboxes_scores, bboxes_classes

def main(_argv):
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
    input_size = FLAGS.size
    image_path = FLAGS.image

    if not os.path.exists(FLAGS.save_result_img):
        os.makedirs(FLAGS.save_result_img)

    images_path, images_data = utils.load_multiple_img(image_path, input_size, FLAGS.save_multiple)

    if FLAGS.framework == 'tflite':
        interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
        interpreter.allocate_tensors()
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()

        interpreter.set_tensor(input_details[0]['index'], images_data)
        interpreter.invoke()
        pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]

        if FLAGS.model == 'yolov3' and FLAGS.tiny == True:
            boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
        else:
            boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
    else:
        saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
        print("saved_model_loaded : {}".format(FLAGS.weights))

        infer = saved_model_loaded.signatures['serving_default']
        batch_data = tf.constant(images_data)
        pred_bbox = infer(batch_data)

        for key, value in pred_bbox.items():
            print("key:{}".format(key))
            # for custom model
            if FLAGS.use_trainres == True  and key[:5] == "tf_op":  #_layer_concat10":   # key는 model.summary의 마지막 layer # and key == "tf.concat_9": #
                if FLAGS.save_multiple == True:
                    for idx in range(len(images_data)):
                        bboxes_coord, bboxes_scores, bboxes_classes = nms(tf.expand_dims(value[idx], axis=0), num_classes=cfg.YOLO.NUM_CLASSES)
                        thres_score = 0.1
                        print(bboxes_classes)
                        image = utils.draw_bbox_trainw(images_data[idx], input_size, bboxes_coord, bboxes_classes, bboxes_scores, thres_score)

                        # check result image and save
                        cv2.imshow("result", image)
                        cv2.waitKey(0)
                        cv2.destroyAllWindows()
                        cv2.imwrite(os.path.join(FLAGS.save_result_img, "{}".format(os.path.basename(images_path[idx]))), image * 255.0)
                else:
                    bboxes_coord, bboxes_scores, bboxes_classes = nms(value, num_classes=cfg.YOLO.NUM_CLASSES)
                    image = utils.draw_bbox_trainw(images_data, input_size, bboxes_coord , bboxes_classes,  bboxes_scores)
                    # image = np.reshape(image, (input_size, input_size, 3))

                    # check result image and save
                    cv2.imshow("result", image)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()
                    cv2.imwrite(os.path.join(FLAGS.save_result_img, os.path.basename(images_path)), image)

            elif FLAGS.use_trainres == False:  # for pretrained weights (.weights -> .pb by saved_model.py) 그 결과
                if FLAGS.save_multiple == True:
                    print("이상함;;")
                    # print("value for multiple:{}".format(value))
                    # for idx in range(len(images_data)):
                    #     boxes = value[idx:idx+1, :, 0:4]
                    #     pred_conf = value[idx:idx+1, :, 4:]
                    #
                    #     boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression \
                    #             (
                    #             boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                    #             scores=tf.reshape(
                    #                 pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                    #             max_output_size_per_class=50,
                    #             max_total_size=50,
                    #             iou_threshold=FLAGS.iou,
                    #             score_threshold=FLAGS.score
                    #         )
                    #     pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
                    #     image = utils.draw_bbox(images_data[idx], pred_bbox)
                    #     cv2.imshow("result", image)
                    #     cv2.waitKey(0)
                    #     cv2.destroyAllWindows()
                else:
                    boxes = value[:, :, 0:4]
                    pred_conf = value[:, :, 4:]

                    boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression \
                            (
                            boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                            scores=tf.reshape(
                                pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                            max_output_size_per_class=50,
                            max_total_size=50,
                            iou_threshold=FLAGS.iou,
                            score_threshold=FLAGS.score
                        )
                    pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
                    image = utils.draw_bbox(images_data[0], pred_bbox)
                    cv2.imshow("result", image)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()

                    # bboxes_coord, bboxes_scores, bboxes_classes = nms_convertedw(value,
                    #                                                              num_classes=cfg.YOLO.NUM_CLASSES)
                    # image = utils.draw_bbox_convertedw(images_data, input_size, bboxes_coord, bboxes_classes,
                    #                                    bboxes_scores)
                    # image = np.reshape(image, (input_size, input_size, 3))
                    #
                    ## check result image and save
                    # cv2.imshow("result", image)
                    # cv2.waitKey(0)
                    # cv2.destroyAllWindows()
                    # cv2.imwrite(os.path.join(FLAGS.save_result_img, os.path.basename(images_path)), image)

                """
                ## 기존 - nms 이상하게 나옴 (stitch pretrained yolov4)

                boxes = value[:, :, 0:4]
                pred_conf = value[:, :, 4:]

                boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression \
                        (
                        boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
                        scores=tf.reshape(
                            pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
                        max_output_size_per_class=50,
                        max_total_size=50,
                        iou_threshold=FLAGS.iou,
                        score_threshold=FLAGS.score
                    )
                pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
                image = utils.draw_bbox(images_data[0], pred_bbox)
                cv2.imshow("result", image)
                cv2.waitKey(0)
                cv2.destroyAllWindows()

            print("valid_detections:{}".format(valid_detections))
            pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
            print("pred_bbox:{}".format(pred_bbox))
            image = utils.draw_bbox(original_image, pred_bbox)
            # cv2.imshow("result", image)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
            image = Image.fromarray(image.astype(np.uint8))
            image.show()
            image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
            cv2.imwrite(FLAGS.output, image)
            """

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


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

from __future__ import print_function

import argparse
import cv2
import numpy as np
import os
import shutil
import sys

FIXEXP = True  # Flag to fix underexposition
MINFACE = 8  # Minimum face size ratio; too low and we get false positives
INCREMENT = 0.06
GAMMA_THRES = 0.001
GAMMA = 0.90
FACE_RATIO = 6  # Face / padding ratio
QUESTION_OVERWRITE = "Overwrite image files?"
FILETYPES = ['.jpg', '.jpeg', '.bmp', '.dib', '.jp2',
             '.png', '.webp', '.pbm', '.pgm', '.ppm',
             '.sr', '.ras', '.tiff', '.tif']
INPUT_FILETYPES = FILETYPES + [s.upper() for s in FILETYPES]

# Load XML Resource
cascFile = 'haarcascade_frontalface_default.xml'
cascPath = os.path.join("./", cascFile)

# Define simple gamma correction fn
def gamma(img, correction):
    img = cv2.pow(img/255.0, correction)
    return np.uint8(img*255)

def crop_positions(imgh, imgw, x, y, w, h,
                   fheight, fwidth, facePercent,
                   padUp, padDown, padLeft, padRight):
    # Check padding values
    padUp = 50 if (padUp is False or padUp < 0) else padUp
    padDown = 50 if (padDown is False or padDown < 0) else padDown
    padLeft = 50 if (padLeft is False or padLeft < 0) else padLeft
    padRight = 50 if (padRight is False or padRight < 0) else padRight

    # enfoce face percent
    facePercent = 100 if facePercent > 100 else facePercent
    facePercent = 50 if facePercent <= 0 else facePercent

    # Adjust output height based on Face percent
    height_crop = h * 100.0 / facePercent

    # Ensure height is within boundaries
    height_crop = imgh if height_crop > imgh else height_crop

    aspect_ratio = float(fwidth) / float(fheight)
    # Calculate width based on aspect ratio
    width_crop = aspect_ratio * float(height_crop)

    # Calculate padding by centering face
    xpad = (width_crop - w) / 2
    ypad = (height_crop - h) / 2

    # Calc. positions of crop
    h1 = float(x - (xpad * padLeft / (padLeft + padRight)))
    h2 = float(x + w + (xpad * padRight / (padLeft + padRight)))
    v1 = float(y - (ypad * padUp / (padUp + padDown)))
    v2 = float(y + h + (ypad * padDown / (padUp + padDown)))

    # Move crop inside photo boundaries
    while h1 < 0:
        h1 = h1 + 1
        h2 = h2 + 1
    while v1 < 0:
        v1 = v1 + 1
        v2 = v2 + 1
    while v2 > imgh:
        v2 = v2 - 1
        h2 = h2 - 1 * aspect_ratio
    while h2 > imgw:
        h2 = h2 - 1
        v2 = v2 - 1 / aspect_ratio

    return [int(v1), int(v2), int(h1), int(h2)]

def crop(image, fheight=500, fwidth=500, facePercent=50,
         padUp=False, padDown=False, padLeft=False, padRight=False):
    """Given a ndarray image with a face, returns cropped array.

    Arguments:
        - image, the numpy array of the image to be processed.
        - fwidth, the final width (px) of the cropped img. Default: 500
        - fheight, the final height (px) of the cropped img. Default: 500
        - facePercent, percentage of face height to image height. Default: 50
        - padUp, Padding from top
        - padDown, Padding to bottom
        - padLeft, Padding from left
        - padRight, Padding to right
    Returns:
        - image, a cropped numpy array

    ndarray, int, int -> ndarray
    """
    # Some grayscale color profiles can throw errors, catch them
    try:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    except cv2.error:
        gray = image

    # Scale the image
    height, width = (image.shape[:2])
    minface = int(np.sqrt(height**2 + width**2) / MINFACE)

    # Create the haar cascade
    faceCascade = cv2.CascadeClassifier(cascPath)

    # ====== Detect faces in the image ======
    faces = faceCascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(minface, minface),
        flags=cv2.CASCADE_FIND_BIGGEST_OBJECT | cv2.CASCADE_DO_ROUGH_SEARCH,
    )

    # Handle no faces
    if len(faces) == 0:
        return None

    # Make padding from biggest face found
    x, y, w, h = faces[-1]

    pos = crop_positions(height, width, x, y, w, h,
                         fheight, fwidth, facePercent,
                         padUp, padDown, padLeft, padRight)
    # Actual cropping
    image = image[pos[0]:pos[1], pos[2]:pos[3]]

    # Resize
    image = cv2.resize(image, (fwidth, fheight),
                       interpolation=cv2.INTER_AREA)

    # ====== Dealing with underexposition ======
    if FIXEXP:
        # Check if under-exposed
        uexp = cv2.calcHist([gray], [0], None, [256], [0, 256])
        if sum(uexp[-26:]) < GAMMA_THRES * sum(uexp):
            image = gamma(image, GAMMA)
    return image

def main(input_d,
         output_d,
         fheight=500,
         fwidth=500,
         facePercent=50,
         padUp=False,
         padDown=False,
         padLeft=False,
         padRight=False):
    """Crops folder of images to the desired height and width if a face is found

    If input_d == output_d or output_d is None, overwrites all files
    where the biggest face was found.

    Args:
        input_d (str): Directory to crop images from.
        output_d (str): Directory where cropped images are placed.
        fheight (int): Height (px) to which to crop the image.
                       Default: 500px
        fwidth (int): Width (px) to which to crop the image.
                       Default: 500px
        facePercent (int) : Percentage of face from height,
                       Default: 50
    Side Effects:
        Creates image files in output directory.

    str, str, (int), (int) -> None
    """
    errors = 0
    files = [os.path.join(input_d, f) for f in os.listdir(input_d)
             if any(f.endswith(t) for t in INPUT_FILETYPES)]

    # Guard against calling the function directly
    assert len(files) > 0

    if output_d is not None:
        filenames = [os.path.basename(f) for f in files]
        target_files = [os.path.join(output_d, fn) for fn in filenames]
        for i, o in zip(files, target_files):
            shutil.copyfile(i, o)
        files = target_files
    else:
        output_d = input_d

    for f in files:
        try:
            filename = os.path.basename(f)

            # Perform the actual crop
            input_img = cv2.imread(f)
            image = crop(input_img,
                        fheight,
                        fwidth,
                        facePercent,
                        padUp,
                        padDown,
                        padLeft,
                        padRight)

            # Make sure there actually was a face in there
            if isinstance(image, type(None)):
                print('No faces can be detected in {}.'.format(filename))
                errors += 1
                continue

            # Write cropfile
            output_filename = os.path.join(output_d, filename)
            cv2.imwrite(output_filename, image)
        except:
            pass

    # Stop and print status
    print(' {} files have been cropped'.format(len(files) - errors))

def input_path(p):
    """Returns absolute path, only if input is a valid directory"""
    no_folder = 'Input folder does not exist'
    no_images = 'Input folder does not contain any image files'
    p = os.path.abspath(p)
    if not os.path.isdir(p):
        raise argparse.ArgumentTypeError(no_folder)
    filetypes = set(os.path.splitext(f)[-1] for f in os.listdir(p))
    if not any(t in INPUT_FILETYPES for t in filetypes):
        raise argparse.ArgumentTypeError(no_images)
    else:
        return p

def output_path(p):
    """Returns absolute path, if input is a valid directory name.
    If directory doesn't exist, creates it."""
    p = os.path.abspath(p)
    if not os.path.isdir(p):
        os.makedirs(p)
    return p

def size(i):
    """Returns valid only if input is a positive integer under 1e5"""
    error = 'Invalid pixel size'
    try:
        i = int(i)
    except TypeError:
        raise argparse.ArgumentTypeError(error)
    if i > 0 and i < 1e5:
        return i
    else:
        raise argparse.ArgumentTypeError(error)

def compat_input(s=''):
    """Compatibility function to permit testing for Python 2 and 3"""
    try:
        return raw_input(s)
    except NameError:
        return input(s)

def confirmation(question, default=True):
    """Ask a yes/no question via standard input and return the answer.

    If invalid input is given, the user will be asked until
    they acutally give valid input.

    Args:
        question(str):
            A question that is presented to the user.
        default(bool|None):
            The default value when enter is pressed with no value.
            When None, there is no default value and the query
            will loop.
    Returns:
        A bool indicating whether user has entered yes or no.

    Side Effects:
        Blocks program execution until valid input(y/n) is given.
    """
    yes_list = ["yes", "y"]
    no_list = ["no", "n"]

    default_dict = {  # default => prompt default string
        None: "[y/n]",
        True: "[Y]/n",
        False: "y/[N]",
    }

    default_str = default_dict[default]
    prompt_str = "%s %s " % (question, default_str)

    while True:
        choice = compat_input(prompt_str).lower()

        if not choice and default is not None:
            return default
        if choice in yes_list:
            return True
        if choice in no_list:
            return False

        notification_str = "Please respond with 'y' or 'n'"
        print(notification_str)

def parse_args(args):
    help_d = {
            'desc': 'Automatically crops faces from batches of pictures',
            'input': '''Folder where images to crop are located.
Default: current working directory''',
            'output': '''Folder where cropped images will be placed.
Default: current working directory''',
            'width': 'Width of cropped files in px. Default=500',
            'height': 'Height of cropped files in px. Default=500',
            'y': 'Bypass any confirmation prompts',
            'facePercent': 'Percentage of face to image height',
            'padUp': 'Add padding up to face cropped',
            'padDown': 'Add padding down to face cropped',
            'padLeft': 'Add padding left to face cropped',
            'padRight': 'Add padding right to face cropped',
            }

    parser = argparse.ArgumentParser(description=help_d['desc'])
    parser.add_argument('-o', '--output', '-p', '--path', type=output_path,
                        default=None, help=help_d['output'])
    parser.add_argument('-i', '--input', default='.', type=input_path,
                        help=help_d['input'])
    parser.add_argument('-w', '--width', type=size,
                        default=500, help=help_d['width'])
    parser.add_argument('-H', '--height',
                        type=size, default=500, help=help_d['height'])
    parser.add_argument('--no-confirm', action='store_true', help=help_d['y'])
    parser.add_argument('--padUp', type=size,
                        default=False, help=help_d['padUp'])
    parser.add_argument('--padDown', type=size,
                        default=False, help=help_d['padDown'])
    parser.add_argument('--padLeft', type=size,
                        default=False, help=help_d['padLeft'])
    parser.add_argument('--padRight', type=size,
                        default=False, help=help_d['padRight'])
    parser.add_argument('--facePercent', type=size,
                        default=50, help=help_d['facePercent'])

    return parser.parse_args()

def cli():
    args = parse_args(sys.argv[1:])
    if not args.no_confirm:
        if args.output is None or args.input == args.output:
            if not confirmation(QUESTION_OVERWRITE):
                sys.exit()
    if args.input == args.output:
        args.output = None
    print('Processing images in folder:', args.input)
    main(args.input,
         args.output,
         args.height,
         args.width,
         args.facePercent,
         args.padUp,
         args.padDown,
         args.padLeft,
         args.padRight)

import os
import sys

if __name__ == '__main__':
    cli()


# -*- coding: utf-8 -*-
"""03.05-Hierarchical-Indexing.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/jakevdp/PythonDataScienceHandbook/blob/master/notebooks/03.05-Hierarchical-Indexing.ipynb

*This notebook contains an excerpt from the [Python Data Science Handbook](http://shop.oreilly.com/product/0636920034919.do) by Jake VanderPlas; the content is available [on GitHub](https://github.com/jakevdp/PythonDataScienceHandbook).*

*The text is released under the [CC-BY-NC-ND license](https://creativecommons.org/licenses/by-nc-nd/3.0/us/legalcode), and code is released under the [MIT license](https://opensource.org/licenses/MIT). If you find this content useful, please consider supporting the work by [buying the book](http://shop.oreilly.com/product/0636920034919.do)!*

< [Handling Missing Data](03.04-Missing-Values.ipynb) | [Contents](Index.ipynb) | [Combining Datasets: Concat and Append](03.06-Concat-And-Append.ipynb) >

# Hierarchical Indexing

Up to this point we've been focused primarily on one-dimensional and two-dimensional data, stored in Pandas ``Series`` and ``DataFrame`` objects, respectively.
Often it is useful to go beyond this and store higher-dimensional data–that is, data indexed by more than one or two keys.
While Pandas does provide ``Panel`` and ``Panel4D`` objects that natively handle three-dimensional and four-dimensional data (see [Aside: Panel Data](#Aside:-Panel-Data)), a far more common pattern in practice is to make use of *hierarchical indexing* (also known as *multi-indexing*) to incorporate multiple index *levels* within a single index.
In this way, higher-dimensional data can be compactly represented within the familiar one-dimensional ``Series`` and two-dimensional ``DataFrame`` objects.

In this section, we'll explore the direct creation of ``MultiIndex`` objects, considerations when indexing, slicing, and computing statistics across multiply indexed data, and useful routines for converting between simple and hierarchically indexed representations of your data.

We begin with the standard imports:
"""

import pandas as pd
import numpy as np

"""## A Multiply Indexed Series

Let's start by considering how we might represent two-dimensional data within a one-dimensional ``Series``.
For concreteness, we will consider a series of data where each point has a character and numerical key.

### The bad way

Suppose you would like to track data about states from two different years.
Using the Pandas tools we've already covered, you might be tempted to simply use Python tuples as keys:
"""

index = [('California', 2000), ('California', 2010),
         ('New York', 2000), ('New York', 2010),
         ('Texas', 2000), ('Texas', 2010)]
populations = [33871648, 37253956,
               18976457, 19378102,
               20851820, 25145561]
pop = pd.Series(populations, index=index)
pop

"""With this indexing scheme, you can straightforwardly index or slice the series based on this multiple index:"""

pop[('California', 2010):('Texas', 2000)]

"""But the convenience ends there. For example, if you need to select all values from 2010, you'll need to do some messy (and potentially slow) munging to make it happen:"""

pop[[i for i in pop.index if i[1] == 2010]]

"""This produces the desired result, but is not as clean (or as efficient for large datasets) as the slicing syntax we've grown to love in Pandas.

### The Better Way: Pandas MultiIndex
Fortunately, Pandas provides a better way.
Our tuple-based indexing is essentially a rudimentary multi-index, and the Pandas ``MultiIndex`` type gives us the type of operations we wish to have.
We can create a multi-index from the tuples as follows:
"""

index = pd.MultiIndex.from_tuples(index)
index

"""Notice that the ``MultiIndex`` contains multiple *levels* of indexing–in this case, the state names and the years, as well as multiple *labels* for each data point which encode these levels.

If we re-index our series with this ``MultiIndex``, we see the hierarchical representation of the data:
"""

pop = pop.reindex(index)
pop

"""Here the first two columns of the ``Series`` representation show the multiple index values, while the third column shows the data.
Notice that some entries are missing in the first column: in this multi-index representation, any blank entry indicates the same value as the line above it.

Now to access all data for which the second index is 2010, we can simply use the Pandas slicing notation:
"""

pop[:, 2010]

"""The result is a singly indexed array with just the keys we're interested in.
This syntax is much more convenient (and the operation is much more efficient!) than the home-spun tuple-based multi-indexing solution that we started with.
We'll now further discuss this sort of indexing operation on hieararchically indexed data.

### MultiIndex as extra dimension

You might notice something else here: we could easily have stored the same data using a simple ``DataFrame`` with index and column labels.
In fact, Pandas is built with this equivalence in mind. The ``unstack()`` method will quickly convert a multiply indexed ``Series`` into a conventionally indexed ``DataFrame``:
"""

pop_df = pop.unstack()
pop_df

"""Naturally, the ``stack()`` method provides the opposite operation:"""

pop_df.stack()

"""Seeing this, you might wonder why would we would bother with hierarchical indexing at all.
The reason is simple: just as we were able to use multi-indexing to represent two-dimensional data within a one-dimensional ``Series``, we can also use it to represent data of three or more dimensions in a ``Series`` or ``DataFrame``.
Each extra level in a multi-index represents an extra dimension of data; taking advantage of this property gives us much more flexibility in the types of data we can represent. Concretely, we might want to add another column of demographic data for each state at each year (say, population under 18) ; with a ``MultiIndex`` this is as easy as adding another column to the ``DataFrame``:
"""

pop_df = pd.DataFrame({'total': pop,
                       'under18': [9267089, 9284094,
                                   4687374, 4318033,
                                   5906301, 6879014]})
pop_df

"""In addition, all the ufuncs and other functionality discussed in [Operating on Data in Pandas](03.03-Operations-in-Pandas.ipynb) work with hierarchical indices as well.
Here we compute the fraction of people under 18 by year, given the above data:
"""

f_u18 = pop_df['under18'] / pop_df['total']
f_u18.unstack()

"""This allows us to easily and quickly manipulate and explore even high-dimensional data.

## Methods of MultiIndex Creation

The most straightforward way to construct a multiply indexed ``Series`` or ``DataFrame`` is to simply pass a list of two or more index arrays to the constructor. For example:
"""

df = pd.DataFrame(np.random.rand(4, 2),
                  index=[['a', 'a', 'b', 'b'], [1, 2, 1, 2]],
                  columns=['data1', 'data2'])
df

"""The work of creating the ``MultiIndex`` is done in the background.

Similarly, if you pass a dictionary with appropriate tuples as keys, Pandas will automatically recognize this and use a ``MultiIndex`` by default:
"""

data = {('California', 2000): 33871648,
        ('California', 2010): 37253956,
        ('Texas', 2000): 20851820,
        ('Texas', 2010): 25145561,
        ('New York', 2000): 18976457,
        ('New York', 2010): 19378102}
pd.Series(data)

"""Nevertheless, it is sometimes useful to explicitly create a ``MultiIndex``; we'll see a couple of these methods here.

### Explicit MultiIndex constructors

For more flexibility in how the index is constructed, you can instead use the class method constructors available in the ``pd.MultiIndex``.
For example, as we did before, you can construct the ``MultiIndex`` from a simple list of arrays giving the index values within each level:
"""

pd.MultiIndex.from_arrays([['a', 'a', 'b', 'b'], [1, 2, 1, 2]])

"""You can construct it from a list of tuples giving the multiple index values of each point:"""

pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1), ('b', 2)])

"""You can even construct it from a Cartesian product of single indices:"""

pd.MultiIndex.from_product([['a', 'b'], [1, 2]])

"""Similarly, you can construct the ``MultiIndex`` directly using its internal encoding by passing ``levels`` (a list of lists containing available index values for each level) and ``labels`` (a list of lists that reference these labels):"""

#pd.MultiIndex(levels=[['a', 'b'], [1, 2]],  labels=[[0, 0, 1, 1], [0, 1, 0, 1]])

"""Any of these objects can be passed as the ``index`` argument when creating a ``Series`` or ``Dataframe``, or be passed to the ``reindex`` method of an existing ``Series`` or ``DataFrame``.

### MultiIndex level names

Sometimes it is convenient to name the levels of the ``MultiIndex``.
This can be accomplished by passing the ``names`` argument to any of the above ``MultiIndex`` constructors, or by setting the ``names`` attribute of the index after the fact:
"""

pop.index.names = ['state', 'year']
pop

"""With more involved datasets, this can be a useful way to keep track of the meaning of various index values.

### MultiIndex for columns

In a ``DataFrame``, the rows and columns are completely symmetric, and just as the rows can have multiple levels of indices, the columns can have multiple levels as well.
Consider the following, which is a mock-up of some (somewhat realistic) medical data:
"""

# hierarchical indices and columns
index = pd.MultiIndex.from_product([[2013, 2014], [1, 2]],
                                   names=['year', 'visit'])
columns = pd.MultiIndex.from_product([['Bob', 'Guido', 'Sue'], ['HR', 'Temp']],
                                     names=['subject', 'type'])

# mock some data
data = np.round(np.random.randn(4, 6), 1)
data[:, ::2] *= 10
data += 37

# create the DataFrame
health_data = pd.DataFrame(data, index=index, columns=columns)
health_data

"""Here we see where the multi-indexing for both rows and columns can come in *very* handy.
This is fundamentally four-dimensional data, where the dimensions are the subject, the measurement type, the year, and the visit number.
With this in place we can, for example, index the top-level column by the person's name and get a full ``DataFrame`` containing just that person's information:
"""

health_data['Guido']

"""For complicated records containing multiple labeled measurements across multiple times for many subjects (people, countries, cities, etc.) use of hierarchical rows and columns can be extremely convenient!

## Indexing and Slicing a MultiIndex

Indexing and slicing on a ``MultiIndex`` is designed to be intuitive, and it helps if you think about the indices as added dimensions.
We'll first look at indexing multiply indexed ``Series``, and then multiply-indexed ``DataFrame``s.

### Multiply indexed Series

Consider the multiply indexed ``Series`` of state populations we saw earlier:
"""

pop

"""We can access single elements by indexing with multiple terms:"""

pop['California', 2000]

"""The ``MultiIndex`` also supports *partial indexing*, or indexing just one of the levels in the index.
The result is another ``Series``, with the lower-level indices maintained:
"""

pop['California']

"""Partial slicing is available as well, as long as the ``MultiIndex`` is sorted (see discussion in [Sorted and Unsorted Indices](#Sorted-and-unsorted-indices)):"""

pop.loc['California':'New York']

"""With sorted indices, partial indexing can be performed on lower levels by passing an empty slice in the first index:"""

pop[:, 2000]

"""Other types of indexing and selection (discussed in [Data Indexing and Selection](03.02-Data-Indexing-and-Selection.ipynb)) work as well; for example, selection based on Boolean masks:"""

pop[pop > 22000000]

"""Selection based on fancy indexing also works:"""

pop[['California', 'Texas']]

"""### Multiply indexed DataFrames

A multiply indexed ``DataFrame`` behaves in a similar manner.
Consider our toy medical ``DataFrame`` from before:
"""

health_data

"""Remember that columns are primary in a ``DataFrame``, and the syntax used for multiply indexed ``Series`` applies to the columns.
For example, we can recover Guido's heart rate data with a simple operation:
"""

health_data['Guido', 'HR']

"""Also, as with the single-index case, we can use the ``loc``, ``iloc``, and ``ix`` indexers introduced in [Data Indexing and Selection](03.02-Data-Indexing-and-Selection.ipynb). For example:"""

health_data.iloc[:2, :2]

"""These indexers provide an array-like view of the underlying two-dimensional data, but each individual index in ``loc`` or ``iloc`` can be passed a tuple of multiple indices. For example:"""

health_data.loc[:, ('Bob', 'HR')]

"""Working with slices within these index tuples is not especially convenient; trying to create a slice within a tuple will lead to a syntax error:"""

#health_data.loc[(:, 1), (:, 'HR')]

"""You could get around this by building the desired slice explicitly using Python's built-in ``slice()`` function, but a better way in this context is to use an ``IndexSlice`` object, which Pandas provides for precisely this situation.
For example:
"""

idx = pd.IndexSlice
health_data.loc[idx[:, 1], idx[:, 'HR']]

"""There are so many ways to interact with data in multiply indexed ``Series`` and ``DataFrame``s, and as with many tools in this book the best way to become familiar with them is to try them out!

## Rearranging Multi-Indices

One of the keys to working with multiply indexed data is knowing how to effectively transform the data.
There are a number of operations that will preserve all the information in the dataset, but rearrange it for the purposes of various computations.
We saw a brief example of this in the ``stack()`` and ``unstack()`` methods, but there are many more ways to finely control the rearrangement of data between hierarchical indices and columns, and we'll explore them here.

### Sorted and unsorted indices

Earlier, we briefly mentioned a caveat, but we should emphasize it more here.
*Many of the ``MultiIndex`` slicing operations will fail if the index is not sorted.*
Let's take a look at this here.

We'll start by creating some simple multiply indexed data where the indices are *not lexographically sorted*:
"""

index = pd.MultiIndex.from_product([['a', 'c', 'b'], [1, 2]])
data = pd.Series(np.random.rand(6), index=index)
data.index.names = ['char', 'int']
data

"""If we try to take a partial slice of this index, it will result in an error:"""

try:
    data['a':'b']
except KeyError as e:
    print(type(e))
    print(e)

"""Although it is not entirely clear from the error message, this is the result of the MultiIndex not being sorted.
For various reasons, partial slices and other similar operations require the levels in the ``MultiIndex`` to be in sorted (i.e., lexographical) order.
Pandas provides a number of convenience routines to perform this type of sorting; examples are the ``sort_index()`` and ``sortlevel()`` methods of the ``DataFrame``.
We'll use the simplest, ``sort_index()``, here:
"""

data = data.sort_index()
data

"""With the index sorted in this way, partial slicing will work as expected:"""

data['a':'b']

"""### Stacking and unstacking indices

As we saw briefly before, it is possible to convert a dataset from a stacked multi-index to a simple two-dimensional representation, optionally specifying the level to use:
"""

pop.unstack(level=0)

pop.unstack(level=1)

"""The opposite of ``unstack()`` is ``stack()``, which here can be used to recover the original series:"""

pop.unstack().stack()

"""### Index setting and resetting

Another way to rearrange hierarchical data is to turn the index labels into columns; this can be accomplished with the ``reset_index`` method.
Calling this on the population dictionary will result in a ``DataFrame`` with a *state* and *year* column holding the information that was formerly in the index.
For clarity, we can optionally specify the name of the data for the column representation:
"""

pop_flat = pop.reset_index(name='population')
pop_flat

"""Often when working with data in the real world, the raw input data looks like this and it's useful to build a ``MultiIndex`` from the column values.
This can be done with the ``set_index`` method of the ``DataFrame``, which returns a multiply indexed ``DataFrame``:
"""

pop_flat.set_index(['state', 'year'])

"""In practice, I find this type of reindexing to be one of the more useful patterns when encountering real-world datasets.

## Data Aggregations on Multi-Indices

We've previously seen that Pandas has built-in data aggregation methods, such as ``mean()``, ``sum()``, and ``max()``.
For hierarchically indexed data, these can be passed a ``level`` parameter that controls which subset of the data the aggregate is computed on.

For example, let's return to our health data:
"""

health_data

"""Perhaps we'd like to average-out the measurements in the two visits each year. We can do this by naming the index level we'd like to explore, in this case the year:"""

data_mean = health_data.mean(level='year')
data_mean

"""By further making use of the ``axis`` keyword, we can take the mean among levels on the columns as well:"""

data_mean.mean(axis=1, level='type')

"""Thus in two lines, we've been able to find the average heart rate and temperature measured among all subjects in all visits each year.
This syntax is actually a short cut to the ``GroupBy`` functionality, which we will discuss in [Aggregation and Grouping](03.08-Aggregation-and-Grouping.ipynb).
While this is a toy example, many real-world datasets have similar hierarchical structure.

## Aside: Panel Data

Pandas has a few other fundamental data structures that we have not yet discussed, namely the ``pd.Panel`` and ``pd.Panel4D`` objects.
These can be thought of, respectively, as three-dimensional and four-dimensional generalizations of the (one-dimensional) ``Series`` and (two-dimensional) ``DataFrame`` structures.
Once you are familiar with indexing and manipulation of data in a ``Series`` and ``DataFrame``, ``Panel`` and ``Panel4D`` are relatively straightforward to use.
In particular, the ``ix``, ``loc``, and ``iloc`` indexers discussed in [Data Indexing and Selection](03.02-Data-Indexing-and-Selection.ipynb) extend readily to these higher-dimensional structures.

We won't cover these panel structures further in this text, as I've found in the majority of cases that multi-indexing is a more useful and conceptually simpler representation for higher-dimensional data.
Additionally, panel data is fundamentally a dense data representation, while multi-indexing is fundamentally a sparse data representation.
As the number of dimensions increases, the dense representation can become very inefficient for the majority of real-world datasets.
For the occasional specialized application, however, these structures can be useful.
If you'd like to read more about the ``Panel`` and ``Panel4D`` structures, see the references listed in [Further Resources](03.13-Further-Resources.ipynb).

< [Handling Missing Data](03.04-Missing-Values.ipynb) | [Contents](Index.ipynb) | [Combining Datasets: Concat and Append](03.06-Concat-And-Append.ipynb) >

"""


from functools import partial

import numpy as np
from aiohttp import web

from .worker import predict

class MainHandler:
    def __init__(self, loop, executor, slack_client, giphy_client):
        self.loop = loop
        self.executor = executor
        self.slack_client = slack_client
        self.giphy_client = giphy_client
        self._toxicity_index = 0.4
        self._run = partial(self.loop.run_in_executor, self.executor)

    async def listen_message(self, request):
        post = await request.json()

        if "challenge" in post:
            return web.Response(body=post["challenge"])

        if post["event"]["type"] == "message":
            await self._message_handler(post["event"])

        return web.Response()

    async def _respond(self, event):
        result = await self.giphy_client.get(
            "random",
            params={
                "tag": "funny cat"
            },
        )
        image_url = result["data"]["image_url"]

        text = (f"Hey <@{event['user']}>, please be polite! "
                f"Here is a funny cat GIF for you {image_url}")

        await self.slack_client.chat.post_message(
            event["channel"],
            text=text,
        )

    async def _message_handler(self, event):
        scores = await self._run(predict, event["text"])
        if np.average([scores.toxic, scores.insult]) >= self._toxicity_index:
            await self._respond(event)


from os import listdir
from os.path import isfile, join
import librosa
import numpy as np
import tensorflow as tf
import pandas as pd

def generate_input_from_audio_file(path_to_audio_file, resample_to=8000):
    """
    function to create input for our neural network from an audio file.
    The function loads the audio file using librosa, resamples it, and creates spectrogram form it
    :param path_to_audio_file (string): path to the audio file
    :param resample_to (int): sampling rate
    :return: spectrogram corresponding to the input file
    """

    signal, sample_rate = librosa.core.load(path_to_audio_file)
    if signal.shape[0] == 2:
        signal = np.mean(signal, axis=0)
    signal_resampled = librosa.core.resample(signal, sample_rate, resample_to)

    mfcc1 = librosa.feature.mfcc(
        y=signal_resampled, sr=resample_to, n_fft=200, hop_length=80
    )
    lenmfcc = int(mfcc1.shape[1])
    return mfcc1.T, lenmfcc

def generate_target_output_from_text(target_text):
    """
    Target output is an array of indices for each character in your string.
    The indices comes from a mapping that will
    be used while decoding the ctc output.
    :param target_text (string): target string
    :return: array of indices for each character in the string
    """
    gujarati_alphabet = [
        "\u0A81",
        "\u0A82",
        "\u0A83",
        "\u0A85",
        "\u0A86",
        "\u0A87",
        "\u0A88",
        "\u0A89",
        "\u0A8A",
        "\u0A8B",
        "\u0A8C",
        "\u0A8D",
        "\u0A8F",
        "\u0A90",
        "\u0A91",
        "\u0A93",
        "\u0A94",
        "\u0A95",
        "\u0A96",
        "\u0A97",
        "\u0A98",
        "\u0A99",
        "\u0A9A",
        "\u0A9B",
        "\u0A9C",
        "\u0A9D",
        "\u0A9E",
        "\u0A9F",
        "\u0AA0",
        "\u0AA1",
        "\u0AA2",
        "\u0AA3",
        "\u0AA4",
        "\u0AA5",
        "\u0AA6",
        "\u0AA7",
        "\u0AA8",
        "\u0AAA",
        "\u0AAB",
        "\u0AAC",
        "\u0AAD",
        "\u0AAE",
        "\u0AAF",
        "\u0AB0",
        "\u0AB2",
        "\u0AB3",
        "\u0AB5",
        "\u0AB6",
        "\u0AB7",
        "\u0AB8",
        "\u0AB9",
        "\u0ABC",
        "\u0ABD",
        "\u0ABE",
        "\u0ABF",
        "\u0AC0",
        "\u0AC1",
        "\u0AC2",
        "\u0AC3",
        "\u0AC4",
        "\u0AC5",
        "\u0AC7",
        "\u0AC8",
        "\u0AC9",
        "\u0ACB",
        "\u0ACC",
        "\u0ACD",
        "\u0AD0",
        "\u0AE0",
        "\u0AE1",
        "\u0AE2",
        "\u0AE3",
        "\u0AF1",
    ]
    space_token = " "
    end_token = ">"
    blank_token = "%"
    alphabet = gujarati_alphabet + [space_token, end_token, blank_token]
    char_to_index = {}
    for idx, char in enumerate(alphabet):
        char_to_index[char] = idx

    y = []
    for char in target_text:
        try:
            y.append(char_to_index[char])
        except:
            continue

    y = np.array(y)
    leny = y.shape[0]

    return y, leny

def encode_input_and_output(t_data, lsx, lsy, lenmfcc, leny, PathDataAudios):
    """
    this function will encode all the audios and their respective
    transcripts in format required for training model
    :param t_data (DataFrame): pandas dataframe containing the audio file names and their respective transcripts
    :param lsx (list): list of all encoded input audio
    :param lsy (list): list of all encoded input transcript
    :param lenmfcc (list): list of length of all encoded input audio
    :param leny (list): list of length of all encoded transcripts
    :param PathDataAudios (int): Path to folder where all audio files are stored
    """
    for d in t_data.itertuples():
        id = str(d.Id).zfill(9)
        text = str(d.Text).lower()
        AudioFileExtention = ".wav"
        path = PathDataAudios + id + AudioFileExtention
        X, lenmfcc_per = generate_input_from_audio_file(path)
        y, leny_per = generate_target_output_from_text(text)
        lsx.append(X)
        lsy.append(y)
        lenmfcc.append(lenmfcc_per)
        leny.append(leny_per)
    return lsx, lsy, lenmfcc, leny

def feature_extraction(
    train_val_split=0.8, pad_value_audio_feature=0.0, pad_value_character_index=75
):
    """
    this function will extract the features from the audio stored at the data directory (here "./Data/Train/Audios/")
    :param train_val_split (float): percent of data in training (range: 0-1)
    :param pad_value_audio_feature (float): padding value for input features (X)
    :param pad_value_character_index (int): padding value for the formatted references (Y)
    """
    lsx = []
    lsy = []
    lenmfcc = []
    leny = []
    PathDataAudios = (
        "./Data Files/Train/Audios/"  # path to the training data audios (X)
    )

    PathDataTranscripts = "./Data Files/Train/transcription.txt"  # path to the training data transcripts (Y)
    # Read number of audio files in the ./Data/Train/Audios
    total_files = [
        f for f in listdir(PathDataAudios) if isfile(join(PathDataAudios, f))
    ]

    # total number of files
    Totald = len(total_files)
    # 0 till Traind will be training data and from Traind till end it will be validation data
    Traind = int(Totald * train_val_split)
    t_data = pd.read_csv(PathDataTranscripts, sep="\t", names=["Id", "Text"])
    t_data.head()

    # generate input feature from audio files
    lsx, lsy, lenmfcc, leny = encode_input_and_output(
        t_data[:Totald], lsx, lsy, lenmfcc, leny, PathDataAudios
    )

    # lsx = lsx[: 2000 + 428]
    # lsy = lsy[: 2000 + 428]
    # lenmfcc = lenmfcc[: 2000 + 428]
    # leny = leny[: 2000 + 428]

    # Validation data
    lsx_v, lsy_v, lenmfcc_v, leny_v = (
        lsx[Traind:],
        lsy[Traind:],
        lenmfcc[Traind:],
        leny[Traind:],
    )
    # trining data
    lsx, lsy, lenmfcc, leny = (
        lsx[:Traind],
        lsy[:Traind],
        lenmfcc[:Traind],
        leny[:Traind],
    )

    """
        maxlen parameter of "tf.keras.preprocessing.sequence.pad_sequences" removes all
        the audio feature having length greater than its given value. It helps to limit length of the feature and is
        helpful in case of memory constraint

        If you have no RAM constraint you can use dynamic padding, that pads to max feature length. Remove 'maxlen' parameter.
    """

    lsx = tf.keras.preprocessing.sequence.pad_sequences(
        lsx, padding="post", dtype="float32", value=pad_value_audio_feature, maxlen=1000
    )

    lsy = tf.keras.preprocessing.sequence.pad_sequences(
        lsy, padding="post", value=pad_value_character_index, maxlen=100
    )

    lsx_v = tf.keras.preprocessing.sequence.pad_sequences(
        lsx_v,
        padding="post",
        dtype="float32",
        value=pad_value_audio_feature,
        maxlen=1000,
    )
    lsy_v = tf.keras.preprocessing.sequence.pad_sequences(
        lsy_v, padding="post", value=pad_value_character_index, maxlen=100
    )

    # convert to tensor so it is trainable by tensorflow
    lsx = tf.convert_to_tensor(lsx)
    lsy = tf.convert_to_tensor(lsy)

    lenmfcc = tf.convert_to_tensor(lenmfcc)
    leny = tf.convert_to_tensor(leny)

    lsx_v = tf.convert_to_tensor(lsx_v)
    lsy_v = tf.convert_to_tensor(lsy_v)

    lenmfcc_v = tf.convert_to_tensor(lenmfcc_v)
    leny_v = tf.convert_to_tensor(leny_v)

    # print("Training")
    # print("X : ", lsx.shape)
    # print("Y : ", lsy.shape)
    # print("Len_X : ", lenmfcc.shape)
    # print("Len_Y : ", leny.shape)
    # print("Validation")
    # print("X_val : ", lsx_v.shape)
    # print("Y_val : ", lsy_v.shape)
    # print("Len_X_val : ", lenmfcc_v.shape)
    # print("Len_Y_val : ", leny_v.shape)

    batch_size = 8

    train_dataset = tf.data.Dataset.from_tensor_slices((lsx, lsy, lenmfcc, leny))
    train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

    validation_dataset = tf.data.Dataset.from_tensor_slices(
        (lsx_v, lsy_v, lenmfcc_v, leny_v)
    )
    validation_dataset = validation_dataset.shuffle(buffer_size=1024).batch(batch_size)

    # flush unnecessary variables
    del lsx
    del lsy
    del lenmfcc
    del leny
    del lsx_v
    del lsy_v
    del lenmfcc_v
    del leny_v
    del t_data

    return train_dataset, validation_dataset


"""

Training RetinaNet

"""
import os
import tqdm
import argparse
import numpy as np
import torch.nn as nn
import torch.optim as optim
from models import retina
from datasets import coco, voc, minibatch
from torch.utils.data import DataLoader
from lib.det_ops.loss import SigmoidFocalLoss, SmoothL1Loss
from IPython import embed
from torch.nn.utils import clip_grad
import tensorboardX
from utils import logger
from cfgs import config as cfg

def initialize(config, args):

    logdir = config['logdir']
    if not os.path.exists(logdir):
        os.mkdir(logdir)
    if not os.path.exists(os.path.join(logdir, args.experiment)):
        os.mkdir(os.path.join(logdir, args.experiment))

    model_dump_dir = os.path.join(logdir, args.experiment, 'model_dump')
    tb_dump = os.path.join(logdir, args.experiment, 'tb_dump')

    if not os.path.exists(model_dump_dir):
        os.mkdir(model_dump_dir)

    if not os.path.exists(tb_dump):
        os.mkdir(tb_dump)

    config['tb_dump_dir'] = tb_dump
    config['model_dump_dir'] = model_dump_dir

def learning_rate_decay(optimizer, step, config):
    base_lr = config['base_lr']
    lr = base_lr
    if step <= config['warmup']:
        lr = (lr - 1e-4)*step/config['warmup'] + 1e-4
    if step >= config['lr_decay'][0]:
        lr = base_lr * 0.1
    if step >= config['lr_decay'][1]:
        lr = base_lr * 0.01

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return lr

def train(args, config):
    anchor_scales = config['anchor_sizes']
    anchor_apsect_ratios = config['aspect_ratios']
    num_anchors = len(anchor_scales) * len(anchor_apsect_ratios)

    model = retina.RetinaNet(config['num_classes']-1, num_anchors, config['basemodel_path']).cuda()
    model = nn.DataParallel(model, device_ids=list(range(args.device)))

    if args.dataset == 'COCO':
        train_dataset = coco.COCODetection(dataroot=config['data_dir'], imageset='train2017', config=config)
    elif args.dataset == 'VOC':
        train_dataset = voc.VOC2012(dataroot=config['data_dir'], imageset='train', config=config)
    else:
        raise NotImplemented()

    collate_minibatch = minibatch.create_minibatch_func(config)

    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=args.batch_size*args.device,
        shuffle=True,
        num_workers=config['workers'],
        collate_fn=collate_minibatch
    )

    writer = tensorboardX.SummaryWriter(config['tb_dump_dir'])
    # torch model

    optimizer = optim.SGD(lr=config['base_lr'], params=model.parameters(),
                          weight_decay=config['weight_decay'], momentum=0.9)

    cls_criterion = SigmoidFocalLoss().cuda()
    box_criterion = SmoothL1Loss().cuda()

    start_epoch = 0
    global_step = 0

    # Load state dict from saved model
    if len(args.continue_path) > 0:
        model_state, optimizer_state, epoch, step = logger.load_checkpoints(args.continue_path)
        model.module.load_state_dict(model_state)
        optimizer.load_state_dict(optimizer_state)
        global_step = step+1
        start_epoch = epoch + 1
    for p in model.module.modules():
        if p.__class__.__name__ == 'BatchNorm2d':
            p.eval()
    for epoch in range(start_epoch, config['epochs']):
        losses = []
        data_iter = iter(train_loader)
        pbar = tqdm.tqdm(range(len(train_loader)))
        for i in pbar:
            lr = learning_rate_decay(optimizer, global_step, config)
            img, labels, boxes = next(data_iter)
            img = img.cuda()
            labels = labels.long().cuda()
            boxes = boxes.cuda()
            cls_outputs, bbox_outputs = model(img)
            cls_loss = cls_criterion(cls_outputs, labels)
            box_loss = box_criterion(bbox_outputs, boxes, labels)
            loss = cls_loss + box_loss

            optimizer.zero_grad()
            loss.backward()
            clip_grad.clip_grad_norm_(model.parameters(), 30)

            optimizer.step()
            writer.add_scalar('train/box_loss', box_loss.item(), global_step)
            writer.add_scalar('train/cls_loss', cls_loss.item(), global_step)
            global_step += 1
            pbar.set_description('e:{} i:{} loss:{:.3f} cls_loss:{:.3f} box_loss:{:.3f} lr:{}'.format(
                epoch, i + 1, loss.item(), cls_loss.item(), box_loss.item(), lr
            ))
            losses.append(loss.item())

            # learning rate decay

        print("e:{} loss: {}".format(epoch, np.mean(losses)))
        logger.save_checkpoints(model.module, optimizer, epoch, global_step,
                                path=os.path.join(config['model_dump_dir'],
                                                  'epoch-{}-iter-{}.pth'.format(epoch, global_step)))

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--device', type=int, default=1, help='training with ? GPUs')
    parser.add_argument('-b', '--batch_size', type=int, default=4, help='training batch size per GPU')
    parser.add_argument('-c', '--continue_path', type=str, default='', help='continue model parameters')
    parser.add_argument('-e', '--experiment', type=str, default='voc_baseline',
                        help='experiment name, correspond to `config.py`')
    parser.add_argument('-ds', '--dataset', type=str, default='VOC', help='dataset')

    _args = parser.parse_args()
    config = cfg.config[_args.experiment]
    initialize(config, _args)
    train(_args, config)


import numpy as np
############################## Amplicon Format ################################
#
# This pipeline expects DNA sequences containing a non-degenerate head,
# followed by a sequence identifying the sgRNA (enumerated in sgRNA_info.csv),
# followed by a random nucleotide barcode, followed by a non-degenerate tail.
#
###############################################################################
                                                            #  Lengths
head = 'GCGCACGTCTGCCGCGCTGTTCTCCTCTTCCTCATCTCCGGGACCCGGA'      #  49
sgID = '........'                                               #   8
random_barcode = 'AA.....TT.....AA.....'                        #  21
tail = 'ATGCCCAAGAAGAAGAGGAAGGTGTCCAATTTACTGACCGTACACCAAAATTTGCCTGCATTACCGGTCGATGCAACGAGTGATGAGGTTCGCAAGAACCT'[:22]
                                                                #  22
                                                                #  --
master_read = head + sgID + random_barcode + tail               # 100

######################## Preprocessing Parameters #############################

min_align_score_frac = 0.6      # lost N cassettes are at ~0.6 - 0.67
allowable_deviation = 4         # Flexibility in length of degenerate region
alignment_flank = 22
training_flank = 18             # Number of flanking nucleotides to use for error estimation
cluster_flank = allowable_deviation # Number of flanking nucleotides for clustering

preprocessed_dir = 'preprocessed/'
training_dir = 'training/'
original_dir = 'original/'
fastq_handle = '.fastq'

# Dictionary of reduction operators to use for aggregating DADA2 clusters that are deemed identical (n10 = n0 + n1):
merge_rules = dict(abundance=np.sum, n0=np.sum, n1=np.sum, nunq=np.sum, pval=np.prod, birth_pval=np.prod, birth_ham=np.min)

sgID_length = len(sgID)
random_barcode_length = len(random_barcode)
barcode_length = sgID_length + random_barcode_length

sample_meta_data_file = ""


# -*- coding: utf-8 -*-
"""
Created on Mon Apr 27 13:37:22 2015

@author: hoseung

Tools in this module should have only minimum dependence, so that
even the most lowlevel modules (like simBase) can import and use this module.

"""
import numpy as np
def get_pt_id_targets2(part, h, halo_inds, r_frac=1):
    """
    Any doc string? what is this for?
    """
    assert(np.size(halo_inds) > 0)
    buffer_frac = 1.2 * np.sqrt(r_frac)
    # search radii larger than r_vir will result in more particles.
    npart_tot_guess = np.math.ceil(sum(h.data.np[halo_inds]) * buffer_frac)
    idlist = np.zeros(npart_tot_guess, dtype=np.int32)
    x = part.x
    y = part.y
    z = part.z
    print("Expected # of particles:", npart_tot_guess)
    nhalo = len(halo_inds)
    halo_range_list = np.zeros((nhalo) + 1, dtype=np.int32)
    len_idlist = len(idlist)
    for i, ihalo in enumerate(halo_inds):
        print("# part in this halo", h.data.np[ihalo])
        xr = [h.data.x[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.x[ihalo] + h.data.rvir[ihalo] * r_frac]
        yr = [h.data.y[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.y[ihalo] + h.data.rvir[ihalo] * r_frac]
        zr = [h.data.z[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.z[ihalo] + h.data.rvir[ihalo] * r_frac]

        ind_z = np.where((x > xr[0]) & (x < xr[1]) &
                       (y > yr[0]) & (y < yr[1]) &
                       (z > zr[0]) & (z < zr[1]))[0]
        if len(ind_z) > 0:
            print("good")
            # If array is not large enough, append it.
            halo_range_list[i + 1] = halo_range_list[i] + len(ind_z)
            if halo_range_list[i + 1] > len_idlist:
                if i + 1 == nhalo:
                    # If it's the last halo, append by exact difference.
                    npart_more = halo_range_list[i + 1] - len_idlist
                else:
                    # Otherwise, guess from previous halos.
                    # 1.5 * mean npart so far * number of remaining halos
                    npart_more = int(1.5 * (halo_range_list[i] / (i + 1)) * (nhalo - i))
                print("increase the array size by {:d} from {:d}".format(npart_more, len_idlist))
                idlist = np.append(idlist, np.zeros(npart_more, dtype=np.int32))
                len_idlist= len(idlist)

            idlist[halo_range_list[i]:halo_range_list[i+1]] = part.id[ind_z]
            print(halo_range_list[i+1], len(idlist))

    return halo_range_list, idlist

def get_pt_id_targets(part, h, halo_inds, r_frac=1):
    """
    Returns id list of particles in each halo in a SINGLE array.
    Also returned is the list of range of id list for each halo.

    So,
    idlist1 = idslist[ind[0] : ind[1]]
    idlist2 = idslist[ind[1] : ind[2]]

    .. note::
        You don't know how many particles will be found within halos you are interested.
        A good guess would be the sum of halo.np which is number of member particles
        defined by the halo finder.

        So, I allocate the idlist array with length of 1.2 * sum(halo.np[ind]).
    """
    assert(np.size(halo_inds) > 0)
    buffer_frac = 1.3 * np.sqrt(r_frac)
    # search radii larger than r_vir will result in more particles.
    npart_tot_guess = np.math.ceil(sum(h.data.np[halo_inds]) * buffer_frac)
    idlist = np.zeros(npart_tot_guess, dtype=np.int32)
    x = part.x
    y = part.y
    z = part.z
    print("Expected # of particles:", npart_tot_guess)
    nhalo = len(halo_inds)
    halo_range_list = np.zeros((nhalo) + 1, dtype=np.int32)
    len_idlist = len(idlist)
    for i, ihalo in enumerate(halo_inds):
        print("# part in this halo", h.data.np[ihalo])
        xr = [h.data.x[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.x[ihalo] + h.data.rvir[ihalo] * r_frac]
        yr = [h.data.y[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.y[ihalo] + h.data.rvir[ihalo] * r_frac]
        zr = [h.data.z[ihalo] - h.data.rvir[ihalo] * r_frac,
              h.data.z[ihalo] + h.data.rvir[ihalo] * r_frac]
        ind_x = np.where( (x > xr[0]) & (x < xr[1]))[0]
        # out of 8.1GB particle information()
        if len(ind_x) > 0:
            ind_y = np.where((y[ind_x] > yr[0]) & (y[ind_x] < yr[1]))[0]
            if len(ind_y) > 0:
                ind_z = np.where((z[ind_x[ind_y]] > zr[0]) & (z[ind_x[ind_y]] < zr[1]))[0]
                # If array is not large enough, append it.
                halo_range_list[i + 1] = halo_range_list[i] + len(ind_z)
                if halo_range_list[i + 1] > len_idlist:
                    if i + 1 == nhalo:
                        # If it's the last halo, append by exact difference.
                        npart_more = halo_range_list[i + 1] - len_idlist
                    else:
                        # Otherwise, guess from previous halos.
                        # 1.5 * mean npart so far * number of remaining halos
                        npart_more = int(1.5 * (halo_range_list[i] / (i + 1)) * (nhalo - i))
                    print("increase the array size by {:d} from {:d}".format(npart_more, len_idlist))
                    idlist = np.append(idlist, np.zeros(npart_more, dtype=np.int32))
                    len_idlist= len(idlist)

                idlist[halo_range_list[i]:halo_range_list[i+1]] = part.id[ind_x[ind_y[ind_z]]]
                print(halo_range_list[i+1], len(idlist))

    return halo_range_list, idlist

def particle_distribution(part, idlist, buffer = 0):
    from utils import match
    """
    returns center of mass of given particles, xr,yr,zr,
    and a,b,c (in the future)
    """
    ind = match.match_list_ind(part.id, idlist)
    # To increase speed, search ind at one go and slice them.
    print(" Number of matching particles: {} out of {}".format(len(ind), len(idlist)))
    xr = [part["x"][ind].min(), part["x"][ind].max()]
    yr = [part["y"][ind].min(), part["y"][ind].max()]
    zr = [part["z"][ind].min(), part["z"][ind].max()]
    radius = max([xr[1] - xr[0], yr[1] - yr[0], zr[1] - zr[0]])
    xc = sum(xr) * 0.5
    yc = sum(yr) * 0.5
    zc = sum(zr) * 0.5

    return xc, yc, zc, radius

class Allrefine():
    def __init__(self, hid, data, aexp):
        self.id = hid
        self.x = data['x_refine']
        self.y = data['y_refine']
        self.z = data['z_refine']
        self.r = data['r_refine']
        self.aexp = aexp

def match_diff_arr(x1, x2, y1, y2, z1=None, z2=None, tolerance=None, window=0.1, reverse=False):
    """
    first array is smaller.
    """
    import numpy as np
    n1 = len(x1)
    n2 = len(x2)
    n21 = n2/n1
    if max([n1, n2]) < 200:
        window = 0.5
    # window size depends on the length of the longer array.
    window = int(n2 * window)

    if tolerance is None:
        tolerance = 0.5 * x1.ptp() / n1 # half of the mean separation in 0-th axis.

    print("window:", window)
    print("tolerance:", tolerance)
    print("array lengths:", n1, n2)

    if reverse:
        # Suppose n1 is smaller
        sort_ind1 = np.argsort(x1)[::-1]
        sort_ind2 = np.argsort(x2)[::-1]
    else:
        sort_ind1 = np.argsort(x1)
        sort_ind2 = np.argsort(x2)

    print("n1 = {}, n2 = {}".format(n1,n2))
    # initialize matched list
    match = np.zeros(n1, dtype=int)
    in21 = [np.math.ceil(i * n21) for i in np.arange(n1)]
    for i in range(n1):
        if i < window:
            ind_1 = sort_ind2[0:in21[i] + window]
        elif i > n1 - window:
            ind_1 = sort_ind2[in21[i] - window :]
        else:
            ind_1 = sort_ind2[in21[i] - window : in21[i] + window]
        xx2 = x2[ind_1]
        xx1 = x1[sort_ind1[i]]
        d1 = np.abs(xx2 - xx1)
        indx = np.where(d1 < tolerance)[0]
        if len(indx) > 0:
            if z1 is not None and z2 is not None:
                yx2 = y2[ind_1] # first N smallest elements
                yx1 = y1[sort_ind1[i]]
                zx2 = z2[ind_1] # first N smallest elements
                zx1 = z1[sort_ind1[i]]
                d2 = (xx2[indx] - xx1)**2 \
                    + (yx2[indx] - yx1)**2 + (zx2[indx] - zx1)**2
                if len(np.where(d2 < tolerance**2)[0]) > 0:
                    print(np.where(d2 < tolerance**2))
                    match[i] = sort_ind2[indx[np.argmin(d2)]] # index of minimum value
                else:
                    match[i]=-1
            elif y1 is not None and y2 is not None:
                d2 = (x2[sort_ind2[indx]] - x1[sort_ind1[i]])**2 \
                            + (y2[sort_ind2[indx]] - y1[sort_ind1[i]])**2
                if len(np.where(d2 < tolerance**2)[0]) > 0:
                    print(np.where(d2 < tolerance**2))
                    match[i] = sort_ind2[indx[np.argmin(d2)]] # index of minimum value
                else:
                    match[i]=-1
                    print("Missed!", i)
        else:
            match[i]=-1
            print("Missed!", i)

    # n-th element in the list match is the index of 2nd array corresponds to
    #        the n-th smallest element of the 1st array.
    # Or, arr1[sort_ind1] = arr2[match]
    # an index to revert sort_ind1 into the original order
    reverse_arr1 = np.argsort(sort_ind1)

    return match[reverse_arr1]

def distance3d(x1,y1,z1, x2,y2,z2):
    from numpy import sqrt
    return sqrt((x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)

def search_zoomin_region(Asim, *args, **kwargs):
    """
    Moved from sim

    Determine Zoomin region.

    Returns a spherical region encompassing maximally refined cells.

    Notes
    -----
    If part is not given, load particles.

    Not only part, but also hydro or amr can be used to find the zoomin region!
    But, don't want to write a new code.
    """
    if hasattr(Asim, 'part'):
        print("Have part")
        set_zregion(Asim.part.search_zoomin( *args, **kwargs))

    if hasattr(Asim, 'amr'):
        print("Have amr")
        set_zregion(Asim.amr.search_zoomin( *args, **kwargs))

def set_zregion(Asim, zregion):
    """
    Set zoom-in region for Zoom-in simulations.
    """
    Asim.zregion = zregion
    Asim.info.zregion = zregion

def get_zoomin(self, scale=1.0):
    """
    Returns a spherical region encompassing maximally refined cells.
    Moved from Amr class.
    What should it do??

    Parameters
    ----------
    scale : float
        The radius of the returned sphere is scaled by 'scale'.

    """
    from ..utils import sampling
    imin = np.where(self.dm['m'] == self.dm['m'].min())
    xr = [self.dm['px'][imin].min(), self.dm['px'][imin].max()]
    yr = [self.dm['py'][imin].min(), self.dm['py'][imin].max()]
    zr = [self.dm['pz'][imin].min(), self.dm['pz'][imin].max()]

    xc = 0.5 * sum(xr)
    yc = 0.5 * sum(yr)
    zc = 0.5 * sum(zr)

    radius = 0.5 * max([xr[1]-xr[0], yr[1]-yr[0], zr[1]-zr[0]]) * scale
    #print(radius)
    return sampling.set_region(centers=[xc, yc, zc], radius=radius)

def search_zoomin(self, scale=1.0, load=False):
    """
    Moved from The Class Part.
    """
    from ..utils import sampling
    if not hasattr(self, 'dm'):
        if load:
            print("Let's load particle data first!")
            self.load()
        else:
            print("Couldn't find dm data. \n \
                Make sure you load it first, \n \
                or use load=True option")

    imin = np.where(self.dm['m'] == self.dm['m'].min())
    xr = [self.dm['x'][imin].min(), self.dm['x'][imin].max()]
    yr = [self.dm['y'][imin].min(), self.dm['y'][imin].max()]
    zr = [self.dm['z'][imin].min(), self.dm['z'][imin].max()]

    xc = 0.5 * sum(xr)
    yc = 0.5 * sum(yr)
    zc = 0.5 * sum(zr)

    radius = 0.5 * max([xr[1]-xr[0], yr[1]-yr[0], zr[1]-zr[0]]) * scale

    return(sampling.set_region(centers=[xc, yc, zc], radius=radius))


import argparse
import gym
import d4rl
import numpy as np

from rlkit.envs import ENVS
import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.env_replay_buffer import EnvReplayBuffer
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.launchers.launcher_util import setup_logger
from rlkit.samplers.data_collector import MdpPathCollector
from rlkit.torch.sac.policies import TanhGaussianPolicy, MakeDeterministic
from rlkit.torch.sac.sac import SACTrainer
from rlkit.torch.networks import ConcatMlp
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm

def load_hdf5(dataset, replay_buffer, max_size):
    all_obs = dataset['observations']
    all_act = dataset['actions']
    N = min(all_obs.shape[0], max_size)

    _obs = all_obs[:N-1]
    _actions = all_act[:N-1]
    _next_obs = all_obs[1:]
    _rew = np.squeeze(dataset['rewards'][:N-1])
    _rew = np.expand_dims(np.squeeze(_rew), axis=-1)
    _done = np.squeeze(dataset['terminals'][:N-1])
    _done = (np.expand_dims(np.squeeze(_done), axis=-1)).astype(np.int32)

    max_length = 1000
    ctr = 0
    ## Only for MuJoCo environments
    ## Handle the condition when terminal is not True and trajectory ends due to a timeout
    for idx in range(_obs.shape[0]):
        if ctr >= max_length - 1:
            ctr = 0
        else:
            replay_buffer.add_sample_only(_obs[idx], _actions[idx], _rew[idx], _next_obs[idx], _done[idx])
            ctr += 1
            if _done[idx][0]:
                ctr = 0

def experiment(variant):
    eval_env = gym.make(variant['env_name'])
    expl_env = eval_env
    obs_dim = expl_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size

    M = variant['layer_size']
    qf1 = ConcatMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    qf2 = ConcatMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf1 = ConcatMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf2 = ConcatMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    policy = TanhGaussianPolicy(
        obs_dim=obs_dim,
        action_dim=action_dim,
        hidden_sizes=[M, M],
    )
    eval_policy = MakeDeterministic(policy)
    eval_path_collector = MdpPathCollector(
        eval_env,
        eval_policy,
    )
    expl_path_collector = MdpPathCollector(
        expl_env,
        policy,
    )
    replay_buffer = EnvReplayBuffer(
        variant['replay_buffer_size'],
        expl_env,
    )
    load_hdf5(eval_env.unwrapped.get_dataset(), replay_buffer, max_size=variant['replay_buffer_size'])

    trainer = SACTrainer(
        env=eval_env,
        policy=policy,
        qf1=qf1,
        qf2=qf2,
        target_qf1=target_qf1,
        target_qf2=target_qf2,
        **variant['trainer_kwargs']
    )
    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        batch_rl=True,  # this is for offline RL (False if online)
        **variant['algorithm_kwargs']
    )
    algorithm.to(ptu.device)
    algorithm.train()

if __name__ == "__main__":
    # noinspection PyTypeChecker
    parser = argparse.ArgumentParser(description='OfflineSAC-runs')
    parser.add_argument("--env", type=str, default='halfcheetah-random-v0')
    args = parser.parse_args()

    variant = dict(
        algorithm="SAC",
        version="normal",
        layer_size=256,
        replay_buffer_size=int(1E6),
        env_name=args.env,
        algorithm_kwargs=dict(
            num_epochs=1000,
            num_eval_steps_per_epoch=5000,
            num_trains_per_train_loop=1000,
            num_expl_steps_per_train_loop=1000,
            min_num_steps_before_training=1000,
            max_path_length=1000,
            batch_size=256,
        ),
        trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=1,
            policy_lr=3E-4,
            qf_lr=3E-4,
            reward_scale=1,
            use_automatic_entropy_tuning=True,
        ),
    )
    setup_logger(exp_prefix='OfflineSAC-' + args.env,
                 variant=variant,
                 text_log_file="debug.log",
                 variant_log_file="variant.json",
                 tabular_log_file="progress.csv",
                 snapshot_mode="gap_and_last",
                 snapshot_gap=200,
                 log_tabular_only=False,
                 log_dir=None,
                 git_infos=None,
                 script_name=None,
                 # **create_log_dir_kwargs
                 base_log_dir='./data',
                 exp_id=0,
                 seed=0)
    ptu.set_gpu_mode(True)  # optionally set the GPU (default=False)
    experiment(variant)


import sys
from dataset import Dataset
import tqdm
import time
from autoencodeur import AutoEncoder
import torch
import torch.nn as nn
import numpy as np
import logging
import argparse
import torchvision.models as models
from torchvision import datasets
from tensorboardX import SummaryWriter

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def getArgs():
    parser = argparse.ArgumentParser(description='Setting model')
    parser.add_argument('-m','--model', type=int, choices=[32,256], default=32, help='Choose the model')
    parser.add_argument('-l','--learning',type=str, choices=['images','landmarks','images_landmarks'], default='images', help='Select the learning type')
    args = parser.parse_args()
    return args.model,args.learning

def save_checkpoint(state, is_best, checkpoint_path,best_model_path):
    torch.save(state, checkpoint_path)
    if is_best:
        torch.save(state, best_model_path)

def startLearning(model_type:int ,mode:str):
    print(device)
    #cudnn.Benchmark = True
    resnet18 = models.resnet18(False)
    writer = SummaryWriter()
    batch_size = 16
    max_epochs = 100
    print("load dataset")
    training_loader = torch.utils.data.DataLoader(dataset=Dataset('training_dataset_pack.h5',model_type,"std_training.png","mean_training.png"), batch_size=batch_size, shuffle=True,num_workers=0)
    validation_loader = torch.utils.data.DataLoader(dataset=Dataset('validation_dataset_pack.h5',model_type,"std_validation.png","mean_validation.png"), batch_size=batch_size, shuffle=True,num_workers=0)

    print("Done")

    if model_type == 32:
        model = AutoEncoder(training_loader.dataset.getInputSize()).cuda()
    else:
        print("Model doesn't exist")
        return

    criterion_expressions = nn.CrossEntropyLoss().cuda()
    criterion_landmarks = nn.MSELoss().cuda()

    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9,0.999))
    #déclarer scheduler reduce plateau
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='min',factor=0.1,patience=10)
    best_loss = None
    ######

    ######
    for epoch in range(1,max_epochs+1):  # loop over the dataset multiple times

        is_best= False
        print("Training...")
        pbart = tqdm.tqdm(total=int(len(training_loader.dataset)/batch_size),postfix={"loss":None,"accuracy":None},desc="Epoch: {}/{}".format(epoch,max_epochs))
        loss = 0.
        loss_lm = 0.
        loss_expr = 0.
        acc = 0.
        val_loss = 0.
        val_acc =0.
        val_loss_lm = 0.
        val_loss_expr = 0.
        for i, data in enumerate(training_loader, 0):
            optimizer.zero_grad()
            # get the inputs
            images, landmarks, expressions = data
            images = images.to(device)
            landmarks = landmarks.to(device).float()
            expressions = expressions.to(device).long()
            # zero the parameter gradients

            outputs = model(images)
            loss_landmarks = criterion_landmarks(outputs[0], landmarks.float())
            loss_expressions = criterion_expressions(outputs[1], expressions)
            current_loss = loss_landmarks + loss_expressions
            current_loss.backward()
            optimizer.step()

            # print statistics
            loss_expr += loss_expressions.item()
            loss_lm += loss_landmarks.item()
            loss += current_loss.item()
            _,predicted_expressions = torch.max(outputs[1],1)

            acc += (predicted_expressions == expressions).sum().float()/batch_size
            pbart.update(1)
            pbart.set_postfix({"loss": loss/(i+1),"e_loss": loss_expr/(i+1),"l_loss": loss_lm/(i+1),"acc_e":acc.item()/(i+1)})
        pbart.close()
        loss /= (len(training_loader.dataset)/batch_size)
        acc /= (len(training_loader.dataset)/batch_size)

        with open("log/training"+"_"+str(model_type)+"_"+mode+".log","a") as f:
            f.write("epoch: {} / {} loss: {} accuracy: {}\n".format(epoch,max_epochs,loss,acc))
        f.close()

        print("Validation...")
        pbarv = tqdm.tqdm(total=int(len(validation_loader.dataset)/batch_size),postfix={"loss":None,"accuracy":None},desc="Epoch: {}/{}".format(epoch,max_epochs))

        with torch.no_grad():
            for i, data in enumerate(validation_loader, 0):
                # get the inputs
                images, landmarks, expressions = data
                images = images.to(device)
                landmarks = landmarks.to(device).float()
                expressions = expressions.to(device).long()

                outputs = model(images)
                loss_landmarks = criterion_landmarks(outputs[0], landmarks.float())
                loss_expressions = criterion_expressions(outputs[1], expressions)
                loss = loss_landmarks + loss_expressions
                val_loss += loss.item()
                val_loss_expr += loss_expressions.item()
                val_loss_lm += loss_landmarks.item()
                _,predicted_expressions = torch.max(outputs[1],1)
                val_acc += (predicted_expressions == expressions).sum().float()/batch_size

                pbarv.update(1)
                pbarv.set_postfix({"loss": val_loss/(i+1),"e_loss": val_loss_expr/(i+1),"l_loss": val_loss_lm/(i+1),"acc_e":val_acc.item()/(i+1)})

        pbarv.close()
        val_loss /= (len(validation_loader.dataset)/batch_size)
        val_acc /= (len(validation_loader.dataset)/batch_size)

        if best_loss == None or val_loss < best_loss:
            best_loss = val_loss
            is_best = True
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'loss': val_loss,
            'accuracy': val_acc,
            'optimizer' : optimizer.state_dict()}, is_best,"save_model/checkpoint"+"_"+str(model_type)+"_"+mode+".pth","save_model/best_model"+"_"+str(model_type)+"_"+mode+".pth")
        is_best = False

        scheduler.step(val_loss)
        #loss accuracy training et validation dans les logs

        with open("log/validation"+"_"+str(model_type)+"_"+mode+".log","a") as f:
            f.write("epoch: {} / {} loss: {} accuracy: {}\n".format(epoch,max_epochs,val_loss,val_acc))
        f.close()

        writer.add_scalar('data/Loss training', loss, epoch)
        writer.add_scalar('data/Loss validation', val_loss, epoch)
        writer.add_scalar('data/Accuracy training', acc, epoch)
        writer.add_scalar('data/Accuracy validation', val_acc, epoch)

        training_loader.dataset.shuffle()

if __name__ == "__main__":
    print(sys.version)
    args =getArgs()
    model_type = args[0]
    mode = args[1]
    with open("log/training"+"_"+str(model_type)+"_"+mode+".log","w") as f:
        f.close()
    with open("log/validation"+"_"+str(model_type)+"_"+mode+".log","w") as f:
        f.close()

    startLearning(model_type,mode)