bl_info = {
    "name": "Framsticks POV import & camera manipulator",
    "author": "Szymon Ulatowski, jm soler",
    "version": (1, 0, 1),
    "blender": (2, 80, 0),
    "location": "File > Import-Export",
    "category": "Import-Export",
    "description": "Imports POV-Ray files generated by Framsticks and exports Blender camera to POV-Ray files",
    "wiki_url": "http://www.framsticks.com/3d-animations-in-blender"
}

import os, math, bpy, re, copy, functools
from bpy_extras.io_utils import ImportHelper
from bpy.props import (CollectionProperty, StringProperty, BoolProperty, EnumProperty, FloatProperty)

re_field=re.compile('^#declare field_([^_]+)_([^=]+)=(.*)$')
re_object=re.compile('^BeginObject\(\)$')
re_part=re.compile('^BeginObject\(\)$')
re_partgeo=re.compile('^PartGeometry\(<([^>]+)>,<([^>]+)>\)$')
re_joint=re.compile('^BeginJoint\(([0-9]+),([0-9]+)\)$')
re_jointgeo=re.compile('^JointGeometry\(<([^>]+)>,<([^>]+)>,<([^>]+)>\)$')
re_neuro=re.compile('^BeginNeuro\(([^)])\)$')

Folder = "" 
FirstPOV=-1 
LastPOV=-1 
CREATURES={}

# assumes the filename is "*_number.pov"
def numFromFilename(f):
    try:
        return int(f[f.find('_')+1:f.find('.')])
    except:
        return -1

# creates a global dictionary of filenames ("*_number.pov"), where key=number, updates FirstPOV/LastPOV
def scanDir(startfile):
    global Files, Folder, FileName
    global FirstPOV, LastPOV

    Files={}
    Folder, FileName=os.path.split(startfile)

    #print("startfile=",startfile,"Folder=",Folder,"FileName=",FileName)

    underscore=FileName.find("_")
    if underscore==-1: 
        return
    if FileName.find(".pov")==-1:
        return

    FirstPOV=numFromFilename(FileName)
    LastPOV=FirstPOV

    for f in os.listdir(Folder): 
        if not f.endswith('.pov'):
            continue
        if not f.startswith(FileName[:underscore+1]):
            continue
        num=numFromFilename(f)
        if num<0:
            continue
        Files[num]=f
        if num>LastPOV:
            LastPOV=num
    #print("N=",len(Files))

def extractValue(val): 
    if val[0]=='"':
        # "string"
        return val[1:-1]
    if val.find('.')!=-1 or val.find('e')!=-1:
       # floating point
        return float(val) 
    else:
       # integer
       return int(val)

def floatList(str):
    return [float(x) for x in str.split(',')]

def analysePOV(fname):
    global re_field,re_object,re_part,re_partgeo,re_joint,re_jointgeo,re_neuro
    f=open(fname,'r',encoding='latin-1')
    POVlines=f.readlines()
    f.close()
    tmpfields={}
    objects={}
    for line in POVlines:
        m=re_field.match(line)
        if m:
            value=m.group(3)
            if value.endswith(';'):
                value=value[:-1]
            value=extractValue(value)
            objname=m.group(1) # creature,m,p,j,n
            fieldname=m.group(2)
            if not objname in tmpfields:
                tmpfields[objname]={}
            tmpfields[objname][fieldname]=value
            #print("obj=",m.group(1)," name=",m.group(2)," value=",value)
        m=re_object.match(line)
        if m:
            objkey=tmpfields['Creature']['name']+'_'+tmpfields['Creature']['uid']
            objkey=objkey.replace(' ','_')
            recentobj={'fields':copy.deepcopy(tmpfields['Creature']),'parts':[],'joints':[],'neurons':[]}
            objects[objkey]=recentobj
        m=re_jointgeo.match(line)
        if m:
            joint={'StickLoc1':floatList(m.group(1)),
                   'StickLoc2':floatList(m.group(2)),
                   'StickRot':floatList(m.group(3)),
                   'fields':copy.deepcopy(tmpfields['j'])}
            #print(joint)
            recentobj['joints'].append(joint)
        m=re_partgeo.match(line)
        if m:
            part={'Loc':floatList(m.group(1)),
                   'Rot':floatList(m.group(2)),
                   'fields':copy.deepcopy(tmpfields['p'])}
            #print(joint)
            recentobj['parts'].append(part)
    #print(tmpfields)
    #print(objects)
    #print(json.dumps(objects,indent=4))
    return objects

