/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.commons.math3.ode;

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.commons.math3.exception.DimensionMismatchException;
import org.apache.commons.math3.exception.MathIllegalArgumentException;
import org.apache.commons.math3.exception.MaxCountExceededException;
import org.apache.commons.math3.exception.util.LocalizedFormats;

This class defines a set of secondary equations to compute the Jacobian matrices with respect to the initial state vector and, if any, to some parameters of the primary ODE set. 
 It is intended to be packed into an ExpandableStatefulODE in conjunction with a primary set of ODE, which may be: 

a FirstOrderDifferentialEquations
a MainStateJacobianProvider

In order to compute Jacobian matrices with respect to some parameters of the
primary ODE set, the following parameter Jacobian providers may be set:

a ParameterJacobianProvider
a ParameterizedODE


See Also:
ExpandableStatefulODE
FirstOrderDifferentialEquations
MainStateJacobianProvider
ParameterJacobianProvider
ParameterizedODE
Since:
3.0/**
 * This class defines a set of {@link SecondaryEquations secondary equations} to
 * compute the Jacobian matrices with respect to the initial state vector and, if
 * any, to some parameters of the primary ODE set.
 * <p>
 * It is intended to be packed into an {@link ExpandableStatefulODE}
 * in conjunction with a primary set of ODE, which may be:
 * <ul>
 * <li>a {@link FirstOrderDifferentialEquations}</li>
 * <li>a {@link MainStateJacobianProvider}</li>
 * </ul>
 * In order to compute Jacobian matrices with respect to some parameters of the
 * primary ODE set, the following parameter Jacobian providers may be set:
 * <ul>
 * <li>a {@link ParameterJacobianProvider}</li>
 * <li>a {@link ParameterizedODE}</li>
 * </ul>
 * </p>
 *
 * @see ExpandableStatefulODE
 * @see FirstOrderDifferentialEquations
 * @see MainStateJacobianProvider
 * @see ParameterJacobianProvider
 * @see ParameterizedODE
 *
 * @since 3.0
 */
