import torch
import torch.nn as nn
import numpy as np
from lib.config import cfg
import lib.utils.kitti_utils as kitti_utils
import lib.utils.roipool3d.roipool3d_utils as roipool3d_utils
import lib.utils.iou3d.iou3d_utils as iou3d_utils

class ProposalTargetLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input_dict):
        roi_boxes3d, gt_boxes3d = input_dict['roi_boxes3d'], input_dict['gt_boxes3d']

        batch_rois, batch_gt_of_rois, batch_roi_iou = self.sample_rois_for_rcnn(roi_boxes3d, gt_boxes3d)

        rpn_xyz, rpn_features = input_dict['rpn_xyz'], input_dict['rpn_features']
        if cfg.RCNN.USE_INTENSITY:
            pts_extra_input_list = [input_dict['rpn_intensity'].unsqueeze(dim=2),
                                    input_dict['seg_mask'].unsqueeze(dim=2)]
        else:
            pts_extra_input_list = [input_dict['seg_mask'].unsqueeze(dim=2)]

        if cfg.RCNN.USE_DEPTH:
            pts_depth = input_dict['pts_depth'] / 70.0 - 0.5
            pts_extra_input_list.append(pts_depth.unsqueeze(dim=2))
        pts_extra_input = torch.cat(pts_extra_input_list, dim=2)

        # point cloud pooling
        pts_feature = torch.cat((pts_extra_input, rpn_features), dim=2)
        pooled_features, pooled_empty_flag = \
            roipool3d_utils.roipool3d_gpu(rpn_xyz, pts_feature, batch_rois, cfg.RCNN.POOL_EXTRA_WIDTH,
                                          sampled_pt_num=cfg.RCNN.NUM_POINTS)

        sampled_pts, sampled_features = pooled_features[:, :, :, 0:3], pooled_features[:, :, :, 3:]

        # data augmentation
        if cfg.AUG_DATA:
            # data augmentation
            sampled_pts, batch_rois, batch_gt_of_rois = \
                self.data_augmentation(sampled_pts, batch_rois, batch_gt_of_rois)

        # canonical transformation
        batch_size = batch_rois.shape[0]
        roi_ry = batch_rois[:, :, 6] % (2 * np.pi)
        roi_center = batch_rois[:, :, 0:3]
        sampled_pts = sampled_pts - roi_center.unsqueeze(dim=2)  # (B, M, 512, 3)
        batch_gt_of_rois[:, :, 0:3] = batch_gt_of_rois[:, :, 0:3] - roi_center
        batch_gt_of_rois[:, :, 6] = batch_gt_of_rois[:, :, 6] - roi_ry

        for k in range(batch_size):
            sampled_pts[k] = kitti_utils.rotate_pc_along_y_torch(sampled_pts[k], batch_rois[k, :, 6])
            batch_gt_of_rois[k] = kitti_utils.rotate_pc_along_y_torch(batch_gt_of_rois[k].unsqueeze(dim=1),
                                                                      roi_ry[k]).squeeze(dim=1)

        # regression valid mask
        valid_mask = (pooled_empty_flag == 0)
        reg_valid_mask = ((batch_roi_iou > cfg.RCNN.REG_FG_THRESH) & valid_mask).long()

        # classification label
        batch_cls_label = (batch_roi_iou > cfg.RCNN.CLS_FG_THRESH).long()
        invalid_mask = (batch_roi_iou > cfg.RCNN.CLS_BG_THRESH) & (batch_roi_iou < cfg.RCNN.CLS_FG_THRESH)
        batch_cls_label[valid_mask == 0] = -1
        batch_cls_label[invalid_mask > 0] = -1
        output_dict = {'sampled_pts': sampled_pts.view(-1, cfg.RCNN.NUM_POINTS, 3),
                       'pts_feature': sampled_features.view(-1, cfg.RCNN.NUM_POINTS, sampled_features.shape[3]),
                       'cls_label': batch_cls_label.view(-1),
                       'reg_valid_mask': reg_valid_mask.view(-1),
                       'gt_of_rois': batch_gt_of_rois.view(-1, 7),
                       'gt_iou': batch_roi_iou.view(-1),
                       'roi_boxes3d': batch_rois.view(-1, 7)}

        return output_dict

    def sample_rois_for_rcnn(self, roi_boxes3d, gt_boxes3d):
        """
        :param roi_boxes3d: (B, M, 7)
        :param gt_boxes3d: (B, N, 8) [x, y, z, h, w, l, ry, cls]
        :return
            batch_rois: (B, N, 7)
            batch_gt_of_rois: (B, N, 8)
            batch_roi_iou: (B, N)
        """
        batch_size = roi_boxes3d.size(0)

        fg_rois_per_image = int(np.round(cfg.RCNN.FG_RATIO * cfg.RCNN.ROI_PER_IMAGE))

        batch_rois = gt_boxes3d.new(batch_size, cfg.RCNN.ROI_PER_IMAGE, 7).zero_()
        batch_gt_of_rois = gt_boxes3d.new(batch_size, cfg.RCNN.ROI_PER_IMAGE, 7).zero_()
        batch_roi_iou = gt_boxes3d.new(batch_size, cfg.RCNN.ROI_PER_IMAGE).zero_()

        for idx in range(batch_size):
            cur_roi, cur_gt = roi_boxes3d[idx], gt_boxes3d[idx]

            k = cur_gt.__len__() - 1
            while cur_gt[k].sum() == 0:
                k -= 1
            cur_gt = cur_gt[:k + 1]

            # include gt boxes in the candidate rois
            iou3d = iou3d_utils.boxes_iou3d_gpu(cur_roi, cur_gt[:, 0:7])  # (M, N)

            max_overlaps, gt_assignment = torch.max(iou3d, dim=1)

            # sample fg, easy_bg, hard_bg
            fg_thresh = min(cfg.RCNN.REG_FG_THRESH, cfg.RCNN.CLS_FG_THRESH)
            fg_inds = torch.nonzero((max_overlaps >= fg_thresh)).view(-1)

            # TODO: this will mix the fg and bg when CLS_BG_THRESH_LO < iou < CLS_BG_THRESH
            # fg_inds = torch.cat((fg_inds, roi_assignment), dim=0)  # consider the roi which has max_iou with gt as fg

            easy_bg_inds = torch.nonzero((max_overlaps < cfg.RCNN.CLS_BG_THRESH_LO)).view(-1)
            hard_bg_inds = torch.nonzero((max_overlaps < cfg.RCNN.CLS_BG_THRESH) &
                                         (max_overlaps >= cfg.RCNN.CLS_BG_THRESH_LO)).view(-1)

            fg_num_rois = fg_inds.numel()
            bg_num_rois = hard_bg_inds.numel() + easy_bg_inds.numel()

            if fg_num_rois > 0 and bg_num_rois > 0:
                # sampling fg
                fg_rois_per_this_image = min(fg_rois_per_image, fg_num_rois)

                rand_num = torch.from_numpy(np.random.permutation(fg_num_rois)).type_as(gt_boxes3d).long()
                fg_inds = fg_inds[rand_num[:fg_rois_per_this_image]]

                # sampling bg
                bg_rois_per_this_image = cfg.RCNN.ROI_PER_IMAGE - fg_rois_per_this_image
                bg_inds = self.sample_bg_inds(hard_bg_inds, easy_bg_inds, bg_rois_per_this_image)

            elif fg_num_rois > 0 and bg_num_rois == 0:
                # sampling fg
                rand_num = np.floor(np.random.rand(cfg.RCNN.ROI_PER_IMAGE) * fg_num_rois)
                rand_num = torch.from_numpy(rand_num).type_as(gt_boxes3d).long()
                fg_inds = fg_inds[rand_num]
                fg_rois_per_this_image = cfg.RCNN.ROI_PER_IMAGE
                bg_rois_per_this_image = 0
            elif bg_num_rois > 0 and fg_num_rois == 0:
                # sampling bg
                bg_rois_per_this_image = cfg.RCNN.ROI_PER_IMAGE
                bg_inds = self.sample_bg_inds(hard_bg_inds, easy_bg_inds, bg_rois_per_this_image)

                fg_rois_per_this_image = 0
            else:
                raise NotImplementedError

            # augment the rois by noise
            roi_list, roi_iou_list, roi_gt_list = [], [], []
            if fg_rois_per_this_image > 0:
                fg_rois_src = cur_roi[fg_inds]
                gt_of_fg_rois = cur_gt[gt_assignment[fg_inds]]
                iou3d_src = max_overlaps[fg_inds]
                fg_rois, fg_iou3d = self.aug_roi_by_noise_torch(fg_rois_src, gt_of_fg_rois, iou3d_src,
                                                                aug_times=cfg.RCNN.ROI_FG_AUG_TIMES)
                roi_list.append(fg_rois)
                roi_iou_list.append(fg_iou3d)
                roi_gt_list.append(gt_of_fg_rois)

            if bg_rois_per_this_image > 0:
                bg_rois_src = cur_roi[bg_inds]
                gt_of_bg_rois = cur_gt[gt_assignment[bg_inds]]
                iou3d_src = max_overlaps[bg_inds]
                aug_times = 1 if cfg.RCNN.ROI_FG_AUG_TIMES > 0 else 0
                bg_rois, bg_iou3d = self.aug_roi_by_noise_torch(bg_rois_src, gt_of_bg_rois, iou3d_src,
                                                                aug_times=aug_times)
                roi_list.append(bg_rois)
                roi_iou_list.append(bg_iou3d)
                roi_gt_list.append(gt_of_bg_rois)

            rois = torch.cat(roi_list, dim=0)
            iou_of_rois = torch.cat(roi_iou_list, dim=0)
            gt_of_rois = torch.cat(roi_gt_list, dim=0)

            batch_rois[idx] = rois
            batch_gt_of_rois[idx] = gt_of_rois
            batch_roi_iou[idx] = iou_of_rois

        return batch_rois, batch_gt_of_rois, batch_roi_iou

    def sample_bg_inds(self, hard_bg_inds, easy_bg_inds, bg_rois_per_this_image):
        if hard_bg_inds.numel() > 0 and easy_bg_inds.numel() > 0:
            hard_bg_rois_num = int(bg_rois_per_this_image * cfg.RCNN.HARD_BG_RATIO)
            easy_bg_rois_num = bg_rois_per_this_image - hard_bg_rois_num

            # sampling hard bg
            rand_idx = torch.randint(low=0, high=hard_bg_inds.numel(), size=(hard_bg_rois_num,)).long()
            hard_bg_inds = hard_bg_inds[rand_idx]

            # sampling easy bg
            rand_idx = torch.randint(low=0, high=easy_bg_inds.numel(), size=(easy_bg_rois_num,)).long()
            easy_bg_inds = easy_bg_inds[rand_idx]

            bg_inds = torch.cat([hard_bg_inds, easy_bg_inds], dim=0)
        elif hard_bg_inds.numel() > 0 and easy_bg_inds.numel() == 0:
            hard_bg_rois_num = bg_rois_per_this_image
            # sampling hard bg
            rand_idx = torch.randint(low=0, high=hard_bg_inds.numel(), size=(hard_bg_rois_num,)).long()
            bg_inds = hard_bg_inds[rand_idx]
        elif hard_bg_inds.numel() == 0 and easy_bg_inds.numel() > 0:
            easy_bg_rois_num = bg_rois_per_this_image
            # sampling easy bg
            rand_idx = torch.randint(low=0, high=easy_bg_inds.numel(), size=(easy_bg_rois_num,)).long()
            bg_inds = easy_bg_inds[rand_idx]
        else:
            raise NotImplementedError

        return bg_inds

    def aug_roi_by_noise_torch(self, roi_boxes3d, gt_boxes3d, iou3d_src, aug_times=10):
        iou_of_rois = torch.zeros(roi_boxes3d.shape[0]).type_as(gt_boxes3d)
        pos_thresh = min(cfg.RCNN.REG_FG_THRESH, cfg.RCNN.CLS_FG_THRESH)

        for k in range(roi_boxes3d.shape[0]):
            temp_iou = cnt = 0
            roi_box3d = roi_boxes3d[k]

            gt_box3d = gt_boxes3d[k].view(1, 7)
            aug_box3d = roi_box3d
            keep = True
            while temp_iou < pos_thresh and cnt < aug_times:
                if np.random.rand() < 0.2:
                    aug_box3d = roi_box3d  # p=0.2 to keep the original roi box
                    keep = True
                else:
                    aug_box3d = self.random_aug_box3d(roi_box3d)
                    keep = False
                aug_box3d = aug_box3d.view((1, 7))
                iou3d = iou3d_utils.boxes_iou3d_gpu(aug_box3d, gt_box3d)
                temp_iou = iou3d[0][0]
                cnt += 1
            roi_boxes3d[k] = aug_box3d.view(-1)
            if cnt == 0 or keep:
                iou_of_rois[k] = iou3d_src[k]
            else:
                iou_of_rois[k] = temp_iou
        return roi_boxes3d, iou_of_rois

    @staticmethod
    def random_aug_box3d(box3d):
        """
        :param box3d: (7) [x, y, z, h, w, l, ry]
        random shift, scale, orientation
        """
        if cfg.RCNN.REG_AUG_METHOD == 'single':
            pos_shift = (torch.rand(3, device=box3d.device) - 0.5)  # [-0.5 ~ 0.5]
            hwl_scale = (torch.rand(3, device=box3d.device) - 0.5) / (0.5 / 0.15) + 1.0  #
            angle_rot = (torch.rand(1, device=box3d.device) - 0.5) / (0.5 / (np.pi / 12))  # [-pi/12 ~ pi/12]
            aug_box3d = torch.cat([box3d[0:3] + pos_shift, box3d[3:6] * hwl_scale, box3d[6:7] + angle_rot], dim=0)
            return aug_box3d
        elif cfg.RCNN.REG_AUG_METHOD == 'multiple':
            # pos_range, hwl_range, angle_range, mean_iou
            range_config = [[0.2, 0.1, np.pi / 12, 0.7],
                            [0.3, 0.15, np.pi / 12, 0.6],
                            [0.5, 0.15, np.pi / 9, 0.5],
                            [0.8, 0.15, np.pi / 6, 0.3],
                            [1.0, 0.15, np.pi / 3, 0.2]]
            idx = torch.randint(low=0, high=len(range_config), size=(1,))[0].long()

            pos_shift = ((torch.rand(3, device=box3d.device) - 0.5) / 0.5) * range_config[idx][0]
            hwl_scale = ((torch.rand(3, device=box3d.device) - 0.5) / 0.5) * range_config[idx][1] + 1.0
            angle_rot = ((torch.rand(1, device=box3d.device) - 0.5) / 0.5) * range_config[idx][2]

            aug_box3d = torch.cat([box3d[0:3] + pos_shift, box3d[3:6] * hwl_scale, box3d[6:7] + angle_rot], dim=0)
            return aug_box3d
        elif cfg.RCNN.REG_AUG_METHOD == 'normal':
            x_shift = np.random.normal(loc=0, scale=0.3)
            y_shift = np.random.normal(loc=0, scale=0.2)
            z_shift = np.random.normal(loc=0, scale=0.3)
            h_shift = np.random.normal(loc=0, scale=0.25)
            w_shift = np.random.normal(loc=0, scale=0.15)
            l_shift = np.random.normal(loc=0, scale=0.5)
            ry_shift = ((torch.rand() - 0.5) / 0.5) * np.pi / 12

            aug_box3d = np.array([box3d[0] + x_shift, box3d[1] + y_shift, box3d[2] + z_shift, box3d[3] + h_shift,
                                  box3d[4] + w_shift, box3d[5] + l_shift, box3d[6] + ry_shift], dtype=np.float32)
            aug_box3d = torch.from_numpy(aug_box3d).type_as(box3d)
            return aug_box3d
        else:
            raise NotImplementedError

    def data_augmentation(self, pts, rois, gt_of_rois):
        """
        :param pts: (B, M, 512, 3)
        :param rois: (B, M. 7)
        :param gt_of_rois: (B, M, 7)
        :return:
        """
        batch_size, boxes_num = pts.shape[0], pts.shape[1]

        # rotation augmentation
        angles = (torch.rand((batch_size, boxes_num), device=pts.device) - 0.5 / 0.5) * (np.pi / cfg.AUG_ROT_RANGE)

        # calculate gt alpha from gt_of_rois
        temp_x, temp_z, temp_ry = gt_of_rois[:, :, 0], gt_of_rois[:, :, 2], gt_of_rois[:, :, 6]
        temp_beta = torch.atan2(temp_z, temp_x)
        gt_alpha = -torch.sign(temp_beta) * np.pi / 2 + temp_beta + temp_ry  # (B, M)

        temp_x, temp_z, temp_ry = rois[:, :, 0], rois[:, :, 2], rois[:, :, 6]
        temp_beta = torch.atan2(temp_z, temp_x)
        roi_alpha = -torch.sign(temp_beta) * np.pi / 2 + temp_beta + temp_ry  # (B, M)

        for k in range(batch_size):
            pts[k] = kitti_utils.rotate_pc_along_y_torch(pts[k], angles[k])
            gt_of_rois[k] = kitti_utils.rotate_pc_along_y_torch(gt_of_rois[k].unsqueeze(dim=1), angles[k]).squeeze(dim=1)
            rois[k] = kitti_utils.rotate_pc_along_y_torch(rois[k].unsqueeze(dim=1), angles[k]).squeeze(dim=1)

            # calculate the ry after rotation
            temp_x, temp_z = gt_of_rois[:, :, 0], gt_of_rois[:, :, 2]
            temp_beta = torch.atan2(temp_z, temp_x)
            gt_of_rois[:, :, 6] = torch.sign(temp_beta) * np.pi / 2 + gt_alpha - temp_beta

            temp_x, temp_z = rois[:, :, 0], rois[:, :, 2]
            temp_beta = torch.atan2(temp_z, temp_x)
            rois[:, :, 6] = torch.sign(temp_beta) * np.pi / 2 + roi_alpha - temp_beta

        # scaling augmentation
        scales = 1 + ((torch.rand((batch_size, boxes_num), device=pts.device) - 0.5) / 0.5) * 0.05
        pts = pts * scales.unsqueeze(dim=2).unsqueeze(dim=3)
        gt_of_rois[:, :, 0:6] = gt_of_rois[:, :, 0:6] * scales.unsqueeze(dim=2)
        rois[:, :, 0:6] = rois[:, :, 0:6] * scales.unsqueeze(dim=2)

        # flip augmentation
        flip_flag = torch.sign(torch.rand((batch_size, boxes_num), device=pts.device) - 0.5)
        pts[:, :, :, 0] = pts[:, :, :, 0] * flip_flag.unsqueeze(dim=2)
        gt_of_rois[:, :, 0] = gt_of_rois[:, :, 0] * flip_flag
        # flip orientation: ry > 0: pi - ry, ry < 0: -pi - ry
        src_ry = gt_of_rois[:, :, 6]
        ry = (flip_flag == 1).float() * src_ry + (flip_flag == -1).float() * (torch.sign(src_ry) * np.pi - src_ry)
        gt_of_rois[:, :, 6] = ry

        rois[:, :, 0] = rois[:, :, 0] * flip_flag
        # flip orientation: ry > 0: pi - ry, ry < 0: -pi - ry
        src_ry = rois[:, :, 6]
        ry = (flip_flag == 1).float() * src_ry + (flip_flag == -1).float() * (torch.sign(src_ry) * np.pi - src_ry)
        rois[:, :, 6] = ry

        return pts, rois, gt_of_rois


