import numpy as np
import os,xml.etree.ElementTree as ET
import cv2
import matplotlib.pyplot as plt
from glob import glob
from concurrent.futures import ThreadPoolExecutor
import tensorflow as tf
from PIL import Image,ImageDraw
from math import *
anchor_scale = 16
#
IOU_NEGATIVE = 0.3
IOU_POSITIVE = 0.7
IOU_SELECT = 0.7

RPN_POSITIVE_NUM = 150
RPN_TOTAL_NUM = 300

# bgr  can find from  here https://github.com/fchollet/deep-learning-models/blob/master/imagenet_utils.py
IMAGE_MEAN = [123.68, 116.779, 103.939]

DEBUG = True
def rotate(img_path,degree=90):

    img=cv2.imread(img_path)
    height, width = img.shape[:2]

    # 旋转后的尺寸
    heightNew = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))  # 这个公式参考之前内容
    widthNew = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))

    matRotation = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)

    matRotation[0, 2] += (widthNew - width) / 2  # 因为旋转之后,坐标系原点是新图像的左上角,所以需要根据原图做转化
    matRotation[1, 2] += (heightNew - height) / 2

    imgRotation = cv2.warpAffine(img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))

    return Image.fromarray(imgRotation)

def drawRect(gtboxes,img):
    draw=ImageDraw.Draw(img)
    for i in gtboxes:
        draw.rectangle((i[0],i[1],i[2],i[3]),outline='red')
    img.show()

def readxml(path):
    gtboxes=[]
    tree=ET.parse(path)
    root=tree.getroot()
    image_name = tree.find('filename').text
    for obj in root.iter('object'):
        try:
            bnd=obj.find('bndbox')
            xmin=int(bnd.find('xmin').text)
            ymin=int(bnd.find('ymin').text)
            xmax=int(bnd.find('xmax').text)
            ymax=int(bnd.find('ymax').text)
            gtboxes.append((xmin,ymin,xmax,ymax))
        except Exception as e:
            raise Exception(e+" when {} loading xml file, there are errors".format(path))
    return np.array(gtboxes),image_name

def gen_anchor(featuresize, scale):
    """
    gen base anchor from feature map [HXW][10][4]
    reshape  [HXW][10][4] to [HXWX10][4]
    生成的锚框是相对于原图的，即原图中每16像素就有10个锚框
    """

    heights = [11, 16, 23, 33, 48, 68, 97, 139, 198, 283]
    widths = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16]

    # gen k=9 anchor size (h,w)
    heights = np.array(heights).reshape(len(heights), 1)
    widths = np.array(widths).reshape(len(widths), 1)

    # 锚框大小为16像素
    base_anchor = np.array([0, 0, 15, 15])
    # center x,y
    xt = (base_anchor[0] + base_anchor[2]) * 0.5
    yt = (base_anchor[1] + base_anchor[3]) * 0.5

    # x1 y1 x2 y2
    x1 = xt - widths * 0.5
    y1 = yt - heights * 0.5
    x2 = xt + widths * 0.5
    y2 = yt + heights * 0.5

    # 一组十个锚框
    base_anchor = np.hstack((x1, y1, x2, y2))

    h, w = featuresize
    shift_x = np.arange(0, w) * scale
    shift_y = np.arange(0, h) * scale
    # apply shift
    anchor = []
    for i in shift_y:
        for j in shift_x:
            anchor.append(base_anchor + [j, i, j, i])
    return np.array(anchor).reshape((-1, 4))

def cal_iou(box1, box1_area, boxes2, boxes2_area):
    """
    box1 [x1,y1,x2,y2]
    boxes2 [Msample,x1,y1,x2,y2]
    """
    x1 = np.maximum(box1[0], boxes2[:, 0])
    x2 = np.minimum(box1[2], boxes2[:, 2])
    y1 = np.maximum(box1[1], boxes2[:, 1])
    y2 = np.minimum(box1[3], boxes2[:, 3])

    intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0)
    iou = intersection / (box1_area + boxes2_area[:] - intersection[:])
    return iou

def cal_overlaps(boxes1, boxes2):
    """
    boxes1 [Nsample,x1,y1,x2,y2]  anchor
    boxes2 [Msample,x1,y1,x2,y2]  grouth-box

    """
    area1 = (boxes1[:, 0] - boxes1[:, 2]) * (boxes1[:, 1] - boxes1[:, 3])  # (Nsample, 1)
    area2 = (boxes2[:, 0] - boxes2[:, 2]) * (boxes2[:, 1] - boxes2[:, 3])  # (Msample, 1)

    overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0]))  # (Nsample, Msample)

    # calculate the intersection of  boxes1(anchor) and boxes2(GT box)
    for i in range(boxes1.shape[0]):
        overlaps[i][:] = cal_iou(boxes1[i], area1[i], boxes2, area2)

    return overlaps

def bbox_transfrom(anchors, gtboxes):
    """
    anchors: (Nsample, 4)
    gtboxes: (Nsample, 4)
     compute relative predicted vertical coordinates Vc ,Vh
        with respect to the bounding box location of an anchor
    """
    Cy = (gtboxes[:, 1] + gtboxes[:, 3]) * 0.5  # (Nsample, )
    Cya = (anchors[:, 1] + anchors[:, 3]) * 0.5  # (Nsample, )
    h = gtboxes[:, 3] - gtboxes[:, 1] + 1.0  # (Nsample, )
    ha = anchors[:, 3] - anchors[:, 1] + 1.0  # (Nsample, )

    Vc = (Cy - Cya) / ha  # (Nsample, )
    Vh = np.log(h / ha)  # (Nsample, )

    ret = np.vstack((Vc, Vh))

    return ret.transpose()  # (Nsample, 2)

def bbox_transfor_inv(anchor, regr):
    """
    anchor: (NSample, 4)
    regr: (NSample, 2)
    根据锚框和偏移量反向得到GTBox
    """

    Cya = (anchor[:, 1] + anchor[:, 3]) * 0.5  # 锚框y中心点
    ha = anchor[:, 3] - anchor[:, 1] + 1

    Vcx = regr[..., 0]  # y中心点偏移
    Vhx = regr[..., 1]  # 高度偏移

    Cyx = Vcx * ha + Cya  # GTBox y中心点
    hx = np.exp(Vhx) * ha  # GTBox 高
    xt = (anchor[:, 0] + anchor[:, 2]) * 0.5  # 锚框x中心点

    x1 = xt - 16 * 0.5
    y1 = Cyx - hx * 0.5
    x2 = xt + 16 * 0.5
    y2 = Cyx + hx * 0.5
    bbox = np.vstack((x1, y1, x2, y2)).transpose()

    return bbox

def clip_box(bbox, im_shape):
    # x1 >= 0
    bbox[:, 0] = np.maximum(np.minimum(bbox[:, 0], im_shape[1] - 1), 0)
    # y1 >= 0
    bbox[:, 1] = np.maximum(np.minimum(bbox[:, 1], im_shape[0] - 1), 0)
    # x2 < im_shape[1]
    bbox[:, 2] = np.maximum(np.minimum(bbox[:, 2], im_shape[1] - 1), 0)
    # y2 < im_shape[0]
    bbox[:, 3] = np.maximum(np.minimum(bbox[:, 3], im_shape[0] - 1), 0)

    return bbox

def filter_bbox(bbox, minsize):
    ws = bbox[:, 2] - bbox[:, 0] + 1
    hs = bbox[:, 3] - bbox[:, 1] + 1
    keep = np.where((ws >= minsize) & (hs >= minsize))[0]
    return keep

def cal_rpn(imgsize, featuresize, scale, gtboxes):
    """
    gtboxes: (Msample, 4)
    """
    imgh, imgw = imgsize

    # gen base anchor
    base_anchor = gen_anchor(featuresize, scale)  # (Nsample, 4)

    # calculate iou
    overlaps = cal_overlaps(base_anchor, gtboxes)  # (Nsample, Msample)

    # init labels -1 don't care  0 is negative  1 is positive
    labels = np.empty(base_anchor.shape[0])
    labels.fill(-1)  # (Nsample,)

    # for each GT box corresponds to an anchor which has highest IOU
    gt_argmax_overlaps = overlaps.argmax(axis=0)  # (Msample, )

    # the anchor with the highest IOU overlap with a GT box
    anchor_argmax_overlaps = overlaps.argmax(axis=1)  # (Nsample, )
    anchor_max_overlaps = overlaps[range(overlaps.shape[0]), anchor_argmax_overlaps]  # (Nsample, )
    print('anchor max overlaps')
    for i in anchor_max_overlaps:
        if i>IOU_POSITIVE:
            print(i)

    # IOU > IOU_POSITIVE
    labels[anchor_max_overlaps > IOU_POSITIVE] = 1
    # IOU <IOU_NEGATIVE
    labels[anchor_max_overlaps < IOU_NEGATIVE] = 0
    # ensure that every GT box has at least one positive RPN region
    labels[gt_argmax_overlaps] = 1
    # print('内部的label',labels)
    # for i in labels:
    #     print(i)

    # only keep anchors inside the image
    outside_anchor = np.where(
        (base_anchor[:, 0] < 0) |
        (base_anchor[:, 1] < 0) |
        (base_anchor[:, 2] >= imgw) |
        (base_anchor[:, 3] >= imgh)
    )[0]
    labels[outside_anchor] = -1

    # 剔除掉多余的正负样例
    # subsample positive labels ,if greater than RPN_POSITIVE_NUM(default 128)
    fg_index = np.where(labels == 1)[0]
    if (len(fg_index) > RPN_POSITIVE_NUM):
        labels[np.random.choice(fg_index, len(fg_index) - RPN_POSITIVE_NUM, replace=False)] = -1

    # subsample negative labels
    bg_index = np.where(labels == 0)[0]
    num_bg = RPN_TOTAL_NUM - np.sum(labels == 1)
    if (len(bg_index) > num_bg):
        # print('bgindex:',len(bg_index),'num_bg',num_bg)
        labels[np.random.choice(bg_index, len(bg_index) - num_bg, replace=False)] = -1

    # calculate bbox targets
    # debug here
    bbox_targets = bbox_transfrom(base_anchor, gtboxes[anchor_argmax_overlaps, :])
    # bbox_targets=[]

    return [labels, bbox_targets], base_anchor

def get_session(gpu_fraction=0.6):
    '''''Assume that you have 6GB of GPU memory and want to allocate ~2GB'''

    num_threads = os.environ.get('OMP_NUM_THREADS')
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)

    if num_threads:
        return tf.Session(config=tf.ConfigProto(
            gpu_options=gpu_options, intra_op_parallelism_threads=num_threads, allow_soft_placement=True))
    else:
        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True))

