func_utils

# -*- coding: utf-8 -*-
from pyfbsdk import *


def offset_character_animation_at_frame(frame=0, all_chars=True,
                                        frame_anim=True):

    # get list of current/all characters
    if all_chars:
        char_list = FBSystem().Scene.Characters
    else:
        char_list = [FBApplication().CurrentCharacter]

    # get initial timespan
    lStartFrame = FBSystem().CurrentTake.LocalTimeSpan.GetStart().GetFrame()
    lEndFrame = FBSystem().CurrentTake.LocalTimeSpan.GetStop().GetFrame()

    # turn on Story mode
    FBStory().Mute = False

    # process character list
    for char in char_list:
        # set timespan
        FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
            FBTime(0, 0, 0, lStartFrame, 0), FBTime(0, 0, 0, lEndFrame, 0))
        # set current character
        FBApplication().CurrentCharacter = char
        # insert character animation track
        track = FBStoryTrack(FBStoryTrackType.kFBStoryTrackCharacter,
                             FBStory().RootFolder)
        track.Name = '{}_charAnimTrack'.format(
            FBApplication().CurrentCharacter.Name)
        track.Details.append(FBApplication().CurrentCharacter)
        # insert take in story mode
        take = FBSystem().CurrentTake
        inserted_clip = track.CopyTakeIntoTrack(take.LocalTimeSpan, take)
        # move inserted clip to given frame
        inserted_clip.Start = FBTime(0, 0, 0, frame)
        # frame new timespan
        FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
            FBTime(0, 0, 0, inserted_clip.Start.GetFrame(), 0),
            FBTime(0, 0, 0, inserted_clip.Stop.GetFrame(), 0))
        # defining plot options and plot to current take
        PlotOptions = FBPlotOptions()
        PlotOptions.ConstantKeyReducerKeepOneKey = True
        PlotOptions.PlotAllTakes = False
        PlotOptions.PlotOnFrame = True
        PlotOptions.PlotPeriod = FBTime(0, 0, 0, 1)
        PlotOptions.PlotTranslationOnRootOnly = True
        PlotOptions.PreciseTimeDiscontinuities = True
        PlotOptions.RotationFilterToApply = FBRotationFilter.kFBRotationFilterGimbleKiller
        PlotOptions.UseConstantKeyReducer = True
        char.PlotAnimation(FBCharacterPlotWhere.kFBCharacterPlotOnSkeleton,
                           PlotOptions)
        # empty Story mode
        for track in FBStory().RootFolder.Tracks:
            for clip in track.Clips:
                clip.FBDelete()
            track.FBDelete()

    # set back original timespan if specified
    if not frame_anim:
        FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
            FBTime(0, 0, 0, lStartFrame, 0), FBTime(0, 0, 0, lEndFrame, 0))

    # turn off Story mode
    FBStory().Mute = True


def offset_generic_animation_at_frame(frame=0, frame_anim=True):
    # get selected components
    lModelList = FBModelList()
    FBGetSelectedModels(lModelList)
    if not lModelList:
        raise ValueError("Select at least one component")

    # get initial timespan
    lStartFrame = FBSystem().CurrentTake.LocalTimeSpan.GetStart().GetFrame()
    lEndFrame = FBSystem().CurrentTake.LocalTimeSpan.GetStop().GetFrame()

    # turn on Story mode
    FBStory().Mute = False

    # set timespan
    FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
        FBTime(0, 0, 0, lStartFrame, 0), FBTime(0, 0, 0, lEndFrame, 0))
    # insert generic animation track and add selected components to it
    track = FBStoryTrack(FBStoryTrackType.kFBStoryTrackAnimation,
                         FBStory().RootFolder)
    track.Name = 'genericAnimTrack'
    for comp in lModelList:
        track.Details.append(comp)
    # insert take in story mode
    take = FBSystem().CurrentTake
    inserted_clip = track.CopyTakeIntoTrack(take.LocalTimeSpan, take)
    # move inserted clip to given frame
    inserted_clip.Start = FBTime(0, 0, 0, frame)
    # frame new timespan
    FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
        FBTime(0, 0, 0, inserted_clip.Start.GetFrame(), 0),
        FBTime(0, 0, 0, inserted_clip.Stop.GetFrame(), 0))
    # plot selected take
    lOptions = FBPlotOptions()
    lOptions.ConstantKeyReducerKeepOneKey = False
    lOptions.PlotAllTakes = False
    lOptions.PlotOnFrame = True
    lOptions.PlotPeriod = FBTime(0, 0, 0, 1)
    lOptions.PlotTranslationOnRootOnly = False
    lOptions.PreciseTimeDiscontinuities = True
    lOptions.RotationFilterToApply = FBRotationFilter.kFBRotationFilterGimbleKiller
    lOptions.UseConstantKeyReducer = False
    FBSystem().CurrentTake.PlotTakeOnSelected(lOptions)

    # empty Story mode
    for track in FBStory().RootFolder.Tracks:
        for clip in track.Clips:
            clip.FBDelete()
        track.FBDelete()

    # set back original timespan if specified
    if not frame_anim:
        FBSystem().CurrentTake.LocalTimeSpan = FBTimeSpan(
            FBTime(0, 0, 0, lStartFrame, 0), FBTime(0, 0, 0, lEndFrame, 0))

    # turn off Story mode
    FBStory().Mute = True


