"""
3D IoU Calculation and Rotated NMS
Written by Shaoshuai Shi
All Rights Reserved 2019-2020.
"""
import torch

from ...utils import common_utils
from . import iou3d_nms_cuda

import numpy as np
from shapely.geometry import Polygon
from pcdet.utils.box_utils import boxes_to_corners_3d

def soft_nms_torch(dets, box_scores, iou_thresh=0.1, sigma=0.5, thresh=0.001):
    """
    # 0.001
    # 应该要使用bev_iou 还是 iou3d ？
    reference https://github.com/DocF/Soft-NMS.git
    Build a pytorch implement of Soft NMS algorithm.
    # Augments
        dets:        3D boxes coordinate tensor
        box_scores:  box score tensors
        sigma:       variance of Gaussian function
        iou_thresh   iou_thresh if use method2 else 0
        thresh:      score thresh
    # Return
        the sorted index of the selected boxes
    """
    # print('==> input dets.shape: ', dets.shape)
    # print('==> input box_scores: ', box_scores)

    N = dets.shape[0]  # the number of boxes

    # Indexes concatenate boxes with the last column
    indexes = torch.arange(0, N, dtype=torch.float).cuda().view(N, 1)
    dets = torch.cat((dets, indexes), dim=1)
    # print('==> cated dets.shape: ', dets.shape)
    # Sort the scores of the boxes from largest to smallest
    box_scores, conf_sort_index = torch.sort(box_scores, descending=True)
    # 这里顺序都改变了？外面的顺序是否还能保持一致呢？
    # print('==> sorted box_scores: ', box_scores)
    dets = dets[conf_sort_index]

    for i in range(N):
    # for i in range(N-1):

        pos=i+1

        #iou calculate
        # ious = box_iou(dets[i][0:4].view(-1,4), dets[pos:,:4])
        # TODO: 选择3Diou还是beviou？
        ious = boxes_iou_bev(dets[i][0:7].view(-1, 7), dets[pos:, :7])
        # ious = boxes_iou3d_gpu(dets[i][0:7].view(-1, 7), dets[pos:, :7])

        # method1
        # Gaussian decay
        box_scores[pos:] = torch.exp(-(ious * ious) / sigma) * box_scores[pos:]

        # method2
        # zero = torch.zeros_like(ious)
        # ious  = torch.where(ious < iou_thresh, zero , ious)
        # box_scores[pos:] = torch.exp(-(ious * ious) / sigma) * box_scores[pos:]

        # method3 original nms
        # weight = torch.ones(ious.shape)
        # weight[ious > iou_thresh] = 0
        # box_scores[pos:] = weight * box_scores[pos:]

        # box_scores[pos:] = box_scores[pos:]
        # box_scores[pos:], arg_sort = torch.sort(box_scores[pos:], descending=True)
        # print('==> box_scores[pos:].shape: ', box_scores[pos:].shape)
        box_scores_sorted, arg_sort = torch.sort(box_scores[pos:], descending=True)
        box_scores[pos:] = box_scores_sorted
        # print('==> in loop, sorted box_scores: ', box_scores)

        a=dets[pos:]

        dets[pos:] = a[arg_sort]

    # select the boxes and keep the corresponding indexes
    # 这里是因为之前把index cat到了box的最后一维度上去，所以这里应该是选择box的index
    # keep = dets[:,4][box_scores>thresh].long()
    # print('==> final dets.shape: ', dets.shape)
    # print('==> final box_scores.shape: ', box_scores.shape)
    # print('==> final box_scores: ', box_scores)
    valid = box_scores>thresh
    # print('==> valid.shape: ', valid.shape)
    det_indexes = dets[:,7]
    # print('==> det_indexes.shape: ', det_indexes.shape)
    print('==> det_indexes: ', det_indexes)

    # keep = dets[:,7][box_scores>thresh].long()
    keep = (det_indexes[valid]).long()

    # print('==> output keep.shape: ', keep.shape)
    return keep

def soft_nms_torch_1(dets, box_scores, iou_thresh=0.1, sigma=0.009, thresh=0.001, cuda=1):
    """
    Build a pytorch implement of Soft NMS algorithm.
    # Augments
        dets:        boxes coordinate tensor (format:[y1, x1, y2, x2])
        box_scores:  box score tensors
        sigma:       variance of Gaussian function
        thresh:      score thresh
        cuda:        CUDA flag
    # Return
        the index of the selected boxes
    """
    # print('0. ==> dets.shape: ', dets.shape)
    # print('1. ==> box_scores.shape: ', box_scores.shape)

    # Indexes concatenate boxes with the last column
    N = dets.shape[0]
    if cuda:
        indexes = torch.arange(0, N, dtype=torch.float).cuda().view(N, 1)
    else:
        indexes = torch.arange(0, N, dtype=torch.float).view(N, 1)
    dets = torch.cat((dets, indexes), dim=1)
    # print('2. ==> dets.shape: ', dets.shape)

    scores = box_scores
    for i in range(N):
        # print('==> i: ', i)
        # intermediate parameters for later parameters exchange
        tscore = scores[i].clone()
        pos = i + 1

        if i != N - 1:
            maxscore, maxpos = torch.max(scores[pos:], dim=0)
            if tscore < maxscore:
                dets[i], dets[maxpos.item() + i + 1] = dets[maxpos.item() + i + 1].clone(), dets[i].clone()
                scores[i], scores[maxpos.item() + i + 1] = scores[maxpos.item() + i + 1].clone(), scores[i].clone()

        # IoU calculate
        if dets[pos:, 0].shape[0] != 0:
            ovr = boxes_iou_bev(dets[i, :7].view(-1, 7), dets[pos:, :7].view(-1, 7))
            # ovr = boxes_iou3d_gpu(dets[i, :7].view(-1, 7), dets[pos:, :7].view(-1, 7))
            # 这里是计算BEV IoU还是3D IoU
            ovr = ovr.squeeze(dim=0)
            # print('==> ovr.shape: ', ovr.shape)
            # print('==> ovr', ovr)
        else:
            ovr = torch.tensor([]).to(dets.device)
            # print('==> ovr.shape: ', ovr.shape)
            # print('==> ovr', ovr)
        # print('==> ovr: ', ovr)
        # Gaussian decay
        # '''
        weight = torch.exp(-(ovr * ovr) / sigma)
        # print('==> weight.shape: ', weight.shape)
        # print('==> weight: ', weight)
        scores[pos:] = weight * scores[pos:]
        # print('==> scores.shape: ', scores.shape)
        # '''

        # original nms
        '''
        weight = torch.ones(ovr.shape).to(dets.device)
        weight[ovr > iou_thresh] = 0
        scores[pos:] = weight * scores[pos:]
        '''

    # select the boxes and keep the corresponding indexes
    # print('==> final scores.shape: ', scores.shape)
    # keep = dets[:, 4][scores > thresh].int()
    keep = dets[:, -1][scores > thresh].long()
    # print('==> keep: ', keep)
    return keep

