import sys
import os

import argparse
import pandas as pd
import numpy as np
import igl

def halfedges(f):
    assert f.shape[0] > 0
    assert f.shape[1] == 3

    return np.block([[f[:,1][:,None], f[:,2][:,None]],
        [f[:,2][:,None], f[:,0][:,None]],
        [f[:,0][:,None], f[:,1][:,None]]])

def edges(f,
    return_boundary_indices=False,
    return_interior_indices=False,
    return_nonmanifold_indices=False):
    assert f.shape[0] > 0
    assert f.shape[1] == 3

    #Sort halfedges. Remove duplicates.
    he = halfedges(f)
    sorted_he = np.sort(he, axis=1)
    unique_hes, unique_indices, unique_count = np.unique(sorted_he, axis=0,
        return_index=True, return_counts=True)

    #Construct edge arrays by preserving the unique indices of boundary and
    # picking sorted orientation for interior edges.
    #Boundary edges have only one halfedge.
    bdry_edges = he[unique_indices[unique_count==1],:]
    #Interior edges have two ore more halfedges.
    interior_edges = unique_hes[unique_count>1,:]
    edges = np.block([[bdry_edges],[interior_edges]])
    assert edges.shape == unique_hes.shape

    if return_boundary_indices or return_interior_indices or return_nonmanifold_indices:
        retval = [edges]
        if return_boundary_indices:
            retval.append(np.arange(0, bdry_edges.shape[0]))
        if return_interior_indices:
            retval.append(np.where(unique_count>1))
        if return_nonmanifold_indices:
            #Nonmanifold edges have three or more halfedges.
            retval.append(np.where(unique_count>2))
        return retval
    else:
        return edges

def boundary_edges(f):
    E,b = edges(f, return_boundary_indices=True)
    return E[b,:]

def boundary_loops(f, allow_wrong_orientations=True):
    assert f.shape[0] > 0
    assert f.shape[1] == 3

    bE = boundary_edges(f)

    #Loop through each boundary, edge, marking them as seen, until all have
    # been seen.
    unseen = np.full(bE.shape[0], True)
    which_loop = []
    loopind = -1
    loop_vertices = []
    while np.any(unseen):
        current_b = np.argmax(unseen)
        current_bE = bE[current_b,:]
        start = current_bE[0]
        unseen[current_b] = False

        loopind += 1
        which_loop.append(loopind)
        loop_vertices.append(start)

        head = current_bE[1]
        while head != start:
            which_loop.append(loopind)
            loop_vertices.append(head)

            if allow_wrong_orientations:
                current_b_0 = np.where((bE[:,0] == head) & unseen)[0]
                current_b_1 = np.where((bE[:,1] == head) & unseen)[0]
                if len(current_b_0)>len(current_b_1):
                    current_b = current_b_0
                    head_ind = 1
                else:
                    current_b = current_b_1
                    head_ind = 0
            else:
                current_b = np.where((bE[:,0] == head) & unseen)[0]
                head_ind = 1
            assert len(current_b) == 1
            current_b = current_b.item(0)
            current_bE = bE[current_b,:]
            unseen[current_b] = False
            head = current_bE[head_ind]

    assert len(loop_vertices) == len(which_loop)
    assert len(loop_vertices) == bE.shape[0]
    return np.array(loop_vertices), np.array(which_loop)

def count_boundary_segments(f):
    _,which_loop = boundary_loops(f)
    if len(which_loop>0):
        return np.amax(which_loop)+1
    else:
        return 0

def get_dataset_characteristics(dataset_folder, output_file):
    df_columns = ["Cut", "Filename", "Object Number", "Mesh Name", "Chart Number", \
                  "Vertices", "Faces", "Manifold", \
                  "Genus", "Boundary Segments", \
                  "Boundary Faces", "Interior Faces", \
                  "Closed"]
    
    artist_folder = os.path.join(dataset_folder, "Artist_UVs")
    dataset_tags_file = os.path.join(dataset_folder, "dataset_tags.csv")
    dataset_tags = pd.read_csv(dataset_tags_file)
    cut_choices = ["Cut", "Uncut"]
    df = pd.DataFrame(columns=df_columns)

    for cut_choice in cut_choices:
        dataset_files = os.listdir(os.path.join(artist_folder, cut_choice))
        for i, fname in enumerate(dataset_files):
            print(i+1, "/", len(dataset_files), fname, end="\r\n")
            fpath = os.path.join(artist_folder, cut_choice, fname)
            full_name, ext = os.path.splitext(fname)
            if os.path.isfile(fpath) and ext == ".obj" and not fname.endswith("_all.obj"):
                split_name = full_name.split("_")
                object_number = int(split_name[1])
                mesh_name = "_".join(split_name[2:-1])
                chart_number = int(split_name[-1])

                v, uv, n, f, ftc, fn = igl.read_obj(fpath)
                f = f.reshape((-1, 3))

                tags_row = dataset_tags.loc[(dataset_tags['Filename'] == fname) & (dataset_tags['Cut'] == cut_choice)]

                is_manifold = tags_row.iloc[0]['Manifold']
                assert f.shape[1] == 3

                if is_manifold:
                    assert np.all(igl.is_edge_manifold(f))
                    boundary_segments = count_boundary_segments(f)
                    boundary_triangles = np.sum(np.any(igl.triangle_triangle_adjacency(f)[0].reshape((-1, 3)) == -1, axis=1))
                    interior_triangles = f.shape[0] - boundary_triangles
                    is_closed = boundary_triangles == 0
                    #Eq. (1) from http://people.csail.mit.edu/jsolomon/assets/vecinterp.pdf
                    genus = 0.5 * (2. - boundary_segments - v.shape[0] + edges(f).shape[0] - f.shape[0])
                    assert genus >= 0
                else:
                    boundary_segments = None
                    boundary_triangles = None
                    interior_triangles = None
                    is_closed = None
                    genus = None

                row = [cut_choice, fname, object_number, mesh_name, chart_number, \
                       v.shape[0], f.shape[0], is_manifold, \
                       genus, boundary_segments, \
                       boundary_triangles, interior_triangles, \
                       is_closed]

                df_row = pd.DataFrame({ col: [row[i]] for i, col in enumerate(df_columns)})
                df = pd.concat([df, df_row], ignore_index=True)
    df.to_csv(output_file, index=False)
    
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Collect dataset statistics")
    parser.add_argument("-d", "--dataset", type=str, required=True, dest="dataset")
    parser.add_argument("-o", "--output", dest="output", default="dataset_characteristics.csv")
    
    args = parser.parse_args()
    dataset_folder = os.path.abspath(args.dataset)
    output_file = os.path.abspath(args.output)
    
    get_dataset_characteristics(dataset_folder, output_file)