#!python

import torch
from os import path
dir_script = path.dirname(path.realpath(__file__))
import sys
sys.path.append(dir_script+'/../')
import pydssp
import tqdm

from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
argparse = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
argparse.add_argument('pdbs', nargs='+', type=str, help='input PDB file')
argparse.add_argument('-d', '--device', type=str, choices=['cpu', 'cuda'], default=None, help='device')
argparse.add_argument('-o', '--output-file', type=str, default=None, help='output file')
argparse.add_argument('--ignore-proline-donor', type=bool, default=True, help='ignore N-H of proline')
args = argparse.parse_args()

if args.device == 'cuda':
    assert torch.cuda.is_available(), "CUDA is not available"
elif args.device is None:
    args.device = 'cuda' if torch.cuda.is_available() else 'cpu'

fh = open(args.output_file, 'w') if args.output_file is not None else sys.stdout

for pdb in tqdm.tqdm(args.pdbs):
    # read pdb file
    coord, sequence = pydssp.read_pdbtext(open(pdb, 'r').read(), return_sequence=True)
    coord = torch.Tensor(coord).to(args.device)
    donor_mask = sequence != 'PRO' if args.ignore_proline_donor else None
    # main calcuration
    dsspline = ''.join(pydssp.assign(coord, donor_mask=donor_mask))
    # write file or STDOUT
    fh.write(f"{dsspline} {pdb}\n")
