/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2025 Univ. Grenoble Alpes, CNRS, Grenoble INP - UGA, TIMC, 38000 Grenoble, France
 *
 * Visit http://camitk.imag.fr for more information
 *
 * This file is part of CamiTK.
 *
 * CamiTK is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * CamiTK is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with CamiTK.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $CAMITK_LICENCE_END$
 ****************************************************************************/

#include <math.h>
#include <iostream>
#include <string>
using namespace std;

#include <pml/PhysicalModel.h>
#include <pml/MultiComponent.h>
#include <pml/StructuralComponent.h>
#include <pml/Atom.h>

/// global variables
// -------------------- the physical model being processed ------------------------
PhysicalModel* pm;

// -------------------- all cluster ------------------------
class Cluster;
vector <Cluster*> allClusters;

// first atom is now replaced by 2nd atom
map<Atom*, Atom*> replacement;

double precision;
unsigned int nrOfNewAtoms;

/** class Cluster
 *
 *  a cluster is a set of point that are nearly at the same place (at a given precision)
 */
class Cluster {
public:
    /// create a cluster of this atom, using the given precision
    Cluster(double distance);

    ~Cluster();

    /// get composing atom by index
    Atom* getAtom(unsigned int id);

    /// add an atom to the cluster if in the given precision (true if added) (update center)
    bool addAtom(Atom*);

    /// get the center (create if necessary)
    Atom* getCenter();

    /// get nr of atoms in the cluster
    unsigned int nrOfAtoms();

    /// current atom index for new atoms
    static unsigned int atomIndex;

private:
    /// cluster center
    Atom* center;

    /// the cluster cell
    Cell* myCell;

    /// distance precision
    double distance;

    /// update center position
    void updateCenter();
};

unsigned int Cluster::atomIndex = 10000;


// -------------------- constructor/destructor ------------------------
Cluster::Cluster(double distance) {
    this->distance = distance;
    myCell = new Cell(pm, StructureProperties::POLY_VERTEX);
    center = new Atom(pm, Cluster::atomIndex++);
    Cluster::atomIndex++;
    nrOfNewAtoms++;
    //pm->addAtom(center);
}

Cluster::~Cluster() {
    delete myCell;
    delete center;
}

// -------------------- updateCenter ------------------------
void Cluster::updateCenter() {
    double pos[3];
    double posC[3];
    unsigned int size = myCell->getNumberOfStructures();

    for (unsigned int j = 0; j < 3; j++) {
        posC[j] = 0.0;
    }

    for (unsigned int i = 0; i < size; i++) {
        ((Atom*)myCell->getStructure(i))->getPosition(pos);
        for (unsigned int j = 0; j < 3; j++) {
            posC[j] += pos[j];
        }
    }

    for (unsigned int j = 0; j < 3; j++) {
        posC[j] /= ((double)size);
    }

    center->setPosition(posC);
}

// -------------------- getCenter ------------------------
Atom* Cluster::getCenter() {
    return center;
}

// -------------------- getAtom ------------------------
Atom* Cluster::getAtom(unsigned int id) {
    return ((Atom*)myCell->getStructure(id));
}

// -------------------- nrOfAtoms ------------------------
unsigned int Cluster::nrOfAtoms() {
    return myCell->getNumberOfStructures();
}

// -------------------- addAtom ------------------------
bool Cluster::addAtom(Atom* a) {
    double posC[3];
    center->getPosition(posC);
    double pos[3];
    a->getPosition(pos);
    double dist = 0.0;
    for (unsigned int i = 0; i < 3; i++) {
        dist += (posC[i] - pos[i]) * (posC[i] - pos[i]);
    }
    dist = sqrt(dist);

    if (nrOfAtoms() == 0 || dist < distance) {
        myCell->addStructure(a);
        updateCenter();
        return true;
    }
    else {
        return false;
    }
}


// -------------------- insert ------------------------
Cluster* insert(Atom* a) {
    vector <Cluster*>::iterator it;

    // look into allClusters for equivalence
    it = allClusters.begin();
    while (it != allClusters.end() && !(*it)->addAtom(a)) {
        ++it;
    }

    // not found => insert
    if (it == allClusters.end()) {
        // no cluster found: create a new one
        Cluster* c = new Cluster(precision);
        c->addAtom(a);
        allClusters.push_back(c);
        return c;
    }
    else  {
        return (*it);
    }

}

