bl_info = {
    "name": "PokeCorrect",
    "description": "Creates an IK rig from existing ripped pokemon models",
    "author": "NeoEmberArts",
    "location": "View 3D > Properties Panel",
    "version": (1, 3),
    "blender": (2, 80, 0),
    "support": "COMMUNITY",
    "category": "Animation",
}

import bpy, os
from bpy.props import CollectionProperty
from bpy_extras.io_utils import ImportHelper
from bpy.types import Operator
import zipfile
import re

page = 1


def parent_bones(armature, child_name, parent_name):
    # Get bone objects
    child_bone = armature.edit_bones.get(child_name)
    parent_bone = armature.edit_bones.get(parent_name)

    # Check if both bones exist
    if child_bone and parent_bone:
        # Clear existing parent relationships
        child_bone.parent = None

        # Set new parent relationship
        child_bone.parent = parent_bone
def create_hierarchy():
    # Check if there is an active object and it's an armature
    if bpy.context.active_object and bpy.context.active_object.type == 'ARMATURE':
        armature = bpy.context.active_object.data       
        ik_complexity = bpy.context.scene.poke_correct_IKComplexity
        if ik_complexity == 'Advanced':
            # Parent FootIK to FootRoll
            parent_bones(armature, 'L.FootIK', 'L.FootRoll')
            parent_bones(armature, 'R.FootIK', 'R.FootRoll')

            # Parent FootRoll to FootControl
            parent_bones(armature, 'L.FootRoll', 'L.FootControl')
            parent_bones(armature, 'R.FootRoll', 'R.FootControl')

            # Parent ToeControl to FootControl
            parent_bones(armature, 'L.ToeControl', 'L.FootControl')
            parent_bones(armature, 'R.ToeControl', 'R.FootControl')

            # Parent FootControl to Origin  ## TODO: NEEDS TO HAVE RENAME FOR ORIGIN ## BUG
            parent_bones(armature, 'L.FootControl', 'Origin')
            parent_bones(armature, 'R.FootControl', 'Origin')
            parent_bones(armature, 'L.FootControl', 'origin')
            parent_bones(armature, 'R.FootControl', 'origin')

            # Parent HandIK to HandRoll
            parent_bones(armature, 'L.HandIK', 'L.HandRoll')
            parent_bones(armature, 'R.HandIK', 'R.HandRoll')

            # Parent HandRoll to HandControl
            parent_bones(armature, 'L.HandRoll', 'L.HandControl')
            parent_bones(armature, 'R.HandRoll', 'R.HandControl')

            # Parent FingerControl to HandControl
            parent_bones(armature, 'L.FingerControl', 'L.HandControl')
            parent_bones(armature, 'R.FingerControl', 'R.HandControl')

            # Parent HandControl to Origin
            parent_bones(armature, 'L.HandControl', 'Origin')
            parent_bones(armature, 'R.HandControl', 'Origin')
            parent_bones(armature, 'L.HandControl', 'origin')
            parent_bones(armature, 'R.HandControl', 'origin')

            # Ensure spine1 and hip are parented to the waist, and the waist to the origin bone
            parent_bones(armature, 'spine1', 'waist')
            parent_bones(armature, 'hip', 'waist')
            parent_bones(armature, 'waist', 'Origin')
            parent_bones(armature, 'waist', 'origin')
        else:
            parent_bones(armature, 'L.FootIK', 'Origin')
            parent_bones(armature, 'R.FootIK', 'Origin')
            parent_bones(armature, 'L.FootIK', 'origin')
            parent_bones(armature, 'R.FootIK', 'origin')
            # Parent HandControl to Origin
            parent_bones(armature, 'L.HandIK', 'Origin')
            parent_bones(armature, 'R.HandIK', 'Origin')
            parent_bones(armature, 'L.HandIK', 'origin')
            parent_bones(armature, 'R.HandIK', 'origin')
            # Ensure spine1 and hip are parented to the waist, and the waist to the origin bone
            parent_bones(armature, 'spine1', 'waist')
            parent_bones(armature, 'hip', 'waist')
            parent_bones(armature, 'waist', 'Origin')
            parent_bones(armature, 'waist', 'origin')

def rename_bones(armature_data, naming_option, separator_option):
    separator = separator_option  # "." or "_"
    for b in armature_data.bones:
        # Adjust regex pattern based on the naming_option
        if naming_option == "Lx/Rx" and re.match(r'^[LR](?=[A-Z])', b.name):
            new_name = re.sub(r'^([LR])(?=[A-Z])', fr'\g<1>{separator}', b.name)
            setattr(b, 'name', new_name)
            print(f'Renamed: {b.name} to {new_name}')
        elif naming_option == "Leftx/Rightx" and re.match(r'^(Left|Right)(?=[A-Z])', b.name):
            new_name = re.sub(r'^(Left|Right)(?=[A-Z])', fr'\g<1>{separator}', b.name)
            setattr(b, 'name', new_name)
            print(f'Renamed: {b.name} to {new_name}')
        elif (naming_option == "Both"):
            if re.match(r'^[LR](?=[A-Z])', b.name) or re.match(r'^(Left|Right)(?=[A-Z])', b.name):
                new_name = re.sub(r'^(Left|Right|L|R)(?=[A-Z])', fr'\g<1>{separator}', b.name)
                setattr(b, 'name', new_name)
                print(f'Renamed: {b.name} to {new_name}')
    
