loop_representations.py

# module for extracting various loop representations for use as input and output
# formats for models to receive/predict
from typing import Tuple, Optional
import numpy as np
import warnings
warnings.filterwarnings('error', category=np.VisibleDeprecationWarning)

import polyscope as ps
import os
import sys

import get_loops
import geometry_utils
import utils
import meshlab_poisson_reco
from thlog import *

thlog = Thlogger(LOG_INFO, VIZ_DEBUG, "loop_repr", 
    imports=[get_loops.thlog, meshlab_poisson_reco.thlog])


# This module involves two directions of conversions:
# The LoopRepresentation class and the get_oriented_point_cloud method describes
# going from (possibly network-predicted) loop representations (i.e. ellipse
# coefficient arrays) to the standard MeshOneSlice segmented representation,
# then to oriented point cloud format.
# This is to visualize and convert the predictions back to meshes via meshlab
# poisson reconstruction.

# And `load_loop_repr_from_mesh`, on the other hand, will slice a given mesh at
# given planes, and produce loop representations (e.g. also ellipse
# coefficients) for use as training data ground truth.

###############################################################################
# LOADING MESH SLICES FROM (PREDICTED) LOOP REPRESENTATIONS
###############################################################################
# helper to sample from a list of mesh slices generated by LoopRepresentation objects
# This function is used in the direction from "model-predicted" => "oriented pcloud"
# for visualization.
def sample_from_list_of_mesh_slices(mesh_slices: list, n_points=50, 
    distribute_points_by_loop_length=True, 
    min_n_points_per_loop=-1, 
    normal_lerping=True):
    """

    Renders a list of mesh slices into points and normals for visualization.
    Parameters:
    - mesh_slices: list of MeshOneSlice objects
    - n_points: number of points in total/per loop; depends on the below args.
    - distribute_points_by_loop_length (default=True): if True, try to meet
        n_points in TOTAL by spreading out points by length of the closed loop. 
        If False, each closed loop gets exactly n_points.
    - min_n_points_per_loop: (default=-1, disabled) 
        to prevent short/small loops to be undersampled in the reconstruction,
        we can impose a minimum number of points per closed loop here,
        regardless of distribution by length.
    - normal_lerping: (default=True) if True, the normals of sampled points 
        are interpolated between the normals of the endpoints of the edge that
        these sampled points come from. If False, the nearest endpoint normal is
        used.

    Returns: (points, normals), each of shape (n, 3).
    """
    all_points= None
    all_normals = None 

    # colors for the viz

    for viz_i, mesh_slice in enumerate(mesh_slices):
        thlog.do(VIZ_DEBUG, 
            lambda: ps.register_curve_network(f"cunet{viz_i}", 
                mesh_slice.loop_pts, mesh_slice.loop_conn_arr, enabled=True, color=utils.PS_COLOR_CURVE_NETWORK)
            , needs_polyscope=True)
    
        # assume normals have already been estimated
        if thlog.guard(VIZ_DEBUG, needs_polyscope=True):
            # mesh_slice = estimate_loop_segment_normals_simple(mesh_slice)
            ps.get_curve_network(f"cunet{viz_i}").add_vector_quantity(
                f"simple_norm_est{viz_i}", 
                mesh_slice.loop_segment_normals,
                defined_on="nodes", enabled=False, length=0.1)

        # try sampling from predicted reprs 
        # (which we will feed to poisson reconstructor)
        sampled_points, sampled_normals = \
            get_loops.sample_from_loops_and_normals(mesh_slice, n_points,
            distribute_points_by_loop_length=distribute_points_by_loop_length,
            min_n_points_per_loop=min_n_points_per_loop,
            normal_lerping=normal_lerping)

        if thlog.guard(VIZ_INFO, needs_polyscope=True): 
            pcloud = ps.register_point_cloud(f"sampled_pc{viz_i}", sampled_points,
                enabled=thlog.guard(VIZ_TRACE, needs_polyscope=True))
            pcloud.add_vector_quantity(f"sampled_pc_norms{viz_i}", sampled_normals, 
                enabled=thlog.guard(VIZ_TRACE, needs_polyscope=True), length=0.11)

        # accumulate all the sub-pointclouds sampled from each slice into
        # a big oriented point cloud we can then later give to the poisson
        # reconstructor
        if all_points is None or all_normals is None:
            all_points = sampled_points
            all_normals = sampled_normals
        else:
            all_points = np.vstack((all_points, sampled_points))
            all_normals = np.vstack((all_normals, sampled_normals))
    return all_points, all_normals


def save_npz_file_from_list_of_mesh_slices(mesh_slices: list, fname: str = ""):
    """
    Save a list of mesh slices to a packed .npz format with the following arrays
    The format requires packing with delimiters because each slice may have a
    different number of points. 
    Where n_i is the number of vertices of slice plane i, p_i is the plane
    corresponding to slice_plane i:
    - loop_pts_packed: array of size (sum([n_i + 1 for each slice i]), 3).
        Each slice's section in the array is preceded by a row of 
        np.NINF (negative infinity).
    - loop_conn_arr_packed: array of size (sum([n_i + 1 for each slice i]), 2).
        Each slice is preceded by a row of -1. (negative numbers are
        never present in these integer connectivity arrays, so -1 is used as
        delimiter.)
    - loop_segment_normals_packed: array of size (sum([n_i + 1 for each slice i]), 3)
        Each slice is preceded by a row of np.NINF.
    The only non-packed array is 'planes' which is just (n_slices, 4, 3).
    - planes: contains all the planes corresponding to these slices. See
        geometry_utils and get_loops for the plane data format.
    """

    loop_pts_packed = None
    loop_conn_arr_packed = None
    loop_segment_normals_packed = None

    for mesh_slice in mesh_slices:

        # first add the delimiting row of np.NINF (negative infinity) or (-1)s
        loop_pts_packed = np.full((1,3), np.NINF) if loop_pts_packed is None else \
            np.concatenate((loop_pts_packed, np.full((1,3), np.NINF)))
        loop_conn_arr_packed = np.full((1, 2), -1) if loop_conn_arr_packed is None else \
            np.concatenate((loop_conn_arr_packed, np.full((1, 2), -1)))
        loop_segment_normals_packed = np.full((1, 3), np.NINF) if loop_segment_normals_packed is None else \
            np.concatenate((loop_segment_normals_packed, np.full((1, 3), np.NINF)))

        loop_pts_packed = np.concatenate((loop_pts_packed, mesh_slice.loop_pts), axis=0)
        loop_conn_arr_packed = np.concatenate((loop_conn_arr_packed, mesh_slice.loop_conn_arr), axis=0)
        loop_segment_normals_packed = np.concatenate((loop_segment_normals_packed, mesh_slice.loop_segment_normals), axis=0)
        
    # then save the planes out. We don't have to do anything fancy with this
    # because all planes have shape (4,3) so we can just list them in an array
    planes = np.array(list(map(lambda ms: ms.plane, mesh_slices)))
    # write out
    if fname and loop_pts_packed is not None and loop_conn_arr_packed is not None and loop_segment_normals_packed is not None:
        np.savez_compressed(fname 
            , loop_pts_packed = loop_pts_packed 
            , loop_conn_arr_packed = loop_conn_arr_packed
            , loop_segment_normals_packed = loop_segment_normals_packed
            , planes = planes
        )
        return fname
    
    else:
        # no fname, just return all the arrays rather than saving to file
        return loop_pts_packed, loop_conn_arr_packed, loop_segment_normals_packed


