package cz.fidentis.analyst.procrustes;

import cz.fidentis.analyst.face.HumanFace;
import cz.fidentis.analyst.feature.FeaturePoint;
import cz.fidentis.analyst.icp.IcpTransformation;
import cz.fidentis.analyst.mesh.core.MeshPoint;
import cz.fidentis.analyst.procrustes.exceptions.ProcrustesAnalysisException;
import cz.fidentis.analyst.procrustes.utils.ProcrustesAnalysisUtils;
import org.ejml.data.Matrix;
import org.ejml.simple.SimpleMatrix;
import org.ejml.simple.SimpleSVD;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import javax.vecmath.Point3d;


/**
 * @author Jakub Kolman
 */
public class ProcrustesAnalysis {

    private ProcrustesAnalysisFaceModel faceModel1;
    private ProcrustesAnalysisFaceModel faceModel2;

    private Point3d modelCentroid1;
    private Point3d modelCentroid2;

    // based on some example models this is an approximate centroid of feature points
//    private Point3d centroid  = new Point3d(0, 0, 0);

    /**
     * Constructor
     *
     * @param humanFace1
     * @param humanFace2
     * @throws ProcrustesAnalysisException
     */
    public ProcrustesAnalysis(
            HumanFace humanFace1,
            HumanFace humanFace2) throws ProcrustesAnalysisException {

        ProcrustesAnalysisFaceModel model1 = new ProcrustesAnalysisFaceModel(humanFace1);
        ProcrustesAnalysisFaceModel model2 = new ProcrustesAnalysisFaceModel(humanFace2);

        if (model1.getFeaturePointsMap().values().size() != model2.getFeaturePointsMap().values().size()) {
            throw new ProcrustesAnalysisException("Lists of feature points do not have the same size");
        }

//        this.orderedFeaturePointList1 = new ArrayList<>(model1.getFeaturePointsMap().values());
//        this.orderedFeaturePointList2 = new ArrayList<>(model2.getFeaturePointsMap().values());

        if (!featurePointTypesEquivalence(
                model1.getFeaturePointValues(),
                model2.getFeaturePointValues())) {
            throw new ProcrustesAnalysisException("Lists of feature points do not have the same feature point types");
        }

        this.faceModel1 = model1;
        this.faceModel2 = model2;

        this.modelCentroid1 = ProcrustesAnalysisUtils.findCentroidOfFeaturePoints(model1.getFeaturePointValues());
        this.modelCentroid2 = ProcrustesAnalysisUtils.findCentroidOfFeaturePoints(model2.getFeaturePointValues());

    }

    public ProcrustesAnalysisFaceModel getFaceModel1() {
        return faceModel1;
    }

    public ProcrustesAnalysisFaceModel getFaceModel2() {
        return faceModel2;
    }

    public void analyze() {
        List<SimpleMatrix> transformation = new ArrayList<>();

        this.superImpose();
        this.rotate();
    }

    private Point3d computeCentroidsDistance(Point3d centroid1, Point3d centroid2) {
        double x = (centroid1.x - centroid2.x);
        double y = (centroid1.y - centroid2.y);
        double z = (centroid1.z - centroid2.z);

        Point3d computedCentroid = new Point3d(x, y, z);
        return computedCentroid;
    }

    /**
     * Imposes two face models (lists of feature points and vertices) over each other
     */
    private void superImpose() {

        // moves vertices of face models
//        moveModelToSameOrigin(this.faceModel2.getVertices(), this.faceModel2.getFeaturePointsMap());

        normalize(this.faceModel1, this.modelCentroid1);
        normalize(this.faceModel2, this.modelCentroid2);

//        moveVertices(this.modelCentroid1, featurePointModel1.getVertices());
//        moveModelToPoint(this.modelCentroid1, faceModel1.getVertices(), faceModel1.getFeaturePointsMap());
//        moveModelToPoint(this.modelCentroid1, faceModel2.getVertices(), faceModel2.getFeaturePointsMap());
    }

