import os
import re

import pkg_resources


def write_vmd(
    colors,  # array of rgb values
    trajectories,  # array of filenames
    states,
    opt_in={},  # dictionary of options
):

    """Writes vmd script to visualize trajectories

    Params:
        colors, np.ndarray (nstates,3), array of rgb values for each state
        trajectories, list (ntrajectories), list of xyzs filenames
        states, dictionary (ntrajectories), trajectory index to state label (0 based)
        opt_in, dictionary (options to override defaults)

    """

    header_path = pkg_resources.resource_filename(__name__, "ref/header.vmd")
    footer_path = pkg_resources.resource_filename(__name__, "ref/footer.vmd")
    molecule_path = pkg_resources.resource_filename(__name__, "ref/molecule.vmd")

    # Set User Options
    options = {
        "fn_header": header_path,
        "fn_footer": footer_path,
        "fn_mol": molecule_path,
        "fn_vmd": "vis_fms.vmd",
        "render": True,
        "exit_vmd": True,
        "opacities": True,
    }
    for key, val in list(opt_in.items()):
        options[key] = val

    fn_header = options["fn_header"]
    fn_footer = options["fn_footer"]
    fn_mol = options["fn_mol"]
    fn_vmd = options["fn_vmd"]
    options["render"]
    options["opacities"]
    exit_vmd = options["exit_vmd"]  # whether to render and exit vmd

    # Lines from vmd Templates
    header_lines = open(fn_header).readlines()
    footer_lines = open(fn_footer).readlines()
    mol_lines = open(fn_mol).readlines()

    # Number of trajectories
    ntrajs = len(trajectories)

    # Write the header
    # Write new material for each opacity
    # default to 1.0/ntrajs?
    fh = open(fn_vmd, "w")
    for lh in header_lines:
        mobj = re.match(r"^  # Add materials!\s*", lh)
        mobj1 = re.match(
            r"^  set mlist { Opaque Transparent BrushedMetal Diffuse Ghost Glass1 Glass2 Glass3 Glossy HardPlastic MetallicPastel Steel Translucent Edgy EdgyShiny EdgyGlass Goodsell AOShiny AOChalky AOEdgy BlownGlass GlassBubble RTChrome }",
            lh,
        )
        if mobj:
            fh.write(lh)
            for n in range(ntrajs):
                fh.write("  material change ambient Transp%d 0.000000\n" % (n + 1))
                fh.write("  material change diffuse Transp%d 0.920000\n" % (n + 1))
                fh.write("  material change specular Transp%d 0.290000\n" % (n + 1))
                fh.write("  material change shininess Transp%d 0.470000\n" % (n + 1))
                fh.write("  material change mirror Transp%d 0.000000\n" % (n + 1))
                fh.write("  material change opacity Transp%d 0.500000\n" % (n + 1))
                fh.write("  material change outline Transp%d 1.510000\n" % (n + 1))
        if mobj1:
            fh.write(
                "  set mlist { Opaque Transparent BrushedMetal Diffuse Ghost Glass1 Glass2 Glass3 Glossy HardPlastic MetallicPastel Steel Translucent Edgy EdgyShiny EdgyGlass Goodsell AOShiny AOChalky AOEdgy BlownGlass GlassBubble RTChrome "
            )
            for n in range(ntrajs):
                fh.write("Transp%d " % (n + 1))
            fh.write("}\n")
        else:
            fh.write(lh)

    # Looping through trajectories
    1.0 / ntrajs
    for n in range(ntrajs):
        # Writing vmd segment for each trajecotry
        for ml in mol_lines:
            # Read in filename
            mobj1 = re.sub(r"filename", trajectories[n], ml)
            # Assign color ID number
            mobj2 = re.sub(
                r"mol color ColorID \d+", r"mol color ColorID %d" % (32 - n), mobj1
            )
            # Assign material to representation
            mobj3 = re.sub(
                r"mol material Transp\d+", r"mol material Transp%d" % (n + 1), mobj2
            )
            # Reassign RGB value to color ID
            mobj4 = re.sub(
                r"color change rgb \d+ \d\.\d+ \d\.\d+ \d\.\d+",
                r"color change rgb %d %6.5f %6.5f %6.5f"
                % (
                    32 - n,
                    colors[states[n]][0],
                    colors[states[n]][1],
                    colors[states[n]][2],
                ),
                mobj3,
            )
            fh.write("%s" % (mobj4))
        fh.write("# Finished traj %d\n" % (n + 1))

    # Writing footer for vmd file
    for lf in footer_lines:
        fh.write(lf)
    if exit_vmd:
        fh.write("source render.tcl\n")
        fh.write("exit")


