// ===========================================================================
// Copyright 2018 Autodesk, Inc. All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy form.
// ===========================================================================


// Load plugin containing description of custom rig retargeter nodes
{
    string $name = "retargeterNodes";
    if ( `pluginInfo -query -loaded $name` == false )
        loadPlugin -quiet $name;
}

//
// This file contains MEL functions that wrap python retargetting code found in retargeter.py
// The functions should have no external dependencies.
// Descriptions for each function follow.
//

// local function to check whether an attribute is locked
proc int isLocked( string $attr, string $dest )
{
    return ( eval( `format -s $dest -s $attr "getAttr -lock ^1s.^2s"` ) ||
             eval( `format -s $dest -s $attr "getAttr -lock ^1s.^2sX"` ) ||
             eval( `format -s $dest -s $attr "getAttr -lock ^1s.^2sY"` ) ||
             eval( `format -s $dest -s $attr "getAttr -lock ^1s.^2sZ"` ) );
}

// local function to unlock an attribute
proc unlock( string $attr, string $dest )
{
    eval( `format -s $dest -s $attr "setAttr -lock 0 ^1s.^2s"` );
    eval( `format -s $dest -s $attr "setAttr -lock 0 ^1s.^2sX"` );
    eval( `format -s $dest -s $attr "setAttr -lock 0 ^1s.^2sY"` );
    eval( `format -s $dest -s $attr "setAttr -lock 0 ^1s.^2sZ"` );
}

proc int lockedAttributeDialogChoice( string $dest, string $args[] )
{
    // Make sure that the dialog message is grammatically correct. This will depend
    // on the number of arguments provided by the user.
    string $msg;
    if( size( $args ) > 1 )
        $msg = `format -s $args[0] -s $args[1] (uiRes("m_retargeter.kPlural"))`;
    else
        $msg = `format -s $args[0] (uiRes("m_retargeter.kSingular"))`;

    // Show dialog ...
    string $result = `confirmDialog
        -title (uiRes("m_retargeter.kTitle"))
        -message $msg
        -button (uiRes("m_retargeter.kIgnore"))
        -button (uiRes("m_retargeter.kUnlockAll"))
        -button (uiRes("m_retargeter.kCancel"))`;

    // .. and act according to the user's choice

    // User cancels the mapping operation
    if( $result == uiRes( "m_retargeter.kCancel" ) )
        return 0;
    // User choses to unlock all attributes
    else if( $result == uiRes( "m_retargeter.kUnlockAll" ) )
    {
        unlock( "translate", $dest );
        unlock( "rotate", $dest );
        return 3;
    }
    // User choses to ignore the locked attributes and map only to those that are unlocked
    else if( $result == uiRes( "m_retargeter.kIgnore" ) )
    {
        if( size( $args ) > 1 )
            return 0;
        else if( $args[0] == uiRes( "m_retargeter.kTranslation" ) )
            return 1;
        else if( $args[0] == uiRes( "m_retargeter.kRotation" ) )
            return 2;
    }
}

proc string getConnectedAttr( string $node, string $attr )
{
    string $s = ( $node + "." + $attr );
    if ( objExists( $s ) )
    {
        string $conns[] = `listConnections $s`;
        if ( size( $conns ) == 1 )
            return $conns[0];
    }
    return "";
}

proc string MappingGetBody( string $mapping )
{
    return getAttr( $mapping + ".bodyPart" );
}

proc string MappingGetType( string $mapping )
{
    int $i = getAttr( $mapping + ".type" );
    return ( $i == 0 ) ? "T" : "R" ;
}

proc int    MappingGetId( string $mapping )
{
    return getAttr( $mapping + ".id" );
}

proc string MappingGetDestinationRig( string $mapping )
{
    return getConnectedAttr( $mapping, "destinationRig" );
}

proc       MappingSetOffset( string $mapping, float $x, float $y, float $z )
{
    setAttr ( $mapping + ".offsetX" ) $x;
    setAttr ( $mapping + ".offsetY" ) $y;
    setAttr ( $mapping + ".offsetZ" ) $z;
}

