import java.util.Random;

// ATTENTION
// Matrices and vectors are stored as Jama.Matrix objects.
// The documentation of the Jama matrix package can be found at
// http://math.nist.gov/javanumerics/jama/doc/
import Jama.*;

// ATTENTION
// You can create matrices on the fly, i.e. by writing down terms for coefficients
// with the following or similar expressions:
// new Matrix (new double[][] 
//               {{ a00, a01, a02}, 
//                {a10, a11, a12}}
//             )
// A shortcut for vectors, i.e. n*1 matrices is
// newVector (new double[] {v0, v1, v2})


class UnscentedKalmanFilter3D 
{
    //! The parameter of the flying cart used in the filter
    CoalCartParameter parameter;

    
    // ****** AUXILIARY FUNCTION ***************
    // Feel free to use.

    //! Auxiliary function that creates a column vector
    //! as a Jama.Matrix object from a 1D double array
    static Matrix newVector (double[] v)
    {
	return new Jama.Matrix (v, v.length);
    }

    //! Auxiliary function that returns column j from matrix A
    static Matrix copyColumn (Matrix A, int j)
    {
	Matrix result = new Matrix (A.getRowDimension(), 1);
	for (int i=0; i<result.getRowDimension(); i++) 
	    result.set(i, 0, A.get(i, j));
	return result;	
    }    

    //! Converts from matrix exponential to matrix representation.
    /*! Returns a 3*3 rotation matrix that rotates around \c axis by an
      angle of \c |axis|. This implements the Rodriguez formula. 
    */
    Matrix rot (Matrix axis)
    {
	double x=axis.get(0,0), y=axis.get(1,0), z=axis.get(2,0);
	double angle = Math.sqrt(x*x+y*y+z*z);
	if (angle>0) {
	    x/=angle;
	    y/=angle;
	    z/=angle;
	}	
	double s=Math.sin(angle), c=Math.cos(angle);
	double v = 1-c;
	double xyv = x*y*v;
	double yzv = y*z*v;
	double xzv = x*z*v;
	return new Matrix ( new double[][]
	    {{x*x*v + c, xyv - z*s, xzv + y*s},
	     {xyv + z*s, y*y*v + c, yzv - x*s},
	     {xzv - y*s, yzv + x*s, z*z*v + c
	     }});
    }

    //! Converts from matrix to matrix exponential representation
    /*! \c Q must be an orthonormal 3*3 matrix. The function returns
	the unique axis of rotation with the length being equal to the
	angle of rotation. If \c Q is identity the 0 vector is
	returned.
     */
    Matrix aRot (Matrix Q)
    {
	double c = (Q.trace()-1)/2;
	if (c<-1) c = -1;
	if (c>1) c = 1;
	double angle = Math.acos(c);
  
	if(angle!=0.0){
	    // next is to norm axis
	    double factor = angle/(2.0*Math.sin(angle));
	    // next is to determine axis and direction of axis
	    return newVector (new double[]{factor*(Q.get(2,1) - Q.get(1,2)),
					   factor*(Q.get(0,2) - Q.get(2,0)),
					   factor*(Q.get(1,0) - Q.get(0,1))});	    
	}
	else{
	    // Angle is zero, matrix is the identity, no unique axis,
	    // but we return (0,0,0) anyway
	    return newVector (new double[]{0,0,0});
	}	
    }
    
    //! Makes A exactly symmetric
    void symmetrize (Matrix A)
    {
	for (int i=0; i<A.getRowDimension(); i++)
	    for (int j=0;j<i;j++) {
		double x = (A.get(i,j)+A.get(j,i))/2;
		A.set (i, j, x);
		A.set (j, i, x);		
	    }		
    }

    //! Returns the Cholesky decomposition of A
    Matrix cholesky (Matrix A)
    {
	Matrix A2 = A.copy();
	symmetrize (A2);	
	CholeskyDecomposition chol = A2.chol();
	if (!chol.isSPD()) {
	    A.print (8,4);	    
	    throw new Error ("Matrix is not SPD");
	}	
	return chol.getL();
    }


    // ********** STATE REPRESENTATION ************

    //! Describes the state to be estimated
    /*! In our case the state is just a single orientation represented by a
        3*3 orthonormal matrix. So this class mainly provides routines
	for converting between this representation and a flat 3-D vector
	representation.
     */
    class State 
    {
	//! The orientation (not translation) of the body to be tracked
	/*! Represented as a 3*3 orthonormal matrix mapping a
	  direction vector in body coordinates to world
	  coordinates. While this matrix contains 9 numbers it
	  represents 3 degree of freedom.
        */
	Matrix body2World;

	State () {body2World = Matrix.identity (3, 3);}

	State (Matrix m) { body2World = m.copy(); }

	State (State s) { body2World = s.body2World.copy(); }	