def remove_zero_length_bones(obj):
    if obj and obj.type == 'ARMATURE':
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.armature.select_all(action='DESELECT')
        [setattr(b.bone, 'select', True) for b in obj.data.bones if b.length == 0]
        bpy.ops.armature.delete()
        bpy.ops.object.mode_set(mode='OBJECT')
        print("Zero-length bones removed.")
    else:
        print("Please select an armature.")

def set_bones_roll_to_zero(armature_obj):
    bpy.context.view_layer.objects.active = armature_obj
    bpy.ops.object.mode_set(mode='EDIT')
    # Set the roll of all bones to zero
    for bone in armature_obj.data.edit_bones:
        bone.roll = 0

    bpy.ops.object.mode_set(mode='OBJECT')

def set_user_transforms_as_rest_pose(armature_obj):
    # Set user transforms as the default rest pose
    bpy.ops.object.mode_set(mode='POSE')
    bpy.ops.pose.armature_apply(selected=False)
    bpy.ops.object.mode_set(mode='OBJECT')

    # Apply transforms to all bones
    bpy.ops.object.mode_set(mode='POSE')
    bpy.ops.pose.select_all(action='SELECT')
    bpy.ops.pose.transforms_clear()
    bpy.ops.object.mode_set(mode='OBJECT')

def changepage(pagenum):
    if globals()["page"] == pagenum or globals()["page"] == 3:
        if globals()["page"] == 3:
            globals()["page"] = pagenum = 4
        else:
            globals()["page"] = pagenum+1
    else:
        globals()["page"] = pagenum

def hide_elements(context):
    context.scene.poke_correct_visible = False
    #changepage(1)



def show_elements(context):
    context.scene.poke_correct_visible = True
    #changepage(2)

def scale_bones():
    # Check if in Object mode
    if bpy.context.mode != 'EDIT_ARMATURE':
        bpy.ops.object.mode_set(mode='EDIT')

    armature = bpy.context.object.data

    # Loop through selected edit bones
    for bone in armature.edit_bones:
        if bone.select:
            # Calculate the center of the bone
            center = (bone.head + bone.tail) / 2.0

            # Scale the bone around its individual origin
            scale_factor = 0.5
            #bone.head = center + scale_factor * (bone.head - center)
            bone.tail = center + scale_factor * (bone.tail - center)


def align_bones_upwards():
    # Check if in Object mode
    if bpy.context.mode != 'EDIT_ARMATURE':
        bpy.ops.object.mode_set(mode='EDIT')

    armature = bpy.context.object.data
    # Loop through selected edit bones
    for bone in armature.edit_bones:
        if bone.select:
            # Calculate the distance from head to tail
            distance = (bone.head - bone.tail).length

            # Move the head to the tail, adjusting along the Z-axis
            bone.tail = bone.head
            bone.tail.z += distance

def snap_bones_to_bottom_mesh():
    # Check if in Object mode
    if bpy.context.mode != 'EDIT_ARMATURE':
        bpy.ops.object.mode_set(mode='EDIT')

    armature = bpy.context.object.data

    # Loop through selected edit bones
    for bone in armature.edit_bones:
        if bone.select:
            # Set the head position to the bottom of the foot
            bone.head.z = 0

            # Move the tail position to the bottom of the foot
            bone.tail.z = 0

def select_digit_bones(type):
    if bpy.context.mode != 'EDIT_ARMATURE':
        bpy.ops.object.mode_set(mode='EDIT')

    armature = bpy.context.object.data

    # Deselect all bones initially
    bpy.ops.armature.select_all(action='DESELECT')

    if (type == "toe"):
        # Define regex patterns
        patterns = [
            #re.compile(r'^[LR]\.|^[LR]_'),
            re.compile(r'ToeB1$'),
            re.compile(r'ToeB$'),
            re.compile(r'Toe$'),
        ]
        for bone in armature.edit_bones:
            # Check if the bone name matches the specified regex patterns
            if any(pattern.search(bone.name) for pattern in patterns):
                # Select the bone
                bone.select = True
                print(bone.name)
    else:
        # Define regex patterns
        patterns = [
            #re.compile(r'^[LR]\.|^[LR]_'),
            re.compile(r'FingerB1$'),
            re.compile(r'FingerB$'),
            re.compile(r'Finger$'),
        ]
        inc = 0
        for bone in armature.edit_bones:
            # Check if the bone name matches the specified regex patterns
            if any(pattern.search(bone.name) for pattern in patterns):
                # Select the bone
                bone.select = True
                inc = inc + 1
                print(bone.name)
        if inc == 0:
            # Define regex patterns
            patterns = [
                #re.compile(r'^[LR]\.|^[LR]_'),
                re.compile(r'HandB1$'),
                re.compile(r'HandB$'),
                re.compile(r'Hand$'),
            ]
            inc = 0
            for bone in armature.edit_bones:
                # Check if the bone name matches the specified regex patterns
                if any(pattern.search(bone.name) for pattern in patterns):
                    # Select the bone
                    bone.select = True
                    inc = inc + 1
                    print(bone.name)