def iou_weighted_nms_cpu(
                        boxes, scores, iou_preds, labels, anchors,
                        suppressed_thresh = 0.3, # 0.3 in CIA-SSD OR 0.1?
                        cnt_thresh = 0.5, #2.6, # from CIA-SSD
                        match_thresh = 0.3,
                        nms_sigma_dist_interval=[0, 20, 40, 60],
                        nms_sigma_square = [0.0009, 0.009, 0.1, 1]):
    """
    必须是单张点云内的数据
    thresh: 是suppressed thresh
    boxes: (N, 7),
    scores: (N, )
    iou_preds: (N, )
    labels: (N, )
    anchors: (N, 7), 与boxes对齐的
    """
    scores_ret = []
    boxes_ret = []
    labels_ret = []

    boxes_r = boxes.clone().cpu().numpy()
    scores_r = scores.clone().cpu().numpy()
    IOU_preds_r = iou_preds.clone().cpu().numpy()
    labels_r = labels.clone().cpu().numpy()
    anchors_r = anchors.clone().cpu().numpy()
    box_corners_r = boxes_to_corners_3d(boxes).clone().cpu().numpy() # (N, 8, 3)
    nms_sigma_dist_interval_r = nms_sigma_dist_interval
    nms_sigma_square_r = nms_sigma_square
    # print('==> boxes_r.shape: ', boxes_r.shape)
    # print('==> scores_r.shape: ', scores_r.shape)
    # print('==> IOU_preds_r.shape: ', IOU_preds_r.shape)
    # print('==> labels_r.shape: ', labels_r.shape)
    # print('==> anchors_r.shape: ', anchors_r.shape)
    # print('==> box_corners_r.shape: ', box_corners_r.shape)

    ndets = boxes.shape[0]
    suppressed_rw = np.zeros((ndets))
    # weight_pos = np.zeros((7))
    # avg_pos = np.zeros((7))

    standup_iou_r = boxes_iou3d_gpu(boxes, boxes).clone().cpu().numpy() # (N, N)

    # 假设用bev_dist(box, anchor)
    # TODO 这里是否需要根据这个来调制呢？
    # dist = torch.pow((boxes[:,0]-anchors[:,0]), 2) + torch.pow( (boxes[:,1]-anchors[:,1]), 2)
    # dist_nrom = (1 - torch.softmax(dist, dim=0)).cpu().numpy()
    # 对score进行调制
    # scores_rw = scores_r * dist_nrom

    scores_rw = scores_r * 1

    # 对调制后的score_rw取最大值，然后归一化
    score_max4norm = np.max(scores_rw)
    scores_rw /= (score_max4norm + 1e-6)

    assert not np.array_equal(scores_rw, scores_r)

    while (True):
        score_max = -1
        idx_max = -1
        flag_all_checked = True
        # find out the box with the maximum score
        for i in range(ndets):
            if suppressed_rw[i] == 1:
                continue
            flag_all_checked = False
            # 找到原始分数最大的那一个
            if scores_r[i] > score_max:
                score_max = scores_r[i]
                idx_max = i
        if flag_all_checked:
            break
        # 计算当前box到原点的距离, bev距离上的
        dist2origin = np.linalg.norm(boxes_r[idx_max, 0:2], ord=2)
        suppressed_rw[idx_max] = 1

        # 对weight_pos avg_pos重新置零
        weight_pos = np.zeros((1))
        avg_pos = np.zeros((7))
        score_box = -1
        cnt = 0
        recover_list = []
        merge_cnt = 0
        merge_box_list = []
        for j in range(ndets):
            # 当前idx_max的box去和所有的box计算bev_iou
            poly = Polygon(box_corners_r[idx_max, 0:4, 0:2])
            qpoly = Polygon(box_corners_r[j, 0:4, 0:2])
            if poly.intersects(qpoly):
                inter_area = poly.intersection(qpoly).area
                # 计算unions
                union_area = poly.union(qpoly).area
                overlap = inter_area / union_area
                # 计算重叠的数量，用于过滤false_positives
                if (overlap > 0) and (labels_r[j] == labels_r[idx_max]):
                    # TODO: 选择？
                    # cnt += overlap * IOU_preds_r[j]
                    cnt += overlap * scores_r[j]
                # 选择和当前box的iou>0.3的这些辅助box来重新加权
                if (overlap >= match_thresh) and (labels_r[j] == labels_r[idx_max]):
                    if score_box < scores_rw[j]:
                        # 始终保留分数最大的一个
                        # TODO: 这里box是否能够加权呢？
                        score_box = scores_rw[j]
                    IOU_weight = 0.
                    # 和当前box的iou越大的辅助box的weight应该更大
                    # 根据到原点的距离分段进行
                    for k in range( len(nms_sigma_dist_interval_r) - 1):
                        dist_l = nms_sigma_dist_interval_r[k]
                        dist_r = nms_sigma_dist_interval_r[k+1]
                        # print('==> dist_l: ', dist_l)
                        # print('==> dist_r: ', dist_r)
                        if (dist2origin >= dist_l) and (dist2origin < dist_r):
                            tmp = -np.power( (1-overlap), 2 ) / nms_sigma_square_r[k]
                            IOU_weight = np.exp(tmp)
                            # print('==> dist2origin: ', dist2origin)

                    # 使用IOU_weight对box进行加权
                    # avg_pos += IOU_weight * IOU_preds_r[j] * boxes_r[j]
                    # weight_pos += IOU_weight * IOU_preds_r[j]
                    avg_pos += IOU_weight * scores_r[j] * boxes_r[j]
                    weight_pos += IOU_weight * scores_r[j]
                    assert avg_pos.shape[0] == 7
                    merge_cnt += 1
                    merge_box_list.append(boxes_r[j])

                # suppress the box whose IOU with box[idx_max] > suppressed_thresh
                if (suppressed_rw[j] != 1) and (standup_iou_r[idx_max, j] > 0):
                    if (overlap > suppressed_thresh):
                        suppressed_rw[j] = 1
                        recover_list.append(j)

        # print('==> merge_cnt: ', merge_cnt)
        # print('==> merge_box_list: ', merge_box_list)
        # print('==> cnt: ', cnt)
        if (cnt > cnt_thresh):
            scores_ret.append(score_box * score_max4norm)
            avg_pos /= (weight_pos + 1e-6)
            boxes_ret.append(avg_pos)
            labels_ret.append(labels_r[idx_max])
            # print('==> merge valid.')
            if avg_pos[3] + avg_pos[4] + avg_pos[5] == 0:
                scores_ret.pop()
                boxes_ret.pop()
                labels_ret.pop()
                for k in range(len(recover_list)):
                    suppressed_rw[recover_list[k]] = 0
        else:
            # 如果当前这个box是一个false pos，那么就恢复(释放)之前被它抑制掉的其它box
            for k in range(len(recover_list)):
                suppressed_rw[recover_list[k]] = 0

    if len(boxes_ret):
        # not empty:
        ret_boxes_np = np.stack(boxes_ret, axis=0)
        ret_scores_np = np.stack(scores_ret, axis=0)
        ret_labels_np = np.stack(labels_ret, axis=0)
    else:
        ret_boxes_np = np.zeros( [0, 7] )
        ret_scores_np = np.zeros( [0, 1] )
        ret_labels_np = np.zeros( [0, 1] )

    ret_boxes = torch.from_numpy(ret_boxes_np).to(boxes.device)
    ret_scores = torch.from_numpy(ret_scores_np).to(boxes.device)
    ret_labels = torch.from_numpy(ret_labels_np).to(boxes.device)
    # boxes, scores, iou_preds, labels
    '''
    print('==> input boxes: ', boxes)
    print('==> input scores: ', scores)
    print('==> input iou_preds: ', iou_preds)
    print('==> input labels: ', labels)
    print('==> ret_boxes.shape: ', ret_boxes)
    print('==> ret_scores.shape: ', ret_scores)
    print('==> ret_labels.shape: ', ret_labels)
    '''
    return ret_boxes, ret_scores, ret_labels

def matched_boxes_bevdist(boxes_a, boxes_b):
    """
    以bev上的中心点距离作为certainty
    """
    dist = torch.pow((boxes_a[:,0]-boxes_b[:,0]), 2) + torch.pow( (boxes_a[:,1]-boxes_b[:,1]), 2)
    dist_nrom = 1 - torch.softmax(dist, dim=0)

    return dist_nrom.view(-1, 1)

def matched_boxes_iou3d_cpu(boxes_a, boxes_b):
    """
    # 前提是一一对齐的两组box
    Args:
        boxes_a: tensor, (N, 7) [x,y,z,dx,dy,dz,heading]
        boxes_b: tensor, (N, 7) [x,y,z,dx,dy,dz,heading]
    Returns:
        ans_iou: tensor, (N, 1)
    """
    assert boxes_a.shape == boxes_b.shape

    # 1. 计算高度方向上的overlap
    # height overlap, (N, )
    boxes_a_height_max = (boxes_a[:, 2] + boxes_a[:, 5] / 2)
    boxes_a_height_min = (boxes_a[:, 2] - boxes_a[:, 5] / 2)
    boxes_b_height_max = (boxes_b[:, 2] + boxes_b[:, 5] / 2)
    boxes_b_height_min = (boxes_b[:, 2] - boxes_b[:, 5] / 2)
    max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min)
    min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max)
    overlaps_h = torch.clamp(min_of_max - max_of_min, min=0)

    # 2. 转为corner, (N, 8, 3)
    corners_a = boxes_to_corners_3d(boxes_a)
    corners_b = boxes_to_corners_3d(boxes_b)

    # 3. 准备numpy格式的polygon 2D points
    np_2dcorners_a = corners_a.clone().cpu().numpy()
    np_2dcorners_b = corners_b.clone().cpu().numpy()
    np_overlaps_bev = np.zeros( (np_2dcorners_a.shape[0]) )
    for i in range(np_2dcorners_a.shape[0]):
        poly_points_a = np_2dcorners_a[i, 0:4, 0:2]
        poly_points_b = np_2dcorners_b[i, 0:4, 0:2]
        poly_a = Polygon(poly_points_a)
        poly_b = Polygon(poly_points_b)
        if poly_a.is_valid and poly_b.is_valid:
            # check is valid,  A valid Polygon may not possess any overlapping exterior or interior rings.
            overlap = poly_a.intersection(poly_b).area
        else:
            overlap = 0.
        np_overlaps_bev[i] = overlap
    # move back to torch
    overlaps_bev = torch.from_numpy(np_overlaps_bev).to(boxes_a.device)

    # 4. 计算3d overlaps
    overlaps_3d = overlaps_bev * overlaps_h
    assert overlaps_3d.shape == overlaps_bev.shape == overlaps_h.shape

    # 5. 计算3d iou
    vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5])
    vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5])

    iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-6)
    return iou3d.view(-1, 1)

def boxes_bev_iou_cpu(boxes_a, boxes_b):
    """
    Args:
        boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
        boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading]

    Returns:

    """
    boxes_a, is_numpy = common_utils.check_numpy_to_torch(boxes_a)
    boxes_b, is_numpy = common_utils.check_numpy_to_torch(boxes_b)
    assert not (boxes_a.is_cuda or boxes_b.is_cuda), 'Only support CPU tensors'
    assert boxes_a.shape[1] == 7 and boxes_b.shape[1] == 7
    ans_iou = boxes_a.new_zeros(torch.Size((boxes_a.shape[0], boxes_b.shape[0])))
    iou3d_nms_cuda.boxes_iou_bev_cpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou)

    return ans_iou.numpy() if is_numpy else ans_iou

def boxes_iou_bev(boxes_a, boxes_b):
    """
    Args:
        boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
        boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading]

    Returns:
        ans_iou: (N, M)
    """
    assert boxes_a.shape[1] == boxes_b.shape[1] == 7
    ans_iou = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_()

    iou3d_nms_cuda.boxes_iou_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), ans_iou)

    return ans_iou

def boxes_iou3d_gpu(boxes_a, boxes_b):
    """
    Args:
        boxes_a: (N, 7) [x, y, z, dx, dy, dz, heading]
        boxes_b: (N, 7) [x, y, z, dx, dy, dz, heading]

    Returns:
        ans_iou: (N, M)
    """
    assert boxes_a.shape[1] == boxes_b.shape[1] == 7

    # height overlap
    boxes_a_height_max = (boxes_a[:, 2] + boxes_a[:, 5] / 2).view(-1, 1)
    boxes_a_height_min = (boxes_a[:, 2] - boxes_a[:, 5] / 2).view(-1, 1)
    boxes_b_height_max = (boxes_b[:, 2] + boxes_b[:, 5] / 2).view(1, -1)
    boxes_b_height_min = (boxes_b[:, 2] - boxes_b[:, 5] / 2).view(1, -1)

    # bev overlap
    overlaps_bev = torch.cuda.FloatTensor(torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_()  # (N, M)
    iou3d_nms_cuda.boxes_overlap_bev_gpu(boxes_a.contiguous(), boxes_b.contiguous(), overlaps_bev)

    max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min)
    min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max)
    overlaps_h = torch.clamp(min_of_max - max_of_min, min=0)

    # 3d iou
    overlaps_3d = overlaps_bev * overlaps_h

    vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1, 1)
    vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(1, -1)

    iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-6)

    # 还要进行限制一下，小于0或者大于1的都设置为0
    iou3d[iou3d < 0] = 0
    iou3d[iou3d > 1] = 1

    return iou3d

def nms_gpu(boxes, scores, thresh, pre_maxsize=None, **kwargs):
    """
    :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
    :param scores: (N)
    :param thresh:
    :return:
    """
    assert boxes.shape[1] == 7
    order = scores.sort(0, descending=True)[1]
    if pre_maxsize is not None:
        order = order[:pre_maxsize]

    boxes = boxes[order].contiguous()
    keep = torch.LongTensor(boxes.size(0))
    num_out = iou3d_nms_cuda.nms_gpu(boxes, keep, thresh)
    return order[keep[:num_out].cuda()].contiguous(), None

def nms_normal_gpu(boxes, scores, thresh, **kwargs):
    """
    :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
    :param scores: (N)
    :param thresh:
    :return:
    """
    assert boxes.shape[1] == 7
    order = scores.sort(0, descending=True)[1]

    boxes = boxes[order].contiguous()

    keep = torch.LongTensor(boxes.size(0))
    num_out = iou3d_nms_cuda.nms_normal_gpu(boxes, keep, thresh)
    return order[keep[:num_out].cuda()].contiguous(), None