	State copy() { return new State(this); }
	
	//! Prints the state, i.e. the orientation as a matrix
	void print () { body2World.print (8, 4); }
	

	//! Modifies a state by a flat vector representing the change
	/*! I.e. performs a circled plus (+) operation. 

            In our concrete case the \c deltaState is a 3-vector
            representing the relative rotation chained with \c
            body2World but with the axis converted to world
            coordinates. Hence \c deltaState is interpreted as a matrix
            exponential representation and converted to a matrix which is
            then multiplied to the right of \c body2World.

	    This replaces the operation + on ordinary vector states.
	*/
	State plus (Matrix deltaState)
	{
	    // TODO: Implement
	    return null;	    
	}
	

	//! Computes the flat vector representing the difference between this and state2
	/*! I.e. performs a circled (-) operation, i.e. s2.minus(s1) returns s2 \oplus s1. 
            The result if the result was added two \c state2 we would get this again. 

	    In our concrete case the result is a 3-vector representing
	    the relative rotation between this and \c state2 in matrix
	    exponential representation with the axis in world coordinates.

	    This replaces the operation - on ordinary vector states.
	*/
	Matrix minus (State state1)
	{
	    // TODO: Implement
	    return null;	    
	}

	//! Degrees of freedom of the state space
	/*! I.e. this is the dimension of the relative vectors and matrices. */
	final static int DOF = -1; // TODO: Change
    };
    

    // **************** MAIN CODE ****************

   
    //! The filter's belief on the current state
    State muX;

    //! and state covariance
    Matrix sigmaX;

    
    
    UnscentedKalmanFilter3D (CoalCartParameter p)
    {
	initialize (p);	
    }
    
    //! Initialize the mean and covariance to roughly identity
    /*! I.e. X, Y, Z of the coal cart are aligned with the 
        respective world axes.
    */
    void initialize (CoalCartParameter p)
    {
	    // TODO: Implement
	checkConsistency ();	
    }



    


    //! Dynamic model mapping old to new state
    /*! 
        Function f returning the state obtained after starting in
        state \c x and moving with dynamic measurement \c u.
	Without taking uncertainty into account.
    */
    State dynamicModel (State x, Matrix u)
    {
	    // TODO: Implement
	    return null;	    
    }    


    //! Generates an array of sigma points corresponding to \c mean
    //! and \c cov by using the Cholesky decomposition
    State[] generateSigmaPoints (State mean, Matrix cov)
    {
	    // TODO: Implement
	    return null;	    
    }


    //! Computes the mean state of a list of state sigma points
    State meanOfSigmaPoints (State[] sigmaPX)
    {
	    // TODO: Implement
	    return null;	    
    }

    //! Computes the covariance matrix of a list of state sigma points
    Matrix covarianceOfSigmaPoints (State[] sigmaPX)
    {
	    // TODO: Implement
	    return null;	    
    }

    //! Computes the mean of a list of vector sigma points
    /*! In our case these are the measurements. */
    Matrix meanOfSigmaPoints (Matrix[] sigmaPZ)
    {
	    // TODO: Implement
	    return null;	    
    }
    

    //! Computes the covariance of a set of vector sigma points
    /*! In our case these are the measurements. */
    Matrix covarianceOfSigmaPoints (Matrix[] sigmaPZ)
    {
	    // TODO: Implement
	    return null;	    
    }    

    //! Computes the covariance matrix of two corresponding lists of sigma points
    /*! \c sigmaPX, the first list of sigma points is a list of states X, the second
        list, a list of vectors, is the list of corresponding function results Z. The
	routine computes cov(Z,X).
    */
    Matrix covarianceOfSigmaPoints (Matrix[] sigmaPZ, State[] sigmaPX)
    {
	    // TODO: Implement
	    return null;	    
    }        
    
    
    //! Returns the covariance matrix of the dynamic measurement noise sigma_eps
    /*! Note, that this matrix represents state space uncertainty, but in our
        scenario we have measurement uncertainty. However this uncertainty
	can be converted to state space by simply scaling.
    */
    Matrix dynamicNoiseCovariance ()
    {
	    // TODO: Implement
	    return null;	    
    }
    
    
    //! Update the UKF after one dynamic step with a
    //! dynamic measurement u
    void dynamic (Matrix u)
    {
	    // TODO: Implement
	printState ("Mean body2World and covariance after dynamic step:");    
	checkConsistency ();
    }

    //! Wrapper for calling \c dynamic with three doubles
    //! namely the x,y,z gyro measurements [rad/s]
    void dynamic (double wx, double wy, double wz)
    {
	dynamic (newVector (new double[] {wx, wy, wz}));
    }

   
    //! Measurement function g for the observation 
    //! This function returns the value that should have been measured
    //! when the robot were in state \c x. In our case this is, what
    //! acceleration vector should have been measured when the robot
    //! was in state x.
    Matrix measurementModel (State x)
    {
	    // TODO: Implement
	    return null;	    
    }


