/*
 * PreconditionHandlerParser.java
 *
 * Copyright (c) 2002-2017 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.inferencexml.operators.hmc;

import dr.inference.hmc.GradientWrtParameterProvider;
import dr.inference.model.PriorPreconditioningProvider;
import dr.inference.operators.hmc.MassPreconditionScheduler;
import dr.inference.operators.hmc.MassPreconditioner;
import dr.inference.operators.hmc.MassPreconditioningOptions;
import dr.inference.operators.hmc.PreconditionHandler;
import dr.util.Transform;
import dr.xml.*;

import static dr.util.Transform.Util.parseTransform;

/**
 * @author Xiang Ji
 * @author Marc A. Suchard
 */
public class PreconditionHandlerParser extends AbstractXMLObjectParser {

    private final static String PRECONDITIONING = "preconditioning";
    private final static String PRECONDITIONING_SCHEDULE = "preconditioningSchedule";
    private final static String PRECONDITIONING_UPDATE_FREQUENCY = "preconditioningUpdateFrequency";
    final static String PRECONDITIONING_MAX_UPDATE = "preconditioningMaxUpdate";
    final static String PRECONDITIONING_DELAY = "preconditioningDelay";
    private final static String PRECONDITIONING_MEMORY = "preconditioningMemory";
    private final static String PRECONDITIONER = "preconditioner";
    private final static String PRECONDITIONING_GUESS_INIT_MASS = "guessInitialMass";
    private final static String PRECONDITION_HANDLER = "preconditionHandler";

    public static PreconditionHandler parsePreconditionHandler(XMLObject xo) throws XMLParseException {
        MassPreconditioner.Type preconditioningType;
        if (xo.hasChildNamed(PRECONDITIONER)) {
            preconditioningType = MassPreconditioner.Type.PRIOR_DIAGONAL;
        } else {
            preconditioningType = parsePreconditioning(xo);
        }

        MassPreconditionScheduler.Type preconditionSchedulerType = parsePreconditionScheduler(xo, preconditioningType);

        int preconditioningUpdateFrequency = xo.getAttribute(PRECONDITIONING_UPDATE_FREQUENCY, 0);

        int preconditioningMaxUpdate = xo.getAttribute(PRECONDITIONING_MAX_UPDATE, 0);

        int preconditioningDelay = xo.getAttribute(PRECONDITIONING_DELAY, 0);

        int preconditioningMemory = xo.getAttribute(PRECONDITIONING_MEMORY, 0);

        boolean guessInitialMass = xo.getAttribute(PRECONDITIONING_GUESS_INIT_MASS, false);

        Transform transform = parseTransform(xo);

        GradientWrtParameterProvider derivative =
                (GradientWrtParameterProvider) xo.getChild(GradientWrtParameterProvider.class);

        MassPreconditioningOptions preconditioningOptions = new MassPreconditioningOptions.Default(preconditioningUpdateFrequency, preconditioningMaxUpdate, preconditioningDelay, preconditioningMemory, guessInitialMass);
        MassPreconditioner preconditioner;

        if (xo.hasChildNamed(PRECONDITIONER)) {
            PriorPreconditioningProvider priorPreconditioningProvider = (PriorPreconditioningProvider) xo.getChild(PRECONDITIONER).getChild(PriorPreconditioningProvider.class);
            if (priorPreconditioningProvider !=  null) {
                preconditioner = new MassPreconditioner.PriorPreconditioner(priorPreconditioningProvider, transform);
            } else {
                throw new XMLParseException("Unknown preconditioner specified");
            }
        } else {
            preconditioner = preconditioningType.factory(derivative, transform, preconditioningOptions);
        }
        PreconditionHandler preconditionHandler = new PreconditionHandler(preconditioner, preconditioningOptions, preconditionSchedulerType);

        return preconditionHandler;
    }
    @Override
    public Object parseXMLObject(XMLObject xo) throws XMLParseException {
        return parsePreconditionHandler(xo);
    }

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

    protected final XMLSyntaxRule[] rules = {
            new ElementRule(Transform.MultivariableTransformWithParameter.class, true),
            new ElementRule(GradientWrtParameterProvider.class),
            new ElementRule(PRECONDITIONER, new XMLSyntaxRule[]{
                    new XORRule(
                            new ElementRule(MassPreconditioner.class),
                            new ElementRule(PriorPreconditioningProvider.class)
                    ),
            }, true),
    };

    static MassPreconditioner.Type parsePreconditioning(XMLObject xo) throws XMLParseException {

        return MassPreconditioner.Type.parseFromString(
                xo.getAttribute(PRECONDITIONING, MassPreconditioner.Type.NONE.getName())
        );
    }

    static MassPreconditionScheduler.Type parsePreconditionScheduler(XMLObject xo,
                                                                     MassPreconditioner.Type preconditioningType) throws XMLParseException {
        if (preconditioningType == MassPreconditioner.Type.NONE) {
            return MassPreconditionScheduler.Type.NONE;
        } else {
            return MassPreconditionScheduler.Type.parseFromString(
                    xo.getAttribute(PRECONDITIONING_SCHEDULE, MassPreconditionScheduler.Type.DEFAULT.getName()));
        }
    }

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

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

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

@@ -2253,86 +2253,8 @@
 ss%22;
-%0A    private final static String PRECONDITION_HANDLER = %22preconditionHandler%22;
 %0A%0A  
@@ -4695,28 +4695,13 @@
 urn 
-new XMLSyntaxRule%5B0%5D
+rules
 ;%0A  
@@ -4715,15 +4715,13 @@
   pr
-otec
+iva
 te
-d
  fin
@@ -4749,16 +4749,17 @@
 les = %7B%0A
+%0A
         
@@ -4834,24 +4834,58 @@
 ass, true),%0A
+%0A            new XORRule(%0A        
             
@@ -4941,32 +4941,40 @@
 s),%0A            
+     
+   
 new ElementRule(
@@ -5022,32 +5022,40 @@
                 
+        
 new XORRule(%0A   
@@ -5071,32 +5071,40 @@
                 
+ 
+       
 new ElementRule(
@@ -5122,32 +5122,40 @@
 itioner.class),%0A
+        
                 
@@ -5242,32 +5242,56 @@
     
-),%0A            %7D, true),
+        ),%0A                    %7D)%0A            )%0A
 %0A   
@@ -6382,16 +6382,11 @@
 TION
-_HANDLER
+ING
 ;%0A  

modify parser rules