import java.util.HashMap;
import java.util.Iterator;

/**
 * This class stores a unitary matrix, either
 * in dense (all entries stored), or sparse (
 * only non-zero entries stored) format. The class
 * allows you to premultiply the stored unitary
 * matrix by various elementary gates.
 */
public class UnitaryMat {
    
    private int _numBits;
    private int _dim;//dimension of matrix (number of rows)
    private int _dimSq;//dimension squared
    private boolean _isSparse;
    //two containers for storing entries of matrix,
    //one for dense case and one for sparse case
    private double[] _denseEntries = null;
    private HashMap<Integer, Complex>[] _sparseRows = null;
    
/**
 *
 */
    public UnitaryMat(int numBits, boolean isSparse) {
    	Throw.If((numBits<0) || (numBits>8), "numBits out of bounds");

        _numBits = numBits;

        //dimension of matrix, between 1 and 256
        _dim = (1<<_numBits);
        //an int has 32 bits
        _dimSq = (1<<(2*_numBits));//between 1 and 256^2
        _isSparse = isSparse;

        if(!_isSparse){
            _denseEntries = new double[2*_dimSq];
         }else{
            for(int row=0; row<_dim; row++){
                _sparseRows[row] = 
                    new HashMap<Integer, Complex>(2/*=intialCapacity*/);
            }
        }
        setToIdentity();  

    }
    
/**
 *
 */
    public int getNumBits(){
        return _numBits;
    }
/**
 * 
 */
    public int getDim(){
        return _dim;
    }
    
/**
 *
 */
    public boolean isSparse(){
        return _isSparse;
    }
    
/**
 *
 */
    public Complex getEntry(int index) {
        //Throw.IfOut(index, _dimSq);
        if(!_isSparse){
            //Complex matrix is stored in successive columns order.          
            //Even slots carry real part and next slot carries imag part
            //Fortran and Clapack also stores matrices 
        	//in successive columns order.
            return new Complex(
                    _denseEntries[index<<1],//index<<1 = 2*index
                    _denseEntries[(index<<1) + 1]);
        }else{
            //index = col*_dim + row
            return _sparseRows[index%_dim].get(index/_dim);
        }        
    }
    
/**
 *
 */    
    public Complex getEntry(int row, int col) {
        //Throw.IfOut(row, _dim);
        //Throw.IfOut(col, _dim);
        return getEntry(col*_dim + row);//index = col*_dim + row
    }
    
/**
 *
 */
    public void setEntry(int index, Complex z){
        //Throw.IfOut(index, _dimSq);
        if(!_isSparse){
            _denseEntries[index<<1] = z.re;
            _denseEntries[(index<<1) + 1] = z.im;
        }else{
            //index = col*_dim + row
            _sparseRows[index%_dim].put(index/_dim, z);
        }
    }

    
/**
 *
 */
    public void setEntry(int row, int col, Complex z){
        //Throw.IfOut(row, _dim);
        //Throw.IfOut(col, _dim);
        setEntry(col*_dim + row, z);//index = col*_dim + row
    }
    
/**
 *
 */
    public void setToIdentity(){
        for(int row = 0; row< _dim; row++){
            setEntry(row, row, new Complex(1.0));
        }     
    }

/**
 *
 */   
    public void premultiplyByControlledBitSwap(
    		int bit1, 
    		int bit2,
    		Controls 	controls
    ){
        Throw.If(bit1==bit2, "swapped bits must be different");
        Throw.IfOut(bit1, _numBits);
        Throw.IfOut(bit2, _numBits);
        //number of controls may be zero
        int numControls = controls.getNumControls();
        int c;
        
        for(c=0; c<numControls; c++){
        	int cbit = controls.getBitPos(c);
            Throw.If(cbit == bit1 || cbit == bit2,
            	"neither swapped bit can be a control bit");
        }
        controls.checkBitPosAreDecreasing();

        Complex[] colVec = new Complex[_dim];
        Complex z;
        int bit1Mask = (1<<bit1);
        int bit2Mask = (1<<bit2);
        int col, row, newRow;
        for(col = 0; col<_dim; col ++){
            for(row = 0; row<_dim; row ++){
                if( ((row>>bit1)&1)  !=  ((row>>bit2)&1)  ){
                    newRow  = (row^bit1Mask^bit2Mask);
                }else{
                    newRow = row;
                }           
        		boolean changeVectorComponent = true;
        		for(c=0; c<numControls; c++){
                    boolean rowBitType = ((row>>controls.getBitPos(c))&1)==1;
                    if(controls.getType(c) != rowBitType){
                        changeVectorComponent = false;
                        break;
                    }                    
                }
        		if(!changeVectorComponent){
        			colVec[row] = getEntry(row, col);
                }else{//changeVectorComponent==true
                	colVec[row] = getEntry(newRow, col);
                }// end of changeVectorComponent == true
            }//first row loop                
            for(row = 0; row<_dim; row++){//second row loop
                if(!_isSparse){
                    setEntry(row, col, colVec[row]);
                }else{
                    z = colVec[row];
                    if(Complex.abs(z)>Globals.TOLERANCE){
                        setEntry(row, col, z);
                    }
                }
            }//second row loop
        }//col loop
    }
/**
*
*/   
   public void premultiplyByBitSwap(
   		int bit1, 
   		int bit2)
   {
	   premultiplyByControlledBitSwap(bit1, bit2, 
			   new Controls(_numBits, 0)/*length zero array*/);  	   
   }
     
/**
 *
 */   
    public void premultiplyByControlledGate(
        int   targetBitPosition,
        Controls   controls,
        OneBitGate oneBitGate 
    ){
        //number of controls may be zero
        int numControls = controls.getNumControls();
        int c;
        for(c=0; c<numControls; c++){
            Throw.If(controls.getBitPos(c) == targetBitPosition,
              "target bit cannot be a control bit");
        }
        controls.checkBitPosAreDecreasing();
        

        Complex[] colVec = new Complex[_dim];
        Complex z;
        int col, row;
        int targetMask = (1<<targetBitPosition);
        for(col = 0; col<_dim; col ++){
        	for(row = 0; row<_dim; row ++){
        		boolean changeVectorComponent = true;
        		for(c=0; c<numControls; c++){
                    boolean rowBitType = ((row>>controls.getBitPos(c))&1)==1;
                    if(controls.getType(c) != rowBitType){
                        changeVectorComponent = false;
                        break;
                    }                    
                }
                if(!changeVectorComponent){
                    colVec[row] = getEntry(row, col);
                }else{//changeVectorComponent==true
                    if(oneBitGate instanceof OneBitPhaseFactor){
                         colVec[row] = Complex.times(
                        	((OneBitPhaseFactor)oneBitGate).getPhaseFactor(),
                            getEntry(row, col));
                    }else if(oneBitGate instanceof OneBitP0PhaseFactor){
                        if((row&targetMask)==targetMask){   
                            //row has 1 at targetBitPosition
                            //row = ...1...
                        	colVec[row] = getEntry(row, col);
                        }else{//row = ...0...
                            colVec[row] = Complex.times(
                            	((OneBitP0PhaseFactor)oneBitGate).getPhaseFactor(),
                                getEntry(row, col));                      
                        }
                         
                    }else if(oneBitGate instanceof OneBitP1PhaseFactor){
                        if((row&targetMask)==targetMask){   
                            //row has 1 at targetBitPosition
                            //row = ...1...
                        	colVec[row] = Complex.times(
                            	((OneBitP1PhaseFactor)oneBitGate).getPhaseFactor(),
                                getEntry(row, col));  
	                    }else{//row = ...0...
	                            colVec[row] = getEntry(row, col);                    
	                    }
                    }else if(oneBitGate instanceof SigX){
                        colVec[row] = getEntry(row^targetMask, col);
                    }else if(oneBitGate instanceof SigY){
                        z = getEntry(row^targetMask, col);
                        if((row&targetMask)==targetMask){   
                            //row has 1 at targetBitPosition
                            //row = ...1...
                        	colVec[row] = new Complex(-z.im, z.re);//i*z
                        }else{//row = ...0...
                            colVec[row] = new Complex(z.im, -z.re);//-i*z                        
                        }

                    }else if(oneBitGate instanceof SigZ){
                        z = getEntry(row, col);
                        if((row&targetMask)==targetMask){   
                            //row has 1 at targetBitPosition
                            //row = ...1...
                            colVec[row] = new Complex(-z.re, -z.im);//-z  
                        }else{//row = ...0...
                            colVec[row] = z;                    
                        }

                    }else if(oneBitGate instanceof OneBitRot||
                        oneBitGate instanceof OneBitHadamard){
                        int smallRow, row0, row1;
                        if((row & targetMask) == targetMask){ 
                            //row has 1 at targetBitPosition
                            //row =...1...
				            smallRow = 1;
				            row0 = (row & ~targetMask); //row0 =...0...		
				            row1 = row; //row1 =...1... 
                        }else{
                            //row has 0 at targetBitPosition
                            //row =...0...
                            smallRow = 0;
                            row0 = row; //row0 =...0...
                            row1 = (row | targetMask); //row1 =...1...         
                        }
                        colVec[row] = Complex.plus(
                            Complex.times(
                            oneBitGate.getEntry(smallRow, 0), 
                            getEntry(row0, col)
                            ),
                            Complex.times(
                            oneBitGate.getEntry(smallRow, 1), 
                            getEntry(row1, col)
                            )
                            );

                    }else{
                        throw new RuntimeException(
                        "premultiplication by unsupported" +
                        (numControls==0? " OneBitGate": " ControlledGate"));
                    }//end of instanceof
                }// end of changeVectorComponent == true
            }//first row loop                
            for(row = 0; row<_dim; row++){//second row loop
                if(!_isSparse){
                    setEntry(row, col, colVec[row]);
                }else{
                    z = colVec[row];
                    if(Complex.abs(z)>Globals.TOLERANCE){
                        setEntry(row, col, z);
                    }
                }
            }//second row loop
        }//col loop
    }

/**
 *
 */   
    public void premultiplyByOneBitGate(
        int   targetBitPosition,
        OneBitGate oneBitGate 
    ){
        premultiplyByControlledGate(
        targetBitPosition,
        new Controls(_numBits, 0),//0 controls
        oneBitGate);
    }
/**
 * multiplies this times e(i*angle)
 */   
    public void multiplyByGlobalPhaseFactor(double angle){
        premultiplyByControlledGate(
        0, // targetBitPosition, anyone will do
        new Controls(_numBits, 0),//0 controls
        new OneBitPhaseFactor(angle));
    }
    
/**
 *
 */ 
 /*  
    public static UnitaryMat multiplySparseMats(
        UnitaryMat  aMat, 
        UnitaryMat  bMat)
    {
    	Throw.If(!aMat._isSparse || !bMat._isSparse,
            "One of the matrices being multiplied is not sparse");
    	Throw.If(aMat.getNumBits() != bMat.getNumBits(),
            "Matrices being multiplied must have same dimension");            
        int numBits = aMat.getNumBits();
        int matrixDim = (1<<numBits);
        UnitaryMat prodMat = new UnitaryMat(
            numBits, true);        
   
        //prodMat(j1, j3) = sum_j2 aMat(j1, j2) bMat(j2, j3)
        //aIter is iterator on columns of aMat
        //bIter is iterator on columns of bMat
        
        int j1, j2, j3;
        
        for(j1 = 0; j1<matrixDim; j1++){
            Iterator<Integer> aIter = (Iterator<Integer>)
                (aMat._sparseRows[j1].keySet().iterator());
            while(aIter.hasNext()){
                j2 = aIter.next();
                Iterator<Integer> bIter = (Iterator<Integer>)
                    (bMat._sparseRows[j2].keySet().iterator());
                 while(bIter.hasNext()){
                    j3 = bIter.next();
                    Complex z = Complex.times(
                        aMat._sparseRows[j1].get(j2),
                        bMat._sparseRows[j2].get(j3));
                    z = Complex.plus(z, prodMat._sparseRows[j1].get(j3));
                    prodMat._sparseRows[j1].put(j3, z);                    
                }
            }        
        }
        return prodMat;

    }
*/

/**
 *
 */
 
