package dr.evomodel.continuous.hmc;

import dr.evolution.tree.Tree;
import dr.evolution.tree.TreeTrait;
import dr.evomodel.treedatalikelihood.TreeDataLikelihood;
import dr.evomodel.treedatalikelihood.continuous.ContinuousDataLikelihoodDelegate;
import dr.evomodel.treedatalikelihood.continuous.IntegratedFactorAnalysisLikelihood;
import dr.evomodel.treedatalikelihood.preorder.WrappedMeanPrecision;
import dr.evomodel.treedatalikelihood.preorder.WrappedTipFullConditionalDistributionDelegate;
import dr.inference.hmc.GradientWrtParameterProvider;
import dr.inference.model.Likelihood;
import dr.inference.model.Parameter;
import dr.xml.*;

import java.util.List;

import static dr.evomodel.continuous.hmc.LinearOrderTreePrecisionTraitProductProvider.castTreeTrait;

/**
 * @author Marc A. Suchard
 * @author Andrew Holbrook
 */
public class IntegratedLoadingsGradient implements GradientWrtParameterProvider {

    private final TreeTrait<List<WrappedMeanPrecision>> fullConditionalDensity;
    private final IntegratedFactorAnalysisLikelihood factorAnalysisLikelihood;
    private final int dim;
    private final Tree tree;

    IntegratedLoadingsGradient(TreeDataLikelihood treeDataLikelihood,
                               ContinuousDataLikelihoodDelegate likelihoodDelegate,
                               IntegratedFactorAnalysisLikelihood factorAnalysisLikelihood) {

        this.factorAnalysisLikelihood = factorAnalysisLikelihood;

        String traitName = factorAnalysisLikelihood.getModelName(); // TODO Is this correct?

        String fcdName = WrappedTipFullConditionalDistributionDelegate.getName(traitName);
        if (treeDataLikelihood.getTreeTrait(fcdName) == null) {
            likelihoodDelegate.addWrappedFullConditionalDensityTrait(traitName);
        }

        this.fullConditionalDensity = castTreeTrait(treeDataLikelihood.getTreeTrait(fcdName));
        this.tree = treeDataLikelihood.getTree();
        this.dim = factorAnalysisLikelihood.getTraitDimension();
    }

    @Override
    public Likelihood getLikelihood() {
        return null;    // TODO Return CompoundLikelihood{factorAnalysisLikelihood, treeDataLikelihood}
    }

    @Override
    public Parameter getParameter() {
        return factorAnalysisLikelihood.getLoadings(); // TODO May need to work with vec(L)
    }

    @Override
    public int getDimension() {
        return dim * dim; // TODO May need to work with vec(L)
    }

    @Override
    public double[] getGradientLogDensity() {

        for (int i = 0; i < tree.getExternalNodeCount(); ++i) {
            // TODO Work with fullConditionalDensity
            List<WrappedMeanPrecision> statistics = fullConditionalDensity.getTrait(tree, tree.getExternalNode(i));

            for (WrappedMeanPrecision meanPrecision : statistics) {
                meanPrecision.getMean();
                meanPrecision.getPrecision();
            }
        }

        return new double[0];
    }

    private static final String PARSER_NAME = "integratedFactorAnalysisLoadingsGradient";

    public static AbstractXMLObjectParser PARSER = new AbstractXMLObjectParser() {
        @Override
        public Object parseXMLObject(XMLObject xo) throws XMLParseException {
            return null;
        }

        @Override
        public XMLSyntaxRule[] getSyntaxRules() {
            return rules;
        }

        @Override
        public String getParserDescription() {
            return "Generates a gradient provider for the loadings matrix when factors are integrated out";
        }

        @Override
        public Class getReturnType() {
            return IntegratedLoadingsGradient.class;
        }

        @Override
        public String getParserName() {
            return PARSER_NAME;
        }

        private final XMLSyntaxRule[] rules = new XMLSyntaxRule[] {
                new ElementRule(IntegratedFactorAnalysisLikelihood.class),
                new ElementRule(TreeDataLikelihood.class),
        };
    };
}

@@ -97,16 +97,78 @@
 eTrait;%0A
+import dr.evomodel.treedatalikelihood.DataLikelihoodDelegate;%0A
 import d
@@ -673,24 +673,68 @@
 .Parameter;%0A
+import dr.math.matrixAlgebra.WrappedVector;%0A
 import dr.xm
@@ -1006,16 +1006,28 @@
 Provider
+, Reportable
  %7B%0A%0A    
@@ -1238,24 +1238,32 @@
  tree;%0A%0A    
+private 
 IntegratedLo
@@ -2609,24 +2609,81 @@
 ensity() %7B%0A%0A
+        double%5B%5D gradient = new double%5BgetDimension()%5D;%0A%0A
         for 
@@ -2838,16 +2838,36 @@
 istics =
+%0A                   
  fullCon
@@ -3116,35 +3116,497 @@
     
-return new double%5B0%5D;%0A    %7D
+if (DEBUG) %7B%0A            System.err.println(getReport(gradient));%0A        %7D%0A%0A        return gradient;%0A    %7D%0A    %0A    @Override%0A    public String getReport() %7B%0A        return getReport(getGradientLogDensity());%0A    %7D%0A%0A    private String getReport(double%5B%5D gradient) %7B%0A%0A        String result = new WrappedVector.Raw(gradient).toString();%0A%0A         if (DEBUG) %7B%0A             result += %22Debug info: %5Cn%22;%0A         %7D%0A%0A         return result;%0A    %7D%0A%0A    private static final boolean DEBUG = true;
 %0A%0A  
@@ -3873,16 +3873,17 @@
 ption %7B%0A
+%0A
         
@@ -3886,27 +3886,964 @@
         
+T
 re
-turn null
+eDataLikelihood treeDataLikelihood = (TreeDataLikelihood)%0A                    xo.getChild(TreeDataLikelihood.class);%0A%0A            IntegratedFactorAnalysisLikelihood factorAnalysis = (IntegratedFactorAnalysisLikelihood)%0A                    xo.getChild(IntegratedFactorAnalysisLikelihood.class);%0A%0A            DataLikelihoodDelegate likelihoodDelegate = treeDataLikelihood.getDataLikelihoodDelegate();%0A%0A            if (!(likelihoodDelegate instanceof ContinuousDataLikelihoodDelegate)) %7B%0A                throw new XMLParseException(%22TODO%22);%0A            %7D%0A%0A            ContinuousDataLikelihoodDelegate continuousDataLikelihoodDelegate =%0A                    (ContinuousDataLikelihoodDelegate) likelihoodDelegate;%0A%0A            // TODO Check dimensions, parameters, etc.%0A%0A            return new IntegratedLoadingsGradient(%0A                    treeDataLikelihood,%0A                    continuousDataLikelihoodDelegate,%0A                    factorAnalysis)
 ;%0A      

add simple parser for loadings gradient