import numpy as np
import os
import cadquery as cq
from cadquery import exporters as et
from scipy.spatial.distance import pdist, squareform
from items import item_list
from collections import defaultdict 


nut_width = 5.5
nut_thickness = 2.5
r_hole = 3.1/2
nut_depth = 8
m_length = 15


def save_stls(objs, base_path, prefix = '', tolerance = 0.1):
    try:
        os.mkdir(base_path)
    except:
        pass

    num = 0 

    for itx in objs:            
        if not isinstance(itx, item_list):
            itc = item_list(itx)
        else:
            itc = itx;

        for fkn, ob in itc.as_dict().items():
            num += 1
            fpath = os.path.join(base_path, prefix + '_' + fkn + '_'+   str(num) + '.stl')
            with open(fpath,'w') as wpx:
                et.exportShape( ob, et.ExportTypes.STL, wpx, tolerance=tolerance)




def n_gon(r, n_sides, xoff = 0, yoff = 0):
    angles = np.linspace(0, 2*np.pi, n_sides+1)
    pts = list()
    for ang in angles[:-1]:
        pts.append(( (r*np.cos(ang))  +xoff, (r*np.sin(ang)) + yoff))


    return pts



def bottom_points(w, d, offw = 50, offd = 50, grid = None):
    if grid is None:
        grid = [2,2]
    pts = list()

    tl_corners = offw, offd
    br_corner = w-offw, d-offd

    for xoff in np.linspace(tl_corners[0],br_corner[0],grid[0]):
        for yoff in np.linspace(tl_corners[1],br_corner[1],grid[1]):
            pts.append( (xoff,yoff))
    
    
    return pts


def copy(obj):
    return obj.translate((0, 0, 0))

def intersect(wp1, wp2):
    """
    Return geometric intersection between 2 cadquery.Workplane instances by
    exploiting.
    A n B = (A u B) - ((A - B) u (B - A))
    """
    neg1 = copy(wp1).cut(wp2)
    neg2 = copy(wp2).cut(wp1)
    negative = neg1.union(neg2)
    return copy(wp1).union(wp2).cut(negative)




def add_rect(wplane, w, d, offx = 0, offy =0):
    return wplane.center(offx, offy).lineTo(0,d).lineTo(w,d).lineTo(w,0).close()