def get_selected_constraints():
    result = []
    for constraint in FBSystem().Scene.Constraints:
        if constraint.Selected:
            result.append(constraint)
    return result


def get_connection_constrait(node):
    constraints = []
    for m in range(node.GetSrcCount()):
        src_node = node.GetSrc(m)
        if type(src_node) == FBConstraint:
            constraints.append(src_node)
    return constraints


def get_time_range():
    take = FBSystem().CurrentTake
    start = take.LocalTimeSpan.GetStart().GetFrame()
    end = take.LocalTimeSpan.GetStop().GetFrame()
    return (start, end)


def get_animation_node_in_relation(box, attr='', flow='in'):
    result = None
    if flow == 'in':
        parent = box.AnimationNodeInGet()
    elif flow == 'out':
        parent = box.AnimationNodeOutGet()
    else:
        return
    for node in parent.Nodes:
        if node.Name == attr:
            result = node
            break
    return result


def get_animation_curve(obj_node):
    all_nodes = []
    for property in obj_node.PropertyList:
        if property and property.IsAnimatable() and property.IsAnimated():
            animation_node = property.GetAnimationNode()
            if animation_node.Nodes:
                for node in animation_node.Nodes:
                    all_nodes.append(node.FCurve)
            if animation_node.FCurve:
                all_nodes.append(animation_node.FCurve)
    return [node for node in all_nodes]


def clear_framekey(obj, start, end):
    ani_curve_list = get_animation_curve(obj)
    for ani_curve in ani_curve_list:
        ani_curve.KeyDeleteByTimeRange(FBTime(0, 0, 0, start),
                                       FBTime(0, 0, 0, end))


def ShiftAnimaion(variable):
    gFilterManager = FBFilterManager()
    gFilter = gFilterManager.CreateFilter("Time Shift And Scale")
    gFilter.PropertyList.Find("Shift").Data = variable
    for lComponent in FBSystem().Scene.Components:
        if lComponent and lComponent.Is(FBModel_TypeInfo()):
            for lNode in lComponent.PropertyList:
                if lNode and lNode.IsAnimatable() and lNode.IsAnimated():
                    gFilter.Apply(lNode.GetAnimationNode(), True)
    print("Animation Shifted")


def remove_ticket(object):
    ani_curve_list = get_animation_curve(object)
    for ani_curve in ani_curve_list:
        for key in ani_curve.Keys:
            key.Time = FBTime(0, 0, 0, key.Time.GetFrame(), 0)


def get_object_by_name(name, no_recursion=False, namespace=''):
    component_list = FBComponentList()
    FBFindObjectsByName(name, component_list, no_recursion, False)
    if namespace:
        component_list = [i for i in component_list if
                          (namespace + ':') in i.FullName]
    return list(component_list)


def create_parent_constraint(name='constraint', source_list=[], target=''):
    constraint = FBConstraintManager().TypeCreateConstraint('Parent/Child')
    constraint.Name = name
    if target:
        constraint.ReferenceAdd(0, target)
    for source in source_list:
        constraint.ReferenceAdd(1, source)
    constraint.Active = True
    constraint.Lock = True
    return constraint


def create_int_property(model, name, value):
    prop = model.PropertyCreate(name, FBPropertyType.kFBPT_int, 'Integer',
                                True, True, None)
    prop.Data = value
    prop.SetAnimated(True)


def create_enum_property(model, name, items):
    prop = model.PropertyCreate(name, FBPropertyType.kFBPT_enum,
                                'Enum', True, True, None)
    enum_list = prop.GetEnumStringList(True)
    for i in items:
        enum_list.Add(i)
    prop.NotifyEnumStringListChanged()
    prop.SetAnimated(True)