#!/usr/bin/env python
#--coding:utf-8 --
"""
pipe.py
2017-03-22: finished basically
2017-06-26: runsecond model refined
2017-06-27: modified runsecond model, as extend the second time result
2017-06-30: modified the combining method
2017-07-20: modified some output.
2017-07-28: modified overlaps, change to the samll one.
2017-08-03: re-deisgn the underlying data structure, improve the speed as 2 fold and reduce the memory as 1/2
2017-08-07: modified pipeline, as accecpt multiple eps.
2017-08-10: to do, re-design data structure as using HDF5 as botoom structure to reduce memory usage
2017-08-18: changed combine method for loops when runing a series of eps.
2018-03-13: modified combining method for multiple eps
2018-03-18: modified pre-processing bedpe files.
2018-03-27: modified merging clustering by remove close PETs
2018-03-30: multiple minPts mode added
2018-03-31: modified self-ligation and inter-ligation distance cutoff selection.
2018-09-26: updated estimated distance cutoff to avoid 0 size self-ligated PETs
2018-10-09: updated usage for self-ligation and inter-ligation distance cutoff.
"""
__author__ = "CAO Yaqiang"
__date__ = ""
__modified__ = ""
__email__ = "caoyaqiang0410@gmail.com"

#python library
import os, time, sys, shutil, gzip, copy
from glob import glob
from datetime import datetime

