from __future__ import print_function
import maya.cmds as mc
import os.path
import string

"""
author: Johan Schreurs
author-email: johan.schreurs@gmail.com
license: MIT license, see below for full license text.
          The gist of it is that an MIT license is open and very permissive.
maya-versions: only tested on Maya 2013, but likely to run on 2014 and 2015.

These tools are your little "monkeys" to to your repetitive tasks.

At the moment these are more like a collection of little tools for quite
different purposes, but I hope this collection will grow big enough that
it will need to be split into several packages.


Tools currently available
=========================

- Some basic tools to help finalize camera tracking scenes.

- Basic Sorting tool to keep the outliner tidier when you have a lot of objects
    (i.e. many camera tracking points)

- A UV snapshot that can export multiple object in one go.

- Tool to setup a shader network quickly, using another net as a template.
    Each (important) node in the network contains one string that will be
    replaced with a networks "base" name.
    Keep names unique and keeps the network readable.
    Nodes that don't have the placeholder just get default names with a number
    at the end

Enjoy!




=====================
The MIT License (MIT)
=====================

Copyright (c) 2014 Johan Schreurs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

http://opensource.org/licenses/MIT
"""

#TODO:  Think of ways to automate testing where possible, perhaps using mayapy
#       or non-interactive maya commandline.
#       This may tell you how:
#       http://tech-artists.org/forum/showthread.php?1904-Unit-testing-DCC-code
#       Research this on stackoverflow.com as well.


# Metadata
# ========
# Initialize this before any imports so we can at least get this info even if
# there is a problem with any imports.
#
# version info:
# * tuple as a machine readable option
# * regular "a.b.c" version (derived from tuple for consistency)

__version_tupple__ = (0, 1, 0)
__version__ = ".".join([str(s) for s in __version_tupple__])
__release_date__ = "2014-11-04"


# Module level variable to turn logging on/off for the whole module.
do_logging = False


def _log_line(*messages):
    """
    Write a message to the standard output, with a simple mechanism to turn
    this logging on/off globally.

    If the module variable DO_LOGGING is True log lines will be written
    otherwise not.
    """
    if do_logging:
        # Python 2.6 has not *tuple syntax, so unpack it the hard way:
        # print each string individually with a space as a fake end of line
        # character.
        for msg in messages:
            print(msg, end=" ")

        # new line
        print()


def is_group(transform_node):
    """
    Checks if transform_node is a group. Groups don't have a
    shape node.

    Parameters
    ----------
    transform_node: the maya DAG node to be checked

    Test procedure
    --------------
    Now implemented as unit test.
    - Create a 2 primitive geometries (e.g. polyCube)
    - Put one in a group but not the other.
    - Check result of function on each group and standalone primitive.
    """

    # Find out if the transform node has a shape
    # If not it's a group.
    shapes = mc.listRelatives(transform_node, shapes=True)

    if shapes:
        return False
    else:
        return True


def sort_in_outliner(transform_nodes, no_groups=True, reverse=False):
    """
    Sorts the given objects <transform_nodes> alphabetically in the outliner.

    Parameters
    ----------
    nodes:
        List of transform nodes of objects to sort

    no_groups:
        boolean, default True
        Leave groups out of the sorting.
        This helps to prevent accidental selection of the parent group of your
        nodes.

    reverse:
       boolean, default False
       Reverse the sorting order, A-Z becomes Z-A.

    Test procedure
    --------------
    Now partially implemented as unit test.
    - Create a number of primitive geometries (e.g. polyCube)
    - Sort them in reverse Z-A
    - Sort them A-Z
    - Really mix them more randomly [not yet in unit test]
    - Sort them again, ascending and descending.

    """

    sorted_objects = sorted(transform_nodes)
    for i, obj in enumerate(sorted_objects):
        if no_groups and is_group(obj):
            _log_line(obj, " is a group")
        else:
            _log_line(i, ": ", obj)
            if reverse:
                mc.reorder(obj, front=True)
            else:
                mc.reorder(obj, back=True)


def get_selected_geo_recursive(output_transforms=True, output_shapes=True):
    """
    Get all selected transform and/or shape nodes in the hierarchy of the
    elected DAG node.

    Parameters
    ----------
    output_transforms:
        Boolean, default True
        Set to True to include transform nodes in the result.

    output_shapes:
        Boolean, default True
        Set to True to include shape nodes in the result.

    Test procedure
    --------------
    - Create a known hierarchy of (primitive) geometries
    - Select the root node.
    - Check result of function. (flat list)

    - Select an intermediate node.
    - Check result of function. (flat list)

    - Select a lead node
    - Check result of function. (should return only one node's transform/shape

    """
    # TODO: Can you find a better name, it says what it does, but still sounds
    #       a bit weird and that's confusing.
    # TODO: opportunity for generalization: param for list of nodes and bool
    #       for use the selection yes/no
    selection = mc.ls(sl=True, transforms=True)
    geometries = []

    for sel in selection:
        children = get_geom_children(sel, output_transforms, output_shapes)
        geometries.extend(children)

    _log_line("@get_selected_geo_recursive, geometries: ", geometries)
    return geometries


