import bpy
import bmesh
import os, sys
import re
from mathutils import Matrix, Vector
from math import sin, cos, pi, sqrt


from bpy_lattice import slicer
from bpy_lattice import materials


#bpy.context.scene.render.engine = 'CYCLES'


def ele_material(ele):
    key = ele['key']
    name = key+'_material'
    if name in bpy.data.materials:
        return bpy.data.materials[name]
    return materials.diffuse_material(name, color=ele_color(ele)+tuple([1]))
   

def map_table_dict(line):
    print('map_table_dict: ', line)
    d = {}
    vals = line.split(',')[0:14]
    #d['layer']  = vals[0].strip()
    d['name']   = vals[0].strip()
    d['index']  = int(vals[1])
    d['x']      = float(vals[2])
    d['y']      = float(vals[3])
    d['z']      = float(vals[4])
    d['theta']  = float(vals[5])
    d['phi']    = float(vals[6])
    d['psi']    = float(vals[7])
    d['key']    = vals[8].strip().upper() 
    d['L']      = float(vals[9])  
    if d['key']=='SBEND':
        d['angle'] = float(vals[10]) 
        d['e1'] = float(vals[11]) 
        d['e2'] = float(vals[12])  
    if d['key']=='PIPE':
        d['radius_x'] =  float(vals[10])  
        d['radius_y'] =  float(vals[11])
        d['thickness'] =  float(vals[12])  
    if d['key']=='WIGGLER':
        d['radius_x'] =  float(vals[10])  
        d['radius_y'] =  float(vals[11])
        d['xray_line_len'] =  float(vals[12])          
    d['descrip'] = vals[13]  
    return d

def blendfile(ele):
    match = re.search('3DMODEL=(.+?).blend', ele['descrip'])
    if match:
        return match.group(1)+'.blend'
    else:
        return None

def import_lattice(file):
    f = open(file, 'r')
    header=f.readline()
    lat = [map_table_dict(line) for line in f]
    f.close()
    return lat


ELE_X_SCALE_FACTOR = 50


ELE_X_SCALE = {
    'E_GUN' :0.4,
    'PIPE': 0.01, 
    'DRIFT': 0.01,
    'ECOLLIMATOR': 0.03,
    'RCOLLIMATOR': 0.03,
    'SBEND': 0.02,
    'QUADRUPOLE': 0.03 , 
    'SEXTUPOLE': 0.02,
    'LCAVITY': 0.1,
    'RFCAVITY':    0.5,
    'SOLENOID':    0.1,
    'WIGGLER':  0.04  ,
    'EM_FIELD': 0.1,
    'INSTRUMENT': 0.05   }  


ELE_COLOR = {
    'PIPE':(1,1,1),
    'E_GUN':(0.5, 0.5, 0.5),
    'SBEND':(1,0,0),
    'QUADRUPOLE':(0,0,1),    
    'LCAVITY': (0,1,0), 
    'RFCAVITY': (0,1,0) ,
    'SOLENOID':(1,0,1), #purple  
    'EM_FIELD': (1,0,0.5), #??
    'SEXTUPOLE': (1,1,0),     # yellow   
    'WIGGLER':(1,0.4,0) ,    # orange?  
    'INSTRUMENT':      (0, 0, 0)} # black       
    
    
def ele_x_scale(ele):
    """
    Scale factor for an element
    """
    key = ele['key']
    scale = 1
    if key in ELE_X_SCALE:
        scale =  ELE_X_SCALE[key]
    else:
        print(f'missing {key} in ELE_X_SCALE ')

    return ELE_X_SCALE_FACTOR*scale


def ele_color(ele):
    """
    Color for an element
    """
    color = (0,0,0)
    key = ele['key']
    if key in ELE_COLOR :
        color  = ELE_COLOR[key]

    return color

def ele_section():
    return 0
    
def faces_from(sections, closed=True):
    """
    A section is a list of vertices that defines a cross-section of an element
    This makes rectangle faces
    """
    nix = [len(s) for s in sections]
    if len(set(nix)) >1:
        print('ERROR: sections must have the same number of points')
        return
    n = nix[0]
    faces = []
    for i0 in range(len(sections) -1):
        for j in range(n-1):
            faces.append(  (i0*n + j, (i0+1)*n +j, (i0+1)*n +j+1, i0*n +j+1)   )
        faces.append(      (i0*n +n-1, (i0+1)*n +n-1, (i0+1)*n, i0*n)         )
    if closed:
        #faces.append(range(n)) #first section
        faces.append( list(reversed(range(n)))) #first section
        faces.append(range((len(sections)-1)*n, (len(sections)-1)*n+n)) # Last section
    return faces  
    
def box_section(X, haperture, vaperture):
    return ( (X, haperture, vaperture), (X, -haperture, vaperture), (X, -haperture, -vaperture), (X, haperture, -vaperture))

def ellipse_section(X, haperture, vaperture, n=30):
    angles = [2*pi*i/n for i in range(n)]
    return [ (X, haperture*cos(a), vaperture*sin(a)) for a in angles]