#3rd library
import numpy as np
import pandas as pd
from joblib import Parallel, delayed

#cLoops
from cLoops.settings import *
from cLoops.utils import getLogger, callSys, cFlush, mainHelp
from cLoops.io import parseRawBedpe, parseRawBedpe2, txt2jd, parseJd, loops2washU, loops2juice
from cLoops.cDBSCAN2 import cDBSCAN as DBSCAN
#from cLoops.blockDBSCAN import blockDBSCAN as DBSCAN
from cLoops.ests import estFragSize, estIntSelCutFrag
from cLoops.cPlots import plotFragSize, plotIntSelCutFrag
from cLoops.cModel import getIntSig, markIntSig, markIntSigHic

#global settings
global logger

def singleDBSCAN(f, eps, minPts, cut=0):
    """
    Run DBSCAN to detect interactions for one chromosome.
    #mat is list, every is [ pointId,x,y ]
    """
    dataI, readI, dataS, readS, dis, dss = [], [], [], [], [], []
    key, mat = parseJd(f, cut=0)
    if cut > 0:
        d = mat[:, 2] - mat[:, 1]
        p = np.where(d >= cut)[0]
        mat = mat[p, :]
        dss.extend(list(d[d < cut]))
    if len(mat) == 0:
        return key, f, dataI, dataS, list(dis), list(dss)
    #data for interaction records, read for readId
    report = "Clustering %s and %s using eps as %s, minPts as %s,pre-set distance cutoff as > %s\n" % (
        key[0], key[1], eps, minPts, cut)
    sys.stderr.write(report)
    db = DBSCAN(mat, eps, minPts)
    labels = pd.Series(db.labels)
    mat = np.array(mat)
    mat = pd.DataFrame(mat[:, 1:].astype("float"),
                       index=mat[:, 0],
                       columns=["X", "Y"])
    nlabels = set(labels.values)
    #collect clusters
    for label in nlabels:
        los = list(labels[labels == label].index)
        sub = mat.loc[los, :]
        #BEDPE format,+1 to escape the error that exact the same start and end
        #2017-05-18, changed to remove such interactions
        if int(np.min(sub["X"])) == int(np.max(sub["X"])) or int(
                np.min(sub["Y"])) == int(np.max(sub["Y"])):
            continue
        r = [
            key[0],
            int(np.min(sub["X"])),
            int(np.max(sub["X"])),
            key[1],
            int(np.min(sub["Y"])),
            int(np.max(sub["Y"])),
            #sub.shape[0],
            #",".join(map(str, los)),
            #los
        ]
        if r[2] < r[4]:
            dataI.append(r)
            readI.extend(los)
        else:
            dataS.append(r)
            readS.extend(los)
    report = "Clustering %s and %s finished. Estimated %s self-ligation reads and %s inter-ligation reads\n" % (
        key[0], key[1], len(readS), len(readI))
    sys.stderr.write(report)
    if len(dataI) > 0:
        dis = mat.loc[readI, "Y"] - mat.loc[readI, "X"]
    if len(dataS) > 0:
        dss.extend(list(mat.loc[readS, "Y"] - mat.loc[readS, "X"]))
    return key, f, dataI, dataS, list(dis), list(dss)

def runDBSCAN(fs, eps, minPts, cut=0, cpu=1):
    """
    Run DBSCAN to detect interactions for all chromosomes.
    """
    ds = Parallel(n_jobs=cpu)(delayed(singleDBSCAN)(f, eps, minPts, cut)
                              for f in fs)
    dataI, dataS, dis, dss = {}, [], [], []
    for d in ds:
        if len(d[2]) == 0:
            continue
        dataI[d[0]] = {"f": d[1], "records": d[2]}
        dataS.extend(d[3])
        dis.extend(d[4])
        dss.extend(d[5])
    return dataI, dataS, dis, dss