#TODO:  There are probably better ways to get different types of output without
#       having an explosion in the number of methods and functions.
#       "Model" functions might add confusing, though it is what many MEL
#       command seem to do.
def get_geom_children(transform_node,
                      output_transforms=True,
                      output_shapes=True):
    """
    Gets all geometries below in the hierarchy as a list of pairs, which
    contain the geometry's transform and shape node in that order.

    This method is meant to get descendants on just one transform node.

    Parameters
    ----------
    output_transforms:
        Boolean, default True
        Set to True to include transform nodes in the result.

    output_shapes:
        Boolean, default True
        Set to True to include shape nodes in the result.

    Test procedure
    --------------
    Similar to get_selected_geo_recursive, but pass name of a (transform) node
    instead of selecting it.

    """
    shapes = mc.listRelatives(transform_node, fullPath=True, ad=True,
                              type='shape')
    _log_line("shapes: ", shapes)

    if not output_transforms:
        return shapes
    else:
        geometries = []

        for shp in shapes:
            xforms = mc.listRelatives(shp, parent=True, fullPath=True)
            _log_line("xforms: ", xforms)
            _log_line("len(xforms): ", len(xforms))

            if output_shapes:
                geometries.append((xforms[0], shp))
            else:
                geometries.append(xforms[0])

    _log_line("@get_geom_children, geometries: ", geometries)
    return geometries


def multi_uv_snapshot(geometries=None, directory=None, file_format="tif",
                      x_res=1024, y_res=1024, color=(255, 0, 0)):
    """
    This will create a separate UV snapshot for all given geometries

    If geometries is None, it will use the ones that are currently selected.
    Unfortunately it has to change your selection because the uvSnapshot is
    based on the selection. So in order to make separate snapshots we need to
    select the geometries one by one in a loop.

    returns:
        type str
        directory where the images were saved.

    Parameters
    -----------

    geometries:
        list of DAG nodes, or None
        If value is None: export UV snapshots for the selected objects in the
        viewport
        If value is a list of DAG nodes: export UV snapshots for these
        geometries.

    directory:
        Path. Default = None,
        Path of directory to write UV snapshot files to.
        If None, use the standard locations which is the directory <images>
        in the Maya project.

    file_format:
        extension of image format to use for the UV snapshot.
        The formats are the ones supported by the MEL command 'uvSnapshot'.


    Test Procedure
    --------------
    - Create a fresh Maya project so there can be no UV snapshots in images.
    - Create a few primitive geometries.
    - Select them.
    - Call multi_uv_snapshot (w/ default params)
    - Check if the expected UV snapshots are there.

    - Delete all UV shapshots created above.
    - Call multi_uv_snapshot with non-default params,
        e.g.: file_format="png", different resolution, different directory.
    - Check if the expected snapshots are there.

    """

    geos_to_export = geometries
    if not geometries:
        geos_to_export = get_selected_geo_recursive(output_shapes=False)
        _log_line("geos_to_export: ", geos_to_export)

    # If directory was not specified (None) use the default location.
    img_dir = directory
    if not directory:
        #set directory to <project dir>/images
        root_dir = mc.workspace(q=True, rd=True)
        # Path separator has to be "/" in Maya, whatever the OS.
        img_dir = root_dir + "/images"

    # We need to change the selection to output the UV snapshots one by one.
    # Store original selection so we can restore it.
    orig_selection = mc.ls(sl=True)

    try:
        for g in geos_to_export:
            mc.select(g)
            selection = mc.ls(sl=True)
            _log_line("selection: ", selection)

            short_name = selection[0]
            short_name = short_name.replace("|", "_")
            _log_line("short_name: ", short_name)

            # Joining path elements with a explicit "/" because Maya paths
            # have to use "/", whatever the OS
            uv_filename = "uv_" + short_name + "." + file_format
            uv_filename = img_dir + "/" + uv_filename
            _log_line("uv_filename: ", uv_filename)

            # The actual UV snapshot
            mc.uvSnapshot(overwrite=True, name=uv_filename,
                          fileFormat=file_format,
                          r=color[0], g=color[1], b=color[2],
                          xResolution=x_res, yResolution=y_res)

    finally:
        # Restore original selection
        mc.select(orig_selection)

    return img_dir