# vector length
def vecLength(p1,p2): 
    p2[0]=p2[0]-p1[0] 
    p2[1]=p2[1]-p1[1] 
    p2[2]=p2[2]-p1[2] 
    return (p2[0]**2+p2[1]**2+p2[2]**2)**0.5

def vecSub(p1,p2): 
    return [p1[0]-p2[0], p1[1]-p2[1], p1[2]-p2[2]]

def vecMul(p1,m): 
    return [p1[0]*m, p1[1]*m, p1[2]*m]

# create an object containing reference to blender cylinder object
class cylindre: 
    def __init__(self, name='cylinderModel', r1=0.1, r2=0.1, h=1.0, n=8, caps=True):
        self.mesh=bpy.data.meshes.new(name=name)
        r=[r1,r2] 
        verts=[]
        for i in range(0,2): 
            for j in range(0,n): 
                z=math.sin(j*math.pi*2/(n))*r[i] 
                y=math.cos(j*math.pi*2/(n))*r[i] 
                x=float(i)*h 
                verts.append((x,y,z))
                
        vlist=[v for v in range(n)] 
        vlist.append(0)
        faces=[]
        for i in range(n): 
            faces.append((vlist[i],vlist[i+1],vlist[i+1]+n,vlist[i]+n))

        if caps:
            faces.append(tuple(range(n-1,-1,-1)))
            faces.append(tuple(range(n,2*n)))

        self.mesh.from_pydata(verts,[],faces)
        self.mesh.update()

    def makeObject(self, name):
        o = bpy.data.objects.new(name, self.mesh)
        bpy.context.collection.objects.link(o)
        return o

# build or update blender objects from a POV file
def updateBlender(SceneParent,num): 
    global Folder, Files, Current, RecentlyCreated, RecentlyDisappeared, CylinderMesh, LinkedMesh
    Incoming=analysePOV(os.path.join(Folder,Files[num]))
    RecentlyCreated=[]
    RecentlyDisappeared=[]
    for oname,obj in Incoming.items():
        if not oname in Current:
            # add object properties
            print('Creature added:',oname)
            RecentlyCreated.append(oname)
            newobj=[] # will contain: [ parent, joint0, joint1, ... ]
            Current[oname]=newobj
            # create new blender objects
            bcrea=bpy.data.objects.new(oname,None)
            bpy.context.collection.objects.link(bcrea)
            bcrea.parent=SceneParent
            newobj.append(bcrea)
            for j in obj['joints']:
                if not LinkedMesh:
                    CylinderMesh = None
                if not CylinderMesh:
                    CylinderMesh = cylindre(name = 'linked cylinder' if LinkedMesh else oname+'_j', r1=0.1, r2=0.1, h=1.0, n=6)
                o = CylinderMesh.makeObject(oname+'_j')
                o.parent=bcrea
                newobj.append(o)

        # update blender loc/rot/scale
        existing_b=Current[oname]
        if len(obj['joints']):
            avg_loc=vecMul(functools.reduce(lambda a,b: [a[0]+b[0],a[1]+b[1],a[2]+b[2]], [j['StickLoc1'] for j in obj['joints']]),1/len(obj['joints']))
        elif len(obj['parts']):
            avg_loc=vecMul(functools.reduce(lambda a,b: [a[0]+b[0],a[1]+b[1],a[2]+b[2]], [p['Loc'] for p in obj['parts']]),1/len(obj['parts']))
        else:
            avg_loc=[0,0,0]
        if len(existing_b)>0:
            existing_b[0].location=avg_loc
            existing_b[0].keyframe_insert(data_path='location',frame=bpy.context.scene.frame_current)
        for i in range(len(obj['joints'])):
            if i>=(len(existing_b)-1):
                continue # number of joints has changed -> ignore
            incoming_geo=obj['joints'][i]
            #print('incoming:',incoming_geo)
            bo=existing_b[i+1] # blender object
            scale=[vecLength(incoming_geo['StickLoc1'],incoming_geo['StickLoc2']), 1.0, 1.0]
            for xyz in [0,1,2]:
                getattr(bo,'location')[xyz]=vecSub(incoming_geo['StickLoc1'],avg_loc)[xyz]
                #getattr(bo,'location')[xyz]=incoming_geo['StickLoc1'][xyz]
                getattr(bo,'rotation_euler')[xyz]=incoming_geo['StickRot'][xyz]
                getattr(bo,'scale')[xyz]=scale[xyz]
            for field in ['location','rotation_euler','scale']:
                bo.keyframe_insert(data_path=field,frame=bpy.context.scene.frame_current)
    for oname,obj in Current.items():
        if not oname in Incoming:
            RecentlyDisappeared.append(oname)
            print('Creature disappeared:',oname)

