utils.py

import copy
import json
import math
import numpy as np
import pandas as pd
import torch
from scipy.spatial import cKDTree
from rdkit import Chem
from rdkit.Chem import RWMol
from rdkit.Chem import Draw, AllChem
from rdkit.Chem import rdDepictor
import matplotlib.pyplot as plt
import re
##################### MolScribe#################################################################################### 
from typing import List
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Circle

     
COLORS = {
    u'c': '0.0,0.75,0.75', u'b': '0.0,0.0,1.0', u'g': '0.0,0.5,0.0', u'y': '0.75,0.75,0',
    u'k': '0.0,0.0,0.0', u'r': '1.0,0.0,0.0', u'm': '0.75,0,0.75'
}

#helper function
def view_box_center(bond_bbox,heavy_centers):
    fig, ax = plt.subplots(figsize=(10, 10))
    # 绘制矩形框 (boxes)
    for box in bond_bbox:
        x1, y1, x2, y2 = box
        width = x2 - x1
        height = y2 - y1
        rect = Rectangle((x1, y1), width, height, linewidth=1, edgecolor='blue', facecolor='none')
        ax.add_patch(rect)

    # 绘制圆形 (centers)
    for center in heavy_centers:
        x, y = center
        circle = Circle((x, y), radius=5, edgecolor='red', facecolor='none', linewidth=1)
        ax.add_patch(circle)

    # 设置坐标轴范围（根据数据自动调整）
    x_min = min(bond_bbox[:, 0].min(), heavy_centers[:, 0].min()) - 10
    x_max = max(bond_bbox[:, 2].max(), heavy_centers[:, 0].max()) + 10
    y_min = min(bond_bbox[:, 1].min(), heavy_centers[:, 1].min()) - 10
    y_max = max(bond_bbox[:, 3].max(), heavy_centers[:, 1].max()) + 10
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)

    # 设置标题和标签
    ax.set_title("Boxes and Centers")
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    # 显示图像
    plt.gca().set_aspect('equal', adjustable='box')  # 保持比例
    plt.grid(True, linestyle='--', alpha=0.7)

def molIDX(mol):
    for i, atom in enumerate(mol.GetAtoms()):
        atom.SetAtomMapNum(i)  #映射
        # print(i)
    return mol

def molIDX_del(mol):
    for i, atom in enumerate(mol.GetAtoms()):
        atom.SetAtomMapNum(0)  #映射
        print(i)
    return mol
from det_engine import ABBREVIATIONS



def Val_extract_atom_info(error_message):
    """
    从错误信息中提取 atomid, atomType 和 valence。
    :param error_message: 错误信息字符串
    :return: (atomid, atomType, valence) 元组
    """
    # 定义正则表达式来提取原子信息
    pattern = r"Explicit valence for atom # (\d+) (\w), (\d+)"
    pattern2 =r"Explicit valence for atom # (\d+) (\w) "
    # print(type(error_message))
    if not isinstance(error_message, type('strs')):
        error_message=str(error_message)
    match = re.search(pattern, error_message)
    match2 = re.search(pattern2, error_message)
    if match:
        # 提取 atomid, atomType 和 valence
        atomid = int(match.group(1))  # 原子索引
        atomType = match.group(2)     # 原子类型
        valence = int(match.group(3)) # 当前价态
        return atomid, atomType, valence
    elif match2:
        atomid = int(match2.group(1))  # 原子索引
        atomType = match2.group(2)     # 原子类型
        # valence = int(match2.group(3)) # 当前价态
        return atomid, atomType, None
        
    else:
        raise ValueError("无法从错误信息中提取原子信息")
    
def calculate_charge_adjustment(atom_symbol, current_valence):
    """
    计算需要调整的电荷，根据反馈的原子符号和当前价态。
    :param atom_symbol: 原子符号（如 "C"）
    :param current_valence: 当前价态（如 5）
    :return: 需要添加的电荷数（正数表示负电荷，负数表示正电荷）
    """
    if atom_symbol not in VALENCES:
        raise ValueError(f"未知的原子符号: {atom_symbol}")

    # 查找该元素的最大价态
    max_valence = max(VALENCES[atom_symbol])
    if current_valence is None:
        current_valence=max_valence
    # 如果当前价态大于最大允许价态，需要调整电荷
    if current_valence > max_valence:
        # 需要添加的负电荷数
        charge_adjustment = current_valence - max_valence
        return charge_adjustment 
    else:
        # 当前价态已经符合最大允许价态，不需要调整
        return 0