def do_fingers(r_side, f_side,  num_tabs = None, swap = False, skip_list = None, axis_step = None, sq_off_mult = 0, sq_dim = None, tab_width = None, cut = True, skip_bolt_1 = False, skip_bolt_2 = False, do_bolt = False, nut_depth_mult = 1, nut_path_mult = 1):
    num_tabs = 9 if num_tabs is None else num_tabs
    sliver = intersect(r_side, f_side)
    pts = list()
    for h in  sliver.vertices().vals():
        pts.append( ( h.X, h.Y, h.Z))
    pt=np.asarray(pts)
    min_pts = np.min(pt, axis=0)    
    max_pts = np.max(pt, axis=0)
    


    r_pts = np.asarray([[a.X, a.Y, a.Z] for a in r_side.vertices().vals()])
    f_pts = np.asarray([[a.X, a.Y, a.Z] for a in f_side.vertices().vals()])


    r_center = np.mean(r_pts, axis=0)
    f_center = np.mean(f_pts, axis=0)

    # print(r_center)
    
    

    if axis_step is None:
        axis_step = np.argmax(max_pts - min_pts)


    axis_sq = list(set([0,1,2]).difference([axis_step]))
    axis_not_step = list(set([0,1,2]).difference([axis_step]))


    if sq_dim is None:
        sq_dim = (max_pts-min_pts)[axis_sq]

    
    total_sz = (max_pts-min_pts)[axis_step]
    
    if tab_width is not None:
        num_tabs = int(np.ceil(total_sz/tab_width))
    


    step_sz = total_sz / num_tabs
    box_sz = max_pts-min_pts
    box_sz[axis_step] /= num_tabs


    if skip_list is None:
        skip_list = [-1]

    tbj = list()
    tbh = list()
    holes = list()
    screw_path = list()
    nut_path = list()
    for i in range(num_tabs):

        xstep = [0,0,0]
        if (axis_sq[1]-axis_sq[0]) > 1:
            xstep[axis_sq[1]] = 1
            flip = True
        else:
            xstep[axis_sq[0]] = 1
            flip = False
        normal = [0,0,0]
        normal[axis_step] = 1
        origin = min_pts.copy()
        origin[axis_step] += i * step_sz
        t_plane = cq.Plane(tuple(origin), tuple(xstep), tuple(normal))

        origin_holes = origin.copy()
        origin_holes[axis_step] += step_sz/2 
        
        


        normal_1 = [0,0,0]
        normal_1[axis_not_step[0]] = 1

        xstep_1 = [0,0,0]
        xstep_1[axis_not_step[1]] = 1

        normal_2 = [0,0,0]  
        normal_2[axis_not_step[1]] = 1

        xstep_2 = [0,0,0]
        xstep_2[axis_not_step[0]] = 1

        if swap:
            xstep_1, xstep_2 = xstep_2, xstep_1
            normal_1, normal_2 = normal_2, normal_1

        t_plane_holes_1 = cq.Plane(tuple(origin_holes), tuple(xstep_1), tuple(normal_1))
        plane_1_normal = [t_plane_holes_1.zDir.x,t_plane_holes_1.zDir.y,t_plane_holes_1.zDir.z]

        
        t_plane_holes_2 = cq.Plane(tuple(origin_holes), tuple(xstep_2), tuple(normal_2))
        plane_2_normal = [t_plane_holes_2.zDir.x,t_plane_holes_2.zDir.y,t_plane_holes_2.zDir.z]

        # if swap:
        #     plane_1_normal, plane_2_normal = plane_2_normal, plane_1_normal

        if flip:
            tbj.append( add_rect(cq.Workplane(t_plane), sq_dim[1], sq_dim[0]).extrude(step_sz) )
        else:
            tbj.append( add_rect(cq.Workplane(t_plane), sq_dim[0], sq_dim[1]).extrude(step_sz) )

    
        if flip:
            mult = -1
        else:
            mult = 1;
  
        if do_bolt:
            offset_nut_1 = [0,0,0]
            offset_nut_2 = [0,0,0]

            offset_nut_1[axis_step] =  -mult*nut_width/2
            offset_nut_2[axis_step] = mult*nut_width/2

    
            offset_nut_1[axis_not_step[0]] = nut_depth_mult* nut_depth
            offset_nut_2[axis_not_step[1]] = nut_depth_mult*nut_depth 

            offset_nut_1 = tuple(offset_nut_1)
            offset_nut_2 = tuple(offset_nut_2)
            
            if swap:
                offset_nut_1, offset_nut_2 = offset_nut_2, offset_nut_1

            def get_closer_dir(origin, dird, vec):
                
                origin = np.asarray(origin)
                dird = np.asarray(dird)
                vec = np.asarray(vec)
                # print(origin)
                # print(dird)
                # print(vec)
                if np.linalg.norm(origin+dird - vec) > np.linalg.norm(origin-dird + vec):
                    return 1
                else:
                    return -1


            h_app = None
            s_app = None
            n_app = None
            if i % 2 == 1 and not skip_bolt_1 :
                h_app = cq.Workplane(t_plane_holes_1).pushPoints([[sq_dim[0]/2 + sq_off_mult * r_hole/2,0]]).circle(r_hole).extrude(25)
                s_app = add_rect(cq.Workplane(t_plane_holes_1),10, r_hole*2, 0, -r_hole).extrude(nut_path_mult* m_length) 
                n_app = add_rect(cq.Workplane(t_plane_holes_1),  100, nut_width, 0, 0).extrude(nut_thickness ) .translate(offset_nut_1)
            elif i % 2 == 0 and not skip_bolt_2:        
                h_app = cq.Workplane(t_plane_holes_2).pushPoints([[sq_dim[1]/2 + sq_off_mult* r_hole/2,0]]).circle(r_hole).extrude(25)
                s_app = add_rect(cq.Workplane(t_plane_holes_2), 10, 2*r_hole, 0, -r_hole).extrude(nut_path_mult* m_length ) 
                n_app = add_rect(cq.Workplane(t_plane_holes_2), 100, nut_width, 0, 0).extrude(nut_thickness ) .translate(offset_nut_2)
            holes.append(h_app)
            screw_path.append(s_app)
            nut_path.append(n_app)

            

    if cut:
        
        
        for i in range(num_tabs):
            if i in skip_list:
                continue

            if i % 2 == 0:
                r_side = r_side.cut(tbj[i])

            if i % 2 == 1:
                f_side = f_side.cut(tbj[i])

        if do_bolt:
            for p in range(num_tabs):
                if holes[p] is not None:
                    if p % 2 == 1:
                        r_side = r_side.cut(holes[p])
                        f_side = f_side.cut(screw_path[p])
                        f_side = f_side.cut( nut_path[p]) 
                    if p % 2 == 0:
                        f_side = f_side.cut(holes[p])
                        r_side = r_side.cut(screw_path[p])
                        r_side = r_side.cut(nut_path[p])



    return r_side, f_side, [tbj, holes, screw_path, nut_path]


