#Jack Frost Hair System
import sys
from mcRig.Utilities.Name import name_count
from mcRig.UI.UITools import AskName
import mcRig.RigTools.Ribbon as RIB
ribbon = RIB()
import mcRig.CtrlTools as CT
my_ctrl = CT.Control()
import mcRig.RigTools.Attribute as ATTR



def convert_VecPth_to_hair_rig(input = False):

    if not input:
        input = cmds.ls(sl=True)

    for every in input:
        #Make a jointFollowChain on the Ribbon
        stuff = cmds.listRelatives(every,c=True)
        trgts = [x for x in stuff if '_POS_' in x]
        trgts.sort()
        print(trgts)
        jnts = []
        last = False
        for e in trgts:
            cmds.select(cl=True)
            jnt2 = cmds.joint(n=e.replace('POS','Path').replace('trgt','Offset_Jnt'))
            jnt = cmds.joint(n=e.replace('POS','Path').replace('trgt','Jnt'),rad = .5)
            cmds.delete(cmds.parentConstraint(e,jnt2,mo=False))
            if last:
                cmds.parent(jnt2,last)
            else:
                grp = cmds.group(em=True, n=e.replace('POS','Path').replace('trgt','RIG'))
                cmds.delete(cmds.parentConstraint(jnt,grp,mo=False))
                cmds.parent(jnt2,grp)
            last=jnt
            jnts.append(jnt)

        cmds.makeIdentity(jnts[0], a = True, r = True, pn = True)

    #Make a control system that follows the ribbon Joints
    last = False
    for e in jnts:
        cmds.parentConstraint(e.replace('Jnt','trgt').replace('Path','POS'),e,mo=True)

        #####
        #ctrl = cst.BuildControl(Type = 'D3Cube', Color = 'pink', Size = 1, Snap = e, Name = e.replace('Jnt','Ctrl'))
        #cst.TopGrp(AddName='Drv', Input = [ctrl])
        #cst.TopGrp(Search='Drv',Replace='Offset')
        #tweak = cst.BuildControl(Type = 'D3Cube', Color = 'pink', Size = .5, Snap = e, Name = e.replace('Jnt','Ctrl'))
        #cst.TopGrp(AddName='Offset', Input = [tweak])
        #####

        my_ctrl.name = e.replace('Jnt','Ctrl')
        my_ctrl.idcode = '2D_Pin_Circle'
        my_ctrl.color = 'yellow'
        my_ctrl.size = 3
        my_ctrl.snapto = e
        ctrl = my_ctrl.make_control()
        my_ctrl.top_group(add='Drv', input=[ctrl])
        my_ctrl.top_group(searchfor = 'Drv',replacewith='Offset')

        my_ctrl.name = e.replace('Jnt','Tweak_Ctrl')
        my_ctrl.idcode = '2D_Simple_Circle'
        my_ctrl.color = 'blue'
        my_ctrl.size = 1
        my_ctrl.snapto = e
        tweak = my_ctrl.make_control()
        my_ctrl.top_group(add='Offset', input=[tweak])

        cmds.parent(tweak+'_Offset_Grp',ctrl)

        cmds.select(cl=True)
        bnd = cmds.joint(n=e.replace('Jnt','Bnd').replace('_Path_','_'))
        cmds.parentConstraint(tweak,bnd,mo=False)
        cmds.scaleConstraint(tweak,bnd,mo=False)
        if not cmds.objExists(jnts[0].replace('Jnt','Bnd')+'_Grp'):
            cmds.group(em=True, n=jnts[0].replace('Jnt','Bnd')+'_Grp')
        cmds.parent(bnd,jnts[0].replace('Jnt','Bnd')+'_Grp')

        if last:
            cmds.parent(ctrl+'_Offset_Grp',last)

        cmds.connectAttr(e+'.r',ctrl+'_Drv_Grp.r')
        cmds.connectAttr(e+'.t',ctrl+'_Drv_Grp.t')

        last = ctrl

    return(jnts[0].replace('Jnt','Bnd')+'_Grp', jnts[0].replace('Jnt','Ctrl_Offset')+'_Grp', jnts[0].replace('Jnt','RIG'))





