#! /usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import re
import pathlib, os
import shutil
from mrdna import get_resource_path

parser = argparse.ArgumentParser(prog="mrdna",
				 description="""Easy multiresolution simulations of DNA nanotechnology objects using ARBD""")

parser.add_argument('-o','--output-prefix', type=str, default=None,
                    help="Name for your job's output")
parser.add_argument('-d','--directory', type=str,  default=None,
                    help='Directory for simulation; does not need to exist yet')

parser.add_argument('--sequence-file', type=str, default=None,
                    help='Sequence of longest strand.')

parser.add_argument('--fill-sequence', choices=['A','T','C','G','random'], default='T',
                    help='Fill sequence where otherwise unset.')

parser.add_argument('--crossover-to-intrahelical-cutoff', type=float, default=-1,
                    help='Set the distance (in Angstroms) beyond which crossovers between helix ends are converted to intrahelical connections; a negative value means no crossovers will be converted')

parser.add_argument('--backbone-scale', type=float, default=1.0,
                    help='Factor to scale DNA backbone in atomic model; try 0.25 to avoid clashes for atomistic simulations')

parser.add_argument('--enrg-md-steps', type=float, default=1e6,
                    help='Number of ENRG-MD steps')

parser.add_argument('--write-pqr', action='store_true',
                    help='Write PQR file in addition PDB?')

parser.add_argument('input_file', type=str,
                    help="""Any of the following:
                    (1) a cadnano JSON file;
                    (2) a vHelix Maya (.ma) file;
                    (3) an atomic PDB file""")

args = parser.parse_args()

if __name__ == '__main__':
    
    infile = pathlib.Path(args.input_file)
    try:
        prefix = infile.stem
        extension = infile.suffix
    except:
        raise Exception("Unrecognized input file '{}'".format(infile))
        
    if extension == '.json':
        from mrdna.readers import read_cadnano as read_model
    elif extension == '.ma':
        from mrdna.readers import read_vhelix as read_model
    elif extension == '.pdb':
        from mrdna.readers import read_atomic_pdb as read_model
    else:
        raise Exception("Unrecognized input file '{}'".format(infile))

    if args.output_prefix is not None:
        prefix = args.output_prefix
    directory = args.directory

    if prefix == infile.stem and extension == ".pdb":
        print("Are you sure you want to overwrite '{}'? (you may provide '-o' option) [y/N]".format(infile))
        if input() not in ("y","Y","yes","Y","Yes","YES"):
            import sys
            sys.exit(1)
                    
    model = read_model( str(infile) )

    if args.crossover_to_intrahelical_cutoff >= 0:
        model.convert_crossovers_at_ends_beyond_cutoff_to_intrahelical(
            args.crossover_to_intrahelical_cutoff)

    if args.sequence_file is not None:
        from mrdna.model.dna_sequence import read_sequence_file
        seq = read_sequence_file(args.sequence_file)
        model.set_sequence( seq, args.fill_sequence )

    d_orig = os.getcwd()
    try:
        if directory is None:
            directory='.'
        elif not os.path.exists(directory):
            os.makedirs(directory)
        os.chdir(directory)

        print("""Generating your ENRG MD simulation files...
If your simulation results in any publications, please cite the following:
http://dx.doi.org/10.1093/nar/gkw155
""")
        model._clear_beads()
        model._generate_atomic_model(scale=args.backbone_scale)
        model.write_atomic_ENM( prefix )
        try:
            shutil.copytree( get_resource_path("charmm36.nbfix"), "charmm36.nbfix" )
        except FileExistsError:
            ...
        try:
            os.makedirs('output')
        except FileExistsError:
            ...
        model.atomic_simulate( output_name = prefix, num_steps=args.enrg_md_steps, write_pqr = args.write_pqr, dry_run=True )
        
    except:
        raise
    finally:
        os.chdir(d_orig)