def filterClusterByDis(data, cut):
    """
    Filter inter-ligation clusters by distances
    """
    for key in data:
        nr = []
        #print len(data[key]["records"])
        for r in data[key]["records"]:
            d = (r[4] + r[5]) / 2 - (r[1] + r[2]) / 2
            if d >= cut:
                nr.append(r)
        data[key]["records"] = nr
        #print len(data[key]["records"])
    return data

def checkSameLoop(ra, rb):
    """
    check if two anchors are exact same
    """
    if ra[1] == rb[1] and ra[2] == rb[2] and ra[4] == rb[4] and ra[5] == rb[5]:
        return True
    return False

def combineTwice(dataI, dataI_2):
    """
    Combining multiple clustering result.
    """
    for key in dataI_2.keys():
        if key not in dataI:
            dataI[key] = {
                "f": dataI_2[key]["f"],
                "records": dataI_2[key]["records"]
            }
        else:
            ds = set()
            for r in dataI[key]["records"]:
                t = [r[1], r[2], r[4], r[5]]
                ds.add(tuple(t))
            for r in dataI_2[key]["records"]:
                t = tuple([r[1], r[2], r[4], r[5]])
                if t not in ds:
                    dataI[key]["records"].append(r)
    return dataI

def runStat(dataI, minPts, cut, cpu, fout, hichip=0):
    """
    Calling p-values of interactions for all chromosomes.
    """
    logger.info(
        "Starting estimate significance for interactions using distance cutoff as %s"
        % cut)
    ds = Parallel(n_jobs=cpu)(
        delayed(getIntSig)(dataI[key]["f"], dataI[key]["records"], minPts, cut)
        for key in dataI.keys())
    ds = [d for d in ds if d is not None]
    if len(ds) == 0:
        logger.error("Something wrong, no loops found, sorry, bye.")
        return 1
    ds = pd.concat(ds)
    try:
        if hichip:
            ds = markIntSigHic(ds)
        else:
            ds = markIntSig(ds)
        ds.to_csv(fout + ".loop", sep="\t", index_label="loopId")
    except:
        logger.warning(
            "Something wrong happend to significance estimation, only output called loops"
        )
        ds.to_csv(fout + "_raw.loop", sep="\t", index_label="loopId")
    return 0

def pipe(fs,
         fout,
         eps,
         minPts,
         chroms="",
         cpu=1,
         tmp=0,
         hic=0,
         washU=0,
         juice=0,
         cut=0,
         plot=0,
         max_cut=False):
    if chroms == "":
        chroms = []
    else:
        chroms = set(chroms.split(","))
    #1.pre-processing data as merge replicates, seperate into cis and trans PETs
    #cis and trans PETs files
    if os.path.isdir(fout):
        r = "working directory %s exists, return." % fout
        logger.error(r)
        return
    os.mkdir(fout)
    #2.read in PETs, collect cis PETs and convert to .jd format, if not assigned eps (eps=0),estimate eps
    if eps == 0:
        cfs, ds = parseRawBedpe(fs, fout, chroms, cut, logger)
    else:
        cfs = parseRawBedpe2(fs, fout, chroms, cut, logger)
    cfs = Parallel(n_jobs=cpu)(delayed(txt2jd)(f) for f in cfs)
    #3. run the DBSCAN
    if eps == 0:
        frags = estFragSize(ds)
        eps = [frags * 2]
    #3.1. run DBSCAN for multiple times
    dataI = {}
    dis, dss = [], []
    cuts = [
        cut,
    ]
    #multiple eps
    for ep in eps:
        #multiple minPts, for minPts like 50,40,30,20, the bigger minPts always get smaller distance cutoff
        for m in minPts:
            dataI_2, dataS_2, dis_2, dss_2 = runDBSCAN(cfs, ep, m, cut, cpu)
            if len(dataI_2) == 0:
                logger.info(
                    "ERROR: no inter-ligation PETs detected for eps %s minPts %s,can't model the distance cutoff,continue anyway"
                    % (ep, m))
                continue
            if len(dis_2) == 0 or len(dss_2) == 0:
                dataI = combineTwice(dataI, dataI_2)
            else:
                cut_2, frags = estIntSelCutFrag(np.array(dis_2),
                                                np.array(dss_2))
                if plot:
                    plotIntSelCutFrag(dis_2,
                                      dss_2,
                                      cut_2,
                                      frags,
                                      prefix=fout +
                                      "_eps%s_minPts%s_disCutoff" % (ep, m))
                logger.info(
                    "Estimated inter-ligation and self-ligation distance cutoff as %s for eps=%s,minPts=%s"
                    % (cut_2, ep, m))
                #experimental
                cuts.append(cut_2)
                #cut = max(cuts)
                cut = cut_2
                dataI = combineTwice(dataI, dataI_2)
    cuts = [c for c in cuts if c > 0]
    if max_cut:
        cut = np.max(cuts)
    else:
        cut = np.min(cuts)
    dataI = filterClusterByDis(dataI, cut)
    #5.estimate the significance
    #e = runStat(dataI, minPts, cut, cpu, fout, hic)
    e = runStat(dataI, minPts, 0, cpu, fout, hic)
    if e:
        shutil.rmtree(fout)
        return
    #6.convert loops and cis-PETs into washU visulization format or juicebox
    if washU:
        loops2washU(fout + ".loop", fout + "_loops_washU.txt", logger)
    if juice:
        loops2juice(fout + ".loop", fout + "_loops_juicebox.txt", logger)
    #7.remove temple files
    if tmp == False:
        shutil.rmtree(fout)

def main():
    start = datetime.now()
    global logger
    fn = os.path.join(os.getcwd(), "cLoops.log")
    logger = getLogger(fn)
    op = mainHelp()
    report = "Command line: cLoops -f {} -o {} -m {} -eps {} -minPts {} -p {} -w {} -j {} -s {} -c {} -hic {} -cut {} -plot {} -max_cut {}".format(
        op.fnIn, op.fnOut, op.mode, op.eps, op.minPts, op.cpu, op.washU,
        op.juice, op.tmp, op.chroms, op.hic, op.cut, op.plot, op.max_cut)
    logger.info(report)
    if op.mode == 0:
        #parse eps
        if "," in str(op.eps):
            eps = map(int, op.eps.split(","))
            eps.sort(reverse=False)
        else:
            eps = int(op.eps)
            if eps != 0:
                eps = [eps]
        #parse minPts
        if "," in str(op.minPts):
            minPts = map(int, op.minPts.split(","))
            minPts.sort(reverse=True)
        else:
            minPts = int(op.minPts)
            if minPts != 0:
                minPts = [minPts]
            else:
                logger.error("minPts not assigned!")
                return
        hic = op.hic
    if op.mode == 1:
        eps = [500, 1000, 2000]
        minPts = [5]
        hic = 0
    if op.mode == 2:
        eps = [1000, 2000, 5000]
        minPts = [5]
        hic = 0
    if op.mode == 3:
        eps = [5000, 7500, 10000]
        minPts = [50, 40, 30, 20]
        hic = 1
    if op.mode == 4:
        eps = [2500, 5000, 7500, 10000]
        minPts = [30, 20]
        hic = 1
    report = "mode:%s\t eps:%s\t minPts:%s\t hic:%s\t" % (op.mode, eps, minPts,
                                                          hic)
    logger.info(report)
    pipe(op.fnIn.split(","), op.fnOut, eps, minPts, op.chroms, op.cpu, op.tmp,
         hic, op.washU, op.juice, op.cut, op.plot, op.max_cut)
    usedtime = datetime.now() - start
    r = "cLoops finished. Used CPU time: %s Bye!\n\n\n" % usedtime
    logger.info(r)

if __name__ == "__main__":
    main()


from thinc.extra.wrappers import PyTorchWrapper, xp2torch, torch2xp
from transformers.optimization import AdamW
import transformers
import torch.autograd
import torch.nn.utils.clip_grad
import torch
from typing import Tuple, Callable, Any
from thinc.neural.optimizers import Optimizer
import numpy
import contextlib
from thinc.compat import BytesIO

from .util import get_model, Dropout
from .activations import RaggedArray, Activations

FINE_TUNE = True
CONFIG = {"output_hidden_states": True, "output_attentions": True}