def setHideWithKeyframe(o,hide):
    o.hide_render = hide
    o.hide_viewport = hide
    o.keyframe_insert(data_path="hide_viewport",frame=bpy.context.scene.frame_current)
    o.keyframe_insert(data_path="hide_render",frame=bpy.context.scene.frame_current)
            
# import a sequence of POV files, create object hiererchy, animate
def framsImport(startfile): 
       global  FirstPOV, LastPOV, Files
       global  Folder, FileName
       global  Current, FirstFrame, FrameCount
       global  RecentlyCreated, RecentlyDisappeared
       global  SceneParent
       global  SkipFrames
       global  CylinderMesh

       scanDir(startfile)

       if len(Files)<1:
           print("No files found")
           return

       bpy.context.scene.frame_end=max(bpy.context.scene.frame_end,FirstFrame+FrameCount-1)

       CylinderMesh = None
       SceneParent=bpy.data.objects.new("Framsticks_"+str(FirstFrame),None)
       bpy.context.collection.objects.link(SceneParent)
       SceneParent.framspov_file=startfile
       SceneParent.framspov_frame=FirstFrame
       SceneParent.framspov_count=FrameCount

       Current={}
       NextSkip=0
       for k in sorted(Files.keys()):
           if k<NextSkip:
               continue
           NextSkip=k+SkipFrames
           bpy.context.scene.frame_set(FirstFrame-FirstPOV+k)
           if bpy.context.scene.frame_current >= FirstFrame+FrameCount:
               break
           print("Frame %d - loading POV %s" % (bpy.context.scene.frame_current,Files[k]))
           updateBlender(SceneParent,k)
           if len(RecentlyDisappeared)>0 or len(RecentlyCreated)>0:
               bpy.context.scene.frame_set(FirstFrame-FirstPOV+k-1)
               for oname in RecentlyCreated:
                   obj=Current[oname]
                   for bo in obj:
                       setHideWithKeyframe(bo,True)
               for oname in RecentlyDisappeared:
                   obj=Current[oname]
                   for bo in obj:
                       setHideWithKeyframe(bo,False)
               bpy.context.scene.frame_set(FirstFrame-FirstPOV+k)
               for oname in RecentlyCreated:
                   obj=Current[oname]
                   for bo in obj:
                       setHideWithKeyframe(bo,False)
               for oname in RecentlyDisappeared:
                   obj=Current[oname]
                   for bo in obj:
                       setHideWithKeyframe(bo,True)
                   Current.pop(oname)


###############################

def povUpdateFile(filename,cam):
    f=open(filename,'r',encoding='latin-1') 
    lines=f.readlines() 
    f.close() 
    for i in range(len(lines)):
        line=lines[i]
        if line.startswith('Camera('):
            line='Camera(<%g,%g,%g>,<%g,%g,%g>)\n' % tuple(cam);
            lines[i]=line
    f=open(filename,'w',encoding='latin-1') 
    f.writelines(lines)
    f.close()

def framsCameraFromObj(obj):
    #print(obj.location.x)
    m=obj.matrix_local
    return [obj.location.x, obj.location.y, obj.location.z, obj.location.x-m[0][2], obj.location.y-m[1][2], obj.location.z-m[2][2]]

