import bpy
import mathutils
from mathutils import Matrix, Vector
from math import pi
import numpy as np
from time import time

def batch_reproject(obj, proj_type='VIEW', all_strokes=True, restore_frame=False):
    '''Reproject - ops method
    :all_stroke: affect hided, locked layers
    '''

    if restore_frame:
        oframe = bpy.context.scene.frame_current
    omode = bpy.context.mode

    # FIXME : if all_stroke is False, might be better to still store>set>restore "lock_frame"
    if all_strokes:
        layers_state = [[l, l.hide, l.lock, l.lock_frame] for l in obj.data.layers]
        for l in obj.data.layers:
            l.hide = False
            l.lock = False
            l.lock_frame = False

    bpy.ops.object.mode_set(mode='EDIT_GPENCIL')
    # TODO : need to toggle visible and unlock all layers before attacking, and store restore.

    for l in obj.data.layers:
        for f in l.frames:
            bpy.context.scene.frame_set(f.frame_number)
            # switch to edit to reproject through ops
            bpy.ops.gpencil.select_all(action='SELECT')
            bpy.ops.gpencil.reproject(type=proj_type) # default is VIEW
            bpy.ops.gpencil.select_all(action='DESELECT')
    
    # restore
    if all_strokes:
        for layer, hide, lock, lock_frame in layers_state:
            layer.hide = hide
            layer.lock = lock
            layer.lock_frame = lock_frame

    bpy.ops.object.mode_set(mode=omode)
    
    if restore_frame:
        bpy.context.scene.frame_current = oframe


def align_global(reproject=True, ref=None, all_strokes=True):

    if not ref:
        ref = bpy.context.scene.camera

    o = bpy.context.object
    # if o.matrix_basis != o.matrix_world and not o.parent:

    ref = bpy.context.scene.camera
    ref_mat = ref.matrix_world
    ref_loc, ref_rot, ref_scale = ref_mat.decompose()

    if o.parent:
        mat = o.matrix_world
    else:
        mat = o.matrix_basis

    o_loc, o_rot, o_scale = mat.decompose()

    mat_90 = Matrix.Rotation(-pi/2, 4, 'X')

    loc_mat = Matrix.Translation(o_loc)
    rot_mat = ref_rot.to_matrix().to_4x4() @ mat_90
    scale_mat = get_scale_matrix(o_scale)

    new_mat = loc_mat @ rot_mat @ scale_mat

    # world_coords = []
    for l in o.data.layers:
        for f in l.frames:
            for s in f.strokes:
                ## foreach
                coords = [p.co @ mat.inverted() @ new_mat for p in s.points]
                # print('coords: ', coords)
                # print([co for v in coords for co in v])
                s.points.foreach_set('co', [co for v in coords for co in v])
                # s.points.update() # seem to works # but adding/deleting a point is "safer"
                ## force update
                s.points.add(1)
                s.points.pop()

                # for p in s.points:
                    ## GOOD :
                    # world_co = mat @ p.co
                    # p.co = new_mat.inverted() @ world_co

                    ## GOOD :
                    # p.co = p.co @ mat.inverted() @ new_mat

    if o.parent:
        o.matrix_world = new_mat
    else:
        o.matrix_basis = new_mat

    if reproject:
        batch_reproject(o, proj_type='FRONT', all_strokes=all_strokes)


def align_all_frames(reproject=True, ref=None, all_strokes=True):
    
    print('aligning all frames...')

    o = bpy.context.object
    if not ref:
        ref = bpy.context.scene.camera

    # get all rot
    chanel = 'rotation_quaternion' if o.rotation_mode == 'QUATERNION' else 'rotation_euler'

    ## double list keys
    rot_keys = [k.co.x for fcu in o.animation_data.action.fcurves for k in fcu.keyframe_points if fcu.data_path == chanel]

    ## normal iter
    # for fcu in o.animation_data.action.fcurves:
    #     if fcu.data_path != chanel :
    #         continue
    #     for k in fcu.keyframe_points():
    #         rot_keys.append(k.co.x)

    rot_keys = list(set(rot_keys))

    # TODO # TOTHINK
    # for now the rotation of the object is adjusted at every frame....
    # might be better to check camera rotation of the current frame only (stored as copy).
    # else the object rotate following the cameraview vector (not constant)...

    mat_90 = Matrix.Rotation(-pi/2, 4, 'X')

    for l in o.data.layers:
        for f in l.frames:
            # set the frame to dedicated
            bpy.context.scene.frame_set(f.frame_number)

            ref_mat = ref.matrix_world
            ref_loc, ref_rot, ref_scale = ref_mat.decompose()

            if o.parent:
                mat = o.matrix_world
            else:
                mat = o.matrix_basis

            o_loc, o_rot, o_scale = mat.decompose()
            loc_mat = Matrix.Translation(o_loc)
            rot_mat = ref_rot.to_matrix().to_4x4() @ mat_90
            scale_mat = get_scale_matrix(o_scale)
            new_mat = loc_mat @ rot_mat @ scale_mat

            for s in f.strokes:
                ## foreach
                coords = [p.co @ mat.inverted() @ new_mat for p in s.points]
                # print('coords: ', coords)
                # print([co for v in coords for co in v])
                s.points.foreach_set('co', [co for v in coords for co in v])
                # s.points.update() # seem to works
                ## force update
                s.points.add(1)
                s.points.pop()

    for fnum in rot_keys:
        bpy.context.scene.frame_set(fnum)
        #/update calculation block
        ref_mat = ref.matrix_world
        ref_loc, ref_rot, ref_scale = ref_mat.decompose()

        if o.parent:
            mat = o.matrix_world
        else:
            mat = o.matrix_basis

        o_loc, o_rot, o_scale = mat.decompose()
        loc_mat = Matrix.Translation(o_loc)
        rot_mat = ref_rot.to_matrix().to_4x4() @ mat_90
        scale_mat = get_scale_matrix(o_scale)
        new_mat = loc_mat @ rot_mat @ scale_mat
        # update calculation block/

        if o.parent:
            o.matrix_world = new_mat
        else:
            o.matrix_basis = new_mat

        o.keyframe_insert(chanel, index=-1, frame=bpy.context.scene.frame_current, options={'INSERTKEY_AVAILABLE'})


    if reproject:
        batch_reproject(o, proj_type='FRONT', all_strokes=all_strokes)

    return