def batch_boxes_iou3d_gpu(boxes_a, boxes_b):
    """
    Args:
        boxes_a: (B, N, 7) [x, y, z, dx, dy, dz, heading]
        boxes_b: (B, M, 7) [x, y, z, dx, dy, dz, heading], GT, 可能会补零，所以先去掉0

    Returns:
        ans_iou: (B, N, 1)
    """
    assert boxes_a.shape[0] == boxes_b.shape[0]
    assert boxes_a.shape[-1] == boxes_b.shape[-1]

    # for each batch
    boxiou_a_list = []
    batch_size = boxes_a.shape[0]
    for i in range(batch_size):
        boxes_a_single = boxes_a[i, ...] # (max_objs, 7)
        boxes_b_single = boxes_b[i, ...] # (M, 7)
        # print('==> 0. src_box_target_single.size(): ', src_box_target_single.size()) # (M, 7)
        # valid = ( torch.sum(boxes_b_single, dim=1) != 0 )
        # print('==> valid.size(): ', valid.size()) # (M)
        # boxes_b_single = boxes_b_single[valid]
        # print('==> 1. src_box_target_single.size(): ', src_box_target_single.size()) # (Nb(=M))
        iou3d = boxes_iou3d_gpu(boxes_a=boxes_a_single,
                                boxes_b=boxes_b_single)  # (Na, Nb)
        # print('==> iou3d.size(): ', iou3d.size()) # (max_objs(=Na), Nb)
        max_overlaps, gt_assignment = torch.max(iou3d, dim=1) # roi side iou
        # print('==> max_overlaps.size(): ', max_overlaps.size()) # (max_objs(=Na))
        boxiou_a_list.append( max_overlaps.view(-1, 1) )

    # stack for batch axis, batch (max_objs, 1) -> (batch, max_objs, 1)
    boxiou_a = torch.stack(boxiou_a_list, dim=0)

    return  boxiou_a


#! /usr/bin/env python
# coding: utf-8

import multiprocessing as mp
from ctypes import c_double, c_int
import numpy as np
import os
import sys
import signal

from multiprocessing.sharedctypes import Value, Array
from multiprocessing import Lock
from multiprocessing.managers import SyncManager

import cProfile
import time
import os

class CheckPoint(Exception):
    pass

def sighandler(signal, frame):
    raise CheckPoint

class CPNest(object):
    """
    Class to control CPNest sampler
    cp = CPNest(usermodel,nlive=100,output='./',verbose=0,seed=None,maxmcmc=100,nthreads=None,balanced_sampling = True)

    Input variables:
    usermodel : an object inheriting cpnest.model.Model that defines the user's problem
    nlive : Number of live points (100)
    poolsize: Number of objects in the sampler pool (100)
    output : output directory (./)
    verbose: Verbosity, 0=silent, 1=progress, 2=diagnostic, 3=detailed diagnostic
    seed: random seed (default: 1234)
    maxmcmc: maximum MCMC points for sampling chains (100)
    nthreads: number of parallel samplers. Default (None) uses mp.cpu_count() to autodetermine
    nhamiltomnian: number of sampler threads using an hamiltonian samplers. Default: 0
    resume: determines whether cpnest will resume a run or run from scratch. Default: False.
    proposal: dictionary/list with custom jump proposals. key 'mhs' for the
    Metropolis-Hastings sampler, 'hmc' for the Hamiltonian Monte-Carlo sampler. Default: None
    n_periodic_checkpoint: int
        checkpoint the sampler every n_periodic_checkpoint iterations
        Default: 100000

    """
    def __init__(self,
                 usermodel,
                 nlive        = 100,
                 poolsize     = 100,
                 output       = './',
                 verbose      = 0,
                 seed         = None,
                 maxmcmc      = 100,
                 nthreads     = None,
                 nhamiltonian = 0,
                 resume       = False,
                 proposals     = None,
                 n_periodic_checkpoint = 1000):
        if nthreads is None:
            self.nthreads = mp.cpu_count()
        else:
            self.nthreads = nthreads

        print('Running with {0} parallel threads'.format(self.nthreads))
        from .sampler import HamiltonianMonteCarloSampler, MetropolisHastingsSampler
        from .NestedSampling import NestedSampler
        from .proposal import DefaultProposalCycle, HamiltonianProposalCycle
        if proposals is None:
            proposals = dict(mhs=DefaultProposalCycle,
                             hmc=HamiltonianProposalCycle)
        elif type(proposals) == list:
            proposals = dict(mhs=proposals[0],
                             hmc=proposals[1])
        self.nlive    = nlive
        self.verbose  = verbose
        self.output   = output
        self.poolsize = poolsize
        self.posterior_samples = None
        self.manager = RunManager(nthreads=self.nthreads)
        self.manager.start()
        self.user     = usermodel
        self.resume = resume

        if seed is None: self.seed=1234
        else:
            self.seed=seed

        self.process_pool = []

        # instantiate the nested sampler class
        resume_file = os.path.join(output, "nested_sampler_resume.pkl")
        if not os.path.exists(resume_file) or resume == False:
            self.NS = NestedSampler(self.user,
                        nlive          = nlive,
                        output         = output,
                        verbose        = verbose,
                        seed           = self.seed,
                        prior_sampling = False,
                        manager        = self.manager,
                        n_periodic_checkpoint = n_periodic_checkpoint)
        else:
            self.NS = NestedSampler.resume(resume_file, self.manager, self.user)

        # instantiate the sampler class
        for i in range(self.nthreads-nhamiltonian):
            resume_file = os.path.join(output, "sampler_{0:d}.pkl".format(i))
            if not os.path.exists(resume_file) or resume == False:
                sampler = MetropolisHastingsSampler(self.user,
                                  maxmcmc,
                                  verbose     = verbose,
                                  output      = output,
                                  poolsize    = poolsize,
                                  seed        = self.seed+i,
                                  proposal    = proposals['mhs'](),
                                  resume_file = resume_file,
                                  manager     = self.manager
                                  )
                sampler.checkpoint()
            else:
                sampler = MetropolisHastingsSampler.resume(resume_file,
                                                           self.manager,
                                                           self.user)

            p = mp.Process(target=sampler.produce_sample)
            self.process_pool.append(p)

        for i in range(self.nthreads-nhamiltonian,self.nthreads):
            resume_file = os.path.join(output, "sampler_{0:d}.pkl".format(i))
            if not os.path.exists(resume_file) or resume == False:
                sampler = HamiltonianMonteCarloSampler(self.user,
                                  maxmcmc,
                                  verbose     = verbose,
                                  output      = output,
                                  poolsize    = poolsize,
                                  seed        = self.seed+i,
                                  proposal    = proposals['hmc'](model=self.user),
                                  resume_file = resume_file,
                                  manager     = self.manager
                                  )
            else:
                sampler = HamiltonianMonteCarloSampler.resume(resume_file,
                                                              self.manager,
                                                              self.user)
            p = mp.Process(target=sampler.produce_sample)
            self.process_pool.append(p)

    def run(self):
        """
        Run the sampler
        """
        if self.resume:
            signal.signal(signal.SIGTERM, sighandler)
            signal.signal(signal.SIGQUIT, sighandler)
            signal.signal(signal.SIGINT, sighandler)
            signal.signal(signal.SIGUSR1, sighandler)
            signal.signal(signal.SIGUSR2, sighandler)

        #self.p_ns.start()
        for each in self.process_pool:
            each.start()
        try:
            self.NS.nested_sampling_loop()
            for each in self.process_pool:
                each.join()
        except CheckPoint:
            self.checkpoint()
            sys.exit()

        self.posterior_samples = self.get_posterior_samples(filename=None)
        if self.verbose>1: self.plot()

        #TODO: Clean up the resume pickles

    def get_nested_samples(self, filename='nested_samples.dat'):
        """
        returns nested sampling chain
        Parameters
        ----------
        filename : string
                   If given, file to save nested samples to

        Returns
        -------
        pos : :obj:`numpy.ndarray`
        """
        import numpy.lib.recfunctions as rfn
        self.nested_samples = rfn.stack_arrays(
                [s.asnparray()
                    for s in self.NS.nested_samples]
                ,usemask=False)
        if filename:
            np.savetxt(os.path.join(
                self.NS.output_folder,'nested_samples.dat'),
                self.nested_samples.ravel(),
                header=' '.join(self.nested_samples.dtype.names),
                newline='\n',delimiter=' ')
        return self.nested_samples

    def get_posterior_samples(self, filename='posterior.dat'):
        """
        Returns posterior samples

        Parameters
        ----------
        filename : string
                   If given, file to save posterior samples to

        Returns
        -------
        pos : :obj:`numpy.ndarray`
        """
        import numpy as np
        import os
        from .nest2pos import draw_posterior_many
        nested_samples = self.get_nested_samples()
        posterior_samples = draw_posterior_many([nested_samples],[self.nlive],verbose=self.verbose)
        posterior_samples = np.array(posterior_samples)
        # TODO: Replace with something to output samples in whatever format
        if filename:
            np.savetxt(os.path.join(
                self.NS.output_folder,'posterior.dat'),
                self.posterior_samples.ravel(),
                header=' '.join(posterior_samples.dtype.names),
                newline='\n',delimiter=' ')
        return posterior_samples

    def plot(self, corner = True):
        """
        Make diagnostic plots of the posterior and nested samples
        """
        pos = self.posterior_samples
        from . import plot
        for n in pos.dtype.names:
            plot.plot_hist(pos[n].ravel(),name=n,filename=os.path.join(self.output,'posterior_{0}.png'.format(n)))
        for n in self.nested_samples.dtype.names:
            plot.plot_chain(self.nested_samples[n],name=n,filename=os.path.join(self.output,'nschain_{0}.png'.format(n)))
        import numpy as np
        plotting_posteriors = np.squeeze(pos.view((pos.dtype[0], len(pos.dtype.names))))
        if corner: plot.plot_corner(plotting_posteriors,labels=pos.dtype.names,filename=os.path.join(self.output,'corner.png'))

    def worker_sampler(self, producer_pipe, logLmin):
        cProfile.runctx('self.sampler.produce_sample(producer_pipe, logLmin)', globals(), locals(), 'prof_sampler.prof')

    def worker_ns(self):
        cProfile.runctx('self.NS.nested_sampling_loop(self.consumer_pipes)', globals(), locals(), 'prof_nested_sampling.prof')

    def profile(self):
        for i in range(0,self.NUMBER_OF_PRODUCER_PROCESSES):
            p = mp.Process(target=self.worker_sampler, args=(self.queues[i%len(self.queues)], self.NS.logLmin ))
            self.process_pool.append(p)
        for i in range(0,self.NUMBER_OF_CONSUMER_PROCESSES):
            p = mp.Process(target=self.worker_ns, args=(self.queues, self.port, self.authkey))
            self.process_pool.append(p)
        for each in self.process_pool:
            each.start()

    def checkpoint(self):
        self.manager.checkpoint_flag=1