def norm_vector(vec):
    return vec / np.linalg.norm(vec)

def vertices_to_array(verts):
    pts = list()
    for h in verts:
        pts.append( (h.X, h.Y, h.Z))
    
    return np.asarray(pts)

def intersect(wp1,wp2):
    neg1 = wp1.cut(wp2)
    neg2 = wp2.cut(wp1)
    negative = neg1.union(neg2)
    intersected = wp1.union(wp2).cut(negative)
    return intersected

def pround(vals, decimals= 5):
    return np.around(vals, decimals = decimals)

def vector_to_array(vector):
    return np.asarray([vector.x, vector.y, vector.z])

def workplane_along_vector(v_vec):
    most_sig_vec = norm_vector(v_vec)
    orth_vector = np.cross(most_sig_vec, np.random.randn(3))
    c_work_vec = cq.Workplane(cq.Plane(origin=(0,0,0), xDir=tuple(most_sig_vec), normal = tuple(orth_vector)))
    return c_work_vec

def make_fingers(wp1, wp2, num_tabs = 3, r_hole = 3.1/2, slot_depth = 15, nut_width = 5.5, nut_thickness = 2.5, nut_depth = 8, skip_nut_slots_A = False, skip_nut_slots_B = False, swap_ab = False):

    sliver = intersect(wp1,wp2)
    out = vertices_to_array(sliver.vertices().vals())

    dist_mat = squareform(pdist(out))
    distances = defaultdict(lambda: dict())

    for x1 in range(len(out)):
        for x2 in range(len(out)):
            distances[np.around(dist_mat[x1,x2], decimals=5)][(x1,x2)] = out[x2] - out[x1]


    
    long_vec = None
    long_idces = None
    long_arr = None
    for dist, ou in distances.items():
        ds = np.asarray([y for x,y in ou.items()])
        avg_vec = pround(np.mean(np.abs(ds), axis=0))
        cross_products = np.cross(ds, avg_vec)
        total_mag = np.linalg.norm(np.abs(cross_products))
        if dist > 0:
            if pround(total_mag) == 0:
                long_vec = avg_vec
                long_idces = [x for x in ou]
                long_arr = ds

    vec_dir = norm_vector(long_vec)

    wp1_com = cq.Shape.centerOfMass(wp1.objects[0])
    wp2_com = cq.Shape.centerOfMass(wp2.objects[0])

    selector = cq.DirectionMinMaxSelector(cq.Vector(tuple(vec_dir)))
    wp1_top_face_vert = vertices_to_array(wp1.faces(selector).vertices().vals())
    wp2_top_face_vert = vertices_to_array(wp2.faces(selector).vertices().vals())


    orientations = np.dot(long_arr, long_vec)
    face_group_1 = list()
    face_group_2 = list()
    for idc, val in zip(long_idces, orientations):
        if val < 0:
            face_group_1.append(idc[0])
            face_group_2.append(idc[1])


    face_group_1_verts = out[face_group_1]
    face_group_2_verts = out[face_group_2]

    face_group_1_mean = np.average(face_group_1_verts, axis=0);
    face_group_2_mean = np.average(face_group_2_verts, axis=0);


    wp1_top_mean = np.average(wp1_top_face_vert, axis=0)
    wp2_top_mean = np.average(wp2_top_face_vert, axis=0)

    face_to_wp1_top = wp1_top_mean - face_group_1_mean 
    face_to_wp2_top = wp2_top_mean - face_group_1_mean 


    dir_to_wp1 = norm_vector( face_to_wp1_top - (np.dot(vec_dir, face_to_wp1_top))*vec_dir )
    dir_to_wp2 = norm_vector( face_to_wp2_top - (np.dot(vec_dir, face_to_wp2_top))*vec_dir )

    step_d = (face_group_2_mean - face_group_1_mean)/(num_tabs)
    origins = list()
    for x in range(num_tabs+1):
        origins.append(np.average(face_group_1_verts, axis=0) + step_d*x)
    
    objs = list()
    step_mag = float(np.linalg.norm(step_d))
    csel = cq.DirectionMinMaxSelector(cq.Vector(tuple(step_d)))
    cut_objs = list()
    for i_st in range(num_tabs):
        offset = -step_mag
        new_obj = sliver.faces(selector=csel).workplane(offset).split(keepTop=True, keepBottom=True)
        new_obj_0 = sliver.newObject([new_obj.objects[0]])
        new_obj_1 = sliver.newObject([new_obj.objects[1]])
        sliver = new_obj_1
        cut_objs.append(new_obj_0)
        


    csel_wp1 = cq.DirectionMinMaxSelector(cq.Vector(tuple(dir_to_wp1)), directionMax=False)
    csel_wp2 = cq.DirectionMinMaxSelector(cq.Vector(tuple(dir_to_wp2)), directionMax=False)

    csel_wp1_max = cq.DirectionMinMaxSelector(cq.Vector(tuple(dir_to_wp1)), directionMax=True)
    csel_wp2_max = cq.DirectionMinMaxSelector(cq.Vector(tuple(dir_to_wp2)), directionMax=True)

    def cut_finger(wp1, cobj):
        wp1 = wp1.cut(cobj)
        return wp1

    def cut_screw_slots(wp1, wp2, csel_wp1, csel_wp1_max, cobj, step_d):
        os1 = cobj.faces(csel_wp1).workplane()
        os1 = os1.cut(cobj)

        if np.dot(vector_to_array(os1.plane.zDir), step_d) == 0:
            do_swap = True 
        else:
            do_swap = False
        
        os_s = os1.pushPoints([[0,0]]).circle(r_hole).extrude(-10)
        wp2 = wp2.cut(os_s)
        

        
        pa = nut_width
        pb = 25
        if do_swap:
            pa, pb = pb, pa
            
        os1 = cut_objs[idx].faces(csel_wp1_max).workplane()
        nut_slot = add_rect(os1.workplane(nut_depth),pa,pb,offx=-pa/2, offy=-pb/2).extrude(nut_thickness)
        wp1 = wp1.cut(nut_slot)

        na = 2*r_hole
        nb = 25
        if do_swap:
            na, nb = nb, na

        os1a = cut_objs[idx].faces(csel_wp1_max).workplane()
        screw_slot = add_rect(os1a,na,nb,offx=-na/2, offy=-nb/2).extrude(slot_depth)
        wp1 = wp1.cut(screw_slot)    
        
        return wp1, wp2, os_s, screw_slot, nut_slot, [os1, os1a]


    if swap_ab:
        offset = 1
    else:
        offset = 0
    for idx in range(len(cut_objs)):
        if (idx+offset) % 2:
            wp1 = cut_finger(wp1, cut_objs[idx])
            if not skip_nut_slots_A:
                wp1, wp2, hole, screw_slot, nut_slot, os1_planes = cut_screw_slots(wp1, wp2, csel_wp1, csel_wp1_max, cut_objs[idx], step_d)
            
        else: 
            wp2 = cut_finger(wp2, cut_objs[idx])    
            if not skip_nut_slots_B:
                wp2, wp1, hole, screw_slot, nut_slot, os2_planes = cut_screw_slots(wp2, wp1, csel_wp2, csel_wp2_max, cut_objs[idx], step_d)

    return wp1, wp2