from rdkit.Chem import rdchem, RWMol, CombineMols

def expandABB(mol,ABBREVIATIONS, placeholder_atoms):
    mols = [mol]
    # **第三步: 替换 * 并合并官能团**
    # 逆序遍历 placeholder_atoms，确保删除后不会影响后续索引
    for idx in sorted(placeholder_atoms.keys(), reverse=True):
        group = placeholder_atoms[idx]  # 获取官能团名称
        # print(idx, group)
        submol = Chem.MolFromSmiles(ABBREVIATIONS[group].smiles)  # 获取官能团的子分子
        submol_rw = RWMol(submol)  # 让 submol 变成可编辑的 RWMol
        anchor_atom_idx = 0  # 选择 `submol` 的第一个原子作为连接点 as defined in ABBREVIATIONS
        # **1. 复制主分子**
        new_mol = RWMol(mol)
        # **2. 计算 `*` 在 `new_mol` 中的索引**
        placeholder_idx = idx
        # **3. 记录 `*` 原子的邻居**
        neighbors = [nb.GetIdx() for nb in new_mol.GetAtomWithIdx(placeholder_idx).GetNeighbors()]
        # **4. 断开 `*` 的所有键**
        bonds_to_remove = []  # 记录要断开的键
        for bond in new_mol.GetBonds():
            if bond.GetBeginAtomIdx() == placeholder_idx or bond.GetEndAtomIdx() == placeholder_idx:
                bonds_to_remove.append((bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()))
        for bond in bonds_to_remove:
            new_mol.RemoveBond(bond[0], bond[1])
        # **5. 删除 `*` 原子**
        new_mol.RemoveAtom(placeholder_idx)
        # **6. 重新计算 `neighbors`（删除后索引变化）**
        new_neighbors = []
        for neighbor in neighbors:
            if neighbor < placeholder_idx:
                new_neighbors.append(neighbor)
            else:
                new_neighbors.append(neighbor - 1)  # 因为删除了一个原子，所有索引 -1
        # **7. 合并 `submol`**
        new_mol = RWMol(CombineMols(new_mol, submol_rw))

        # **8. 计算 `submol` 的第一个原子在合并后的位置**
        new_anchor_idx = new_mol.GetNumAtoms() - len(submol_rw.GetAtoms()) + anchor_atom_idx

        # **9. 重新连接官能团**
        for neighbor in new_neighbors:
            # print(neighbor, new_anchor_idx, "!!")
            new_mol.AddBond(neighbor, new_anchor_idx, Chem.BondType.SINGLE)
            a1=new_mol.GetAtomWithIdx(neighbor)
            a2=new_mol.GetAtomWithIdx(new_anchor_idx)
            a1.SetNumRadicalElectrons(0)
            a2.SetNumRadicalElectrons(0)## 将自由基电子数设为 0,as has added new bond
        # **10. 更新主分子**
        mol = new_mol
        mols.append(mol)
    # # 遍历分子中的每个原子
    # for atom in mols[-1].GetAtoms(): NOTE considering original image has the RadicalElectrons
    #     atom_idx = atom.GetIdx()  # 原子索引
    #     radical_electrons = atom.GetNumRadicalElectrons()  # 自由基电子数
    #     if radical_electrons > 0:
    #         # print(f"原子 {atom_idx} 存在自由基，自由基电子数: {radical_electrons}\n current NumExplicitHs: {atom.GetNumExplicitHs()}")
    #         # 消除自由基：通过添加氢原子调整价态
    #         atom.SetNumRadicalElectrons(0)  # 将自由基电子数设为 0,as has added bond
    #         # atom.SetNumExplicitHs(atom.GetNumExplicitHs() + radical_electrons) 
    Chem.SanitizeMol(mols[-1])
    # 输出修改后的分子 SMILES
    modified_smiles = Chem.MolToSmiles(mols[-1])
    # print(f"修改后的分子 SMILES: {modified_smiles}")            
    return mols[-1], modified_smiles