def SelectIK():
    # Check if there is an active object and it's an armature
    if bpy.context.active_object and bpy.context.active_object.type == 'ARMATURE':
        armature = bpy.context.active_object.data
        # List of possible bone names
        bone_patterns = [
            re.compile(r'^[LR]\.FootIK$'),
            re.compile(r'^[LR]\.FootControl$'),
            re.compile(r'^[LR]\.FootRoll$'),
            re.compile(r'^[LR]\.ToeControl$'),
            re.compile(r'^[LR]\.HandIK$'),
            re.compile(r'^[LR]\.HandControl$'),
            re.compile(r'^[LR]\.HandRoll$'),
            re.compile(r'^[LR]\.FingerControl$'),
        ]

        # Switch to Edit Mode
        bpy.ops.object.mode_set(mode='EDIT')

        # Deselect all bones
        bpy.ops.armature.select_all(action='DESELECT')

        # Iterate through bone patterns and select matching bones
        for pattern in bone_patterns:
            for bone in armature.edit_bones:
                if pattern.match(bone.name):
                    bone.select = True
        else:
            print("No matching bones found.")




def select_foot_bones_in_edit_mode():
    #-------------------Reaname bypass
    bpy.ops.object.editmode_toggle()
    bpy.ops.object.editmode_toggle()
    armature = bpy.context.object.data
    incorrect = ["left_hand", "right_hand", "left_arm_02", "right_arm_02", "left_foot_index", "right_foot_index", "left_foot", "right_foot", "left_leg_02", "right_leg_02", "left_leg_01", "right_leg_01"]
    bpy.ops.armature.select_all(action='DESELECT')
    for aaaa in incorrect:
        if aaaa in armature.edit_bones:
            bone = armature.edit_bones[aaaa]
            bone.select = True
    bpy.ops.object.editmode_toggle()
    bpy.ops.object.editmode_toggle()
    armature = bpy.context.object.data
    for bone in armature.edit_bones:
        if bone.select:
            base_name = bone.name
            new_name = base_name.replace("left_hand", "L.Hand").replace("right_hand", "R.Hand").replace("left_arm_02", "L.ForeArm").replace("right_arm_02", "R.ForeArm").replace("left_foot_index", "L.Toe").replace("right_foot_index", "R.Toe").replace("left_foot", "L.Foot").replace("right_foot", "R.Foot").replace("left_leg_02", "L.Leg").replace("right_leg_02", "R.Leg").replace("right_leg_01", "R.UpperLeg").replace("left_leg_01", "L.UpperLeg")
            bone.name = new_name
    bpy.ops.armature.select_all(action='DESELECT')
    # Check if there is an active object and it's an armature
    if bpy.context.active_object and bpy.context.active_object.type == 'ARMATURE':
        bpy.ops.object.editmode_toggle()
        bpy.ops.object.editmode_toggle()
        armature = bpy.context.active_object.data
        # List of possible bone names
        bone_names = ["L.Foot", "R.Foot", "L_Foot", "R_Foot"]

        # Switch to Edit Mode
        bpy.ops.object.mode_set(mode='EDIT')

        # Deselect all bones
        
        bpy.ops.armature.select_all(action='DESELECT')
        bonecount = 0
        # Iterate through bone names and select the first one found
        
        for bone_name in bone_names:
            if bone_name in armature.edit_bones:
                bonecount = bonecount + 1
                bone = armature.edit_bones[bone_name]
                bone.select = True
                bpy.context.active_object.data.bones.active = bpy.context.active_object.data.bones[bone_name]
        else:
            print("No matching bones found.")
        if bonecount == 2:
            ik_complexity = bpy.context.scene.poke_correct_IKComplexity
            changepage(3)
            #--------------------Foot IK
            bpy.ops.object.editmode_toggle()
            bpy.ops.object.editmode_toggle()
            bpy.ops.armature.duplicate()
            armature = bpy.context.object.data
            for bone in armature.edit_bones:
                if bone.select:
                    # Get the bone name and remove any trailing numbers
                    base_name = bone.name.rstrip("0123456789.")
                    new_name = base_name + "IK"
                    bone.name = new_name
            scale_bones()
            #align_bones_upwards()
            #--------------------Toe control
            if ik_complexity == 'Advanced':
                select_digit_bones("toe")
                bpy.ops.object.editmode_toggle()
                bpy.ops.object.editmode_toggle()
                bpy.ops.armature.duplicate()
                armature = bpy.context.object.data
                for bone in armature.edit_bones:
                    if bone.select:
                        # Get the bone name and remove any trailing numbers
                        base_name = bone.name.rstrip("0123456789.")
                        new_name = base_name.replace("ToeB", "Toe") + "Control"
                        bone.name = new_name
                scale_bones()
            #-------------------Foot Roll
            if ik_complexity == 'Advanced':
                bpy.ops.object.editmode_toggle()
                bpy.ops.object.editmode_toggle()
                bpy.ops.armature.duplicate()
                armature = bpy.context.object.data
                for bone in armature.edit_bones:
                    if bone.select:
                        # Get the bone name and remove any trailing numbers
                        base_name = bone.name.rstrip("0123456789.")
                        new_name = base_name.replace("Toe", "Foot").replace("Control", "Roll").replace("FootB", "Foot")
                        bone.name = new_name
                scale_bones()
                #align_bones_upwards()
                #------------------Foot Control
                bpy.ops.armature.select_all(action='DESELECT')
                select_digit_bones("toe")
                bpy.ops.object.editmode_toggle()
                bpy.ops.object.editmode_toggle()
                bpy.ops.armature.duplicate()
                armature = bpy.context.object.data
                for bone in armature.edit_bones:
                    if bone.select:
                        # Get the bone name and remove any trailing numbers
                        base_name = bone.name.rstrip("0123456789.")
                        new_name = base_name.replace("Toe", "FootControl").replace("B", "")
                        bone.name = new_name
                #scale_bones()
                snap_bones_to_bottom_mesh()
            ##########################################################################HANDS
            armature = bpy.context.object.data
            fingers = ["L.Hand", "R.Hand", "L_Hand", "R_Hand"]
            bpy.ops.armature.select_all(action='DESELECT')
            # Iterate through bone names and select the first one found
            for finger in fingers:
                if finger in armature.edit_bones:
                    bone = armature.edit_bones[finger]
                    bone.select = True
            #--------------------Hand IK
            bpy.ops.object.editmode_toggle()
            bpy.ops.object.editmode_toggle()
            bpy.ops.armature.duplicate()
            armature = bpy.context.object.data
            for bone in armature.edit_bones:
                if bone.select:
                    # Get the bone name and remove any trailing numbers
                    base_name = bone.name.rstrip("0123456789.")
                    new_name = base_name + "IK"
                    bone.name = new_name
            scale_bones()
            align_bones_upwards()
            #--------------------Finger control
            if ik_complexity == 'Advanced':
                #####if 4 legs
                select_digit_bones("finger")
                ##if 2 legs
                bpy.ops.object.editmode_toggle()
                bpy.ops.object.editmode_toggle()
                bpy.ops.armature.duplicate()
                armature = bpy.context.object.data
                for bone in armature.edit_bones:
                    if bone.select:
                        # Get the bone name and remove any trailing numbers
                        base_name = bone.name.rstrip("0123456789.")
                        new_name = base_name.replace("FingerB", "Finger") + "Control"
                        bone.name = new_name
                scale_bones()
            #-------------------Hand Roll
            if ik_complexity == 'Advanced':
                bpy.ops.object.editmode_toggle()
                bpy.ops.object.editmode_toggle()
                bpy.ops.armature.duplicate()
                armature = bpy.context.object.data
                for bone in armature.edit_bones:
                    if bone.select:
                        # Get the bone name and remove any trailing numbers
                        base_name = bone.name.rstrip("0123456789.")
                        new_name = base_name.replace("Finger", "Hand").replace("Control", "Roll").replace("HandB", "Hand").replace("hand", "Hand")
                        bone.name = new_name
                scale_bones()
            #align_bones_upwards()
            #------------------Hand Control
            select_digit_bones("finger")
            bpy.ops.object.editmode_toggle()
            bpy.ops.object.editmode_toggle()
        
            bpy.ops.armature.duplicate()
            armature = bpy.context.object.data
            for bone in armature.edit_bones:
                if bone.select:
                    # Get the bone name and remove any trailing numbers
                    base_name = bone.name.rstrip("0123456789.")
                    new_name = base_name.replace("Finger", "HandControl").replace("B", "")
                    bone.name = new_name
            scale_bones()
            #align_bones_upwards()
            #--------------------ORIGIN OPTIONS
            armature = bpy.context.object.data
            origin = ["origin", "Origin"]
            bpy.ops.armature.select_all(action='DESELECT')
            # Iterate through bone names and select the first one found
            for bon in origin:
                if bon in armature.edit_bones:
                    bone = armature.edit_bones[bon]
                    bone.select = True
            ######
            bpy.ops.object.editmode_toggle()
            bpy.ops.object.editmode_toggle()
            armature = bpy.context.object.data
            for bone in armature.edit_bones:
                if bone.select:
                    # Get the bone name and remove any trailing numbers
                    base_name = bone.name.rstrip("0123456789.")
                    new_name = base_name.replace("origin", "Origin")
                    bone.name = new_name
            scale_bones()
            parent_bones(armature, "hips", "waist")
            parent_bones(armature, "Hips", "waist")
            parent_bones(armature, "spine.01", "waist")
            parent_bones(armature, "spine_01", "waist")
            parent_bones(armature, "waist", "Origin")
            #bpy.ops.transform.resize(value=(0.5, 0.5, 0.5), orient_type='GLOBAL', orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), orient_matrix_type='GLOBAL', mirror=False, snap=False, snap_elements={'INCREMENT'}, use_snap_project=False, snap_target='CLOSEST', use_snap_self=True, use_snap_edit=True, use_snap_nonedit=True, use_snap_selectable=False, alt_navigation=True)
            #origin scale by .5
            #split waist
            #bpy.ops.wm.call_menu(name="TEXT_MT_view")
            #parent to origin disconeccted
            #bpy.ops.armature.parent_set(type='OFFSET')
            #parent to listed connected
            #bpy.ops.armature.parent_set(type='CONNECTED')
            #hips to waist
            #spine.01 to waist

            #.....
            failurebypass()
            SelectIK()
            bpy.ops.object.editmode_toggle()
            bpy.ops.object.editmode_toggle()
            selected_armature = bpy.context.active_object
            set_bones_roll_to_zero(selected_armature)
            bpy.ops.transform.transform(mode='BONE_SIZE', value=(2, 2, 2, 0))
            selected_armature = bpy.context.active_object
            set_user_transforms_as_rest_pose(selected_armature)
        else:
            changepage(2)
    else:
        print("Please select an armature in Object Mode.")




