package dr.util;


/**
 * @author Zhenyu Zhang
 */
public class LKJCholeskyTransformConstrainedWithDiag extends LKJCholeskyTransformConstrained {

    private int dimCPC;
    private final int totalDimension;

    LKJCholeskyTransformConstrainedWithDiag(int dim) {
        super(dim);
        this.dimCPC = dim * (dim - 1) / 2;
        this.totalDimension = dim * (dim + 1) / 2;
    }

    @Override
    protected double[] transform(double[] values) {

        double[] choleskyFactor = subsetCholeskyOrCPCs(values);
        double[] diagonals = subsetDiagonals(values);
        double[] CPCs = super.transform(choleskyFactor, 0, dimCPC);

        return pasteTogether(CPCs, diagonals);
    }

    @Override
    //values: CPCs appended with log-transformed diagonals
    protected double[] inverse(double[] values) {

        assert values.length == totalDimension : "The transform function can only be applied to CPCs appended " +
                "with diagonals";
        for (int k = 0; k < dimCPC; k++) {
            assert values[k] <= 1.0 && values[k] >= -1.0 : "CPCs must be between -1.0 and 1.0";
        }

        double[] CPCs = subsetCholeskyOrCPCs(values);
        double[] diagonals = subsetDiagonals(values);
        double[] choleskyFactor = super.inverse(CPCs, 0, CPCs.length);
        return pasteTogether(choleskyFactor, diagonals);
    }

    @Override
    public double[][] computeJacobianMatrixInverse(double[] values) {

        double[] CPCs = subsetCholeskyOrCPCs(values);
        double[][] jacobian = super.computeJacobianMatrixInverse(CPCs);
        return appendIdentityMatrix(jacobian);
    }

    public double[] getGradientLogJacobianInverse(double[] values) {

        double[] CPCs = subsetCholeskyOrCPCs(values);
        double[] gradientLogJacobianInverse = super.getGradientLogJacobianInverse(CPCs);
        return pasteTogether(gradientLogJacobianInverse, new double[dim]);
    }

    private double[] subsetCholeskyOrCPCs(double[] values) { //todo: to ensure Cholesky factor comes first in "values"

        assert values.length == totalDimension;
        double[] choleskyOrCPC = new double[dimCPC];
        System.arraycopy(values, 0, choleskyOrCPC, 0, dimCPC);
        return choleskyOrCPC;
    }

    private double[] subsetDiagonals(double[] values) { //todo: to ensure Cholesky factor comes first in "values"

        assert values.length == totalDimension;
        double[] diagonals = new double[dim];
        System.arraycopy(values, dimCPC, diagonals, 0, dim);
        return diagonals;
    }

    private double[][] appendIdentityMatrix(double[][] jacobian) {

        assert jacobian.length == dimCPC;
        int length = dimCPC + dim;
        double[][] appendedJacobian = new double[length][length];

        for (int i = 0; i < length; i++) {
            for (int j = 0; j < length; j++) {
                if (i >= dimCPC || j >= dimCPC) {
                    if (i == j) {
                        appendedJacobian[i][j] = 1;
                    }
                } else {
                    appendedJacobian[i][j] = jacobian[i][j];
                }
            }
        }

        return appendedJacobian;
    }

    private double[] pasteTogether(double[] choleskyOrCPCs, double[] diagonals) {

        double[] concatenatedArray = new double[dimCPC + dim];
        System.arraycopy(choleskyOrCPCs, 0, concatenatedArray, 0, dimCPC);
        System.arraycopy(diagonals, 0, concatenatedArray, dimCPC, dim);
        return concatenatedArray;
    }

    public double[] updateDiagonalHessianLogDensity(double[] diagonalHessian, double[] gradient, double[] value, int from, int to) {
        throw new RuntimeException("Not yet implemented");
    }
}

@@ -394,32 +394,29 @@
 erride%0A    p
-rotected
+ublic
  double%5B%5D tr
@@ -1898,16 +1898,22 @@
 uble%5Bdim
+Vector
 %5D);%0A    
@@ -2442,16 +2442,22 @@
 uble%5Bdim
+Vector
 %5D;%0A     
@@ -2505,24 +2505,30 @@
 nals, 0, dim
+Vector
 );%0A        r
@@ -2694,16 +2694,22 @@
 PC + dim
+Vector
 ;%0A      
@@ -3329,16 +3329,22 @@
 PC + dim
+Vector
 %5D;%0A     
@@ -3405,32 +3405,32 @@
 ay, 0, dimCPC);%0A
-
         System.a
@@ -3482,16 +3482,22 @@
 CPC, dim
+Vector
 );%0A     

call the right dim