################################################################################################################################################################
def output_to_smiles(output,idx_to_labels,bond_labels,result):#this will output * without abbre version
    #only output smiles with *
    x_center = (output["boxes"][:, 0] + output["boxes"][:, 2]) / 2
    y_center = (output["boxes"][:, 1] + output["boxes"][:, 3]) / 2

    center_coords = torch.stack((x_center, y_center), dim=1)

    output = {'bbox':         output["boxes"].to("cpu").numpy(),
                'bbox_centers': center_coords.to("cpu").numpy(),
                'scores':       output["scores"].to("cpu").numpy(),
                'pred_classes': output["labels"].to("cpu").numpy()}
    

    atoms_list, bonds_list,charge = bbox_to_graph_with_charge(output,
                                                idx_to_labels=idx_to_labels,
                                                bond_labels=bond_labels,
                                                result=result)
    smiles, mol= mol_from_graph_with_chiral(atoms_list, bonds_list,charge)
    abc=[atoms_list, bonds_list,charge ]
    
    if isinstance(smiles, type(None)):
        print(f"get atoms_list problems")
        # smiles, mol=None,None
    elif isinstance(atoms_list,type(None)):
        print(f"get atoms_list problems")
        # smiles, mol=None,None
    # else:
    #     smiles, mol=smiles_mol
    return abc,smiles,mol,output


def output_to_smiles2(output,idx_to_labels,bond_labels,result):#this will output * without abbre version
    #only output smiles with *
    x_center = (output["boxes"][:, 0] + output["boxes"][:, 2]) / 2
    y_center = (output["boxes"][:, 1] + output["boxes"][:, 3]) / 2

    center_coords = torch.stack((x_center, y_center), dim=1)

    output = {'bbox':         output["boxes"].to("cpu").numpy(),
                'bbox_centers': center_coords.to("cpu").numpy(),
                'scores':       output["scores"].to("cpu").numpy(),
                'pred_classes': output["labels"].to("cpu").numpy()}
    

    atoms_list, bonds_list,charge = bbox_to_graph_with_charge(output,
                                                idx_to_labels=idx_to_labels,
                                                bond_labels=bond_labels,
                                                result=result)
    smiles, mol= mol_from_graph_with_chiral(atoms_list, bonds_list,charge)
    abc=[atoms_list, bonds_list,charge ]
    if isinstance(smiles, type(None)):
        print(f"get atoms_list problems")
        # smiles, mol=None,None
    elif isinstance(atoms_list,type(None)):
        print(f"get atoms_list problems")
        # smiles, mol=None,None
    # else:
    #     smiles, mol=smiles_mol
    return abc,smiles,mol,output