class OBJECT_OT_PrintStatus(bpy.types.Operator):
    bl_idname = "pokecorrect.print_status"
    bl_label = "Execute"

    def execute(self, context):
        # Check if an armature is selected
        if context.active_object and context.active_object.type == 'ARMATURE':
            armature_data = context.active_object.data
            checkboxes = [context.scene.poke_correct_checkbox1, context.scene.poke_correct_checkbox2, context.scene.poke_correct_checkbox3]
            naming_option = context.scene.poke_correct_naming_option
            separator_option = context.scene.poke_correct_separator_option

            if checkboxes[0]:
                rename_bones(armature_data, naming_option, separator_option)
            if checkboxes[1]:
                remove_zero_length_bones(context.active_object)
            if checkboxes[2]:
                selected_armature = bpy.context.active_object
                set_user_transforms_as_rest_pose(selected_armature)
                selected_armature = bpy.context.active_object
                set_bones_roll_to_zero(selected_armature)

            for i, checkbox in enumerate(checkboxes, start=1):
                print(f"Checkbox {i}: {checkbox}")

        else:
            # Display a warning if an armature is not selected
            self.report({'WARNING'}, "No armature selected. Please select an armature.")

        return {'FINISHED'}


class OBJECT_OT_Setup(bpy.types.Operator):
    bl_idname = "pokecorrect.setup"
    bl_label = "Setup"

    def execute(self, context):
        bpy.context.object.data.display_type = 'BBONE'
        bpy.ops.object.mode_set(mode='EDIT')
        hide_elements(context)
        context.scene.poke_correct_show_setup_elements = True
        changepage(2)
        select_foot_bones_in_edit_mode()
        SelectIK()
        bpy.ops.object.editmode_toggle()
        bpy.ops.object.editmode_toggle()
        for area in bpy.context.screen.areas:
            if area.type == 'VIEW_3D':
                ctx = bpy.context.copy()
                ctx['area'] = area
                ctx['region'] = area.regions[-1]
                try:
                    bpy.ops.view3d.view_selected(ctx) 
                except:
                    print("unable to focus view")
        return {'FINISHED'}