 /*  
    public void multiplySparseMats(UnitaryMat  bMat)
    {
    	Throw.If(!_isSparse || !bMat._isSparse,
            "One of the matrices being multiplied is not sparse");
    	Throw.If(_numBits != bMat.getNumBits(),
            "Matrices being multiplied must have same dimension");            
         
        //this(j1, j3) = sum_j2 this(j1, j2) bMat(j2, j3)
        //aIter is iterator on columns of this
        //bIter is iterator on columns of bMat
        
        int j1, j2, j3;
        
        for(j1 = 0; j1<_dim; j1++){
            Iterator<Integer> aIter = (Iterator<Integer>)
                (_sparseRows[j1].keySet().iterator());
            while(aIter.hasNext()){
                j2 = aIter.next();
                Iterator<Integer> bIter = (Iterator<Integer>)
                    (bMat._sparseRows[j2].keySet().iterator());
                 while(bIter.hasNext()){
                    j3 = bIter.next();
                    Complex z = Complex.times(
                        _sparseRows[j1].get(j2),
                        bMat._sparseRows[j2].get(j3));
                    z = Complex.plus(z, _sparseRows[j1].get(j3));
                    _sparseRows[j1].put(j3, z);                    
                }
            }        
        }

    }
*/

/**
 *
 */    
    public void makeDense(){
        if(!_isSparse){
            int index = 0;
            for(int col=0; col<_dim; col++){
                for(int row=0; row<_dim; row++){
                    Complex value = _sparseRows[row].get(col);
                    if(value==null){
                        _denseEntries[index<<1] = 0.0;
                        _denseEntries[index<<1 + 1] = 0.0;
                    }else{
                        _denseEntries[index<<1] = value.re;
                        _denseEntries[index<<1 + 1] = value.im;
                    }
                    //this gives index=col*_dim + row 
                    //if col loop before row loop 
                    index++;
                }
                //_sparseRows are usually indexed by row
                //but it doesn't matter if index by col 
                //for the purpose of releasing memory
                _sparseRows[col] = null ;
            }
            _isSparse = false;
        }
    }
    
/**
 *
 */    
    public String toString(){
        String output = "";
        output += "numBits= "+ _numBits + '\n';
        
        output += "Real part:" + '\n';
        int row, col;
        for(row=0; row<_dim; row++){
            for(col=0; col<_dim; col++){
                output += getEntry(row, col).re + " ";
            }
            output += '\n';
        }
 
        output += "Imag part:" + '\n';
        for(row=0; row<_dim; row++){
            for(col=0; col<_dim; col++){
                output += getEntry(row, col).im + " ";            
            }
            output += '\n';
        }
        
        return output;
    }
/**
 * 
 */
    public static void testSwap(){
    	UnitaryMat u;

    	u = new UnitaryMat(
    			2/*numBits*/, 
    			false/*isSparse*/);
    	u.premultiplyByBitSwap(0,1);
    	System.out.println("test1\n" + u.toString());

    	u = new UnitaryMat(
    			3/*numBits*/, 
    			false/*isSparse*/);
    	u.premultiplyByBitSwap(0,1);
    	System.out.println("test2\n" + u.toString());

    	u = new UnitaryMat(
    			3/*numBits*/, 
    			false/*isSparse*/);
    	u.premultiplyByBitSwap(1,2);
    	System.out.println("test3\n" + u.toString());
    	
    	Controls  controls = new Controls(3/*numBits*/, 1);    	
    	controls.setControl(0, 
    			2/*bit position*/, 
    			true/*type*/);
    	u = new UnitaryMat(
    			3/*numBits*/, 
    			false/*isSparse*/);
    	u.premultiplyByControlledBitSwap(0, 1, controls);
    	System.out.println("test4\n" + u.toString());

    	controls.setControl(0,
    			0,/*bit position*/
    			true/*type*/);
    	u = new UnitaryMat(
    			3/*numBits*/, 
    			false/*isSparse*/);
    	u.premultiplyByControlledBitSwap(1, 2, controls);
    	System.out.println("test5\n" + u.toString());

    }
/**
 * 
 */
    public static void testOneBitGates(){
    	int targetBitPosition = 1;
    	int numBits = 3;
    	
    	UnitaryMat u;


    	u = new UnitaryMat(
    			numBits, 
    			false/*isSparse*/);
    	u.premultiplyByOneBitGate(
    		targetBitPosition,
    		new SigX()
    	);
    	System.out.println("test1\n" + u.toString());
    	u.premultiplyByOneBitGate(
        		targetBitPosition,
        		new SigX()
        	);
        System.out.println("test1'\n" + u.toString());

    	u = new UnitaryMat(
    			numBits, 
    			false/*isSparse*/);
    	u.premultiplyByOneBitGate(
        	targetBitPosition,
    		new SigY()
    	);
    	System.out.println("test2\n" + u.toString());
    	u.premultiplyByOneBitGate(
            	targetBitPosition,
        		new SigY()
        	);
        System.out.println("test2'\n" + u.toString());

    	u = new UnitaryMat(
    			numBits, 
    			false/*isSparse*/);
    	u.premultiplyByOneBitGate(
            targetBitPosition,
    		new SigZ()
    	);
    	System.out.println("test3\n" + u.toString());
    	u.premultiplyByOneBitGate(
                targetBitPosition,
        		new SigZ()
        	);
        System.out.println("test3'\n" + u.toString());


    	u = new UnitaryMat(
    			numBits, 
    			false/*isSparse*/);
    	u.premultiplyByOneBitGate(
            targetBitPosition,
    		new OneBitHadamard()
    	);
    	System.out.println("test4\n" + u.toString());
    	u.premultiplyByOneBitGate(
                targetBitPosition,
        		new OneBitHadamard()
        	);
        System.out.println("test4'\n" + u.toString());

    	u = new UnitaryMat(
    			numBits, 
    			false/*isSparse*/);
    	u.premultiplyByOneBitGate(
            targetBitPosition,
    		new OneBitRot(.1, .2, .3)
    	);
    	System.out.println("test5\n" + u.toString());
    	u.premultiplyByOneBitGate(
                targetBitPosition,
                new OneBitRot(-.1, -.2, -.3)
        	);
        System.out.println("test5'\n" + u.toString());



    }
/**
 * 
 */    
    public static void main(String[] args){
    	//testSwap();
    	testOneBitGates();
    }
   
}