class TransformersWrapper(PyTorchWrapper):
    """Wrap a Transformers model for use in Thinc.

    The model will take as input a spacy_transformers.util.RaggedArray
    object that will specify the input IDs and optionally the segment IDs. The
    RaggedArray is basically a tuple (ids, lengths), where ids is concatenated
    for a whole batch (this format allows the data to be contiguous even if
    the sequences are different lengths). The segment IDs should be coded as
    the different models expect them -- see
    https://github.com/huggingface/transformers/blob/master/examples/utils_glue.py
    """

    _model: Any
    _optimizer: Any
    cfg: dict

    @classmethod
    def from_pretrained(cls, name):
        model_cls = get_model(name)
        model = model_cls.from_pretrained(name, **CONFIG)
        self = cls(name, model.config.to_dict(), model)
        self.cfg.update(self.transformers_model.config.to_dict())
        return self

    def __init__(self, name, config, model):
        PyTorchWrapper.__init__(self, model)
        self.cfg = dict(config)

    @property
    def nO(self):
        if "hidden_size" in self.cfg:
            # BERT
            return self.cfg["hidden_size"]
        elif "hidden_dim" in self.cfg:
            # DistilBERT
            return self.cfg["hidden_dim"] // 4
        elif "n_embd" in self.cfg:
            # GPT2
            return self.cfg["n_embd"]
        elif "d_model" in self.cfg:
            # XLNet
            return self.cfg["d_model"]
        elif hasattr(self.transformers_model, "dim"):
            # XLM
            return self.transformers_model.dim
        else:
            keys = ", ".join(self.cfg.keys())
            raise ValueError(f"Unexpected config. Keys: {keys}")

    @property
    def transformers_model(self):
        return self._model

    @property
    def max_length(self):
        return self.cfg.get("max_position_embeddings", 128)

    def predict(self, inputs: RaggedArray):
        self._model.eval()
        model_kwargs = self.get_model_kwargs(inputs)
        with torch.no_grad():
            if hasattr(self._optimizer, "swap_swa_sgd"):
                self._optimizer.swap_swa_sgd()
            y_var = self._model(**model_kwargs)
            if hasattr(self._optimizer, "swap_swa_sgd"):
                self._optimizer.swap_swa_sgd()
        return self.make_activations(y_var, inputs.lengths)

    def begin_update(
        self, inputs: RaggedArray, drop: Dropout = 0.0
    ) -> Tuple[Activations, Callable[..., None]]:
        if drop is None:
            # "drop is None" indicates prediction. It's one of the parts of
            # Thinc's API I'm least happy with...
            return self.predict(inputs), lambda dY, sgd=None: None
        max_original = max(inputs.lengths, default=0)
        model_kwargs = self.get_model_kwargs(inputs)
        self._model.train()
        # Prepare all the model arguments, including the attention mask
        y_var = self._model(**model_kwargs)
        output = self.make_activations(y_var, inputs.lengths)
        assert output.lh.data.shape[0] == inputs.data.shape[0], (
            output.lh.data.shape,
            inputs.data.shape,
        )

        def backward_pytorch(d_output: Activations, sgd: Optimizer = None) -> None:
            y_for_bwd = []
            dy_for_bwd = []
            if d_output.has_lh:
                assert d_output.lh.data.shape[0] == sum(d_output.lh.lengths)
                d_lh = d_output.lh.to_padded(to=max_original)
                if self.max_length and d_lh.shape[1] >= self.max_length:
                    d_lh = d_lh[:, : self.max_length]
                dy_for_bwd.append(xp2torch(d_lh))
                y_for_bwd.append(y_var[0])
            if d_output.has_po:
                dy_for_bwd.append(xp2torch(d_output.po.data))
                y_for_bwd.append(y_var[1])
            if FINE_TUNE:
                torch.autograd.backward(y_for_bwd, grad_tensors=dy_for_bwd)
                if sgd is not None:
                    if self._optimizer is None:
                        self._optimizer = self._create_optimizer(sgd)
                    if sgd.max_grad_norm:
                        torch.nn.utils.clip_grad.clip_grad_norm_(
                            self._model.parameters(), sgd.max_grad_norm
                        )
                    optimizer = self._optimizer
                    for group in optimizer.param_groups:
                        group["lr"] = getattr(sgd, "trf_lr", sgd.alpha)
                    optimizer.step()
                    optimizer.zero_grad()
                    self._update_pytorch_averages(sgd)
            return None

        self._model.eval()
        return output, backward_pytorch

    @contextlib.contextmanager
    def use_params(self, params):
        key_prefix = f"pytorch_{self.id}_"
        state_dict = {}
        for k, v in params.items():
            if hasattr(k, "startswith") and k.startswith(key_prefix):
                state_dict[k.replace(key_prefix, "")] = xp2torch(v)
        if state_dict:
            backup = {k: v.clone() for k, v in self._model.state_dict().items()}
            self._model.load_state_dict(state_dict)
            yield
            self._model.load_state_dict(backup)
        else:
            yield

    def make_activations(self, fields, lengths) -> Activations:
        """Create Activations from the output tuples produced by PyTorch Transformers.
        Includes converting torch tensors to xp, and handling missing values.
        """
        fields = list(fields)
        fields[0] = torch2xp(fields[0])
        fields[0] = RaggedArray.from_padded(fields[0], lengths)
        assert fields[0].data.shape[0] == sum(lengths)
        # lh: last hidden
        # po: pooler_output
        # ah: all_hidden
        # aa: all_attention
        if len(fields) != 4:
            lh = fields[0]
            po = RaggedArray.blank()
        else:
            if isinstance(fields[1], tuple):
                fields[1] = RaggedArray.blank()
            else:
                fields[1] = RaggedArray(torch2xp(fields[1]), [1] * len(lengths))
            lh, po, _, _2 = fields
        # Convert last_hidden_state to xp
        return Activations(lh, po)

    def get_model_kwargs(self, inputs):
        padded = inputs.to_padded()
        if padded.ndim == 2:
            padded = padded.reshape(padded.shape + (1,))
        if self.max_length:
            padded = padded[:, : self.max_length]
        ids = padded[:, :, 0]
        neg_idx = ids < 0
        ids[neg_idx] = 0
        ids = torch.as_tensor(ids, dtype=torch.int64)
        if padded.shape[2] == 2:
            segment_ids = padded[:, :, 1]
            segment_ids = torch.as_tensor(segment_ids, dtype=torch.int64)
        else:
            segment_ids = torch.zeros_like(ids)
        # Calculate "attention mask" for BERT and  XLNet, but not GPT2 (sigh)
        if isinstance(self._model, (transformers.BertModel, transformers.XLNetModel)):
            mask = self.ops.xp.ones(ids.shape, dtype=numpy.int_)
            mask[neg_idx] = 0
            mask = xp2torch(mask)
            return {
                "input_ids": ids,
                "attention_mask": mask,
                "token_type_ids": segment_ids,
            }
        elif isinstance(self._model, (transformers.DistilBertModel)):
            # Mask, but no token type IDs for DistilBert (sigh again...)
            mask = self.ops.xp.ones(ids.shape, dtype=numpy.int_)
            mask[neg_idx] = 0
            mask = xp2torch(mask)
            return {"input_ids": ids, "attention_mask": mask}
        else:
            return {"input_ids": ids, "token_type_ids": segment_ids}

    def _create_optimizer(self, sgd):
        optimizer = AdamW(
            self._model.parameters(),
            lr=getattr(sgd, "trf_lr", sgd.alpha),
            eps=sgd.eps,
            betas=(sgd.b1, sgd.b2),
            weight_decay=getattr(sgd, "trf_weight_decay", 0.0),
        )
        optimizer.zero_grad()
        return optimizer

    def _update_pytorch_averages(self, sgd, *, init_steps=1):
        if sgd.averages is None:
            return
        # Collect parameters if we don't have them
        for name, param in self._model.state_dict().items():
            key = f"pytorch_{self.id}_{name}"
            sgd.nr_update[key] += 1
            xp_param = torch2xp(param)
            if key in sgd.averages:
                self.ops.update_averages(
                    sgd.averages[key], xp_param, sgd.nr_update[key]
                )
            else:
                sgd.averages[key] = xp_param.copy()
                sgd.nr_update[key] = init_steps

    def to_disk(self, path):
        torch.save(self._model.state_dict(), str(path))

    def from_disk(self, path):
        if self.ops.device == "cpu":
            map_location = "cpu"
        else:
            map_location = "cuda:0"
        self._model.load_state_dict(torch.load(path, map_location=map_location))
        self._model.to(map_location)

    def to_bytes(self):
        filelike = BytesIO()
        torch.save(self._model.state_dict(), filelike)
        filelike.seek(0)
        return filelike.getvalue()

    def from_bytes(self, data):
        filelike = BytesIO(data)
        filelike.seek(0)
        if self.ops.device == "cpu":
            map_location = "cpu"
        else:
            map_location = "cuda:0"
        self._model.load_state_dict(torch.load(filelike, map_location=map_location))
        self._model.to(map_location)