# def do_fingers(r_side, f_side, num_tabs = 9, tab_width = None, cut = True):
    # sliver = intersect(r_side, f_side)
    # pts = list()
    # for h in  sliver.vertices().vals():
    #     pts.append( ( h.X, h.Y, h.Z))
    # pt=np.asarray(pts)
    # min_pts = np.min(pt, axis=0)    
    # max_pts = np.max(pt, axis=0)
    # axis_step = np.argmax(max_pts - min_pts)
    # axis_sq = list(set([0,1,2]).difference([axis_step]))
    # sq_dim = (max_pts-min_pts)[axis_sq]

    
    
    # total_sz = (max_pts-min_pts)[axis_step]
    
    # if tab_width is not None:
    #     num_tabs = int(np.ceil(total_sz/tab_width))
    
    # step_sz = total_sz / num_tabs
    # box_sz = max_pts-min_pts
    # box_sz[axis_step] /= num_tabs
    # tbj = list()
    # tbh = list()
    # for i in range(num_tabs):
    #     xstep = [0,0,0]
    #     if (axis_sq[1]-axis_sq[0]) > 1:
    #         xstep[axis_sq[1]] = 1
    #         flip = True
    #     else:
    #         xstep[axis_sq[0]] = 1
    #         flip = False
    #     normal = [0,0,0]
    #     normal[axis_step] = 1
    #     origin = min_pts.copy()
    #     origin[axis_step] += i * step_sz
    #     t_plane = cq.Plane(tuple(origin), tuple(xstep), tuple(normal))
        
    #     if flip:
    #         tbj.append( add_rect(cq.Workplane(t_plane), sq_dim[1], sq_dim[0]).extrude(step_sz) )
    #     else:
    #         tbj.append( add_rect(cq.Workplane(t_plane), sq_dim[0], sq_dim[1]).extrude(step_sz) )
  
    # if cut:
    #     for i in range(num_tabs):
    #         if i % 2 == 0:
    #             r_side = r_side.cut(tbj[i])

    #         if i % 2 == 1:
    #             f_side = f_side.cut(tbj[i])
    
    # return r_side, f_side, tbj