proc float[] MappingGetOffset( string $mapping )
{
    return getAttr( $mapping + ".offset" );
}

proc string GenerateUUID()
{
    return python( "import uuid; uuid.uuid4().hex;" );
}

global proc string[] RetargeterGetMappings( string $retargeter )
{
    string $rhs[];
    if ( RetargeterExists( $retargeter ) )
        $rhs = `listConnections ( $retargeter  + ".mappings" )` ;
    return $rhs;
}

global proc string RetargeterGetSource( string $retargeter )
{
    return ( RetargeterExists( $retargeter ) ) ? 
        getConnectedAttr( $retargeter, "source" ) :
        "" ;
}

global proc string RetargeterGetDestination( string $retargeter )
{
    return ( RetargeterExists( $retargeter ) ) ? 
        getConnectedAttr( $retargeter, "destination" ) :
        "" ;
}


global proc int   RetargeterGetNumMappings( string $retargeter )
{
    return ( RetargeterExists( $retargeter ) ) ?
        size( `listConnections ( $retargeter + ".mappings" )` ) : 0 ;
}

global proc string RetargeterTypeEnumToString( int $i )
{
    if ( $i < 0 || $i > 1 )
        error( (uiRes("m_retargeter.kInvalidTypeEnum")) );
    return ( $i == 0 ) ? "T" : "R";
}

// 0 is for failure
// 1 is for rotation success
// 2 is for translation success
// 3 is for success in both
global proc int RetargeterValidateMapping( string $name, string $body, string $type, string $dest )
{
    if ( RetargeterExists( $name ) == false )
        return 0;

    // If no destination is specified try to figure it out based on the destination of the current retargeter
    if ( size($dest) == 0 )
    {
        string $other = ( $type == "T" ) ? "R" : "T";
        $dest = python( `format -s $name -s $body -s $other "^1s.getMapping( bodyPart='^2s', type='^3s' ).getDestinationRig()"` );
    }

    string $args[] = {};

    int $tLocked = isLocked( "translate", $dest );
    int $rLocked = isLocked( "rotate",    $dest );

    // Start by determining if the locked state will interfere with
    // the desired mapping the user wishes to perform. If there is
    // no problem, return the appropriate success type.
    if( $tLocked == false && $rLocked == false )
        return 3;    // client can proceed with any desired mapping
    if( $rLocked == false && $type == "R" )
        return 1;    // client can proceed with rotation mapping
    if( $tLocked == false && $type == "T" )
        return 2;    // client can proceed with translation mapping

    // Something appears to be locked. To figure out what to do,
    // build up the array of arguments that will be passed to the
    // helper function.
    if ( $tLocked && ( $type == "TR" || $type == "T" ) )
        $args[ size($args) ] = (uiRes("m_retargeter.kTranslation"));

    if ( $rLocked && ( $type == "TR" || $type == "R" ) )
        $args[ size($args) ] = (uiRes("m_retargeter.kRotation"));

    return lockedAttributeDialogChoice( $dest, $args );
}

// Return the name of the retargeter associated with a given character
global
proc string RetargeterGetName( string $character )
{
    //print( _NOL10N( "===> RetargeterGetName\n" ) );

    if ( !size($character) || 
         !objExists( $character ) || 
          nodeType( $character ) != "HIKCharacterNode" )
        return "";

    string $conns[] = `listConnections -d 1 -type CustomRigRetargeterNode ( $character + ".message" )`;
    return ( size( $conns ) != 1 ) ?
        "" :
        $conns[0];

}

// Return the name of the character associated with a given retargeter
global
proc string RetargeterGetCharacter( string $retargeter )
{
    return RetargeterGetDestination( $retargeter );
}

// Determine if the named retargeter node exists in the scene
global
proc int RetargeterExists( string $retargeter )
{
    return  ( ( size( $retargeter ) == 0 ) ||
            ( objExists( $retargeter ) == false ) ||
            ( nodeType( $retargeter ) != "CustomRigRetargeterNode" ) ) ?
        false :
        true  ;
}

