# -*- coding: utf-8 -*-

"""

Created on Mon May  1 19:41:07 2023

@author: Sen

"""

import os

import sys

import subprocess

import hashlib

import warnings

import platform

import csv

import numpy as np

from tqdm import tqdm

import argparse

import torch

from transformers import AutoTokenizer, GPT2LMHeadModel

import shutil

from openbabel import openbabel

import logging

import time

import subprocess

import threading

import os

import signal

import psutil

import os

class Command(object):

    def __init__(self, cmd):

        self.cmd = cmd

        self.process = None

    def run(self, timeout):

        def target():

            try:

                if os.name == 'posix':  # Unix/Linux/Mac

                    self.process = subprocess.Popen(self.cmd, shell=True, stderr=subprocess.DEVNULL,preexec_fn=os.setsid)

                else:  # Windows

                    self.process = subprocess.Popen(self.cmd, shell=True, stderr=subprocess.DEVNULL)

                self.process.communicate()

            except Exception:

                pass

        thread = threading.Thread(target=target)

        thread.start()

        thread.join(timeout)

        if thread.is_alive():

            if os.name == 'posix':  # Unix/Linux/Mac

                os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)

            else:  # Windows

                parent = psutil.Process(self.process.pid)

                for child in parent.children(recursive=True):

                    child.kill()

                parent.kill()

            thread.join()

        return self.process.returncode if self.process else None

class LigandPostprocessor:

    def __init__(self, path):

        self.hash_ligand_mapping = {}

        self.output_path = path  # Output directory for SDF files

        self.load_mapping()

    def load_mapping(self):

        mapping_file = os.path.join(output_path, 'hash_ligand_mapping.csv')

        if os.path.exists(mapping_file):

            print("Found existed mapping file, now reading ...")

            with open(mapping_file, 'r') as f:

                reader = csv.reader(f)

                for row in reader:

                    self.hash_ligand_mapping[row[0]] = row[1]

    # Define a function to save the hash-ligand mapping to a file

    def save_mapping(self):

        mapping_file = os.path.join(output_path, 'hash_ligand_mapping.csv')

        with open(mapping_file, 'w', newline='') as f:

            writer = csv.writer(f)

            for ligand_hash, ligand in self.hash_ligand_mapping.items():

                writer.writerow([ligand_hash, ligand])

    # Define a function to filter out empty SDF files

    def filter_sdf(self, hash_ligand_mapping_per_batch):

        print("Filtering sdf ...")

        ligand_hash_list = list(hash_ligand_mapping_per_batch.keys())

        mapping_per_match = hash_ligand_mapping_per_batch.copy()

        for ligand_hash in tqdm(ligand_hash_list):

            filepath = os.path.join(self.output_path, ligand_hash + '.sdf')            

            if os.path.getsize(filepath) < 2*1024:  #2kb

                try:

                    os.remove(filepath)

                    #mapping_per_match.pop(ligand_hash)

                except Exception:

                    print(filepath)

                mapping_per_match.pop(ligand_hash)    

        return mapping_per_match

    # Define a function to generate SDF files from a list of ligand SMILES using OpenBabel

    def to_sdf(self, ligand_list_per_batch):

        print("Converting to sdf ...")

        hash_ligand_mapping_per_batch = {}

        for ligand in tqdm(ligand_list_per_batch):  

            obConversion = openbabel.OBConversion()

            obConversion.SetInAndOutFormats("smi", "smi")

            mol = openbabel.OBMol()

            if not obConversion.ReadString(mol, ligand):

                continue  # Skip invalid SMILES

            num_atoms = sum(1 for atom in openbabel.OBMolAtomIter(mol) if atom.GetAtomicNum() != 1)

            if min_atoms is not None and num_atoms < min_atoms:

                continue  # Skip molecules with too few non-hydrogen atoms

            if max_atoms is not None and num_atoms > max_atoms:

                continue  # Skip molecules with too many non-hydrogen atoms

            ligand_hash = hashlib.sha1(ligand.encode()).hexdigest()

            if ligand_hash not in self.hash_ligand_mapping.keys():

                filepath = os.path.join(self.output_path , ligand_hash + '.sdf')

                if platform.system() == "Windows":

                    cmd = "obabel -:" + ligand + " -osdf -O " + filepath + " --gen3d --forcefield mmff94"

                elif platform.system() == "Linux":

                    obabel_path = shutil.which('obabel')

                    cmd = f"{obabel_path} -:'{ligand}' -osdf -O '{filepath}' --gen3d --forcefield mmff94"

                else:pass

                try:

                    command = Command(cmd)

                    return_code = command.run(timeout=10)

                    if return_code != 0:  # Check the return value

                        #print(f"Command execution failed with return code: {return_code}")

                        continue  # Skip the current iteration if the command execution failed

                except Exception:

                    time.sleep(1)

                    continue

                if os.path.exists(filepath):

                    hash_ligand_mapping_per_batch[ligand_hash] = ligand  # Add the hash-ligand mapping to the dictionary

        self.hash_ligand_mapping.update(self.filter_sdf(hash_ligand_mapping_per_batch))

    def delete_empty_files(self):

    # 遍历指定目录及其子目录中的所有文件

        for foldername, subfolders, filenames in os.walk(self.output_path):

            for filename in filenames:

                file_path = os.path.join(foldername, filename)

                # 如果文件大小为0，则删除该文件

                if os.path.getsize(file_path) < 2*1024:  #2kb

                    try:

                        os.remove(file_path)

                        print(f'Deleted {file_path}')

                    except Exception:

                        pass 

    def check_sdf(self):

        file_list = os.listdir(self.output_path)

        sdf_file_list = [x for x in file_list if x[-4:]=='sdf']

        for filename in sdf_file_list:

            hash_ = filename[:-4]

            if hash_ not in self.hash_ligand_mapping.keys():

                filepath = os.path.join(self.output_path,filename)

                try:

                    os.remove(filepath)

                    print('remove ' + filepath)

                except Exception:

                    pass

            else:pass    