def bbox_to_graph(output, idx_to_labels, bond_labels,result):
    
    # calculate atoms mask (pred classes that are atoms/bonds)
    atoms_mask = np.array([True if ins not in bond_labels else False for ins in output['pred_classes']])

    # get atom list
    atoms_list = [idx_to_labels[a] for a in output['pred_classes'][atoms_mask]]

    # if len(result) !=0 and 'other' in atoms_list:
    #     new_list = []
    #     replace_index = 0
    #     for item in atoms_list:
    #         if item == 'other':
    #             new_list.append(result[replace_index % len(result)])
    #             replace_index += 1
    #         else:
    #             new_list.append(item)
    #     atoms_list = new_list

    atoms_list = pd.DataFrame({'atom': atoms_list,
                            'x':    output['bbox_centers'][atoms_mask, 0],
                            'y':    output['bbox_centers'][atoms_mask, 1]})

    # in case atoms with sign gets detected two times, keep only the signed one
    for idx, row in atoms_list.iterrows():
        if row.atom[-1] != '0':
            if row.atom[-2] != '-':#assume charge value -9~9
                overlapping = atoms_list[atoms_list.atom.str.startswith(row.atom[:-1])]
            else:
                overlapping = atoms_list[atoms_list.atom.str.startswith(row.atom[:-2])]

            kdt = cKDTree(overlapping[['x', 'y']])
            dists, neighbours = kdt.query([row.x, row.y], k=2)
            if dists[1] < 7:
                atoms_list.drop(overlapping.index[neighbours[1]], axis=0, inplace=True)

    bonds_list = []

    # get bonds
    for bbox, bond_type, score in zip(output['bbox'][np.logical_not(atoms_mask)],
                                    output['pred_classes'][np.logical_not(atoms_mask)],
                                    output['scores'][np.logical_not(atoms_mask)]):
         
        # if idx_to_labels[bond_type] == 'SINGLE':
        if idx_to_labels[bond_type] in ['-','SINGLE', 'NONE', 'ENDUPRIGHT', 'BEGINWEDGE', 'BEGINDASH', 'ENDDOWNRIGHT']:
            _margin = 5
        else:
            _margin = 8

        # anchor positions are _margin distances away from the corners of the bbox.
        anchor_positions = (bbox + [_margin, _margin, -_margin, -_margin]).reshape([2, -1])
        oposite_anchor_positions = anchor_positions.copy()
        oposite_anchor_positions[:, 1] = oposite_anchor_positions[:, 1][::-1]

        # Upper left, lower right, lower left, upper right
        # 0 - 1, 2 - 3
        anchor_positions = np.concatenate([anchor_positions, oposite_anchor_positions])

        # get the closest point to every corner
        atoms_pos = atoms_list[['x', 'y']].values
        kdt = cKDTree(atoms_pos)
        dists, neighbours = kdt.query(anchor_positions, k=1)

        # check corner with the smallest total distance to closest atoms
        if np.argmin((dists[0] + dists[1], dists[2] + dists[3])) == 0:
            # visualize setup
            begin_idx, end_idx = neighbours[:2]
        else:
            # visualize setup
            begin_idx, end_idx = neighbours[2:]

        #NOTE  this proces may lead self-bonding for one atom
        if begin_idx != end_idx:# avoid self-bond
            bonds_list.append((begin_idx, end_idx, idx_to_labels[bond_type], idx_to_labels[bond_type], score))
        else:
            continue
    # return atoms_list.atom.values.tolist(), bonds_list
    return atoms_list, bonds_list


def calculate_distance(coord1, coord2):
    # Calculate Euclidean distance between two coordinates
    return math.sqrt((coord1[0] - coord2[0])**2 + (coord1[1] - coord2[1])**2)

def assemble_atoms_with_charges(atom_list, charge_list):
    used_charge_indices=set()
    atom_list = atom_list.reset_index(drop=True)
    # atom_list['atom'] = atom_list['atom'] + '0'
    kdt = cKDTree(atom_list[['x','y']])
    for i, charge in charge_list.iterrows():
        if i in used_charge_indices:
            continue
        charge_=charge['charge']
        # if charge_=='1':charge_='+'
        dist, idx_atom=kdt.query([charge_list.x[i],charge_list.y[i]], k=1)
        # atom_str=atom_list.loc[idx_atom,'atom'] 
        if idx_atom not in atom_list.index:
            print(f"Warning: idx_atom {idx_atom} is out of range for atom_list.")
            continue  # 跳过当前循环迭代
        atom_str = atom_list.iloc[idx_atom]['atom']
        if atom_str=='*':
            atom_=atom_str + charge_
        else:
            try:
                atom_ = re.findall(r'[A-Za-z*]+', atom_str)[0] + charge_
            except Exception as e:
                print(atom_str,charge_,charge_list)
                print(f"@assemble_atoms_with_charges\n {e}\n{atom_list}")
                atom_=atom_str + charge_
        atom_list.loc[idx_atom,'atom']=atom_

    return atom_list



def assemble_atoms_with_charges2(atom_list, charge_list, max_distance=10):
    used_charge_indices = set()

    for idx, atom in atom_list.iterrows():
        atom_coord = atom['x'],atom['y']
        atom_label = atom['atom']
        closest_charge = None
        min_distance = float('inf')

        for i, charge in charge_list.iterrows():
            if i in used_charge_indices:
                continue

            charge_coord = charge['x'],charge['y']
            charge_label = charge['charge']

            distance = calculate_distance(atom_coord, charge_coord)
            #NOTE how t determin this max_distance, dependent on image size??
            if distance <= max_distance and distance < min_distance:
                closest_charge = charge
                min_distance = distance

        
        if closest_charge is not None:
            if closest_charge['charge'] == '1':
                charge_ = '+'
            else:
                charge_ = closest_charge['charge']
            atom_ = atom['atom'] + charge_

            # atom['atom'] = atom_
            atom_list.loc[idx,'atom'] = atom_
            used_charge_indices.add(tuple(charge))

        else:
            # atom['atom'] = atom['atom'] + '0'
            atom_list.loc[idx,'atom'] = atom['atom'] + '0'

    return atom_list