// Query the system to see if the named retargeter has the destination rig connected to the source
global
proc int RetargeterIsConnected( string $retargeter )
{
    //print( _NOL10N( "===> RetargeterIsConnected\n" ) );

    return ( RetargeterExists( $retargeter ) ) ? 
        getAttr( $retargeter + ".connected" ) :
        false ;

}

// Connect the named retargeter source to the destination elements
global
proc int RetargeterConnect( string $retargeter )
{
    //print( _NOL10N( "===> RetargeterConnect\n" ) );

    if ( RetargeterExists( $retargeter ) && RetargeterIsConnected( $retargeter ) == false )
    {
        string $attr = ( $retargeter + ".destination" );

        string $dest = "";
        if ( objExists( $attr ) )
        {
            string $conns[] = `listConnections $attr`;
            if ( size( $conns ) )
                $dest = "'" + $conns[0] + "'" ;
        }

        string $fmt = stringArrayToString( { 
            "import maya.app.hik.retargeter",
            "^1s = maya.app.hik.retargeter.HIKRetargeter(^2s)",
            "if ^1s.getDestination() is not None:",
            "    ^1s.fromGraph( '^3s' )",
            "    ^1s.connect()" }, "\n" );

        string $var = ( "HIKRetargeter_" + GenerateUUID() );
        python( `format -s $var -s $dest -s $retargeter $fmt` );

        setAttr                ( $retargeter + ".connected" ) true ;
        setAttr -type "string" ( $retargeter + ".pythonVar" ) $var ;
    }

    return RetargeterIsConnected( $retargeter );
}

// Disconnect the named retargeter source from the destination elements
global
proc int RetargeterDisconnect( string $retargeter )
{
    if ( RetargeterExists( $retargeter ) && RetargeterIsConnected( $retargeter ) == true )
    {
        string $var = getAttr( $retargeter + ".pythonVar" );

        // $var will not exist if a scene is first saved with the 
        // retargeter in the connected state and then reloaded into a 
        // instance of Maya. Therefore, before trying to cleanup the 
        // node n/w, we initialize a python oject (with the same 
        // name as the old variable) using the scene graph to provide
        // needed information
        //
        {
            string $attr = ( $retargeter + ".destination" );
            string $dest = "";
            if ( objExists( $attr ) )
            {
                string $conns[] = `listConnections $attr`;
                if ( size( $conns ) )
                    $dest = "'" + $conns[0] + "'" ;
            }

            string $fmt = stringArrayToString( {
                "import maya.app.hik.retargeter",
                "^1s = maya.app.hik.retargeter.HIKRetargeter(^2s)",
                "if ^1s.getDestination() is not None:",
                "   ^1s.fromGraph( '^3s' )" }, "\n" );

            python( `format -s $var -s $dest -s $retargeter $fmt` );
        }

        // Now that we know $var is valid, delete it. This will also 
        // disconnect the node network connecting the source to the 
        // destination character
        {
            string $fmt = stringArrayToString( {
                "try:",
                "    del ^1s",
                "except:",
                "    pass" }, "\n" );

            python( `format -s $var $fmt` );
        }

        setAttr                ( $retargeter + ".connected" ) false;
        setAttr -type "string" ( $retargeter + ".pythonVar" ) "";
    }

    return RetargeterIsConnected( $retargeter );
}

// Create a new retargeter connected to the specified destination.
// $dest must be the name of a valid HIKCharacterNode
global
proc string RetargeterCreate( string $dest )
{
    if ( size( $dest ) == 0 || nodeType( $dest ) != "HIKCharacterNode" )
        error( (uiRes("m_retargeter.kInvalidDestination")) );
    
    string $retargeter = createNode( "CustomRigRetargeterNode" );
    if ( size( $retargeter ) )
    {
        string $source = hikGetStateToGlobalSk( $dest, 1 );
        connectAttr ( $source + ".message" ) ( $retargeter + ".source" );
        connectAttr ( $dest   + ".message" ) ( $retargeter + ".destination" );

        setAttr                ( $retargeter + ".connected" ) false;
        setAttr -type "string" ( $retargeter + ".pythonVar" ) ""   ;
    }
    return $retargeter;
}