"""YOLO Estimator."""
# pylint: disable=logging-format-interpolation,abstract-method
import os
import math
import time
import warnings
import pandas as pd
import numpy as np
import mxnet as mx
from mxnet import gluon
from mxnet import autograd
from mxnet.contrib import amp

from .... import utils as gutils
from ....data.batchify import Tuple, Stack, Pad
from ....data.transforms.presets.yolo import YOLO3DefaultValTransform
from ....data.transforms.presets.yolo import load_test, transform_test
from ....model_zoo import get_model
from ....model_zoo import custom_yolov3
from ....utils.metrics.voc_detection import VOC07MApMetric, VOCMApMetric
from ....utils import LRScheduler, LRSequential

from ..base_estimator import BaseEstimator, set_default
from .utils import _get_dataloader
from ...data.dataset import ObjectDetectionDataset
from ..conf import _BEST_CHECKPOINT_FILE

try:
    import horovod.mxnet as hvd
except ImportError:
    hvd = None

from .default import YOLOv3Cfg

__all__ = ['YOLOv3Estimator']

@set_default(YOLOv3Cfg())
class YOLOv3Estimator(BaseEstimator):
    """Estimator implementation for YOLOv3.

    Parameters
    ----------
    config : dict
        Config in nested dict.
    logger : logging.Logger, default is None
        Optional logger for this estimator, can be `None` when default setting is used.
    reporter : callable, default is None
        If set, use reporter callback to report the metrics of the current estimator.

    Attributes
    ----------
    _logger : logging.Logger
        The customized/default logger for this estimator.
    _logdir : str
        The temporary dir for logs.
    _cfg : autocfg.dataclass
        The configurations.
    """
    Dataset = ObjectDetectionDataset
    def __init__(self, config, logger=None, reporter=None):
        super(YOLOv3Estimator, self).__init__(config, logger, reporter)
        self.last_train = None
        if self._cfg.yolo3.amp:
            amp.init()
        if self._cfg.horovod:
            if hvd is None:
                raise SystemExit("Horovod not found, please check if you installed it correctly.")
            hvd.init()

    def _fit(self, train_data, val_data, time_limit=math.inf):
        """Fit YOLO3 model."""
        tic = time.time()
        self._best_map = 0
        self.epoch = 0
        self._time_elapsed = 0
        if max(self._cfg.train.start_epoch, self.epoch) >= self._cfg.train.epochs:
            return {'time', self._time_elapsed}
        if not isinstance(train_data, pd.DataFrame):
            self.last_train = len(train_data)
        else:
            self.last_train = train_data
        self.net.collect_params().reset_ctx(self.ctx)
        self._init_trainer()
        self._time_elapsed += time.time() - tic
        return self._resume_fit(train_data, val_data, time_limit=time_limit)

    def _resume_fit(self, train_data, val_data, time_limit=math.inf):
        tic = time.time()
        if max(self._cfg.train.start_epoch, self.epoch) >= self._cfg.train.epochs:
            return {'time', self._time_elapsed}
        if not self.classes or not self.num_class:
            raise ValueError('Unable to determine classes of dataset')

        # training dataset
        train_dataset = train_data.to_mxnet()
        val_dataset = val_data.to_mxnet()
        # training dataloader
        self.batch_size = self._cfg.train.batch_size // hvd.size() if self._cfg.horovod else self._cfg.train.batch_size
        train_loader, val_loader, train_eval_loader = _get_dataloader(
            self.async_net, train_dataset, val_dataset, self._cfg.yolo3.data_shape,
            self.batch_size, self._cfg.num_workers, self._cfg)

        if self._cfg.train.no_wd:
            for _, v in self.net.collect_params('.*beta|.*gamma|.*bias').items():
                v.wd_mult = 0.0
        if self._cfg.train.label_smooth:
            self.net._target_generator._label_smooth = True
        self._time_elapsed += time.time() - tic
        return self._train_loop(train_loader, val_loader, train_eval_loader, time_limit=time_limit)

    def _train_loop(self, train_data, val_data, train_eval_data, time_limit=math.inf):
        start_tic = time.time()
        # fix seed for mxnet, numpy and python builtin random generator.
        gutils.random.seed(self._cfg.train.seed)

        # metrics
        obj_metrics = mx.metric.Loss('ObjLoss')
        center_metrics = mx.metric.Loss('BoxCenterLoss')
        scale_metrics = mx.metric.Loss('BoxScaleLoss')
        cls_metrics = mx.metric.Loss('ClassLoss')
        trainer = self.trainer
        self._logger.info('Start training from [Epoch %d]', max(self._cfg.train.start_epoch, self.epoch))
        mean_ap = [-1]
        cp_name = ''
        self._time_elapsed += time.time() - start_tic
        for self.epoch in range(max(self._cfg.train.start_epoch, self.epoch), self._cfg.train.epochs):
            epoch = self.epoch
            if self._best_map >= 1.0:
                self._logger.info('[Epoch {}] Early stopping as mAP is reaching 1.0'.format(epoch))
                break
            tic = time.time()
            last_tic = time.time()
            if self._cfg.train.mixup:
                # TODO(zhreshold): more elegant way to control mixup during runtime
                try:
                    train_data._dataset.set_mixup(np.random.beta, 1.5, 1.5)
                except AttributeError:
                    train_data._dataset._data.set_mixup(np.random.beta, 1.5, 1.5)
                if epoch >= self._cfg.train.epochs - self._cfg.train.no_mixup_epochs:
                    try:
                        train_data._dataset.set_mixup(None)
                    except AttributeError:
                        train_data._dataset._data.set_mixup(None)

            mx.nd.waitall()
            self.net.hybridize()
            for i, batch in enumerate(train_data):
                btic = time.time()
                if self._time_elapsed > time_limit:
                    self._logger.warning(f'`time_limit={time_limit}` reached, exit early...')
                    return {'train_map': float(mean_ap[-1]), 'valid_map': self._best_map,
                            'time': self._time_elapsed, 'checkpoint': cp_name}
                data = gluon.utils.split_and_load(batch[0], ctx_list=self.ctx, batch_axis=0, even_split=False)
                # objectness, center_targets, scale_targets, weights, class_targets
                fixed_targets = [gluon.utils.split_and_load(batch[it], ctx_list=self.ctx,
                                                            batch_axis=0, even_split=False) for it in range(1, 6)]
                gt_boxes = gluon.utils.split_and_load(batch[6], ctx_list=self.ctx, batch_axis=0, even_split=False)
                sum_losses = []
                obj_losses = []
                center_losses = []
                scale_losses = []
                cls_losses = []
                with autograd.record():
                    for ix, x in enumerate(data):
                        obj_loss, center_loss, scale_loss, cls_loss = self.net(x, gt_boxes[ix],
                                                                               *[ft[ix] for ft in fixed_targets])
                        sum_losses.append(obj_loss + center_loss + scale_loss + cls_loss)
                        obj_losses.append(obj_loss)
                        center_losses.append(center_loss)
                        scale_losses.append(scale_loss)
                        cls_losses.append(cls_loss)
                    if self._cfg.yolo3.amp:
                        with amp.scale_loss(sum_losses, trainer) as scaled_loss:
                            autograd.backward(scaled_loss)
                    else:
                        autograd.backward(sum_losses)
                trainer.step(self.batch_size)
                if (not self._cfg.horovod or hvd.rank() == 0):
                    obj_metrics.update(0, obj_losses)
                    center_metrics.update(0, center_losses)
                    scale_metrics.update(0, scale_losses)
                    cls_metrics.update(0, cls_losses)
                    if self._cfg.train.log_interval and not (i + 1) % self._cfg.train.log_interval:
                        name1, loss1 = obj_metrics.get()
                        name2, loss2 = center_metrics.get()
                        name3, loss3 = scale_metrics.get()
                        name4, loss4 = cls_metrics.get()
                        self._logger.info(
                            '[Epoch {}][Batch {}], LR: {:.2E}, Speed: {:.3f} samples/sec,'
                            ' {}={:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}'.format(
                                epoch, i, trainer.learning_rate, self._cfg.train.batch_size / (time.time() - last_tic),
                                name1, loss1, name2, loss2, name3, loss3, name4, loss4))
                        last_tic = time.time()
                    self._time_elapsed += time.time() - btic

            post_tic = time.time()
            if (not self._cfg.horovod or hvd.rank() == 0):
                name1, loss1 = obj_metrics.get()
                name2, loss2 = center_metrics.get()
                name3, loss3 = scale_metrics.get()
                name4, loss4 = cls_metrics.get()
                self._logger.info('[Epoch {}] Training cost: {:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}'.format(
                    epoch, (time.time() - tic), name1, loss1, name2, loss2, name3, loss3, name4, loss4))
                if not (epoch + 1) % self._cfg.valid.val_interval:
                    # consider reduce the frequency of validation to save time
                    map_name, mean_ap = self._evaluate(val_data)
                    val_msg = '\n'.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)])
                    self._logger.info('[Epoch {}] Validation: \n{}'.format(epoch, val_msg))
                    current_map = float(mean_ap[-1])
                    if current_map > self._best_map:
                        cp_name = os.path.join(self._logdir, _BEST_CHECKPOINT_FILE)
                        self._logger.info('[Epoch %d] Current best map: %f vs previous %f, saved to %s',
                                          self.epoch, current_map, self._best_map, cp_name)
                        self.save(cp_name)
                        self._best_map = current_map
                if self._reporter:
                    self._reporter(epoch=epoch, map_reward=current_map)
            self._time_elapsed += time.time() - post_tic

        # map on train data
        tic = time.time()
        map_name, mean_ap = self._evaluate(train_eval_data)
        self._time_elapsed += time.time() - tic
        return {'train_map': float(mean_ap[-1]), 'valid_map': self._best_map,
                'time': self._time_elapsed, 'checkpoint': cp_name}

    def _evaluate(self, val_data):
        """Evaluate the current model on dataset."""
        if not isinstance(val_data, gluon.data.DataLoader):
            if hasattr(val_data, 'to_mxnet'):
                val_data = val_data.to_mxnet()
            val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
            val_data = gluon.data.DataLoader(
                val_data.transform(YOLO3DefaultValTransform(self._cfg.yolo3.data_shape, self._cfg.yolo3.data_shape)),
                self._cfg.valid.batch_size, False, batchify_fn=val_batchify_fn, last_batch='keep',
                num_workers=self._cfg.num_workers)
        if self._cfg.valid.metric == 'voc07':
            eval_metric = VOC07MApMetric(iou_thresh=self._cfg.valid.iou_thresh, class_names=self.classes)
        elif self._cfg.valid.metric == 'voc':
            eval_metric = VOCMApMetric(iou_thresh=self._cfg.valid.iou_thresh, class_names=self.classes)
        else:
            raise ValueError(f'Invalid metric type: {self._cfg.valid.metric}')
        self.net.collect_params().reset_ctx(self.ctx)
        # set nms threshold and topk constraint
        self.net.set_nms(nms_thresh=0.45, nms_topk=400)
        mx.nd.waitall()
        self.net.hybridize()
        for batch in val_data:
            val_ctx = self.ctx
            if batch[0].shape[0] < len(val_ctx):
                val_ctx = val_ctx[:batch[0].shape[0]]
            data = gluon.utils.split_and_load(batch[0], ctx_list=val_ctx, batch_axis=0, even_split=False)
            label = gluon.utils.split_and_load(batch[1], ctx_list=val_ctx, batch_axis=0, even_split=False)
            det_bboxes = []
            det_ids = []
            det_scores = []
            gt_bboxes = []
            gt_ids = []
            gt_difficults = []
            for x, y in zip(data, label):
                # get prediction results
                ids, scores, bboxes = self.net(x)
                det_ids.append(ids)
                det_scores.append(scores)
                # clip to image size
                det_bboxes.append(bboxes.clip(0, batch[0].shape[2]))
                # split ground truths
                gt_ids.append(y.slice_axis(axis=-1, begin=4, end=5))
                gt_bboxes.append(y.slice_axis(axis=-1, begin=0, end=4))
                gt_difficults.append(y.slice_axis(axis=-1, begin=5, end=6) if y.shape[-1] > 5 else None)

            # update metric
            eval_metric.update(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults)
        return eval_metric.get()

    def _predict(self, x):
        """Predict an individual example."""
        short_size = int(self._cfg.yolo3.data_shape)
        if isinstance(x, str):
            x = load_test(x, short=short_size, max_size=1024)[0]
        elif isinstance(x, np.ndarray):
            return self._predict(mx.nd.array(x))
        elif isinstance(x, mx.nd.NDArray):
            if len(x.shape) != 3 or x.shape[-1] != 3:
                raise ValueError('array input with shape (h, w, 3) is required for predict')
            x = transform_test(x, short=short_size, max_size=1024)[0]
        elif isinstance(x, pd.DataFrame):
            assert 'image' in x.columns, "Expect column `image` for input images"
            def _predict_merge(x):
                y = self._predict(x)
                y['image'] = x
                return y
            return pd.concat([_predict_merge(xx) for xx in x['image']]).reset_index(drop=True)
        elif isinstance(x, (list, tuple)):
            return pd.concat([self._predict(xx) for xx in x]).reset_index(drop=True)
        else:
            raise ValueError('Input is not supported: {}'.format(type(x)))
        height, width = x.shape[2:4]
        x = x.as_in_context(self.ctx[0])
        ids, scores, bboxes = [xx[0].asnumpy() for xx in self.net(x)]
        bboxes[:, (0, 2)] /= width
        bboxes[:, (1, 3)] /= height
        bboxes = np.clip(bboxes, 0.0, 1.0).tolist()
        df = pd.DataFrame({'predict_class': [self.classes[int(id)] for id in ids], 'predict_score': scores.flatten(),
                           'predict_rois': [{'xmin': bbox[0], 'ymin': bbox[1], 'xmax': bbox[2], 'ymax': bbox[3]} \
                                for bbox in bboxes]})
        # filter out invalid (scores < 0) rows
        valid_df = df[df['predict_score'] > 0].reset_index(drop=True)
        return valid_df

    def _init_network(self):
        if not self.num_class:
            raise ValueError('Unable to create network when `num_class` is unknown. \
                It should be inferred from dataset or resumed from saved states.')
        assert len(self.classes) == self.num_class

        # training contexts
        if self._cfg.horovod:
            self.ctx = [mx.gpu(hvd.local_rank())]
        else:
            ctx = [mx.gpu(int(i)) for i in self._cfg.gpus]
            self.ctx = ctx if ctx else [mx.cpu()]

        if self._cfg.yolo3.transfer is None:
            # use sync bn if specified
            if self._cfg.yolo3.syncbn and len(self.ctx) > 1:
                self.net = custom_yolov3(base_network_name=self._cfg.yolo3.base_network,
                                         filters=self._cfg.yolo3.filters,
                                         anchors=self._cfg.yolo3.anchors,
                                         strides=self._cfg.yolo3.strides,
                                         classes=self.classes,
                                         dataset=self._cfg.dataset,
                                         pretrained_base=True,
                                         norm_layer=gluon.contrib.nn.SyncBatchNorm,
                                         norm_kwargs={'num_devices': len(self.ctx)})
                self.async_net = custom_yolov3(base_network_name=self._cfg.yolo3.base_network,
                                               filters=self._cfg.yolo3.filters,
                                               anchors=self._cfg.yolo3.anchors,
                                               strides=self._cfg.yolo3.strides,
                                               classes=self.classes,
                                               dataset=self._cfg.dataset,
                                               pretrained_base=False)
            else:
                self.net = custom_yolov3(base_network_name=self._cfg.yolo3.base_network,
                                         filters=self._cfg.yolo3.filters,
                                         anchors=self._cfg.yolo3.anchors,
                                         strides=self._cfg.yolo3.strides,
                                         classes=self.classes,
                                         dataset=self._cfg.dataset,
                                         pretrained_base=True)
                self.async_net = self.net
        else:
            assert isinstance(self._cfg.yolo3.transfer, str)
            self._logger.info(
                f'Using transfer learning from {self._cfg.yolo3.transfer}, the other network parameters are ignored.')
            # use sync bn if specified
            if self._cfg.yolo3.syncbn and len(self.ctx) > 1:
                self.net = get_model(self._cfg.yolo3.transfer, pretrained=True,
                                     norm_layer=gluon.contrib.nn.SyncBatchNorm,
                                     norm_kwargs={'num_devices': len(self.ctx)})
                self.async_net = get_model(self._cfg.yolo3.transfer, pretrained=True)  # used by cpu worker
            else:
                self.net = get_model(self._cfg.yolo3.transfer, pretrained=True)
                self.async_net = self.net
            self.net.reset_class(self.classes,
                                 reuse_weights=[cname for cname in self.classes if cname in self.net.classes])

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            self.net.initialize()
            self.async_net.initialize()
        self.net.collect_params().reset_ctx(self.ctx)

    def _init_trainer(self):
        if self.last_train is None:
            raise RuntimeError('Cannot init trainer without knowing the size of training data')
        if isinstance(self.last_train, pd.DataFrame):
            train_size = len(self.last_train)
        elif isinstance(self.last_train, int):
            train_size = self.last_train
        else:
            raise ValueError("Unknown type of self.last_train: {}".format(type(self.last_train)))

        if self._cfg.train.lr_decay_period > 0:
            lr_decay_epoch = list(range(self._cfg.train.lr_decay_period,
                                        self._cfg.train.epochs,
                                        self._cfg.train.lr_decay_period))
        else:
            lr_decay_epoch = [int(i) for i in self._cfg.train.lr_decay_epoch]
        lr_decay_epoch = [e - self._cfg.train.warmup_epochs for e in lr_decay_epoch]
        num_batches = train_size // self._cfg.train.batch_size
        lr_scheduler = LRSequential([
            LRScheduler('linear', base_lr=0, target_lr=self._cfg.train.lr,
                        nepochs=self._cfg.train.warmup_epochs, iters_per_epoch=num_batches),
            LRScheduler(self._cfg.train.lr_mode, base_lr=self._cfg.train.lr,
                        nepochs=self._cfg.train.epochs - self._cfg.train.warmup_epochs,
                        iters_per_epoch=num_batches,
                        step_epoch=lr_decay_epoch,
                        step_factor=self._cfg.train.lr_decay, power=2),
        ])

        if self._cfg.horovod:
            hvd.broadcast_parameters(self.net.collect_params(), root_rank=0)
            self.trainer = hvd.DistributedTrainer(
                self.net.collect_params(), 'sgd',
                {'wd': self._cfg.train.wd, 'momentum': self._cfg.train.momentum, 'lr_scheduler': lr_scheduler})
        else:
            self.trainer = gluon.Trainer(
                self.net.collect_params(), 'sgd',
                {'wd': self._cfg.train.wd, 'momentum': self._cfg.train.momentum, 'lr_scheduler': lr_scheduler},
                kvstore='local', update_on_kvstore=(False if self._cfg.yolo3.amp else None))

        if self._cfg.yolo3.amp:
            amp.init_trainer(self.trainer)


