#!/usr/bin/env python3

import sys, textwrap
from pathlib import Path
sys.path.append( str(Path(__file__).resolve().parents[1] / "src") )
import gen_top_ENM, gen_top_Go, json_to_top, map_to_cg, map_traj
import wat_to_polar, setup_lmp, setup_gmx, gen_lmp_input, gen_gmx_input
import mod_ion
from argparse import ArgumentParser, RawTextHelpFormatter

def get_option():
    argparser = ArgumentParser(add_help=False, formatter_class=RawTextHelpFormatter,
                               usage=textwrap.dedent("""\
    cg_spica tool_name 

    # list of tool_names #------------------------------------------------
    json2top  : make top file from json file  
    map2cg    : map AA configuration to CG (only PDB format)
    maptraj   : map AA MD trajectory to CG (MDAnalysis module is required)
    modion    : adjust NaCl salt concentration in CG configuration
    ENM       : generate top file of protein with elastic network
    Go        : generate top file of protein with Go model
    wat2polar : convert SPICA CG water to pSPICA polar CG water
    setup_lmp : generate input files to run CG-MD with LAMMPS
    setup_gmx : generate input files to run CG-MD with GROMACS
    gen_lmpin : generate LAMMPS input file for SPICA or pSPICA
    gen_gmxin : generate GROMACS input file for SPICA or pSPICA
    ----------------------------------------------------------------------
                                                 """))
    argparser.add_argument('tool_name', type=str,
                            help='input tool_name.')
    return argparser

def main():
    ap        = get_option()
    args, sub = ap.parse_known_args()
    if args.tool_name == "ENM":
        args    = gen_top_ENM.get_option_script(sub)
        cgpdb   = args.cgpdb
        outfile = args.output
        kENM    = args.kENM
        MAXdr   = args.maxr
        pspica  = args.pspica
        v1      = args.v1
        dssp    = args.dssp
        aapdb   = args.aapdb
        gen     = gen_top_ENM.gen_top_ENM(cgpdb, outfile, kENM, MAXdr, pspica, v1, dssp, aapdb)
        gen.run()
    elif args.tool_name == "Go":
        args    = gen_top_Go.get_option_script(sub)
        cgpdb   = args.cgpdb
        outfile = args.output
        eps     = args.eps
        MAXdr   = args.maxr
        pspica  = args.pspica
        v1      = args.v1
        dssp    = args.dssp
        aapdb   = args.aapdb
        gen     = gen_top_Go.gen_top_Go(cgpdb, outfile, eps, MAXdr, pspica, v1, dssp, aapdb)
        gen.run()
    elif args.tool_name == "json2top":
        args  = json_to_top.get_option_script(sub)
        res   = args.resname
        top   = f"{res}.top"
        jfile = args.json
        da    = args.dupang
        jt    = json_to_top.json_to_top(jfile, da)
        jt.run(res, top)
    elif args.tool_name == "map2cg":
        args     = map_to_cg.get_option_script(sub)
        infile   = args.input
        outfile  = args.output
        jsonfile = args.json
        nodelwat = args.nodelwat
        verbose  = args.verbose
        mapCG    = map_to_cg.map_to_cg(infile, outfile, jsonfile, nodelwat, verbose)
        mapCG.run()
    elif args.tool_name == "maptraj":
        args     = map_traj.get_option_script(sub)
        inPDB    = args.inputAAPDB
        inTRAJ   = args.inputAATRAJ
        outTRAJ  = args.outputCGTRAJ
        outPDB   = args.outpdb
        jsonfile = args.json
        nodelwat = args.nodelwat
        beg      = args.begin
        last     = args.last
        mt       = map_traj.map_traj(inPDB, inTRAJ, outTRAJ, outPDB, jsonfile, nodelwat, beg, last)
        mt.run()
    elif args.tool_name == "wat2polar":
        args     = wat_to_polar.get_option_script(sub)
        infile   = args.input
        outfile  = args.output
        wat2polar = wat_to_polar.wat_to_polar(infile, outfile)
        wat2polar.run()
    elif args.tool_name == "modion":
        args     = mod_ion.get_option_script(sub)
        infile   = args.input
        outfile  = args.output
        conc     = args.conc
        soluq    = args.soluq
        pspica   = args.pspica
        verbose  = args.verbose
        modion   = mod_ion.mod_ion(infile, outfile, conc, soluq, pspica, verbose)
        modion.run()
    elif args.tool_name == "setup_lmp":
        args = setup_lmp.get_option_script(sub)
        setup_lmp.run(args)
    elif args.tool_name == "setup_gmx":
        args = setup_gmx.get_option_script(sub)
        setup_gmx.run(args)
    elif args.tool_name == "gen_lmpin":
        args = gen_lmp_input.get_option_script(sub)
        datafile = args.data
        parmfile = args.parm
        outfile = args.output
        temp = args.T
        press = args.P
        pspica = args.pspica
        gen_lmpin = gen_lmp_input.gen_lmp_inp(datafile, parmfile, outfile, temp, press, pspica)
        gen_lmpin.run()
    elif args.tool_name == "gen_gmxin":
        args = gen_gmx_input.get_option_script(sub)
        conf = args.pdb
        ndxf = args.ndx
        outf = args.output
        temp = args.T
        press = args.P
        pspica = args.pspica
        em = args.em
        gen_gmxin = gen_gmx_input.gen_gmx_inp(conf, ndxf, outf, temp, press, pspica, em)
        gen_gmxin.run()
    else:
        sys.exit(f"ERROR: invalid tool name {args.tool_name}")

if __name__ == "__main__":
    main()