def failurebypass():
    armature = bpy.context.object
    bones_to_delete = []
    for bone in armature.data.edit_bones:
        if (bone.name.endswith((".001", ".002")) or bone.name.startswith("End")):
            armature.data.edit_bones.remove(bone)
            


class OBJECT_OT_Cancel(bpy.types.Operator):
    bl_idname = "pokecorrect.cancel"
    bl_label = "Cancel"

    def execute(self, context):
        show_elements(context)
        context.scene.poke_correct_show_setup_elements = False
        changepage(1)
        return {'FINISHED'}


# needs to add extra bbon scale for all other bones # BUG 
bbonescale = 2
def update(self, context):
        #bpy.ops.object.mode_set(mode='EDIT')
        SelectIK()
        bpy.ops.object.editmode_toggle()
        bpy.ops.object.editmode_toggle()
        bbs = globals()["bbonescale"]
        bpy.ops.transform.transform(mode='BONE_SIZE', value=(1/bbs, 1/bbs, 1/bbs, 0))
        newbbs = context.scene.poke_correct_bbonesize
        bpy.ops.transform.transform(mode='BONE_SIZE', value=(newbbs,newbbs,newbbs, 0))
        globals()["bbonescale"] = newbbs
        print(context.scene.poke_correct_bbonesize)
        bpy.ops.object.mode_set(mode='POSE')
edited = 2
def fix(self, context):
        #bpy.ops.object.mode_set(mode='EDIT')
        SelectIK()
        bpy.ops.object.editmode_toggle()
        bpy.ops.object.editmode_toggle()
        bpy.ops.armature.select_all(action='INVERT')
        bbs = globals()["edited"]
        bpy.ops.transform.transform(mode='BONE_SIZE', value=(1/bbs, 1/bbs, 1/bbs, 0))
        newbbs = context.scene.poke_correct_bbonecorrect
        bpy.ops.transform.transform(mode='BONE_SIZE', value=(newbbs,newbbs,newbbs, 0))
        globals()["edited"] = newbbs
        print(context.scene.poke_correct_bbonecorrect)
        bpy.ops.object.mode_set(mode='POSE')

def CopyRotation(leg1, leg2):
    # Ensure you are in Pose Mode
    bpy.ops.object.mode_set(mode='POSE')
    bpy.ops.pose.select_all(action='DESELECT')
    
    # Select R.FootIK and R.Leg
    bone1 = [bone.name for bone in bpy.context.object.pose.bones if re.search(leg1, bone.name)]
    bone2 = [bone.name for bone in bpy.context.object.pose.bones if re.search(leg2, bone.name)]

    if bone1 and bone2:
        # Set active bones
        bpy.context.view_layer.objects.active.data.bones.active = bpy.context.view_layer.objects.active.data.bones.get(bone2[0])
        bpy.ops.pose.constraint_add(type='COPY_ROTATION')
        copy_rotation_constraint = bpy.context.object.pose.bones[bone2[0]].constraints["Copy Rotation"]
        copy_rotation_constraint.target = bpy.context.object
        copy_rotation_constraint.subtarget = bone1[0]
        # Add IK constraint
        
        # Set chain count for the last selected bone
    else:
        print("Bones not found: ")

def IK(leg1, leg2, chain):
    # Ensure you are in Pose Mode
    bpy.ops.object.mode_set(mode='POSE')
    bpy.ops.pose.select_all(action='DESELECT')
    
    # Select R.FootIK and R.Leg
    bone1 = [bone.name for bone in bpy.context.object.pose.bones if re.search(leg1, bone.name)]
    bone2 = [bone.name for bone in bpy.context.object.pose.bones if re.search(leg2, bone.name)]

    if bone1 and bone2:
        # Set active bones
        bpy.context.view_layer.objects.active.data.bones.active = bpy.context.view_layer.objects.active.data.bones.get(bone1[0])
        bpy.context.view_layer.objects.active.data.bones.active = bpy.context.view_layer.objects.active.data.bones.get(bone2[0])

        # Add IK constraint
        bpy.ops.pose.ik_add(with_targets=True)
        
        # Set chain count for the last selected bone
        bpy.context.object.pose.bones[bone2[0]].constraints["IK"].chain_count = chain
    else:
        print("Bones not found: ")