// Delete the named retargeter
global
proc RetargeterDelete( string $retargeter )
{
    if ( !RetargeterExists( $retargeter ) )
        return;

    RetargeterDisconnect( $retargeter );

    for ( $m in RetargeterGetMappings( $retargeter ) )
        if ( nodeType( $m ) == "CustomRigDefaultMappingNode" )
            delete $m;

    if ( objExists( $retargeter ) )
        delete $retargeter;
}

// Get a list of the attributes on the destination rig being controlled by the retargeter
global
proc string[] RetargeterDestinationAttributes( string $retargeter )
{
    // print( ( "===> RetargeterDestinationAttributes\n" ) );

    string $attrs[];

    for ( $m in RetargeterGetMappings( $retargeter ) )
    {
        string $attr = ( MappingGetType( $m ) == "T" ) ? ".translate" : ".rotate" ;
        string $node = MappingGetDestinationRig( $m );
        if ( size( $node ) && objExists( $node ) )
            $attrs[ size( $attrs ) ] = ( $node + $attr ) ;
    }

    return $attrs;
}

// Get a list of indices for all assigned retargeters
global
proc int[] RetargeterAssignedElementIndices( string $retargeter )
{
    //print( _NOL10N( "===> RetargeterAssignedElementIndices\n" ) );

    int $inx[];
    for ( $m in RetargeterGetMappings( $retargeter ) )
        $inx[ size( $inx ) ] = MappingGetId( $m );
    return $inx;
}

// Get the name of the destination objects being used by the translate and rotation mappings (if they exist)
global
proc string[] RetargeterGetDestinationsForMapping( string $retargeter, string $body )
{
    // print( _NOL10N( "===> RetargeterGetDestinationsForMapping\n" ) );

    string $names[] = { "", "" };

    if ( size( $body ) )
    {
        for ( $m in RetargeterGetMappings( $retargeter ) )
        {
            if ( MappingGetBody( $m ) != $body )
                continue;

            string $type =  MappingGetType( $m );
            if      ( $type == "T" )
                $names[0] = MappingGetDestinationRig( $m );
            else if ( $type == "R" )
                $names[1] = MappingGetDestinationRig( $m );
        }
    }

    return $names;
}

// Read the set of mappings for the named retargeter from the user specified file
global
proc int RetargeterReadFromFile( string $retargeter, string $filename )
{
    //print( _NOL10N( "===> RetargeterReadFromFile\n" ) );
    
    if ( size( $filename ) && RetargeterExists( $retargeter ) )
    {
        // Before importing the new retargeter, check if we are connected to a source.
        // If so, disconnect from the source temporarily.
        int $connected = RetargeterIsConnected( $retargeter );
        if ( $connected )
            RetargeterDisconnect( $retargeter );
        
        // Next generate the new retargeter specified in the file
        string $dst = RetargeterGetDestination( $retargeter );
        string $var = ( "temporary_" + GenerateUUID() );
        string $fmt = stringArrayToString( {
            "import maya.app.hik.retargeter",
            "import maya.cmds as cmds",
            "^4s = maya.app.hik.retargeter.HIKRetargeter()",
            "if ^4s.read( '^1s', '^2s' ):",
            "    ^4sNode = ^4s.toGraph()",
            "    del ^4s",
            "    cmds.rename( '^3s', '^4s' )",
            "    cmds.rename( ^4sNode, '^3s' )" }, "\n" );

        python( `format -s $filename -s $dst -s $retargeter -s $var $fmt` );

        // Remove the old retargeter
        RetargeterDelete( $var );

        // Restore the connection state
        if ( $connected )
            RetargeterConnect( $retargeter );

        return true;
    }
    return false;
}