class GPTB_OT_realign(bpy.types.Operator):
    bl_idname = "gp.realign"
    bl_label = "Realign GP"
    bl_description = "Realign the grease pencil front axis with active camera"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'GPENCIL'

    reproject : bpy.props.BoolProperty(
        name='Reproject', default=True,
        description='Reproject stroke on the new alignment')
    
    all_strokes : bpy.props.BoolProperty(
        name='All Strokes', default=False,
        description='Hided and locked layer will also be reprojected')

    ## add option to bake strokes if rotation anim is not constant ? might generate too many Keyframes

    def invoke(self, context, event):
        if context.object.data.use_multiedit:
            self.report({'ERROR'}, 'Does not work in Multi-edit')
            return {"CANCELLED"}
        
        self.alert = ''
        o = context.object
        if o.animation_data and o.animation_data.action:
            act = o.animation_data.action
            for chan in ('rotation_euler', 'rotation_quaternion'):
                if act.fcurves.find(chan):
                    self.alert = 'Animated Rotation (CONSTANT interpolation)'
                    interpos = [p for fcu in act.fcurves if fcu.data_path == chan for p in fcu.keyframe_points if p.interpolation != 'CONSTANT']
                    if interpos:
                        self.alert = f'Animated Rotation ! ({len(interpos)} key not constant)'
                    break

        return context.window_manager.invoke_props_dialog(self, width=450)

    def draw(self, context):
        layout = self.layout
        if self.alert:
            layout.label(text=self.alert, icon='INFO')
            layout.label(text='(rotations key will be overwritten to face camera)')

        layout.prop(self, "reproject")
        layout.label(text='Realign the GP front axis with active camera')
        if self.reproject:
            if not context.region_data.view_perspective == 'CAMERA':
                layout.label(text='Not in camera ! (reprojection is made from view)', icon='ERROR')
            layout.prop(self, "all_strokes")

    def execute(self, context):
        t0 = time()
        oframe = context.scene.frame_current

        o = bpy.context.object
        if o.animation_data and o.animation_data.action:
            if o.animation_data.action.fcurves.find('rotation_euler') or o.animation_data.action.fcurves.find('rotation_quaternion'):
                align_all_frames(reproject=self.reproject)
                
                bpy.context.scene.frame_current = oframe
                print(f'\nAnim realign ({time()-t0:.2f}s)')
                return

        align_global(reproject=self.reproject)
        
        context.scene.frame_current = oframe

        print(f'\nGlobal Realign ({time()-t0:.2f}s)')

        return {"FINISHED"}


class GPTB_OT_batch_reproject_all_frames(bpy.types.Operator):
    bl_idname = "gp.batch_reproject_all_frames"
    bl_label = "Reproject All Frames"
    bl_description = "Reproject all frames of active object."
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'GPENCIL'

    # reproject : bpy.props.BoolProperty(name='Reproject', default=True)
    
    all_strokes : bpy.props.BoolProperty(
        name='All Strokes', default=False,
        description='Hided and locked layer will also be reprojected')
 
    type: bpy.props.EnumProperty(name='Type',
        items=(('FRONT', "Front", ""),
               ('SIDE', "Side", ""),
               ('TOP', "Top", ""),
               ('VIEW', "View", ""),
               ('SURFACE', "Surface", ""),
               ('CURSOR', "Cursor", ""),
               ),
        default='FRONT')

    def invoke(self, context, event):
        if context.object.data.use_multiedit:
            self.report({'ERROR'}, 'Does not work in Multi-edit')
            return {"CANCELLED"}
        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        layout = self.layout
        if not context.region_data.view_perspective == 'CAMERA':
            layout.label(text='Not in camera ! (reprojection is made from view)', icon='ERROR')
        layout.prop(self, "all_strokes")
        layout.prop(self, "type")
        

    def execute(self, context):
        t0 = time()

        batch_reproject(context.object, proj_type=self.type, all_strokes=self.all_strokes, restore_frame=True)

        self.report({'INFO'}, f'Reprojected in ({time()-t0:.2f}s)' )

        return {"FINISHED"}


classes = (
# GPTB_OT_realign,
GPTB_OT_batch_reproject_all_frames,
)

def register():
    for cl in classes:
        bpy.utils.register_class(cl)

def unregister():
    for cl in reversed(classes):
        bpy.utils.unregister_class(cl)