def multipole_section(X, aperture, n):
    angles = [pi*i/n + pi/(2*n) for i in range(2*n)]
    return [ (X, aperture*cos(a), aperture*sin(a)) for a in angles]

def ele_section(s_rel, ele):
    """
    Make sections relatice to center of element
    """
    sc = ele_x_scale(ele)
    if ele['key'] == 'QUADRUPOLE':
        return multipole_section(s_rel, sc, 4) 
    if ele['key'] == 'SEXTUPOLE':
        return multipole_section(s_rel, sc, 6)     
    elif ele['key'] == 'SBEND':
        
        a = ele['angle']
        if abs(a) < 1e-5 or abs(ele['L']) < 1e-5:
            return box_section(s_rel, sc, sc) 
        L = ele['L']
        rho = L/a
        # Baseline section
        s0 = box_section(0, sc, sc)
        # Edge angle
        f = s_rel/L+0.5
        edge =  ele['e2']*f + (-1)*ele['e1']*(1-f)
        m0=Matrix.Rotation(edge, 4, 'Z')
        m1=Matrix.Translation((0, rho, 0))
        m2=Matrix.Rotation(-s_rel/rho, 4, 'Z')
        m3=Matrix.Translation((0, -rho, 0))
        #m=m3*m2*m1*m0 Old syntax
        m=m3 @ m2 @ m1 @ m0
        sec=[]
        for p in s0:
            v = m @ Vector(p)
            sec.append(v[:])
        return sec
        
    elif ele['key'] == 'WIGGLER':
        return box_section(s_rel, sc, 2*sc)    
    elif ele['key'] == 'PIPE':
        rx = ele['radius_x'] 
        ry = ele['radius_y'] 
        t = ele['thickness']
        if rx== 0 or ry==0:
            return ellipse_section(s_rel, sc, sc)
        else:
            return ellipse_section(s_rel, rx+t, ry+t)
    else:
        return ellipse_section(s_rel, sc, sc)

def ele_mesh(ele):
    name = ele['name']
    print('Mesh: ', name)
    L = ele['L']
    if ele['key'] == 'SBEND':
        n=20
        slist = [L*i/(n-1) - L/2 for i in range(n)]
    else:
        slist = [-L/2, L/2]
    sections = [ele_section(s_rel, ele) for s_rel in slist]
    faces = faces_from(sections)
    verts = []
    for s in sections:
        for p in s:
            verts.append(p)
            
    mesh = bpy.data.meshes.new(name)
    mesh.from_pydata(verts, [], faces)
    # Fix normals
    #bpy.ops.object.editmode_toggle()
    #bpy.ops.mesh.normals_make_consistent(inside=False)
    #bpy.ops.object.editmode_toggle()
    #print('mesh update')
    mesh.update(calc_edges=True)
    #print('calc_normals')
    mesh.calc_normals()
    #print('calc_normals DONE')
    return mesh


# ------ Pipe stuff


# basic pipe:
# ele = {'name':'xxx', 'L':0.1, 'key':'PIPE', 'radius_x':0.3, 'radius_y':0.3, 'thickness':0.01}

    
def rotate_mesh(mesh, mat=Matrix.Rotation(pi/2, 4, 'Y')): 
    for v in mesh.vertices:
        v.co = mat*v.co
        
def new_ellipse(radius_x=1, radius_y=1, length=1, vertices=32):
   
    bpy.ops.mesh.primitive_cylinder_add(radius=radius_x, depth = length, vertices=vertices)
    #bpy.ops.transform.rotate(value=pi/2, axis=(0, 1, 0))
    ob = bpy.context.scene.objects.active
    mesh = ob.data
    for v in mesh.vertices:
        v.co[0] *= radius_y/radius_x
        
    #ob.rotation_euler[1] = pi/2 # Align to x-axis 
    
    return ob

def old_new_pipe(name = 'pipe', length=1, radius_x=0.1, radius_y=1, material=None, thickness = 0.01):
    #pipe = new_ellipse(radius_x=radius_x + thickness, radius_y=radius_y+ thickness, length=length, vertices=64)
    #pipe.name = name
    #aperture = new_ellipse(radius_x=radius_x, radius_y=radius_y, length=length, vertices=32)
    #punch_hole(pipe, aperture)
    #bpy.context.scene.objects.active = aperture
    #bpy.ops.object.delete()
    #rotate_mesh(pipe.data)
    if material:
        pipe.data.materials.append(material)
    return pipe

