main.py

import time
from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
import subprocess
from typing import List
from fragment_processor import fragmentize_molecule
from fastapi.responses import StreamingResponse
import torch
import pandas as pd
from rdkit import Chem
from rdkit.Chem import Descriptors, QED, Draw
from io import BytesIO
from PIL import Image
import io
from generate import GenerateRunner
from dataset import Dataset
from combine_mol import connect_constVar_try
import sascorer
app = FastAPI()


class Fragment(BaseModel):
    variable_smiles: str
    constant_smiles: str
    record_id: str
    normalized_smiles: str
    attachment_order: int

class FragmentResponse(BaseModel):
    fragments: List[Fragment]



class GenerateRequest(BaseModel):
    constSmiles: str
    varSmiles: str
    mainCls: str
    minorCls: str
    deltaValue: str
    num: int

class MoleculeOutput(BaseModel):
    smile: str
    molwt: float
    tpsa: float
    slogp: float
    sa: float
    qed: float

class Options:
    def __init__(self, **entries):
        self.__dict__.update(entries)

def calculate_descriptors(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    
    # molwt = Descriptors.MolWt(mol)
    # tpsa = Descriptors.TPSA(mol)
    # slogp = Descriptors.MolLogP(mol)
    # sa = sascorer.calculateScore(mol) 
    # qed = QED.qed(mol)
    molwt = round(Descriptors.MolWt(mol), 1)  # 保留 1 位小数
    tpsa = round(Descriptors.TPSA(mol), 1)
    slogp = round(Descriptors.MolLogP(mol), 1)
    sa = round(sascorer.calculateScore(mol), 1)
    qed = round(QED.qed(mol), 1)

    # 检查除法前是否为 0
    if tpsa == 0:
        print("Warning: TPSA is zero, skipping division.")
        some_ratio = None  
    else:
        some_ratio = molwt / tpsa  # 安全的除法操作
    
    return {"molwt": molwt, "tpsa": tpsa, "slogp": slogp, "sa": sa, "qed": qed}

def run_generate_runner(const_smiles, var_smiles, main_cls, minor_cls, delta_value, num_samples):
    # 初始化生成器的配置选项
    opt = {
        'batch_size': num_samples,
        'data_path' :  './',
        'decode_type' :  'multinomial',
        'dev_no' :  0,
        'epoch' :  20,
        'model_choice' :  'transformer',
        'model_path' :  './raw_pretrain_frag/checkpoint',
        'num_samples' :  50,
        'overwrite' :  True,
        # 'save_directory' :  './demo_gen',
        'test_file_name' :  'test_cut',
        'vocab_path' :  './'
    }

    # 将 opt 字典转换为 Options 对象
    opt = Options(**opt)
    print("--------------opt---------------")
    print(opt)
    runner = GenerateRunner(opt)

    # 创建数据
    data = {
        "constantSMILES": const_smiles,
        "fromVarSMILES": var_smiles,
        "main_cls": main_cls,
        "minor_cls": minor_cls,
        "Delta_Value": delta_value
    }
    
    # 创建 Dataset 实例
    test_data = pd.DataFrame([data])
    dataset = Dataset(test_data, vocabulary=runner.vocab, tokenizer=runner.tokenizer, prediction_mode=True)

    # 生成 SMILES
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=num_samples, shuffle=False, collate_fn=Dataset.collate_fn)
    result = []

    for batch in dataloader:
        src, source_length, _, src_mask, _, _, df = batch
        src = src.to(runner.device)
        src_mask = src_mask.to(runner.device)
        source_length = source_length.to(runner.device)  # 将 source_length 也移到同一设备
        smiles_list = runner.sample(
            model_choice="transformer", 
            model=runner.model, 
            src=src, 
            src_mask=src_mask, 
            source_length=source_length, 
            decode_type="multinomial", 
            num_samples=num_samples
        )
        
        # 计算每个 SMILES 的化学性质
        for smiles_group in smiles_list:
            for smile in smiles_group:  # smiles_group 是一个子列表
                #链接新分子
                newsmile=connect_constVar_try(const_smiles,smile)
                descriptors = calculate_descriptors(newsmile)
                if descriptors:
                    result.append({
                        "smile": newsmile,
                        "molwt": descriptors['molwt'],
                        "tpsa": descriptors['tpsa'],
                        "slogp": descriptors['slogp'],
                        "sa": descriptors['sa'],
                        "qed": descriptors['qed']
                    })


    return result


@app.get("/smiles2img")
async def smiles2img(smiles: str = Query(..., description="SMILES string of the molecule")):
    
    # 生成分子对象
    print("---开始生成分子图像--")
    print(smiles)
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {"error": "Invalid SMILES string"}

    # 创建绘图对象
    drawer = Draw.MolDraw2DCairo(200, 200)
    drawer.SetFontSize(1.0)
    drawer.DrawMolecule(mol)
    drawer.FinishDrawing()

    # 将绘制的图像转换为PIL图像对象
    pil_image = Image.open(io.BytesIO(drawer.GetDrawingText()))

    # 创建字节流
    strIO = BytesIO()
    pil_image.save(strIO, "PNG")
    strIO.seek(0)

    # 返回图像作为流
    return StreamingResponse(strIO, media_type="image/png")

@app.get("/fragmentize", response_model=FragmentResponse)
async def fragmentize(smiles: str = Query(..., description="SMILES string of the molecule")):
    try:
        fragment_df = fragmentize_molecule(smiles)
        fragments = fragment_df.to_dict(orient="records")
        return FragmentResponse(fragments=fragments)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"发生错误: {str(e)}")

@app.post("/generate", response_model=List[MoleculeOutput])
async def generate_molecules(request: GenerateRequest):
    start_time = time.time()  # 记录请求接受的时间
    try:
        # 调用 SMILES 生成逻辑
        print("/generate请求开始时间:", start_time)
        print("--------------/generate start---------------")
        result = run_generate_runner(request.constSmiles, request.varSmiles, request.mainCls, request.minorCls, request.deltaValue, request.num)
        end_time = time.time()  # 记录生成结束的时间
        duration = end_time - start_time  # 计算用时

        print("/generate请求结束时间:", end_time)
        print(f"请求处理用时: {duration:.2f}秒，本次处理分子数量为 {request.num}")
        return result
    except Exception as e:
        # 捕获异常并记录详细的错误信息，包括堆栈追踪
        error_message = f"Error occurred: {str(e)}"
        print(error_message)  # 打印到控制台，或者使用 logging 模块记录到日志文件
        raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")