/************************************************************************************/ /**
// Copyright (c) 2006-2024 Advanced Micro Devices, Inc. All rights reserved.
/// \author AMD Developer Tools Team
/// \file
****************************************************************************************/
#include "tootlepch.h"

#include "soup.h"
#include "mesh.h"
#include "souptomesh.h"
#include "clustering.h"
#include "error.h"

//=================================================================================================================================
//
//          Clustering implementation
//
//=================================================================================================================================

const float BIGFLOAT = FLT_MAX;

class QNode
{
public:
    float weight;
    int   value;
    bool  operator<(const QNode& a) const
    {
        return weight < a.weight;
    }
    bool operator>(const QNode& a) const
    {
        return weight > a.weight;
    }
    QNode()
    {
    }
    QNode(float w, int v)
    {
        weight = w;
        value  = v;
    }
};

static float EdgeCost(Vector3 cn, Vector3 fn)
{
    float w = cn[0] * cn[0] + cn[1] * cn[1] + cn[2] * cn[2];

    if (w > 0)
    {
        cn /= sqrt(w);
    }

    float c = 1.f - Dot(cn, fn);
    assert(c < BIGFLOAT);
    return .00001f + (std::max)(0.f, c);
}

static void AddNeibToQueue(std::priority_queue<QNode, std::vector<QNode>, std::greater<QNode> >& q,
                           Mesh& /*mesh*/,
                           std::vector<int>&     fixed,
                           std::vector<float>&   cost,
                           std::vector<int>&     cluster,
                           std::vector<Vector3>& clusterNormal,
                           int                   f,
                           int                   ff,
                           std::vector<Vector3>& tn)
{
    if (fixed[ff])
    {
        return;
    }

    if (cost[f] < cost[ff])
    {
        float c = cost[f];
        assert(c < BIGFLOAT);
        c += EdgeCost(clusterNormal[cluster[f]], tn[ff]);

        if (c < cost[ff])
        {
            fixed[ff]   = false;
            cluster[ff] = cluster[f];
            cost[ff]    = c;
            assert(c > 0.f);
            q.push(QNode(c, ff));
        }
    }
}

static float FaceDistance(int a, int b, std::vector<Vector3>& tc)
{
    Vector3 d = tc[a] - tc[b];
    return sqrt(Dot(d, d));
}

static int MoveFaces(Mesh& mesh, std::vector<int>& seeds, std::vector<int>& cluster, std::vector<int>& fixed, std::vector<Vector3>& tn, float& fAvgDist)
{
    float fMaxDist = 0.f;
    int   distcnt  = 0;

    std::priority_queue<QNode, std::vector<QNode>, std::greater<QNode> > q;

    std::vector<int>     vis;
    std::vector<float>   cost;
    std::vector<int>     seed;
    std::vector<Vector3> clusterNormal;

    int nseeds = (int)seeds.size();

    //flood from seeds here
    for (int i = 0; i < static_cast<int>(mesh.t().size()); i++)
    {
        cluster[i] = nseeds;
        fixed[i]   = false;

        seed.push_back(0);
        cost.push_back(BIGFLOAT);
        vis.push_back(-1);
    }

    for (int i = 0; i < nseeds; i++)
    {
        int f      = seeds[i];
        cluster[f] = i;
        cost[f]    = 0.f;
        fixed[f]   = true;
        seed[f]    = 1;
        clusterNormal.push_back(tn[f]);
    }

    for (int i = 0; i < nseeds; i++)
    {
        int f = seeds[i];

        for (unsigned int j = 0; j < mesh.ae(f).size(); j++)
        {
            int ff = mesh.ae(f)[j];

            if (ff != -1)
            {  //not mesh boundary
                AddNeibToQueue(q, mesh, fixed, cost, cluster, clusterNormal, f, ff, tn);
            }
        }
    }

    int nvis  = 0;
    int f     = -1;
    int count = 0;

    while (!q.empty())
    {
        QNode n = q.top();
        f       = n.value;
        q.pop();

        if (fixed[f])
        {
            continue;
        }

        fAvgDist = (fAvgDist * (distcnt) + cost[f]) / (distcnt + 1);
        distcnt++;

        count++;

        assert(seed[f] == 0);

        assert(cluster[f] != nseeds);

        if (vis[f] == -1)
        {
            vis[f] = nvis++;
        }

        fixed[f] = true;
        clusterNormal[cluster[f]] += tn[f];

        for (unsigned int j = 0; j < mesh.ae(f).size(); j++)
        {
            int ff = mesh.ae(f)[j];

            if (ff != -1)
            {  //not mesh boundary
                AddNeibToQueue(q, mesh, fixed, cost, cluster, clusterNormal, f, ff, tn);
            }
        }
    }

    if (f == -1)
    {
        // this means that none of the seeds had any neighboring faces.
        // if this happens, then we should find the first face not assigned to a cluster, and return that one
        // this will cause us to add a new seed to that face, effectively jumping across the connectivity gap

        // if the entire input mesh happens to be disconnected, then this should converge to one face per cluster
        UINT i;

        for (i = 0; i < mesh.t().size(); i++)
        {
            if (cluster[i] == nseeds)
            {
                // use this face
                break;
            }
        }

        // we shouldn't get to the end
        assert(i != mesh.t().size());

        return i;
    }

    fMaxDist = cost[f];
    debugf(("Max: %.6f, Avg: %.6f", fMaxDist, fAvgDist));

    int lowvisi = nvis;

    int lowvisf = -1;

    for (int i = 0; i < static_cast<int>(mesh.t().size()); i++)
    {
        if (fixed[i] == false && vis[i] < lowvisi)
        {
            lowvisi = vis[i];
            lowvisf = i;
        }
    }

    if (lowvisf != -1)
    {
        fMaxDist = BIGFLOAT / 2.f;
        return lowvisf;
    }
    else
    {
        return f;
    }

    return 0;
}