import argparse
import datetime
import gym
import envs

import numpy as np
import torch
import imageio

from torch.utils.data import ConcatDataset
from torch.utils.data import DataLoader

from transformer_split.util import getGraphStructure
from transformer_split.replay_memory import ReplayMemoryDataset
from transformer_split.vae_model import VAE_Model

from torch.nn import functional as F
from transformer_vae import util

parser = argparse.ArgumentParser(description='PyTorch Soft Actor-Critic Args')
parser.add_argument('--env1-name', default="ant",
                    help='Mujoco Gym environment (default: HalfCheetah-v2)')
parser.add_argument('--env2-name', default="ant3",
                    help='Mujoco Gym environment (default: HalfCheetah-v2)')
parser.add_argument('--model_path', default="runs/2021-05-19_13-46-41_VAE_ant-v0_both/",
                    help='model path')
parser.add_argument('--seed', type=int, default=123456, metavar='N',
                    help='random seed (default: 123456)')
parser.add_argument('--n_nodes', type=int, default=17,
                    help='max number of nodes in decoder')
parser.add_argument('--hidden_dim', type=int, default=512, metavar='N',
                    help='hidden size (default: 256)')
parser.add_argument('--latent_dim', type=int, default=64,
                    help='Encoder latent dimension')
parser.add_argument('--cuda', action="store_true",
                    help='run on CUDA (default: False)')
