def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('pkg', help='Package name', nargs=None, type=str)
    parser.add_argument('--debug',
                        help='Enable debugging output',
                        action='store_true')
    args = parser.parse_args()
    pkg = getattr(dualdata, args.pkg)

    # p = multiprocessing.Pool(8)
    p = multiprocessing.Pool(1)
    for i, e in enumerate(pkg.STICKS_X_DESC):
        e['debug'] = args.debug
        e['bar_id'] = i
    for i, e in enumerate(pkg.STICKS_Y_DESC):
        e['debug'] = args.debug
        e['bar_id'] = i
    print(pkg.STICKS_X_DESC)
    print(pkg.STICKS_Y_DESC)
    '''
    meshes_x = []
    for desc in pkg.STICKS_X_DESC:
        meshes_x.append(build_stick_x(desc))
    '''
    meshes_x = p.map(build_stick_x, pkg.STICKS_X_DESC)
    # print('Mpos {}'.format(meshes_x[0][0]))
    # print('Mneg {}'.format(meshes_x[0][1]))
    # meshes_x = []
    '''
    meshes_y = []
    for desc in pkg.STICKS_Y_DESC:
        meshes_y.append(build_stick_y(desc))
    '''
    meshes_y = p.map(build_stick_y, pkg.STICKS_Y_DESC)
    # meshes_y = []
    meshes = meshes_x + meshes_y
    if args.debug:
        Mpos = []
        Mneg = []
        for i, e in enumerate(meshes):
            #print("{}: {}".format(i, e))
            pmesh, nmesh = e
            Mpos += pmesh
            Mneg += nmesh
        # print(Mneg)
        multi_save(Mpos, 'dual-pos.obj', color=np.array([0.0, 0.0, 0.75]))
        multi_save(Mneg, 'dual-neg.obj', color=np.array([0.75, 0.0, 0.0]))
        return
    # meshes = meshes[0:1] # Debug
    print(len(meshes))
    while len(meshes) > 1:
        meshes_next = p.map(merge_mesh, zip(meshes[0::2], meshes[1::2]))
        if len(meshes) % 2 == 1:
            meshes_next.append(meshes[-1])
        print(len(meshes_next))
        meshes = meshes_next
    V, F = meshes[0]
    print("Final V F {} {}".format(V.shape, F.shape))
    pyosr.save_obj_1(V, F, 'dual.obj')