def material_from_template_ui():
    """
    Sets up a shader network quickly using the selected network as a template.

    Each (important) node in the network contains one string that will be
    replaced with a networks "base" name.
    For example "TEMPLATE" could be replaced with "wood", "concrete" or
    "chrome", etc.

    This keeps names unique and keeps the network readable. Nodes that don't
    have the placeholder just get default names with a number at the end.


    Test procedure
    --------------
    If material_from_template passes test, then this is only a thin UI wrapper
    around it. However you can repeat the same steps via the UI.

    """

    # get the selection
    selection = mc.ls(sl=True)
    _log_line("selection:", selection)

    # Sanity checks: one and only one template shader should be selected.
    if not selection:
        mc.confirmDialog(title="Nothing selected",
                         message="Nothing selected.\nSelect material that serves as your template material",
                         button="OK",
                         defaultButton="OK")
    elif len(selection) > 1:
        mc.confirmDialog(title="More than one object selected",
                         message="You can only select one template material at a time for this tool.",
                         button="OK",
                         defaultButton="OK")

        # Bail out after error.
        # Keeps the if-else tree from getting too deep
        return

    # Template selection OK, continue
    # The first item in <selection> is the shader node that serves as the
    # template. Its network will be duplicated.
    orig_shader_node = selection[0]
    _log_line("origShaderNode: ", orig_shader_node)
    
    
    # Get a name for the new shader, this also serves as a base name for
    # its network nodes.
    # TO DO: needs a nicer UI with just one dialog, not 2 separate ones.
    # Find out how. Either MEL or Qt.
    prompt_result = mc.promptDialog(title="Search-replace",
                                    message="Base name for material:",
                                    button=("OK", "Cancel"),
                                    defaultButton="OK",
                                    cancelButton="Cancel",
                                    dismissString="Cancel")
    base_name = mc.promptDialog(query=True)
    _log_line("base_name: ", base_name)
        
    if prompt_result == "OK":    
        # Check if input isn't empty or just whitespace.
        if not base_name.strip():
            mc.confirmDialog(title="Base name not filled in",
                                     message="You have to fill in a value for base_name.",
                                     button="OK",
                                     defaultButton="OK")
            # Bail out to keep if-else tree shallow.
            return
        
        mc.promptDialog(title="Search-replace",
                        message="Word/text to replace:",
                        button=("OK", "Cancel"),
                        defaultButton="OK",
                        cancelButton="Cancel",
                        dismissString="Cancel")
        old_string = mc.promptDialog(query=True)
        _log_line("old_string: ", old_string)
        
        if prompt_result == "OK":
            _log_line(old_string)
            # Check if input isn't empty or just whitespace.
            if not old_string.strip():
                mc.confirmDialog(title="Word/text to replace not filled in",
                                         message="You have to fill in a value for word/text to replace.",
                                         button="OK",
                                         defaultButton="OK")
            elif old_string == base_name:
                mc.confirmDialog(title="Base name identical to string to replace",
                                 message="The base name and the string to replace should not be identical. Please try again with different values.",
                                 button="OK",
                                 defaultButton="OK")
            else:
                # Everything OK, go ahead and create a new material.
                material_from_template(orig_shader_node, old_string, base_name)


def material_from_template(material_node, string2replace, base_name):
    """
    Sets up a shader network quickly using the selected network as a template.

    Each (important) node in the network contains one string that will be
    replaced with a networks "base" name.
    For example "TEMPLATE" could be replaced with "wood", "concrete" or
    "chrome", etc.

    This keeps names unique and keeps the network readable. Nodes that don't
    have the placeholder just get default names with a number at the end.

    Parameters
    ----------
    material_node:
        name fo DAG node that's the material node in the shader network

    string2replace:
        Type: str/unicode
        The string in the node names that should be replaced with the value
        of <base_name>

    base_name:
        For practical intents the name of the material,
        More accurately, the part of the name that is used in all its network
        nodes to make it clear they belong to that material.
        For example the material node may be called "wood_material", where
        "wood" is the base name, and other nodes may be called,
        - wood_diffuseClr_file
        - wood_diff_place2dTex
        - wood_bumb_file
        - wood_bump_place2dTex
        et cetera


    Test procedure
    --------------

    - Create a shader network or load one from file that has been prepared.
        (create one en in HyperShade use menu File > Export Selected Network,
         then save it as a *.ma file)
    - call material_from_template with various values for base_name.
        - Special cases:
            - Call twice with same base_name,
              --> Should get two networks but with "1" appended to each node
                  name.
            - Call with string2replace that does not occur in the template.
              --> Should get copy of template with "1" appended to each node
                  name.
            - Call with base_name = string2replace
              --> Should get copy of template with "1" appended to each node
                  name.
    """


    # Duplicate and rename shader nodes: replace the string <string2replace>
    # with the name of the geometry and remove the automatically
    # added number at the end.
    shader_network_nodes = mc.duplicate(material_node,
                                        upstreamNodes=True)
    _log_line("shader_network_nodes: ", shader_network_nodes)

    for nd in shader_network_nodes:
        # Remove auto-added number at end of new name
        name_wo_number = nd
        last_digit = len(nd)-1
        for i in range(len(nd)-1, 2, -1):
            _log_line("i: ", i)
            if not nd[i] in string.digits:
                # Found where the number starts
                break
            else:
                last_digit = i
        _log_line("last_digit: ", last_digit)

        # Chop the number off the name to create a clean new name,
        # then rename the node.
        name_wo_number = nd[:last_digit]
        _log_line("name_wo_number: ", name_wo_number)
        new_name = name_wo_number.replace(string2replace, base_name)
        _log_line("new_name: ", new_name)
        mc.rename(nd, new_name)