def bbox_to_graph_with_charge(output, idx_to_labels, bond_labels,result):
    
    bond_labels_pre=bond_labels
    # charge_labels = [18,19,20,21,22]#make influence
    atoms_mask = np.array([True if ins not in bond_labels and ins not in charge_labels else False for ins in output['pred_classes']])

    try:
        # print(atoms_mask.shape)
        # print(output['pred_classes'].shape)
        atoms_list = [idx_to_labels[a] for a in output['pred_classes'][atoms_mask]]
        if isinstance(atoms_list, pd.Series) and atoms_list.empty:
            return None, None, None
        else:
            atoms_list = pd.DataFrame({'atom': atoms_list,
                                    'x':    output['bbox_centers'][atoms_mask, 0],
                                    'y':    output['bbox_centers'][atoms_mask, 1],
                                    'bbox':  output['bbox'][atoms_mask].tolist() ,#need this for */other converting
                                    'scores': output['scores'][atoms_mask].tolist(),
                                    })
    except Exception as e:
        print(output['pred_classes'][atoms_mask].dtype,output['pred_classes'][atoms_mask])#int64 [ 1  1  1  1  1  2  1 29]
        print(e)
        print(idx_to_labels)
        # print(output['pred_classes'][atoms_mask],"output['pred_classes'][atoms_mask]")
    
        
        # confict_atompaire=[]
        # # 如果你想计算所有边界框之间的IOU，考虑2个原子box 重叠 是否要删掉一个？？ TODO gmy version most box larger then normal mix the rules
        # for i in range(len(atoms_list)):
        #     for j in range(i + 1, len(atoms_list)):
        #         iou_value = calculate_iou(atoms_list.bbox[i], atoms_list.bbox[j])
        #         if iou_value !=0:
        #             # print(f"IOU between box {i} and box {j}: {iou_value}")
        #             if i !=j : confict_atompaire.append([i,j])
        # if len(confict_atompaire)>0:
        #     need_del=[]
        #     for i,j in confict_atompaire:
        #         ij_lab=[atoms_list.loc[i].atom,atoms_list.loc[j].atom ]
        #         ij_score=[atoms_list.loc[i].scores,atoms_list.loc[j].scores]
        #         # print(ij_lab,ij_score)
        #         if ij_lab==['C','N'] or ij_lab==['N','C']:
        #             if atoms_list.loc[i].atom =='C':
        #                 need_del.append(i)
        #             else:
        #                 need_del.append(j)
                # elif atoms_list.loc[i].scores> atoms_list.loc[j].scores:
                #         need_del.append(j)
                # elif atoms_list.loc[j].scores> atoms_list.loc[i].scores:
                #         need_del.append(i)  
            # print(need_del)          
            # atoms_list= atoms_list.drop(need_del)

    charge_mask = np.array([True if ins in charge_labels else False for ins in output['pred_classes']])
    charge_list = [idx_to_labels[a] for a in output['pred_classes'][charge_mask]]
    charge_list = pd.DataFrame({'charge': charge_list,
                        'x':    output['bbox_centers'][charge_mask, 0],
                        'y':    output['bbox_centers'][charge_mask, 1],
                        'scores':    output['scores'][charge_mask],
                        
                        })
    
    # print(charge_list,'\n@bbox_to_graph_with_charge')
    try:
        atoms_list['atom'] = atoms_list['atom']+'0'#add 0 
    except Exception as e:
        print(e)
        print(atoms_list['atom'],'atoms_list["atom"] @@ adding 0 ')
        

    if len(charge_list) > 0:
        atoms_list = assemble_atoms_with_charges(atoms_list,charge_list)
    # else:#Note Most mols are not formal charged 
        # atoms_list['atom'] = atoms_list['atom']+'0'
    # print(atoms_list,"after @@assemble_atoms_with_charges ")
    
    # in case atoms with sign gets detected two times, keep only the signed one
    for idx, row in atoms_list.iterrows():
        if row.atom[-1] != '0':
            try:
                if row.atom[-2] != '-':#assume charge value -9~9
                    overlapping = atoms_list[atoms_list.atom.str.startswith(row.atom[:-1])]
            except Exception as e:
                print(row.atom,"@rin case atoms with sign gets detected two times")
                print(e)
            else:
                overlapping = atoms_list[atoms_list.atom.str.startswith(row.atom[:-2])]

            kdt = cKDTree(overlapping[['x', 'y']])
            dists, neighbours = kdt.query([row.x, row.y], k=2)
            if dists[1] < 7:
                atoms_list.drop(overlapping.index[neighbours[1]], axis=0, inplace=True)

    bonds_list = []
    # get bonds
    # bond_mask=np.logical_not(np.logical_not(atoms_mask) | np.logical_not(charge_mask))
    bond_mask=np.logical_not(atoms_mask) & np.logical_not(charge_mask)
    for bbox, bond_type, score in zip(output['bbox'][bond_mask],  #NOTE also including the charge part
                                    output['pred_classes'][bond_mask],
                                    output['scores'][bond_mask]):
         
        # if idx_to_labels[bond_type] == 'SINGLE':
        if len(idx_to_labels)==23:
            if idx_to_labels[bond_type] in ['-','SINGLE', 'NONE', 'ENDUPRIGHT', 'BEGINWEDGE', 'BEGINDASH', 'ENDDOWNRIGHT']:
                _margin = 5
            else:
                _margin = 8
        elif len(idx_to_labels)==30:
            _margin=0#ad this version bond dynamicaly changed
        elif len(idx_to_labels)==24:
            _margin=0#ad this version bond dynamicaly changed
        # anchor positions are _margin distances away from the corners of the bbox.
        anchor_positions = (bbox + [_margin, _margin, -_margin, -_margin]).reshape([2, -1])
        oposite_anchor_positions = anchor_positions.copy()
        oposite_anchor_positions[:, 1] = oposite_anchor_positions[:, 1][::-1]

        # Upper left, lower right, lower left, upper right
        # 0 - 1, 2 - 3
        anchor_positions = np.concatenate([anchor_positions, oposite_anchor_positions])

        # get the closest point to every corner
        atoms_pos = atoms_list[['x', 'y']].values
        kdt = cKDTree(atoms_pos)
        dists, neighbours = kdt.query(anchor_positions, k=1)

        # check corner with the smallest total distance to closest atoms
        if np.argmin((dists[0] + dists[1], dists[2] + dists[3])) == 0:
            # visualize setup
            begin_idx, end_idx = neighbours[:2]
        else:
            # visualize setup
            begin_idx, end_idx = neighbours[2:]

        #NOTE  this proces may lead self-bonding for one atom
        if begin_idx != end_idx: 
            if bond_type in bond_labels:# avoid self-bond
                bonds_list.append((begin_idx, end_idx, idx_to_labels[bond_type], idx_to_labels[bond_type], score))
            else:
                print(f'this box may be charges box not bonds {[bbox, bond_type, score ]}')
        else:
            continue
    # return atoms_list.atom.values.tolist(), bonds_list
    # print(f"@box2graph: atom,bond nums:: {len(atoms_list)}, {len(bonds_list)}")
    return atoms_list, bonds_list,charge_list#dataframe, list