class OBJECT_OT_fixgeometry(bpy.types.Operator):
    bl_idname = "pokecorrect.fixgeometry"
    bl_label = "Fix Geometry"
    def execute(seld, content):
        #Select Mesh that the selected armature is parented to
        # Check if there is an active object
        if bpy.context.active_object is not None:
            # Check if the active object is a mesh
            if bpy.context.active_object.type == 'MESH':
                # If it's a mesh, select it and go to edit mode
                bpy.context.active_object.select_set(True)
                bpy.ops.object.mode_set(mode='EDIT')
            else:
                # If it's not a mesh, deselect it
                bpy.context.active_object.select_set(False)

        # Check if there is a selected armature
        if bpy.context.active_object is not None and bpy.context.active_object.type == 'ARMATURE':
            # Find the parent mesh of the selected armature
            parent_mesh = None
            for obj in bpy.context.scene.objects:
                if obj.type == 'MESH' and obj.parent == bpy.context.active_object:
                    parent_mesh = obj
                    break

            if parent_mesh is not None:
                # Select the parent mesh
                parent_mesh.select_set(True)

                # Go into edit mode
                bpy.context.view_layer.objects.active = parent_mesh
                bpy.ops.object.mode_set(mode='EDIT')
        else:
            # If there is no selected armature, just go into edit mode
            bpy.ops.object.mo
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.mesh.remove_doubles()
        return {'FINISHED'}
    
#bpy.ops.mesh.remove_doubles(threshold=1.1001)

class OBJECT_OT_zip(bpy.types.Operator):
    bl_idname = "pokecorrect.zip"
    bl_label = "Zip"
    def execute(self, context):
        bpy.ops.import_test.some_data('INVOKE_DEFAULT')      
        return {'FINISHED'}  
        

class ImportSomeData(Operator, ImportHelper):
    bl_idname = "import_test.some_data"
    bl_label = ".DAE or .ZIP"
    
    files = CollectionProperty(type=bpy.types.PropertyGroup)
    # Define the filter for .dae files
    filter_glob: bpy.props.StringProperty(default='*.dae;*.zip', options={'HIDDEN'})

    def execute(self, context):
        if self.filepath.endswith(".zip"):
            # Get the directory of the file path
            directory = os.path.dirname(self.filepath)
            # Get the base name of the file path
            base_name = os.path.basename(self.filepath)
            # Create a new directory name by removing the '.zip' extension
            new_dir_name = base_name[:-4]
            # Create the final path for the new directory
            finalpath = os.path.join(directory, new_dir_name)
            # Create the new directory
            os.makedirs(finalpath, exist_ok=True)
            # Open the ZIP file
            datazip = zipfile.ZipFile(self.filepath)
            # Extract the ZIP file directly into the finalpath folder 
            datazip.extractall(path=finalpath)
            # Close the ZIP file
            datazip.close()
            # Remove the original ZIP file
            os.remove(self.filepath)
            
            # Reopen the file browser with the new folder selected and filter for .dae files
            bpy.ops.import_test.some_data('INVOKE_DEFAULT')
        elif self.filepath.endswith(".dae"):
            bpy.ops.wm.collada_import(filepath=self.filepath, fix_orientation=True, find_chains=True, auto_connect=True)
        return{'FINISHED'}


class OBJECT_OT_fixmat(bpy.types.Operator):
    bl_idname = "pokecorrect.fixmaterials"
    bl_label = "Fix Geometry"
    def execute(seld, content):
        print("Fixing Materials")
        # Check if there is an active object
        if bpy.context.active_object is not None:
            # Check if the active object is a mesh
            if bpy.context.active_object.type == 'MESH':
                # If it's a mesh, select it and go to edit mode
                bpy.context.active_object.select_set(True)

                # Loop through all materials of the selected object
                for material_slot in bpy.context.active_object.material_slots:
                    material = material_slot.material

                    # Set the active material index to the current material
                    #bpy.context.object.active_material = material

                    # Access the Principled BSDF node and set the value of the Base Color input to 0
                    principled_bsdf = material.node_tree.nodes.get("Principled BSDF")
                    if principled_bsdf:
                        principled_bsdf.inputs[1].default_value = 0

                # Go into edit mode
                bpy.ops.object.mode_set(mode='EDIT')
            else:
                # If it's not a mesh, deselect it
                bpy.context.active_object.select_set(False)
        else:
            print("No active object selected.")
        return {'FINISHED'}