// Write the set of mappings for the named retargeter to the user specified file
global
proc RetargeterWriteToFile( string $retargeter, string $filename )
{
    //print( _NOL10N( "===> RetargeterWriteToFile\n" ) );

    if ( size( $filename ) && RetargeterExists( $retargeter ) )
    {
        // In order to serialize the retargeter to a file, we will create a
        // an in-memory version of the retargeter, deserialize its state from the
        // DG and use available methods to write the object to file. 
        //
        // To avoid negative side-effects on networks of connected mapping noddes,
        // we first check if the retargeter is connected and diconnect mappings.
        int $connected = RetargeterIsConnected( $retargeter );
        if ( $connected )
            RetargeterDisconnect( $retargeter );

        // Next we dump the object contents to file (see above)
        string $var = ( "temporary_" + GenerateUUID() );
        string $fmt = stringArrayToString( {
            "import maya.app.hik.retargeter",
            "^4s = maya.app.hik.retargeter.HIKRetargeter( '^1s' )",
            "^4s.fromGraph( '^2s' )",
            "^4s.write( '^3s' )",
            "del ^4s" }, "\n" );

        string $dst = RetargeterGetDestination( $retargeter );
        python( `format -s $dst -s $retargeter -s $filename -s $var $fmt` );

        // Lastly we reconnect mappings (if needed)
        if ( $connected )
            RetargeterConnect( $retargeter );
    }
}

// Determine the type of mappings that exist on the named retargeter for the named body
global
proc int RetargeterHasMapping( string $retargeter, string $body )
{
    //print( _NOL10N( "===> RetargeterHasMapping\n" ) );

    for ( $m in RetargeterGetMappings( $retargeter ) )
        if ( MappingGetBody( $m ) == $body )
            return true;

    return false;
}

// Add a mapping to the named retargeter, with options provided as arguments
global
proc int RetargeterAddMapping( string $retargeter, string $body, string $type, string $rig, string $id )
{
    //print( _NOL10N( "===> RetargeterAddMapping\n" ) );

    if ( !RetargeterExists( $retargeter ) )
        return false;

    int $connected = RetargeterIsConnected( $retargeter );
    if ( $connected )
        RetargeterDisconnect( $retargeter );

    // If no destination is specified try to figure it out based on the destination of the current retargeter
    if ( size($rig)==0 )
    {
        string $other = ( $type == "T" )? "R" : "T";
        for ( $m in RetargeterGetMappings( $retargeter ) )
            if ( MappingGetBody( $m ) == $body && MappingGetType( $m ) == $other )
            {
                $rig = getConnectedAttr( $m, "destinationRig" );
                break;
            }
    }

    // Before proceeding to add the mapping, validate that it is valid and that
    // target attribute channels are not locked.
    int $response = RetargeterValidateMapping( $retargeter, $body, $type, $rig );
    if ( ( $response == 3 ) || 
         ( $response == 1 && $type == "R" ) || 
         ( $response == 2 && $type == "T" ) )
    {
        // If the new mapping appears valid ...

        // Determine the link number of bodies with non-zero link #s
        string $linkNum = "0";
        {
            int $n = size("Spine");
            int $m = size( $body ); 
            if ( startsWith( $body, "Spine" ) && ( $m > $n ) )
                $linkNum = 1 + substring( $body, $n+1, $m );
        }

        // Use the python interface to create a new mapping
        string $var = ( "temporary_" + GenerateUUID() );
        string $fmt = stringArrayToString( {
            "import maya.app.hik.retargeter as r",
            "^8s  = r.HIKRetargeter.createDefaultMapping( '^1s', '^2s', '^3s', '^4s', '^5s', ^6s, ^7s )",
            "r.DefaultRetargeter.toGraph( ^8s, '^1s' )",
            "del ^8s" }, "\n" );

        string $src  = RetargeterGetSource( $retargeter );
        string $dst  = RetargeterGetDestination( $retargeter );

        string $node = python( `format -s $src -s $dst -s $body -s $rig -s $type -s $id -s $linkNum -s $var $fmt` );

        // Connect the new mapping to the parent retargeter
        int $num = RetargeterGetNumMappings( $retargeter );
        connectAttr ( $node + ".message" ) ( $retargeter + ".mappings[" + $num + "]" );
    }

    // Reconnect the n/w of nodes if it was connected prior to adding the 
    // new mapping
    if ( $connected )
        RetargeterConnect( $retargeter );

    return true;
}

