import gradio as gr
import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms
import torchvision.transforms.v2 as T
import io
import rdkit
from rdkit import Chem
from rdkit.Chem import Draw
from postprocessor import RTDETRPostProcessor
from utils import bbox_to_graph_with_charge, mol_from_graph_with_chiral

bond_labels = [13,14,15,16,17]
idx_to_labels = {0:'other',1:'C',2:'O',3:'N',4:'Cl',5:'Br',6:'S',7:'F',8:'B',
            9:'I',10:'P',11:'*',12:'Si',13:'NONE',14:'BEGINWEDGE',15:'BEGINDASH',
            16:'=',17:'#',18:'-4',19:'-2',20:'-1',21:'1',22:'+2',} #NONE is single ?


def image_to_tensor(image_path):
    # Open the image using PIL
    image = Image.open(image_path)
    w, h = image.size
    # print("width: {}, height: {}".format(w, h))
    # Define a transform to convert the image to a tensor and normalize it
    transform = transforms.Compose([
        # transforms.Grayscale(num_output_channels=1),  # Convert to grayscale (1 channel)
        T.Resize((640, 640)),  # Resize the image to 224x224
        T.ToImageTensor(),  # Convert to Tensor (C x H x W)
        T.ConvertDtype(dtype=torch.float32)
        # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Optional normalization for pretrained models
    ])
    
    # Apply the transform to the image
    tensor = transform(image)
    
    return tensor,w,h


def visualize_molecule(smiles):
    """
    使用RDKit将SMILES转换为分子结构图
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        img = Draw.MolToImage(mol)
        return img
    except:
        return None

def predict(input_image):
    """
    主要的推理函数
    """
    try:
        # 加载和初始化ONNX模型
        session = ort.InferenceSession("model.onnx")  # 替换为实际模型路径
        
        # 预处理图片
         # Example usage: #change thie image
        tensor,w,h = image_to_tensor(input_image)
        processed_image=tensor.unsqueeze(0)
        
        # 获取模型输入输出名称
        input_name = session.get_inputs()[0].name
        output_name = session.get_outputs()[0].name
        
        # 进行推理
        outputs = session.run([output_name], {input_name: processed_image})
        ori_size=torch.Tensor([w,h]).long().unsqueeze(0)
        postprocessor = RTDETRPostProcessor()
        result_ = postprocessor(outputs, ori_size)
        score_=result_[0]['scores']
        boxe_=result_[0]['boxes']
        label_=result_[0]['labels']
        selected_indices =score_ > 0.5
        output={
            'labels': label_[selected_indices],
            'boxes': boxe_[selected_indices],
            'scores': score_[selected_indices]
        }
        
        filtered_output_dict={image_path: output
        }
        
        
        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_df, bonds_list,charge_list =bbox_to_graph_with_charge(output, idx_to_labels=idx_to_labels,
                                                    bond_labels=bond_labels,  result=[])
        smiles,mol_rebuit=mol_from_graph_with_chiral(atoms_df, bonds_list,charge_list )
        
        # 使用RDKit生成分子结构图
        mol_image = visualize_molecule(smiles)
        
        if mol_image is None:
            return "无效的SMILES字符串", None
        
        return smiles, mol_image
        
    except Exception as e:
        return f"发生错误: {str(e)}", None

# 创建Gradio界面
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Text(label="SMILES字符串"),
        gr.Image(label="分子结构图")
    ],
    title="化学结构OCR",
    description="上传一张包含化学结构的图片，获取对应的SMILES表示和分子结构图。",
    examples=[
        ["example.jpg"]
    ]
)

# 启动应用
if __name__ == "__main__":
    iface.launch()