static void MoveSeeds(Mesh& mesh, std::vector<int>& seeds, std::vector<int>& cluster, std::vector<int>& fixed, std::vector<Vector3>& tc)
{
    //flood from boundaries here

    int                nseeds = (int)seeds.size();
    std::vector<float> cost;
    std::vector<int>   visited;

    std::priority_queue<QNode, std::vector<QNode>, std::greater<QNode> > q;

    for (int i = 0; i < static_cast<int>(mesh.t().size()); i++)
    {
        cost.push_back(BIGFLOAT);
        visited.push_back(0);

        for (int j = 0; j < static_cast<int>(mesh.ae(i).size()); j++)
        {
            int ff = mesh.ae(i)[j];

            if (ff != -1 &&  //not mesh boundary
                (cluster[ff] != cluster[i] || fixed[i] == false || fixed[ff] == false))
            {
                cost[i] = 0;
                q.push(QNode(.0f, i));
                break;
            }
        }
    }

    while (!q.empty())
    {
        QNode n = q.top();
        int   f = n.value;
        q.pop();

        if (visited[f] == 1)
        {
            continue;
        }

        visited[f] = 1;

        for (unsigned int j = 0; j < mesh.ae(f).size(); j++)
        {
            int ff = mesh.ae(f)[j];

            if (ff != -1 &&  //not mesh boundary
                cost[f] < cost[ff])
            {
                float c = cost[f] + FaceDistance(f, ff, tc);

                if (c < cost[ff])
                {
                    cost[ff] = c;
                    q.push(QNode(c, ff));
                }
            }
        }
    }

    // find the face in each cluster with highest edge cost
    // this will be the last face that was reached during the flood

    std::vector<float> seeddist;

    for (int i = 0; i < nseeds; i++)
    {
        seeddist.push_back(cost[seeds[i]]);
    }

    for (int i = 0; i < static_cast<int>(mesh.t().size()); i++)
    {
        if (cost[i] > seeddist[cluster[i]])
        {
            // do not change seeds if the seed we have chosen is the seed of another cluster
            bool bMove = true;

            for (int j = 0; j < nseeds; j++)
            {
                if (i == seeds[j])
                {
                    bMove = false;
                }
            }

            if (bMove)
            {
                seeddist[cluster[i]] = cost[i];
                seeds[cluster[i]]    = i;
            }
        }
    }
}

static int FingerPrint(Mesh& mesh, std::vector<int>& cluster)
{
    int ret = 0;

    for (int i = 0; i < static_cast<int>(mesh.t().size()); i++)
    {
        ret = (ret + i * cluster[i]) % 40923840;
    }

    return ret;
}