def multi_save(meshes, fn, color):
    Vall = None
    Fall = None
    for V, F in meshes:
        Vall = V if Vall is None else np.concatenate((Vall, V), axis=0)
        # print("FT {}".format(F))
        F = np.copy(F) + Vall.shape[0] - V.shape[0]
        # print("F delta {}".format(Vall.shape[0] - V.shape[0]))
        Fall = F if Fall is None else np.concatenate((Fall, F), axis=0)
    empty = np.empty(shape=(0, 0))
    color = color.reshape((1, 3))
    Vc = np.repeat(color, Vall.shape[0], axis=0)
    print(Vall.shape)
    print(Vc.shape)
    V = np.concatenate((Vall, Vc), axis=1)
    F = Fall
    pyosr.save_obj_1(V, F, fn)
 def _dump_conf_common(self, fn_func, key_str):
     uw = self._uw
     vert_id = self._args.vert_id
     conf_sel = np.array(aux.range2list(self._args.conf_sel), dtype=np.int)
     print("Printing {}".format(conf_sel))
     io_dir = self._args.in_dir
     out_obj_dir = self._args.out_dir
     # Iterate through all ReTouchQ files
     batch_id = 0
     conf_base = 0
     largest = np.max(conf_sel)
     while conf_base <= largest:
         fn = fn_func(io_dir, vert_id=vert_id, batch_id=batch_id)
         if not os.path.exists(fn):
             print("Fatal: Cannot find file {}".format(fn))
             return
         Q = np.load(fn)[key_str]
         indices = np.where(
             np.logical_and(conf_sel >= conf_base,
                            conf_sel < conf_base + len(Q)))
         indices = indices[0]
         # print("Indices {}".format(indices))
         Vall_list = None
         Fall_list = None
         Fall_base = 0
         for index in indices:
             conf_id = conf_sel[index]
             # print("conf_sel index {}".format(index))
             # print("conf_id {}".format(conf_id))
             rem = conf_id - conf_base
             # print("rem {}".format(rem))
             q = Q[rem]
             print("Getting geometry from file {}, offset {}".format(
                 fn, rem))
             print("tq {}".format(q))
             V1, F1 = uw.get_robot_geometry(q, True)
             out_obj = "{}/rob-vert-{}-conf-{}.obj".format(
                 out_obj_dir, vert_id, conf_id)
             pyosr.save_obj_1(V1, F1, out_obj)
             out_obj = "{}/env-vert-{}-conf-{}.obj".format(
                 out_obj_dir, vert_id, conf_id)
             V2, F2 = uw.get_scene_geometry(q, True)
             pyosr.save_obj_1(V2, F2, out_obj)
             out_obj = "{}/union-vert-{}-conf-{}.obj".format(
                 out_obj_dir, vert_id, conf_id)
             Va = np.concatenate((V1, V2), axis=0)
             F2 += V1.shape[0]
             Fa = np.concatenate((F1, F2), axis=0)
             if Vall_list is None:
                 # Include the environment
                 Vall_list = [V1, V2]
                 Fall_list = [F1, F2]  # F2 has been rebased
                 Fall_base = V1.shape[0] + V2.shape[0]
             else:
                 # Exclude the environment, so env geo only included once
                 Vall_list.append(V1)
                 F1prime = F1 + Fall_base
                 Fall_list.append(F1prime)
                 Fall_base += V1.shape[0]
             #print("Dumping to {}".format(out_obj))
             pyosr.save_obj_1(Va, Fa, out_obj)
         Vall = np.concatenate(Vall_list, axis=0)
         Fall = np.concatenate(Fall_list, axis=0)
         out_obj = "{}/union-all.obj".format(out_obj_dir)
         pyosr.save_obj_1(Vall, Fall, out_obj)
         batch_id += 1
         conf_base += len(Q)
 def _common_sample(self, enum_axis=False):
     uw = self._uw
     task_id = self._args.task_id
     batch_size = self._args.batch_size
     io_dir = self._args.io_dir
     tp = TaskPartitioner(io_dir,
                          None,
                          batch_size,
                          tunnel_v=self._get_tunnel_v())
     rob_sampler = atlas.AtlasSampler(tp, 'rob', uw.GEO_ROB, task_id)
     env_sampler = atlas.AtlasSampler(tp, 'env', uw.GEO_ENV, task_id)
     pcloud1 = []
     pn1 = []
     pcloud1x = []
     pn1x = []
     pcloud2 = []
     pn2 = []
     conf = []
     signature_conf = []
     SANITY_CHECK = False
     if not SANITY_CHECK:
         for i in progressbar(range(batch_size)):
             if enum_axis:
                 tup1 = rob_sampler.sample(uw)
                 tup2 = env_sampler.sample(uw)
                 qs_raw = uw.enum_free_configuration(tup1[0],
                                                     tup1[1],
                                                     tup2[0],
                                                     tup2[1],
                                                     1e-6,
                                                     denominator=64)
                 qs = [q for q in qs_raw if not uw.is_disentangled(q)
                       ]  # Trim disentangled state
                 for q in qs:
                     conf.append(q)
                     # break # debug: only one sample per touch
                 if len(qs) > 0:
                     pcloud2.append(tup2[0])
                     signature_conf.append(qs[0])
             else:
                 while True:
                     tup1 = rob_sampler.sample(uw)
                     tup2 = env_sampler.sample(uw)
                     q = uw.sample_free_configuration(tup1[0],
                                                      tup1[1],
                                                      tup2[0],
                                                      tup2[1],
                                                      1e-6,
                                                      max_trials=16)
                     if uw.is_valid_state(q):
                         break
                 conf.append(q)
         # print("tqre_fn {}".format(tp.get_tqre_fn(task_id)))
         np.savez(tp.get_tqre_fn(task_id),
                  ReTouchQ=conf,
                  SigReTouchQ=signature_conf)
         np.savetxt('pc2.txt', pcloud2, fmt='%.17g')
     else:
         #
         # Sanity check code
         #
         fail = 0
         for i in progressbar(range(32)):
             while True:
                 tup1 = rob_sampler.sample(uw)
                 tup2 = env_sampler.sample(uw)
                 q = uw.sample_free_configuration(tup1[0],
                                                  tup1[1],
                                                  tup2[0],
                                                  tup2[1],
                                                  1e-6,
                                                  max_trials=16)
                 fail += 1
                 if uw.is_valid_state(q):
                     break
                 print("Reject {}".format(q))
             print("Accept {}".format(q))
             pcloud1.append(tup1[0])
             pcloud2.append(tup2[0])
             conf.append(q)
             '''
             # Sanity check
             # Test if sample_free_configuration aligns the sample pair
             tup1 = rob_sampler.sample(uw, unit=False)
             tup2 = env_sampler.sample(uw, unit=False)
             pcloud1.append(tup1[0])
             pcloud2.append(tup2[0])
             pn1.append(tup1[1])
             pn2.append(tup2[1])
             q = uw.sample_free_configuration(tup1[0], tup1[1], tup2[0], tup2[1], 1e-6, free_guarantee=False)
             tr,rot = pyosr.decompose_2(q)
             new_point = rot.dot(tup1[0].reshape((3,1)))+tr.reshape((3,1))
             pcloud1x.append(new_point.reshape((3)))
             new_normal = rot.dot(tup1[1].reshape((3,1)))
             pn1x.append(new_normal.reshape((3)))
             conf.append(q)
             '''
         pyosr.save_obj_1(pcloud1, [], 'pc1.obj')
         pyosr.save_obj_1(pcloud2, [], 'pc2.obj')
         np.savez('cf1.npz', Q=conf)
         print("Success ratio {} out of {} = {}".format(
             len(conf), fail,
             len(conf) / float(fail)))
         '''
    def uvrender(self):
        uw = self._uw
        tunnel_v = self._get_tunnel_v()
        uw.avi = False
        geo_type = self._args.geo_type
        TYPE_TO_FLAG = {
            'rob': uw.BARY_RENDERING_ROBOT,
            'env': uw.BARY_RENDERING_SCENE
        }
        geo_flag = TYPE_TO_FLAG[geo_type]

        vert_id = self._args.vert_id
        io_dir = self._args.io_dir
        iq = uw.translate_to_unit_state(tunnel_v[vert_id])
        afb = None
        afb_nw = None

        tq_gen = TouchQGenerator(in_dir=io_dir, vert_id=vert_id)
        obj_gen = UVObjGenerator(in_dir=io_dir,
                                 geo_type=geo_type,
                                 vert_id=vert_id)
        i = 0
        DEBUG_UVRENDER = False
        for tq, is_inf in tq_gen:
            # print('tq {} is_inf {}'.format(tq, is_inf))
            IBV, IF = next(obj_gen)  # Must, otherwise does not pair
            if is_inf:
                continue
            if IBV is None or IF is None:
                print('IBV {}'.format(None))
                continue
            print('{}: IBV {} IF {}'.format(i, IBV.shape, IF.shape))
            if DEBUG_UVRENDER:
                svg_fn = '{}.svg'.format(i)
                # Paint everything ...
                if i == 0 and geo_type == 'rob':
                    V, F = uw.get_robot_geometry(tq, True)
                    print("V {}\nF {}".format(V.shape, F.shape))
                    IF, IBV = uw.intersecting_to_robot_surface(tq, True, V, F)
                    print("IF {}\nIBV {}\n{}".format(IF.shape, IBV.shape,
                                                     IBV[:5]))
                    '''
                    NPICK=3000
                    IF = IF[:NPICK]
                    IBV = IBV[:NPICK*3]
                    '''
            else:
                svg_fn = ''
            uw.clear_barycentric(geo_flag)
            uw.add_barycentric(IF, IBV, geo_flag)
            if DEBUG_UVRENDER and i == 2:
                print("BaryF {}".format(IF))
                print("Bary {}".format(IBV))
            fb = uw.render_barycentric(geo_flag,
                                       np.array([ATLAS_RES, ATLAS_RES],
                                                dtype=np.int32),
                                       svg_fn=svg_fn)
            #np.clip(fb, 0, 1, out=fb) # Clip to binary
            nw = texture_format.texture_to_file(fb.astype(np.float32))
            distance = np.clip(pyosr.distance(tq, iq), 1e-4, None)
            w = nw * (1.0 / distance)
            if afb is None:
                afb = w
                afb_nw = nw
            else:
                afb += w
                afb_nw += nw
                np.clip(afb_nw, 0, 1.0,
                        out=afb_nw)  # afb_nw is supposed to be binary
            # Debugging code
            if DEBUG_UVRENDER:
                print('afb shape {}'.format(afb.shape))
                print('distance {}'.format(distance))
                rgb = np.zeros(list(afb.shape) + [3])
                rgb[..., 1] = w
                imsave(_fn_atlastex(io_dir, geo_type, vert_id, i), rgb)
                np.savez(atlas_fn(io_dir, geo_type, vert_id, i), w)
                print('NW NZ {}'.format(nw[np.nonzero(nw)]))
                V1, F1 = uw.get_robot_geometry(tq, True)
                pyosr.save_obj_1(V1, F1, '{}.obj'.format(i))
                V2, F2 = uw.get_scene_geometry(tq, True)
                pyosr.save_obj_1(V2, F2, '{}e.obj'.format(i))
                if i >= 16:
                    return
            i += 1
        rgb = np.zeros(list(afb.shape) + [3])
        rgb[..., 1] = afb
        # FIXME: Give savez an explicity array name
        imsave(_fn_atlastex(io_dir, geo_type, vert_id, None), rgb)
        np.savez(atlas_fn(io_dir, geo_type, vert_id, None), afb)
        rgb[..., 1] = afb_nw
        imsave(_fn_atlastex(io_dir, geo_type, vert_id, None, nw=True), rgb)
        np.savez(atlas_fn(io_dir, geo_type, vert_id, None, nw=True), afb)
def _debug_notch_worker(wag):
    ws = util.create_workspace_from_args(wag.args)
    ws.current_trial = wag.current_trial
    kpp = pygeokey.KeyPointProber(wag.refined_env_fn)
    os.makedirs(ws.local_ws('debug'), exist_ok=True)
    kps_fn = ws.local_ws('debug', f'notch-{ws.current_trial}.npz')
    dic = dict()
    SKV = kpp.get_all_skeleton_points()
    SKE = kpp.get_skeleton_edges()
    SKVd = np.concatenate([SKV, SKV], axis=0)
    SKF = np.zeros((SKE.shape[0], 3), dtype=SKE.dtype)
    SKF[:, 0] = SKE[:, 0]
    SKF[:, 1] = SKE[:, 1]
    SKF[:, 2] = SKE[:, 0] + SKV.shape[0]
    import pyosr
    pyosr.save_obj_1(SKVd, SKF,
                     ws.local_ws('debug', f'ske-{ws.current_trial}.obj'))
    # pyosr.save_obj_1(SKV, SKE, ws.local_ws('debug', f'ske-{ws.current_trial}.obj'))
    # pyosr.save_ply_2(SKV, SKE, np.zeros([]), np.array([]), ws.local_ws('debug', f'ske-{ws.current_trial}.ply'))
    kpp.local_min_thresh = 0.10
    kpp.group_tolerance_epsilon = 0.05
    for i in range(1):
        npts = kpp.probe_notch_points(seed=0, keep_intermediate_data=True)
        ZF = np.zeros(shape=(0, 0))
        for it in range(3):
            P1 = kpp.get_intermediate_data(it, 0)
            P2 = kpp.get_intermediate_data(it, 1)
            P3 = kpp.get_intermediate_data(it, 2)
            pyosr.save_obj_1(
                P1, ZF, ws.local_ws('debug',
                                    f'P1-I{it}-{ws.current_trial}.obj'))
            pyosr.save_obj_1(
                P2, ZF, ws.local_ws('debug',
                                    f'P2-I{it}-{ws.current_trial}.obj'))
            pyosr.save_obj_1(
                P3, ZF, ws.local_ws('debug',
                                    f'P3-I{it}-{ws.current_trial}.obj'))
            PE = np.concatenate((P2, P3, P3), axis=0)
            D = P2.shape[0]
            edges = [(i, i + D, i + D * 2) for i in range(D)]
            pyosr.save_obj_1(
                PE, edges,
                ws.local_ws('debug', f'PE-I{it}-{ws.current_trial}.obj'))
        npt1 = npts[:, :3]
        npt2 = npts[:, 3:6]
        D = npt1.shape[0]
        nptb = np.concatenate((npt1, npt2, npt2), axis=0)
        edges = [(i, i + D, i + D * 2) for i in range(D)]
        pyosr.save_obj_1(nptb, edges,
                         ws.local_ws('debug', f'NPTE-{ws.current_trial}.obj'))
        matio.savetxt(ws.local_ws('debug', f'EL-{ws.current_trial}.txt'),
                      np.linalg.norm(npt1 - npt2, axis=1))
        dic[str(i)] = npts
    np.savez(kps_fn, **dic)