# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os 
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = '2'

import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader
from torch.autograd import Variable

import math, random, sys
import numpy as np
import argparse
from collections import deque
# import cPickle as pickle
import pickle

# from fast_jtnn import *
from hgraph import *
import rdkit

from cnn_finetune import make_model
import torch.nn as nn

from albumentations.pytorch import ToTensorV2
import albumentations as A
from albumentations import (
    HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,
    IAASharpen, IAAEmboss, RandomContrast, RandomBrightness, Flip, OneOf, Compose,
    RandomCrop, Normalize, Resize
)
import cv2
from tqdm import tqdm

lg = rdkit.RDLogger.logger() 
lg.setLevel(rdkit.RDLogger.CRITICAL)


# +
class ARGS:
    def __init__(self):
        self.train = './train_processed_bms_ver2/mol/'
        self.vocab = 'data/chembl/vocab.txt'
        self.atom_vocab = common_atom_vocab
        self.save_dir = 'vae_model/'
        self.load_model = None
        self.seed = 7
        
        self.rnn_type = 'LSTM'
        self.hidden_size = 250
        self.embed_size = 250
        self.batch_size = 32
        self.latent_size = 32
        self.depthT = 15
        self.depthG = 15
        self.diterT = 1
        self.diterG = 3
        self.dropout = 0.0
        
        self.lr = 1e-3
        self.clip_norm = 5.0
#         self.beta = 0.0
        self.step_beta = 0.001
        self.max_beta = 1.0
        self.warmup = 10000
        self.kl_anneal_iter = 2000
        
        self.epoch = 100
        self.anneal_rate = 0.9
        self.anneal_iter = 25000
        self.kl_anneal_iter = 2000
        self.print_iter = 50
        self.save_iter = 5000
        
        
args = ARGS()
# -

# ## Model

# +
if not torch.cuda.is_available():
    device = torch.device('cpu')
else:
    device = torch.device('cuda:0')
    
print (device)


# +
# class JTNNVAE_edit(JTNNVAE):
#     def __init__(self, vocab, hidden_size, latent_size, depthT, depthG):
#         super().__init__(vocab, hidden_size, latent_size, depthT, depthG)
    
#     def forward_edit(self, x_batch, beta, x_tree_vecs, x_mol_vecs):
#         x_batch, x_jtenc_holder, x_mpn_holder, x_jtmpn_holder = x_batch
#         _, x_tree_mess, _ = self.encode(x_jtenc_holder, x_mpn_holder)
#         z_tree_vecs,tree_kl = self.rsample(x_tree_vecs, self.T_mean, self.T_var)
#         z_mol_vecs,mol_kl = self.rsample(x_mol_vecs, self.G_mean, self.G_var)

#         kl_div = tree_kl + mol_kl
#         word_loss, topo_loss, word_acc, topo_acc = self.decoder(x_batch, z_tree_vecs)
#         assm_loss, assm_acc = self.assm(x_batch, x_jtmpn_holder, z_mol_vecs, x_tree_mess)

#         return word_loss + topo_loss + assm_loss + beta * kl_div, kl_div.item(), word_acc, topo_acc, assm_acc

def make_cuda(tensors):
    tree_tensors, graph_tensors = tensors
    make_tensor = lambda x: x if type(x) is torch.Tensor else torch.tensor(x)
    tree_tensors = [make_tensor(x).cuda().long() for x in tree_tensors[:-1]] + [tree_tensors[-1]]
    graph_tensors = [make_tensor(x).cuda().long() for x in graph_tensors[:-1]] + [graph_tensors[-1]]
    return tree_tensors, graph_tensors

class HierVAE_edit(HierVAE):

    def __init__(self, args):
        super().__init__(args)

#     def rsample(self, z_vecs, W_mean, W_var, perturb=True):
#         batch_size = z_vecs.size(0)
#         z_mean = W_mean(z_vecs)
#         z_log_var = -torch.abs( W_var(z_vecs) )
#         kl_loss = -0.5 * torch.sum(1.0 + z_log_var - z_mean * z_mean - torch.exp(z_log_var)) / batch_size
#         epsilon = torch.randn_like(z_mean).cuda()
#         z_vecs = z_mean + torch.exp(z_log_var / 2) * epsilon if perturb else z_mean
#         return z_vecs, kl_loss

#     def sample(self, batch_size, greedy):
#         root_vecs = torch.randn(batch_size, self.latent_size).cuda()
#         return self.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=greedy, max_decode_step=150)