// -------------------- update ------------------------
void update(MultiComponent* mc) {
    if (mc == nullptr) {
        cout << "0 cells)" << endl;
        return;
    }

    cout << mc->getNumberOfCells() << " cells)" << endl;
    unsigned int degeneratedCount = 0;

    // WARNING non-rentrant loop
    for (unsigned int i = 0; i < mc->getNumberOfCells(); i++) {
        Cell* c = mc->getCell(i);
        StructuralComponent* sc = c->getAtoms();
        c->deleteAllStructures();
        cout << '\r' << i << flush;
        // remove previous list
        for (unsigned int j = 0; j < sc->getNumberOfStructures(); j++) {
            c->removeStructure(sc->getStructure(j));
        }
        // check it was done properly
        if (c->getNumberOfStructures() > 0)  {
            cout << "Cell #" << c->getIndex() << " cannot delete all atoms (still has " << c->getNumberOfStructures() << ")" << endl;
        }
        // add new atoms
        map<Atom*, Atom*>::iterator proxy;
        for (unsigned int j = 0; j < sc->getNumberOfStructures(); j++) {
            proxy = replacement.find((Atom*)sc->getStructure(j));
            if (proxy == replacement.end()) {
                // not found?!
                cout << "Cell #" << c->getIndex() << " cannot find replacement for Atom#" << sc->getStructure(i)->getIndex() << endl;
            }
            else {
                c->addStructureIfNotIn(proxy->second);
            }
        }
        // check that the cell did not degenerate
        StructureProperties::GeometricType t = c->getType();
        int normalSize = ((t == StructureProperties::LINE) ? 2 :
                          ((t == StructureProperties::TRIANGLE) ? 3 :
                           ((t == StructureProperties::QUAD) ? 4 :
                            ((t == StructureProperties::TETRAHEDRON) ? 4 :
                             ((t == StructureProperties::WEDGE) ? 6 :
                              ((t == StructureProperties::HEXAHEDRON) ? 8 : -1))))));
        if (normalSize != -1 && normalSize != (int)c->getNumberOfStructures()) {
            cout << "deleting degenerated Cell #" << c->getIndex() << ": it is a " <<
                 StructureProperties::toString(t) << " and should have " << normalSize <<
                 " atoms, but has now " << c->getNumberOfStructures() << flush;
            degeneratedCount++;
            c->getStructuralComponent(0)->removeStructure(c);
            // a cell was deleted, as it is a non re-entrant loop, get back one index
            i--;
        }
    }
    cout << flush;
    cout << degeneratedCount << " degenerated cells removed                                                        " << endl;
}

// -------------------- saveClusters ------------------------
void saveClusters(string clusterFile) {
    stringstream clusterOSS;
    clusterOSS << precision;


    string nonClusterFile = clusterFile;
    unsigned int pLast = clusterFile.rfind(".");

    if (pLast != string::npos) {
        clusterFile.erase(pLast);
        clusterFile += "-clusters-" + clusterOSS.str() + ".txt";
    }
    else {
        clusterFile = "clusters-" + clusterOSS.str() + ".txt";
    }

    pLast = nonClusterFile.rfind(".");
    if (pLast != string::npos) {
        nonClusterFile.erase(pLast);
        nonClusterFile += "-non-clusters-" + clusterOSS.str() + ".txt";
    }
    else {
        nonClusterFile = "non-clusters-" + clusterOSS.str() + ".txt";
    }

    cout << "-> please wait while saving cluster informations in " << clusterFile << " (" << pm->getNumberOfAtoms() - allClusters.size() << " clusters) and non-clusters in " << nonClusterFile << endl;
    ofstream clusterOutputFile(clusterFile.c_str());
    ofstream nonClusterOutputFile(nonClusterFile.c_str());
    for (vector <Cluster*>::iterator it = allClusters.begin(); it != allClusters.end(); ++it) {
        if ((*it)->nrOfAtoms() > 1) {
            clusterOutputFile << (*it)->getCenter()->getIndex() << " : ";
            for (unsigned int i = 0; i < (*it)->nrOfAtoms(); i++) {
                clusterOutputFile << (*it)->getAtom(i)->getIndex() << " ";
            }
            clusterOutputFile << endl;
        }
        else {
            nonClusterOutputFile << (*it)->getAtom(0)->getIndex() << endl;
        }
    }
    clusterOutputFile.close();
    nonClusterOutputFile.close();
}