class random_uniform_num():
    """
    uniform random
    """

    def __init__(self, total, start=0):
        self.total = total
        self.range = [i for i in range(total)]
        np.random.shuffle(self.range)
        self.index = start

    def get(self, batch_size):
        ret = []
        if self.index + batch_size > self.total:
            piece1 = self.range[self.index:]
            np.random.shuffle(self.range)
            self.index = (self.index + batch_size) - self.total
            piece2 = self.range[0:self.index]
            ret.extend(piece1)
            ret.extend(piece2)
        else:
            ret = self.range[self.index:self.index + batch_size]
            self.index = self.index + batch_size
        return ret

def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]  # Sort from high to low

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep

def gen_sample(xmlpath, imgpath, batchsize=1):
    """
    由于图像大小不定，批处理大小只能为1
    """

    # list xml
    xmlfiles = glob(xmlpath + '/*.xml')
    rd = random_uniform_num(len(xmlfiles))
    xmlfiles = np.array(xmlfiles)

    while True:
        shuf = xmlfiles[rd.get(1)]
        gtbox, imgfile = readxml(shuf[0])
        img = cv2.imread(imgpath + "\\" + imgfile)
        h, w, c = img.shape

        # clip image
        if np.random.randint(0, 100) > 50:
            img = img[:, ::-1, :]
            newx1 = w - gtbox[:, 2] - 1
            newx2 = w - gtbox[:, 0] - 1
            gtbox[:, 0] = newx1
            gtbox[:, 2] = newx2

        [cls, regr], _ = cal_rpn((h, w), (int(h / 16), int(w / 16)), 16, gtbox)
        # zero-center by mean pixel
        m_img = img - IMAGE_MEAN
        m_img = np.expand_dims(m_img, axis=0)

        regr = np.hstack([cls.reshape(cls.shape[0], 1), regr])

        #
        cls = np.expand_dims(cls, axis=0)
        cls = np.expand_dims(cls, axis=1)
        # regr = np.expand_dims(regr,axis=1)
        regr = np.expand_dims(regr, axis=0)

        yield m_img, {'rpn_class_reshape': cls, 'rpn_regress_reshape': regr}

def rpn_test():
    xmlpath = 'G:\data\VOCdevkit\VOC2007\Annotations\img_4375.xml'
    imgpath = 'G:\data\VOCdevkit\VOC2007\JPEGImages\img_4375.jpg'
    gtbox, _ = readxml(xmlpath)
    img = cv2.imread(imgpath)
    h, w, c = img.shape
    [cls, regr], base_anchor = cal_rpn((h, w), (int(h / 16), int(w / 16)), 16, gtbox)
    print(cls.shape)
    print(regr.shape)

    regr = np.expand_dims(regr, axis=0)
    inv_anchor = bbox_transfor_inv(base_anchor, regr)
    anchors = inv_anchor[cls == 1]
    anchors = anchors.astype(int)
    for i in anchors:
        cv2.rectangle(img, (i[0], i[1]), (i[2], i[3]), (255, 0, 0), 3)
    plt.imshow(img)

def threshold(coords, min_, max_):
    return np.maximum(np.minimum(coords, max_), min_)

def clip_boxes(boxes, im_shape):
    """
    Clip boxes to image boundaries.
    """
    boxes[:, 0::2]=threshold(boxes[:, 0::2], 0, im_shape[1]-1)
    boxes[:, 1::2]=threshold(boxes[:, 1::2], 0, im_shape[0]-1)
    return boxes

class Graph:
    def __init__(self, graph):
        self.graph=graph

    def sub_graphs_connected(self):
        sub_graphs=[]
        for index in range(self.graph.shape[0]):
            if not self.graph[:, index].any() and self.graph[index, :].any():
                v=index
                sub_graphs.append([v])
                while self.graph[v, :].any():
                    v=np.where(self.graph[v, :])[0][0]
                    sub_graphs[-1].append(v)
        return sub_graphs
class TextLineCfg:
    SCALE=600
    MAX_SCALE=1200
    TEXT_PROPOSALS_WIDTH=16
    MIN_NUM_PROPOSALS = 2
    MIN_RATIO=0.5
    LINE_MIN_SCORE=0.9
    MAX_HORIZONTAL_GAP=60
    TEXT_PROPOSALS_MIN_SCORE=0.7
    TEXT_PROPOSALS_NMS_THRESH=0.3
    MIN_V_OVERLAPS=0.6
    MIN_SIZE_SIM=0.6

class TextProposalGraphBuilder:
    """
        Build Text proposals into a graph.
    """
    def get_successions(self, index):
            box=self.text_proposals[index]
            results=[]
            for left in range(int(box[0])+1, min(int(box[0])+TextLineCfg.MAX_HORIZONTAL_GAP+1, self.im_size[1])):
                adj_box_indices=self.boxes_table[left]
                for adj_box_index in adj_box_indices:
                    if self.meet_v_iou(adj_box_index, index):
                        results.append(adj_box_index)
                if len(results)!=0:
                    return results
            return results

    def get_precursors(self, index):
        box=self.text_proposals[index]
        results=[]
        for left in range(int(box[0])-1, max(int(box[0]-TextLineCfg.MAX_HORIZONTAL_GAP), 0)-1, -1):
            adj_box_indices=self.boxes_table[left]
            for adj_box_index in adj_box_indices:
                if self.meet_v_iou(adj_box_index, index):
                    results.append(adj_box_index)
            if len(results)!=0:
                return results
        return results

    def is_succession_node(self, index, succession_index):
        precursors=self.get_precursors(succession_index)
        if self.scores[index]>=np.max(self.scores[precursors]):
            return True
        return False

    def meet_v_iou(self, index1, index2):
        def overlaps_v(index1, index2):
            h1=self.heights[index1]
            h2=self.heights[index2]
            y0=max(self.text_proposals[index2][1], self.text_proposals[index1][1])
            y1=min(self.text_proposals[index2][3], self.text_proposals[index1][3])
            return max(0, y1-y0+1)/min(h1, h2)

        def size_similarity(index1, index2):
            h1=self.heights[index1]
            h2=self.heights[index2]
            return min(h1, h2)/max(h1, h2)

        return overlaps_v(index1, index2)>=TextLineCfg.MIN_V_OVERLAPS and \
               size_similarity(index1, index2)>=TextLineCfg.MIN_SIZE_SIM

    def build_graph(self, text_proposals, scores, im_size):
        self.text_proposals=text_proposals
        self.scores=scores
        self.im_size=im_size
        self.heights=text_proposals[:, 3]-text_proposals[:, 1]+1

        boxes_table=[[] for _ in range(self.im_size[1])]
        for index, box in enumerate(text_proposals):
            boxes_table[int(box[0])].append(index)
        self.boxes_table=boxes_table

        graph=np.zeros((text_proposals.shape[0], text_proposals.shape[0]), np.bool)

        for index, box in enumerate(text_proposals):
            successions=self.get_successions(index)
            if len(successions)==0:
                continue
            succession_index=successions[np.argmax(scores[successions])]
            if self.is_succession_node(index, succession_index):
                # NOTE: a box can have multiple successions(precursors) if multiple successions(precursors)
                # have equal scores.
                graph[index, succession_index]=True
        return Graph(graph)
class TextProposalConnector:
    def __init__(self):
        self.graph_builder=TextProposalGraphBuilder()

    def group_text_proposals(self, text_proposals, scores, im_size):
        graph=self.graph_builder.build_graph(text_proposals, scores, im_size)
        return graph.sub_graphs_connected()

    def fit_y(self, X, Y, x1, x2):
        len(X)!=0
        # if X only include one point, the function will get line y=Y[0]
        if np.sum(X==X[0])==len(X):
            return Y[0], Y[0]
        p=np.poly1d(np.polyfit(X, Y, 1))
        return p(x1), p(x2)

    def get_text_lines(self, text_proposals, scores, im_size):
        # tp=text proposal
        tp_groups=self.group_text_proposals(text_proposals, scores, im_size)
        text_lines=np.zeros((len(tp_groups), 5), np.float32)

        for index, tp_indices in enumerate(tp_groups):
            text_line_boxes=text_proposals[list(tp_indices)]

            x0=np.min(text_line_boxes[:, 0])
            x1=np.max(text_line_boxes[:, 2])

            offset=(text_line_boxes[0, 2]-text_line_boxes[0, 0])*0.5

            lt_y, rt_y=self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 1], x0+offset, x1-offset)
            lb_y, rb_y=self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 3], x0+offset, x1-offset)

            # the score of a text line is the average score of the scores
            # of all text proposals contained in the text line
            score=scores[list(tp_indices)].sum()/float(len(tp_indices))

            text_lines[index, 0]=x0
            text_lines[index, 1]=min(lt_y, rt_y)
            text_lines[index, 2]=x1
            text_lines[index, 3]=max(lb_y, rb_y)
            text_lines[index, 4]=score

        text_lines=clip_boxes(text_lines, im_size)

        text_recs = np.zeros((len(text_lines), 9), np.float)
        index = 0
        for line in text_lines:
            xmin,ymin,xmax,ymax=line[0],line[1],line[2],line[3]
            text_recs[index, 0] = xmin
            text_recs[index, 1] = ymin
            text_recs[index, 2] = xmax
            text_recs[index, 3] = ymin
            text_recs[index, 4] = xmin
            text_recs[index, 5] = ymax
            text_recs[index, 6] = xmax
            text_recs[index, 7] = ymax
            text_recs[index, 8] = line[4]
            index = index + 1

        return text_recs

