package com.srbenoit.math.solvers;

import com.srbenoit.math.linear.DMatrix;
import com.srbenoit.math.linear.MatrixNotCompleteException;
import com.srbenoit.math.linear.SingularMatrixException;
import com.srbenoit.math.linear.TupleN;
import com.srbenoit.util.EvaluationException;

/**
 * This solver can take a system of nonlinear equations and find local solutions given a starting
 * guess, using the Newton-Raphson method. We are given a set of functions that take a vector of
 * values to generate an output. We seek to find the set of values such that all functions yield
 * zero.
 */
public class NewtonRaphsonMethod {

    /** tolerance used to detect when the solver has converged */
    public static final double TOLERANCE = 1E-16;

    /** tolerance used to detect when changes have grown too small */
    public static final double DELTA_TOLERANCE = 1E-100;

    /** the list of functions to be solved */
    private final transient SystemOfFunctions functions;

    /**
     * Constructs a new <code>NewtonRaphsonMethod</code>.
     *
     * @param  functionsToSolve  the list of functions to be solved.\
     */
    public NewtonRaphsonMethod(final SystemOfFunctions functionsToSolve) {

        this.functions = functionsToSolve;
    }

    /**
     * Finds a solution based on an initial guess.
     *
     * @param   maxIterations  the maximum number of iterations
     * @param   initialGuess   the initial guess
     * @return  the solution, if found
     * @throws  SolverFailedException  if the solver failed to find a solution
     */
    public TupleN solve(final int maxIterations, final TupleN initialGuess)
        throws SolverFailedException {

        TupleN coords;
        double err;
        TupleN rhs;
        DMatrix jacobian;
        DMatrix inverse;
        TupleN deltas;

        coords = initialGuess.copy();
        rhs = new TupleN(initialGuess.getDimension(), false); 
        jacobian = new DMatrix(initialGuess.getDimension(), initialGuess.getDimension()); 

        for (int iter = 0; iter < maxIterations; iter++) {

            // Do any pre-computation based on current coordinates
            this.functions.precompute(coords);

            // See if we are done based on error being below tolerance.
            err = computeError(coords);

            if (err < TOLERANCE) {
                break; // Solver has finished.
            }

            // Compute the approximate Jacobian at the current point
            computeJacobian(coords, jacobian);

            // Invert the Jacobian matrix
            try {
                inverse = jacobian.inverse(); 
            } catch (EvaluationException e1) {
                throw new SolverFailedException("Unable to invert matrix", e1);
            } catch (SingularMatrixException e2) {
                throw new SolverFailedException("Unable to invert matrix", e2);
            } catch (MatrixNotCompleteException e2) {
                throw new SolverFailedException("Unable to invert matrix", e2);
            }

            // Set up right-hand side of dX = - J^-1 dot F(X)
            for (int n = 0; n < coords.getDimension(); n++) {
                rhs.set(n, -this.functions.getFunction(n).evaluate(this.functions, coords));
            }

            // Compute dX
            try {
                deltas = inverse.transform(rhs);
            } catch (EvaluationException e1) {
                throw new SolverFailedException("Unable to compute deltas", e1);
            }

            // Shrink deltas to avoid overshoot, check for convergence based on
            // Deltas being to small, and at the same time, update the values
            // based on the deltas.
            err = 0;

            for (int i = 0; i < coords.getDimension(); i++) {
                deltas.set(i, deltas.get(i) * 0.1);
                err += Math.abs(deltas.get(i));
                coords.set(i, coords.get(i) + deltas.get(i));
            }

            if (err < DELTA_TOLERANCE) {
                break; // Solver has finished.
            }
        }

        return coords;
    }

    /**
     * Computes the error at the current set of coordinates.
     *
     * @param   coords  the coordinates
     * @return  the error
     * @throws  SolverFailedException  if the error cannot be evaluated or is infinite
     */
    private double computeError(final TupleN coords) throws SolverFailedException {

        double err;

        err = 0;

        for (int i = 0; i < coords.getDimension(); i++) {
            err += Math.abs(this.functions.getFunction(i).evaluate(this.functions, coords));
        }

        if (Double.isNaN(err)) {
            throw new SolverFailedException("Failed to evaluate");
        }

        if (Double.isInfinite(err)) {
            throw new SolverFailedException("Error is infinite");
        }

        return err;
    }

    /**
     * Computes the Jacobian matrix of the functions at the current point.
     *
     * @param  coordinates  the point at which to evaluate the Jacobian
     * @param  jacobian     the resulting Jacobian
     */
    private void computeJacobian(final TupleN coordinates, final DMatrix jacobian) {

        AbstractFunction func;
        double partial;

        // Go through each function, get its value at the current point
        for (int i = 0; i < this.functions.numFunctions(); i++) {
            func = this.functions.getFunction(i); 

            // Now loop through each value and take the partial derivative in
            // that value's direction
            for (int j = 0; j < coordinates.getDimension(); j++) {

                partial = func.derivative(this.functions, coordinates, j);
                jacobian.set(i, j, partial);
            }
        }
    }
}