#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :qsar_3d.py
@Description:       : 3D-QSAR modeling with SMILES and MIC data
@Date               :2024/09/28
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''

import pandas as pd
import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import rdPartialCharges
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression, SGDRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.tree import DecisionTreeRegressor
from sklearn.neural_network import MLPRegressor
from xgboost import XGBRegressor
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
import joblib
import os

# Function to calculate 3D-QSAR descriptors using Coulomb Matrix
def calculate_3dqsar_repr(SMILES, max_atoms=100, three_d=False):
    mol = Chem.MolFromSmiles(SMILES)  # 从SMILES表示创建分子对象
    mol = Chem.AddHs(mol)  # 添加氢原子
    if three_d:
        AllChem.EmbedMolecule(mol, AllChem.ETKDG())  # 计算3D坐标
    else:
        AllChem.Compute2DCoords(mol)  # 计算2D坐标
    natoms = mol.GetNumAtoms()  # 获取原子数量
    rdPartialCharges.ComputeGasteigerCharges(mol)  # 计算分子的Gasteiger电荷
    charges = np.array([float(atom.GetProp("_GasteigerCharge")) for atom in mol.GetAtoms()])  # 获取电荷值
    coords = mol.GetConformer().GetPositions()  # 获取原子坐标
    coulomb_matrix = np.zeros((max_atoms, max_atoms))  # 初始化库仑矩阵
    n = min(max_atoms, natoms)
    for i in range(n):  # 遍历原子
        for j in range(i, n):
            if i == j:
                coulomb_matrix[i, j] = 0.5 * charges[i] ** 2
            if i != j:
                delta = np.linalg.norm(coords[i] - coords[j])  # 计算原子间距离
                if delta != 0:
                    coulomb_matrix[i, j] = charges[i] * charges[j] / delta  # 计算库仑矩阵的元素值
                    coulomb_matrix[j, i] = coulomb_matrix[i, j]
    coulomb_matrix = np.where(np.isinf(coulomb_matrix), 0, coulomb_matrix)  # 处理无穷大值
    coulomb_matrix = np.where(np.isnan(coulomb_matrix), 0, coulomb_matrix)  # 处理NaN值
    return coulomb_matrix.reshape(max_atoms*max_atoms).tolist()  # 将库仑矩阵转换为列表并返回

if __name__ == '__main__':
    # Load your dataset
    df = pd.read_csv('/mnt/c/project/qsar/data/A_85.csv', sep=';')
    df_with_MIC = df[df['Standard Type'] == 'MIC']

    # Select relevant columns
    qsar_df = df_with_MIC[['Smiles', 'Standard Value']].copy()
    qsar_df.rename(columns={'Smiles': 'SMILES', 'Standard Value': 'TARGET'}, inplace=True)

    # Calculate 3D-QSAR descriptors for each molecule
    qsar_df['3dqsar_mr'] = qsar_df['SMILES'].apply(calculate_3dqsar_repr)

    # Remove rows where descriptor calculation failed
    qsar_df = qsar_df.dropna()

    # Split the data into training and testing sets
    train_df, test_df = train_test_split(qsar_df, test_size=0.2, random_state=42)

    # Convert the data to NumPy arrays
    train_x = np.array(train_df['3dqsar_mr'].tolist())
    train_y = np.array(train_df['TARGET'].tolist())
    test_x = np.array(test_df['3dqsar_mr'].tolist())
    test_y = np.array(test_df['TARGET'].tolist())

    # Define regressors to use
    regressors = [
        ("Linear Regression", LinearRegression()),
        ("Stochastic Gradient Descent", SGDRegressor(random_state=42)),
        ("K-Nearest Neighbors", KNeighborsRegressor()),
        ("Decision Tree", DecisionTreeRegressor(random_state=42)),
        ("Random Forest", RandomForestRegressor(random_state=42)),
        ("XGBoost", XGBRegressor(eval_metric="rmse", random_state=42)),
        ("Multi-layer Perceptron", MLPRegressor(hidden_layer_sizes=(128, 64, 32), activation='relu', solver='adam', max_iter=10000, random_state=42)),
    ]

    # Initialize results dictionary
    results = {}

    # Train and evaluate each regressor
    for name, regressor in regressors:
        regressor.fit(train_x, train_y)
        pred_y = regressor.predict(test_x)
        mse = mean_squared_error(test_y, pred_y)
        r2 = r2_score(test_y, pred_y)
        results[f"3D-QSAR-{name}"] = {"MSE": mse, "R2": r2}
        print(f"[3D-QSAR][{name}]\tMSE:{mse:.4f}\tR2:{r2:.4f}")
        
        # Save the trained model
        model_filename = f"3d_qsar_{name.replace(' ', '_').lower()}_model.pkl"
        joblib.dump(regressor, os.path.join(os.getcwd(), model_filename))
        print(f"Model saved to {model_filename}")

    # Sort results by R2 score
    sorted_results = sorted(results.items(), key=lambda x: x[1]["R2"], reverse=True)

    # Plot R2 Scores for Regressors
    plt.figure(figsize=(10, 6), dpi=300)
    plt.title("R2 Scores for 3D-QSAR Regressors")
    plt.barh([name for name, _ in sorted_results], [result['R2'] for _, result in sorted_results])
    plt.xlabel("R2 Score")
    plt.savefig('qsar_3D_r2_scores.png', dpi=300)