def parse_atom(node):
    s10 = [str(x) for x in range(10)]
    # Determine atom and formal charge
    if 'other' in node:
        a = '*'
        if '-' in node or '+' in node:
            fc = -1 if node[-1] == '-' else 1
        else:
            fc = int(node[-2:]) if node[-2:] in s10 else 0
    elif node[-1] in s10:
        if '-' in node or '+' in node:
            fc = -1 if node[-1] == '-' else 1
            a = node[:-1]
        else:
            a = node[:-1]
            fc = int(node[-1])
    elif node[-1] == '+':
        a = node[:-1]
        fc = 1
    elif node[-1] == '-':
        a = node[:-1]
        fc = -1
    else:
        a = node
        fc = 0
    return a, fc

#from engine

def iou_(box1, box2):
    """
    计算两个框的 IoU（Intersection over Union）。
    参数:
        box1, box2: [x1, y1, x2, y2] 格式的框坐标
        
    返回:
        float: IoU 值
    """
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - intersection
    return intersection / union if union > 0 else 0


def calculate_iou(bbox1, bbox2):
    # 提取坐标
    x_min1, y_min1, x_max1, y_max1 = bbox1
    x_min2, y_min2, x_max2, y_max2 = bbox2
    
    # 计算交集坐标
    x_min_inter = max(x_min1, x_min2)
    y_min_inter = max(y_min1, y_min2)
    x_max_inter = min(x_max1, x_max2)
    y_max_inter = min(y_max1, y_max2)
    
    # 计算交集面积
    inter_width = max(0, x_max_inter - x_min_inter)
    inter_height = max(0, y_max_inter - y_min_inter)
    inter_area = inter_width * inter_height
    
    # 计算两个框的面积
    area1 = (x_max1 - x_min1) * (y_max1 - y_min1)
    area2 = (x_max2 - x_min2) * (y_max2 - y_min2)
    
    # 计算并集面积
    union_area = area1 + area2 - inter_area
    
    # 计算 IoU
    iou = inter_area / union_area if union_area > 0 else 0
    
    # 判断关系并记录
    result = []
    if iou == 0:
        result.append("无重叠")
    elif iou > 0:
        result.append("有重叠")
        if iou == 1:
            result.append("完全重合")
        elif inter_area == area2:
            result.append("bbox1 包含 bbox2")
        elif inter_area == area1:
            result.append("bbox2 包含 bbox1")
    
    return iou, result, inter_area, union_area