def create_ik_constraint(namespace=''):
    """
    创建基本的IK约束，如果有有不完整的则删除
    左手：hand_l
    右手：hand_r
    左脚：foot_l
    右脚：foot_r
    左手IK：ik_hand_l
    右手IK：ik_hand_r
    左脚IK：ik_foot_l
    右脚IK；ik_foot_r
    gun IK：ik_hand_gun

    6个约束：
     IK_Hd_2_HdSk_L：hand_l → ik_hand_l
     IK_Hd_2_HdSk_R：hand_r → ik_hand_r
     IK_Ft_2_FtSk_L：foot_l → ik_foot_l
     IK_Ft_2_FtSk_R：foot_r → ik_foot_r
     IK_Hd_2_gun：hand_r → ik_hand_gun
    Returns:

    """
    ik_constraint_list = ['IK_Hd_2_HdSk_L', 'IK_Hd_2_HdSk_R',
                          'IK_Ft_2_FtSk_L', 'IK_Ft_2_FtSk_R',
                          'IK_Hd_2_gun']
    ik_constraint_data_list = [['IK_Hd_2_HdSk_L', 'hand_l', 'ik_hand_l'],
                               ['IK_Hd_2_HdSk_R', 'hand_r', 'ik_hand_r'],
                               ['IK_Ft_2_FtSk_L', 'foot_l', 'ik_foot_l'],
                               ['IK_Ft_2_FtSk_R', 'foot_r', 'ik_foot_r'],
                               ['IK_Hd_2_gun', 'hand_r', 'ik_hand_gun']]

    _delete = False
    for ik_constraint in ik_constraint_list:
        if ik_constraint not in [i.Name for i in
                                 FBSystem().Scene.Constraints]:
            _delete = True
            break
    if _delete:
        for ik_constraint in ik_constraint_list:
            componentList = get_object_by_name(ik_constraint,
                                               namespace=namespace)
            for i in componentList:
                if type(i) == FBConstraint:
                    FBComponent.FBDelete(i)

    for constraint_data in ik_constraint_data_list:
        constraint_name, bone, ik_bone = constraint_data
        create_parent_constraint(constraint_name, [bone], ik_bone, namespace)


def selected_clear():
    for comp in FBSystem().Scene.Components:
        comp.Selected = False


def export_selected(take, file_path):
    file_options = FBFbxOptions(False)
    file_options.SaveSelectedModelsOnly = True
    file_options.ShowFileDialog = False
    file_options.ShowOptionsDialog = False
    file_options.UpdateRecentFiles = False

    for index in range(file_options.GetTakeCount()):
        if file_options.GetTakeName(index) == take.Name:
            file_options.SetTakeSelect(index, True)
        else:
            file_options.SetTakeSelect(index, False)

    FBApplication().FileSave(file_path, file_options)


def plot_anim_cr2sk(char):
    options = FBPlotOptions()
    options.ConstantKeyReducerKeepOneKey = False
    options.PlotAllTakes = False
    options.PlotOnFrame = True
    options.PreciseTimeDiscontinuities = False
    options.PlotLockedProperties = False
    options.RotationFilterToApply = \
        FBRotationFilter.kFBRotationFilterUnroll
    options.PlotTranslationOnRootOnly = False
    options.UseConstantKeyReducer = False
    char.PlotAnimation(FBCharacterPlotWhere.
                       kFBCharacterPlotOnSkeleton,
                       options)


def select_skeleton_hierarchy(top_model):
    if FBModel.FbxGetObjectSubType(top_model) != 'FBModel':
        top_model.Selected = True

    for child_model in top_model.Children:
        select_skeleton_hierarchy(child_model)


def get_char_skeleton_root(char):
    ls_characters = []
    ls_char_roots = FBComponentList()
    FBFindObjectsByName('root', ls_char_roots, False, True)

    char_root = char.GetModel(FBBodyNodeId.kFBHipsNodeId).Parent
    if char_root in ls_char_roots:
        char_namespace = ''
        if ":" in char_root.LongName:
            char_namespace = char_root.LongName.rsplit(':', 2)[0]
            if char_namespace:
                char_namespace = char_namespace + ':'

        ls_characters.append([char, char_root, char_namespace])
        skel_root = ls_characters[0][1]

    else:
        skel_root = char_root

    return skel_root


def get_selected():
    selected_list = FBModelList()
    topModel = None  # 搜索所有模型，而不仅仅是一个特定的分支
    selectionState = True  # 返回选择物体
    sortBySelectOrder = True  # 按照选择顺序排序
    FBGetSelectedModels(selected_list,
                        topModel,
                        selectionState,
                        sortBySelectOrder)
    return selected_list