package dr.inference.operators;

import dr.evomodel.continuous.LatentFactorModel;
import dr.inference.distribution.DistributionLikelihood;
import dr.inference.model.DiagonalMatrix;
import dr.inference.model.MatrixParameter;
import dr.math.MathUtils;
import dr.math.distributions.GammaDistribution;

/**
 * Created by max on 6/12/14.
 */
public class LatentFactorModelPrecisionGibbsOperator extends SimpleMCMCOperator implements GibbsOperator {
    private double[] FacXLoad;
    private double[] residual;
    private int factorDimension;
    private LatentFactorModel LFM;
    private GammaDistribution prior;
    private boolean randomScan;

    public LatentFactorModelPrecisionGibbsOperator(LatentFactorModel LFM, DistributionLikelihood prior, double weight, boolean randomScan) {
        setWeight(weight);
        this.LFM=LFM;
        this.prior=(GammaDistribution)prior.getDistribution();
        this.randomScan=randomScan;
        setupParameters();
        FacXLoad=new double[LFM.getFactors().getColumnDimension()];
        residual=new double[LFM.getFactors().getColumnDimension()];
    }

    private void setupParameters(){
        factorDimension=LFM.getFactorDimension();
    }


    private void setPrecision(int i){
        MatrixParameter factors=LFM.getFactors();
        MatrixParameter loadings=LFM.getLoadings();
        DiagonalMatrix precision= (DiagonalMatrix) LFM.getColumnPrecision();
        MatrixParameter data=LFM.getScaledData();
        double shape=prior.getShape()+factors.getColumnDimension()/2;
        double sum;
        for (int j = 0; j <factors.getColumnDimension(); j++) {
            sum=0;
            for (int k = 0; k < factorDimension; k++) {
                sum+=factors.getParameterValue(k,j)*loadings.getParameterValue(k,i);
            }
            FacXLoad[j]=sum;
            residual[j]=data.getParameterValue(i,j)-FacXLoad[j];
        }
        sum=0;
        for (int j = 0; j <factors.getColumnDimension() ; j++) {
            sum+=residual[j]*residual[j];
        }
        double scale=1/(1/prior.getScale()+sum/2);
        double nextPrecision=GammaDistribution.nextGamma(shape, scale);
        precision.getParameter(0).setParameterValueQuietly(i, nextPrecision);
    }

    @Override
    public int getStepCount() {
        return 0;
    }

    @Override
    public String getPerformanceSuggestion() {
        return null;
    }

    @Override
    public String getOperatorName() {
        return null;
    }

    @Override
    public double doOperation() throws OperatorFailedException {
        if(LFM.getFactorDimension()!=factorDimension){
            setupParameters();
        }
        if(!randomScan){
        for (int i = 0; i <LFM.getColumnPrecision().getColumnDimension(); i++) {
            setPrecision(i);
        }}
        else{
            int i= MathUtils.nextInt(LFM.getColumnPrecision().getColumnDimension());
            setPrecision(i);
        }
        LFM.getColumnPrecision().getParameter(0).fireParameterChangedEvent();
        
        
        return 0;
    }
}

@@ -437,16 +437,18 @@
 rator %7B%0A
+//
     priv
@@ -470,16 +470,18 @@
 cXLoad;%0A
+//
     priv
@@ -507,41 +507,8 @@
 al;%0A
-    private int factorDimension;%0A
     
@@ -606,16 +606,42 @@
 domScan;
+%0A    private double shape;
 %0A%0A    pu
@@ -927,35 +927,11 @@
 an;%0A
-        setupParameters();%0A
+%0A//
     
@@ -990,24 +990,26 @@
 mension()%5D;%0A
+//
         resi
@@ -1072,84 +1072,66 @@
     
-%7D%0A%0A
     
-private void setupParameters()%7B%0A        factorDimension=LFM.getFactor
+shape=this.prior.getShape()+LFM.getFactors().getColumn
 Dime
@@ -1133,24 +1133,27 @@
 nDimension()
+*.5
 ;%0A    %7D%0A%0A%0A  
@@ -1150,16 +1150,18 @@
     %7D%0A%0A%0A
+%0A%0A
     priv
@@ -1438,81 +1438,12 @@
 ble 
-shape=prior.getShape()+factors.getColumnDimension()/2;%0A        double sum
+di=0
 ;%0A  
@@ -1508,32 +1508,39 @@
 ) %7B%0A            
+double 
 sum=0;%0A         
@@ -1564,24 +1564,32 @@
 ; k %3C factor
+s.getRow
 Dimension; k
@@ -1585,16 +1585,18 @@
 imension
+()
 ; k++) %7B
@@ -1687,32 +1687,118 @@
 ;%0A            %7D%0A
+            double temp=data.getParameterValue(i,j)-sum;%0A            di+=temp*temp;%0A//
             FacX
@@ -1806,24 +1806,26 @@
 oad%5Bj%5D=sum;%0A
+//
             
@@ -1873,24 +1873,26 @@
 acXLoad%5Bj%5D;%0A
+//
         %7D%0A  
@@ -1889,24 +1889,26 @@
       %7D%0A
+//
         
 sum=0;%0A 
@@ -1899,23 +1899,32 @@
         
+double 
 sum=0;%0A
+//
         
@@ -1972,32 +1972,34 @@
 sion() ; j++) %7B%0A
+//
             sum+
@@ -2060,11 +2060,15 @@
 le=1
+.0
 /(1
+.0
 /pri
@@ -2085,13 +2085,13 @@
 e()+
-sum/2
+di*.5
 );%0A 
@@ -2569,103 +2569,8 @@
 n %7B%0A
-        if(LFM.getFactorDimension()!=factorDimension)%7B%0A            setupParameters();%0A        %7D
 %0A   

Removes some unnecessary computations