/// Takes a triangle soup and builds face clusters
/// \param soup  The soup to cluster
/// \param nClusters The number of clusters to create.  If 0, then automatic clustering is used, and nClusters is set
///                  to the number of clusters that are generated
/// \param cluster  An array that will receive the cluster ID to assign to each face
/// \return  One of the ClusterResult return codes
ClusterResult Cluster(Soup* soup, UINT& nClusters, std::vector<int>& cluster)
{
    // compute face normals
    if (!soup->ComputeNormals())
    {
        return CLUSTER_OUT_OF_MEMORY;
    }

    const int nFaces = static_cast<int>(soup->t().size());

    Mesh mesh;

    SoupToMesh(soup, &mesh);

    // compute the set of triangles which use each vertex
    VTArray meshVT;

    if (!mesh.ComputeVT(meshVT))
    {
        return CLUSTER_OUT_OF_MEMORY;
    }

    // compute per-face adjacency
    // allocate the AE array ahead of time so that we can detect out-of-memory conditions
    mesh.ae().resize(mesh.t().size());

    // computeAE should now only fail if the mesh is non-manifold
    if (!mesh.ComputeAE(meshVT))
    {
        return CLUSTER_OUT_OF_MEMORY;
    }

    // compute face normals
    std::vector<Vector3> tn;

    if (!mesh.ComputeTriNormals(tn))
    {
        return CLUSTER_OUT_OF_MEMORY;
    }

    // compute face centers
    std::vector<Vector3> tc;

    if (!mesh.ComputeTriCenters(tc))
    {
        return CLUSTER_OUT_OF_MEMORY;
    }

    cluster.clear();

    for (int i = 0; i < nFaces; i++)
    {
        cluster.push_back(0);
    }

    int   nCurClusters = 0;
    float fAvgDist     = BIGFLOAT;
    float fAvgDistOld  = BIGFLOAT;

    std::vector<int> seeds;
    std::vector<int> fixed;
    std::vector<int> fp;
    fixed.resize(mesh.t().size());

    // if cluster count is fixed, then clamp it
    if ((int)nClusters > mesh.t().size())
    {
        nClusters = (UINT)mesh.t().size();
    }

    srand(982748);  //in order to get reproducible results
    int last = rand() % mesh.t().size();

    bool bHasUnassignedFaces = true;

    for (int i = 0;; i++)
    {
        fp.push_back(FingerPrint(mesh, cluster));

        if (bHasUnassignedFaces || nCurClusters < (int)nClusters || nClusters == 0)
        {
            seeds.push_back(last);
            nCurClusters++;
        }
        else
        {
            nClusters = nCurClusters;
            int done  = 0;

            for (int j = i - 1; j >= 0; j--)
            {
                if (fp[i] == fp[j])
                {
                    done = 1;
                    break;
                }
            }

            if (done)
            {
                break;
            }
        }

        // if we have nClusters == nFaces, then we need to bail out
        if (nCurClusters == nFaces)
        {
            nClusters = nFaces - 1;
            return CLUSTER_OK;
        }

        if (nCurClusters > 1)
        {
            MoveSeeds(mesh, seeds, cluster, fixed, tc);
        }

        if (fAvgDistOld == BIGFLOAT)
        {
            fAvgDistOld = fAvgDist;
        }
        else
        {
            fAvgDistOld = (fAvgDistOld + fAvgDist) / 2.f;
        }

        last = MoveFaces(mesh, seeds, cluster, fixed, tn, fAvgDist);

        // do not stop adding seeds until EVERY face has been assigned a cluster
        // if the mesh has lots of connected components then we can get a case like this
        bHasUnassignedFaces = false;
        UINT nSeeds         = (UINT)seeds.size();

        // @TODO Verify this code - previously, it used a local i
        for (UINT j = 0; j < cluster.size(); j++)
        {
            if (cluster[j] == static_cast<int>(nSeeds))
            {
                bHasUnassignedFaces = true;
                break;
            }
        }

        //heuristic for picking # of clusters, if in automatic mode
        if (nClusters == 0)
        {
            if (!bHasUnassignedFaces)
            {
                if (nCurClusters >= 128 && nFaces / nCurClusters < 250)
                {
                    nClusters = nCurClusters;
                }

                if (nCurClusters >= 16 && fAvgDist > .97f * fAvgDistOld)
                {
                    nClusters = nCurClusters;
                }

                //comparison: debugf(("%f < %f", fAvgDist, .97f * fAvgDistOld));
            }
        }
    }

    debugf(("Final # of clusters: %i", nClusters));
    return CLUSTER_OK;
}

// array of clusters for each face, needed for cluster sorting
static int* g_pCluster;

// comparison function to sort by cluster ID
int __cdecl SortByClusterID(const void* a, const void* b)
{
    int* ia = (int*)a;
    int* ib = (int*)b;

    if (g_pCluster[*ia] < g_pCluster[*ib])
    {
        return -1;
    }
    else if (g_pCluster[*ia] > g_pCluster[*ib])
    {
        return 1;
    }
    else
    {
        return 0;
    }
}

/// \param soup   A triangle soup to sort
/// \param clusterIDs  Per-face cluster IDs
/// \param pRemapArray  An array that will receive the face re-mapping.  May NOT be NULL as it is used internally for temporary storage
/// \return True if successful, false if out of memory
bool SortFacesByCluster(Soup& soup, std::vector<int>& clusterIDs, UINT* pRemapArray)
{
    // sort faces by cluster ID
    std::vector<Soup::Triangle> t;
    std::vector<int>            c;

    for (int i = 0; i < static_cast<int>(soup.t().size()); i++)
    {
        pRemapArray[i] = i;
    }

    t = soup.t();
    c = clusterIDs;

    g_pCluster = &clusterIDs[0];
    qsort(pRemapArray, soup.t().size(), sizeof(int), SortByClusterID);

    for (int i = 0; i < static_cast<int>(soup.t().size()); i++)
    {
        soup.t(i)     = t[pRemapArray[i]];
        clusterIDs[i] = c[pRemapArray[i]];
    }

    return true;
}