def ret_front_pts(bd, bh, spread = 96):
    pts = list()
    if spread == 0:
        pts.append( (bd/2, bh))
    else:
        pts.append( (bd/2 - spread/2, bh))
        pts.append( (bd/2 + spread/2, bh))

    return pts
    
    

def ret_pts(bd, bh, center = False, max_num = None, spacing = 32):
    pts = list()
    if center:
            ix = int(np.ceil(bd/spacing)/2)
            for i in range((-ix+1), ix ):
                pts.append((bd/2 + i*spacing, bh))
    else:
        for i in range(1,int(np.ceil(bd/spacing)) - 1):
            pts.append((i*spacing, bh))
    
    pts = tuple(pts)
    return pts


def slide_pts(bd, bh, offset_from_back = 54, spacing = (192, 192), reverse = False, which='push-open'):
    if which == 'push-open':
        pass
    elif which == '24inch':
        # spacing = (278,201)
        spacing = (224,256)


    pts = list()
    if reverse:
        curr_pt = bd-offset_from_back
        pts.append([curr_pt, bh])
        for spc in spacing:
            curr_pt -= spc
            pts.append([ curr_pt, bh])

        return pts


    else:
        curr_pt = offset_from_back
        pts.append([curr_pt, bh])
        for spc in spacing:
            curr_pt += spc
            pts.append([ curr_pt, bh])
        return pts