import argparse
import h5py

import numpy as np
import torch
import torch.nn.functional as F
from model import DeepDrug3D

def load_data(path):
    with h5py.File(path, 'r') as f:
        X = f['X'][()]
    if len(X.shape) == 4:
        X = np.expand_dims(X, axis=0)
    X = torch.FloatTensor(X)
    return X

def predict(path, model_path):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print('Current device: ' + str(device))
    
    net = DeepDrug3D(14)
    net.load_state_dict(torch.load(model_path))
    net.to(device)
    data = load_data(path)
    data = data.to(device)
    output = net(data)
    proba = F.softmax(output, dim=1)
    score = proba.data.cpu().numpy()[0]
    print('The probability of pocket provided binds with ATP ligands: {:.4f}'.format(score[0]))
    print('The probability of pocket provided binds with Heme ligands: {:.4f}'.format(score[1]))
    print('The probability of pocket provided binds with other ligands: {:.4f}'.format(score[2]))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--f', type=str, required=True, help='Input h5 file name')
    parser.add_argument('--m', type=str, required=True, help='Path to the trained model weights')
    opt = parser.parse_args()

    predict(opt.f, opt.m)