"""
아래의 코드는 이전에 dataframe_generation_by_group에서 생성된 그룹별 데이터프레임을 
input으로 하여 train image set을 생성하는 역할을 합니다. 생성된 이미지는 sequence의
길이별로 생성된 폴더에 길이에 따라 저장이 되며 모델학습의 gernalization을 위해 사용됩니다.

"""

"""
The code below serves to generate a train image set by inputting dataframes 
by group previously generated in the dataframe_generation_by_group. 
The images generated are stored in folders generated by sequence length, 
and are used for the gernalization of model learning.
"""


import rdkit
import os
import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem import Draw
from tqdm import tqdm
import click

import warnings
warnings.filterwarnings(action = 'ignore')
from rdkit import RDLogger
RDLogger.DisableLog('rdApp.*')

# path
path = '/new_images/' # Saving new image
data_path = '/train_dataset/'
if os.path.exists(path) == False : 
    os.mkdir(path)
else:
    pass


@click.command()
@click.option('--group', default=1, help='group number')

# Oragainizing the group by the number of core
# The number of data sameple for one group calculated as 
# the number of total data sample / the number of core
# ex) 111307682 / 31 = 3700000
# The number of core can be different by each environment 

def making_data(group):
    print("group number:", group)

    filtered_df = pd.read_csv(data_path +'filtered_df_group{}.csv'.format(group))
    data_len = len(filtered_df)
    # for idx in tqdm(range(len(filtered_df[filtered_df['group'] == group]))):
    for idx in range(data_len):
        # idx += 3700000 * (group-1)
        # print('idx:',idx, end='\r')
        # smiles = filtered_df[filtered_df['group'] == group]['SMILES'][idx]
        smiles = filtered_df['SMILES'][idx]
        smiles_g = Chem.MolFromSmiles(smiles)
        try:
            smile_plt = Draw.MolToImage(smiles_g, size = (300,300))

            # converting color
            # switching r channel with b channel
            im = np.array(smile_plt)
            r = im[:,:,0]
            g = im[:,:,1]
            b = im[:,:,2]
            convert = np.stack((b,g,r),axis=-1)
          
            
            # Making directory by the length of sequnece and saving 
            # making filename as "the length of smiels"_train_"index"
            # ex) 0020_train_4
            dir_name = len(filtered_df['SMILES'][idx])
            dir_name = str(dir_name).zfill(4)
        
            np.save(path  + str(dir_name) + '/' +'{0}_train_{1}_{2}.npy'.format(dir_name, group, idx), arr = convert)
            
        
            del (smile_plt)
        except ValueError:
            pass

        # checking the completion
        if idx % 100000 == 0 :
            print('group : {0}, index : {1}'.format(group, idx))

    del(filtered_df)

if __name__ == '__main__':
    making_data()
    