public class JacobianMatrices {

    
Expandable first order differential equation. /** Expandable first order differential equation. */
    private ExpandableStatefulODE efode;

    
Index of the instance in the expandable set. /** Index of the instance in the expandable set. */
    private int index;

    
FODE with exact primary Jacobian computation skill. /** FODE with exact primary Jacobian computation skill. */
    private MainStateJacobianProvider jode;

    
FODE without exact parameter Jacobian computation skill. /** FODE without exact parameter Jacobian computation skill. */
    private ParameterizedODE pode;

    
Main state vector dimension. /** Main state vector dimension. */
    private int stateDim;

    
Selected parameters for parameter Jacobian computation. /** Selected parameters for parameter Jacobian computation. */
    private ParameterConfiguration[] selectedParameters;

    
FODE with exact parameter Jacobian computation skill. /** FODE with exact parameter Jacobian computation skill. */
    private List<ParameterJacobianProvider> jacobianProviders;

    
Parameters dimension. /** Parameters dimension. */
    private int paramDim;

    
Boolean for selected parameters consistency. /** Boolean for selected parameters consistency. */
    private boolean dirtyParameter;

    
State and parameters Jacobian matrices in a row. /** State and parameters Jacobian matrices in a row. */
    private double[] matricesData;

    
Simple constructor for a secondary equations set computing Jacobian matrices.
 Parameters must belong to the supported ones given by Parameterizable.getParametersNames(), so the primary set of differential equations must be Parameterizable. 
Note that each selection clears the previous selected parameters.
Params:
fode – the primary first order differential equations set to extend
hY – step used for finite difference computation with respect to state vector
parameters – parameters to consider for Jacobian matrices processing
(may be null if parameters Jacobians is not desired)
Throws:
DimensionMismatchException – if there is a dimension mismatch between the steps array hY and the equation dimension/** Simple constructor for a secondary equations set computing Jacobian matrices.
     * <p>
     * Parameters must belong to the supported ones given by {@link
     * Parameterizable#getParametersNames()}, so the primary set of differential
     * equations must be {@link Parameterizable}.
     * </p>
     * <p>Note that each selection clears the previous selected parameters.</p>
     *
     * @param fode the primary first order differential equations set to extend
     * @param hY step used for finite difference computation with respect to state vector
     * @param parameters parameters to consider for Jacobian matrices processing
     * (may be null if parameters Jacobians is not desired)
     * @exception DimensionMismatchException if there is a dimension mismatch between
     * the steps array {@code hY} and the equation dimension
     */
    public JacobianMatrices(final FirstOrderDifferentialEquations fode, final double[] hY,
                            final String... parameters)
        throws DimensionMismatchException {
        this(new MainStateJacobianWrapper(fode, hY), parameters);
    }

    
Simple constructor for a secondary equations set computing Jacobian matrices.
 Parameters must belong to the supported ones given by Parameterizable.getParametersNames(), so the primary set of differential equations must be Parameterizable. 
Note that each selection clears the previous selected parameters.
Params:
jode – the primary first order differential equations set to extend
parameters – parameters to consider for Jacobian matrices processing
(may be null if parameters Jacobians is not desired)/** Simple constructor for a secondary equations set computing Jacobian matrices.
     * <p>
     * Parameters must belong to the supported ones given by {@link
     * Parameterizable#getParametersNames()}, so the primary set of differential
     * equations must be {@link Parameterizable}.
     * </p>
     * <p>Note that each selection clears the previous selected parameters.</p>
     *
     * @param jode the primary first order differential equations set to extend
     * @param parameters parameters to consider for Jacobian matrices processing
     * (may be null if parameters Jacobians is not desired)
     */
    public JacobianMatrices(final MainStateJacobianProvider jode,
                            final String... parameters) {

        this.efode = null;
        this.index = -1;

        this.jode = jode;
        this.pode = null;

        this.stateDim = jode.getDimension();

        if (parameters == null) {
            selectedParameters = null;
            paramDim = 0;
        } else {
            this.selectedParameters = new ParameterConfiguration[parameters.length];
            for (int i = 0; i < parameters.length; ++i) {
                selectedParameters[i] = new ParameterConfiguration(parameters[i], Double.NaN);
            }
            paramDim = parameters.length;
        }
        this.dirtyParameter = false;

        this.jacobianProviders = new ArrayList<ParameterJacobianProvider>();

        // set the default initial state Jacobian to the identity
        // and the default initial parameters Jacobian to the null matrix
        matricesData = new double[(stateDim + paramDim) * stateDim];
        for (int i = 0; i < stateDim; ++i) {
            matricesData[i * (stateDim + 1)] = 1.0;
        }

    }

    
Register the variational equations for the Jacobians matrices to the expandable set.
Params:
expandable – expandable set into which variational equations should be registered
Throws:
DimensionMismatchException – if the dimension of the partial state does not
match the selected equations set dimension
MismatchedEquations – if the primary set of the expandable set does
not match the one used to build the instance
See Also:
ExpandableStatefulODE.addSecondaryEquations(SecondaryEquations)/** Register the variational equations for the Jacobians matrices to the expandable set.
     * @param expandable expandable set into which variational equations should be registered
     * @throws DimensionMismatchException if the dimension of the partial state does not
     * match the selected equations set dimension
     * @exception MismatchedEquations if the primary set of the expandable set does
     * not match the one used to build the instance
     * @see ExpandableStatefulODE#addSecondaryEquations(SecondaryEquations)
     */
    public void registerVariationalEquations(final ExpandableStatefulODE expandable)
        throws DimensionMismatchException, MismatchedEquations {

        // safety checks
        final FirstOrderDifferentialEquations ode = (jode instanceof MainStateJacobianWrapper) ?
                                                    ((MainStateJacobianWrapper) jode).ode :
                                                    jode;
        if (expandable.getPrimary() != ode) {
            throw new MismatchedEquations();
        }

        efode = expandable;
        index = efode.addSecondaryEquations(new JacobiansSecondaryEquations());
        efode.setSecondaryState(index, matricesData);

    }

    
Add a parameter Jacobian provider.
Params:
provider – the parameter Jacobian provider to compute exactly the parameter Jacobian matrix/** Add a parameter Jacobian provider.
     * @param provider the parameter Jacobian provider to compute exactly the parameter Jacobian matrix
     */
    public void addParameterJacobianProvider(final ParameterJacobianProvider provider) {
        jacobianProviders.add(provider);
    }

    
Set a parameter Jacobian provider.
Params:
parameterizedOde – the parameterized ODE to compute the parameter Jacobian matrix using finite differences/** Set a parameter Jacobian provider.
     * @param parameterizedOde the parameterized ODE to compute the parameter Jacobian matrix using finite differences
     */
    public void setParameterizedODE(final ParameterizedODE parameterizedOde) {
        this.pode = parameterizedOde;
        dirtyParameter = true;
    }

    
Set the step associated to a parameter in order to compute by finite
 difference the Jacobian matrix.
 Needed if and only if the primary ODE set is a ParameterizedODE. 
 Given a non zero parameter value pval for the parameter, a reasonable value for such a step is pval * FastMath.sqrt(Precision.EPSILON). 

A zero value for such a step doesn't enable to compute the parameter Jacobian matrix.

Params:
parameter – parameter to consider for Jacobian processing
hP – step for Jacobian finite difference computation w.r.t. the specified parameter
Throws:
UnknownParameterException – if the parameter is not supported
See Also:
ParameterizedODE/** Set the step associated to a parameter in order to compute by finite
     *  difference the Jacobian matrix.
     * <p>
     * Needed if and only if the primary ODE set is a {@link ParameterizedODE}.
     * </p>
     * <p>
     * Given a non zero parameter value pval for the parameter, a reasonable value
     * for such a step is {@code pval * FastMath.sqrt(Precision.EPSILON)}.
     * </p>
     * <p>
     * A zero value for such a step doesn't enable to compute the parameter Jacobian matrix.
     * </p>
     * @param parameter parameter to consider for Jacobian processing
     * @param hP step for Jacobian finite difference computation w.r.t. the specified parameter
     * @see ParameterizedODE
     * @exception UnknownParameterException if the parameter is not supported
     */
    public void setParameterStep(final String parameter, final double hP)
        throws UnknownParameterException {

        for (ParameterConfiguration param: selectedParameters) {
            if (parameter.equals(param.getParameterName())) {
                param.setHP(hP);
                dirtyParameter = true;
                return;
            }
        }

        throw new UnknownParameterException(parameter);

    }

    
Set the initial value of the Jacobian matrix with respect to state.

If this method is not called, the initial value of the Jacobian
matrix with respect to state is set to identity.

Params:
dYdY0 – initial Jacobian matrix w.r.t. state
Throws:
DimensionMismatchException – if matrix dimensions are incorrect/** Set the initial value of the Jacobian matrix with respect to state.
     * <p>
     * If this method is not called, the initial value of the Jacobian
     * matrix with respect to state is set to identity.
     * </p>
     * @param dYdY0 initial Jacobian matrix w.r.t. state
     * @exception DimensionMismatchException if matrix dimensions are incorrect
     */
    public void setInitialMainStateJacobian(final double[][] dYdY0)
        throws DimensionMismatchException {

        // Check dimensions
        checkDimension(stateDim, dYdY0);
        checkDimension(stateDim, dYdY0[0]);

        // store the matrix in row major order as a single dimension array
        int i = 0;
        for (final double[] row : dYdY0) {
            System.arraycopy(row, 0, matricesData, i, stateDim);
            i += stateDim;
        }

        if (efode != null) {
            efode.setSecondaryState(index, matricesData);
        }

    }

    
Set the initial value of a column of the Jacobian matrix with respect to one parameter.

If this method is not called for some parameter, the initial value of
the column of the Jacobian matrix with respect to this parameter is set to zero.

Params:
pName – parameter name
dYdP – initial Jacobian column vector with respect to the parameter
Throws:
UnknownParameterException – if a parameter is not supported
DimensionMismatchException – if the column vector does not match state dimension/** Set the initial value of a column of the Jacobian matrix with respect to one parameter.
     * <p>
     * If this method is not called for some parameter, the initial value of
     * the column of the Jacobian matrix with respect to this parameter is set to zero.
     * </p>
     * @param pName parameter name
     * @param dYdP initial Jacobian column vector with respect to the parameter
     * @exception UnknownParameterException if a parameter is not supported
     * @throws DimensionMismatchException if the column vector does not match state dimension
     */
    public void setInitialParameterJacobian(final String pName, final double[] dYdP)
        throws UnknownParameterException, DimensionMismatchException {

        // Check dimensions
        checkDimension(stateDim, dYdP);

        // store the column in a global single dimension array
        int i = stateDim * stateDim;
        for (ParameterConfiguration param: selectedParameters) {
            if (pName.equals(param.getParameterName())) {
                System.arraycopy(dYdP, 0, matricesData, i, stateDim);
                if (efode != null) {
                    efode.setSecondaryState(index, matricesData);
                }
                return;
            }
            i += stateDim;
        }

        throw new UnknownParameterException(pName);

    }

    
Get the current value of the Jacobian matrix with respect to state.
Params:
dYdY0 – current Jacobian matrix with respect to state./** Get the current value of the Jacobian matrix with respect to state.
     * @param dYdY0 current Jacobian matrix with respect to state.
     */
    public void getCurrentMainSetJacobian(final double[][] dYdY0) {

        // get current state for this set of equations from the expandable fode
        double[] p = efode.getSecondaryState(index);

        int j = 0;
        for (int i = 0; i < stateDim; i++) {
            System.arraycopy(p, j, dYdY0[i], 0, stateDim);
            j += stateDim;
        }

    }

    
Get the current value of the Jacobian matrix with respect to one parameter.
Params:
pName – name of the parameter for the computed Jacobian matrix
dYdP – current Jacobian matrix with respect to the named parameter/** Get the current value of the Jacobian matrix with respect to one parameter.
     * @param pName name of the parameter for the computed Jacobian matrix
     * @param dYdP current Jacobian matrix with respect to the named parameter
     */
    public void getCurrentParameterJacobian(String pName, final double[] dYdP) {

        // get current state for this set of equations from the expandable fode
        double[] p = efode.getSecondaryState(index);

        int i = stateDim * stateDim;
        for (ParameterConfiguration param: selectedParameters) {
            if (param.getParameterName().equals(pName)) {
                System.arraycopy(p, i, dYdP, 0, stateDim);
                return;
            }
            i += stateDim;
        }

    }

    
Check array dimensions.
Params:
expected – expected dimension
array – (may be null if expected is 0)
Throws:
DimensionMismatchException – if the array dimension does not match the expected one/** Check array dimensions.
     * @param expected expected dimension
     * @param array (may be null if expected is 0)
     * @throws DimensionMismatchException if the array dimension does not match the expected one
     */
    private void checkDimension(final int expected, final Object array)
        throws DimensionMismatchException {
        int arrayDimension = (array == null) ? 0 : Array.getLength(array);
        if (arrayDimension != expected) {
            throw new DimensionMismatchException(arrayDimension, expected);
        }
    }

    
Local implementation of secondary equations.
 This class is an inner class to ensure proper scheduling of calls by forcing the use of JacobianMatrices.registerVariationalEquations(ExpandableStatefulODE). 
/** Local implementation of secondary equations.
     * <p>
     * This class is an inner class to ensure proper scheduling of calls
     * by forcing the use of {@link JacobianMatrices#registerVariationalEquations(ExpandableStatefulODE)}.
     * </p>
     */
    private class JacobiansSecondaryEquations implements SecondaryEquations {

        
{@inheritDoc} /** {@inheritDoc} */
        public int getDimension() {
            return stateDim * (stateDim + paramDim);
        }

        
{@inheritDoc} /** {@inheritDoc} */
        public void computeDerivatives(final double t, final double[] y, final double[] yDot,
                                       final double[] z, final double[] zDot)
            throws MaxCountExceededException, DimensionMismatchException {

            // Lazy initialization
            if (dirtyParameter && (paramDim != 0)) {
                jacobianProviders.add(new ParameterJacobianWrapper(jode, pode, selectedParameters));
                dirtyParameter = false;
            }

            // variational equations:
            // from d[dy/dt]/dy0 and d[dy/dt]/dp to d[dy/dy0]/dt and d[dy/dp]/dt

            // compute Jacobian matrix with respect to primary state
            double[][] dFdY = new double[stateDim][stateDim];
            jode.computeMainStateJacobian(t, y, yDot, dFdY);

            // Dispatch Jacobian matrix in the compound secondary state vector
            for (int i = 0; i < stateDim; ++i) {
                final double[] dFdYi = dFdY[i];
                for (int j = 0; j < stateDim; ++j) {
                    double s = 0;
                    final int startIndex = j;
                    int zIndex = startIndex;
                    for (int l = 0; l < stateDim; ++l) {
                        s += dFdYi[l] * z[zIndex];
                        zIndex += stateDim;
                    }
                    zDot[startIndex + i * stateDim] = s;
                }
            }

            if (paramDim != 0) {
                // compute Jacobian matrices with respect to parameters
                double[] dFdP = new double[stateDim];
                int startIndex = stateDim * stateDim;
                for (ParameterConfiguration param: selectedParameters) {
                    boolean found = false;
                    for (int k = 0 ; (!found) && (k < jacobianProviders.size()); ++k) {
                        final ParameterJacobianProvider provider = jacobianProviders.get(k);
                        if (provider.isSupported(param.getParameterName())) {
                            provider.computeParameterJacobian(t, y, yDot,
                                                              param.getParameterName(), dFdP);
                            for (int i = 0; i < stateDim; ++i) {
                                final double[] dFdYi = dFdY[i];
                                int zIndex = startIndex;
                                double s = dFdP[i];
                                for (int l = 0; l < stateDim; ++l) {
                                    s += dFdYi[l] * z[zIndex];
                                    zIndex++;
                                }
                                zDot[startIndex + i] = s;
                            }
                            found = true;
                        }
                    }
                    if (! found) {
                        Arrays.fill(zDot, startIndex, startIndex + stateDim, 0.0);
                    }
                    startIndex += stateDim;
                }
            }

        }
    }

    
Wrapper class to compute jacobian matrices by finite differences for ODE
 which do not compute them by themselves.
/** Wrapper class to compute jacobian matrices by finite differences for ODE
     *  which do not compute them by themselves.
     */
    private static class MainStateJacobianWrapper implements MainStateJacobianProvider {

        
Raw ODE without jacobians computation skill to be wrapped into a MainStateJacobianProvider. /** Raw ODE without jacobians computation skill to be wrapped into a MainStateJacobianProvider. */
        private final FirstOrderDifferentialEquations ode;

        
Steps for finite difference computation of the jacobian df/dy w.r.t. state. /** Steps for finite difference computation of the jacobian df/dy w.r.t. state. */
        private final double[] hY;

        
Wrap a FirstOrderDifferentialEquations into a MainStateJacobianProvider. 
Params:
ode – original ODE problem, without jacobians computation skill
hY – step sizes to compute the jacobian df/dy
Throws:
DimensionMismatchException – if there is a dimension mismatch between the steps array hY and the equation dimension/** Wrap a {@link FirstOrderDifferentialEquations} into a {@link MainStateJacobianProvider}.
         * @param ode original ODE problem, without jacobians computation skill
         * @param hY step sizes to compute the jacobian df/dy
         * @exception DimensionMismatchException if there is a dimension mismatch between
         * the steps array {@code hY} and the equation dimension
         */
        MainStateJacobianWrapper(final FirstOrderDifferentialEquations ode,
                                 final double[] hY)
            throws DimensionMismatchException {
            this.ode = ode;
            this.hY = hY.clone();
            if (hY.length != ode.getDimension()) {
                throw new DimensionMismatchException(ode.getDimension(), hY.length);
            }
        }

        
{@inheritDoc} /** {@inheritDoc} */
        public int getDimension() {
            return ode.getDimension();
        }

        
{@inheritDoc} /** {@inheritDoc} */
        public void computeDerivatives(double t, double[] y, double[] yDot)
            throws MaxCountExceededException, DimensionMismatchException {
            ode.computeDerivatives(t, y, yDot);
        }

        
{@inheritDoc} /** {@inheritDoc} */
        public void computeMainStateJacobian(double t, double[] y, double[] yDot, double[][] dFdY)
            throws MaxCountExceededException, DimensionMismatchException {

            final int n = ode.getDimension();
            final double[] tmpDot = new double[n];

            for (int j = 0; j < n; ++j) {
                final double savedYj = y[j];
                y[j] += hY[j];
                ode.computeDerivatives(t, y, tmpDot);
                for (int i = 0; i < n; ++i) {
                    dFdY[i][j] = (tmpDot[i] - yDot[i]) / hY[j];
                }
                y[j] = savedYj;
            }
        }

    }

    
Special exception for equations mismatch.
Since:
3.1/**
     * Special exception for equations mismatch.
     * @since 3.1
     */
    public static class MismatchedEquations extends MathIllegalArgumentException {

        
Serializable UID. /** Serializable UID. */
        private static final long serialVersionUID = 20120902L;

        
Simple constructor. /** Simple constructor. */
        public MismatchedEquations() {
            super(LocalizedFormats.UNMATCHED_ODE_IN_EXPANDED_SET);
        }

    }

}