    private void normalize(ProcrustesAnalysisFaceModel faceModel, Point3d centroid) {
//        float size = ProcrustesAnalysisUtils.countSize(centroid, faceModel.getFeaturePointsMap());
        centerToOrigin(faceModel, centroid);
    }

    private void centerToOrigin(ProcrustesAnalysisFaceModel faceModel, Point3d centroid) {
        for (FeaturePoint fp: faceModel.getFeaturePointsMap().values()) {
            fp.getPosition().x -= centroid.x;
            fp.getPosition().y -= centroid.y;
            fp.getPosition().z -= centroid.z;
        }
        SimpleMatrix centeredVertices = new SimpleMatrix(
                faceModel.getVerticesMatrix().numRows(), faceModel.getVerticesMatrix().numCols());
        for (int i = 0; i < faceModel.getVerticesMatrix().numRows(); i++) {
            centeredVertices.set(i, 0, faceModel.getVerticesMatrix().get(i, 0) - centroid.x);
            centeredVertices.set(i, 1, faceModel.getVerticesMatrix().get(i, 1) - centroid.y);
            centeredVertices.set(i, 2, faceModel.getVerticesMatrix().get(i, 2) - centroid.z);
        }
        faceModel.setVerticesMatrix(centeredVertices);
    }

    private void moveModelToSameOrigin(List<MeshPoint> vertices, HashMap<Integer, FeaturePoint> featurePointsMap) {
        Point3d diff = computeCentroidsDistance(this.modelCentroid1, this.modelCentroid2);
        moveModelToPoint(diff, vertices, featurePointsMap);
    }

    private void moveModelToPoint(Point3d centroid, List<MeshPoint> vertices, HashMap<Integer, FeaturePoint> featurePointsMap) {
        moveVertices(centroid, vertices);
        moveFeaturePointsToVertices(centroid, featurePointsMap);
    }

    private void moveFeaturePointsToVertices(Point3d centroid, HashMap<Integer, FeaturePoint> featurePointsMap) {
        for (FeaturePoint fp : featurePointsMap.values()) {
            fp.getPosition().x = fp.getPosition().x + centroid.x;
            fp.getPosition().y = fp.getPosition().y + centroid.y;
            fp.getPosition().z = fp.getPosition().z + centroid.z;
        }
    }

    /**
     * By rotation of matrices solves orthogonal procrustes problem
     */
    private void rotate() {

        // There is no reason trying to rotate less than 3 elements
        if (this.getFaceModel1().getFeaturePointsMap().size() < 3) {
            throw new ProcrustesAnalysisException("To do procrustes analysis models have to have at least 3 feature points");
        }
        SimpleMatrix primaryMatrix = ProcrustesAnalysisUtils.createFeaturePointMatrix(this.faceModel2.getFeaturePointValues());
        SimpleMatrix transposedMatrix = ProcrustesAnalysisUtils.createFeaturePointMatrix(
                this.faceModel1.getFeaturePointValues()).transpose();

        SimpleMatrix svdMatrix = transposedMatrix.mult(primaryMatrix);
        SimpleSVD<SimpleMatrix> singularValueDecomposition = svdMatrix.svd();
        SimpleMatrix transposedU = singularValueDecomposition.getU().transpose();
        SimpleMatrix r = singularValueDecomposition.getV().mult(transposedU);
        primaryMatrix = primaryMatrix.mult(r);

        this.faceModel2.setFeaturePointsMap(
                ProcrustesAnalysisUtils.createFeaturePointMapFromMatrix(
                        primaryMatrix, this.faceModel2));

//        for (int i = 0; i < this.faceModel2.getFeaturePointTypeCorrespondence().size(); i++) {
//            FeaturePoint fp = new FeaturePoint(
//                    primaryMatrix.get(i, 0),
//                    primaryMatrix.get(i, 1),
//                    primaryMatrix.get(i, 2),
//                    this.faceModel2.getFeaturePointsMap().get(
//                            this.faceModel2.getFeaturePointTypeCorrespondence().get(i)).getFeaturePointType()
//            );
//            this.faceModel2.getFeaturePointsMap().put(i, fp);
//        }

//        rotateVertices(this.faceModel2, primaryMatrix);

        SimpleMatrix rotatedVertices = this.faceModel2.getVerticesMatrix().mult(r);
        createListFromMatrix(rotatedVertices, this.faceModel2);
    }