def write_render(
    opt_in={},
):

    """Renders each frame of trajectory in vmd

    Params:
        opt_in, dictionary (options to override defaults)

    """

    options = {
        # visualization options
        "opacities": [],
        "end_frame": 1000,
        "increment_frame": 10,
        "sleep": 1,
        "ntrajs": 0,
        "display_size": None,
        "translation": None,
        "axis_on": False,
        "scale": 0.25,
    }
    for key, val in list(opt_in.items()):
        options[key] = val
    opacities = options["opacities"]
    end_frame = options["end_frame"]
    increment_frame = options["increment_frame"]
    ntrajs = options["ntrajs"]
    sleep = options["sleep"]
    display_size = options["display_size"]  # (1000, 600)
    translation = options["translation"]  # [-0.5, 0.0, 0.0]
    axis_on = options["axis_on"]
    scale = options["scale"]

    # TODO: Create options for:
    # Length of trajectory
    # Interval of frames (steps)
    # Translations and rotations
    # image scale and size
    # Render sleep settings

    # Number of trajectories
    fh = open("render.tcl", "w")

    # Writing Opacities based on amplitude for all trajectories
    # Setting the amplitude to zero for the time before Spawning
    # For groundstate trajectories that dropped, setting the amplitude
    # to it's last amplitude until the end of the trajectory
    for n in range(ntrajs):
        fh.write("set opa%d [list   " % (n + 1))
        for val in opacities[n]:
            fh.write("%4.3f   " % (val))
        fh.write("]\n")

    # Writing remainder of display settings
    if not axis_on:
        fh.write("axes location off\n")
    # TODO: include custom image size settings!
    fh.write("scale to %1.3f\n" % scale)
    if display_size is not None:
        fh.write("display resize %d %d\n" % (display_size[0], display_size[1]))
    fh.write("display culling off\n")
    fh.write("animate goto start\n")
    if translation is not None:
        fh.write(
            "translate 1 %1.3f %1.3f %1.3f\n"
            % (translation[0], translation[1], transpation[2])
        )

    # Writing commands to loop through the trajectories
    # TODO: include frames and increment settings
    fh.write(
        "for {set frame 0} {$frame < %d} {incr frame %d} {\n"
        % (end_frame, increment_frame)
    )
    fh.write("    animate goto $frame\n")

    # Writing commands to reset opacity
    for n in range(ntrajs):
        fh.write("    set var%d [lindex $opa%d $frame]\n" % (n + 1, n + 1))
        fh.write("    material change opacity Transp%d $var%d\n" % (n + 1, n + 1))

    # Writing commands to render tga files
    fh.write('    set fname_tga snap.[format "%04d" $frame].tga\n')
    fh.write("    render TachyonInternal snapshots/$fname_tga\n")
    fh.write("    draw delete all\n")
    # TODO: custom sleep settings
    fh.write("    sleep %1.2f\n" % sleep)
    fh.write("}\n")


# TODO: add more options
def run_vmd_render():

    """Loads VMD and renders images, also composites into gif format"""

    os.system("mkdir snapshots")
    os.system("module load VMD/1.9.2")
    print("Running VMD")
    os.system("vmd -dispdev text -e vis_fms.vmd")


def get_mp4():
    print("Writing mp4")
    print("ffmpeg must be installed - not necessarily capable of running on fire!")
    os.system(
        "ffmpeg -r 4 -i snapshots/snap.%04d.tga -c:v libx264 -pix_fmt yuv420p -r 10 snap.mp4"
    )


def get_gif():

    print("Writing gif")
    os.system("convert -delay 20 -loop 1 snapshots/snap*tga snap.gif")