#     def reconstruct(self, batch):
#         graphs, tensors, _ = batch
#         tree_tensors, graph_tensors = tensors = make_cuda(tensors)
#         root_vecs, tree_vecs, _, graph_vecs = self.encoder(tree_tensors, graph_tensors)

#         root_vecs, root_kl = self.rsample(root_vecs, self.R_mean, self.R_var, perturb=False)
#         return self.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=True, max_decode_step=150)
       
    def forward(self, graphs, tensors, orders, beta, perturb_z=True):
        tree_tensors, graph_tensors = tensors = make_cuda(tensors)

        root_vecs, tree_vecs, _, graph_vecs = self.encoder(tree_tensors, graph_tensors)
        root_vecs, root_kl = self.rsample(root_vecs, self.R_mean, self.R_var, perturb_z)
        kl_div = root_kl

        loss, wacc, iacc, tacc, sacc = self.decoder((root_vecs, root_vecs, root_vecs), graphs, tensors, orders)
        return loss + beta * kl_div, kl_div.item(), wacc, iacc, tacc, sacc
    
    def forward_edit(self, graphs, tensors, orders, root_vecs, beta, perturb_z=True):
        tree_tensors, graph_tensors = tensors = make_cuda(tensors)

#         root_vecs, tree_vecs, _, graph_vecs = self.encoder(tree_tensors, graph_tensors)
        root_vecs, root_kl = self.rsample(root_vecs, self.R_mean, self.R_var, perturb_z)
        kl_div = root_kl

        loss, wacc, iacc, tacc, sacc = self.decoder((root_vecs, root_vecs, root_vecs), graphs, tensors, orders)
        return loss + beta * kl_div, kl_div.item(), wacc, iacc, tacc, sacc


# +
vocab = [x.strip("\r\n ").split() for x in open(args.vocab)] 
args.vocab = PairVocab(vocab)

model = HierVAE_edit(args).cuda()
# print (model)

print("Model #Params: %dK" % (sum([x.nelement() for x in model.parameters()]) / 1000,))

for param in model.parameters():
    if param.dim() == 1:
        nn.init.constant_(param, 0)
    else:
        nn.init.xavier_normal_(param)

optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.ExponentialLR(optimizer, args.anneal_rate)

if args.load_model:
    print('continuing from checkpoint ' + args.load_model)
    model_state, optimizer_state, total_step, beta = torch.load(args.load_model)
    model.load_state_dict(model_state)
    optimizer.load_state_dict(optimizer_state)
else:
    total_step = beta = 0

param_norm = lambda m: math.sqrt(sum([p.norm().item() ** 2 for p in m.parameters()]))
grad_norm = lambda m: math.sqrt(sum([p.grad.norm().item() ** 2 for p in m.parameters() if p.grad is not None]))

meters = np.zeros(6)

# +
# model.load_state_dict(torch.load('./fast_molvae/vae_model/model.iter-1980'))
# model.load_state_dict(torch.load('./ckpt/chembl-pretrained/model.ckpt'))
model_state, optimizer_state, total_step, beta = torch.load('./ckpt/bms-pretrained/model.ckpt.5160')
model.load_state_dict(model_state)
model = model.to(device)

optimizer_vae = optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.ExponentialLR(optimizer_vae, args.anneal_rate)

# model.eval()
# for param in model.parameters():
#     param.requires_grad = False

# +
model_K = make_model('inception_v4', num_classes=10, pretrained=True)

for param in model_K.parameters():
    param.requires_grad = False

model_K._classifier.weight.requires_grad = True
model_K._classifier.bias.requires_grad = True


class InceptionV4Bottom(nn.Module):
            def __init__(self, original_model):
                super(InceptionV4Bottom, self).__init__()
                self.features = nn.Sequential(
                    # stop at conv4
                    *list(original_model.children())[:-1]
                )
#                 self.features = nn.Sequential(
#                     # stop at conv4
#                     *list(original_model.children())
#                 )
                dim1 = 512
                dim2 = 250
                self.num_ftrs = original_model._classifier.in_features
                self.classifier1 = nn.Sequential(nn.Linear(self.num_ftrs, dim1), nn.BatchNorm1d(dim1), nn.ReLU(), 
                                                nn.Linear(dim1, dim2))

            def forward(self, x1):
                x1 = self.features(x1)
#                 print (x1.shape)
                x1 = x1.view(-1, self.num_ftrs)
#                 x2 = self.dense1(x2)
#                 print (x2.shape, x1.shape)
#                 x = torch.cat((x1, x2), dim=1)
                o1 = self.classifier1(x1)
                return o1

model_cnn = InceptionV4Bottom(model_K)