class OBJECT_OT_NextStep(bpy.types.Operator):
    bl_idname = "pokecorrect.next_step"
    bl_label = "Next Step"

    def execute(self, context):
        changepage(2)
        if (page ==4):
            ik_complexity = bpy.context.scene.poke_correct_IKComplexity
            if ik_complexity == 'Advanced':
                create_hierarchy()
                IK(r'^R.*FootIK$', r'^R.*Leg$', 2)
                IK(r'^L.*FootIK$', r'^L.*Leg$', 2)
                IK(r'^R.*HandIK$', r'^R.*ForeArm$', 2)
                IK(r'^L.*HandIK$', r'^L.*ForeArm$', 2)
                CopyRotation(r'^L.*FootIK$', r'^L.*Foot$')
                CopyRotation(r'^R.*FootIK$', r'^R.*Foot$')
                CopyRotation(r'^L.*HandIK$', r'^L.*Hand$')
                CopyRotation(r'^R.*HandIK$', r'^R.*Hand$')
                CopyRotation(r'^R.*ToeControl$', r'^R.*Toe$')
                CopyRotation(r'^R.*ToeControl$', r'^R.*ToeB$')
                CopyRotation(r'^L.*ToeControl$', r'^L.*Toe$')
                CopyRotation(r'^L.*ToeControl$', r'^L.*ToeB$')
                CopyRotation(r'^L.*FingerControl$', r'^L.*FingerB$')
                CopyRotation(r'^L.*FingerControl$', r'^L.*FingerB1$')
                CopyRotation(r'^L.*FingerControl$', r'^L.*Finger$')
                CopyRotation(r'^R.*FingerControl$', r'^R.*FingerB$')
                CopyRotation(r'^R.*FingerControl$', r'^R.*FingerB1$')
                CopyRotation(r'^R.*FingerControl$', r'^R.*Finger$')
                selected_armature = bpy.context.active_object
                set_user_transforms_as_rest_pose(selected_armature)
            elif ik_complexity == 'Simple':
                create_hierarchy()
                IK(r'^R.*FootIK$', r'^R.*Leg$', 2)
                IK(r'^L.*FootIK$', r'^L.*Leg$', 2)
                IK(r'^R.*HandIK$', r'^R.*ForeArm$', 2)
                IK(r'^L.*HandIK$', r'^L.*ForeArm$', 2)
                CopyRotation(r'^L.*FootIK$', r'^L.*Foot$')
                CopyRotation(r'^R.*FootIK$', r'^R.*Foot$')
                CopyRotation(r'^L.*HandIK$', r'^L.*Hand$')
                CopyRotation(r'^R.*HandIK$', r'^R.*Hand$')
                selected_armature = bpy.context.active_object
                set_user_transforms_as_rest_pose(selected_armature)
            else:
                print("Error A")

        return {'FINISHED'}



class PokeCorrectPanel(bpy.types.Panel):

    bl_label = "PokeCorrect"
    bl_idname = "PT_PokeCorrectPanel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'Armature'

    def draw(self, context):
        layout = self.layout

        # Check if the elements should be shown
        if page == 1:  #context.scene.poke_correct_visible:
            # Checkbox 1
            layout.operator(OBJECT_OT_zip.bl_idname, text="Import model [zip/dae]", icon='FILE_FOLDER')
            # Title
            layout.separator()
            box = layout.box()
            box.label(text="Correction", icon='ALIGN_LEFT')  # Change the icon as needed
            # Checkbox 1
            layout.prop(context.scene, "poke_correct_checkbox1", text="Correct Naming Convention")

            if context.scene.poke_correct_checkbox1:
                # First Dropdown
                layout.prop(context.scene, "poke_correct_naming_option", text="Naming Option", expand=True,
                            icon='BLENDER')

                # Second Dropdown
                layout.prop(context.scene, "poke_correct_separator_option", text="Separator", expand=True, icon='BLENDER')
            # Checkbox 2
            layout.prop(context.scene, "poke_correct_checkbox2", text="Remove Zero Length Bones")
            # Checkbox 3
            layout.prop(context.scene, "poke_correct_checkbox3", text="Set Rest Pose and Bone Roll")
            # Separator line
            # Button
            layout.operator(OBJECT_OT_PrintStatus.bl_idname, text="Execute")
            # Create Rig Section
            layout.separator()
            box = layout.box()
            box.label(text="Miscelanous fixing", icon='ARMATURE_DATA')  # Change the icon as neede
            layout.operator(OBJECT_OT_fixgeometry.bl_idname, text="Fix Geometry")
            layout.operator(OBJECT_OT_fixmat.bl_idname, text="Fix Materials")
            # Dropdown for legs
            
            box = layout.box()
            #.label(text="", icon='ARMATURE_DATA')
            layout.prop(context.scene, "poke_correct_IKComplexity", text="Complexity", expand=True, icon='BLENDER')
            box.label(text="Create Rig", icon='ARMATURE_DATA')
            #layout.prop(context.scene, "poke_correct_legs_option", text="Legs", expand=True, icon='BLENDER')
            # Setup button
            layout.operator(OBJECT_OT_Setup.bl_idname, text="Begin Setup")
        elif page==2:
            # Dropdown for legs (Hidden)
            layout.operator(OBJECT_OT_Cancel.bl_idname, text="Cancel")
            layout.label(text="Error 4 - unable to proccess armature", icon='ARMATURE_DATA')
            # Begin Setup button (Hidden)
            ##layout.operator(OBJECT_OT_NextStep.bl_idname, text="Next Step")
        elif page==3:
            #layout.operator(OBJECT_OT_Cancel.bl_idname, text="Cancel")
            #####################################################
            box = layout.box()
            box.label(text="Constraints", icon='ARMATURE_DATA')  # Change the icon as needed
            #.label(text="", icon='ARMATURE_DATA')
            #layout.prop(context.scene, "poke_correct_IKComplexity", text="Complexity", expand=True, icon='BLENDER')
            #box = layout.box()
            #box.label(text="Extra settings", icon='ARMATURE_DATA')
            #layout.prop(context.scene, "poke_correct_auto_origin", text="Auomatic Origin Bone")
            #layout.prop(context.scene, "poke_correct_auto_hips", text="Automatic Hip Bone")
            layout.operator(OBJECT_OT_NextStep.bl_idname, text="Continue")
        elif page==4:
            box = layout.box()
            box.label(text="Display Settings", icon='ARMATURE_DATA')
            #layout.prop(context.scene, "poke_correct_auto_hips", text="Auto Bone Layers")
            layout.prop(context.scene, "poke_correct_bbonecorrect", text="All BBone Scale")
            layout.prop(context.scene, "poke_correct_bbonesize", text="BBone Scale")
            layout.operator(OBJECT_OT_Cancel.bl_idname, text="Done")