parser.add_argument('--agent_memory1', default='data/ant_jump.memory',
                    help='Path for saved replay memory')
parser.add_argument('--video_file_name', default="ant_transfered.mp4",
                    help='output file name')
parser.add_argument('--msg_dim', type=int, default=32,
                        help='run on CUDA (default: False)')
parser.add_argument('--batch_size', type=int, default=1,
                        help='run on CUDA (default: False)')
parser.add_argument('--root_size', type=int, default=11,
                        help='root dimension')
parser.add_argument('--lr', type=float, default=1e-3, metavar='N',
                        help='random seed (default: 123456)')
parser.add_argument(
    "--transformer_norm", default=0, type=int, help="Use layernorm",
)
parser.add_argument(
    "--attention_layers",
    default=3,
    type=int,
    help="How many attention layers to stack",
)
parser.add_argument(
    "--attention_heads",
    default=2,
    type=int,
    help="How many attention heads to stack",
)
parser.add_argument(
    "--attention_hidden_size",
    type=int,
    default=128,
    help="Hidden units in an attention block",
)

parser.add_argument(
    "--attention_embedding_size",
    type=int,
    default=128,
    help="Hidden units in an attention block",
)

parser.add_argument(
    "--dropout_rate",
    type=float,
    default=0.0,
    help="How much to drop if drop in transformers",
)

parser.add_argument(
        "--beta",
        type=float,
        default=1e-2,
        help="beta coefficient of KL divergence",
    )
parser.add_argument(
        "--discriminator_limiting_accuracy",
        type=float,
        default=0.6,
        help="beta coefficient of KL divergence",
    )
parser.add_argument(
        "--gradient_penalty",
        type=float,
        default=10,
        help="beta coefficient of KL divergence",
    )
args = parser.parse_args()

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

# Agent
device = torch.device("cuda" if args.cuda else "cpu")

env_names = ["ant4-walk-v0", "ant3-walk-v0", "ant5-walk-v0", "ant6-walk-v0"]
train_envs = [gym.make(n) for n in env_names]
graphs = [getGraphStructure(e.xml) for e in train_envs]
# All environments have the same dimension per limb.
num_limbs = len(graphs[0])  #torso + body limbs
body_limbs = num_limbs - 1
dim_per_limb = int((train_envs[0].observation_space.shape[0] - args.root_size) / (body_limbs - 1))
max_num_limbs = max(len(g) for g in graphs)

args.dim_per_limb = dim_per_limb
args.max_num_limbs = max_num_limbs

root_dir = util.get_project_root()
render_env = train_envs[0]
render_graph = graphs[0]

memory_file = root_dir / f"data/ant3-walk.memory"
env = gym.make('ant4-walk-v0')
dataset = ReplayMemoryDataset(
            memory_file,
            getGraphStructure(env.xml),
            args)

vae_model = VAE_Model(args)

vae_model.load_model(args.model_path)

# Evaluation loop
total_numsteps = 0
avg_reward = 0.

state = render_env.reset()

dataloader = DataLoader(dataset, batch_size=1, num_workers=0)
with imageio.get_writer(args.video_file_name, fps=30) as video:
    for idx, data, in enumerate(dataloader):

        x, src_top = data

        tgt_top = torch.full((max_num_limbs,), -1, dtype=torch.int32)
        tgt_top[:len(render_graph)] = torch.tensor(render_graph, dtype=torch.int32)
        tgt_top = tgt_top.unsqueeze(0)

        x_hat = vae_model.transfer(x, tgt_top)
        error = (x_hat - x).pow(2)
        render_env.set_to_observation(x_hat.detach().numpy().squeeze(0))
        video.append_data(render_env.render('rgb_array'))

        if idx > 1000:
            break