#!/usr/bin/env python3

import argparse
from os.path import dirname
from warnings import warn

import numpy as np
from AaronTools.fileIO import FileReader
from AaronTools.geometry import Geometry
from AaronTools.pathway import Pathway
from AaronTools.utils.utils import get_filename, get_outfile


def width(n):
    """use to determine 0-padding based on number of files we're printing"""
    return np.ceil(np.log10(n))


def parse_mode_str(s, t):
    """split mode string into modes and mode combos
    e.g.
    t=int, 1,2+3,4 -> [[0], [1,2], [3]]
    t=float 0.1,0.05+0.03,0.07 -> [[0.1], [0.05, 0.03], [0.07]]"""

    # the way this is being used is if t is int, we are changing 1-indexed things to 0-index
    # if t is float, were going to use the result to scale a normal mode (don't subtract 1)

    if t is not int and t is not float:
        raise TypeError(
            "can only parse mode string into ints or floats, not %s" % repr(t)
        )

    modes = s.split(",")
    out_modes = []
    for mode in modes:
        out_modes.append([])
        for combo in mode.split("+"):
            if t is int:
                out_modes[-1].append(int(combo) - 1)
            elif t is float:
                out_modes[-1].append(float(combo))

    return out_modes


follow_parser = argparse.ArgumentParser(
    description="move the structure along a normal mode",
    formatter_class=argparse.RawTextHelpFormatter,
)
follow_parser.add_argument(
    "input_file",
    metavar="input file",
    type=str,
    help='input frequency file (i.e. Gaussian output where "freq" was specified)',
)

follow_parser.add_argument(
    "-m",
    "--mode",
    type=str,
    nargs="+",
    required=False,
    default=None,
    dest="mode",
    metavar=("mode 1", "mode 2"),
    help="comma-separated list of modes to follow (1-indexed)",
)

follow_parser.add_argument(
    "-r",
    "--reverse",
    action="store_const",
    const=True,
    default=False,
    required=False,
    dest="reverse",
    help="follow the normal mode in the opposite direction",
)

follow_parser.add_argument(
    "-a",
    "--animate",
    type=int,
    nargs=1,
    default=None,
    required=False,
    dest="animate",
    metavar="frames",
    help="print specified number of structures to make an animation",
)

follow_parser.add_argument(
    "-rt",
    "--roundtrip",
    action="store_const",
    const=True,
    default=False,
    required=False,
    dest="roundtrip",
    help="make animation roundtrip",
)

follow_parser.add_argument(
    "-s",
    "--scale",
    type=str,
    nargs=1,
    default=None,
    required=False,
    dest="scale",
    metavar="max displacement",
    help="scale the normal mode so that this is the maximum amount an \natom is displaced",
)

follow_parser.add_argument(
    "-o",
    "--output-destination",
    type=str,
    nargs="+",
    default=False,
    required=False,
    dest="outfile",
    help="output destination\n" +
    "$i in file name will be replaced with zero-padded numbers if --animate is used\n" +
    "$INFILE will be replaced with the name of the input file\n" +
    "Default: stdout",
)

args = follow_parser.parse_args()

fr = FileReader(args.input_file, just_geom=False)
geom = Geometry(fr)

if args.mode is None:
    modes = [
        [i]
        for i, freq in enumerate(fr.other["frequency"].data)
        if freq.frequency < 0
    ]
else:
    modes = parse_mode_str(args.mode[0], int)

if not modes:
    raise RuntimeError("no vibrational mode specified and no imaginary modes found")

# copy the list of output files or set all output files to False (print all to stdout)
if args.outfile is not False:
    outfiles = [f for f in args.outfile]
else:
    outfiles = [False for m in modes]

if len(outfiles) != len(modes):
    warn(
        "number of output files does not match number of modes: %i files, %i modes"
        % (len(outfiles), len(modes))
    )

if args.scale is None:
    scale = [[0.35] * len(mode) for mode in modes]
else:
    scale = parse_mode_str(args.scale[0], float)

for i, mode in enumerate(modes):
    if outfiles[i]:
        if "$i" not in outfiles[i]:
            append = True
    else:
        append = False
    dX = np.zeros((len(geom.atoms), 3))
    # figure out how much we'll have to scale each mode
    for j, combo in enumerate(mode):
        max_norm = 0
        for k, v in enumerate(
                fr.other["frequency"].data[combo].vector
        ):
            n = np.linalg.norm(v)
            if n > max_norm:
                max_norm = n

        # scale this mode by 0.35 (or whatever the user asked for)/max_norm
        x_factor = scale[i][j] / max_norm

        if args.reverse:
            x_factor *= -1

        dX += x_factor * fr.other["frequency"].data[combo].vector

    if args.animate is not None:
        # animate by setting up 3 geometries: -, 0, and +
        # then create a Pathway to interpolate between these
        # if roundtrip, - -> 0 -> + -> 0 -> -
        # if not roundtrip, - -> 0 -> +
        w = width(args.animate[0])
        fmt = "%0" + "%i" % w + "i"

        Xf = geom.coords + dX
        X = geom.coords
        Xr = geom.coords - dX

        # make a scales(t) function so we can see the animation progress in the XYZ file comment
        if args.roundtrip:
            other_vars = {}
            for i, mode_scale in enumerate(scale[i]):
                other_vars["scale %i"] = [
                    mode_scale,
                    0,
                    -mode_scale,
                    0,
                    mode_scale,
                ]
            pathway = Pathway(np.array([Xf, X, Xr, X, Xf]), other_vars=other_vars)

        else:
            other_vars = {}
            for i, mode_scale in enumerate(scale[i]):
                other_vars["scale %i"] = [mode_scale, 0, -mode_scale]
            pathway = Pathway(np.array([Xf, X, Xr]), other_vars=other_vars)

        # print animation frames
        for k, t in enumerate(np.linspace(0, 1, num=args.animate[0])):
            if outfiles[i] is not False:
                outfile = outfiles[i].replace("$i", fmt % k)
            else:
                outfile = outfiles[i]

            followed_geom = pathway.interpolate_geometry(t, geom)
            followed_geom.comment = (
                "animating mode %s scaled to displace at most [%s]"
                % (
                    repr(mode),
                    ", ".join(str(pathway.interpolate_other_var(t, key)) for key in other_vars),
                )
            )
            if args.outfile:
                outfile = get_outfile(
                    args.outfile,
                    INFILE=get_filename(args.input_file, include_parent_dir=False),
                    INDIR=dirname(args.input_file),
                )
                followed_geom.write(
                    append=False,
                    outfile=outfile,
                )
            else:
                print(followed_geom.write(outfile=False))

    else:
        w = width(len(modes))
        fmt = "%0" + "%i" % w + "i"

        followed_geom = geom.copy()
        followed_geom.update_geometry(geom.coords + dX)
        followed_geom.comment = "following mode %s scaled to displace at most %s" % (
            repr(mode),
            repr(scale[i]),
        )

        outfile = outfiles[i]
        if args.outfile:
            outfile = get_outfile(
                args.outfile,
                INFILE=get_filename(args.input_file, include_parent_dir="$INDIR" not in args.outfile),
                INDIR=dirname(args.input_file),
            )
            followed_geom.write(append=True, outfile=outfile)
        else:
            print(followed_geom.write(outfile=False))