class TextProposalConnectorOriented:
    """
        Connect text proposals into text lines
    """

    def __init__(self):
        self.graph_builder = TextProposalGraphBuilder()

    def group_text_proposals(self, text_proposals, scores, im_size):
        graph = self.graph_builder.build_graph(text_proposals, scores, im_size)
        return graph.sub_graphs_connected()

    def fit_y(self, X, Y, x1, x2):
        len(X) != 0
        # if X only include one point, the function will get line y=Y[0]
        if np.sum(X == X[0]) == len(X):
            return Y[0], Y[0]
        p = np.poly1d(np.polyfit(X, Y, 1))
        return p(x1), p(x2)

    def get_text_lines(self, text_proposals, scores, im_size):
        """
        text_proposals:boxes

        """
        # tp=text proposal
        tp_groups = self.group_text_proposals(text_proposals, scores, im_size)  # 首先还是建图，获取到文本行由哪几个小框构成

        text_lines = np.zeros((len(tp_groups), 8), np.float32)

        for index, tp_indices in enumerate(tp_groups):
            text_line_boxes = text_proposals[list(tp_indices)]  # 每个文本行的全部小框
            X = (text_line_boxes[:, 0] + text_line_boxes[:, 2]) / 2  # 求每一个小框的中心x，y坐标
            Y = (text_line_boxes[:, 1] + text_line_boxes[:, 3]) / 2

            z1 = np.polyfit(X, Y, 1)  # 多项式拟合，根据之前求的中心店拟合一条直线（最小二乘）

            x0 = np.min(text_line_boxes[:, 0])  # 文本行x坐标最小值
            x1 = np.max(text_line_boxes[:, 2])  # 文本行x坐标最大值

            offset = (text_line_boxes[0, 2] - text_line_boxes[0, 0]) * 0.5  # 小框宽度的一半

            # 以全部小框的左上角这个点去拟合一条直线，然后计算一下文本行x坐标的极左极右对应的y坐标
            lt_y, rt_y = self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 1], x0 + offset, x1 - offset)
            # 以全部小框的左下角这个点去拟合一条直线，然后计算一下文本行x坐标的极左极右对应的y坐标
            lb_y, rb_y = self.fit_y(text_line_boxes[:, 0], text_line_boxes[:, 3], x0 + offset, x1 - offset)

            score = scores[list(tp_indices)].sum() / float(len(tp_indices))  # 求全部小框得分的均值作为文本行的均值

            text_lines[index, 0] = x0
            text_lines[index, 1] = min(lt_y, rt_y)  # 文本行上端 线段 的y坐标的小值
            text_lines[index, 2] = x1
            text_lines[index, 3] = max(lb_y, rb_y)  # 文本行下端 线段 的y坐标的大值
            text_lines[index, 4] = score  # 文本行得分
            text_lines[index, 5] = z1[0]  # 根据中心点拟合的直线的k，b
            text_lines[index, 6] = z1[1]
            height = np.mean((text_line_boxes[:, 3] - text_line_boxes[:, 1]))  # 小框平均高度
            text_lines[index, 7] = height + 2.5

        text_recs = np.zeros((len(text_lines), 9), np.float)
        index = 0
        for line in text_lines:
            b1 = line[6] - line[7] / 2  # 根据高度和文本行中心线，求取文本行上下两条线的b值
            b2 = line[6] + line[7] / 2
            x1 = line[0]
            y1 = line[5] * line[0] + b1  # 左上
            x2 = line[2]
            y2 = line[5] * line[2] + b1  # 右上
            x3 = line[0]
            y3 = line[5] * line[0] + b2  # 左下
            x4 = line[2]
            y4 = line[5] * line[2] + b2  # 右下
            disX = x2 - x1
            disY = y2 - y1
            width = np.sqrt(disX * disX + disY * disY)  # 文本行宽度

            fTmp0 = y3 - y1  # 文本行高度
            fTmp1 = fTmp0 * disY / width
            x = np.fabs(fTmp1 * disX / width)  # 做补偿
            y = np.fabs(fTmp1 * disY / width)
            if line[5] < 0:
                x1 -= x
                y1 += y
                x4 += x
                y4 -= y
            else:
                x2 += x
                y2 += y
                x3 -= x
                y3 -= y
            text_recs[index, 0] = x1
            text_recs[index, 1] = y1
            text_recs[index, 2] = x2
            text_recs[index, 3] = y2
            text_recs[index, 4] = x3
            text_recs[index, 5] = y3
            text_recs[index, 6] = x4
            text_recs[index, 7] = y4
            text_recs[index, 8] = line[4]
            index = index + 1

        return text_recs

