/*
 * TransformedCompoundMatrix.java
 *
 * Copyright (c) 2002-2018 Alexei Drummond, Andrew Rambaut and Marc Suchard
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST 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 for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.inference.model;

import dr.math.matrixAlgebra.WrappedMatrix;
import dr.util.Transform;

/**
 * @author Marc Suchard
 * @author Paul Bastide
 */
abstract public class AbstractTransformedCompoundMatrix extends MatrixParameter {

    final Parameter diagonalParameter;
    final Parameter offDiagonalParameter;

    protected final int dim;

    AbstractTransformedCompoundMatrix(Parameter diagonals, Parameter offDiagonal) {
        super(MATRIX_PARAMETER);
        diagonalParameter = diagonals;
        dim = diagonalParameter.getDimension();
        offDiagonalParameter = offDiagonal;
        addParameter(diagonalParameter);
        addParameter(offDiagonalParameter);
    }

    AbstractTransformedCompoundMatrix(Parameter diagonals, Parameter offDiagonal,
                                      Transform.MultivariableTransform transform, Boolean inverse) {
        super(MATRIX_PARAMETER);
        diagonalParameter = diagonals;
        dim = diagonalParameter.getDimension();
        offDiagonalParameter =
                (transform == null) ? offDiagonal: new TransformedMultivariateParameter(offDiagonal, transform, inverse);
        addParameter(diagonalParameter);
        addParameter(offDiagonalParameter);
    }

    @Override
    public int getDimension() {
        return getColumnDimension() * getRowDimension();
    }

    @Override
    public double getParameterValue(int index) {
        final int dim = getColumnDimension();
        return getParameterValue(index / dim, index % dim);
    }

    @Override
    public double getParameterUntransformedValue(int index) {
        final int dim = getColumnDimension();
        return getParameterUntransformedValue(index / dim, index % dim);
    }

    private double getParameterUntransformedValue(int row, int col) {
        if (row != col) {
            return offDiagonalParameter.getParameterUntransformedValue(getUpperTriangularIndex(row, col));
        }
        return diagonalParameter.getParameterUntransformedValue(row);
    }

    @Override
    public void setParameterUntransformedValue(int index, double a) {
        final int dim = getColumnDimension();
        setParameterUntransformedValue(index / dim, index % dim, a);
    }

    private void setParameterUntransformedValue(int row, int col, double a) {
        if (row != col) {
            offDiagonalParameter.setParameterUntransformedValue(getUpperTriangularIndex(row, col), a);
        } else {
            diagonalParameter.setParameterUntransformedValue(row, a);
        }
    }

    @Override
    public int getColumnDimension() {
        return diagonalParameter.getDimension();
    }

    @Override
    public int getRowDimension() {
        return diagonalParameter.getDimension();
    }

    @Override
    public void setParameterValue(int index, double a) {
        throw new RuntimeException("Do not set entries of a TransformedCompoundMatrix directly");
    }

    @Override
    public void setParameterValue(int row, int column, double a) {
        throw new RuntimeException("Do not set entries of a TransformedCompoundMatrix directly");
    }

    public double[] getDiagonal() {
        return diagonalParameter.getParameterValues();
    }

    public Parameter getDiagonalParameter() {
        return diagonalParameter;
    }

    public Parameter getOffDiagonalParameter() {
        return offDiagonalParameter;
    }

    public Parameter getUntransformedOffDiagonalParameter() {
        if (offDiagonalParameter instanceof TransformedMultivariateParameter) {
            return ((TransformedMultivariateParameter) offDiagonalParameter).getUntransformedParameter();
        }
        return offDiagonalParameter;
    }

    public CompoundParameter getUntransformedCompoundParameter(){
        CompoundParameter param = new CompoundParameter(this.getParameterName());
        param.addParameter(getDiagonalParameter());
        param.addParameter(getUntransformedOffDiagonalParameter());
        return param;
    }

    abstract double[] updateGradientDiagonal(double[] gradient);

    abstract double[] updateGradientOffDiagonal(double[] gradient);

    @Override
    public double[] getAttributeValue() {
        double[] stats = new double[dim * dim];
        int index = 0;
        for (int i = 0; i < dim; i++) {
            for (int j = 0; j < dim; j++) {
                stats[index] = getParameterValue(i, j);
                index++;
            }
        }
        return stats;
    }

    @Override
    public String getDimensionName(int index) {
        int dim = getColumnDimension();
        String row = Integer.toString(index / dim);
        String col = Integer.toString(index % dim);

        return getId() + row + col;
    }

    @Override
    public String getReport() {
        return new WrappedMatrix.ArrayOfArray(getParameterAsMatrix()).toString();
    }

    int getUpperTriangularIndex(int i, int j) {
        assert i != j;
        if (i < j) {
            return upperTriangularTransformation(i, j);
        } else {
            return upperTriangularTransformation(j, i);
        }
    }

    private int upperTriangularTransformation(int i, int j) {
        return i * (2 * dim - i - 1) / 2 + (j - i - 1);
    }
}

@@ -1318,16 +1318,77 @@
 meter;%0A%0A
+    final CompoundParameter untransformedCompoundParameter;%0A%0A
     prot
@@ -1408,16 +1408,16 @@
 nt dim;%0A
-
 %0A    Abs
@@ -1733,32 +1733,167 @@
 gonalParameter);
+%0A%0A        untransformedCompoundParameter = new CompoundParameter(this.getParameterName());%0A        setUntransformedCompoundParameter();
 %0A    %7D%0A%0A    Abst
@@ -2373,32 +2373,32 @@
 onalParameter);%0A
-
         addParam
@@ -2420,24 +2420,391 @@
 lParameter);
+%0A%0A        untransformedCompoundParameter = new CompoundParameter(this.getParameterName());%0A        setUntransformedCompoundParameter();%0A    %7D%0A%0A    private void setUntransformedCompoundParameter() %7B%0A        untransformedCompoundParameter.addParameter(getDiagonalParameter());%0A        untransformedCompoundParameter.addParameter(getUntransformedOffDiagonalParameter());
 %0A    %7D%0A%0A    
@@ -5327,222 +5327,45 @@
     
-CompoundParameter param = new CompoundParameter(this.getParameterName());%0A        param.addParameter(getDiagonalParameter());%0A        param.addParameter(getUntransformedOffDiagonalParameter());%0A        return param
+return untransformedCompoundParameter
 ;%0A  

Fix memory leak in getting parameter for PrecisionGradient.