class RunManager(SyncManager):
    def __init__(self, nthreads=None, **kwargs):
        super(RunManager,self).__init__(**kwargs)
        self.nconnected=mp.Value(c_int,0)
        self.producer_pipes = list()
        self.consumer_pipes = list()
        for i in range(nthreads):
            consumer, producer = mp.Pipe(duplex=True)
            self.producer_pipes.append(producer)
            self.consumer_pipes.append(consumer)
        self.logLmin=None
        self.nthreads=nthreads

    def start(self):
        super(RunManager, self).start()
        self.logLmin = mp.Value(c_double,-np.inf)
        self.checkpoint_flag=mp.Value(c_int,0)

    def connect_producer(self):
        """
        Returns the producer's end of the pipe
        """
        with self.nconnected.get_lock():
            n = self.nconnected.value
            pipe = self.producer_pipes[n]
            self.nconnected.value+=1
        return pipe, n


# -*- coding: utf-8 -*-
"""
PiViewer: an open-source tool for automated detection and display of pi-pi interactions.
Created by Hu Ge <gehuchina@gmail.com>
"""

# Copyright Notice
# ================
#
# The PyMOL Plugin source code in this file is copyrighted, but you can
# freely use and copy it as long as you don't change or remove any of
# the copyright notices.
#
# ----------------------------------------------------------------------
# #
#                        All Rights Reserved
#
# Permission to use, copy, modify, distribute, and distribute modified
# versions of this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both the copyright notice and
# this permission notice appear in supporting documentation, and that
# the name(s) of the author(s) not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# THE AUTHOR(S) DISCLAIM ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.  IN
# NO EVENT SHALL THE AUTHOR(S) BE LIABLE FOR ANY SPECIAL, INDIRECT OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF
# USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
# OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.
# ----------------------------------------------------------------------

import openbabel as ob
import pybel
import numpy as np

# Define vecAngle for the degree between two vectors.
def vecAngle(vec1, vec2):
    '''
    vec1,vec2 are ob.vector3 objects,
    return the angle between them in degree.
    '''
    dotprod = vec1.GetX() * vec2.GetX() + vec1.GetY() * vec2.GetY() + vec1.GetZ() * vec2.GetZ()
    deg = np.arccos(dotprod) * 180 / np.pi
    if deg > 90:
        deg = 180 - deg
    return deg

# The main PiPi viewer function
def find_PiPi(pdb_file, lig_name, centroid_distance=5.0, dih_parallel=25, dih_tshape=80, verbose=1):
    """
    Find Pi-Pi interactions around the specified ligand residue from the pdb file.
    :param pdb_file: path of the target file in PDB format.
    :param lig_name: ligand residue name.
    :param centroid_distance: Max ring centroid distance
    :param dih_parallel: Max dihedral (parallel)
    :param dih_tshape: Min dihedral (T-shaped)
    :return: number of Pi-Pi interactions found
    """
    # Get ligand residue and print its name.
    ligAtomList = []
    ligAtomIdList = []
    mol = pybel.readfile('pdb', pdb_file).next()
    if verbose: print "A total of %s residues" % mol.OBMol.NumResidues()
    lig = None
    for res in ob.OBResidueIter(mol.OBMol):
        # print res.GetName()
        if res.GetName() == lig_name:
            lig = res
            if verbose: print "Ligand residue name is:", lig.GetName()
            break
    if not lig:
        if verbose: print "No ligand residue %s found, please confirm." % lig_name
        return -1
    else:
        for atom in ob.OBResidueAtomIter(lig):
            # print atom.GetIdx()
            ligAtomList.append(atom)
            ligAtomIdList.append(atom.GetIdx())

    # Set ring_id
    i = 0
    for ring in mol.sssr:
        ring.ring_id = i
        i += 1
        # print ring.ring_id

    # Determine which rings are from ligand.
    ligRingList = []
    ligAroRingList = []
    ligRingIdList = []
    recRingList = []
    recAroRingList = []
    for ring in mol.sssr:
        for atom in ligAtomList:
            if ring.IsMember(atom):
                if ring not in ligRingList:
                    ligRingList.append(ring)
                    ligRingIdList.append(ring.ring_id)
                    if verbose: print "ligand ring_ID: ", ring.ring_id,
                    if ring.IsAromatic():
                        if verbose: print "aromatic"
                        ligAroRingList.append(ring)
                    else:
                        if verbose: print "saturated"
    for ring in mol.sssr:
        if ring.ring_id not in ligRingIdList:
            recRingList.append(ring)
            if ring.IsAromatic():
                recAroRingList.append(ring)
    if verbose: print "\nReceptor has ", len(recRingList), " rings,",
    if verbose: print " has ", len(recAroRingList), " aromatic rings."

    # Find and show the rings
    ligRingCenter = ob.vector3()
    recRingCenter = ob.vector3()
    ligNorm1 = ob.vector3()
    ligNorm2 = ob.vector3()
    recNorm1 = ob.vector3()
    recNorm2 = ob.vector3()
    count = 0
    lig_ring_index = 0
    for ligRing in ligAroRingList:
        lig_ring_index += 1
        ligRing.findCenterAndNormal(ligRingCenter, ligNorm1, ligNorm2)
        rec_ring_index = 0
        for recRing in recAroRingList:
            rec_ring_index += 1
            recRing.findCenterAndNormal(recRingCenter, recNorm1, recNorm2)
            dist = ligRingCenter.distSq(recRingCenter) ** 0.5
            angle = vecAngle(ligNorm1, recNorm1)
            if (dist < centroid_distance and (angle < dih_parallel or angle > dih_tshape)):  # the criteria
                count += 1
                if verbose: print "Pi-Pi ring pairs: %3s,%3s  Angle(deg.): %5.2f  Distance(A): %.2f" % (recRing.ring_id, ligRing.ring_id, angle, dist)
    if verbose: print "Total Pi-Pi interactions:", count
    return count

if __name__ == '__main__':
    pdb_file = r'C:\CloudStation\Epicat\Git\PiViewer\1ACJ.pdb'
    lig_name = 'THA'
    find_PiPi(pdb_file, lig_name, 5.0, 25, 80)


#!/usr/bin/env python3

import gemmi
import copy
import numpy as np
from scipy.spatial import distance
import string
from rdkit import Chem
from rdkit.Geometry import Point3D
import warnings

from . import xyz2mol

import argparse

def parse_arguments():

    """
    Parse command line arguments.
    """

    parser = argparse.ArgumentParser(
        description="Python script for building supercell that can be used with PBC in MD simulation."
        )

    parser.add_argument(
        '--input',
        "-i",
        type=str,
        help="Input cif file",
        required=True
        )

    parser.add_argument(
        '--prefix',
        "-pre",
        type=str,
        help="Prefix for pdb and csv file.",
        required=True
        )

    parser.add_argument(
        '--a_min_max',
        "-a",
        type=float,
        help="If this is two arguments: Minimum and maximum unit cell replicates along direction `a`. \
If this is a single argument: Minimum length of axis `a` in the supercell. Units [nm]",
        required=False,
        default=[-1,1],
        nargs='+',
        )

    parser.add_argument(
        '--b_min_max',
        "-b",
        type=float,
        help="If this is two arguments: Minimum and maximum unit cell replicates along direction `b`. \
If this is a single argument: Minimum length of axis `b` in the supercell. Units [nm]",
        required=False,
        default=[-1,1],
        nargs='+',
        )

    parser.add_argument(
        '--c_min_max',
        "-c",
        type=float,
        help="If this is two arguments: Minimum and maximum unit cell replicates along direction `c`. \
If this is a single argument: Minimum length of axis `c` in the supercell. Units [nm]",
        required=False,
        default=[-1,1],
        nargs='+',
        )

    parser.add_argument(
        '--addhs',
        "-ah",
        action='store_true',
        help="Remove any existing hydrogen and add protonate molecule. Requires OpenEye Toolkits.",
        required=False,
        default=False,
        )

    parser.add_argument(
        '--addwater',
        "-aw",
        type=int,
        help="Number of water molecules to add",
        required=False,
        default=0,
        )

    parser.add_argument(
        '--use_symmetry_operations',
        "-op",
        action='store_true',
        help="Use symmetry operations in cif file instead of space group.",
        required=False,
        default=False,
        )

    parser.add_argument(
        '--n_protonation_attempts',
        "-np",
        type=int,
        help="Number of attempts to compute protonatation states in  unit cell.",
        required=False,
        default=0,
        )

    parser.add_argument(
        '--use_openeye',
        "-oe",
        action='store_true',
        help="Use openeye-toolkit for topology building. Otherwise use xyz2mol.",
        required=False,
        default=False,
        )

    return parser.parse_args()

def combine_mols(mol_list):

    mol_list  = copy.deepcopy(mol_list)
    N_mol_per_unitcell = len(mol_list)
    mol_combo = Chem.Mol()
    for mol in mol_list:
        mol_combo = Chem.CombineMols(mol_combo, mol)
    return mol_combo