    private void createListFromMatrix(SimpleMatrix rotatedVertices, ProcrustesAnalysisFaceModel faceModel) {
        if (faceModel.getVertices() != null) {
            for (int i = 0; i < rotatedVertices.numRows(); i++) {
                faceModel.getVertices().get(i).getPosition().x = rotatedVertices.get(i, 0);
                faceModel.getVertices().get(i).getPosition().y = rotatedVertices.get(i, 1);
                faceModel.getVertices().get(i).getPosition().z = rotatedVertices.get(i, 2);
            }
        }
    }

    /**
     * Rotates every vertex of model by given matrix
     *
     * @param model
     * @param matrix
     */
    // if rotated vertices are drawn immediately it is better to set them after rotating them all
    // so it would be drawn just once
//    private void rotateVertices(ProcrustesAnalysisFaceModel model, SimpleMatrix matrix) {
//        if (model.getVertices() != null) {
//            for (MeshPoint v : model.getVertices()) {
//                v.setPosition(ProcrustesAnalysisUtils.rotateVertex(v, matrix));
//            }
//        }
//    }

    /**
     * Calculates new vertices adjusted to the centroid by method {@see findCenteroidOfFeaturePoints},
     * so moved featurepoints would correspond with the same place on face as they did before.
     *
     * @param centroid
     * @param vertices
     * @return
     */
    private void moveVertices(Point3d centroid, List<MeshPoint> vertices) {
        if (vertices != null && centroid != null) {
//            List<MeshPoint> movedVertices = null;
            for (MeshPoint v : vertices) {
                v.getPosition().x = (v.getPosition().x + centroid.x);
                v.getPosition().y = (v.getPosition().y + centroid.y);
                v.getPosition().z = (v.getPosition().z + centroid.z);
            }
        } else {
            throw new ProcrustesAnalysisException("Could not compute vertices locations after moving the centroid from model to feature points");
        }
    }

    /**
     * Creates sorted feature point lists and compares them whether they have
     * the same feature point types.
     *
     * @param fpList1
     * @param fpList2
     * @return
     */
    private boolean featurePointTypesEquivalence(List<FeaturePoint> fpList1, List<FeaturePoint> fpList2) {
        return ProcrustesAnalysisUtils.checkFeaturePointsType(fpList1, fpList2);
    }

    /**
     * Calculate scaling ratio of how much the appropriate object corresponding
     * to the second feature point list has to be scale up or shrunk.
     * <p>
     * If returned ratioValue is greater 1 then it means that the second object
     * should be scaled up ratioValue times. If returned ratioValue is smaller 1
     * than the second object should be shrunk.
     *
     * @return ratioValue
     */
//    private double calculateScalingValue() {
//        double[] distancesOfList1 = ProcrustesAnalysisUtils.calculateMeanDistancesFromOrigin(
//                this.faceModel1.getFeaturePointValues());
//        double[] distancesOfList2 = ProcrustesAnalysisUtils.calculateMeanDistancesFromOrigin(
//                this.faceModel2.getFeaturePointValues());
//
//        double[] ratioArray = new double[distancesOfList1.length];
//        double ratioValue = 0;
//
//        for (int i = 0; i < distancesOfList1.length; i++) {
//            ratioArray[i] += distancesOfList1[i] / distancesOfList2[i];
//        }
//
//        for (double ratio : ratioArray) {
//            ratioValue += ratio;
//        }
//
//        return ratioValue / distancesOfList1.length;
//    }

}