// -------------------- optimize ------------------------
void optimize(string clusterFile) {
    StructuralComponent* atoms = pm->getAtoms();

    // create clusters
    for (unsigned int i = 0; i < atoms->getNumberOfStructures(); i++) {
        insert((Atom*) atoms->getStructure(i));
        cout << '\r' <<  i << flush;
    }
    cout << endl;

    cout << "found " << allClusters.size() << " clusters, for " <<  atoms->getNumberOfStructures() << " atoms" << endl;
    cout << "-> please wait while removing " << atoms->getNumberOfStructures() - allClusters.size() << " atoms" << endl;

    // create newAtoms with all atoms that are alone in their clusters and
    // all cluster centers
    StructuralComponent* newAtoms = new StructuralComponent(pm);
    for (vector <Cluster*>::iterator it = allClusters.begin(); it != allClusters.end(); ++it) {
        if ((*it)->nrOfAtoms() == 1) {
            newAtoms->addStructure((*it)->getAtom(0));
            replacement.insert(pair<Atom*, Atom*>((*it)->getAtom(0), (*it)->getAtom(0)));
        }
        else {
            newAtoms->addStructure((*it)->getCenter());
            // memorize replacements
            for (unsigned int i = 0; i < (*it)->nrOfAtoms(); i++) {
                replacement.insert(pair<Atom*, Atom*>((*it)->getAtom(i), (*it)->getCenter()));
            }
        }
    }

    // save the clusters and non cluster
    saveClusters(clusterFile);

    // update all atoms using replacements
    cout << "-> please wait while replacing atoms in exclusive components (";
    update(pm->getExclusiveComponents());
    cout << "-> please wait while replacing atoms in informative components (";
    update(pm->getInformativeComponents());

    // Use new atoms instead of atoms
    cout << "-> please wait while changing atoms" << endl;
    pm->setAtoms(newAtoms, false); // do not delete old atom list!
}

// -------------------- main ------------------------
int main(int argc, char** argv) {

    if (argc != 3) {
        cout << "Usage:" << endl;
        cout << "\t" << argv[0] << " file.pml precision" << endl;
        cout << "output to: file-optimized-precision.pml" << endl;
        cout << "PML " << PhysicalModel::VERSION << endl;
        exit(-1);
    }

    try {
        precision = atof(argv[2]);
        // read the pml
        string filename(argv[1]);
        cout << "-> please wait while opening " << filename << " for optimizing precision " << precision << endl;
        pm = new PhysicalModel(filename.c_str());

        cout << "-> please wait while optimizing " << pm->getNumberOfAtoms() << " atoms..." << endl;

        optimize(filename);
        /*
        // now, we have the right component, let's do some work...
        nrOfNewAtoms = 0;
        StructuralComponent * refined = refine(sc);

        // add the new SC to the exclusive component
        if (refined)
            pm->getExclusiveComponents()->addSubComponent(refined);
        */

        // save the result
        stringstream clusterOSS;
        clusterOSS << precision;
        unsigned int pLast = filename.rfind(".");
        if (pLast != string::npos) {
            filename.erase(pLast);
            filename += "-optimized-" + clusterOSS.str() + ".pml";
        }
        else {
            filename = "optimized-" + clusterOSS.str() + ".pml";
        }

        cout << "-> please wait while saving " << filename << " (" << nrOfNewAtoms << " new atoms)" << endl;

        ofstream outputFile(filename.c_str());
        pm->setName(pm->getName() + " optimized for " + clusterOSS.str());
        // do not optimize output (do not touch cell and atom id)
        pm->xmlPrint(outputFile, false);
        outputFile.close();

        delete pm;
    }
    catch (const PMLAbortException& ae) {
        cout << "AbortException: Physical model aborted:" << endl ;
        cout << ae.what() << endl;
    }
}