class FfWrapper(object):

    def __init__(self, mol, only_hydrogen=True):

        from rdkit.Chem import AllChem as Chem

        self._mol = copy.deepcopy(mol)

        mp  = Chem.MMFFGetMoleculeProperties(mol)
        self._ffm = Chem.MMFFGetMoleculeForceField(mol, mp)

        self._H_atom_list   = list()
        self._all_atom_list = list()
        for atom in mol.GetAtoms():
            if atom.GetAtomicNum() == 1:
                self._H_atom_list.append(atom.GetIdx())
            self._all_atom_list.append(atom.GetIdx())

        self._H_atom_list = np.array(self._H_atom_list)
        self._H_idxs = np.arange(mol.GetNumAtoms()*3, dtype=int)
        self._H_idxs = self._H_idxs.reshape((mol.GetNumAtoms(), 3))
        self._H_idxs = self._H_idxs[self._H_atom_list]
        self._H_idxs_flat = self._H_idxs.flatten()

        self._all_atom_list = np.array(self._all_atom_list)
        self._all_idxs = np.arange(mol.GetNumAtoms()*3, dtype=int)
        self._all_idxs = self._all_idxs.reshape((mol.GetNumAtoms(), 3))
        self._all_idxs_flat = self._H_idxs.flatten()

        self._num_all_atoms = len(self._all_atom_list)
        self._num_H_atoms   = len(self._H_atom_list)

        self.only_hydrogen = only_hydrogen

    @property
    def num_all_atoms(self):
        return self._num_all_atoms

    @property
    def num_H_atoms(self):
        return self._num_H_atoms

    @property
    def H_atom_list(self):
        return self._H_atom_list

    @property
    def H_idxs(self):
        return self._H_idxs

    @property
    def H_idxs_flat(self):
        return self._H_idxs_flat

    @property
    def all_atom_list(self):
        return self._all_atom_list

    @property
    def all_idxs(self):
        return self._all_idxs

    @property
    def all_idxs_flat(self):
        return self._all_idxs_flat

    @property
    def ffm(self):
        return self._ffm

    @property
    def mol(self):
        return self._mol

    @property
    def pos(self):
        pos = np.array(self.ffm.Positions())
        if self.only_hydrogen:
            pos = pos[self.H_idxs_flat].tolist()
        return pos

    @pos.setter
    def pos(self, pos):

        from rdkit.Chem import AllChem as Chem

        conformer      = self._mol.GetConformer()

        if self.only_hydrogen:
            pos_stack = np.array(pos).reshape((self.num_H_atoms, 3))
            for H_atm_idx, atm_idx in enumerate(self.H_atom_list):
                conformer.SetAtomPosition(
                    int(atm_idx),
                    Point3D(
                        *pos_stack[H_atm_idx]
                        )
                    )
        else:
            pos_stack = np.array(pos).reshape((self.num_all_atoms, 3))
            for atm_idx in range(self._mol.GetNumAtoms()):
                conformer.SetAtomPosition(
                    int(atm_idx),
                    Point3D(
                        *pos_stack[atm_idx]
                        )
                    )

        mp  = Chem.MMFFGetMoleculeProperties(self._mol)
        self._ffm = Chem.MMFFGetMoleculeForceField(self._mol, mp)

    def ene(self,x):
        pos = self.ffm.Positions()
        pos = np.array(pos)
        if self.only_hydrogen:
            pos[self.H_idxs_flat] = x
        else:
            pos[:] = x
        pos = pos.tolist()
        return self.ffm.CalcEnergy(pos)

    def grad(self,x):
        pos = self.ffm.Positions()
        pos = np.array(pos)
        if self.only_hydrogen:
            pos[self.H_idxs_flat] = x
        else:
            pos[:] = x
        pos = pos.tolist()
        grad_vals = self.ffm.CalcGrad(pos)
        grad_vals = np.array(grad_vals)[self.H_idxs_flat]
        grad_vals = grad_vals.tolist()
        return grad_vals

def apply_delta_hydrogen(
    mol_list,
    delta_hydrogen_dict
    ):

    import numpy as np
    from rdkit.Chem import AllChem as Chem
    import copy

    mol_combo   = combine_mols(mol_list)
    emol        = Chem.EditableMol(mol_combo)
    offset_list = np.zeros(mol_combo.GetNumAtoms(), dtype=int)
    for atomidx, switch in delta_hydrogen_dict.items():
        atom    = mol_combo.GetAtomWithIdx(atomidx)
        charge  = atom.GetFormalCharge()
        if switch == 1:
            Hatom = Chem.Atom(1)
            idx1  = emol.AddAtom(Hatom)
            Patom = Chem.AtomFromSmarts(atom.GetSmarts())
            Patom.SetFormalCharge(charge+1)
            emol.ReplaceAtom(
                atomidx+int(offset_list[atomidx]),
                Patom
                )
            emol.AddBond(
                idx1,
                atomidx+int(offset_list[atomidx]),
                Chem.BondType.SINGLE
            )

        elif switch == -1:
            for n_atom in atom.GetNeighbors():
                if n_atom.GetAtomicNum() == 1:
                    idx1  = n_atom.GetIdx()
                    Patom = Chem.AtomFromSmarts(atom.GetSmarts())
                    Patom.SetFormalCharge(charge-1)
                    emol.ReplaceAtom(
                        atomidx+int(offset_list[atomidx]),
                        Patom
                        )
                    emol.RemoveAtom(
                        idx1+int(offset_list[idx1])
                        )
                    offset_list[idx1:] -= 1
                    break
        else:
            ### Nix.
            pass
    mol_combo_prot = emol.GetMol()
    Chem.SanitizeMol(mol_combo_prot)
    return Chem.GetMolFrags(mol_combo_prot, asMols=True)

def minimize_H(
    mol_list,
    ):

    import numpy as np
    from rdkit.Chem import AllChem as Chem
    from scipy import optimize
    import copy

    mol_combo = combine_mols(mol_list)
    ff        = FfWrapper(mol_combo, only_hydrogen=True)
    x0        = ff.pos

    result = optimize.minimize(
        fun=ff.ene,
        x0=x0,
        jac=ff.grad,
        method="L-BFGS-B"
        )
    ff.pos = result.x
    energy = ff.ene(result.x)

    return energy, Chem.GetMolFrags(ff.mol, asMols=True)

def assign_protonation_states(
    cell,
    mol_list,
    N_iterations=10):

    """
    Find the energetically best protonation state of the unitcell.
    """

    import numpy as np
    from rdkit.Chem import AllChem as Chem
    import copy

    mol_combo = combine_mols(mol_list)
    polar_heavy_atom = Chem.MolFromSmarts("[#7,#8,#16]")
    matches = mol_combo.GetSubstructMatches(
        polar_heavy_atom,
        uniquify=False
        )
    N_matches = len(matches)
    delta_hydrogen_dict_best = dict()
    for atm_idx in matches:
        atm_idx = int(atm_idx[0])
        delta_hydrogen_dict_best[atm_idx] = 0
    energy_best = 999999999999999999999.
    mol_list_best = copy.deepcopy(mol_list)
    N_mol_per_unitcell = len(mol_list)

    charge0 = Chem.GetFormalCharge(mol_combo)

    #with open(f"./best_init.pdb", "w") as fopen:
    #    fopen.write(
    #        Chem.MolToPDBBlock(
    #            combine_mols(
    #                mol_list_best
    #                )
    #            )
    #        )
    #count = 0
    for _ in range(N_iterations):
        delta_hydrogen_dict = copy.deepcopy(
            delta_hydrogen_dict_best
            )
        for i in range(N_matches):
            atm_idx_i = int(matches[i][0])
            delta0_i  = delta_hydrogen_dict[atm_idx_i]
            for j in range(i+1, N_matches):
                atm_idx_j = int(matches[j][0])
                delta0_j  = delta_hydrogen_dict[atm_idx_j]
                for dij in [[-1,1],[1,-1]]:
                    ###  0: Do nothing
                    ### +1: Add hydrogen
                    ### -1: Remove hydrogen
                    delta_hydrogen_dict[atm_idx_i] = dij[0]
                    delta_hydrogen_dict[atm_idx_j] = dij[1]
                    mol_list_prot = apply_delta_hydrogen(
                        copy.deepcopy(mol_list),
                        delta_hydrogen_dict
                        )
                    charge = Chem.GetFormalCharge(
                        combine_mols(mol_list_prot)
                        )
                    if charge == charge0:
                        mol_list_prot, _, _ = make_supercell(
                            cell,
                            mol_list_prot,
                            0, 2,
                            0, 2,
                            0, 2)
                        energy, mol_list_prot = minimize_H(mol_list_prot)
                        if energy < energy_best:
                            energy_best = energy
                            delta_hydrogen_dict_best = copy.deepcopy(
                                delta_hydrogen_dict
                                )
                            mol_list_best = copy.deepcopy(
                                mol_list_prot[:N_mol_per_unitcell]
                                )
                delta_hydrogen_dict[atm_idx_j] = delta0_j
            delta_hydrogen_dict[atm_idx_i] = delta0_i
        #print(f"Best energy {energy_best:4.2f}")
        #with open(f"./best_{count}.pdb", "w") as fopen:
        #    fopen.write(
        #        Chem.MolToPDBBlock(
        #            combine_mols(
        #                mol_list_best
        #                )
        #            )
        #        )
        #count += 1

    return mol_list_best

def get_nonoverlapping_atoms(atom_crds_ortho, filter_overlapping=False):

    """
    Retrieve indices of rows in `atom_crds_ortho` (N,3) list
    that correspond to non-overlapping atoms.
    If `filter_overlapping==True`, then one of the overlapping (not clear
    which one though) atoms will be retained.

    """

    from scipy.spatial import distance

    atom_crds_ortho_cp = np.copy(atom_crds_ortho)
    dists  = distance.cdist(atom_crds_ortho_cp, atom_crds_ortho_cp)
    np.fill_diagonal(dists, np.inf)
    if filter_overlapping:
        tril_idxs = np.tril_indices(atom_crds_ortho.shape[0])
        dists[tril_idxs] = np.inf
    invalids = np.where(dists < 0.01)[0]
    invalids = np.unique(invalids)
    valids   = np.arange(atom_crds_ortho_cp.shape[0], dtype=int)
    valids   = np.delete(valids, invalids)

    return valids

def random_fill(
    cell,
    mol_list,
    N_per_unitcell,
    radius=1.4,
    smiles="[H]O[H]"):

    """
    Randomly fill unit cell with molecules.
    """

    from rdkit.Chem import AllChem as Chem
    from rdkit.Geometry import Point3D
    import numpy as np
    from scipy.spatial import distance
    import gemmi
    import copy

    vdw_dict = {
        1  : 1.09, #H
        6  : 1.7,  #C
        7  : 1.55, #N
        8  : 1.52, #O
        9  : 1.47, #F
        15 : 1.8,  #P
        16 : 1.8,  #S
        17 : 1.75  #Cl
    }

    mol = Chem.MolFromSmiles(smiles)
    mol = Chem.AddHs(mol)

    Chem.EmbedMolecule(mol)
    Chem.UFFOptimizeMolecule(mol)

    conformer      = mol.GetConformer()
    atom_crds_mol  = conformer.GetPositions()
    atom_crds_mol  = np.array(atom_crds_mol)

    replicated_mol_list, _, _ = make_supercell(
        cell,
        mol_list,
        -1, 1,
        -1, 1,
        -1, 1)

    atom_crds_xtal = list()
    atom_radii_xtal = list()
    for mol_r in replicated_mol_list:
        conf     = mol_r.GetConformer(0)
        conf_pos = conf.GetPositions()
        for atom in mol_r.GetAtoms():
            if atom.GetAtomicNum() == 1:
                continue
            atom_radii_xtal.append(
                vdw_dict[atom.GetAtomicNum()]
            )
            atom_crds_xtal.append(
                conf_pos[atom.GetIdx()]
                )
    atom_crds_xtal = np.array(atom_crds_xtal)
    atom_radii_xtal = np.array(atom_radii_xtal)
    atom_radii_xtal += radius

    grid = list()
    for a in np.linspace(0., 1., 50, True):
        for b in np.linspace(0., 1., 50, True):
            for c in np.linspace(0., 1., 50, True):
                frac  = np.array([a,b,c], dtype=float)
                ortho = cell.orthogonalize(
                    gemmi.Fractional(*frac)
                            ).tolist()
                dists      = distance.cdist(atom_crds_xtal, [ortho])
                is_outside = np.all(dists[:,0] > atom_radii_xtal)
                if is_outside:
                    grid.append(ortho)

    overlap = True
    grid    = np.array(grid)
    print(f"Found {grid.shape[0]} / {50**3} valid grid points.")
    print("Scanning overlap...")
    while overlap:
        mask           = np.arange(grid.shape[0], dtype=int)
        mask_selection = np.random.choice(
            mask,
            size=N_per_unitcell,
            replace=False
            )
        grid_selection = grid[mask_selection]
        grid_selection_query = np.copy(grid_selection).tolist()
        for crd in grid_selection:
            frac = cell.fractionalize(
                gemmi.Position(
                    *crd
                    )
                ).tolist()
            frac  = np.array(frac)
            for a in [-1.,0.,1.]:
                for b in [-1.,0.,1.]:
                    for c in [-1.,0.,1.]:
                        if (a==0) & (b==0) & (c==0):
                            continue
                        frac += [a,b,c]
                        ortho = cell.orthogonalize(
                            gemmi.Fractional(*frac)
                                    ).tolist()
                        grid_selection_query.append(ortho)
                        frac -= [a,b,c]

        grid_selection_query = np.array(grid_selection_query)
        dists = distance.cdist(grid_selection_query, grid_selection_query)
        np.fill_diagonal(dists, np.inf)
        min_dist = np.min(dists)
        if min_dist > 2.*radius:
            overlap = False

    import copy
    mol_list_new = copy.deepcopy(mol_list)
    for crds in grid_selection:
        mol_cp         = copy.deepcopy(mol)
        conformer      = mol_cp.GetConformer()
        atom_crds_mol  = conformer.GetPositions()
        trans          = crds - np.mean(atom_crds_mol, axis=0)
        atom_crds_mol += trans
        for atm_idx in range(mol_cp.GetNumAtoms()):
            conformer.SetAtomPosition(
                atm_idx,
                Point3D(*atom_crds_mol[atm_idx])
            )
        mol_list_new.append(mol_cp)

    return mol_list_new