def read_mesh_slices_npz_file_for_polyscope(fname, consolidate_into_one_curve_network=True):
    """ 
    Reads the file exported by save_npz_file_from_list_of_mesh_slices and
    returns curve network data usable with polyscope.

    If consolidate_into_one_curve_network=True (default), returns a single tuple
    of the three arrays `loop_pts` and `loop_conn_arr` and
    `loop_segment_normals`, which can be directly used with a single call of
    polyscope's register_curve_network() for visualizing loops.

    (sanity note: in this case, these loop_* arrays in the code below are
    different from the same-named ones in MeshOneSlice code (mostly in
    get_loops.py and below here). This function's loop_* arrays are consolidated
    and contain data for ALL slices/planes involved (i.e. we only do one
    ps.register_curve_network call to load it into ps) whereas each MeshOneSlice
    and its values concern only one plane.)

    If consolidate_into_one_curve_network is False, then returns 3 lists
    containing loop_pts, loop_conn_arr, loop_segment_normals arrays, one for
    each slice. Needs n_slices calls to register_curve_network() to fully init
    in polyscope, but allows separate coloring of each slice for nice visuals.
    """
    with np.load(fname) as npz:
        loop_pts = npz["loop_pts_packed"]
        loop_conn_arr = npz["loop_conn_arr_packed"]
        loop_segment_normals = npz["loop_segment_normals_packed"]

    if consolidate_into_one_curve_network:
        # we can just combine all points and conn into one curve network. all we
        # need to do is nudge up the loop_conn_arr indices then remove the delimiter
        # rows.
        curr_conn_offset = 0
        segments_seen_this_slice = 0
        
        for row_i, row in enumerate(loop_conn_arr):
            if np.all(row == -1):
                curr_conn_offset = segments_seen_this_slice
            else:
                loop_conn_arr[row_i] += curr_conn_offset
                segments_seen_this_slice += 1
        
        # mask out the -infs and -1s
        loop_pts = loop_pts[np.all(np.logical_not(np.isneginf(loop_pts)), axis=-1)]
        loop_conn_arr = loop_conn_arr[np.all(loop_conn_arr != -1, axis=-1)]
        loop_segment_normals = loop_segment_normals[np.all(np.logical_not(np.isneginf(loop_segment_normals)), axis=-1)]

        return loop_pts, loop_conn_arr, loop_segment_normals
    else:
        loop_pts_each_slice = []
        loop_conn_arr_each_slice = []
        loop_segment_normals_each_slice = []
        current_slice_starts_at = None
        for i in range(len(loop_pts)): # the three npz-loaded arrays have the same length
            if np.all(loop_conn_arr[i] == -1):
                # is a delimiter row, gather existing slices
                if current_slice_starts_at is not None:
                    start_i = current_slice_starts_at
                    loop_pts_each_slice.append(loop_pts[start_i:i])
                    loop_conn_arr_each_slice.append(loop_conn_arr[start_i:i])
                    loop_segment_normals_each_slice.append(loop_segment_normals[start_i:i])
                    current_slice_starts_at = i+1 # skip this delimiter row, next slice contents start the next row
                else:
                    current_slice_starts_at = 1
        return loop_pts_each_slice, loop_conn_arr_each_slice, loop_segment_normals_each_slice
            

                


def save_contour_file_from_list_of_mesh_slices(mesh_slices: list, fname: str = ""):
    """
    Experimental function (also part of the "model-predicted => reconstructed mesh" 
    direction of the pipeline) that tries to export a list of mesh slices to the ".contour"
    file format described at https://www.cs.wustl.edu/~taoju/lliu/paper/ctr2suf/program.html
    for reconstructing using Taoju's contour reconstruction method
    """
    n_planes = len(mesh_slices)
    contour_file_lines = [f"{n_planes}"]
    def __to_cartesian(plane_arr):
        # using the same plane format as in get_loops, we calculate the
        # cartesian eqn form (ax+by+cz=d) used by the .contour format
        plane_n, plane_p = plane_arr[0],  plane_arr[1]
        return (plane_n[0], plane_n[1], plane_n[2], np.dot(plane_n, plane_p))
    
    __cosine_sim = lambda a, b: (np.dot(a,b) / (np.linalg.norm(a) * np.linalg.norm(b)))

    for mesh_slice in mesh_slices:
        # plane definition line. a b c d
        a, b, c, d = __to_cartesian(mesh_slice.plane)
        contour_file_lines.append(f"{a} {b} {c} {d}")

        # slice vert/edge count line. #vers_1(m) #edges(n)
        n_vertices = len(mesh_slice.loop_pts)
        n_edges = len(mesh_slice.loop_conn_arr)
        contour_file_lines.append(f"{n_vertices} {n_edges}")

        # then list out all the vertices
        for loop_pt in mesh_slice.loop_pts:
            contour_file_lines.append(f"{loop_pt[0]} {loop_pt[1]} {loop_pt[2]}")
        
        # then list out all the edges.

        # to find the material of each edge in loop_conn_arr (required by the file format),
        # we take the cross prod of the edge normal and the plane_n, and
        # if the edge conn IS IN THE DIRECTION of this cross product,  
        # then its left side is the outside material and the right s ide
        # is the mesh's "inside".
        # (these normals already account for disjoint loops etc; the normals
        # always point outwards, but the vertex-index order might not reflect
        # the same outward-pointing normals if the normal were calculated from that alone.)
        # ((note that the estimated normals always point outwards 
        # unless we run into inner-surfaces; to be implemented))
        plane_n = mesh_slice.plane[0]
        for segment in mesh_slice.loop_conn_arr:
            from_i = segment[0]
            to_i = segment[1]
            seg_norm = mesh_slice.loop_segment_normals[from_i] # see get_loops.py for similar usage of loop_segment_normals..
            seg_vec = mesh_slice.loop_segment_vectors[from_i]
            forward_direction = np.cross(seg_norm, plane_n)
            dir_similarity = np.dot(seg_vec, plane_n)
            if dir_similarity < 0:
                # this edge connectivity index order is opposite the 
                # "canonical"  forward direction; we 
                # flip it in the file output so that the
                # "outside" is on the left, "inside" is on the right.
                from_i_towrite = to_i
                to_i_towrite = from_i
            else:
                # the edge connectivity index order implies the same
                # direction as the forward_direction calculated from
                # the plane normal and edge normal, so the left side
                # is already OUtside and the right side Inside
                from_i_towrite = from_i
                to_i_towrite = to_i
            # note that we consider "outside material" = material 1,
            # and inside material = material 0
            contour_file_lines.append(f"{from_i_towrite} {to_i_towrite} 0 1")

    # done with the file format lines; give all lines a newline then write out
    
    if fname:
        with open(fname, "w") as f:
            f.writelines(map(lambda l: l + "\n", contour_file_lines))
    thlog.info(f"Saved .contour file to {fname}")
    return


