# CONFIDENTIAL. (C) 2018 OpenEye Scientific Software Inc.
# All rights reserved.
# ALL SOFTWARE BELOW IS PROPRIETARY AND CONFIDENTIAL TO OPENEYE
# SCIENTIFIC SOFTWARE INC., AND IS SUBJECT TO THE FULL PROTECTION OF
# COPYRIGHT AND TRADESECRET LAW.
# Copying or modifying this software is strictly prohibited and illegal.
# Using this software without a valid license from OpenEye is strictly
# prohibited and illegal.  De-obfuscating, de-minimizing, or any other
# attempt to reverse engineer or discover the workings of this software,
# is strictly prohibited. Even if software is exposed or visible, it is
# still proprietary and confidential and the above prohibitions apply.
# This software is NOT "Sample Code". For purposes of customization or
# interoperation with other software you must use Sample Code
# specifically labeled as such.
# Please contact OpenEye at eyesopen.com/legal if you have any questions
# about this warning.
"""
usage:
    create_meta_data.py [options] PDBFILE
    create_meta_data.py [options] --dir DIR

options:
    --meta-dir DIR  Directory with hand-curated meta files.
    --force         Overwrite existing file
    -v              Verbose output
"""
import docopt
import json
import itertools
import os
import re
import shutil
import sys
import multiprocessing

from spruce.utils.file_utils import get_pdb_files
from spruce.pdb.metadata import OECreateStructureMetadataFromFile
from spruce.utils.progress import OEPyProgress

from openeye import oespruce


def create_one(pdb_file, args):
    dirname, basename = os.path.split(pdb_file)
    meta_name = '{}.json'.format(basename)
    in_meta = re.sub(".pdb", ".meta", pdb_file)

    out_meta = os.path.join(dirname, meta_name)
    if os.path.isfile(out_meta) and not args['--force']:
        print('{} already exists. Overwrite with --force'.format(out_meta))
        return None

    meta_dir = args['--meta-dir']

    existing = None
    if meta_dir and os.path.isdir(meta_dir):
        existing = os.path.join(meta_dir, meta_name)
        if not os.path.isfile(existing):
            existing = None

    if existing:
        shutil.copy(existing, out_meta)
    else:
        loaded_metadata = None
        if os.path.isfile(in_meta):
            with open(in_meta, "r") as f:
                loaded_metadata = json.load(f)
        try:
            meta_data = OECreateStructureMetadataFromFile(pdb_file, loaded_metadata)
        except Exception as e:
            print("failed: {}".format(pdb_file))
            return

        meta_data = json.loads(meta_data)

        # fix up possible blank entries
        if "Revision" in meta_data and len(meta_data["Revision"]) == 0:
            meta_data["Revision"] = "1.0"
        if "RevisionDate" in meta_data and len(meta_data["RevisionDate"]) == 0:
            if "ExperimentDate" in meta_data and len(meta_data["ExperimentDate"]) > 0:
                meta_data["RevisionData"] = meta_data["ExperimentDate"]
            else:
                meta_data.pop("RevisionDate")

        with open(out_meta, 'w') as ofp:
            ofp.write(json.dumps(meta_data))

    return 'created {}'.format(out_meta)


def run(args):
    return create_one(args[0], args[1])


def create_all(pdb_dir, args):
    pdb_files = get_pdb_files(pdb_dir)
    num_files = len(pdb_files)
    progress = OEPyProgress(hide=args["-v"])
    progress.set_task("creating meta data")
    with progress:
        pool = multiprocessing.Pool()
        count = 0
        for res in pool.imap_unordered(run, zip(pdb_files, itertools.repeat(args))):
            count += 1
            progress.set_progress(count/num_files)
            if args['-v'] and res is not None:
                print(res)
        pool.terminate()


def main():
    args = docopt.docopt(__doc__)
    pdb_file = args['PDBFILE']
    pdb_dir = args['DIR']

    if pdb_file is not None:
        create_one(pdb_file, args)
    else:
        create_all(pdb_dir, args)


if __name__ == "__main__":
    sys.exit(main())