def pipe_object(ele):
    print('pipe_object')
    ele0 = ele.copy()
    ele0['thickness'] = 0
    ele0['L'] = ele['L']*1.1 # the punch needs to be slightly longer to work
    #ele['thickness']  = 0.2
    mesh0 = ele_mesh(ele0)
    mesh = ele_mesh(ele)
    print('pipe mesh vertices: ', len(mesh0.vertices), len(mesh.vertices))
    
    fix_mesh(mesh0)
    fix_mesh(mesh)
    object0 = bpy.data.objects.new('dummy', mesh0)
    object = bpy.data.objects.new(ele['name'], mesh)

    # add to scene to
    #bpy.context.scene.objects.link(object0)
    #bpy.context.scene.objects.link(object)
    bpy.context.collection.objects.link(object0)
    bpy.context.collection.objects.link(object)
    #punch_hole(object,object0)
    
    
    slicer.slice_object(object,object0)
    
    
    #o = bpy.data.objects['xxx']
    #p = bpy.data.objects['dummy']
    #slice.slice(o, p)
    
     
    bpy.context.scene.objects.unlink(object0)
     
    bpy.data.objects.remove(object0)
    bpy.data.meshes.remove(mesh0)
    
    #fix_mesh(mesh)
    #for testing: 
    #bpy.context.scene.objects.link(object0)
    return object

# ------------------------------------------

def load_blend(filepath):
    '''
    Load .blend model objects.
    '''
    with bpy.data.libraries.load(filepath, link=False) as (data_from, data_to):    
        onames = [name for name in data_from.objects if True]
        print('Library: ', filepath)
        print('         has objects', onames)
        data_to.objects = data_from.objects
    return data_to.objects



def add_children_from_blend(parent, blendfilepath, libdict):
    if blendfilepath in libdict:
        print('Library already loaded: ', blendfilepath)
        print('Data will be linked')
        # Library has already been loaded. Copy meshes and materials
        children = []
        for o in libdict[blendfilepath]:
            child = bpy.data.objects.new(o.name, o.data)
            child.location = o.location
            child.rotation_euler = o.rotation_euler
            children.append(child)
    else:
        print('New library: ', blendfilepath)
        children = load_blend(blendfilepath)  
        libdict[blendfilepath] = children
    for child in children:
        #bpy.context.scene.objects.link(child) OLD blender
        bpy.context.collection.objects.link(child)
        if child.parent == None:
            child.parent = parent         



def fix_mesh(mesh):
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bmesh.ops.recalc_face_normals(bm, faces=bm.faces)
    bm.to_mesh(mesh)
    bm.free()

def old_fix_mesh(object):
    bpy.context.scene.objects.active = object
    bpy.ops.object.mode_set(mode='EDIT')

    # select all faces
    bpy.ops.mesh.select_all(action='SELECT')

    # recalculate outside normals 
    bpy.ops.mesh.normals_make_consistent(inside=False)

    # go object mode again
    bpy.ops.object.editmode_toggle()




def ele_object(ele, library, use_real_model=False, catalogue=None, hide_real_model=True):
    name = ele['name']
    print('Object: ', name)
    if ele['key']=='PIPE' and ele['thickness'] > 0 and ele['radius_x'] > 0:
      object = pipe_object(ele)
    else:
      mesh = ele_mesh(ele)
      object = bpy.data.objects.new(name, mesh)
      #bpy.context.scene.objects.link(object) old blender
      bpy.context.collection.objects.link(object)
    
    
      #fix_mesh(object)
    object.location = (0, 0, 0) 
    bfile=blendfile(ele)
    if bfile and use_real_model and catalogue:
        f=os.path.join(catalogue, bfile)
        if os.path.isfile(f):
            print('blend file: ', f, 'exists!')
            add_children_from_blend(object, f, library)
            # Hide options for preview
            if hide_real_model:
                for c in object.children:
                    c.hide_set(True)
            else:
                object.hide_set(True)
            object.hide_render = True
        else:
            print('Blend file missing: ', f)
            
    mat = ele_material(ele)
    #print('color: ', color)
    mat.diffuse_color = ele_color(ele)+tuple([1])
    object.data.materials.append(mat)
    

    return object

    #bpy.context.scene.objects.link(object)
    


def ele_objects(eles,
    library={},
    use_real_model=False,
    catalogue=None,
    hide_real_model=True,
    origin = (0, 0, 0)):
    """
    Create multiple objects from a list of eles (a lattice)
    """

    Xcenter, Ycenter, Zcenter = origin



    objects = []
    for ele in eles:
        if ele['L'] == 0:
            continue
        ob = ele_object(ele, library=library, use_real_model=use_real_model, hide_real_model=hide_real_model, catalogue =catalogue )
    
        # Set location and angles 
        ob.rotation_euler.z =  ele['theta']
        ob.rotation_euler.y = -ele['phi']
        ob.rotation_euler.x =  ele['psi']
        ob.location=( ele['z']-Xcenter, ele['x']-Ycenter, ele['y']-Zcenter)

        
        objects.append(ob)

    return objects

def lat_borders(lat, dim='x'):
  xlist = [ele[dim] for ele in lat]
  return min(xlist), max(xlist)
  
  
def test():
    ele = {'name':'test_pipe', 'radius_x': 0.2, 'radius_y': 0.1, 'thickness':0.01, 'L': 1, 'key':'PIPE', 'descrip':''}
    o = ele_object(ele, None)
  

if __name__ == '__main__':
  test()