def make_P1(
    cell,
    atom_crds_ortho,
    atom_num,
    addhs=False,
    use_openeye=False):

    """
    Generate the P1 cell. Return tuple with atomic coordinates (in Ang) and
    atomic numbers of all atoms in P1 cell.
    """

    import networkx as nx
    from rdkit.Chem import AllChem as Chem
    from rdkit.Geometry import Point3D
    import copy

    _atom_crds_ortho = list()
    _atom_num = list()
    N_atoms = len(atom_num)
    for i in range(N_atoms):
        if addhs:
            if atom_num[i] != 1:
                _atom_crds_ortho.append(copy.copy(atom_crds_ortho[i]))
                _atom_num.append(copy.copy(atom_num[i]))
        else:
            _atom_crds_ortho.append(copy.copy(atom_crds_ortho[i]))
            _atom_num.append(copy.copy(atom_num[i]))

    atom_crds_ortho = _atom_crds_ortho
    atom_num = _atom_num

    atom_crds_ortho = np.array(atom_crds_ortho, dtype=float)
    atom_num = np.array(atom_num, dtype=int)
    nonoverlapping_idxs = get_nonoverlapping_atoms(atom_crds_ortho, filter_overlapping=True)
    atom_crds_ortho = atom_crds_ortho[nonoverlapping_idxs].tolist()
    atom_num = atom_num[nonoverlapping_idxs].tolist()
    N_atoms = len(atom_num)

    found_bond = True
    ### Terminate if we haven't found any bonds.
    while found_bond:
        ### Update the frac coordinates
        atom_crds_frac = list()
        for atm_idx in range(N_atoms):
            frac = cell.fractionalize(
                gemmi.Position(
                    *atom_crds_ortho[atm_idx]
                    )
                )
            atom_crds_frac.append(frac.tolist())

        ### Get the "pre-molecules" (adjacency matrix with not much chemistry)
        acmatrix_new, _ = xyz2mol.xyz2AC(
            atom_num,
            atom_crds_ortho,
            0,
            )
        acmatrix_best = acmatrix_new

        ### Find the disconnected graphs from the adjacency matrix
        G           = nx.convert_matrix.from_numpy_matrix(acmatrix_new)
        G_node_list = list(nx.connected_components(G))
        ### Translate molecules to neighboring unit cells in + direction
        ### and check if we can form new bonds. If yes, update `atom_crds_ortho`
        found_bond = False
        for g in G_node_list:
            for a in [0,-1]:
                for b in [0,-1]:
                    for c in [0,-1]:
                        atom_crds_ortho_cp = copy.deepcopy(atom_crds_ortho)
                        for atm_idx in g:
                            frac = copy.deepcopy(atom_crds_frac[atm_idx])
                            frac[0] += a
                            frac[1] += b
                            frac[2] += c
                            ortho = cell.orthogonalize(
                                gemmi.Fractional(*frac)
                            ).tolist()
                            if not ortho in atom_crds_ortho_cp:
                                atom_crds_ortho_cp[atm_idx] = ortho
                        acmatrix_new, _ = xyz2mol.xyz2AC(
                            atom_num,
                            atom_crds_ortho,
                            0,
                        )
                        if not np.all(acmatrix_new == acmatrix_best) and np.sum(acmatrix_new) >= np.sum(acmatrix_best):
                            nonoverlapping_idxs = get_nonoverlapping_atoms(atom_crds_ortho_cp)
                            if nonoverlapping_idxs.size == N_atoms:
                                atom_crds_ortho = copy.deepcopy(atom_crds_ortho_cp)
                                acmatrix_best = acmatrix_new
                                found_bond = True

    acmatrix, _ = xyz2mol.xyz2AC(
        atom_num,
        atom_crds_ortho,
        0,
        )
    G           = nx.convert_matrix.from_numpy_matrix(acmatrix)
    G_node_list = list(nx.connected_components(G))
    atom_num    = np.array(atom_num, dtype=int)
    atom_crds_ortho = np.array(atom_crds_ortho)
    mol_list    = list()
    for g in G_node_list:
        g = list(g)
        _, mol = xyz2mol.xyz2AC(
                atom_num[g].tolist(),
                atom_crds_ortho[g].tolist(),
                0)
        mol_list.append(mol)

    N_mol           = len(mol_list)
    atom_crds_ortho = list()
    atom_num        = list()
    mol_list_new    = list()
    for mol_idx in range(N_mol):
        mol  = mol_list[mol_idx]
        conf = mol.GetConformer()
        conf_pos = conf.GetPositions()
        frac_crds = list()
        for pos in conf_pos:
            frac_crds.append(
                cell.fractionalize(
                    gemmi.Position(
                        *pos
                    )
                ).tolist()
            )
        frac_crds   = np.array(frac_crds)
        valid_atoms = np.where(
            (frac_crds[:,0] > 0.) * (frac_crds[:,0] < 1.) *\
            (frac_crds[:,1] > 0.) * (frac_crds[:,1] < 1.) *\
            (frac_crds[:,2] > 0.) * (frac_crds[:,2] < 1.)
        )[0]
        ### If no atom is in uc, bring the molecule into unit cell
        if valid_atoms.size == 0:
            is_inside = False
            for a in [0,-1,1]:
                for b in [0,-1,1]:
                    for c in [0,-1,1]:
                        frac_crds += [a,b,c]
                        valid_atoms = np.where(
                            (frac_crds[:,0] > 0.) * (frac_crds[:,0] < 1.) *\
                            (frac_crds[:,1] > 0.) * (frac_crds[:,1] < 1.) *\
                            (frac_crds[:,2] > 0.) * (frac_crds[:,2] < 1.)
                        )[0]
                        if valid_atoms.size != 0:
                            is_inside = True
                        else:
                            frac_crds -= [a,b,c]
                        if is_inside:
                            break
                    if is_inside:
                        break
                if is_inside:
                    break
            if not is_inside:
                continue

        frac_crds_query = np.copy(frac_crds[valid_atoms[0]])
        ### Check for overlap
        overlap = False
        for a in [0,-1,1]:
            for b in [0,-1,1]:
                for c in [0,-1,1]:
                    frac_crds_query += [a,b,c]
                    ortho = cell.orthogonalize(
                        gemmi.Fractional(*frac_crds_query)
                    ).tolist()
                    if len(atom_crds_ortho) > 0:
                        dists = distance.cdist([ortho], atom_crds_ortho)
                        valids = np.where(dists < 0.01)[0]
                        if valids.size > 0:
                            overlap = True
                    frac_crds_query -= [a,b,c]
        if not overlap:
            for atm_idx, frac in enumerate(frac_crds):
                ortho = cell.orthogonalize(
                    gemmi.Fractional(*frac)
                ).tolist()
                atom_crds_ortho.append(ortho)
                rd_atom = mol.GetAtomWithIdx(atm_idx)
                atom_num.append(rd_atom.GetAtomicNum())
                conf.SetAtomPosition(
                    atm_idx,
                    Point3D(*ortho)
                )
            mol_list_new.append(mol)

    mol_list = list()
    if addhs:
        if use_openeye:
            import warnings
            warnings.warn("With addhs=True, we automatically set use_openeye=True.")
        from openeye import oechem
        from openeye import oequacpac
        from xtalmdscripts.supercellbuilding.oe_utils import rdmol_from_oemol
        from xtalmdscripts.supercellbuilding.oe_utils import oemol_from_rdmol

        count = 0
        for mol in mol_list_new:
            oemol = oechem.OEMol()
            oemol.SetDimension(3)
            conf_pos = mol.GetConformer(0).GetPositions()
            crds = list()
            for atm_idx in range(mol.GetNumAtoms()):
                atom = mol.GetAtomWithIdx(atm_idx)
                oemol.NewAtom(int(atom.GetAtomicNum()))
                crds.extend(conf_pos[atm_idx])
            oemol.SetCoords(crds)

            oechem.OEAssignAromaticFlags(oemol)
            oechem.OEDetermineConnectivity(oemol)
            oechem.OEFindRingAtomsAndBonds(oemol)
            oechem.OEPerceiveBondOrders(oemol)
            oechem.OE3DToInternalStereo(oemol)
            oechem.OEPerceiveChiral(oemol)
            oechem.OEAssignImplicitHydrogens(oemol)
            oechem.OEAssignFormalCharges(oemol)

            oequacpac.OERemoveFormalCharge(oemol)
            oechem.OEAddExplicitHydrogens(oemol)

            oechem.OEAssignAromaticFlags(oemol)
            mol = rdmol_from_oemol(oemol)
            Chem.AssignStereochemistryFrom3D(mol)
            mol_list.append(mol)

            #with open(f"./test_{count}.pdb", "w") as fopen:
            #    fopen.write(Chem.MolToPDBBlock(mol))
            count += 1

    else:
        if use_openeye:
            from openeye import oechem
            from openeye import oequacpac
            from xtalmdscripts.supercellbuilding.oe_utils import rdmol_from_oemol
            from xtalmdscripts.supercellbuilding.oe_utils import oemol_from_rdmol

            count = 0
            for mol in mol_list_new:
                oemol = oechem.OEMol()
                oemol.SetDimension(3)
                conf_pos = mol.GetConformer(0).GetPositions()
                crds = list()
                for atm_idx in range(mol.GetNumAtoms()):
                    atom = mol.GetAtomWithIdx(atm_idx)
                    oemol.NewAtom(int(atom.GetAtomicNum()))
                    crds.extend(conf_pos[atm_idx])
                oemol.SetCoords(crds)

                oechem.OEDetermineConnectivity(oemol)
                oechem.OEFindRingAtomsAndBonds(oemol)
                oechem.OEPerceiveBondOrders(oemol)
                oechem.OE3DToInternalStereo(oemol)
                oechem.OEPerceiveChiral(oemol)
                oechem.OEAssignFormalCharges(oemol)

                oechem.OEAssignAromaticFlags(oemol)
                mol = rdmol_from_oemol(oemol)
                Chem.AssignStereochemistryFrom3D(mol)
                mol_list.append(mol)

                #with open(f"./test_{count}.pdb", "w") as fopen:
                #    fopen.write(Chem.MolToPDBBlock(mol))
                count += 1

        else:
            acmatrix, _ = xyz2mol.xyz2AC(
                atom_num,
                atom_crds_ortho,
                0,
                )
            G           = nx.convert_matrix.from_numpy_matrix(acmatrix)
            G_node_list = list(nx.connected_components(G))
            atom_num    = np.array(atom_num, dtype=int)
            atom_crds_ortho = np.array(atom_crds_ortho)
            for g in G_node_list:
                g = list(g)
                mol = Chem.GetMolFrags(
                    xyz2mol.xyz2mol(
                        atom_num[g].tolist(),
                        atom_crds_ortho[g].tolist(),
                        charge=0)[0],
                    asMols=True
                )[0]
                mol_list.append(mol)

    #strc_write               = gemmi.Structure()
    #strc_write.spacegroup_hm = "P1"
    #strc_write.cell          = cell
    #with open("./make_p1_test.pdb", "w") as fopen:
    #    fopen.write(get_pdb_block(mol_list, strc_write))

    return mol_list

