package dr.inference.operators.hmc;

import dr.inference.hmc.HessianWrtParameterProvider;
import dr.math.MathUtils;
import dr.math.matrixAlgebra.ReadableVector;
import dr.math.matrixAlgebra.WrappedVector;

/**
 * @author Marc A. Suchard
 */
public interface MassPreconditioner {

    // It looks like my interface is slowly reverting to something close to what Xiang had previously

    WrappedVector drawInitialMomentum();

    double getKineticEnergy(ReadableVector momentum);

    double getVelocity(int index, ReadableVector momentum);

    void storeSecant(ReadableVector gradient, ReadableVector position);

    class NoPreconditioning implements MassPreconditioner {

        final int dim;

        NoPreconditioning(int dim) { this.dim = dim; }

        @Override
        public WrappedVector drawInitialMomentum() {

            double[] momentum = new double[dim];

            for (int i = 0; i < dim; ++i) {
                momentum[i] = MathUtils.nextGaussian();
            }

            return new WrappedVector.Raw(momentum);
        }

        @Override
        public double getKineticEnergy(ReadableVector momentum) {
            double total = 0.0;
            for (int i = 0; i < dim; i++) {
                total += momentum.get(i) * momentum.get(i) / 2.0;
            }
            return total;
        }

        @Override
        public double getVelocity(int i, ReadableVector momentum) {
            return momentum.get(i);
        }

        @Override
        public void storeSecant(ReadableVector gradient, ReadableVector position) { }
    }

    abstract class HessianBased implements MassPreconditioner {
        final protected int dim;
        final protected HessianWrtParameterProvider hessian;
        final protected ReadableVector massTODO;

        HessianBased(HessianWrtParameterProvider hessian) {
            this.dim = hessian.getDimension();
            this.hessian = hessian;
            this.massTODO = null;
        }

        public void storeSecant(ReadableVector gradient, ReadableVector position) {
            // Do nothing
        }
    }

    class DiagonalPreconditioning extends HessianBased {

        WrappedVector mass = null;
        WrappedVector massInverse = null;

        DiagonalPreconditioning(int dim, HessianWrtParameterProvider hessian) {
            super(hessian);
            // TODO Make mass final (an const class ... which is easier to reason about)
        }

        @Override
        public WrappedVector drawInitialMomentum() {

            if (mass == null || massInverse == null) {
                updateMass();  // TODO Compute mass once per class instantiation
            }

            double[] momentum = new double[dim];

            for (int i = 0; i < dim; i++) {
                momentum[i] = MathUtils.nextGaussian() * Math.sqrt(mass.get(i));
            }

            return new WrappedVector.Raw(momentum);
        }

        private void updateMass() {
            double[] diagonalHessian = hessian.getDiagonalHessianLogDensity();
            double[] boundedMassInverse = boundMassInverse(diagonalHessian);
            for (int i = 0; i < dim; i++) {
                massInverse.set(i, boundedMassInverse[i]);
                mass.set(i, 1.0 / massInverse.get(i));
            }
        }

        private double[] boundMassInverse(double[] diagonalHessian) {

            double sum = 0.0;
            final double lowerBound = 1E-2; //TODO bad magic numbers
            final double upperBound = 1E2;
            double[] boundedMassInverse = new double[dim];

            for (int i = 0; i < dim; i++) {
                boundedMassInverse[i] = -1.0 / diagonalHessian[i];
                if (boundedMassInverse[i] < lowerBound) {
                    boundedMassInverse[i] = lowerBound;
                } else if (boundedMassInverse[i] > upperBound) {
                    boundedMassInverse[i] = upperBound;
                }
                sum += 1.0 / boundedMassInverse[i];
            }
            final double mean = sum / dim;
            for (int i = 0; i < dim; i++) {
                boundedMassInverse[i] = boundedMassInverse[i] * mean;
            }
            return boundedMassInverse;
        }

        @Override
        public double getKineticEnergy(ReadableVector momentum) {
            double total = 0.0;
            for (int i = 0; i < dim; i++) {
                total += momentum.get(i) * momentum.get(i) / (2.0 * mass.get(i));
            }
            return total;
        }

        @Override
        public double getVelocity(int i, ReadableVector momentum) {
            return momentum.get(i) / mass.get(i);
        }
    }

    class FullPreconditioning extends HessianBased {

        WrappedVector mass;
        WrappedVector massInverse;

        FullPreconditioning(HessianWrtParameterProvider hessian) {
            super(hessian);
        }

        @Override
        public WrappedVector drawInitialMomentum() {
            return null;
        }

        @Override
        public double getKineticEnergy(ReadableVector momentum) {
            return 0.0;
        }

        @Override
        public double getVelocity(int i, ReadableVector momentum) {
            return 0.0; // M^{-1}_{i} momentum
        }
    }

    abstract // TODO Implement interface
    class SecantPreconditioing extends HessianBased {

        SecantPreconditioing(HessianWrtParameterProvider hessian) {
            super(hessian);
        }

        @Override
        public void storeSecant(ReadableVector gradient, ReadableVector position) {
            // TODO Do something
        }
    }
}

@@ -4649,32 +4649,54 @@
        %7D%0A    %7D%0A%0A
+    // TODO Implement%0A
     class FullPr
@@ -5367,32 +5367,143 @@
 HessianBased %7B%0A%0A
+        private final Secant%5B%5D queue;%0A        private int secantIndex;%0A        private int secantUpdateCount;%0A%0A
         SecantPr
@@ -5542,32 +5542,48 @@
 Provider hessian
+, int secantSize
 ) %7B%0A            
@@ -5589,32 +5589,156 @@
  super(hessian);
+%0A%0A            this.queue = new Secant%5BsecantSize%5D;%0A            this.secantIndex = 0;%0A            this.secantUpdateCount = 0;
 %0A        %7D%0A%0A    
@@ -5851,28 +5851,518 @@
     
-// TODO Do something
+queue%5BsecantIndex%5D = new Secant(gradient, position);%0A            secantIndex = (secantIndex + 1) %25 queue.length;%0A            ++secantUpdateCount;%0A        %7D%0A%0A        private class Secant %7B // TODO Inner class because we may change the storage depending on efficiency%0A%0A            ReadableVector gradient;%0A            ReadableVector position;%0A%0A            Secant(ReadableVector gradient, ReadableVector position) %7B%0A                this.gradient = gradient;%0A                this.position = position;%0A            %7D
 %0A   

add secant circular queue