#include <iMesh_extensions.h>

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <map>
#include <stdexcept>
#include <string>
#include <vector>

class itaps_error : public std::runtime_error
{
public:
    itaps_error(const char *what) : runtime_error(what) {}
};

void check_error(iMesh_Instance mesh,int err)
{
    if(err)
    {
        char descr[120];
        int descr_err;
        iMesh_getDescription(mesh,descr,&descr_err,sizeof(descr)-1);
        if(descr_err)
            strncpy(descr,"Unable to retrieve error description",sizeof(descr));
        throw itaps_error(descr);
    }
}

// Get a mapping of tag values to sets that have these values on them.
template<typename T>
std::map<T,iBase_EntitySetHandle> sets_by_tag(iMesh_Instance mesh,
                                              const std::string &tag_name)
{
    iBase_TagHandle tag;
    int err;
    std::map<T,iBase_EntitySetHandle> mapped;

    iMesh_getTagHandle(mesh,tag_name.c_str(),&tag,&err,tag_name.size());
    check_error(mesh,err);

    iBase_EntitySetHandle rootset;
    iMesh_getRootSet(mesh,&rootset,&err);
    check_error(mesh,err);

    iBase_EntitySetHandle *sets = NULL;
    int sets_alloc = 0,sets_size;
    iMesh_getEntSets(mesh,rootset,0,&sets,&sets_alloc,&sets_size,&err);
    check_error(mesh,err);

    for(int i=0; i<sets_size; i++)
    {
        T key;
        int key_alloc = sizeof(key),key_size;
        T *k = &key;
        iMesh_getEntSetData(mesh,sets[i],tag,(char**)&k,&key_alloc,&key_size,
                            &err);
        if(err == 0)
            mapped[key] = sets[i];
    }

    free(sets);
    return mapped;
}

typedef std::vector<iBase_EntityHandle> entity_vector;

// Get the entities of dimension |dim| from |set|, and all entities of dimension
// |adj_dim| adjacent to those which are also in |adj_set|.
entity_vector select(iMesh_Instance mesh,iBase_EntitySetHandle set,int dim,
                     iBase_EntitySetHandle adj_set,int adj_dim)
{
    int err;

    // First, get the entities of the specified dimension, all adjacent entities
    // of the specified adj dimension, and all entities in |adj| of the
    // specified adj dimension.

    iBase_EntityHandle *ents = NULL;
    int ents_alloc = 0,ents_size;
    iMesh_getEntitiesRec(mesh,set,dim,iMesh_ALL_TOPOLOGIES,true,&ents,
                         &ents_alloc,&ents_size,&err);
    check_error(mesh,err);

    iBase_EntityHandle *adj = NULL;
    int adj_alloc = 0,adj_size;
    int *off = NULL;
    int off_alloc = 0,off_size;
    iMesh_getEntArrAdj(mesh,ents,ents_size,adj_dim,&adj,&adj_alloc,&adj_size,
                       &off,&off_alloc,&off_size,&err);
    check_error(mesh,err);

    iBase_EntityHandle *adj_mask = NULL;
    int adj_mask_alloc = 0,adj_mask_size;
    iMesh_getEntitiesRec(mesh,adj_set,adj_dim,iMesh_ALL_TOPOLOGIES,true,
                         &adj_mask,&adj_mask_alloc,&adj_mask_size,&err);
    check_error(mesh,err);

    std::sort(ents,ents+ents_size);
    std::sort(adj,adj+adj_size);
    std::sort(adj_mask,adj_mask+adj_mask_size);

    // Get the intersection of adj and adj_mask. adj may contain duplicates, but
    // since adj_mask doesn't, the intersection won't have duplicates.

    entity_vector filtered_adj;
    std::back_insert_iterator<entity_vector> i(filtered_adj);
    std::set_intersection(adj,adj+adj_size,
                          adj_ents,adj_ents+adj_size,
                          i);

    free(adj);
    free(adj_ents);

    // Get the union of adj and adj_mask. As long as |dim| and |adj_dim| are
    // different, this should be guaranteed to have no duplicates. (TODO: handle
    // that case.)

    std::vector<iBase_EntityHandle> filtered;
    std::back_insert_iterator<entity_vector> ii(filtered);
    std::set_union(filtered_adj.begin(),filtered_adj.end(),
                   ents,ents+ents_size,
                   ii);

    free(ents);

    return filtered;
}