if __name__=='__main__':
    xmlpath = 'D:\py_projects\data_new\data_new\data\\annotation\\img_calligraphy_00003_bg.xml'
    imgpath = 'D:\py_projects\data_new\data_new\data\\train_img\\img_calligraphy_00003_bg.jpg'
    gtboxes, _ = readxml(xmlpath)
    img = rotate(imgpath,90)
    print(gtboxes)
    ow,oh=img.size
    newbox = []
    for i in gtboxes:
        newbox.append([i[1], ow - i[2], i[3], ow - i[0]])
    gtboxes = np.array(newbox)
    image_shape = (512,256)
    x_scale = image_shape[0] / ow
    y_scale = image_shape[1] / oh
    newbox = []
    for i in range(len(gtboxes)):
        newbox.append(
            [gtboxes[i][0] * x_scale, gtboxes[i][1] * y_scale, gtboxes[i][2] * x_scale,
             gtboxes[i][3] * y_scale]
        )
    img = img.resize(image_shape, Image.ANTIALIAS)
    gtboxes = np.array(newbox)
    print(gtboxes)
    [cls,regr],base_anchor=cal_rpn((256,512),(256//16,512//16),16,gtboxes)
    print(cls)
    regr=np.expand_dims(regr,axis=0)
    anchor=bbox_transfor_inv(base_anchor,regr)
    anchor=anchor[cls==1]
    anchor=anchor.astype(int)
    print(anchor)
    drawRect(anchor,img)


from __future__ import division
from build.cioverlap import *
from qm.dftbplus.dftbplus import DFTBplus
from qm.dftbplus.dftbpar import spin_w, spin_w_lc, onsite_uu, onsite_ud, max_l
from misc import data, eps, eV_to_au, call_name
import os, shutil, re, textwrap
import numpy as np

class DFTB(DFTBplus):
    """ Class for (TD)DFTB method of DFTB+

        :param object molecule: Molecule object
        :param boolean l_scc: Include self-consistent charge (SCC) scheme
        :param double scc_tol: Stopping criteria for the SCC iterations
        :param integer scc_max_iter: Maximum number of SCC iterations
        :param boolean l_onsite: Include onsite correction to SCC term
        :param boolean l_range_sep: Include long-range corrected functional
        :param string lc_method: Algorithms for LC-DFTB
        :param boolean l_spin_pol: Include spin-polarisation scheme
        :param double unpaired_elec: Number of unpaired electrons
        :param string guess: Initial guess method for SCC scheme
        :param string guess_file: Initial guess file for charges
        :param double elec_temp: Electronic temperature in Fermi-Dirac scheme
        :param string mixer: Charge mixing method used in DFTB
        :param string ex_symmetry: Symmetry of excited state in TDDFTB
        :param double e_window: Energy window for TDDFTB. Increases efficiency of NACME calculation.
        :param integer,list k_point: Number of k-point samplings
        :param boolean l_periodic: Use periodicity in the calculations
        :param double,list cell_length: The lattice vectors of periodic unit cell
        :param string sk_path: Path for Slater-Koster files
        :param string install_path: Path for DFTB+ install directory
        :param boolean mpi: Use MPI parallelization
        :param string mpi_path: Path for MPI binary
        :param integer nthreads: Number of threads in the calculations
        :param string version: Version of DFTB+
    """
    def __init__(self, molecule, l_scc=True, scc_tol=1E-6, scc_max_iter=100, l_onsite=False, \
        l_range_sep=False, lc_method="MatrixBased", l_spin_pol=False, unpaired_elec=0., guess="h0", \
        guess_file="./charges.bin", elec_temp=0., mixer="Broyden", ex_symmetry="singlet", e_window=0., \
        k_point=[1, 1, 1], l_periodic=False, cell_length=[0., 0., 0., 0., 0., 0., 0., 0., 0.,], \
        sk_path="./", install_path="./", mpi=False, mpi_path="./", nthreads=1, version="20.1"):
        # Initialize DFTB+ common variables
        super(DFTB, self).__init__(molecule, sk_path, install_path, nthreads, version)

        # Initialize DFTB+ DFTB variables
        self.l_scc = l_scc
        self.scc_tol = scc_tol
        self.scc_max_iter = scc_max_iter

        self.l_onsite = l_onsite

        self.l_range_sep = l_range_sep
        self.lc_method = lc_method.lower()

        self.l_spin_pol = l_spin_pol
        self.unpaired_elec = unpaired_elec

        # Set initial guess for SCC term
        self.guess = guess.lower()
        self.guess_file = guess_file
        if not (self.guess in ["h0", "read"]):
            error_message = "Invalid initial guess for DFTB!"
            error_vars = f"guess = {self.guess}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.elec_temp = elec_temp
        self.mixer = mixer.lower()

        self.ex_symmetry = ex_symmetry.lower()
        self.e_window = e_window

        self.k_point = k_point
        self.l_periodic = l_periodic
        self.a_axis = np.copy(cell_length[0:3])
        self.b_axis = np.copy(cell_length[3:6])
        self.c_axis = np.copy(cell_length[6:9])

        # Check excitation symmetry in TDDFTB
        # TODO : Currently, allows only singlet excited states with TDDFTB
#        if not (self.ex_symmetry in ["singlet", "triplet"]):
        if (not self.ex_symmetry == "singlet"):
            error_message = "Invalid symmetry of excited states for TDDFTB given!"
            error_vars = f"ex_symmetry = {self.ex_symmetry}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.mpi = mpi
        self.mpi_path = mpi_path

        # Set 'l_nacme' and 're_calc' with respect to the computational method
        # TDDFTB do not produce NACs, so we should get NACME from CIoverlap
        # TDDFTB cannot compute the gradient of several states simultaneously.
        molecule.l_nacme = True
        self.re_calc = True

        # Calculate number of basis for current system
        # Set new variable to decide the position of basis functions in terms of atoms
        # DFTB method considers only valence electrons, so core electrons should be removed
        core_elec = 0.
        self.nbasis = 0
        self.check_atom = [0]
        for iat in range(molecule.nat_qm):
            # Check number of basis functions with respect to maximum angular momentum
            max_ang = max_l[molecule.symbols[iat]]
            if (max_ang == 's'):
                self.nbasis += 1
            elif (max_ang == 'p'):
                self.nbasis += 4
            elif (max_ang == 'd'):
                self.nbasis += 9
            else:
                error_message = "Number of basis for f orbital not implemented, see '$PYUNIXMDHOME/src/qm/dftbplus/dftb.py'!"
                error_vars = f"current atom = {molecule.symbols[iat]}, max_ang = {max_ang}"
                raise NotImplementedError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
            self.check_atom.append(self.nbasis)
            # Check number of core electrons with respect to atomic number
            sym_index = list(data.keys()).index(molecule.symbols[iat])
            if (sym_index > 0 and sym_index <= 2):
                core_elec += 0.
            elif (sym_index > 2 and sym_index <= 10):
                core_elec += 2.
            elif (sym_index > 10 and sym_index <= 18):
                core_elec += 10.
            elif (sym_index > 18 and sym_index <= 36):
                core_elec += 18.
            elif (sym_index > 36 and sym_index <= 54):
                core_elec += 36.
            else:
                error_message = "Core electrons for current element not implemented, see '$PYUNIXMDHOME/src/qm/dftbplus/dftb.py'!"
                error_vars = f"current atom = {molecule.symbols[iat]}, sym_index = {sym_index}"
                raise NotImplementedError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Set new variable to decide the position of atoms in terms of basis functions
        self.check_basis = []
        for ibasis in range(self.nbasis):
            for iat in range(molecule.nat_qm):
                ind_a = self.check_atom[iat] + 1
                ind_b = self.check_atom[iat + 1]
                if (ibasis + 1 >= ind_a and ibasis + 1 <= ind_b):
                    self.check_basis.append(iat + 1)

        # Initialize NACME variables
        # There is no core orbitals in TDDFTB (fixed occupations)
        # nocc is number of occupied orbitals and nvirt is number of virtual orbitals
        self.norb = self.nbasis
        self.nocc = int(int(molecule.nelec - core_elec) / 2)
        self.nvirt = self.norb - self.nocc

        # Replace norb by arrays containing the limits of the for loops.
        # For energy window calculations loops will not go from (0 to nocc/nvirt) or (0 to norb)
        # but from (nocc_min to nocc/0 to nvirt_max) or (nocc_min to norb).
        self.orb_ini = np.zeros(1, dtype=np.int32)
        self.orb_final = np.zeros(1, dtype=np.int32)
        self.orb_final[0] = self.norb

        if (self.e_window > eps):
            # Swap minimal/maximal values to replace them in reading of SPX.DAT by the minimal/maximal values.
            self.orb_ini[0] = self.norb
            self.orb_final[0] = 0

        self.ao_overlap = np.zeros((self.nbasis, self.nbasis))
        self.mo_coef_old = np.zeros((self.norb, self.nbasis))
        self.mo_coef_new = np.zeros((self.norb, self.nbasis))
        self.ci_coef_old = np.zeros((molecule.nst, self.nocc, self.nvirt))
        self.ci_coef_new = np.zeros((molecule.nst, self.nocc, self.nvirt))

    def get_data(self, molecule, base_dir, bo_list, dt, istep, calc_force_only):
        """ Extract energy, gradient and nonadiabatic couplings from (TD)DFTB method

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        self.copy_files(molecule, istep, calc_force_only)
        super().get_data(base_dir, calc_force_only)
        self.write_xyz(molecule)
        self.get_input(molecule, istep, bo_list, calc_force_only)
        self.run_QM(molecule, base_dir, istep, bo_list, calc_force_only)
        self.extract_QM(molecule, base_dir, istep, bo_list, dt, calc_force_only)
        self.move_dir(base_dir)

    def copy_files(self, molecule, istep, calc_force_only):
        """ Copy necessary scratch files in previous step

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Copy required files for NACME
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "geometry.xyz"), \
                os.path.join(self.scr_qm_dir, "../geometry.xyz.pre"))
            if (istep == 0):
                shutil.copy(os.path.join(self.scr_qm_dir, "eigenvec.bin"), \
                    os.path.join(self.scr_qm_dir, "../eigenvec.bin.pre"))
                shutil.copy(os.path.join(self.scr_qm_dir, "SPX.DAT"), \
                    os.path.join(self.scr_qm_dir, "../SPX.DAT.pre"))
                shutil.copy(os.path.join(self.scr_qm_dir, "XplusY.DAT"), \
                    os.path.join(self.scr_qm_dir, "../XplusY.DAT.pre"))

        # Copy required files to read initial guess
        if (self.guess == "read" and istep >= 0):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "charges.bin"), \
                os.path.join(self.scr_qm_dir, "../charges.bin.pre"))

    def get_input(self, molecule, istep, bo_list, calc_force_only):
        """ Generate DFTB+ input files: geometry.gen, dftb_in.hsd

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Make 'geometry.gen' file
        os.system("xyz2gen geometry.xyz")
        if (self.l_periodic):
            # Substitute C to S in first line
            file_be = open('geometry.gen', 'r')
            file_af = open('tmp.gen', 'w')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm} S\n'
                    first_row = False
                file_af.write(row)
            # Add gamma-point and cell lattice information
            geom_periodic = textwrap.dedent(f"""\
            {0.0:15.8f} {0.0:15.8f} {0.0:15.8f}
            {self.a_axis[0]:15.8f} {self.a_axis[1]:15.8f} {self.a_axis[2]:15.8f}
            {self.b_axis[0]:15.8f} {self.b_axis[1]:15.8f} {self.b_axis[2]:15.8f}
            {self.c_axis[0]:15.8f} {self.c_axis[1]:15.8f} {self.c_axis[2]:15.8f}
            """)
            file_af.write(geom_periodic)
            file_be.close()
            file_af.close()
            os.rename('tmp.gen', 'geometry.gen')

        # Make 'double.gen' file for CIoverlap in TDDFTB
        # In this case, we do not need to consider periodicity
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # Move previous files to currect directory
            os.rename('../geometry.xyz.pre', './geometry.xyz.pre')
            if (istep == 0):
                os.rename('../eigenvec.bin.pre', './eigenvec.bin.pre')
                os.rename('../SPX.DAT.pre', './SPX.DAT.pre')
                os.rename('../XplusY.DAT.pre', './XplusY.DAT.pre')
            # Open 'geometry.xyz.pre'
            file_af = open('double.xyz', 'w')
            file_be = open('geometry.xyz.pre', 'r')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm * 2}\n'
                    first_row = False
                file_af.write(row)
            file_be.close()
            # Open 'geometry.xyz'
            file_be = open('geometry.xyz', 'r')
            iline = 1
            for row in file_be:
                if (iline > 2):
                    file_af.write(row)
                iline += 1
            file_be.close()
            file_af.close()
            os.system("xyz2gen double.xyz")

        # Make 'dftb_in.hsd' file
        input_dftb = ""

        # Geometry Block
        input_geom = textwrap.dedent(f"""\
        Geometry = GenFormat{{
          <<< 'geometry.gen'
        }}
        """)
        input_dftb += input_geom

        # Hamiltonian Block
        input_ham_init = textwrap.dedent(f"""\
        Hamiltonian = DFTB{{
        """)
        input_dftb += input_ham_init

        # SCC-DFTB option
        if (self.l_scc):
            input_ham_scc = textwrap.indent(textwrap.dedent(f"""\
              SCC = Yes
              SCCTolerance = {self.scc_tol}
              MaxSCCIterations = {self.scc_max_iter}
              Mixer = {self.mixer}{{}}
            """), "  ")
            input_dftb += input_ham_scc

            # Onsite-corrected DFTB (OC-DFTB) option
            if (self.l_onsite):
                onsite_const_uu = ("\n" + " " * 18).join([f"  {itype}uu = {{ {onsite_uu[f'{itype}']} }}" for itype in self.atom_type])
                onsite_const_ud = ("\n" + " " * 18).join([f"  {itype}ud = {{ {onsite_ud[f'{itype}']} }}" for itype in self.atom_type])
                input_ham_oc = textwrap.indent(textwrap.dedent(f"""\
                  OnsiteCorrection = {{
                  {onsite_const_uu}
                  {onsite_const_ud}
                  }}
                """), "  ")
                input_dftb += input_ham_oc

            # Long-range corrected DFTB (LC-DFTB) option
            if (self.l_range_sep):
                input_ham_lc = textwrap.indent(textwrap.dedent(f"""\
                  RangeSeparated = LC{{
                    Screening = {self.lc_method}{{}}
                  }}
                """), "  ")
                input_dftb += input_ham_lc

            # Spin-polarized DFTB option
            if (self.l_spin_pol and molecule.nst == 1):
                input_ham_spin = textwrap.dedent(f"""\
                SpinPolarisation = Colinear{{
                  UnpairedElectrons = {self.unpaired_elec}
                }}
                """)
                input_dftb += input_ham_spin

            # Read atomic spin constants used in spin-polarized DFTB or TDDFTB
            # TODO : Currently, allows only singlet excited states with TDDFTB
#            if (self.l_spin_pol or self.ex_symmetry == "triplet"):
            if (self.l_spin_pol and molecule.nst == 1):
                if (self.l_range_sep):
                    spin_constant = ("\n" + " " * 18).join([f"  {itype} = {{ {spin_w_lc[f'{itype}']} }}" for itype in self.atom_type])
                else:
                    spin_constant = ("\n" + " " * 18).join([f"  {itype} = {{ {spin_w[f'{itype}']} }}" for itype in self.atom_type])
                input_ham_spin_w = textwrap.indent(textwrap.dedent(f"""\
                  SpinConstants = {{
                    ShellResolvedSpin = Yes
                  {spin_constant}
                  }}
                """), "  ")
                input_dftb += input_ham_spin_w

            # Read 'charges.bin' from previous step
            if (self.guess == "read"):
                if (istep == -1):
                    if (os.path.isfile(self.guess_file)):
                        # Copy guess file to currect directory
                        shutil.copy(self.guess_file, os.path.join(self.scr_qm_dir, "charges.bin"))
                        restart = "Yes"
                    else:
                        restart = "No"
                elif (istep >= 0):
                    # Move previous file to currect directory
                    os.rename("../charges.bin.pre", "./charges.bin")
                    restart = "Yes"
            elif (self.guess == "h0"):
                restart = "No"

            # Read 'charges.bin' for surface hopping dynamics when hop occurs
            if (calc_force_only):
                restart = "Yes"

            input_ham_restart = textwrap.indent(textwrap.dedent(f"""\
              ReadInitialCharges = {restart}
            """), "  ")
            input_dftb += input_ham_restart

        # TODO: for QM/MM, point_charge??

        if (self.l_periodic):
            num_k_point = np.sum(self.k_point)
            if (num_k_point == 3):
                # gamma-point sampling
                input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
                  KPointsAndWeights = {{
                    0.0 0.0 0.0 1.0
                  }}
                """), "  ")
            else:
                # K-point sampling
                shift_vector = [0.5 if (ik % 2 == 0) else 0 for ik in self.k_point]
                input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
                  KPointsAndWeights = SupercellFolding{{
                    {self.k_point[0]} 0.0 0.0
                    0.0 {self.k_point[1]} 0.0
                    0.0 0.0 {self.k_point[2]}
                    {shift_vector[0]} {shift_vector[1]} {shift_vector[2]}
                  }}
                """), "  ")
            input_dftb += input_ham_periodic

        angular_momentum = ("\n" + " " * 10).join([f"  {itype} = '{max_l[f'{itype}']}'" for itype in self.atom_type])
        input_ham_basic = textwrap.dedent(f"""\
          Charge = {molecule.charge}
          Filling = Fermi{{
            Temperature[K] = {self.elec_temp}
          }}
          MaxAngularMomentum = {{
          {angular_momentum}
          }}
          SlaterKosterFiles = Type2FileNames{{
            Prefix = '{self.sk_path}'
            Separator = '-'
            Suffix = '.skf'
            LowerCaseTypeName = No
          }}
        }}
        """)
        input_dftb += input_ham_basic

        # Analysis Block
        input_analysis = textwrap.dedent(f"""\
        Analysis = {{
          CalculateForces = Yes
          WriteBandOut = Yes
          WriteEigenvectors = Yes
          MullikenAnalysis = Yes
        }}
        """)
        input_dftb += input_analysis

        # Options Block
        input_options = textwrap.dedent(f"""\
        Options = {{
          WriteDetailedXml = Yes
          WriteDetailedOut = Yes
          TimingVerbosity = -1
        }}
        """)
        input_dftb += input_options

        # ExcitedState Block
        if (molecule.nst > 1):
            # Calculate excited state force for target state
            if (bo_list[0] > 0):
                ex_force = "Yes"
                rst = bo_list[0]
            else:
                ex_force = "No"
                rst = bo_list[0] + 1

            # Set number of excitations in TDDFTB
            # This part can be modified by users
            if (molecule.nat_qm <= 5):
                num_ex = molecule.nst + 2
            elif (molecule.nat_qm > 5 and molecule.nat_qm <= 15):
                num_ex = 2 * molecule.nst + 2
            else:
                num_ex = 3 * molecule.nst + 2

            # Write XplusY data?
            if (self.calc_coupling):
                xpy = "Yes"
            else:
                xpy = "No"

            input_excited = textwrap.dedent(f"""\
            ExcitedState = Casida{{
              NrOfExcitations = {num_ex}
              StateOfInterest = {rst}
              Symmetry = {self.ex_symmetry}
              WriteTransitions = Yes
              WriteSPTransitions = {xpy}
              WriteMulliken = Yes
              WriteXplusY = {xpy}
              EnergyWindow [eV] = {self.e_window}
              ExcitedStateForces = {ex_force}
            }}
            """)

            input_dftb += input_excited

        # ParserOptions Block
        if (self.version == "19.1"):
            parser_version = 7
        elif (self.version == "20.1"):
            parser_version = 8

        input_parseroptions = textwrap.dedent(f"""\
        ParserOptions = {{
          ParserVersion = {parser_version}
        }}
        """)
        input_dftb += input_parseroptions

        # Parallel Block
        if (self.mpi):
            if (self.l_spin_pol and self.nthreads > 1):
                groups = 2
            else:
                groups = 1
            input_parallel = textwrap.dedent(f"""\
            Parallel = {{
              Groups = {groups}
              UseOmpThreads = No
              Blacs = BlockSize {{ 32 }}
            }}
            """)
            input_dftb += input_parallel

        # Write 'dftb_in.hsd.geom' file
        file_name = f"dftb_in.hsd.geom.{bo_list[0]}"
        with open(file_name, "w") as f:
            f.write(input_dftb)

        # Write 'dftb_in.hsd.double' file
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # New input for dftb
            input_dftb = ""

            # Geometry Block
            input_geom = textwrap.dedent(f"""\
            Geometry = GenFormat{{
              <<< 'double.gen'
            }}
            """)
            input_dftb += input_geom
            input_dftb += input_ham_init
            input_dftb += input_ham_basic

            # Options Block
            input_options = textwrap.dedent(f"""\
            Options = {{
              WriteDetailedXml = Yes
              WriteDetailedOut = Yes
              WriteHS = Yes
              TimingVerbosity = -1
            }}
            """)
            input_dftb += input_options

            file_name = "dftb_in.hsd.double"
            with open(file_name, "w") as f:
                f.write(input_dftb)

    def run_QM(self, molecule, base_dir, istep, bo_list, calc_force_only):
        """ Run (TD)DFTB calculation and save the output files to qm_log directory

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Set run command
        qm_command = os.path.join(self.qm_path, "dftb+")
        if (self.mpi):
            # MPI setting
            os.environ["OMP_NUM_THREADS"] = "1"
            mpi_command = os.path.join(self.mpi_path, "mpirun")
            command = f"{mpi_command} -np {self.nthreads} {qm_command} > log"
        else:
            # OpenMP setting
            os.environ["OMP_NUM_THREADS"] = f"{self.nthreads}"
            command = f"{qm_command} > log"

        # Run DFTB+ for calculation of overlap matrix
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            shutil.copy("dftb_in.hsd.double", "dftb_in.hsd")
            os.system(command)

        # Copy dftb_in.hsd for target state
        file_name = f"dftb_in.hsd.geom.{bo_list[0]}"
        shutil.copy(file_name, "dftb_in.hsd")

        # Run DFTB+ method for molecular dynamics
        os.system(command)

        # Copy detailed.out for target state
        file_name = f"detailed.out.{bo_list[0]}"
        shutil.copy("detailed.out", file_name)

        # Copy the output file to 'qm_log' directory
        tmp_dir = os.path.join(base_dir, "qm_log")
        if (os.path.exists(tmp_dir)):
            detailed_out_step = f"detailed.out.{istep + 1}.{bo_list[0]}"
            shutil.copy("detailed.out", os.path.join(tmp_dir, detailed_out_step))
            log_step = f"log.{istep + 1}.{bo_list[0]}"
            shutil.copy("log", os.path.join(tmp_dir, log_step))

    def extract_QM(self, molecule, base_dir, istep, bo_list, dt, calc_force_only):
        """ Read the output files to get BO information

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Read 'detailed.out' file
        # TODO: the qmmm information is written in this file
        file_name = f"detailed.out.{bo_list[0]}"
        with open(file_name, "r") as f:
            detailed_out = f.read()

        # Read 'EXC.DAT' file
        if (molecule.nst > 1):
            file_name = "EXC.DAT"
            with open(file_name, "r") as f:
                exc_out = f.read()

        # Energy
        if (not calc_force_only):
            energy = re.findall('Total energy:\s+([-]\S+) H', detailed_out)
            energy = np.array(energy[0], dtype=np.float64)
            molecule.states[0].energy = energy

            if (molecule.nst > 1):
                tmp_e = f'[=]+\n' + ('\s+([-]*\S+)\s+\S+\s+\d+\s+->\s+\d+\s+\S+\s+\S+\s+[ST]') * molecule.nst
                energy = re.findall(tmp_e, exc_out)
                energy = np.array(energy[0], dtype=np.float64)
                energy *= eV_to_au
                for ist in range(1, molecule.nst):
                    molecule.states[ist].energy = molecule.states[0].energy + energy[ist - 1]

        # Force
        tmp_f = 'Total Forces' + '\n\s+\d*\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
        force = re.findall(tmp_f, detailed_out)
        force = np.array(force[0], dtype=np.float64)
        force = force.reshape(molecule.nat_qm, 3, order='C')
        molecule.states[bo_list[0]].force = np.copy(force)

        # NACME
        if (self.calc_coupling and not calc_force_only):
            if (istep >= 0):
                self.CI_overlap(molecule, istep, dt)

    def CI_overlap(self, molecule, istep, dt):
        """ Read the necessary files and calculate NACME from tdnac.c routine,
            note that only reading of several files is required in this method

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param double dt: Time interval
        """
        # Read upper right block of 'oversqr.dat' file (< t | t+dt >)
        file_name_in = "oversqr.dat"

        self.ao_overlap = np.zeros((self.nbasis, self.nbasis))
        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            row = 0
            iline = 0
            for line in lines:
                # Skip first five lines and read upper block
                if (iline in range(5, 5 + self.nbasis)):
                    col = 0
                    count = False
                    field = line.split()
                    for element in field:
                        # Read right block
                        if (count):
                            ind_a = self.check_basis[row]
                            ind_b = self.check_basis[col]
                            if (ind_a == ind_b):
                                # Choose onsite (same-atom) block
                                # Sometimes NaN or too large values appear in the onsite block due to the slater-koster file
                                # The values set to 1 or 0 regardless of original elements
                                if (row == col):
                                    # Diagonal element in onsite block
                                    new_val = 1.
                                else:
                                    # Off-diagonal element in onsite block
                                    new_val = 0.
                            else:
                                # Choose offsite (different-atom) block
                                new_val = float(element)
                            # Set overlap matrix element
                            self.ao_overlap[row, col] = new_val
                        col += 1
                        # Read right block
                        if (col > self.nbasis - 1):
                            col -= self.nbasis
                            count = True
                    row += 1
                iline += 1
#        np.savetxt("test-over", self.ao_overlap, fmt=f"%6.3f")

        # Read 'eigenvec.bin.pre' file at time t
        if (istep == 0):
            file_name_in = "eigenvec.bin.pre"

            self.mo_coef_old = np.zeros((self.norb, self.nbasis))
            with open(file_name_in, "rb") as f_in:
                dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                for iorb in range(self.norb):
                    dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                    data = np.fromfile(f_in, dtype=np.float64, count=self.nbasis)
                    self.mo_coef_old[iorb] = data
#            np.savetxt("test-mo1", self.mo_coef_old, fmt=f"%12.6f")

        # Read 'eigenvec.bin' file at time t + dt
        file_name_in = "eigenvec.bin"

        self.mo_coef_new = np.zeros((self.norb, self.nbasis))
        with open(file_name_in, "rb") as f_in:
            dummy = np.fromfile(f_in, dtype=np.int32, count=1)
            for iorb in range(self.norb):
                dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                data = np.fromfile(f_in, dtype=np.float64, count=self.nbasis)
                self.mo_coef_new[iorb] = data
#        np.savetxt("test-mo2", self.mo_coef_new, fmt=f"%12.6f")

        # The CI coefficients are arranged in order of single-particle excitations
        # Read 'SPX.DAT.pre' file at time t
        if (istep == 0):
            file_name_in = "SPX.DAT.pre"

            with open(file_name_in, "r") as f_in:
                lines = f_in.readlines()
                # Dimension for CI coefficients (number of excitations)
                ndim_old = int(lines[-2].strip().split()[0])
                get_wij_ind_old = np.zeros((ndim_old, 2), dtype=np.int32)
                iline = 0
                for line in lines:
                    # Skip first five lines
                    if (iline in range(5, 5 + ndim_old)):
                        # Column information: 1st = index, 4th = occ(i), 6th = virt(a)
                        field = line.split()
                        # Determine new limits for for-loops
                        if (int(field[5]) > self.orb_final[0]):
                            self.orb_final[0] = int(field[5])
                        if (int(field[3]) < self.orb_ini[0] + 1):
                            self.orb_ini[0] = int(field[3]) - 1
                        get_wij_ind_old[int(field[0]) - 1] = [int(field[3]), int(field[5])]
                    iline += 1

        # Read 'SPX.DAT' file at time t + dt
        file_name_in = "SPX.DAT"

        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            # Dimension for CI coefficients (number of excitations)
            ndim = int(lines[-2].strip().split()[0])
            get_wij_ind_new = np.zeros((ndim, 2), dtype=np.int32)
            iline = 0
            for line in lines:
                # Skip first five lines
                if (iline in range(5, 5 + ndim)):
                    # Column information: 1st = index, 4th = occ(i), 6th = virt(a)
                    field = line.split()
                    if (int(field[5]) > self.orb_final[0]):
                        self.orb_final[0] = int(field[5])
                    if (int(field[3]) < self.orb_ini[0] + 1):
                        self.orb_ini[0] = int(field[3]) - 1
                    get_wij_ind_new[int(field[0]) - 1] = [int(field[3]), int(field[5])]
                iline += 1

        # Read 'XplusY.DAT.pre' file at time t
        if (istep == 0):
            file_name_in = "XplusY.DAT.pre"

            self.ci_coef_old = np.zeros((molecule.nst, self.nocc, self.nvirt))
            with open(file_name_in, "r") as f_in:
                lines = f_in.readlines()
                iline = 0
                for line in lines:
                    if (iline == 0):
                        field = line.split()
                        assert (int(field[0]) == ndim_old)
                        assert (int(field[1]) >= molecule.nst - 1)
                        # nxply is number of lines for each excited state in 'XplusY.dat'
                        nxply = int(ndim_old / 6) + 1
                        if (ndim_old % 6 != 0):
                            nxply += 1
                    else:
                        field = line.split()
                        if (iline % nxply == 1):
                            ind = 0
                            ist = int(field[0]) - 1
                            # In general, TDDFTB calculate the excited states more than molecule.nst,
                            # so we do not need to read all data for 'XplusY.DAT'
                            if (ist == molecule.nst - 1):
                                break
                        else:
                            # Currently, elements for CI coefficients for S0 state are zero (not used values)
                            for element in field:
                                ind_occ = get_wij_ind_old[ind, 0] - 1
                                ind_virt = get_wij_ind_old[ind, 1] - self.nocc - 1
                                self.ci_coef_old[ist + 1, ind_occ, ind_virt] = float(element)
                                ind += 1
                    iline += 1
#            np.savetxt("test-ci1", self.ci_coef_old[1], fmt=f"%12.6f")

        # Read 'XplusY.DAT' file at time t + dt
        file_name_in = "XplusY.DAT"

        self.ci_coef_new = np.zeros((molecule.nst, self.nocc, self.nvirt))
        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            iline = 0
            for line in lines:
                if (iline == 0):
                    field = line.split()
                    assert (int(field[0]) == ndim)
                    assert (int(field[1]) >= molecule.nst - 1)
                    # nxply is number of lines for each excited state in 'XplusY.dat'
                    nxply = int(ndim / 6) + 1
                    if (ndim % 6 != 0):
                        nxply += 1
                else:
                    field = line.split()
                    if (iline % nxply == 1):
                        ind = 0
                        ist = int(field[0]) - 1
                        # In general, TDDFTB calculate the excited states more than molecule.nst,
                        # so we do not need to read all data for 'XplusY.DAT'
                        if (ist == molecule.nst - 1):
                            break
                    else:
                        # Currently, elements for CI coefficients for S0 state are zero (not used values)
                        for element in field:
                            ind_occ = get_wij_ind_new[ind, 0] - 1
                            ind_virt = get_wij_ind_new[ind, 1] - self.nocc - 1
                            self.ci_coef_new[ist + 1, ind_occ, ind_virt] = float(element)
                            ind += 1
                iline += 1
#        np.savetxt("test-ci2", self.ci_coef_new[1], fmt=f"%12.6f")

        # Calculate wavefunction overlap with orbital scheme
        # Reference: J. Phys. Chem. Lett. 2015, 6, 4200-4203
        wf_overlap(self, molecule, istep, dt)


# -*- coding: UTF-8 -*-
#   \file solvate_maths.py
#   \brief This file provides the specific computations functionality for other parts of solvate.
#
#   Copyright by the Authors and individual contributors. All rights reserved.
#
#   Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
#       1) Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
#       2) Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
#       3) Neither the name of Michal Tykac nor the names of this code's contributors may be used to endorse or promote products derived from this software without specific prior written permission.
#
#   This software is provided by the copyright holder and contributors "as is" and any express or implied warranties, including, but not limitted to, the implied warranties of merchantibility and fitness for a particular purpose are disclaimed. In no event shall the copyright owner or the contributors be liable for any direct, indirect, incidental, special, exemplary, or consequential damages (including, but not limitted to, procurement of substitute goods or services, loss of use, data or profits, or business interuption) however caused and on any theory of liability, whether in contract, strict liability or tort (including negligence or otherwise) arising in any way out of the use of this software, even if advised of the possibility of such damage.
#
#   \author    Michal Tykac
#   \author    Lada Biedermannová
#   \author    Jiří Černý
#   \version   0.1.0
#   \date      DEC 2020
######################################################

######################################################
# Imports

import numpy                                          ### Maths
import solvate.solvate_log as solvate_log             ### For writing the log

######################################################
# residueToResFragmentDistance_backbone ()
def residueToResFragmentDistance_backbone ( res, resFrag, settings ):
    """
    This function computes the procrustes optimal overlay of the backbone atoms of the input residue and
    hydrated residue fragment and then proceeds to obtain the RMSD of these backbone atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : resFrag
        A list containing the positions and names of the hydrated residue frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Reduce residue to N, Ca, C and O (backbone) atoms
    resAts                                            = []
    resAtsNames                                       = []
    for rIt in range( 0, len( res[1] ) ):
        if ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "O" ) or ( res[1][rIt][0] == "N" ):
            resAts.append                             ( [ res[1][rIt][1], res[1][rIt][2], res[1][rIt][3] ] )
            resAtsNames.append                        ( res[1][rIt][0] )

    ### Sanity check
    if len( resAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all backbone atoms (C, CA, N and O) in the residue described by " + str( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( settings )

    ### Reduce hydrated residue fragment to N, Ca, C and O (backbone) atoms
    fraAts                                            = []
    for fIt in range( 0, len( resFrag ) ):
        for rIt in range ( 0, len( resAtsNames ) ):
            if resFrag[fIt][0] == resAtsNames[rIt]:
                fraAts.append                         ( [ resFrag[fIt][1], resFrag[fIt][2], resFrag[fIt][3] ] )

    ### Sanity check
    if len( fraAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all backbone atoms (C, CA, N and O) in the supplied hydrated fragment " + str ( resFrag ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( settings )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAts ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fraAts ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAts, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated residue fragment distance of " + str( rmsdDist ), settings, 5 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# residueToResFragmentDistance_rotamer ()
def residueToResFragmentDistance_rotamer ( res, resFrag, settings ):
    """
    This function computes the procrustes optimal overlay of the side-chain atoms of the input residue and
    hydrated residue fragment and then proceeds to obtain the RMSD of these side-chain atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : resFrag
        A list containing the positions and names of the hydrated residue frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Initialise variables
    resAts                                            = []
    resAtsNames                                       = []
    fraAts                                            = []

    ### Parse out the side-chain atoms for the residue
    for rIt in range( 0, len( res[1] ) ):

        # Deal with various atom names and lengths for different amino acids
        if ( ( ( res[0] == "VAL" ) or ( res[0] == "ILE" ) ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "CG1" ) ) ) or \
           ( ( ( res[0] == "THR" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "OG1" ) ) ) or \
           ( ( ( res[0] == "SER" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "OG"  ) ) ) or \
           ( ( ( res[0] == "CYS" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "SG"  ) ) ) or \
           ( ( ( res[0] == "ARG" ) or ( res[0] == "ASN" ) or ( res[0] == "ASP" ) or ( res[0] == "GLN" ) or ( res[0] == "GLU" ) or ( res[0] == "HIS" ) or ( res[0] == "LEU" ) or \
               ( res[0] == "LYS" ) or ( res[0] == "MET" ) or ( res[0] == "PHE" ) or ( res[0] == "PRO" ) or ( res[0] == "TRP" ) or ( res[0] == "TYR" ) \
                                                          ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "CG"  ) ) ):

            # Save
            resAts.append                             ( [ res[1][rIt][1], res[1][rIt][2], res[1][rIt][3] ] )
            resAtsNames.append                        ( res[1][rIt][0] )

    ### Sanity check
    if len ( resAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all side-chain atoms (C, CA, CB and ?) in the supplied residue " + str ( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Reduce fragment to N, Ca, Cb and X (as above) atoms
    for rIt in range( 0, len( resAtsNames ) ):

        # For each fragment atom
        for fIt in range( 0, len( resFrag ) ):

            # Get only the same atoms
            if resAtsNames[rIt] == resFrag[fIt][0]:

                # Save
                fraAts.append                         ( [ resFrag[fIt][1], resFrag[fIt][2], resFrag[fIt][3] ] )

    ### Sanity check
    if len ( fraAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all side-chain atoms (C, CA, CB and ?) in the supplied hydrated residue fragment " + str ( resFrag ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAts ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fraAts ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAts, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated residue fragment distance of " + str( rmsdDist ), settings, 4 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# residueToResFragmentDistance_direct ()
def residueToResFragmentDistance_direct ( res, fragInfo, settings ):
    """
    This function computes the procrustes optimal overlay of all available atoms
    and then proceeds to obtain the RMSD of these atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : fragInfo
        A list containing the positions and names of the hydrated frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Initialise variables
    resAtoms                                          = []
    fragAtoms                                         = []

    ### Get residue positions in same order as fragment
    for fIt in range ( 0, len ( fragInfo ) ):

        # For each residue atom name
        for atIt in range ( 0, len ( res[1] ) ):

            # Compare the atom names
            if res[1][atIt][0] == fragInfo[fIt][0]:

                # Save atoms in same order
                resAtoms.append                       ( [ res[1][atIt][1], res[1][atIt][2], res[1][atIt][3] ] )
                fragAtoms.append                      ( [ fragInfo[fIt][1], fragInfo[fIt][2], fragInfo[fIt][3] ] )

    ### Check for missing atoms
    if len ( fragInfo ) != len ( resAtoms ):
        solvate_log.writeLog                          ( "!!! ERROR !!! Failed to match the fragment atoms of fragment " + str ( fragInfo ) + " to residue " + str( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAtoms ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fragAtoms ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAtoms, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated fragment distance of " + str( rmsdDist ), settings, 5 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# procrustes ()
def procrustes ( X, Y ):
    """
    A port of MATLAB's "procrustes" function to Numpy.

    Procrustes analysis determines a linear transformation (translation,
    reflection, orthogonal rotation and scaling) of the points in Y to best
    conform them to the points in matrix X, using the sum of squared errors
    as the goodness of fit criterion.

        d, Z, [tform] = procrustes ( X, Y )

    In this particular code, scaling has been disabled (in protein comparisons,
    scale plays a role) as was the reflections ( mirror immage of a molecule is
    a different molecule ).

    Parameters
    ----------
    numpy.ndarray : X, Y
        Matrices of target and input coordinates. They must have equal
        numbers of  points (rows), but Y may have fewer dimensions
        (columns) than X.

    Returns
    -------
    float : d
        The residual sum of squared errors, normalized according to a
        measure of the scale of X, ( ( X - X.mean(0))^2 ).sum( )

    matrix : Z
        The matrix of transformed Y-values

    dictionary : tform
        A dict specifying the rotation, translation and scaling that
        maps X --> Y

    """

    ### Find the dimensions of the input matrices
    nx , mx                                           = X.shape
    ny , my                                           = Y.shape

    ### Find the column means of the input matrices
    muX                                               = X.mean(0)
    muY                                               = Y.mean(0)

    ### Find the x - mean values of the input matrices
    X0                                                = X - muX
    Y0                                                = Y - muY

    ### Sum the squares of the x - mean values of the input matrices
    ssX                                               = ( numpy.power ( X0, 2.0 ) ).sum()
    ssY                                               = ( numpy.power ( Y0, 2.0 ) ).sum()

    ### Compute the centred Frobenius norm of the input matrices (i.e. square root of sum of x - mean squared)
    normX                                             = numpy.sqrt ( ssX )
    normY                                             = numpy.sqrt ( ssY )

    ### Scale x - mean to equal (unit) norm
    X0                                               /= normX
    Y0                                               /= normY

    ### If second matrix has more rows than the first, shorten it ( this should really be error... )
    if my < mx:
        Y0                                            = numpy.concatenate ( ( Y0, numpy.zeros ( nx, mx - my ) ), 0 )

    ### Find the optimal rotation matrix of Y using least squares
    A                                                 = numpy.dot ( X0.T, Y0 )
    U,s,Vt                                            = numpy.linalg.svd ( A, full_matrices=False )
    V                                                 = Vt.T
    T                                                 = numpy.dot(V, U.T)

    ### Does the current solution use a reflection?
    have_reflection                                   = numpy.linalg.det(T) < 0

    ### If that's the case, force the other reflection (i.e. solution with no reflections)
    if have_reflection != False:
        V[:,-1]                                      *= -1
        s[-1]                                        *= -1
        T                                             = numpy.dot(V, U.T)

    ### Find the trace of singular values
    traceTA                                           = s.sum()

    ### With no scaling
    b                                                 = 1
    d                                                 = 1 + ssY/ssX - 2 * traceTA * normY / normX
    Z                                                 = normY * numpy.dot ( Y0, T ) + muX

    ### If second matrix has more rows than the first, shorten it ( this should really be error... )
    if my < mx:
        T                                             = T[:my,:]

    ### Compute the translation vector
    c                                                 = muX - b * numpy.dot ( muY, T )

    ### Save all to transformation list
    tform                                             = {'rotation':T, 'scale':b, 'translation':c}

    ### Done
    return d, Z, tform

######################################################
# getRMSD ()
def getRMSD ( coo1, coo2 ):
    """
    This is a simple function for computing RMSD given two matrices with equal dimmensions (n rows and
    exactly 3 columns) of (presumably atomic) co-ordinates.

    Parameters
    ----------
    matrix : coo1, coo2
        matrices of target and input coordinates. They must have equal
        numbers of points (rows) and equal number of dimensions (columns
        = 3)

    Returns
    -------
    float : dist
        The RMSD distance value between the two matrices.

    """
    ### Initialise sum
    sum                                               = 0.0

    ### For each row
    for iter in range ( 0, len ( coo1 ) ):
        sum                                          += ( numpy.power ( coo1[iter][0] - coo2[iter][0], 2.0 ) +
                                                          numpy.power ( coo1[iter][1] - coo2[iter][1], 2.0 ) +
                                                          numpy.power ( coo1[iter][2] - coo2[iter][2], 2.0 ) )

    ### Done
    return                                            ( numpy.sqrt ( sum  / float ( len ( coo1 ) ) ) )

######################################################
# findSmallestDistance ()
def findSmallestDistance ( distances, threshold ):
    """
    This function takes a list of distances and a threshold and returns the index of the
    smallest distance or -1 if no distances is below the threshold.

    Parameters
    ----------
    list : distances
        The list of the distances to be searched.

    float : threshold
        The threshold below which any passing distance needs to be.

    Returns
    -------
    int : index
        The index of the smallest distance or -1 if no distance is below the threshold.

    """
    ### Find lowest RMSD and check if it is larger than the threshold
    minVal                                            = numpy.Inf
    minInd                                            = -1
    indCtr                                            = 0

    ### Find the best match (lowest distance)
    for val in distances:
        if ( val[0] < minVal ) and ( val[0] < threshold ):
            minVal                                    = val[0]
            minInd                                    = indCtr
        indCtr                                       += 1

    ### Return the index
    return                                            ( minInd )

######################################################
# findPassingDistances ()
def findPassingDistances ( distances, threshold ):
    """
    This function takes a list of distances and a threshold and returns all indices of
    distances which pass the threshold, or -1 if no distances is below the threshold.

    Parameters
    ----------
    list : distances
        The list of the distances to be searched.

    float : threshold
        The threshold below which any passing distance needs to be.

    Returns
    -------
    list : indices
        The indices of all passing distances or -1 if no distance is below the threshold.

    """
    ### Initialise variables
    minInd                                            = []
    indCtr                                            = 0

    ### Find all passing distances
    for val in distances:
        if val[0] < threshold:
            minInd.append                             ( indCtr )
        indCtr                                       += 1

    ### Check for at least one
    if len( minInd ) == 0:
        minInd                                        = -1

    ### Return the index
    return                                            ( minInd )

######################################################
# findPassingDistances ()
def rotateAndTranslateAtom ( atom, transform ):
    """
    This function takes a numpy.array of atom positions (X, Y, Z) and proceeds to
    apply the transform rotation and then transform translation to this position.
    It then returns the new position.

    Warning: Note that rotation is done before translation, if this is not the
    correct order of transformations, then please do not use the function!

    Parameters
    ----------
    numpy.array : atom
        Array with three numbers signifying the current atom position.

    dictionary : transform
        Dictionary containing the "rotation" key with the optimal rotation matrix and
        a "translation" key with the optimal translation. It should be the result of
        the procrustes analysis.

    Returns
    -------
    numpy.array : res
        Array containing the new position resulting from rotating and translation the
        input position.

    """
    ### Apply rotation
    atomPos                                           = numpy.matmul ( atom, transform["rotation"] )

    ### Apply translation
    atomPos                                           = atomPos + transform["translation"]

    ### Return the index
    return                                            ( atomPos )

######################################################
# getWaterClusterCOM ()
def getWaterClusterCOM ( clusterIndex, clusteredWaters ):
    """
    This function takes the index of a cluster, the list of all water molecules and  the
    appropriate labels list and proceeds to compute the occupancy weighted COM position
    for the cluster, which it then returns.

    Parameters
    ----------
    int : clusterIndex
        The index of the cluster for which the computation is to be made.

    list : clusteredWaters
        A list of cluster labels, cluster count, water positions and molecule occupancies and B factors for
        each structure type available in the noClashWaters input variable.

    list : watersNoClash
        List containing a list of water molecules which have no clashes (as defined by the settings)
        to the various versions of the input structure. The versions are: full structure, no hydrogens,
        no waters, no ligands. If some of the versions were not requested, the list for that version will
        be empty.

    Returns
    -------
    list : COM
        This list contains the COM position weighted by the occupancy of the contributing water molecules and the
        average occupancy and B factor as well.

    """
    ### Find all member indices
    members                                           = numpy.where ( clusteredWaters[0] == clusterIndex )[0]

    ### Initialise variables
    COM                                               = numpy.zeros ( 5 )
    occSum                                            = 0.0
    bFacSum                                           = 0.0

    ### For each member, find sum of positions weighted by occupancy
    for mem in members:

        ### Get occupancy weighted COM and occupancy sum
        COM[0]                                       += ( clusteredWaters[2][mem][0] * clusteredWaters[3][mem] )
        COM[1]                                       += ( clusteredWaters[2][mem][1] * clusteredWaters[3][mem] )
        COM[2]                                       += ( clusteredWaters[2][mem][2] * clusteredWaters[3][mem] )
        occSum                                       += clusteredWaters[3][mem]
        bFacSum                                      += clusteredWaters[4][mem]

    ### Average by occupancy sum
    COM[0]                                           /= occSum
    COM[1]                                           /= occSum
    COM[2]                                           /= occSum
    COM[3]                                            = occSum / float ( len ( members ) )
    COM[4]                                            = bFacSum / float ( len ( members ) )

    ### Done
    return ( COM )


#!/usr/bin/python3

import cobra as cb
import pandas as pd
import re
import sys
import getopt
import os.path
import copy
import csv
import math
import cobra.flux_analysis.variability
import subprocess
import shutil, errno
import statistics
from cobra import Reaction
import cometspy as c
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'figure.max_open_warning': 0})
################################################################

def model_modifications(glu1, ac1, o21, glu2, ac2, o22):
    # Single GEMs parameter modifications
  # ===================================

    # 1.1.- [COBRApy] Establish modifications in model 1
    model1=c.model('/home/ana/FLYCOP/Scripts/iJO1366.cmd')
    model1.change_bounds('GLCtex_copy2', 0, 0)
    model1.change_bounds('PFK_3', 0, 0)
    model1.change_bounds('EX_o2_e',-1000,1000)
    model1.change_bounds('EX_glc__D_e',-1000,1000)
    model1.change_bounds('GLCtex_copy1', 0, -glu1)
    model1.change_bounds('O2tex', 0,-o21)
    if ac1<=0:
        model1.change_bounds('ACtex',-1000,-ac1)
        model1.change_bounds('EX_ac_e', ac1, 1000)
    else:
        model1.change_bounds('ACtex', -ac1,1000)
        model1.change_bounds('EX_ac_e',-1000,ac1)
    model1.id='ecoli_1_tmp'
    model1.write_comets_model()
    del(model1)

    # 1.2.- [COBRApy] Establish modifications in model 2
    model2=c.model("/home/ana/FLYCOP/Scripts/iJO1366.cmd")
    model2.change_bounds('GLCtex_copy2', 0, 0)
    model2.change_bounds('PFK_3', 0, 0)
    model2.change_bounds('EX_o2_e',-1000,1000)
    model2.change_bounds('EX_glc__D_e',-1000,1000)
    model2.change_bounds('GLCtex_copy1',0,-glu2)
    model2.change_bounds('O2tex',0,-o22)
    if(ac2<=0):
        model2.change_bounds('ACtex',-1000,-ac2)
        model2.change_bounds('EX_ac_e',ac2,1000)

    else:
        model2.change_bounds('ACtex',-ac2,1000)
        model2.change_bounds('EX_ac_e',-1000,ac2)
    model2.id='ecoli_2_tmp'
    model2.write_comets_model()
    del(model2)
    return (print('Modifications of the model done!'))

def initialize_layout(file):
    #The layout is a file with the stablished format
    layout=c.layout(file)
    return layout

def initialize_params(package, globall):
    """Function to initialize the comets' params class
    it can be initialize by submitting two files, one for the package parameters
    and one for the global ones.
    If you don't submit a file, the params class will be initialize with the values stated below
    which have been tested in this simulation"""

    if package or globall is not None:
        params = c.params(global_params = globall, package_params= package)
    else:
        params = c.params()
        params.all_params['maxCycles']=240
        params.all_params['timeStep']=0.1
        params.all_params['spaceWidth']=0.05
        params.all_params['allowCellOverlap']= True
        params.all_params['deathRate']= 0.01
        params.all_params['numRunThreads']= 4
        params.all_params['maxSpaceBiomass']= 10
        params.all_params['defaultVmax']=20
        params.all_params['showCycleTime']=True
        params.all_params['writeTotalBiomassLog']=True
        params.all_params['writeMediaLog']=True
        params.all_params['writeFluxLog']=True
        params.all_params['useLogNameTimeStamp']=False
        params.all_params['FluxLogRate']=1
        params.all_params['MediaLogRate']=1

    return params

def make_graph(comets, met, met22):

    '''This function creates a figure and saves it to pdf format.
    It also creates the file biomass_vs_met.txt which contais the quantity
    of each strain and metabolite and has the following columns:
    time(h), strain1 ... strainX, met1 ... metX.'''

    df =comets.media
    df=df.drop(columns=['x', 'y'])
    met2 =df[df['metabolite'] == str(met22)]
    met1 =df[df['metabolite'] == str(met)]
    met1 =met1.reset_index()
    met1=met1.drop(columns=['index', 'metabolite'])

    i = 45
    while i < 242:
        met1.loc[i] = [i+2, 0]
        i +=1

    met1=met1.drop(columns='cycle')

    met2=met2.reset_index()
    met2=met2.drop(columns=['metabolite', 'index'])

    j=46

    while j < 240:
        met2.loc[j] = [j+2, 0]
        j=j+1
    met2.loc[-1] = [1, 0]
    met2 =met2.sort_values(by=['cycle'])
    met2=met2.drop(columns='cycle')
    met2=met2.reset_index()
    met2=met2.drop(columns='index')
    df2=comets.total_biomass

    df2[str(met)]=met1
    df2[str(met22)]=met2

    df2.columns=['time(h)','Ecoli1', 'Ecoli2', str(met), str(met22)]
    np.savetxt(r'biomass_vs_glc_D_e_ac_e_template.txt', df2.values, fmt='%s',delimiter='\t')

    #reading the biomass file as a dataframe and adding it the metabolite concentrations
    plt.ioff()
    fig, ax = plt.subplots()
    ax.axis(ymax=1.6)
    ax.set_xlabel('time (h)')
    ax.set_ylabel('biomass (g/L)')
    l1,=ax.plot(df2['time(h)']*0.1, df2['Ecoli1'], label='Ecoli1', color='k')
    l2,=ax.plot(df2['time(h)']*0.1, df2['Ecoli2'], linestyle='--', label='Ecoli2', color='k')
    ax2 = ax.twinx()
    ax2.set_ylabel('metabolite conc (mM)')
    ax2.axis(ymin=0, ymax=10)
    l3,=ax2.plot(df2['time(h)']*0.1, df2[str(met)], label=str(met), color='b')
    l4,=ax2.plot(df2['time(h)']*0.1, df2[str(met22)], label=str(met22), color='r')
    fig.tight_layout()
    plt.legend([l1, l2, l3, l4],["Ecoli1", "Ecoli2", str(met), str(met22)])

    #Saving the graph as a pdf
    plt.savefig('biomass_vs_glc_D_e_ac_e_template_plot1.pdf')
    return df2
    #df2.drop(df2.index, inplace=True)

def end_simulation_cycle(ac1, ac2, df2, met):
    #function that stablishes the endCyle after a certain stop condition

    iniBiomass=float(df2.at[0,'Ecoli1'])+float(df2.at[0, 'Ecoli2'])
    totGlc=float(df2.at[0,'glc__D_e'])
    endGlcCycle=0
    for row in df2.iterrows():
        endCycle=int(row[1][0])
        glcConc=float(row[1][3])
        acConc=float(row[1][4])
        if endGlcCycle==0 and glcConc==0.0:
            endGlcCycle=endCycle
        if glcConc==0.0 and acConc==0.0:
            break
        if glcConc==0.0 and ac1>=0 and ac2>=0:
            break

    return iniBiomass, endCycle, totGlc

def biomass_yield(df2,endCycle):
    # Function that compute final biomass as the maximum biomass of each strain

    finalBiomass1=0
    finalBiomass2=0
    count=0
    for row in df2.iterrows():
        if float(row[1][1])>finalBiomass1:
            finalBiomass1=float(row[1][1])
        if float(row[1][2])>finalBiomass2:
            finalBiomass2=float(row[1][2])
        if count > endCycle:
            break
        count+=1
    finalBiomass=finalBiomass1+finalBiomass2
    return finalBiomass

def compute_fitness(biomassYieldNew, MaximumYield, finalBiomass, iniBiomass):
    #Function that computes the fitness according to the fitness function 'yield'
    fitness=(biomassYieldNew/MaximumYield) # Normalizing Yield
    if(float(finalBiomass-iniBiomass) > 1.03):  # Given that with both strains with WT, total biomass=1.028
        fitness=0
    return fitness

def uptakeMetabolite(comets,layout):
    #Function that calculates the uptake of the metabolite 'ac_e'
    #Changing the metabolite needs changing numRxnMet; the position of the reaction 'Ex_metName' in the model
    flux = comets.fluxes
    numRxnMet=88
    uptakeMet=[]
    i=0
    while i < len(layout.models)+1:
        uptakeMet.append(float(flux.iat[i+2,numRxnMet+3]))
        i+=1
    return uptakeMet

def ecoliLongTermFLYCOP_oneConf(glu1,ac1,o21,glu2,ac2,o22, fitFunc='MaxYield_MinTime', dirPlot='',repeat=10, params_package=None, params_global=None, layout_file='/home/ana/FLYCOP/Scripts/Nuevo_layout.txt',met='glc__D_e', met2='ac_e'):
    """Primarily function of the simulation"""

    print("Fitness function:"+fitFunc)

    if not(os.path.exists('/home/ana/FLYCOP/Scripts/ecoli_1_tmp.cmd')):
        model_modifications(glu1, ac1, o21, glu2, ac2, o22)

    layout=initialize_layout(layout_file)

    params =initialize_params(params_package, params_global)

    #establish comets class
    comets = c.comets(layout, params)

    """You may need to adjust the comets.set_classpath object if your programs aren't in the default comets location"""

    if not(os.path.exists('IndividualRunsResults')):
        os.makedirs('IndividualRunsResults')
    totfitness=0
    sumTotBiomass=0
    sumTotYield=0
    fitnessList=[]

    # To repeat X times, due to random behaviour in COMETS:
    for i in range(repeat):

        comets.run(delete_files=True)

        #obtaining the results and writing them to csv
        comets.total_biomass.to_csv('Total_biomass_log_template.txt',sep='\t',index=False)
        comets.fluxes.to_csv('Flux_log_template.txt',sep='\t', index=False)
        comets.media.to_csv('Media_log_template.txt',sep='\t', index=False)

        #graphic part
        df2 =make_graph(comets, met, met2)

        # 7.- Compute fitness (measure to optimize):
        print('computing fitness...')

        # 7.1.- Determine endCycle: when glucose and acetate are exhausted
        iniBiomass, endCycle, totGlc = end_simulation_cycle(ac1, ac2, df2, met)

        # 7.2.- Compute first element fitness: maximize biomass yield
        finalBiomass = biomass_yield(df2,endCycle)

        biomassYieldNew=float((finalBiomass-iniBiomass)/(totGlc*0.1801559)) # molecular weigth of met1 per mmol

        # For normalizing yield
        MaximumYield=0.6

        # 7.3.- Compute second element fitnes: minimize time
        fitTime=1-(float(endCycle)/float(240))

        # 7.4.- Compute joint fitness, as a 50% each element.
        fitness= compute_fitness(biomassYieldNew, MaximumYield, finalBiomass, iniBiomass)

        # 7.5.- Calculate the uptake of ac_e
        uptakeMet = uptakeMetabolite(comets,layout)

        print(" Total biomass: "+str(round(finalBiomass,6))+" in cycle "+str(endCycle)+". Biomass yield="+str(round(biomassYieldNew,6)))
        totfitness=totfitness+fitness
        fitnessList.append(fitness)
        sumTotBiomass=sumTotBiomass+finalBiomass
        sumTotYield=sumTotYield+biomassYieldNew

        # Copy individual solution
        file='IndividualRunsResults/'+'biomass_vs_glc_D_ac_run'+str(i)+'_'+str(fitness)+'_'+str(endCycle)+'.pdf'
        shutil.move('biomass_vs_glc_D_e_ac_e_template_plot1.pdf',file)
        if(dirPlot != ''):
            file2=dirPlot+'biomass_vs_glc_D_ac_'+str(glu1)+'_'+str(ac1)+'_'+str(o21)+'_'+str(glu2)+'_'+str(ac2)+'_'+str(o22)+'_'+str(round(uptakeMet[0],1))+'_'+str(round(uptakeMet[1],1))+'_run'+str(i)+'_'+str(fitness)+'_'+str(endCycle)+'.pdf'
            shutil.copy(file,file2)

        file='IndividualRunsResults/'+'total_biomass_log_run'+str(i)+'.txt'
        shutil.move('Total_biomass_log_template.txt',file)
        file='IndividualRunsResults/'+'media_log_run'+str(i)+'.txt'
        shutil.move('Media_log_template.txt',file)
        file='IndividualRunsResults/'+'flux_log_run'+str(i)+'.txt'
        shutil.move('Flux_log_template.txt',file)

    avgfitness=totfitness/repeat
    sdfitness=statistics.stdev(fitnessList)
    avgBiomass=sumTotBiomass/repeat
    avgYield=sumTotYield/repeat
    print("Fitness_function\tconfiguration\tfitness\tsd\tavg.Biomass\tavg.Yield\tendCycle")
    print(fitFunc+"\t"+str(glu1)+','+str(ac1)+','+str(o21)+','+str(glu2)+','+str(ac2)+','+str(o22)+','+str(round(uptakeMet[0],1))+','+str(round(uptakeMet[1],1))+"\t"+str(round(avgfitness,6))+"\t"+str(sdfitness)+"\t"+str(round(avgBiomass,6))+"\t"+str(round(avgYield,6))+"\t"+str(endCycle))
    with open(dirPlot+"configurationsResults"+fitFunc+".txt", "a") as myfile:
        myfile.write("Fitness_function\tconfiguration\tfitness\tsd\tavg.Biomass\tavg.Yield\tendCycle\n")
        myfile.write(fitFunc+"\t"+str(glu1)+','+str(ac1)+','+str(o21)+','+str(glu2)+','+str(ac2)+','+str(o22)+','+str(round(uptakeMet[0],1))+','+str(round(uptakeMet[1],1))+"\t"+str(round(avgfitness,6))+"\t"+str(sdfitness)+"\t"+str(round(avgBiomass,6))+"\t"+str(round(avgYield,6))+"\t"+str(endCycle)+"\n")
    print("Avg.fitness(sd):\t"+str(avgfitness)+"\t"+str(sdfitness)+"\n")
    if(sdfitness>0.1):
        avgfitness=0.0
    try:
        os.remove('/home/ana/FLYCOP/Scripts/ecoli_1_tmp.cmd')
    except FileNotFoundError:
        pass
    try:
        os.remove('/home/ana/FLYCOP/Scripts/ecoli_2_tmp.cmd')
    except FileNotFoundError:
        pass
    return avgfitness,sdfitness