package dr.util;

public class LKJCholeskyTransformConstrainedWithDiag extends LKJCholeskyTransformConstrained {

    protected int CPCdimension;

    public LKJCholeskyTransformConstrainedWithDiag(int dim) {
        super(dim);
        this.CPCdimension = dim * (dim - 1) / 2;
    }

    @Override
    public double[] transform(double[] values, int from, int to) {

        assert from == 0 && to == values.length : "The transform function can only be applied to the whole array of " +
                "values.";

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

        return pasteTogether(CPCs, diagonals);
    }

    @Override
    //values: CPCs appended with log-transformed diagonals
    public double[] inverse(double[] values, int from, int to) {

        assert from == 0 && to == values.length : "The transform function can only be applied to the whole array of " +
                "values.";
        assert values.length == dim * (dim + 1) / 2 : "The transform function can only be applied to CPCs appended " +
                "with diagonals";
        for (int k = 0; k < dim * (dim - 1) / 2; 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);

        return pasteTogether(CPCs, diagonals);
    }

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

        double[] choleskyFactor = subsetCholeskyOrCPCs(values);

        double[][] jacobian = super.computeJacobianMatrixInverse(choleskyFactor);

        return appendZeros(jacobian);
    }

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

        assert values.length == dim * (dim + 1) / 2;
        double[] choleskyFactor = new double[CPCdimension];
        System.arraycopy(values, 0, choleskyFactor, 0, CPCdimension);
        return choleskyFactor;
    }

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

        assert values.length == dim * (dim + 1) / 2;
        double[] diagonals = new double[dim];
        System.arraycopy(values, CPCdimension, diagonals, 0, dim);
        return diagonals;
    }


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

        assert jacobian.length == CPCdimension;
        int length = dim * (dim - 1) / 2 + dim;
        double[][] appendedJacobian = new double[length][length];

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

        return appendedJacobian;
    }

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

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

@@ -1618,38 +1618,28 @@
    double%5B%5D 
-choleskyFactor
+CPCs
  = subsetCho
@@ -1651,33 +1651,32 @@
 OrCPCs(values);%0A
-%0A
         double%5B%5D
@@ -1728,26 +1728,282 @@
 rse(
-choleskyFactor
+CPCs);%0A        return appendIdentityMatrix(jacobian);%0A    %7D%0A%0A    public double%5B%5D getGradientLogJacobianInverse(double%5B%5D values) %7B%0A%0A        double%5B%5D CPCs = subsetCholeskyOrCPCs(values);%0A        double%5B%5D gradientLogJacobianInverse = super.getGradientLogJacobianInverse(CPCs
 );%0A
-%0A
     
@@ -2017,28 +2017,65 @@
 urn 
-appendZeros(jacobian
+pasteTogether(gradientLogJacobianInverse, new double%5Bdim%5D
 );%0A 
@@ -2278,22 +2278,21 @@
 cholesky
-Factor
+OrCPC
  = new d
@@ -2348,38 +2348,37 @@
 ues, 0, cholesky
-Factor
+OrCPC
 , 0, CPCdimensio
@@ -2404,22 +2404,21 @@
 cholesky
-Factor
+OrCPC
 ;%0A    %7D%0A
@@ -2732,17 +2732,16 @@
 %0A    %7D%0A%0A
-%0A
     priv
@@ -2765,13 +2765,22 @@
 pend
-Zeros
+IdentityMatrix
 (dou
@@ -2869,35 +2869,28 @@
 ength = 
-dim * (dim - 1) / 2
+CPCdimension
  + dim;%0A
@@ -3075,27 +3075,81 @@
  %3E= 
-dim %7C%7C j %3E= dim) %7B%0A
+CPCdimension %7C%7C j %3E= CPCdimension) %7B%0A                    if(i == j)%7B%0A    
     
@@ -3189,18 +3189,40 @@
 i%5D%5Bj%5D = 
-0;
+1;%0A                    %7D
 %0A       

correct appended jacobian matrix