# +
def unfreeze(model, ct_c = 7):
    for param in model.parameters():
        param.requires_grad = True

    ct = 0
    for child in model.features[0].children():
        ct += 1
        if ct < ct_c:
            for param in child.parameters():
                param.requires_grad = False
                
# unfreeze(model_cnn,3)


# +
model_cnn = model_cnn.to(device)

# Observe that all parameters are being optimized
optimizer = optim.Adam(model_cnn.parameters(), lr=0.001)
# -

IMG_SIZE = 512
transform = A.Compose([
    A.Resize(int(IMG_SIZE), int(IMG_SIZE)),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2()
])

dataset = DataFolder_ver2_BMS(args.train, '/home/quang/working/Theory_of_ML/hgraph2graph/train_processed_bms_ver2/ids', args.batch_size, transform, '/home/quang/working/Theory_of_ML/mbs-molecular-captioning/images/train/')
for batch, images_list, labels in tqdm(dataset):
    break

# +
import matplotlib.pyplot as plt
import torchvision.utils as vutils

from IPython.display import SVG
from rdkit import Chem
from rdkit.Chem import Draw

def show(img, size = 8):
    plt.figure(figsize=(size,size))
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')
    plt.show()

show(vutils.make_grid(images_list[:8], padding=10, normalize=True), size = 12)
plt.show()

smiles_code = labels[7]
m = Chem.MolFromSmiles(smiles_code)
print (Chem.MolToInchi(m))
SVG(Draw.MolsToGridImage([m], molsPerRow=4, subImgSize=(180, 150), useSVG=True))
# -

# ## Training

# +
for epoch in range(args.epoch):
    dataset = DataFolder_ver2_BMS(args.train, '/home/quang/working/Theory_of_ML/hgraph2graph/train_processed_bms_ver2/ids', args.batch_size, transform, '/home/quang/working/Theory_of_ML/mbs-molecular-captioning/images/train/')
    for batch, images_list, labels in tqdm(dataset):
        total_step += 1
#         images_list, tree_batch, jtenc_holder, mpn_holder, (jtmpn_holder,batch_idx) = batch_t
#         batch = tree_batch, jtenc_holder, mpn_holder, (jtmpn_holder,batch_idx)
        images_list = images_list.to(device)
    
        optimizer.zero_grad()
        optimizer_vae.zero_grad()
        mol_vec = model_cnn(images_list)
        loss, kl_div, wacc, iacc, tacc, sacc = model.forward_edit(*batch, mol_vec, beta=beta)
#             loss, kl_div, wacc, tacc, sacc = model.forward_edit(batch, beta, tree_vec, mol_vec)
        loss.backward()
        optimizer_vae.step()
        optimizer.step()
        
        if total_step % args.anneal_iter == 0:
            scheduler.step()
            print("learning rate: %.6f" % scheduler.get_lr()[0])
        
#         try:
#             optimizer.zero_grad()
#             tree_vec, mol_vec = model_cnn(images_list)
#             loss, kl_div, wacc, iacc, tacc, sacc = model.forward_edit(*batch, mol_vec, beta=beta)
# #             loss, kl_div, wacc, tacc, sacc = model.forward_edit(batch, beta, tree_vec, mol_vec)
#             loss.backward()
#             optimizer.step()
#         except Exception as e:
#             print (e)
#             continue

        meters = meters + np.array([kl_div, loss.item(), wacc * 100, iacc * 100, tacc * 100, sacc * 100])

        if total_step % args.print_iter == 0:
            meters /= args.print_iter
            print("[%d] Beta: %.3f, KL: %.2f, loss: %.3f, Word: %.2f, %.2f, Topo: %.2f, Assm: %.2f, PNorm: %.2f, GNorm: %.2f" % (total_step, beta, meters[0], meters[1], meters[2], meters[3], meters[4], meters[5], param_norm(model), grad_norm(model)))
            sys.stdout.flush()
            meters *= 0

#         if total_step % args.save_iter == 0:
#             torch.save(model.state_dict(), args.save_dir + "/model.iter-" + str(total_step))

#         if total_step % args.anneal_iter == 0:
#             scheduler.step()
#             print ("learning rate: %.6f" % scheduler.get_lr()[0])

#         if total_step % args.kl_anneal_iter == 0 and total_step >= args.warmup:
#             beta = min(args.max_beta, beta + args.step_beta)
            
torch.save(model_cnn.state_dict(),  "./model_cnn/model_cnn_hgraph_ver3_good_100eps_512.iter-" + str(total_step))
# -

# ## Validation