int main(int argc,char **argv)
{
    if(argc != 2)
    {
        printf("Usage: %s input\n",argv[0]);
        exit(1);
    }

    iMesh_Instance mesh;
    iBase_EntitySetHandle rootset;
    int err;

    iMesh_newMesh("",&mesh,&err,0);
    check_error(mesh,err);

    iMesh_getRootSet(mesh,&rootset,&err);
    check_error(mesh,err);

    iMesh_load(mesh,rootset,argv[1],0,&err,strlen(argv[1]),0);
    check_error(mesh,err);

    std::map<int,iBase_EntitySetHandle> neumann = sets_by_tag<int>(
        mesh,"NEUMANN_SET");

    entity_vector src;
    src = select(mesh,neumann[100],iBase_FACE,neumann[500],iBase_EDGE);
    printf("# edges/faces = %d\n",src.size());

    iBase_EntityHandle *adj = NULL;
    int adj_alloc = 0,adj_size;
    int *off = NULL;
    int off_alloc = 0,off_size;
    iMesh_getEntArrAdj(mesh,&src[0],src.size(),iBase_VERTEX,
                       &adj,&adj_alloc,&adj_size,
                       &off,&off_alloc,&off_size,
                       &err);
    check_error(mesh,err);

    std::sort(adj,adj+adj_size);
    iBase_EntityHandle *end = std::unique(adj,adj+adj_size);
    printf("# adj verts = %d\n",adj_size);
    printf("# unique adj verts = %d\n",end-adj);

    free(adj);
    free(off);

    iBase_EntitySetHandle set;
    iMesh_createEntSet(mesh,false,&set,&err);
    iMesh_addEntArrToSet(mesh,&src[0],src.size(),set,&err);

    for(int d=iBase_EDGE; d<=iBase_REGION; d++)
    {
        iBase_EntityHandle *tmp = NULL;
        int tmp_alloc = 0,tmp_size;

        iMesh_getEntities(mesh,set,d,iMesh_ALL_TOPOLOGIES,&tmp,&tmp_alloc,
                          &tmp_size,&err);
        check_error(mesh,err);

        iBase_EntityHandle *verts = NULL;
        int verts_alloc = 0,verts_size;
        int *off = NULL;
        int off_alloc = 0,off_size;
        iMesh_getEntArrAdj(mesh,tmp,tmp_size,iBase_VERTEX,
                           &verts,&verts_alloc,&verts_size,
                           &off,&off_alloc,&off_size,
                           &err);
        check_error(mesh,err);

        iMesh_addEntArrToSet(mesh,verts,verts_size,set,&err);
        check_error(mesh,err);

        free(tmp);
        free(verts);
        free(off);
    }

    iBase_EntityHandle *tmp = NULL;
    int tmp_alloc = 0,tmp_size;
    iMesh_getEntities(mesh,set,iBase_ALL_TYPES,iMesh_ALL_TOPOLOGIES,
                      &tmp,&tmp_alloc,&tmp_size,&err);
    check_error(mesh,err);

    printf("\nTotal # entities = %d\n",tmp_size);
    printf("Should be = %d\n",end-adj+src.size());

    std::sort(tmp,tmp+tmp_size);

    printf("\nDuplicates:\n");
    for(int i=1; i<tmp_size; i++)
    {
        if(tmp[i] == tmp[i-1])
            printf("  %p\n",tmp[i]);
    }

    free(tmp);

    iMesh_dtor(mesh,&err);
    check_error(mesh,err);
}