def adjust_bbox1(large_bbox, small_bbox, bond_bbox):
    # 假设调整策略：扣除小的 atom bbox 和 bond box 的区域
    # 这里简单假设从较大 bbox 中移除小的区域，可能需要根据具体需求调整
    x_min_l, y_min_l, x_max_l, y_max_l = large_bbox
    x_min_s, y_min_s, x_max_s, y_max_s = small_bbox
    x_min_b, y_min_b, x_max_b, y_max_b = bond_bbox
    scaled_box= max([x_min_l,x_min_s,x_min_b]),max([y_min_l,y_min_s,y_min_b]),x_max_l, y_max_l
    return large_bbox
    # 示例调整：如果小的 bbox 和 bond box 在较大 bbox 内，缩小较大 bbox
    # if x_min_s > x_min_l and y_min_s > y_min_l:
    #     return [x_min_l, y_min_l, x_min_s, y_min_s]  # 示例：保留左上部分
    # return large_bbox  # 默认不调整


def nms_per_class(labels, boxes, scores, iou_thresh=0.5):
    """
    对每个类别应用 NMS，保留得分最高的框。
    参数:
        labels: numpy array，类别标签
        boxes: numpy array，框坐标 [x1, y1, x2, y2]
        scores: numpy array，得分
        iou_thresh: float，IoU 阈值
    返回:
        dict: 筛选后的输出
    """
    # 按类别分组
    unique_labels = np.unique(labels)
    kept_indices = []
    for label in unique_labels:
        # 筛选当前类别的框
        class_mask = labels == label
        class_indices = np.where(class_mask)[0]
        class_boxes = boxes[class_mask]
        class_scores = scores[class_mask]
        
        # 按得分从高到低排序
        order = np.argsort(class_scores)[::-1]
        class_boxes = class_boxes[order]
        class_scores = class_scores[order]
        class_indices = class_indices[order]
        
        # NMS
        keep = []
        while len(class_scores) > 0:
            # 保留得分最高的框
            keep.append(class_indices[0])
            if len(class_scores) == 1:
                break
            
            # 计算当前框与其他框的 IoU
            ious = np.array([calculate_iou(class_boxes[0], box) for box in class_boxes[1:]])
            # 保留 IoU 低于阈值的框
            keep_mask = ious < iou_thresh
            class_boxes = class_boxes[1:][keep_mask]
            class_scores = class_scores[1:][keep_mask]
            class_indices = class_indices[1:][keep_mask]
        
        kept_indices.extend(keep)
    
    # 根据保留的索引更新输出
    kept_indices = np.array(kept_indices)
    return {
        'labels': labels[kept_indices],
        'boxes': boxes[kept_indices],
        'scores': scores[kept_indices]
    }