# +
# model_cnn.load_state_dict(torch.load('./model_cnn/model_cnn.iter-145'))

# +
# dataset_val = DataFolder_ver2_BMS('./val_processed_bms/mol/', '/home/quang/working/Theory_of_ML/hgraph2graph/val_processed_bms/ids',args.batch_size, transform, '/home/quang/working/Theory_of_ML/mbs-molecular-captioning/images/val/')
# for batch, images_list, labels in tqdm(dataset_val):
# #     total_step += 1
# # #         images_list, tree_batch, jtenc_holder, mpn_holder, (jtmpn_holder,batch_idx) = batch_t
# # #         batch = tree_batch, jtenc_holder, mpn_holder, (jtmpn_holder,batch_idx)
# #     images_list = images_list.to(device)

# #     optimizer.zero_grad()
# #     tree_vec, mol_vec = model_cnn(images_list)
# #     loss, kl_div, wacc, iacc, tacc, sacc = model.forward_edit(*batch, mol_vec, beta=beta)
# # #             loss, kl_div, wacc, tacc, sacc = model.forward_edit(batch, beta, tree_vec, mol_vec)
# #     loss.backward()
# #     optimizer.step()
#     break

dataset_val = DataFolder_ver2_BMS('./val_processed_bms/mol/', '/home/quang/working/Theory_of_ML/hgraph2graph/val_processed_bms/ids',args.batch_size, transform, '/home/quang/working/Theory_of_ML/mbs-molecular-captioning/images/val/')
total_labels = []
total_preds = []
for batch, images_list, labels in tqdm(dataset_val):
    with torch.no_grad():
        images_list = images_list.to(device)
        mol_vec = model_cnn(images_list)

        root_vecs, root_kl = model.rsample(mol_vec, model.R_mean, model.R_var, perturb=True)
        outputs = model.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=True, max_decode_step=150)
    total_labels.extend(labels)
    total_preds.extend(outputs)
# -

# ## Levenshtein distance for SMILES code

# +
import Levenshtein
def get_score(y_true, y_pred):
    scores = []
    for true, pred in zip(y_true, y_pred):
        score = Levenshtein.distance(true, pred)
        scores.append(score)
    avg_score = np.mean(scores)
    return avg_score

score = get_score(total_labels, total_preds)
print (score)

# +
import Levenshtein
def get_score_inchi(y_true, y_pred):
    scores = []
    for true, pred in zip(y_true, y_pred):
        inchi_pred = Chem.MolToInchi(Chem.MolFromSmiles(pred))
        inchi_true = Chem.MolToInchi(Chem.MolFromSmiles(true))
        score = Levenshtein.distance(inchi_true, inchi_pred)
        scores.append(score)
    avg_score = np.mean(scores)
    return avg_score

score = get_score_inchi(total_labels, total_preds)
print (score)
# -

# ## Validity score

# +
invalid_smiles = 0
for smile_temp in tqdm(total_preds):
    m = Chem.MolFromSmiles(smile_temp)
    if m is None:
        invalid_smiles += 1
        
print (1 - invalid_smiles/len(total_preds))
# -

# ## Visualization

# +
images_list = images_list.to(device)
mol_vec = model_cnn(images_list)

root_vecs, root_kl = model.rsample(mol_vec, model.R_mean, model.R_var, perturb=True)
outputs = model.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=True, max_decode_step=150)
# -

for i in range(len(outputs)):
    print ('GT: ',labels[i])
    print ('Pred: ', outputs[i])

i = 5
m = Chem.MolFromSmiles(labels[i])
# print (Chem.MolToInchi(m))
SVG(Draw.MolsToGridImage([m], molsPerRow=4, subImgSize=(180, 150), useSVG=True))

smiles_code = labels[i]
m = Chem.MolFromSmiles(smiles_code)
print (Chem.MolToInchi(m))
SVG(Draw.MolsToGridImage([m], molsPerRow=4, subImgSize=(180, 150), useSVG=True))

smiles_code = 'C1#CCC[CH2:1]CC#CCC1'
m = Chem.MolFromSmiles(smiles_code)
print (Chem.MolToInchi(m))
SVG(Draw.MolsToGridImage([m], molsPerRow=4, subImgSize=(180, 150), useSVG=True))

smiles_code = 'C1#CC[CH2:1][CH2:1][CH2:1]C#CCC1'
m = Chem.MolFromSmiles(smiles_code)
print (Chem.MolToInchi(m))
SVG(Draw.MolsToGridImage([m], molsPerRow=4, subImgSize=(180, 150), useSVG=True))



StarcoderdataJupyterScriptsDedupFiltered