package cz.fidentis.analyst.procrustes;

import cz.fidentis.analyst.face.HumanFace;
import cz.fidentis.analyst.feature.FeaturePoint;
import cz.fidentis.analyst.feature.api.IPosition;
import cz.fidentis.analyst.mesh.core.MeshPoint;
import org.ejml.simple.SimpleMatrix;
import org.ejml.simple.SimpleSVD;

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


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

    private ProcrustesAnalysisFaceModel faceModel1;
    private ProcrustesAnalysisFaceModel faceModel2;

    protected final Point3d modelCentroid1;
    protected final Point3d modelCentroid2;

    private boolean scale = false;

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

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

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

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

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

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

    /**
     * Constructor with variable for scaling option.
     * Default scale option is set to false.
     *
     * @param humanFace1
     * @param humanFace2
     * @param scale
     * @throws DataFormatException
     */
    public ProcrustesAnalysis(
            HumanFace humanFace1,
            HumanFace humanFace2,
            boolean scale) throws DataFormatException {
        this(humanFace1, humanFace2);
        this.scale = scale;
    }

    /**
     * Method called for analysis after creating initial data in constructor. This method causes superimposition
     * and rotation of the faces.
     *
     * @throws DataFormatException if faces have less than 3 feature points in which case analysis doesn't make sense
     */
    public void analyze() throws DataFormatException {
        if (this.faceModel1.getFeaturePointsMap().size() > 3) {
            if (scale) {
                double scaleFactorValue = this.calculateScalingValue(this.faceModel1.getFeaturePointValues(), this.faceModel2.getFeaturePointValues());
                if (scaleFactorValue != 1) {
                    this.scaleFace(this.faceModel2, scaleFactorValue);
                }
            }
            this.superImpose();
            this.rotate();
        } else {
            throw new DataFormatException("Faces have less than 3 feature points.");
        }
    }

    /**
     * Imposes two face models (lists of feature points and vertices) over each other
     */
    protected void superImpose() {
        centerToOrigin(this.faceModel1, this.modelCentroid1);
        centerToOrigin(this.faceModel2, this.modelCentroid2);
    }

    /**
     * Centers given face model to origin that is given centroid.
     * Moves all vertices and feature points by value of difference between vertex/feature point and centroid.
     *
     * @param faceModel
     * @param 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;
        }
        for (MeshPoint v : faceModel.getVertices()) {
            v.getPosition().x = v.getX() - centroid.x;
            v.getPosition().y = v.getY() - centroid.y;
            v.getPosition().z = v.getZ() - centroid.z;
        }
    }

    /**
     * By rotation of matrices solves orthogonal procrustes problem
     */
    private void rotate() {
        SimpleMatrix primaryMatrix = createMatrixFromList(this.faceModel2.getFeaturePointValues());
        SimpleMatrix transposedMatrix = createMatrixFromList(
                this.faceModel1.getFeaturePointValues()).transpose();

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

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

        rotateVertices(this.faceModel2.getVertices(), rotationMatrix);
    }


    /**
     * Rotates all vertices.
     * <p>
     * For more details check out single vertex rotation {@link #rotateVertex}.
     *
     * @param vertices
     * @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(List<MeshPoint> vertices, SimpleMatrix matrix) {
        if (vertices != null) {
            for (int i = 0; i < vertices.size(); i++) {
                rotateVertex(vertices.get(i), matrix);
            }
        }
    }

    /**
     * Rotates vertex v by simulating matrix multiplication with given matrix
     *
     * @param v
     * @param matrix
     */
    private static void rotateVertex(MeshPoint v, SimpleMatrix matrix) {
        double x = ((v.getX() * matrix.get(0, 0))
                + (v.getY() * matrix.get(1, 0))
                + (v.getZ() * matrix.get(2, 0)));
        double y = ((v.getX() * matrix.get(0, 1))
                + (v.getY() * matrix.get(1, 1))
                + (v.getZ() * matrix.get(2, 1)));
        double z = ((v.getX() * matrix.get(0, 2))
                + (v.getY() * matrix.get(1, 2))
                + (v.getZ() * matrix.get(2, 2)));
        v.getPosition().x = x;
        v.getPosition().y = y;
        v.getPosition().z = z;
    }

    /**
     * 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
     */
    protected double calculateScalingValue(List<FeaturePoint> featurePointList1, List<FeaturePoint> featurePointList2) {
        double[] distancesOfList1 = new double[featurePointList1.size()];
        double[] distancesOfList2 = new double[featurePointList2.size()];

        Point3d featurePointCentroid1 = findCentroidOfFeaturePoints(
                featurePointList1);
        Point3d featurePointCentroid2 = findCentroidOfFeaturePoints(
                featurePointList2);

        for (int i = 0; i < featurePointList1.size(); i++) {
            distancesOfList1[i] = calculateDistanceFromPoint(featurePointList1.get(i), featurePointCentroid1);
            distancesOfList2[i] = calculateDistanceFromPoint(featurePointList2.get(i), featurePointCentroid2);
        }

        double[] ratioArray = new double[distancesOfList1.length];
        double ratioValue = 0;

        for (int i = 0; i < distancesOfList1.length; i++) {
            ratioArray[i] += distancesOfList1[i] / distancesOfList2[i];
        }
        for (int i = 0; i < ratioArray.length; i++) {
            ratioValue += ratioArray[i];
        }
        return ratioValue / distancesOfList1.length;
    }

    /**
     * Initiates scaling of feature points and vertices of given face model by scaleFactor.
     *
     * @param faceModel
     * @param scaleFactor
     */
    private void scaleFace(ProcrustesAnalysisFaceModel faceModel, double scaleFactor) {
        calculateScaledList(faceModel.getVertices(), scaleFactor);
        calculateScaledList(faceModel.getFeaturePointValues(), scaleFactor);
    }

    /**
     * Scales each given point from list by multiplying its position coordinates with scaleFactor.
     *
     * @param list
     * @param scaleFactor
     * @param <T> either MeshPoint or FeaturePoint type.
     */
    private <T extends IPosition> void calculateScaledList(List<T> list, double scaleFactor) {
        List<T> scaledList = new ArrayList<>();
        for (T point : list) {
            scaledList.add(scalePointDistance(point, scaleFactor));
        }
    }

    /**
     * Scales position of given point by multiplying its coordinates with given scaleFactor.
     *
     * @param point
     * @param scaleFactor
     * @param <T>
     * @return
     */
    private <T extends IPosition> T scalePointDistance(T point, double scaleFactor) {
        point.getPosition().x = point.getX() * scaleFactor;
        point.getPosition().y = point.getY() * scaleFactor;
        point.getPosition().z = point.getZ() * scaleFactor;
        return point;
    }

    /**
     * Checks if two feature point lists have the same types of feature points.
     * <p>
     * To use this method you need to supply ordered lists by Feature Point type
     * as parameters. Otherwise even if two lists contain the same feature points
     * it will return false.
     * <p>
     * Use sort function sortListByFeaturePointType on feature point lists
     * to get correct results.
     *
     * @param featurePointList1
     * @param featurePointList2
     * @return true if two sorted lists by the feature point type contain the
     * same feature point types
     */
    private boolean checkFeaturePointsType(List<FeaturePoint> featurePointList1, List<FeaturePoint> featurePointList2) {
        for (int i = 0; i < featurePointList1.size(); i++) {
            if (featurePointList1.get(i).getFeaturePointType().getType() != featurePointList2.get(i).getFeaturePointType().getType()) {
                System.out.print(featurePointList1.get(i).getFeaturePointType().getType());
                System.out.print(featurePointList2.get(i).getFeaturePointType().getType());
                return false;
            }
        }
        return true;
    }

    /**
     * Finds centroid from given feature point List
     *
     * @param featurePointList
     * @return centroid of feature points (Vector3F)
     */
    private Point3d findCentroidOfFeaturePoints(List<FeaturePoint> featurePointList) {
        float x = 0;
        float y = 0;
        float z = 0;
        for (FeaturePoint fp : featurePointList) {
            x += fp.getX();
            y += fp.getY();
            z += fp.getZ();
        }
        return new Point3d(x / featurePointList.size(), y / featurePointList.size(), z / featurePointList.size());
    }

    /**
     * Creates feature point map HashMap with key FeaturePoint.type and value FeaturePoint back from matrix.
     *
     * @param matrix
     * @param model
     * @return
     */
    private HashMap<Integer, FeaturePoint> createFeaturePointMapFromMatrix(
            SimpleMatrix matrix, ProcrustesAnalysisFaceModel model) {
        HashMap<Integer, FeaturePoint> map = new HashMap<>();
        for (int i = 0; i < matrix.numRows(); i++) {
            FeaturePoint featurePoint = new FeaturePoint(
                    matrix.get(i, 0),
                    matrix.get(i, 1),
                    matrix.get(i, 2),
                    model.getFeaturePointsMap().get(
                            model.getFeaturePointTypeCorrespondence().get(i)
                    ).getFeaturePointType()
            );
            map.put(model.getFeaturePointTypeCorrespondence().get(i), featurePoint);
        }
        return map;
    }

    /**
     * Creates matrix from given feature point list
     *
     * @param list
     * @return matrix
     */

    private <T extends IPosition> SimpleMatrix createMatrixFromList(List<T> list) {
        SimpleMatrix matrix = new SimpleMatrix(list.size(), 3);
        for (int i = 0; i < list.size(); i++) {
            matrix.set(i, 0, list.get(i).getPosition().x);
            matrix.set(i, 1, list.get(i).getPosition().y);
            matrix.set(i, 2, list.get(i).getPosition().z);
        }
        return matrix;
    }
    /**
     * Calculates distance of one feature point from another
     *
     * @param fp1
     * @param point
     * @return distance
     */
    private double calculateDistanceFromPoint(FeaturePoint fp1, Point3d point) {
        return Math.sqrt(
                (Math.pow(fp1.getX() - point.x, 2))
                        + (Math.pow(fp1.getY()- point.y, 2))
                        + (Math.pow(fp1.getZ()- point.z, 2)));
    }

}