    //! Returns the measurement noise covariance matrix Sigma_delta
    /*! Remember that this matrix has to be much bigger than the
	actual accelerometers sensor noise
	parameter.measurementNoiseSigma because we treat the physical
	acceleration as a random error in this approach. */
    Matrix measurementNoiseCovariance ()
    {
	    // TODO: Implement
	    return null;	    
    }
    

    //! Integrate a single 3D measurement of the accelerometers
    //! into the estimate
    void measurement (Matrix z)
    {	
	    // TODO: Implement

	printState ("Mean body2World and covariance after measurement update:");    
	checkConsistency ();	
    }

    
    //! Wrapper for calling \c measurement with one double
    //! namely the actual angle observed to beacon \c beaconId
    void measurement (double ax, double ay, double az)
    {
	measurement (newVector (new double[] {ax, ay, az}));	
    }
    
    // *********** AUXILIARY CODE FOR THE FRAMEWORK
    // *** The code below is not part of the Unscented Kalman Filter algorithm.

      //! Returns a random sample state from the posterior distribution
      /*! described by \c muX and \c sigmaX
       */
    State drawFromPosterior ()
    {
	if (muX==null) return null; // no information
	if (sigmaX==null) return muX; // only mean	
	Matrix L = cholesky (sigmaX);
	Random rnd = new Random();	
	Matrix v = newVector (new double[]{rnd.nextGaussian(), rnd.nextGaussian(), rnd.nextGaussian()});
	return muX.plus (L.times(v)); // mean plus random uncertainty
    }

    //! Prints the current mean and covariance
    void printState (String header)
    {
	System.out.println (header);
	if (muX!=null) muX.body2World.print (8, 4);
	if (sigmaX!=null) {
	    sigmaX.print (8, 4);	
	    double overallSigma = 180*Math.sqrt(sigmaX.trace())/Math.PI;	    
	    System.out.println ("Overall std deviation (sqrt(trace(cov))) in degree: "+
				overallSigma);
	}	
    }
    
    
    //! Checks, whether sigmaX is symmetric and positive definit
    //! If this routine fails, than usually one of your matrices
    //! or matrix formulas is broken.
    void checkConsistency () 
    {
	if (muX==null || sigmaX==null) return;
	
	// Test matrix format
	int n = State.DOF;	
	if (sigmaX.getRowDimension()!=n || sigmaX.getColumnDimension()!=n)
	    throw new Error ("sigmaX has wrong format");	

	Matrix cov2 = sigmaX.copy();	
	// Test symmetry
	for (int i=0;i<n;i++) for (int j=0;j<n;j++) {
	    cov2.set (i, j, (sigmaX.get(i,j)+sigmaX.get(j,i))/2);	
	    double delta = Math.abs(sigmaX.get(i,j)-sigmaX.get(j,i));
	    if (delta > 0.01) throw new Error ("Matrix is not symmetric");	    
	}	
	      
	// Assert positiv definiteness	
	if (!cov2.chol().isSPD()) throw new Error ("sigmaX is not positive definite");
    }    

    //! Checks, whether rot and aRot a inverse operations
    void checkMatrixExponential ()
    {
	Matrix vTest = newVector (new double[]{0.3, 0.7, -1.2});
	System.out.println ("A scaled axis rotation converted to matrix and back (should be the same)");	
	vTest.print (8, 4);
	rot (vTest).print(8, 4);
	aRot (rot (vTest)).print(8, 4);	
    }    

    //! Checks, whether relative to absolute conversion with State.plus and State.minus is consistent
    void checkStatePlusMinus ()
    {
	Matrix vTest1 = newVector (new double[]{0.3, 0.7, -1.2});
	Matrix vTest2 = newVector (new double[]{-0.2, 1.1, 0.15});
	State S1 = new State (rot (vTest1));
	State S2 = new State (rot (vTest2));
	Matrix delta = S2.minus(S1);
	State SNew = S1.plus (delta);
	System.out.println ("Original State S2");
	S2.print ();
	System.out.println ("Computed: S1+(S2-S1), should be S2");
	SNew.print ();		
    }
    

    //! Checks, whether converting the mean and covariance to sigma points and back changes
    //! something
    void checkSigmaPoints ()
    {
	System.out.println ("Mean and Covariance before conversion");
	muX.print ();	    
	sigmaX.print (8, 4);
	State[] sigmaP = generateSigmaPoints (muX, sigmaX);
	muX = meanOfSigmaPoints (sigmaP);
	sigmaX  = covarianceOfSigmaPoints (sigmaP);

	System.out.println ("Mean and Covariance after conversion (should be the same)");
	muX.print ();	    
	sigmaX.print (8, 4);	
    }    
};