def make_supercell(
    cell,
    mol_list,
    a_min, a_max,
    b_min, b_max,
    c_min, c_max):

    """
    Generate supercell based specified parameters. Assuming that `atom_crds_ortho`
    and `atom_num` are for P1 cell. See method `make_P1`. Returns list of rdkit mol
    objects with all molecules in supercell, list containing an int that is unique
    among the molecules in a unit cell, list containing the frac coordinates of the
    unit cell origins in the basis of the supercell.
    """

    a_replicate = np.arange(a_min,a_max+1, dtype=int)
    b_replicate = np.arange(b_min,b_max+1, dtype=int)
    c_replicate = np.arange(c_min,c_max+1, dtype=int)

    N_mol               = len(mol_list)
    replicated_mol_list = list()
    mol_identifies      = list()
    unitcell_in_supercell_fracs = list()
    for a in a_replicate:
        for b in b_replicate:
            for c in c_replicate:
                for mol_idx in range(N_mol):
                    mol      = copy.deepcopy(mol_list[mol_idx])
                    conf     = mol.GetConformer(0)
                    conf_pos = conf.GetPositions()
                    N_atoms  = mol.GetNumAtoms()

                    for atom_idx in range(N_atoms):
                        pos      = conf_pos[atom_idx]
                        frac_pos = cell.fractionalize(
                            gemmi.Position(*pos)
                        )
                        frac_pos.x += a
                        frac_pos.y += b
                        frac_pos.z += c

                        conf_pos_abc = cell.orthogonalize(frac_pos).tolist()

                        conf.SetAtomPosition(
                            atom_idx,
                            Point3D(*conf_pos_abc)
                        )

                    replicated_mol_list.append(mol)
                    conf = mol.GetConformer(0)
                    mol_identifies.append(mol_idx)
                    unitcell_in_supercell_fracs.append([a,b,c])

    return replicated_mol_list, mol_identifies, unitcell_in_supercell_fracs

def clean_names(mol_list):

    """
    Uniquify atom and residue names. Equalize residue names for
    chemically equal residues.
    """

    import copy

    mol_list_new = list()
    N_mol = len(mol_list)
    for mol_idx in range(N_mol):
        mol = copy.deepcopy(mol_list[mol_idx])

        atom_counts_dict = dict()
        for atom in mol.GetAtoms():
            mi  = Chem.AtomPDBResidueInfo()
            mi.SetIsHeteroAtom(True)
            mi.SetResidueName(f'M{mol_idx}'.ljust(3))
            mi.SetResidueNumber(mol_idx + 1)
            mi.SetOccupancy(1.0)
            mi.SetTempFactor(0.0)
            atomic_num = atom.GetAtomicNum()
            atomic_ele = atom.GetSymbol()
            if not atomic_num in atom_counts_dict:
                atom_counts_dict[atomic_num] = 1
            else:
                atom_counts_dict[atomic_num] += 1
            mi.SetName(
                f"{atomic_ele}{atom_counts_dict[atomic_num]}".ljust(4)
                )
            atom.SetMonomerInfo(mi)

        mol_list_new.append(mol)

    return mol_list_new

def get_unique_mapping(
    mol_list,
    stereochemistry=True
    ):

    """
    Get unique mapping dict and list of unique rdkit mol objects in list of rdkit mol objects.
    if `stereochemistry=True`, the mapping will honor stereochemistry.
    """

    N_mol = len(mol_list)

    smiles_list = [Chem.MolToSmiles(mol, isomericSmiles=stereochemistry) for mol in mol_list]
    smiles_list_unique = set(smiles_list)
    smiles_list_unique = list(smiles_list_unique)

    rdmol_list_unique  = list()
    unique_mapping     = dict()
    for smiles_unique_idx, smiles_unique in enumerate(smiles_list_unique):
        found_unique = False
        for mol_idx in range(N_mol):
            mol    = mol_list[mol_idx]
            smiles = smiles_list[mol_idx]
            if smiles == smiles_unique:
                if not found_unique:
                    rdmol_list_unique.append(mol)
                    found_unique = True
                    unique_mapping[mol_idx] = smiles_unique_idx
                else:
                    unique_mapping[mol_idx] = smiles_unique_idx

    assert len(unique_mapping) == N_mol

    return unique_mapping, rdmol_list_unique

def equalize_rdmols(
    mol_list,
    stereochemistry=True
    ):

    """
    Get list of rdkit mol objects in which all chemically idential mol objects
    have identical topology and pdb monomer info. Only difference are coordinates.
    If `stereochemistry=True` it will honor stereochemistry.
    """

    import copy

    unique_mapping, rdmol_list_unique = get_unique_mapping(mol_list, stereochemistry)

    mol_list_new = copy.deepcopy(mol_list)
    for mol_idx in unique_mapping:
        if mol_idx == unique_mapping[mol_idx]:
            mol_info = copy.deepcopy(rdmol_list_unique[unique_mapping[mol_idx]])
            for mol_info_atm_idx in range(mol_info.GetNumAtoms()):
                mi = mol_info.GetAtomWithIdx(mol_info_atm_idx).GetMonomerInfo()
                mi.SetResidueName(f'M{unique_mapping[mol_idx]}'.ljust(3))
                mi.SetResidueNumber(mol_idx + 1)
                mol_info.GetAtomWithIdx(mol_info_atm_idx).SetMonomerInfo(mi)
        else:
            ### This is the molecule that holds the correct coordinates
            ### and pdb monomer info.
            mol_crds = copy.deepcopy(mol_list[mol_idx])
            ### This is the molecule that holds the correct names, ordering, etc...
            mol_info = copy.deepcopy(rdmol_list_unique[unique_mapping[mol_idx]])
            match   = mol_crds.GetSubstructMatch(mol_info, useChirality=stereochemistry)

            conf_pos_crds = mol_crds.GetConformer(0).GetPositions()
            conf_info     = mol_info.GetConformer(0)
            for mol_info_atm_idx, mol_crds_atm_idx in enumerate(match):
                pos = conf_pos_crds[mol_crds_atm_idx]
                conf_info.SetAtomPosition(
                    mol_info_atm_idx,
                    Point3D(*pos)
                )
                ### Note, we cannot `copy.copy(mi_original)`
                ### or `copy.copy(mol_target.GetAtomWithIdx(atm_idx))`
                mi = mol_info.GetAtomWithIdx(mol_info_atm_idx).GetMonomerInfo()
                mi.SetResidueName(f'M{unique_mapping[mol_idx]}'.ljust(3))
                mi.SetResidueNumber(mol_idx + 1)
                mol_info.GetAtomWithIdx(mol_info_atm_idx).SetMonomerInfo(mi)

        mol_list_new[mol_idx] = mol_info

    return mol_list_new

def generate_replicated_mol_list(
    cell,
    atom_crds_ortho,
    atom_num,
    a_min_max,
    b_min_max,
    c_min_max,
    addhs=False,
    protonate_unitcell=True,
    addwater=0,
    N_iterations_protonation=0,
    use_openeye=False,
    ):

    """
    Generate rdkit mol object list for molecules in supercell. supercell is generated
    according input parameters.
    """

    mol_list = make_P1(cell, atom_crds_ortho, atom_num, addhs, use_openeye)
    if N_iterations_protonation > 0:
        mol_list = assign_protonation_states(
            cell=cell,
            mol_list=mol_list,
            N_iterations=N_iterations_protonation
            )
    if addwater > 0:
        mol_list = random_fill(
            cell,
            mol_list,
            N_per_unitcell=addwater,
            radius=0.5,
            smiles="O"
            )
    mol_list = clean_names(mol_list)

    replicated_mol_list, mol_identifies, unitcell_in_supercell_fracs = make_supercell(
        cell,
        mol_list,
        a_min_max[0], a_min_max[1],
        b_min_max[0], b_min_max[1],
        c_min_max[0], c_min_max[1],
        )

    replicated_mol_list = equalize_rdmols(replicated_mol_list)

    return replicated_mol_list, mol_identifies, unitcell_in_supercell_fracs