def nurbs_to_hair(input = False, name = False, jointCount = False, anti_stretch = True):
    #select all of the ribbons you wish to turn into hair
    if not input:
        input = cmds.ls(sl=True)

    #give a name for the system to name the ribbons for you just to make sure they are proper.
    if not name:
        name = AskName(Message = 'Name for the Hair System', Title = 'Name')
    check = name_count(name + '_01_Ribn')
    count = check[0]+1

    ribn_list=[]
    for each in input:
        new_name = cmds.rename(each, name+'_%02d_Ribn' % count)
        ribn_list.append(new_name)
        count = count+1

    #We will then for each of the ribbons make a vector path system

    vecPath_list = []
    if not jointCount:
        jointCount = int(AskName(Message = 'How many Joints in system', Title = 'Count'))
    for each in ribn_list:
        build = ribbon.ribbon_vector_setup(uiclear=True, input=[each], jointCount=jointCount)
        vecPath_list.append(build['container'])

    #For each of the vecor path systems we will make a hair system
    rigs = []
    ctrls = []
    bnds = []
    for each in vecPath_list:
        build = convert_VecPth_to_hair_rig(input = [each])  # Bnd, Ctrl, Rig
        bnds.append(build[0])
        ctrls.append(build[1])
        rigs.append(build[2])


        anti_stretch_list = [x for x in cmds.listRelatives(each, c=True, type = 'dagContainer') if x.endswith('CrvPth')]
        moPth_list = cmds.container(anti_stretch_list, q=True, nl=True)

        #moPth_list no holds all the nodes that we need to install the anti stretch system to
        #add in the anti stretch system
        if anti_stretch:
            # for each motionPath we will need to find out what curve is attached to it.
            for moPth in moPth_list:
                curve = cmds.listConnections(moPth+'.geometryPath')[0]
                # if set up dose not exists
                ref_info = curve.replace('_Crv', '_AS_Ref_Crv_CI')
                curve_info = curve.replace('_Crv', '_AS_Crv_CI')
                ref_md = curve.replace('_Crv', '_AS_MD')
                if not cmds.objExists(ref_md):
                    # Duplicate that curve and put it in a reference group
                    ref = cmds.duplicate(curve, n=curve.replace('_Crv', '_AS_Ref_Crv'))[0]
                    # Make a curve info node to get the length of both curves.
                    cmds.createNode('curveInfo', n=ref_info)
                    cmds.createNode('curveInfo', n=curve_info)
                    cmds.createNode('multiplyDivide',n=ref_md)
                    # Connect them together with a MD to get the offset.
                    cmds.connectAttr(curve+'.worldSpace', curve_info+'.inputCurve')
                    cmds.connectAttr(ref+'.worldSpace', ref_info+'.inputCurve')
                    cmds.connectAttr(ref_info+'.arcLength', ref_md+'.input1X')
                    cmds.connectAttr(curve_info+'.arcLength', ref_md+'.input2X')
                    cmds.setAttr(ref_md+'.operation', 2)
                    # Set up the control override if wanted
                    cmds.container(each, e = True, an=[ref_info, curve_info, ref_md])

                    if cmds.objExists('Utility_Cns_Grp'):
                        cmds.scaleConstraint('Utility_Cns_Grp', ref, mo=True)


                # Connect the strand driver system to the motion path with the md offset 0-1
                # make a md for each moPth and then set that to the path value, and then drive the other with the ref_md output
                as_md = cmds.createNode('multiplyDivide',n=moPth+'_AS_Offset_MD')
                cmds.container(each, e = True, an=[as_md])
                cmds.setAttr(as_md+'.input1X', cmds.getAttr(moPth+'.uValue'))
                cmds.connectAttr(ref_md+'.outputX', as_md+'.input2X')
                cmds.connectAttr(as_md+'.outputX', moPth+'.uValue')


    cmds.group(ribn_list, vecPath_list, n=name+'_Utl_Grp')
    cmds.group(bnds, n=name+'_Bnd_Jnt_Grp')
    cmds.group(ctrls, n=name+'_Ctrl_Grp')
    cmds.group(rigs, n=name+'_RIG_Grp')
    for rig in rigs:
        cmds.setAttr(rig+'.inheritsTransform', 0)


    attr = ATTR()
    attr.type = 'enum'
    attr.name = 'HairControls'
    attr.default = 1
    attr.keyable = False
    attr.object = 'Control_Ctrl'
    attr.make_attr()
    cmds.connectAttr('Control_Ctrl.HairControls', name+'_Ctrl_Grp.v')

    # set up the chain blends to and from the ribbons
    attr = ATTR()
    for each in ctrls:
        stuff = [x for x in cmds.listRelatives(each, ad=True) if x.endswith('_Ctrl')]
        final = [x for x in stuff if 'Tweak' not in x]
        #these are the controls we need to add a blend to for turning the chain off from the rest.
        my_ctrl.color = 'red'
        my_ctrl.name = final[0]
        my_ctrl.set_coloroverride()

        attr.type = 'float'
        attr.name = 'Follow'
        attr.min = 0
        attr.max = 1
        attr.default = 1
        attr.object = final[0]
        attr.make_attr()

        #adjust the controls and connect them to the blends
        nodes = []
        for ctrl in final:
            if ctrl != final[0]:
                cmds.rotate(90, 0, 0, ctrl+'.cv[0:*]')

            if ctrl != final[-1]:
                # bc for the trans and rot
                drv = cmds.listRelatives(ctrl, p=True)[0]
                tbc = cmds.createNode('blendColors', n=drv.replace('Grp', 'Follow_T_BC'))
                rbc = cmds.createNode('blendColors', n=drv.replace('Grp', 'Follow_R_BC'))
                # set the bc to the current values for trans and rot
                cmds.setAttr(tbc+'.color2R', cmds.getAttr(drv+'.tx'))
                cmds.setAttr(tbc+'.color2G', cmds.getAttr(drv+'.ty'))
                cmds.setAttr(tbc+'.color2B', cmds.getAttr(drv+'.tz'))

                cmds.setAttr(rbc+'.color2R', cmds.getAttr(drv+'.rx'))
                cmds.setAttr(rbc+'.color2G', cmds.getAttr(drv+'.ry'))
                cmds.setAttr(rbc+'.color2B', cmds.getAttr(drv+'.rz'))
                # connect the final[0] to the bc inputs. 1 is follow
                cmds.connectAttr(final[0]+'.Follow', tbc+'.blender')
                cmds.connectAttr(final[0]+'.Follow', rbc+'.blender')

                driver = cmds.listConnections(drv+'.t')[0]
                # connect the incomming data to the bc in color 1
                cmds.connectAttr(driver+'.t', tbc+'.color1')
                cmds.connectAttr(driver+'.r', rbc+'.color1')
                # set the bc color 2 to the current values for trans and rot
                # connect the out of blend to the origanal driver gorup
                cmds.connectAttr(tbc+'.output', drv+'.t', f=True)
                cmds.connectAttr(rbc+'.output', drv+'.r', f=True)

                nodes.append(tbc)
                nodes.append(rbc)

        cmds.container(each.replace('_Path_', '_Ribn_Setup_').replace('_Ctrl_Offset_Grp', '_VecPth'), e = True, an=nodes)
        cmds.setAttr(each+'.inheritsTransform', 0)

        # now we can override the above group inputs with a md system for translate and rotate.

#Make nurbs strips like normal.
#then select them and run this
# or supply input with a list of the nurbs surfaces.
# once built i can have a deformer on the nurbs (skin fk chain), then i can move the deformer and all the hair groups and it shoudl move and stay together
# each chain has a red tip, this is where you can blend the fk chain on and off from the ribbon system, when off only the first control follows the ribbon, so skinning of the ribbon is important here
nurbs_to_hair(input = False, name = False, jointCount = False, anti_stretch = True)