# (C) 2026 Cadence Design Systems, Inc. (Cadence)
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  CADENCE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.

import argparse
import os

import mdtraj as md

DEFAULT_INPUTS = ['trajectory_comb.h5', 'trajectory.h5']


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert an MDTraj .h5 trajectory into .dcd and .pdb files.'
    )
    parser.add_argument(
        '-i', '--input',
        help=f"Path to the input .h5 trajectory file. If omitted, the current directory is searched for: {', '.join(DEFAULT_INPUTS)}",
    )
    parser.add_argument(
        '-o', '--output',
        default='bound',
        help='Output path prefix (without extension). Produces <output>.pdb '
             'and <output>.dcd. Default: %(default)s',
    )
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        help='Print the OpenMM topology after conversion.',
    )
    return parser.parse_args()


def find_input(path):
    if path is not None:
        if not os.path.isfile(path):
            raise FileNotFoundError(f'input file not found: {path}')
        return path
    for candidate in DEFAULT_INPUTS:
        if os.path.isfile(candidate):
            return candidate
    raise FileNotFoundError(
        f"unable to find input file; provide one with --input or run from a directory containing: {', '.join(DEFAULT_INPUTS)}"
    )


def main():
    args = parse_args()
    in_file = find_input(args.input)

    out_prefix = os.path.splitext(args.output)[0]
    out_dir = os.path.dirname(out_prefix)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)

    trj = md.load(in_file)

    trj[0].save(f'{out_prefix}.pdb')
    trj.save(f'{out_prefix}.dcd')
    print(f'Wrote {out_prefix}.pdb and {out_prefix}.dcd from {in_file}')

    if args.verbose:
        print(trj.topology.to_openmm())


if __name__ == '__main__':
    main()