def about():

    print("""

  _____                    _____ _____ _______ 

 |  __ \                  / ____|  __ \__   __|

 | |  | |_ __ _   _  __ _| |  __| |__) | | |   

 | |  | | '__| | | |/ _` | | |_ |  ___/  | |   

 | |__| | |  | |_| | (_| | |__| | |      | |   

 |_____/|_|   \__,_|\__, |\_____|_|      |_|   

                     __/ |                     

                    |___/                      

 A generative drug design model based on GPT2

    """)

# Function to read in FASTA file

def read_fasta_file(file_path):

    with open(file_path, 'r') as f:

        sequence = []

        for line in f:

            line = line.strip()

            if not line.startswith('>'):

                sequence.append(line)

        protein_sequence = ''.join(sequence)

    return protein_sequence

if __name__ == "__main__":

    about()

    warnings.filterwarnings('ignore')

    if platform.system() == "Linux":

        os.environ["TOKENIZERS_PARALLELISM"] = "false"

    #Sometimes, using Hugging Face may require a proxy.

    #os.environ["http_proxy"] = "http://your.proxy.server:port"

    #os.environ["https_proxy"] = "http://your.proxy.server:port"

    # Set up command line argument parsing

    parser = argparse.ArgumentParser()

    parser.add_argument('-p','--pro_seq', type=str, default=None, help='Input a protein amino acid sequence. Default value is None. Only one of -p and -f should be specified.')

    parser.add_argument('-f','--fasta', type=str, default=None, help='Input a FASTA file. Default value is None. Only one of -p and -f should be specified.')

    parser.add_argument('-l','--ligand_prompt', type=str, default='', help='Input a ligand prompt. Default value is an empty string.')

    parser.add_argument('-e','--empty_input', action='store_true', default=False, help='Enable directly generate mode.')

    parser.add_argument('-n','--number',type=int, default=100, help='At least how many molecules will be generated. Default value is 100.')

    parser.add_argument('-d','--device',type=str, default='cuda', help="Hardware device to use. Default value is 'cuda'.")

    parser.add_argument('-o','--output', type=str, default='./ligand_output/', help="Output directory for generated molecules. Default value is './ligand_output/'.")

    parser.add_argument('-b','--batch_size', type=int, default=16, help="How many molecules will be generated per batch. Try to reduce this value if you have low RAM. Default value is 16.")

    parser.add_argument('-t','--temperature', type=float, default=1.0, help="Adjusts the randomness of text generation; higher values produce more diverse outputs. Default value is 1.0.")

    parser.add_argument('--top_k', type=int, default=9, help='The number of highest probability tokens to consider for top-k sampling. Defaults to 9.')

    parser.add_argument('--top_p', type=float, default=0.9, help='The cumulative probability threshold (0.0 - 1.0) for top-p (nucleus) sampling. It defines the minimum subset of tokens to consider for random sampling. Defaults to 0.9.')

    parser.add_argument('--min_atoms', type=int, default=None, help='Minimum number of non-H atoms allowed for generation.')

    parser.add_argument('--max_atoms', type=int, default=35, help='Maximum number of non-H atoms allowed for generation. Default value is 35.')

    parser.add_argument('--no_limit', action='store_true', default=False, help='Disable the default max atoms limit.')

    args = parser.parse_args()

    protein_seq = args.pro_seq

    fasta_file = args.fasta

    ligand_prompt = args.ligand_prompt

    directly_gen = args.empty_input

    num_generated = args.number

    device = args.device

    output_path = args.output

    batch_generated_size = args.batch_size

    temperature_value = args.temperature

    top_k = args.top_k

    top_p = args.top_p

    min_atoms = args.min_atoms

    max_atoms = args.max_atoms

    if args.no_limit:

        max_atoms = None

    if (args.min_atoms is not None) and (args.max_atoms is not None) and (args.min_atoms > args.max_atoms):

        raise ValueError("Error: min_atoms cannot be greater than max_atoms.")

    if args.ligand_prompt:

        args.max_atoms = None

        args.min_atoms = None

        print("Note: --ligand_prompt is specified. --max_atoms and --min_atoms settings will be ignored.")

    logging.basicConfig(level=logging.CRITICAL)

    openbabel.obErrorLog.StopLogging()

    os.makedirs(output_path, exist_ok=True)

    # Check if the input is either a protein amino acid sequence or a FASTA file, but not both

    if directly_gen:

        print("Now in directly generate mode.")

        prompt = "<|startoftext|><P>"

        print(prompt)

    else:

        if (not protein_seq) and (not fasta_file):

            print("Error: Input is empty.")

            sys.exit(1)

        if protein_seq and fasta_file:

            print("Error: The input should be either a protein amino acid sequence or a FASTA file, but not both.")

            sys.exit(1)

        if fasta_file:

            protein_seq = read_fasta_file(fasta_file)

        # Generate a prompt for the model

        p_prompt = "<|startoftext|><P>" + protein_seq + "<L>"

        l_prompt = "" + ligand_prompt

        prompt = p_prompt + l_prompt

        print(prompt)

    # Load the tokenizer and the model

    tokenizer = AutoTokenizer.from_pretrained('liyuesen/druggpt')

    model = GPT2LMHeadModel.from_pretrained("liyuesen/druggpt")

    model.eval()

    device = torch.device(device)

    model.to(device)

    # Create a LigandPostprocessor object

    ligand_post_processor = LigandPostprocessor(output_path)

    # Generate molecules

    generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)

    generated = generated.to(device)

    batch_number = 0

    directly_gen_protein_list = []

    directly_gen_ligand_list = []

    attention_mask = generated.ne(tokenizer.pad_token_id).float()

    while len(ligand_post_processor.hash_ligand_mapping) < num_generated:

        generate_ligand_list = []

        batch_number += 1

        print(f"=====Batch {batch_number}=====")

        print("Generating ligand SMILES ...")

        sample_outputs = model.generate(

            generated,

            do_sample=True,

            top_k=top_k,

            max_length=1024,

            top_p=top_p,

            temperature=temperature_value,

            num_return_sequences=batch_generated_size, 

            attention_mask=attention_mask,

            pad_token_id = tokenizer.eos_token_id

        )

        for sample_output in sample_outputs:

            generate_ligand = tokenizer.decode(sample_output, skip_special_tokens=True).split('<L>')[1]

            generate_ligand_list.append(generate_ligand)

            if directly_gen:

                directly_gen_protein_list.append(tokenizer.decode(sample_output, skip_special_tokens=True).split('<L>')[0])

                directly_gen_ligand_list.append(generate_ligand)

        torch.cuda.empty_cache()

        ligand_post_processor.to_sdf(generate_ligand_list)

        ligand_post_processor.delete_empty_files()

        ligand_post_processor.check_sdf()

    if directly_gen:

        arr = np.array([directly_gen_protein_list, directly_gen_ligand_list])

        processed_ligand_list = ligand_post_processor.hash_ligand_mapping.values()

        with open(os.path.join(output_path, 'generate_directly.csv'), 'w', newline='') as f:

            writer = csv.writer(f)

            for index in range(arr.shape[1]):

                protein, ligand = arr[0, index], arr[1, index]

                if ligand in processed_ligand_list:

                    writer.writerow([protein, ligand])

    print("Saving mapping file ...")

    ligand_post_processor.save_mapping()

    print(f"{len(ligand_post_processor.hash_ligand_mapping)} molecules successfully generated!")

    print("Ligand Energy Minimization")

    result = subprocess.run(['python', 'druggpt_min_multi.py', '-d', output_path])