# template for representations of loosp that may be predicted through a neural
# network..
class LoopRepresentation:
    def get_oriented_point_cloud(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        This is responsible for converting the whole sequence of data into
        points and normals for an oriented point cloud, to be fed into the 
        poisson reconstructor.
        """
        raise NotImplementedError()
        
# this should be created/init'd by inference code, after receiving outputs
# predicted by a model
class EllipseSingle(LoopRepresentation):
    def __init__(self, planes: np.ndarray, predicted_sequence: np.ndarray, 
        segmentize_resolution: int = 30, n_points_to_sample_for_point_cloud: int=64):
        """
        This does not support multiple ellipses per mesh level. 

        Parameters:
        - planes: iterable of np arrays each (4, 3) describing the plane normal, 
            a point on plane, and two orthogonal vectors on/parallel to the
            plane selected to be the x and y axes of a coordinate system with
            the plane normal as the z axis and the plane point as the origin.
            (same format used in MeshOneSlice.) plane[0] is plane_n, plane[1] is
            plane_p, plane[2] and plane[3] are such that np.array([plane[2],
            plane[3],plane[0] ]) form an orthogonal basis for a coordinate
            system with origin being plane_p. (i.e. such that we can throw away
            the z coordinate to obtain "2D local coordinates of the points on
            the plane with respect to a 2D coord system on the plane")
            
        - predicted_sequence: np array (n_timesteps, 6) consisting of the 6 polar
          coefficients (see ellipse_cart2pol

        - segmentize_resolution: to do normal estimation, we go through the step
            of sampling points from the ellipse, and then forming a closed
            polyline out of it. This parameter determines the number of points
            we should sample from the ellipse for this.
        - n_points_to_sample_for_point_cloud: ...the above is to discretize the
            ellipse. We then have to sample FROM this discretized segments thing
            as well. This is the sampling resolution for THAT step.

        """
        self.pred_seq = predicted_sequence
        self.planes = planes
        self.segmentize_resolution = segmentize_resolution
        self.n_points_to_sample_for_point_cloud = n_points_to_sample_for_point_cloud

        self.mesh_slices = self.process_into_mesh_slices()

    def process_into_mesh_slices(self):
        # recall that a "time step" is actually a single slice plane in this
        # loop representation
        mesh_slices = []
        for plane, ellipse_params_polar in zip(self.planes, self.pred_seq):
            try:
                mesh_slice = EllipseSingle.to_mesh_slice(plane, ellipse_params_polar, 
                    segmentize_resolution=self.segmentize_resolution)
            except AssertionError as assert_e:
                thlog.err(assert_e.args)
                continue
            if len(mesh_slice.loop_conn_arr) == 0:
                thlog.trace("loop slice is empty, skipping this plane")
                continue
            mesh_slices.append(mesh_slice)
        
        return mesh_slices
    
    def get_oriented_point_cloud(self):
        return sample_from_list_of_mesh_slices(self.mesh_slices)

    def export_as_contour_file(self, fname: str):
        return save_contour_file_from_list_of_mesh_slices(self.mesh_slices, fname)

    def export_as_npz_file(self, fname: str):
        return save_npz_file_from_list_of_mesh_slices(self.mesh_slices, fname)
    
    @staticmethod
    def to_mesh_slice(plane, ellipse_params_polar, return_just_data_no_slice=False
        , segmentize_resolution=100, run_in_fixed_resolution_polylines_mode=False):
        """
        plane is (4,3) described above,
        ellipse_params_polar should be (n_ellipse_features)
        This is a staticmethod because this is reusable in the other 
        LoopRepresentation classes.

        return_just_data_no_slice: defaults to False; if True, then don't create
        a MeshOneSlice object, but just return the data needed to make one, i.e.
        the loop points and loop point connectivity. (This is useful for when
        we need to collect up a whole bunch of these pairs of arrays (points, conn)
        each one representing a disjoint loop in the mesh slice, then creating
        the actual Slice later at the end in one swoop,.

        the optional arg run_in_fixed_resolution_polylines_mode is there to
        skip the ellipse-sampling step if we already have the actual point data.
        most of the code after that is the same so we can reuse this.
        See EllipseMultiple.get_oriented_point_cloud for the use case.
        """
        plane_n, plane_p, onplane_xaxis, onplane_yaxis, plane_basis, \
            change_to_plane_basis, plane_p_in_plane_basis = \
                geometry_utils.extract_useful_data_from_plane(plane)

        # ellipse_params_polar = geometry_utils.ellipse_cartesian2polar(self.ellipse_cartesian_coeffs)
        if run_in_fixed_resolution_polylines_mode:
            # when this is True, "ellipse_params_polar" is actually the raw
            # point data in a flattened array of shape (n_points * 2, ).
            # we only have to reshape it, but nothing else should change
            ellipse_sampled_points = ellipse_params_polar.reshape((-1, 2))
        else:
            ellipse_sampled_points = geometry_utils.sample_from_ellipse(
                segmentize_resolution, ellipse_params_polar)

        # tack on the z-coordinate of 0 again to each point
        ellipse_sampled_points = np.pad(ellipse_sampled_points, ((0,0),(0,1)))

        # convert back into world coordinates by multiplying the plane_basis with
        # each found-point...
        ellipse_sampled_points_in_world_coords = \
            np.transpose(np.dot(plane_basis, np.transpose(ellipse_sampled_points)))
        
        # and add back the plane_p (i.e. translate back to world coords origin)
        ellipse_sampled_points_in_world_coords += plane_p

        # make fake connectivity
        n_ellipse_points = len(ellipse_sampled_points)
        ellipse_conn = np.array([[i, (i+1) % n_ellipse_points] for i in range(n_ellipse_points)])

        if return_just_data_no_slice:
            return ellipse_sampled_points_in_world_coords, ellipse_conn
        else:
            slice_obj = get_loops.MeshOneSlice(plane, 
                predicted_loops=(ellipse_sampled_points_in_world_coords, ellipse_conn, None))
            slice_obj = get_loops.estimate_loop_segment_normals_simple(slice_obj)
            return slice_obj


class EllipseMultiple(LoopRepresentation):
    def __init__(self, planes, predicted_sequence, binary_flag_threshold=0.5, 
        segmentize_resolution=50, n_points_to_sample_for_point_cloud=64, 
        min_n_points_per_loop=16,
        run_in_fixed_resolution_polylines_mode=False,
        use_eos_token=False,
        postprocessing_heuristics=[]):
        """
        predicted_sequence: np array (n_loops i.e. n_timesteps, n_features 
            (including the binary plane-levelup flag))
        
        binary_flag_threshold = default 0.5; threshold to determine that 
            a step involves a plane level-up.
                - segmentize_resolution: to do normal estimation, we go through the step
            of sampling points from the ellipse, and then forming a closed
            polyline out of it. This parameter determines the number of points
            we should sample from the ellipse for this.
        
        segmentize_resolution: to do normal estimation, we go through the step
            of sampling points from the ellipse, and then forming a closed
            polyline out of it. This parameter determines the number of points
            we should sample from the ellipse for this.
            (No effect when run_in_fixed_resolution_polylines_mode=True. In
            that situation, the segment count is however many points the FRP
            representation is using (i.e. in most of what I've done, 32 points))
        
        n_points_to_sample_for_point_cloud: ...the above is to discretize the
            ellipse. We then have to sample FROM this discretized segments thing
            as well. This is the sampling resolution for THAT step.
            (Only significant for the .get_oriented_point_cloud() method)

        min_n_points_per_loop: the minimum number of points sampled for each
            loop; useful for making sure very small/short loops still get an 
            adequate number of sampled points for better reconstruction/viewing
            (Only significant for the .get_oriented_point_cloud() method)

        run_in_fixed_resolution_polylines_mode: since the reconstruction code
            for this EllipseMultiple representation and the fixed-resolution
            polylines representation are so similar, we can just slot in
            some special code here that handles the pred -> loops conversion
            for that representation  too. Specify True to use this class w/
            that representation. (new 2022/08/07)
        
        use_eos_token: detect the EOS vector (all-zeroes loop params but planeup
            of 1) to stop generation early rather than filling all planes
        
        postprocessing_heuristics: a list of postprocessing heuristics to apply 
            to the closed loops upon obtaining them. A list of currently
            available ones: (new 2022-08-17)

            - 'prevent-thin-loops': expand the loop points out (per closed loop)
            if they're too small, to give them more volume in the reconstruction
            - 'caps': add points to close off the first and last loops manually
            to prevent holes in the reconstruction
        """
        self.pred_seq = predicted_sequence
        self.planes = planes
        self.n_planes = len(planes)
        self.binary_flag_threshold = binary_flag_threshold
        self.segmentize_resolution = segmentize_resolution
        self.n_points_to_sample_for_point_cloud = n_points_to_sample_for_point_cloud
        self.run_in_fixed_resolution_polylines_mode = run_in_fixed_resolution_polylines_mode
        self.min_n_points_per_loop = min_n_points_per_loop
        self.use_eos_token = use_eos_token
        self.postprocessing_heuristics = postprocessing_heuristics

        self.mesh_slices = self.process_into_mesh_slices()
    
    def process_into_mesh_slices(self):
        current_plane_i = None
        increment_plane = lambda curr: (1 if curr is None else curr + 1)

        # to build up a mesh slice
        current_mesh_slice_points = None
        current_mesh_slice_conn = None
        current_mesh_slice_n_points = 0

        # list of mesh slices
        mesh_slices_formed = []
        sigmoid_np = lambda x: 1 / (1 + np.exp(-x))
        for step_i, predicted_step in enumerate(self.pred_seq):
            # 2022/06/16:the model does not convert predicted_step[-1] to
            # sigmoid, but the autoregressive inference function
            # (do_lstm_inference) handles that (so that autoregressive feeding
            # matches the training scenario.)
            
            # =================================================================
            # Step 1: process levelup flag, and move to the right plane
            # (the plane state starts with None, and goes to 0 on first levelup)
            levelup_flag = predicted_step[-1]
            guessing_sigmoid_is_needed = False
            if np.abs(levelup_flag) > 1:
                # 2022/06/16: this is probably a raw, pre-sigmoid value (because 
                # the pure non-autoregressive reconstruction scenario does not
                # go through do_lstm_inference, which applies the sigmoid
                # in the autoregressive outputs/inputs).
                # this doesn't actually matter because a positive value is 
                # obviously > 0.5 and a negative value is < 0.5 so this works
                # with the threshold even without applying sigmoid, but w/e
                levelup_flag = sigmoid_np(levelup_flag)
                guessing_sigmoid_is_needed = True
            levelup = (levelup_flag > self.binary_flag_threshold)
            thlog.debug( f"levelup flag predicted for step {step_i} is {levelup_flag}"
                + (" (auto-applied sigmoid)" if guessing_sigmoid_is_needed else ""))
            if not levelup:
                if step_i == 0:
                    thlog.debug("WARNING: The first step has a non-1 levelup "
                        "flag, which is invalid.")
                    thlog.debug("WARNING: Ignoring this, and treating it as "
                        "using the first plane until the next 1 flag..")
                    current_plane_i = 0
            else:
                thlog.debug(
                    f"Formed a mesh slice, moving up a new plane (plane "
                    f"{current_plane_i} -> {increment_plane(current_plane_i)})")
                
                current_plane_i = increment_plane(current_plane_i)

            assert current_plane_i is not None # sanity check
            if current_plane_i >= self.n_planes:
                thlog.debug(f"there are {self.n_planes} planes; stopping")
                break

            # =================================================================
            # Step 2: Process loop data into a MeshOneSlice, applying heuristics
            # on the way if there are any
            ellipse_params_polar = predicted_step[:-1]
            
            # 2022-09-26: EOS token is 'all zeros params/points, with a planeup flag of 1'
            if self.use_eos_token and current_plane_i >= round(0.6*self.n_planes):
                if self.run_in_fixed_resolution_polylines_mode:
                    # crazy heuristic: distance to centroid of loop, as well as 
                    # magnitude of coords
                    loop_pt_x = ellipse_params_polar[0::2]
                    loop_pt_y = ellipse_params_polar[1::2]
                    loop_pts = np.stack((loop_pt_x, loop_pt_y), axis=-1)
                    avg_loop_pt = np.mean(loop_pts, axis=0)
                    avg_dist_to_avg_loop_pt = np.mean(np.linalg.norm(loop_pts - avg_loop_pt))
                    under_threshold_for_EOS = np.isclose(avg_dist_to_avg_loop_pt, 0, atol=2e-2)
                        # or (np.isclose(np.mean(loop_pts), 0, atol=1e-3) )
                    if under_threshold_for_EOS:
                        thlog.info(f"EOS hit conditions: avg_dist_to_avg_loop_pt {avg_dist_to_avg_loop_pt}, mean coord {np.mean(loop_pts)}")
                else:
                    under_threshold_for_EOS = np.all(np.isclose(ellipse_params_polar, 0, atol=5e-3))


                if (under_threshold_for_EOS and levelup):
                    # EOS
                    thlog.debug(f"hit EOS token at plane #{current_plane_i}; stopping. \n {ellipse_params_polar}")
                    break


            ### convert the repr data to a points-and-conn representation that
            ### is compatible with creating a new MeshOneSlice 

            # (2022/08/07)
            # note that if self.run_in_fixed_resolution_polylines_mode is True,
            # then what is called "ellipse params" here is actually the whole
            # (n_points*2,) flattened array of coordinates, so we treat it 
            # as such.
            # As it happens, the EllipseSingle.to_mesh_slice code is also
            # highly close and adaptable to the code used for recovering a
            # MeshOneSlice from a fixed number of points. So we can just use
            # that directly and just "add a mode" to that function, like here.
            ellipse_sampled_points, ellipse_conn = \
                EllipseSingle.to_mesh_slice(
                    self.planes[current_plane_i]
                    , ellipse_params_polar
                    , return_just_data_no_slice=True
                    , segmentize_resolution=self.segmentize_resolution
                    , run_in_fixed_resolution_polylines_mode=\
                        self.run_in_fixed_resolution_polylines_mode
                    )

            # 2022-08-17: heuristic: some loops spit out by fixed-res-polyline
            # are too thin. We could apply a heuristic here that "balloons"
            # out the points from the loop's center point (avg point) to
            # give it more "volume" in the representation...
            if 'prevent-thin-loops' in self.postprocessing_heuristics:
                average_loop_point = np.mean(ellipse_sampled_points, axis=0)
                __MEANDIST_THRESHOLD = 0.08
                __EXPAND_FACTOR = 1.25
                vectors_to_average_loop_point = ellipse_sampled_points - average_loop_point
                distances_to_average_loop_point = np.linalg.norm(vectors_to_average_loop_point, axis=1)
                mean_distance_to_average_loop_point = np.mean(distances_to_average_loop_point)
                if mean_distance_to_average_loop_point == 0.:
                    continue
                if mean_distance_to_average_loop_point < __MEANDIST_THRESHOLD:
                    ellipse_sampled_points = average_loop_point + \
                        (__EXPAND_FACTOR * vectors_to_average_loop_point * \
                            np.expand_dims(( mean_distance_to_average_loop_point / distances_to_average_loop_point), axis=0).T )


            if current_mesh_slice_points is None:
                current_mesh_slice_points = ellipse_sampled_points
                current_mesh_slice_conn = ellipse_conn
                current_mesh_slice_n_points += len(ellipse_sampled_points)
            else:
                current_mesh_slice_points = np.vstack((current_mesh_slice_points, ellipse_sampled_points))
                
                # nudge the indices in ellipse_conn up to continue from the
                # current number of points in the whole MeshOneSlice object
                # we're building up, because as obtained from to_mesh_slice,
                # they only start at 0 and loop back to 0
                
                ellipse_conn = \
                    list(map(lambda conn_list: \
                        list(map(lambda conn_idx: \
                            (conn_idx + current_mesh_slice_n_points), conn_list))
                        , ellipse_conn))
                current_mesh_slice_conn  = np.vstack((current_mesh_slice_conn, ellipse_conn))
                current_mesh_slice_n_points += len(ellipse_sampled_points)

            # =================================================================
            # Step 3: Gather any pending slice-related arrays into a fully-
            # formed MeshOneSlice, if we're leveling up
            if levelup and current_mesh_slice_points is not None:
                # (the None check is to prevent mesh slices from being created
                # while we've not gathered any mesh slice points at all, 
                # since an empty mesh slice is invalid)
                mesh_slice = get_loops.MeshOneSlice(self.planes[current_plane_i], 
                    predicted_loops=(current_mesh_slice_points, current_mesh_slice_conn, None))
                mesh_slice = get_loops.estimate_loop_segment_normals_simple(mesh_slice)
                mesh_slices_formed.append(mesh_slice)
                # clear out the current_ arrays for the next meshslice-building
                current_mesh_slice_points = None
                current_mesh_slice_conn = None
                current_mesh_slice_n_points = 0


        return mesh_slices_formed
    
    def get_oriented_point_cloud(self):
        all_points, all_normals = sample_from_list_of_mesh_slices(self.mesh_slices, 
            n_points=self.n_points_to_sample_for_point_cloud, 
            distribute_points_by_loop_length=True,
            min_n_points_per_loop=self.min_n_points_per_loop,
            normal_lerping="smooth-normals" in self.postprocessing_heuristics)
        
        # the cap-fill heuristic to prevent holes in poisson reco
        __CAPS_N_BBOX_POINTS = 25
        if "caps" in self.postprocessing_heuristics:
            bottom_cap_pts, bottom_cap_normals = get_loops.sample_points_to_fill_polygons(self.mesh_slices[0], __CAPS_N_BBOX_POINTS, flip_normals=True)
            top_cap_pts, top_cap_normals = get_loops.sample_points_to_fill_polygons(self.mesh_slices[-1], __CAPS_N_BBOX_POINTS, flip_normals=False)

            # then add all this to the raw points and normals for visualization
            all_points = np.concatenate((all_points, bottom_cap_pts, top_cap_pts), axis=0)
            all_normals = np.concatenate((all_normals, bottom_cap_normals, top_cap_normals), axis=0)
            
        return all_points, all_normals

    def export_as_contour_file(self, fname: str):
        return save_contour_file_from_list_of_mesh_slices(self.mesh_slices, fname)

    def export_as_npz_file(self, fname: str):
        return save_npz_file_from_list_of_mesh_slices(self.mesh_slices, fname)





###############################################################################
# LOADING LOOP REPRESENTATIONS FROM MESH
###############################################################################


def sort_loops_by_decreasing_area(mesh_slice: get_loops.MeshOneSlice, change_to_plane_basis, plane_p_in_plane_basis):
    """ sort disjoint loops in a MeshOneSlice in order of descending area
    (so that the biggest loop comes first in the sequence, and so on)
    Modifies the mesh_slice in place but also returns back (a reference to) it.
    """
    loop_pts_in_plane_basis = \
        np.transpose(np.dot(change_to_plane_basis, np.transpose(mesh_slice.loop_pts))) - plane_p_in_plane_basis

    def __calculate_area(loop_conn_arr):
        loop_pts_in_plane_basis_this_loop = loop_pts_in_plane_basis[loop_conn_arr[:, 0]]
        return geometry_utils.polygon_area(loop_pts_in_plane_basis_this_loop[:, 0], loop_pts_in_plane_basis_this_loop[:, 1])

    sorted_indices = sorted(list(range((len(mesh_slice.disjoint_loop_conns)))), 
        key=lambda i: __calculate_area(mesh_slice.disjoint_loop_conns[i]), reverse=True)

    mesh_slice.disjoint_loop_conns = [mesh_slice.disjoint_loop_conns[i] for i in sorted_indices]
    # mesh_slice.disjoint_loop_lengths = [mesh_slice.disjoint_loop_lengths[i] for i in sorted_indices]
    return mesh_slice


def fixed_resolution_polylines_disjoint_loops(mesh_slice : get_loops.MeshOneSlice,
    fixed_size_sampling_howmany_points: int):
    # (2022/08/07)
    """
    Gets the representation data for the "fixed-resolution polylines WITH 
    levelup flags" representation. This is like the ellipse-parameters representation,
    but instead of the 5 ellipse parameters, it would be however many x, y coords
    for a fixed number of points sampled each polyline. This would allow us to
    make arbitrary loop shapes (albeit restricted to a fixed resolution, but
    that's fine). The order of points in this fixed-size array of points would
    be hopefully such that the points go in clockwise order starting from the
    point with the least (x+y) coordinate sum.

    returns repr_points, repr_data where
    - repr_points is (n_points_total, 3), not separated by disjjoint loops
    - repr_data is of shape
        (n_disjoint_loops_this_mesh_slice, 2* fixed_size_sampling_howmany_points + 1)
        where the last number is the level-up binary flag.
    """
    plane = mesh_slice.plane
    plane_n, plane_p, onplane_xaxis, onplane_yaxis, plane_basis, \
        change_to_plane_basis, plane_p_in_plane_basis = \
            geometry_utils.extract_useful_data_from_plane(plane)
    
    # new 2022-12-20; incorporating this from loopfields idea. Sort loops in
    # decreasing area
    mesh_slice = sort_loops_by_decreasing_area(mesh_slice, 
        change_to_plane_basis, plane_p_in_plane_basis)


    # this sampling is sampling the segments. For this representation, we also
    # use this for the actual representation data, which needs a fixed n points
    # for each disjoint loop, hence distribute_points_by_loop_length=False.
    sampled_points_by_disjoint_loops, sampled_normals_by_disjoint_loops = \
        get_loops.sample_from_loops_and_normals(mesh_slice, 
            fixed_size_sampling_howmany_points, 
            separate_by_disjoint_loops=True, 
            distribute_points_by_loop_length=False)
    
    # we need this for the first return item (for passing to polyscope viz etc)
    sampled_points_NOT_separated_by_disjoint_loops = \
        np.concatenate(sampled_points_by_disjoint_loops)

    # this will eventually be the actual representation data array
    arranged_flattened_point_data = []
    
    for sampled_points, sampled_normals in \
        zip(sampled_points_by_disjoint_loops, sampled_normals_by_disjoint_loops):

        sampled_points_2d_coords_on_plane = \
            np.zeros((sampled_points.shape[0], 2), dtype=sampled_points.dtype)

        for i, point in enumerate(sampled_points):
            # get the projection of the sampled point on the plane, as a 2d coordinate
            # Apply the change-of-basis matrix, and grab the x and y coords only
            # (throw away the z coordinate, which should be 0)
            point_in_plane_basis = \
                np.dot(change_to_plane_basis, point) - plane_p_in_plane_basis
            assert np.isclose(point_in_plane_basis[2], 0, atol=1e-5), \
                (point, point_in_plane_basis, plane_p, plane_n, change_to_plane_basis)
            sampled_points_2d_coords_on_plane[i] = point_in_plane_basis[:2]

        # for this representation, we just take the
        # sampled_points_2d_coords_on_plane, and do the following heuristics
        # to hopefully standardize the ordering of the array of points a
        # bit. 
        # 1. rotate the array of points so that the point with the least
        # (x+y) comes first.
        # 2. check the orientation of the 0th, 1st, and 2nd points. If they
        # are counterclockwise, then reverse point_data[1:] so that
        # hopefully the point sequence goes clockwise. (of course we won't
        # be able to enforce this for the whole loop because the orientation
        # may change (concave loops) but at least the start should be
        # standard-ish... 

        # 1. Rotate so that the first point is the one with least (x+y)
        point_idx_with_least_x_plus_y = np.argmin(np.sum(sampled_points_2d_coords_on_plane, axis=-1))
        point_data = np.concatenate((
            sampled_points_2d_coords_on_plane[point_idx_with_least_x_plus_y:]
            , sampled_points_2d_coords_on_plane[:point_idx_with_least_x_plus_y]
        ))

        #2. check orientation of points 0, 1, 2
        first_three_pts_orientation = \
            geometry_utils.orientation_of_3pts_in_2d(
                point_data[0], point_data[1], point_data[2])
        if first_three_pts_orientation < 0: # if CCW
            # reverse point_data[1:] in-place so that the first three points
            # should now be CW (and hopefully the rest should also now be CW)
            point_data[1:] = (point_data[1:])[::-1]

        # then we have to flatten this out (to make this look like 'fixed
        # set of ellipse params', but in actuality 2*n_points numbers, 
        # which are x and y coords alternating)
        point_data = point_data.flatten()

        # append an empty 0 for the "plane-level-up" binary flag (see below)
        point_data = np.append(point_data, 0)
        arranged_flattened_point_data.append(point_data)
    
    # mark the first loop on this mesh slice with a flag (1) indicating "new
    # level" and the rest of the loops with a 0 ("stay on this plane level")
    arranged_flattened_point_data[0][-1] = 1

    return sampled_points_NOT_separated_by_disjoint_loops, \
        np.array(arranged_flattened_point_data)



def many_ellipses_disjoint_loops(mesh_slice  : get_loops.MeshOneSlice):
    """ 
    From one mesh slice, obtain the paramters for many ellipses that closely
    matches the disjoint loops present on the plane level. Like single_ellipse,
    each ellipse involves a center point, and two radii (major and minor),
    and the angle offset from horizontal. However, with this representation, each
    ellipse parameter set also comes with a flag that says whether that ellipse
    is the "first" loop on this plane level. The flag is 1 for the first loop
    on a plane level, and is 0 for the rest of the loops on that same plane
    level.

    Returns the sampled points from the fit ellipses (combined into one array)
    (n_points_total, 3) and the list of ellipse coefficients (with the flag as
    described above), of shape (n_disjoint_loops_this_mesh_slice, 6)
    """

    plane = mesh_slice.plane
    plane_n, plane_p, onplane_xaxis, onplane_yaxis, plane_basis, \
        change_to_plane_basis, plane_p_in_plane_basis = \
            geometry_utils.extract_useful_data_from_plane(plane)


    # this sampling is sampling the segments in order to fit ellipses
    sampled_points_by_disjoint_loops, sampled_normals_by_disjoint_loops = \
        get_loops.sample_from_loops_and_normals(mesh_slice, 50, separate_by_disjoint_loops=True)
    
    ellipse_polar_paramses = []
    ellipse_sampled_pointses = None
    
    for sampled_points, sampled_normals in \
        zip(sampled_points_by_disjoint_loops, sampled_normals_by_disjoint_loops):

            sampled_points_2d_coords_on_plane = \
                np.zeros((sampled_points.shape[0], 2), dtype=sampled_points.dtype)

            for i, point in enumerate(sampled_points):
                # get the projection of the sampled point on the plane, as a 2d coordinate
                # Apply the change-of-basis matrix, and grab the x and y coords only
                # (throw away the z coordinate, which should be 0)
                point_in_plane_basis = \
                    np.dot(change_to_plane_basis, point) - plane_p_in_plane_basis
                assert np.isclose(point_in_plane_basis[2], 0, atol=1e-5), \
                    (point_in_plane_basis, plane_p)
                sampled_points_2d_coords_on_plane[i] = point_in_plane_basis[:2]

            # feed into the fit function, to obtain fits
            ### old fitting code:
            # ellipse_coefficients = geometry_utils.ellipse_fit_to_points(sampled_points_2d_coords_on_plane)
            # ellipse_params_polar = geometry_utils.ellipse_cartesian2polar(ellipse_coefficients)
            # ellipse_sampled_points = geometry_utils.ellipse_sample_points(ellipse_params_polar)
            
            ### new fitting code (SVD, from Haochen)
            try:
                ellipse_params_polar = \
                    np.array(geometry_utils.svd_fit_ellipse(sampled_points_2d_coords_on_plane))
            except AssertionError as e:
                thlog.err("Caught nonfatal error in ellipse fit, skipping this loop \n"
                f"(extra info: current plane_p = {plane_p})")
                continue

            ellipse_sampled_points = geometry_utils.sample_from_ellipse(100, ellipse_params_polar)
            ### end fitting code

            # tack on the z-coordinate of 0 again to each point
            ellipse_sampled_points = np.pad(ellipse_sampled_points, ((0,0),(0,1)))

            # convert back into world coordinates by multiplying the plane_basis
            # with each found-point...
            ellipse_sampled_points_in_world_coords = \
                np.transpose(np.dot(plane_basis, np.transpose(ellipse_sampled_points)))
            
            # and add back the plane_p (i.e. translate back to world coords origin)
            ellipse_sampled_points_in_world_coords += plane_p

            # then sample points from the fit-ellipse, and put on polyscope

            if ellipse_sampled_pointses is None:
                ellipse_sampled_pointses = ellipse_sampled_points_in_world_coords
            else:
                ellipse_sampled_pointses = np.vstack((
                    ellipse_sampled_pointses, 
                    ellipse_sampled_points_in_world_coords))

            # append an empty 0 for the "plane-level-up" binary flag (see below)
            ellipse_params_polar = np.append(ellipse_params_polar, 0)
            ellipse_polar_paramses.append(ellipse_params_polar)
    
    # mark the first loop on this mesh slice with a flag (1) indicating "new
    # level" and the rest of the loops with a 0 ("stay on this plane level")
    ellipse_polar_paramses[0][-1] = 1
    # np.array(ellipse_polar_paramses) should now have shape
    # (n_disjoint_loops_on_this_plane_level, n_polar_params + 1)
    # n_polar_params I think is 6? it has eccentricity included, but ecc. is
    # not strictly necesssary to fully draw the ellipse being specified, but
    # whatever...
    return ellipse_sampled_pointses, np.array(ellipse_polar_paramses)





def single_ellipse_no_disjoint_loops(mesh_slice : get_loops.MeshOneSlice):
    """ 
    From one mesh slice, obtain the paramters for an ellipse that closely
    matches the actual loop. The simplest loop representation for learning, as
    an ellipse only involves a center point, and two radii (major and minor),
    and the angle offset from horizontal.

    Returns the sampled points from the fit ellipse (n_points,3), and the set of
    ellipse coefficients (n_disjoint_loops_this_mesh_slice, 5)
    """


    # obtain coordinates of these points with respect to the slice plane (i.e.
    # such that we have effectively 2D coordinates)
    plane = mesh_slice.plane
    plane_n, plane_p, onplane_xaxis, onplane_yaxis, plane_basis, \
        change_to_plane_basis, plane_p_in_plane_basis = \
            geometry_utils.extract_useful_data_from_plane(plane)

    # this sampling is sampling the segments to make points to fit ellipses to
    sampled_points, sampled_normals = \
        get_loops.sample_from_loops_and_normals(mesh_slice, 50)

    sampled_points_2d_coords_on_plane = \
        np.zeros((sampled_points.shape[0], 2), dtype=sampled_points.dtype)
    for i, point in enumerate(sampled_points):
        # get the projection of the sampled point on the plane, as a 2d coordinate
        # Apply the change-of-basis matrix, and grab the x and y coords only
        # (throw away the z coordinate, which should be 0)
        point_in_plane_basis = \
            np.dot(change_to_plane_basis, point) - plane_p_in_plane_basis
        assert np.isclose(point_in_plane_basis[2], 0, atol=1e-5), \
            (point_in_plane_basis, plane_p)
        sampled_points_2d_coords_on_plane[i] = point_in_plane_basis[:2]

    # feed into the fit function, to obtain fits
    ### old fitting code:
    # ellipse_coefficients = geometry_utils.ellipse_fit_to_points(sampled_points_2d_coords_on_plane)
    # ellipse_params_polar = geometry_utils.ellipse_cartesian2polar(ellipse_coefficients)
    # ellipse_sampled_points = geometry_utils.ellipse_sample_points(ellipse_params_polar)
    
    ### new fitting code (SVD, from Haochen)
    ellipse_params_polar = np.array(geometry_utils.svd_fit_ellipse(sampled_points_2d_coords_on_plane))
    ellipse_sampled_points = geometry_utils.sample_from_ellipse(100, ellipse_params_polar)
    ### end fitting code

    # tack on the z-coordinate of 0 again to each point
    ellipse_sampled_points = np.pad(ellipse_sampled_points, ((0,0),(0,1)))

    # convert back into world coordinates by multiplying the plane_basis with
    # each found-point...
    ellipse_sampled_points_in_world_coords = \
        np.transpose(np.dot(plane_basis, np.transpose(ellipse_sampled_points)))
    
    # and add back the plane_p (i.e. translate back to world coords origin)
    ellipse_sampled_points_in_world_coords += plane_p

    # then sample points from the fit-ellipse, and put on polyscope

    # once that works, we'll use this as training data (and the parameters
    # of the ellipse will form the "targets" for a model to learn)
    return ellipse_sampled_points_in_world_coords, ellipse_params_polar


# Constants for loop representation types
LOOP_REPR_ELLIPSE_SINGLE = 0
LOOP_REPR_ELLIPSE_MULTIPLE = 1
LOOP_REPR_CLOSED_POLYLINE_STROKES = 2  # not supported! 
LOOP_REPR_FIXED_RES_POLYLINE = 3


def loop_repr_uses_binary_levelup_flags(loop_repr_type: int) -> bool:
    """
    A function that some parts of the MeshOneSlice recovery process / inference
    may rely on to get the right outputs (i.e. applying sigmoid after the model)
    depending on the loop representation used. The point of putting this into
    its own function is that many other representations may need this too. 
    """
    return (loop_repr_type == LOOP_REPR_ELLIPSE_MULTIPLE) or \
           (loop_repr_type == LOOP_REPR_FIXED_RES_POLYLINE)
        

def load_loop_repr_from_mesh(fname, planes, reprtype, repr_options=None,
    throw_on_bad_slice_plane=False, append_endofsequence_timestep=False):
    """
    Given a mesh file, and some planes (each specifying a point-on-plane (index
    0), a plane normal (1), and two vectors (index 2, 3) specifying and x+ and
    y+ axis of a coord system on the plane with origin as the provided
    point-on-plane), generate the coefficients/data array of some given
    mesh representation usable as ground truth data.

    current supproted representations:
    - LOOP_REPR_ELLIPSE_SINGLE (0)
        each mesh slice/plane level only has one ellipse.
    - LOOP_REPR_ELLIPSE_MULTIPLE (1)
        each mesh slice/plane level may have multiple ellipses fit to the
        multiple disjoint loops present on the same plane level.
        In the sequence data representation, the first loop of each level gets a
        "new plane level"  flag set to 1, and subsequent loops on the same level
        get a 0 for that flag; the flag stays at 0 until the next level starts.
        Return shape for one mesh: (n_disjoint_loops_across_mesh, 5 ellipse params + 1 planeup flag)
    - LOOP_REPR_CLOSED_POLYLINE_STROKES (2)
        (unsupported; formerly a two-tier RNN representation but this was
        never implemented.)
    - LOOP_REPR_FIXED_RES_POLYLINE:
        uses the binary levelup flag (in fact most of the same machinery, and
        the same model architecture) as LOOP_REPR_ELLIPSE_MULTIPLE, except
        instead of using 5 ellipse parameters as the repr data, this takes
        repr_options as the number of points to sample from each disjoint loop
        and simply encodes the raw point data as teh repr data. 
        This point data is arranged in an order that indicates the sequence of
        segments to draw, too. There are some heuristics to make this ordering
        more consistent across disjoint loops, namely rotating the array so that
        the topmost leftmost point is at the start, and also reversing the tail
        of the array such that the first three points (and by extension
        hopefully the rest) of the array "walks in a CLOCKWISE orientation".
        Return shape (n_disjoint_loops_across_mesh (=n_timesteps), 
            fixed_size_sampling_howmany_points * 2 + 1 planeup flag)

    other optional arguments: 
    - throw_on_bad_slice_plane: bool (default False): if True, throw
        an exception when a plane does NOT successfully slice the mesh for
        whatever reason (for instance, most often due to not intersecting the
        mesh at all.) If False, safely ignore this error and simply not include
        that plane's slice in the representation. (Specify True in order to make
        sure that all meshes loadable have the same amount of planes cutting it)

        NOTE we COULD also make the "nonfatal error" in the SVD ellipse fit
        throw (currently it is caught in many_ellipses_disjoint_loops and
        considered nonfatal as the loop would just be skipped) and it would be
        caught and re-thrown here. But losing a loop is not as bad as losing
        and entire plane so maybe we shouldn't do that.
    - append_endofsequence_timestep: bool (default False): if True, there will
        be one last timestep at the end with all-zeroes coordinates and a plane-
        up flag of 1 to indicate an "end-of-sequence token".
    """
    mesh_dir= os.path.dirname(fname)
    mesh_name = os.path.basename(fname)
    # if thlog.guard(VIZ_INFO, needs_polyscope=True):
    #     polysoup = utils.PolygonSoup.from_obj(fname)
    #     ps.register_surface_mesh("gt", polysoup.vertices, polysoup.indices, color=utils.PS_COLOR_SURFACE_MESH)
    
    all_points = None
    all_normals = None
    all_repr_data = None # coefficients, stroke data, etc etc

    all_ellipse_points = None
    for viz_i, plane in enumerate(planes):
        try:
            mesh_slice = get_loops.MeshOneSlice(plane, mesh_fname=fname)
        except Exception as e:
            thlog.err(f"[load_loop_repr_from_mesh] error at slice plane index {viz_i}:\n{repr(e)}")
            if throw_on_bad_slice_plane:
                thlog.err(f"[load_loop_repr_from_mesh] bad planes are set to be unacceptable. Throwing!")
                raise e # and don't continue to next planes
            continue
        thlog(LOG_TRACE, "final loop connectivity:\n {}".format(mesh_slice.loop_conn_arr))
        if len(mesh_slice.loop_conn_arr) == 0:
            continue
        thlog.do(VIZ_INFO, 
            lambda: ps.register_curve_network(f"cunet{viz_i}", 
                mesh_slice.loop_pts, mesh_slice.loop_conn_arr, enabled=True, color=utils.PS_COLOR_CURVE_NETWORK)
            , needs_polyscope=True)
    
        # then, try estimating the normals just from the plane normal
        mesh_slice = get_loops.estimate_loop_segment_normals_simple(mesh_slice)
        if thlog.guard(VIZ_DEBUG, needs_polyscope=True):
            ps.get_curve_network(f"cunet{viz_i}").add_vector_quantity(
                f"simple_norm_est{viz_i}", 
                mesh_slice.loop_segment_normals,
                defined_on="nodes", enabled=False, length=0.1)

        # this sampling is sampling the segments from the real-mesh-slice
        # (and not sampling from ellipses/our own representation)
        # (only for viewing)
        sampled_points, sampled_normals = \
            get_loops.sample_from_loops_and_normals(mesh_slice, 50)
        if thlog.guard(VIZ_DEBUG, needs_polyscope=True):
            pcloud = ps.register_point_cloud(f"sampled_pc{viz_i}", sampled_points)
            pcloud.add_vector_quantity(f"sampled_pc_norms{viz_i}", sampled_normals, 
                enabled=False, length=0.11)


        ########### LOOP REPRESENTATIONS TESTING HERE 
        
        if reprtype == LOOP_REPR_ELLIPSE_SINGLE:
            #1. Ellipse fit to the loop (bad, but basic and easy to predict  )
            sampled_ellipse_points, ellipse_polar_coefficients = \
                single_ellipse_no_disjoint_loops(mesh_slice)

            # all loop representations should fill up a "repr points" array
            # and a "repr_data" array
            # that represents the sampled points using some loop representation/
            # parameterization technique, as well as the sequence data to learn
            repr_points = sampled_ellipse_points
            repr_data = ellipse_polar_coefficients

        elif reprtype == LOOP_REPR_ELLIPSE_MULTIPLE:
            #2. Ellipse fit to multiple disjoint loops each plane level
            # (same as above, but now supports multiple ellipses per plane)
            sampled_ellipse_points, ellipse_polar_coefficientses = \
                many_ellipses_disjoint_loops(mesh_slice)
            repr_points = sampled_ellipse_points
            repr_data = ellipse_polar_coefficientses 
            # ^^^ shape (n_disjoint_loops_on_this_slice, n_polar_params+1)
            # when this repr_data is vstacked with the rest, we'll get a final
            # all_repr_data array of shape (n_disjoint_loops_across_whole_mesh, n_polar_params+1)

        elif reprtype == LOOP_REPR_FIXED_RES_POLYLINE:
            #2.5. uses levleup flags like EllipseMultiple but rather than encoding
            # implicit ellipse parameters, this repr encodes a fixed number of
            # points that are directly interpreted as the closed loop's nodes.
            # the order of these points in the predicted vector is taken as
            # the order to draw the segments in, too.
            fixed_size_sampling_howmany_points = 32 if repr_options is None else repr_options
            assert isinstance(fixed_size_sampling_howmany_points, int), \
                "repr_option for this loop representation should be of type int"
            sampled_polyline_points, arranged_flattened_point_data = \
                fixed_resolution_polylines_disjoint_loops(mesh_slice, \
                    fixed_size_sampling_howmany_points)
            
            repr_points = sampled_polyline_points  # for viewing, not for data
            repr_data = arranged_flattened_point_data 
            # ^^ shape (n_disjoint_loops_this_slice, fixed_size_sampling_howmany_points*2+1)
            


        elif reprtype == LOOP_REPR_CLOSED_POLYLINE_STROKES:
            raise NotImplementedError()
            
        ####################################################



        # accumulate all the sub-pointclouds sampled from each slice into
        # a big oriented point cloud we can then later give to the poisson
        # reconstructor
        if all_points is None:
            all_points = sampled_points #  points from slicing, not from the ellipse-fit
            all_normals = sampled_normals
            all_repr_points = repr_points
            all_repr_data = repr_data
        else:
            all_points = np.vstack((all_points, sampled_points))
            all_normals = np.vstack((all_normals, sampled_normals))
            all_repr_points = np.vstack((all_repr_points, repr_points))
            # all_repr_data = np.vstack((all_repr_data, repr_data))
            all_repr_data = utils.vstack_with_padding((all_repr_data, repr_data))
    
    # 2022-09-26: add an EOS timestep, which by convention we set to consist
    # of all zeros but a levelup flag of 1. Only applicable for the reprs with 
    # the binary up flag
    if append_endofsequence_timestep and loop_repr_uses_binary_levelup_flags(reprtype):
        # all_repr_data is of shape (n_timesteps, n_features)
        last_EOS_slice_data = np.zeros_like(all_repr_data[0])
        last_EOS_slice_data[-1] = 1 # "binary flag of 1"
        all_repr_data = np.vstack((all_repr_data, np.expand_dims(last_EOS_slice_data, axis=0)))

    # normalize all the normal vectors for output
    all_normals /= np.transpose(np.tile(np.linalg.norm(all_normals, axis=1), (3, 1)))
    if thlog.guard(VIZ_INFO, needs_polyscope=True):     
        # utils.write_oriented_pcloud_to_plyfile(
        #     "scratch-output/" + mesh_name + "-OrPtCloud.ply", mesh_dir + mesh_name, all_points, all_normals)
        ps.reset_camera_to_home_view()
        ps.show()

        # no need for reco here.... or do we??????????
        # reco_vertices, reco_faces = meshlab_poisson_reco.do_poisson_reco(
        #     all_points, all_normals)
        # ps.register_surface_mesh("reconstructed", reco_vertices, reco_faces)

        # flatten all_repr_points in case this is the return from the 
        # LOOP_REPR_CLOSED_POLYLINE_STROKES representation
        all_repr_points_shape = all_repr_points.shape
        all_repr_points_to_show = all_repr_points if len(all_repr_points_shape) == 2 else \
            np.reshape(all_repr_points, (all_repr_points_shape[0] * all_repr_points_shape[1], 3))
        ps.register_point_cloud("ellipses", all_repr_points_to_show)
        ps.show()
    
    # only all_repr_data is actually used in training.. the rest are for 
    # visualizations of the slicing process. 
    # because the representation may not simply be directly resampled source
    # slice segments, "all_repr_points" is a point sampling of that
    # representation (i.e. points sampled from fit curves, ellipses).
    # but in the Fixed Res Polyline case, it's no different from all_points.
    return all_repr_data, all_points, all_normals, all_repr_points



def get_nice_set_of_slice_planes(
    num: int, min_coord:float=-0.5, max_coord:float=0.3, 
    scan_in_direction_xyz=1):
    planes = []
    scan_in_direction= scan_in_direction_xyz # 0 = x, 1 = y, 2 = z
    # for this set,
    # cello-125f.obj, celloshort-125f.obj,  coney-125f.obj coneypointy-125f.obj,
    # hglassblob-125f.obj hglassblob2-125f.obj hglassblob2flat-125f.obj
    # hglassblobflatter-125f.obj wormy-125f.obj wormythick-125f.obj
    # the following min and max values result in a full set of slices.

    # (first experiments single-ellipse-per-plane used min=-0.5, max=0.3)
    for viz_z in np.linspace(min_coord, max_coord, num=num):
        plane_p = [0,0,0]
        plane_p[scan_in_direction] = viz_z
        plane_n = [0,0,0]
        plane_n[scan_in_direction] = 1
        # select a 2D coordinate system on the plane
        # these lines here are just convenience for testing. The only constraint
        # here is that the onplane x and y basis vectors should properly form
        # a basis with the plane normal as the z direction.
        if scan_in_direction == 0:
            # it turns out that [0,2,1] is the right orientation for the other 2
            # axes when the scan axis is x
            remaining_directions = [0, 2, 1]
        else:
            remaining_directions = [0, 1, 2]
        
        remaining_directions.remove(scan_in_direction)
        
        onplane_xaxis = [0,0,0]
        onplane_xaxis[remaining_directions[0]] = 1
        remaining_directions.pop(0)
        onplane_yaxis = [0,0,0]
        onplane_yaxis[remaining_directions[0]] = 1
        # okay finished with selecting a x, y coordinate system on the plane
        plane = np.array([plane_n,plane_p, onplane_xaxis, onplane_yaxis])
        planes.append(plane)
    return planes


# these test main functions are  not very nicely written because they're for
# testing/playground purposes only; lots of duplicate code abound...
def loop_reprs_test_main():
    """
    requires pytorch3d
    usage: 
    python loop_representations.py <planespec file> folder|file <dir or single obj file path>
    """
    import torch
    import pytorch3d.loss
    import pytorch3d.ops
    import pytorch3d.io
    import pytorch3d.structures
    import glob

    def normalize_to_centered_cube_inplace(meshes: pytorch3d.structures.Meshes, cube_side_length: float=1.0):
        """
        normalize to a cube of side length cube_side_length centered at origin
        """
        bounding_boxes = meshes.get_bounding_boxes()  # (n_meshes, 3, 2)
        mesh_to_verts_packed_first_idx = meshes.mesh_to_verts_packed_first_idx()
        
        bounding_boxes_packed = bounding_boxes[meshes.verts_packed_to_mesh_idx()]  # (sum of all vertex counts, 3, 2)
        min_coords_packed = bounding_boxes_packed[:, :, 0]
        max_coords_packed = bounding_boxes_packed[:, :, 1]
        extent_packed, _ = (max_coords_packed - min_coords_packed).max(dim=-1)
        
        scale_per_mesh = cube_side_length / extent_packed[mesh_to_verts_packed_first_idx]  # (n_meshes, )
        center_packed = (min_coords_packed + max_coords_packed) / 2       # (sum of all vertex counts, 3)

        # normalize in-place
        meshes.offset_verts_(-center_packed)
        meshes.scale_verts_(scale_per_mesh)


    loop_repr_to_use = LOOP_REPR_FIXED_RES_POLYLINE
    throw_on_fail = True
    do_poisson_reconstruction = True
    throw_on_bad_slice_plane = False
    use_eos_token = True
    thlog.set_levels(LOG_INFO, VIZ_DEBUG)
    
    reconstructed_save_dir = None

    thlog.init_polyscope()
    # first argument is the planespec file
    scan_dir, num_slices, min_coord, max_coord = utils.read_planespec_file(sys.argv[1])
    planes = get_nice_set_of_slice_planes(num_slices, min_coord=min_coord, max_coord=max_coord, scan_in_direction_xyz=scan_dir)
    # bad_meshes = []
    good_meshes = []

    mode = sys.argv[2]
    if mode == "folder":
        fnames = sorted(glob.glob(os.path.join(sys.argv[3], "*.obj")))
        n_shapes = int(sys.argv[4])
    else:
        fnames = [sys.argv[3]]
        n_shapes = 1

    
    if n_shapes > 0:
        fnames = fnames[:n_shapes]
    else:
        # specify n_shapes argument to be negative to use all obj files in the
        # specified dir in sys.argv[2]
        n_shapes = len(fnames)
    
    
    pytorch3d_gt_meshes = pytorch3d.io.load_objs_as_meshes(fnames, load_textures=False)
    normalize_to_centered_cube_inplace(pytorch3d_gt_meshes)

    chamfer_losses = np.zeros(n_shapes, dtype=float)
    good_meshes_mask = np.zeros(n_shapes, dtype=bool)
    for i, (fname, pytorch3d_gt_mesh) in enumerate(zip(fnames, pytorch3d_gt_meshes)):
        ps.remove_all_structures()
        mesh_dir = os.path.dirname(fname)
        mesh_name = os.path.basename(fname)
        thlog.info(f"loading mesh {mesh_name}")
        
        # now we make the mesh_slice and stuff from there    
        try:
            repr_data, slice_points, slice_normals, repr_sampled_points = \
                load_loop_repr_from_mesh(fname, planes, loop_repr_to_use, 
                repr_options=32, throw_on_bad_slice_plane=throw_on_bad_slice_plane, append_endofsequence_timestep=use_eos_token)
        except Exception as e:
            if isinstance(e, KeyboardInterrupt):
                raise KeyboardInterrupt()
            if not throw_on_fail:
                thlog.info(f"mesh {mesh_name} failed to load with error\n{e.args}\n---continuing..")
                continue
            else:
                thlog.err(f"mesh {mesh_name} failed to load. Here is the exception: \n")
                raise e
        good_meshes.append(mesh_name)
        thlog.debug(f"repr data: first and rest \n {repr_data[0]} all: \n {repr_data}")
        thlog.info(f"{mesh_name}: There are {len(repr_data)} time steps in the loop sequence data; loop data has shape {repr_data.shape}")
        
        if do_poisson_reconstruction:
            # this is going thru the machinery in the LoopRepresentation classes
            # (EllipseMultiple etc) above, same as what is run at inference to
            # recover point cloud from loops.
            # specifically, from the 32-segment loop representation, we
            # - distribute (by loop length) 48 points across all loops on each slice
            # - any loop that has <24 points assigned by above step gets 24 instead
            # - normals are estimated using the original 32-segment representation
            # and then smoothed.
            # - cap points are added for closed reconstruction
            # - very small loops are made into a circle
            loop_repr_object = EllipseMultiple(planes, repr_data,
                segmentize_resolution=64, n_points_to_sample_for_point_cloud=48,
                run_in_fixed_resolution_polylines_mode=True,
                min_n_points_per_loop=24,
                use_eos_token = use_eos_token,
                postprocessing_heuristics=
                    [ 'smooth-normals']
                )
            all_points, all_normals = loop_repr_object.get_oriented_point_cloud()
            reco_vertices, reco_faces = meshlab_poisson_reco.do_poisson_reco(all_points, all_normals, save_filename=None if not reconstructed_save_dir else os.path.join(reconstructed_save_dir, mesh_name + ".obj"))
            
            # bring this to pytorch3d and compare chamfer
            pytorch3d_reco_mesh = pytorch3d.structures.Meshes([torch.from_numpy(reco_vertices).float()], [torch.from_numpy(reco_faces)])
            normalize_to_centered_cube_inplace(pytorch3d_reco_mesh)
            if thlog.guard(VIZ_DEBUG, needs_polyscope=True):
                ps.register_surface_mesh("reconstructed", pytorch3d_reco_mesh.verts_packed().cpu().detach().double().numpy(), pytorch3d_reco_mesh.faces_packed().cpu().detach().numpy(), color=utils.PS_COLOR_SURFACE_MESH)
                ps.register_surface_mesh("gt", pytorch3d_gt_mesh.verts_packed().cpu().detach().double().numpy(), pytorch3d_gt_mesh.faces_packed().cpu().detach().numpy())
                ps.show()
                for viz_i in range(132): # arbitrary number.... for viz purposes
                    try:
                        ps.get_curve_network(f"cunet{viz_i}").set_enabled(False)
                    except:
                        pass
                    try:
                        ps.get_point_cloud(f"ellipses").set_enabled(False)
                    except:
                        pass
                ps.show()


            gt_sampled_pcloud = pytorch3d.ops.sample_points_from_meshes(pytorch3d_gt_mesh)
            reco_sampled_pcloud = pytorch3d.ops.sample_points_from_meshes(pytorch3d_reco_mesh)
            chamfer_loss, _ = pytorch3d.loss.chamfer_distance(gt_sampled_pcloud, reco_sampled_pcloud)
            chamfer_loss_val = chamfer_loss.item()
            thlog.info(f"mesh {os.path.basename(fname)}: {chamfer_loss_val}")
            chamfer_losses[i] = chamfer_loss_val
            good_meshes_mask[i] = True


    thlog.info(f"good meshes: ")
    for name in good_meshes:
        print(name)    
    
    thlog.info(f"avg reconstruction loss (gt loops, poisson reco, vs gt mesh) chamfer distance: {chamfer_losses[good_meshes_mask].mean()} (avg out of {good_meshes_mask.sum()} good sliceable meshes)")

if __name__ ==  "__main__":
    loop_reprs_test_main()