def get_pdb_block(
    replicated_mol_list,
    strc_write):

    """
    Get pdb block as str. strc_write is gemmi structure object and must reflect
    the dimensions of the supercell.
    """

    ### Combine all rdmols in a single big rdmol
    N_mol   = len(replicated_mol_list)
    mol_new = Chem.Mol()
    for mol_idx in range(N_mol):
        mol = copy.deepcopy(replicated_mol_list[mol_idx])
        mol_new = Chem.CombineMols(mol_new, mol)
    header = strc_write.make_pdb_headers()

    ### With the flavor options, one can control what is written
    ### to the pdb block.
    ###
    ### flavor: (optional)
    ### flavor & 1 : Write MODEL/ENDMDL lines around each record
    ### flavor & 2 : Don’t write any CONECT records
    ### flavor & 4 : Write CONECT records in both directions
    ### flavor & 8 : Don’t use multiple CONECTs to encode bond order
    ### flavor & 16 : Write MASTER record
    ### flavor & 32 : Write TER record

    crds_block = Chem.MolToPDBBlock(mol_new, flavor=8|32)
    pdb_block  = header + crds_block

    return pdb_block

def get_pdb_str(
    replicated_mol_list,
    strc,
    a_min_max,
    b_min_max,
    c_min_max):

    """
    Get full pdb file as str.
    """

    import gemmi

    ### Write pdb file
    ### ==============
    a_len = np.max(a_min_max) - np.min(a_min_max) + 1.
    b_len = np.max(b_min_max) - np.min(b_min_max) + 1.
    c_len = np.max(c_min_max) - np.min(c_min_max) + 1.

    strc_write               = gemmi.Structure()
    strc_write.spacegroup_hm = strc.spacegroup_hm
    strc_write.cell          = gemmi.UnitCell(
        strc.cell.a * a_len,
        strc.cell.b * b_len,
        strc.cell.c * c_len,
        strc.cell.alpha,
        strc.cell.beta,
        strc.cell.gamma
    )

    pdb_block = get_pdb_block(replicated_mol_list, strc_write)

    return pdb_block

def parse_cif(
    cif_path,
    use_symmetry_operations=False
    ):

    """
    Parse cif file as gemmis structure object.
    """

    import gemmi

    doc  = gemmi.cif.read(cif_path)[0]
    strc = gemmi.make_small_structure_from_block(doc)

    ### finding it by number is much better
    ### Sometimes the HM name cannot be found by gemmi.
    table_number = -1
    for item in doc:
        if item.pair == None:
            continue
        key, value = item.pair
        if "_symmetry_Int_Tables_number".lower() == key.lower():
            table_number=int(value)
            break
    if table_number > -1:
        strc.spacegroup_hm = gemmi.find_spacegroup_by_number(table_number).hm

    atom_crds_ortho = list()
    atom_num = list()
    if use_symmetry_operations:
        op_list  = doc.find_values('_symmetry_equiv_pos_as_xyz')
        gops     = gemmi.GroupOps([gemmi.Op(o) for o in op_list])
        for site in strc.sites:
            for op in gops:
                pos_frac = op.apply_to_xyz(site.fract.tolist())
                pos = strc.cell.orthogonalize(
                    gemmi.Fractional(
                        *pos_frac
                        )
                    ).tolist()
                atom_crds_ortho.append(pos)
                atom_num.append(site.element.atomic_number)
    else:
        for site in strc.get_all_unit_cell_sites():
            pos = strc.cell.orthogonalize(site.fract)
            atom_crds_ortho.append([pos.x, pos.y, pos.z])
            atom_num.append(site.element.atomic_number)

    return strc, atom_crds_ortho, atom_num

def get_supercell_info_str(
    mol_identifies,
    unitcell_in_supercell_fracs
    ):

    """
    Takes list of in unitcell molecule identifiers and unitcell in supercell
    frac coordinates. See output generated by method `generate_replicated_mol_list`.
    Returns csv formatted info string.
    """

    info_str  = "#Mol_idx,"
    info_str += "mol_in_unitcell,"
    info_str += "unitcell_in_supercell_a,"
    info_str += "unitcell_in_supercell_b,"
    info_str += "unitcell_in_supercell_c\n"

    N_mols = len(mol_identifies)
    for mol_idx in range(N_mols):
        info_str += f"{mol_idx:d},"
        info_str += f"{mol_identifies[mol_idx]:d},"
        info_str += f"{unitcell_in_supercell_fracs[mol_idx][0]:d},"
        info_str += f"{unitcell_in_supercell_fracs[mol_idx][1]:d},"
        info_str += f"{unitcell_in_supercell_fracs[mol_idx][2]:d}\n"

    return info_str

def get_replicated_mol_list_json(replicated_mol_list):

    """
    Returns rdkit json string of collapsed replicated mol_list.
    """

    from rdkit import Chem

    mol_combo = Chem.Mol()
    for mol in replicated_mol_list:
        mol_combo = Chem.CombineMols(mol_combo, mol)
    return Chem.MolToJSON(mol_combo)

def main():

    """
    Run the workflow.
    """

    import gemmi

    args = parse_arguments()
    strc, atom_crds_ortho, atom_num = parse_cif(args.input, args.use_symmetry_operations)

    if len(args.a_min_max) == 2:
        a_min_max = [int(args.a_min_max[0]), int(args.a_min_max[1])]
    elif len(args.a_min_max) == 1:
        uc_length_a = strc.cell.a * 0.1
        a_min_max = [0, np.ceil(args.a_min_max[0] / uc_length_a)]
    else:
        raise ValueError(
            "Argument a_min_max must be either pair (x,y) or single value (z)"
            )

    if len(args.b_min_max) == 2:
        b_min_max = [int(args.b_min_max[0]), int(args.b_min_max[1])]
    elif len(args.b_min_max) == 1:
        uc_length_b = strc.cell.b * 0.1
        b_min_max = [0, np.ceil(args.b_min_max[0] / uc_length_b)]
    else:
        raise ValueError(
            "Argument b_min_max must be either pair (x,y) or single value (z)"
            )

    if len(args.c_min_max) == 2:
        c_min_max = [int(args.c_min_max[0]), int(args.c_min_max[1])]
    elif len(args.c_min_max) == 1:
        uc_length_c = strc.cell.c * 0.1
        c_min_max = [0, np.ceil(args.c_min_max[0] / uc_length_c)]
    else:
        raise ValueError(
            "Argument c_min_max must be either pair (x,y) or single value (z)"
            )

    ### Build the supercell as a set of rdkit molecule objects
    ### ======================================================
    replicated_mol_list, mol_identifies, unitcell_in_supercell_fracs = generate_replicated_mol_list(
        cell=strc.cell,
        atom_crds_ortho=atom_crds_ortho,
        atom_num=atom_num,
        a_min_max=a_min_max,
        b_min_max=b_min_max,
        c_min_max=c_min_max,
        addhs=args.addhs,
        addwater=args.addwater,
        N_iterations_protonation=args.n_protonation_attempts,
        use_openeye=args.use_openeye
        )

    ### Write pdb file
    ### ==============
    strc_write               = gemmi.Structure()
    strc_write.spacegroup_hm = "P1"
    strc_write.cell          = strc.cell
    pdb_str = get_pdb_str(
        replicated_mol_list,
        strc_write,
        a_min_max,
        b_min_max,
        c_min_max
        )
    with open(f"{args.prefix}.pdb", "w") as fopen:
        fopen.write(pdb_str)
    with open(f"{args.prefix}.csv", "w") as fopen:
        info_str = get_supercell_info_str(
            mol_identifies,
            unitcell_in_supercell_fracs
            )
        fopen.write(info_str)

    with open(f"{args.prefix}.json", "w") as fopen:
        json_str = get_replicated_mol_list_json(replicated_mol_list)
        fopen.write(json_str)

    ### Generate list of unique smiles for unique
    ### molecules in UC
    ### =========================================
    mol_list = make_P1(strc.cell, atom_crds_ortho, atom_num, args.addhs, args.use_openeye)
    if args.addwater > 0:
        random_fill(
            strc.cell,
            mol_list,
            N_per_unitcell=args.addwater,
            radius=0.5,
            smiles="O"
            )

    unitcell_mol_list, _, _ = make_supercell(
        strc.cell,
        mol_list,
        0,0,
        0,0,
        0,0,
        )
    from rdkit.Chem import Descriptors
    unitcell_weight = 0.
    smiles_list     = list()
    for mol in unitcell_mol_list:
        unitcell_weight += Descriptors.MolWt(mol)
        smiles_list.append(Chem.MolToSmiles(mol, isomericSmiles=True))

    smiles_list = set(smiles_list)
    smiles_list = list(smiles_list)

    ### Output final summary
    ### ====================
    a_len = np.max(a_min_max) - np.min(a_min_max) + 1.
    b_len = np.max(b_min_max) - np.min(b_min_max) + 1.
    c_len = np.max(c_min_max) - np.min(c_min_max) + 1.

    import gemmi
    doc           = gemmi.cif.read(args.input)[0]
    cif_info_dict = {
        "temperature"  : "Not found",
        "cell_setting" : "Not found",
        "space_group"  : "Not found",
        "density"      : "Not found"
    }
    for item in doc:
        if item.pair == None:
            continue
        key, value = item.pair
        if "_diffrn_ambient_temperature".lower() == key.lower():
            cif_info_dict["temperature"] = value
        elif "_symmetry_cell_setting".lower() == key.lower():
            cif_info_dict["cell_setting"] = value
        elif "_symmetry_space_group_name_H-M".lower() == key.lower():
            cif_info_dict["space_group"] = value
        elif "_exptl_crystal_density_diffrn".lower() == key.lower():
            cif_info_dict["density"] = value

    print(f"""
Summary:
========
Expt:
-----
Temperature [K]               : {cif_info_dict['temperature']},
Cell Setting                  : {cif_info_dict['cell_setting']},
Space Group H-M               : {cif_info_dict['space_group']},
Density [g/cm3]               : {cif_info_dict['density']},

Supercell:
----------
Total number of molecules     : {len(replicated_mol_list)},
Total Length edge a [Ang]     : {strc.cell.a * a_len:4.2f},
Total Length edge b [Ang]     : {strc.cell.b * b_len:4.2f},
Total Length edge c [Ang]     : {strc.cell.c * c_len:4.2f},
Cell angle alpha [deg]        : {strc.cell.alpha:4.2f},
Cell angle beta  [deg]        : {strc.cell.beta:4.2f},
Cell angle gamma [deg]        : {strc.cell.gamma:4.2f},
Total Volume supercell [Ang3] : {strc.cell.volume * a_len * b_len * c_len:4.2f}
Density [g/cm3]               : {unitcell_weight / strc.cell.volume * 1.6605:4.2f}
SMILES for molecules in UC    : {" ".join(smiles_list)}
"""
### 1.6605 conversion g/mol/Ang^3 to g/cm^3
)

    try:
        diff = (unitcell_weight / strc.cell.volume * 1.6605) - float(cif_info_dict["density"])
        if abs(diff) > 0.1:
            warnings.warn(f"Density difference {diff}. Check structure.")
    except:
        pass

def entry_point():

    main()

if __name__ == "__main__":

    entry_point()