def povUpdateScene(obj,writepath):
    global Folder,FirstFrame,FrameCount,Files,FirstPOV
    print("Updating scene %s" % obj.name)
    scanDir(obj.framspov_file)
    if len(Files)<1:
        print("No files found for "+obj.name)
        return

    if writepath:
        f=open(os.path.join(Folder,'camerapath.inc'),'w',encoding='latin-1')
        f.write("#local CameraPathFirst=1;\n")
        f.write("#local CameraPath=array["+str(len(Files))+"*2]\n")
        f.write("{\n")

    FirstFrame=obj.framspov_frame
    FrameCount=obj.framspov_count
    for k in sorted(Files.keys()):
        #bpy.context.scene.frame_current=FirstFrame-FirstPOV+k
        bpy.context.scene.frame_set(FirstFrame-FirstPOV+k)
        if bpy.context.scene.frame_current >= FirstFrame+FrameCount:
            break
        print("Frame %d - updating camera in %s" % (bpy.context.scene.frame_current,Files[k]))
        cam=framsCameraFromObj(bpy.context.scene.camera)
        cam[0]-=obj.location.x
        cam[1]-=obj.location.y
        cam[2]-=obj.location.z
        cam[3]-=obj.location.x
        cam[4]-=obj.location.y
        cam[5]-=obj.location.z
        if writepath:
            f.write(("  <%g,%g,%g>,<%g,%g,%g>," % tuple(cam))+" //"+Files[k]+"\n")
        else:
            povUpdateFile(os.path.join(Folder,Files[k]),cam)

    if writepath:
        f.write('''
}
#if ((AnimFrame>=CameraPathFirst) & ((AnimFrame-CameraPathFirst)<(dimension_size(CameraPath,1)/2)))
  #local i=2*(AnimFrame-CameraPathFirst);
  Camera(CameraPath[i],CameraPath[i+1])
#end
''')
        f.close()

####################################

class FramsticksPOVImporter(bpy.types.Operator, ImportHelper):
    """Load a collection of Framsticks POV files"""
    bl_idname = "framspov.import"
    bl_label = "Import Framsticks POV"
    bl_options = {'UNDO'}

    files: CollectionProperty(name="File Path",
                          description="File path used for importing",
                          type=bpy.types.OperatorFileListElement)
    directory: StringProperty()

    framspov_skip: bpy.props.IntProperty(name="Frame step",min=1,max=100)

    framspov_linkedmesh: bpy.props.BoolProperty(name="Linked cylinder mesh", description="Use a single shared cylinder mesh for all sticks", default=True)

    filename_ext = ".pov"
    filter_glob: StringProperty(default="*.pov", options={'HIDDEN'})

    def execute(self, context):
        global FirstFrame, FrameCount,FileName,SkipFrames,LinkedMesh
        FirstFrame = bpy.context.scene.frame_current
        FrameCount = 9999
        SkipFrames = self.framspov_skip
        LinkedMesh = self.framspov_linkedmesh
        framsImport(os.path.join(self.directory, self.files[0].name))
        return {'FINISHED'}
    

def menu_func_import(self, context):
    self.layout.operator(FramsticksPOVImporter.bl_idname, text="Framsticks POV (.pov)")

class OBJECT_PT_framspov(bpy.types.Panel):
    bl_label = "Framsticks POV"
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "object"
 
    @classmethod
    def poll(cls, context):
        obj=context.object
        return obj.framspov_file!=''
 
    def draw(self, context):
        layout = self.layout
        row = layout.row()
        row.operator("framspov.updatecam",icon='SCRIPT')
        row = layout.row()
        row.operator("framspov.writecamerapath",icon='SCRIPT')


class VIEW3D_OT_UpdatePOVCamera(bpy.types.Operator):
    bl_idname = "framspov.updatecam"
    bl_label = "Update POV camera"
 
    def execute(self, context):
        povUpdateScene(context.object,False)
        return{'FINISHED'}

class VIEW3D_OT_WritePOVCameraPath(bpy.types.Operator):
    bl_idname = "framspov.writecamerapath"
    bl_label = "Write camerapath.inc"
 
    def execute(self, context):
        povUpdateScene(context.object,True)
        return{'FINISHED'}

def register():
    bpy.types.Object.framspov_file=bpy.props.StringProperty(name="Name of the first POV file")
    bpy.types.Object.framspov_frame=bpy.props.IntProperty(name="First frame",min=0,max=999999)
    bpy.types.Object.framspov_count=bpy.props.IntProperty(name="Number of frames",min=0,max=999999)
    bpy.utils.register_class(FramsticksPOVImporter)
    bpy.utils.register_class(VIEW3D_OT_UpdatePOVCamera)
    bpy.utils.register_class(VIEW3D_OT_WritePOVCameraPath)
    bpy.utils.register_class(OBJECT_PT_framspov)
    bpy.types.TOPBAR_MT_file_import.append(menu_func_import)

def unregister():
    bpy.utils.unregister_class(FramsticksPOVImporter)
    bpy.utils.unregister_class(VIEW3D_OT_UpdatePOVCamera)
    bpy.utils.unregister_class(VIEW3D_OT_WritePOVCameraPath)
    bpy.utils.unregister_class(OBJECT_PT_framspov)
    bpy.types.TOPBAR_MT_file_import.remove(menu_func_import)