def register():
    bpy.utils.register_class(ImportSomeData)
    bpy.utils.register_class(OBJECT_OT_PrintStatus)
    bpy.utils.register_class(OBJECT_OT_Setup)
    bpy.utils.register_class(OBJECT_OT_Cancel)
    bpy.utils.register_class(OBJECT_OT_NextStep)
    bpy.utils.register_class(OBJECT_OT_fixgeometry)
    bpy.utils.register_class(OBJECT_OT_zip)
    bpy.utils.register_class(OBJECT_OT_fixmat)
    bpy.utils.register_class(PokeCorrectPanel)
    #del bpy.types.Scene.poke_correct_IKComplexity
    #del bpy.types.Scene.poke_correct_BBone1Size
    #del bpy.types.Scene.poke_correct_BBone2Size
    #del bpy.types.Scene.poke_correct_BBone3Size
    bpy.types.Scene.poke_correct_auto_origin = bpy.props.BoolProperty(name="origin", default=False)
    bpy.types.Scene.poke_correct_auto_hips = bpy.props.BoolProperty(name="hips", default=False)
    bpy.types.Scene.poke_correct_IKComplexity = bpy.props.EnumProperty(
        name="IKComplexity",
        items=[
            ("Simple", "Simple", ""),
            ("Advanced", "Advanced", ""),
        ],
        default="Simple"
    )
    #ahhhhhhhhhhhhh
    bpy.types.Scene.poke_correct_visible = bpy.props.BoolProperty(name="Visible", default=True)
    bpy.types.Scene.poke_correct_show_setup_elements = bpy.props.BoolProperty(name="Show Setup Elements", default=False)
    bpy.types.Scene.poke_correct_checkbox1 = bpy.props.BoolProperty(name="Correct Naming Convention", default=True)
    bpy.types.Scene.poke_correct_naming_option = bpy.props.EnumProperty(
        name="Naming Option",
        items=[
            ("Lx/Rx", "Lx/Rx", ""),
            ("Leftx/Rightx", "Leftx/Rightx", ""),
            ("Both", "Both", ""),
        ],
        default="Lx/Rx"
    )
    bpy.types.Scene.poke_correct_separator_option = bpy.props.EnumProperty(
        name="Separator",
        items=[
            (".", ".", ""),
            ("_", "_", ""),
        ],
        default="."
    )
    bpy.types.Scene.poke_correct_bbonesize = bpy.props.FloatProperty(
        name = "BBone Scale Offset",
        default = 2,
        min = 0.01,
        max = 300,
        update = update,
    )
    bpy.types.Scene.poke_correct_bbonecorrect = bpy.props.FloatProperty(
        name = "BBone Scale Offset",
        default = 2,
        min = 0.01,
        max = 300,
        update = fix,
    )
    bpy.types.Scene.poke_correct_checkbox2 = bpy.props.BoolProperty(name="Remove Zero Length Bones", default=True)
    bpy.types.Scene.poke_correct_checkbox3 = bpy.props.BoolProperty(name="Checkbox 3", default=True)
    bpy.types.Scene.poke_correct_legs_option = bpy.props.EnumProperty(
        name="Legs Option",
        items=[
            ("2", "2 Legs", ""),
            ("4", "4 Legs", ""),
        ],
        default="2"
    )
    


def unregister():
    bpy.utils.register_class(ImportSomeData)
    bpy.utils.unregister_class(OBJECT_OT_PrintStatus)
    bpy.utils.unregister_class(OBJECT_OT_Setup)
    bpy.utils.unregister_class(OBJECT_OT_Cancel)
    bpy.utils.unregister_class(OBJECT_OT_NextStep)
    bpy.utils.unregister_class(OBJECT_OT_fixgeometry)
    bpy.utils.unregister_class(OBJECT_OT_zip)
    bpy.utils.unregister_class(PokeCorrectPanel)
    bpy.utils.unregister_class(OBJECT_OT_fixmat)

    del bpy.types.Scene.poke_correct_visible
    del bpy.types.Scene.poke_correct_show_setup_elements
    del bpy.types.Scene.poke_correct_checkbox1
    del bpy.types.Scene.poke_correct_naming_option
    del bpy.types.Scene.poke_correct_separator_option
    del bpy.types.Scene.poke_correct_checkbox2
    del bpy.types.Scene.poke_correct_checkbox3
    del bpy.types.Scene.poke_correct_legs_option
    #setup
    del bpy.types.Scene.poke_correct_IKComplexity
    del bpy.types.Scene.poke_correct_BBone1Size
    del bpy.types.Scene.poke_correct_BBone2Size
    del bpy.types.Scene.poke_correct_BBone3Size
    del bpy.types.Scene.poke_correct_auto_origin
    del bpy.types.Scene.poke_correct_auto_hips
    del bpy.types.Scene.poke_correct_bbonesize
    del bpy.types.Scene.poke_correct_bbonecorrect


if __name__ == "__main__":
    register()