def rename_shader_nodes_ui():
    """
    Interactively renames the selected nodes of a shader network,
    searching and replacing a specific string.
    This does not work recursively. Only selected nodes are renamed to give
    you more control.

    - Select the any nodes you want to rename.
    - Call function / press button
    - Enter which string you'd like to replace
    - Enter the new text to replace that with.
    The nodes will now be renamed.

    Test Procedure
    --------------
    - Create a shader net or load one from file.
    - Make various selections of nodes in the shader net and execute
      instruction above on those selections.
        - single node
        - all nodes
        - only a few nodes

    """
    # get the selection
    shader_net_nodes = mc.ls(sl=True)
    _log_line("shader_net_nodes:", shader_net_nodes)

    for nd in shader_net_nodes:
        obj_type = mc.objectType(nd)
        _log_line("Node '%s' has object type '%s' " % (nd, obj_type))

    # TO DO: needs a nicer UI with just one dialog, not 2 separate ones.
    # Find out how. Either MEL or Qt.
    mc.promptDialog(title="Search-replace",
                    message="Word/text to replace:",
                    button=("OK", "Cancel"),
                    defaultButton="OK",
                    cancelButton="Cancel",
                    dismissString="Cancel")
    old_string = mc.promptDialog(query=True)
    
    # Check if input isn't empty or just whitespace.
    if not old_string.strip():
        mc.confirmDialog(title="Old string not filled in",
                                 message="You have to fill in a value for old string.",
                                 button="OK",
                                 defaultButton="OK")
        
    else:
        mc.promptDialog(title="Search-replace",
                        message="Replace it with:",
                        button=("OK", "Cancel"),
                        defaultButton="OK",
                        cancelButton="Cancel",
                        dismissString="Cancel")
        new_string = mc.promptDialog(query=True)
        
        if not new_string.strip():
            mc.confirmDialog(title="New string not filled in",
                             message="You have to fill in a value for new string.",
                             button="OK",
                             defaultButton="OK")
                             
        elif old_string == new_string:
            mc.confirmDialog(title="Old and new string identical",
                             message="The old and the new string should not be identical. Please try again with different values.",
                             button="OK",
                             defaultButton="OK")
                             
        else:
            rename_shader_nodes(shader_net_nodes, old_string, new_string)


#TODO:  Add option to replace names in whole shader net, with only top node
#       selected.
def rename_shader_nodes(shader_net_nodes, search_string, replace_with):
    """
    Renames nodes of a shader network, searching and replacing a specific
    string.


    Test Procedure
    --------------
    - Create a shader net or load one from file.
    - Programmatically make various selections of nodes in the shader net
      and call rename_shader_nodes.

    - special cases to test:
        - search_string not present in network
        - search_string = replace_with
        TODO: Clarify what desired behavior should be.

    """
    # TO DO: can this be used generally, like the menu
    #        Modify > Search and Replace Names ?
    # TO DO: #refactor. Could be replaced with MEL command,
    #        e.g.: searchReplaceNames "FOO" "BAR" "selected";

    for nd in shader_net_nodes:
        new_name = nd.replace(search_string, replace_with)
        _log_line("new_name: ", new_name)
        mc.rename(nd, new_name)


class TrackerConePlacer(object):
    """
    Creates cones to visualize camera tracking markers, which are normally
    locators.
    """

    @staticmethod
    def run():
        """
        Run this class as a tool, easily.
        """