// Delete the specified mapping from the named retargeter
global
proc RetargeterDeleteMapping( string $retargeter, string $body, string $type )
{
    //print( _NOL10N( "===> RetargeterDeleteMapping\n" ) );

    if ( !RetargeterExists( $retargeter ) )
        return;

    int $connected = RetargeterIsConnected( $retargeter );
    if ( $connected )
        RetargeterDisconnect( $retargeter );

    // Walk the list of connections until the desired mapping is found.
    // Once found, move remaining connections on the "mappings" plugs
    // so there are no gaps in the sequence. This allows the len($mappings)
    // to be used to determine where new mappings should be inserted in
    // the array.
    int $found = false;
    string $mappings[] = RetargeterGetMappings( $retargeter );
    int $n = size($mappings);
    for ( $i=0; $i<$n; $i++ )
    {
        string $m = $mappings[$i];
        if ( $found == false )
        {
            if ( MappingGetBody( $m ) == $body && MappingGetType( $m ) == $type )
            {
                delete $m;
                $found = true;
            }
        }

        int $j = $i+1;
        if ( $found && $j != $n )
        {
            disconnectAttr ( $mappings[$j] + ".message" ) ( $retargeter + ".mappings[" + $j + "]" );
            connectAttr    ( $mappings[$j] + ".message" ) ( $retargeter + ".mappings[" + $i + "]" );
        }
    }

    if ( $connected )
        RetargeterConnect( $retargeter );

}

// Remove all mappings associated with the named retargeter
global
proc RetargeterDeleteAllMappings( string $retargeter )
{
    int $connected = RetargeterIsConnected( $retargeter );
    if ( $connected )
        RetargeterDisconnect( $retargeter );

    for ( $m in RetargeterGetMappings( $retargeter ) )
        delete $m;

    if ( $connected )
        RetargeterConnect( $retargeter );
}

// Set offsets for the specified mapping
global
proc RetargeterSetMappingOffsets( string $retargeter, string $body, string $type, string $x, string $y, string $z )
{
    int $connected = RetargeterIsConnected( $retargeter );
    if ( $connected )
        RetargeterDisconnect( $retargeter );

    for ( $m in RetargeterGetMappings( $retargeter ) )
        if ( MappingGetBody( $m ) == $body && MappingGetType( $m ) == $type )
        {
            MappingSetOffset( $m, $x, $y, $z );
            break;
        }

    if ( $connected )
        RetargeterConnect( $retargeter );
}

// Get offsets for the specified mapping
global
proc float[] RetargeterGetMappingOffsets( string $retargeter, string $body, string $type )
{
    for ( $m in RetargeterGetMappings( $retargeter ) )
        if ( MappingGetBody( $m ) == $body && MappingGetType( $m ) == $type )
            return MappingGetOffset( $m );
    return {};
}

// Remove the mappings that do not exist
global
proc RetargeterRemoveInvalidMappings( string $name , string $bodies[] )
{
    //print( _NOL10N( "===> RetargeterRemoveInvalidMappings\n" ) );

    string $warnFormat = (uiRes("m_retargeter.kHIKBadMapping"));

    for( $body in $bodies )
    {
        string $mappings[] = RetargeterGetDestinationsForMapping( $name, $body );
        for( $m in $mappings )
        {
            if( size( $m ) == 0 || objExists( $m ) )
                continue;

            warning( `format -s $m $warnFormat` );
            RetargeterDeleteMapping( $name, $body, "T" );
            RetargeterDeleteMapping( $name, $body, "R" );
            break;
        }
    }
}