import java.util.Random;
import java.util.Vector;
import java.util.Comparator;
import java.util.Map;
import java.util.TreeMap;
import java.util.Iterator;

/**
 * <p>
 * Creates a random flock of robots in an n-dimensional box
 * </p>
 * @author     Michael Schuresko
 * @version    %I%, %G%
 * @since      1.0
*/
public class RandomConnectedInit implements IInitializer
{

    int m_nNumAgents, m_nDim;
    int m_nNumAgentsActual;
    double [] m_arrLfMin;
    double [] m_arrLfMax;
    double m_lfConnectRad;
    
    double m_lfMinRad;
    
    double [] m_arrLfIniContState;

    ILogicVarBundle m_logicVarsToClone;
    IControlFunc m_controlFuncToClone;

    Random m_randomGen;

    boolean m_bVerbose;
    
    void setVerbosity() { 
    	m_bVerbose = false;
    }

    /**
     * constructor specifying number of agents and
     * dimensionality in which they live.
     * @param nAgents number of agents
     * @param nDim dimensionality in which they live
     */
    public RandomConnectedInit(int nAgents, int nDim, double lfRad)
    {
    	setVerbosity();
    	m_randomGen = new Random();
    	m_nNumAgents = nAgents;
    	m_nDim = nDim;
    	m_lfConnectRad = lfRad;
    	m_arrLfMin = new double[nDim];
		m_arrLfMax = new double[nDim];
		for(int i = 0; i < nDim; ++i) {
			m_arrLfMin[i] = -0.5;
			m_arrLfMax[i] = 0.5;
		}
		m_logicVarsToClone = new NullLogicVarBundle();
		m_controlFuncToClone = new ControlFuncDoNothing(nDim);
		
		m_lfMinRad = 0.0;
    }

    /**
     * constructor specifying number of agents, the
     * dimensionality in which they live, and the size
     * of the containing box along each dimension.
     * @param nAgents number of agents
     * @param arrLfSizes length, height, width, etc of box
     * in which agents live-arrLfSizes.length is dimensionality.
     */
    public RandomConnectedInit(int nAgents, double [] arrLfSizes, double lfRad)
    {
    	setVerbosity();
    	m_randomGen = new Random();
    	m_nNumAgents = nAgents;
    	m_lfConnectRad = lfRad;
    	m_nDim = arrLfSizes.length;
    	m_arrLfMin = new double[m_nDim];
    	m_arrLfMax = new double[m_nDim];
    	for(int i = 0; i < m_nDim; ++i) {
    		m_arrLfMin[i] = -0.5*arrLfSizes[i];
    		m_arrLfMax[i] = 0.5*arrLfSizes[i];
    	}	
    	m_logicVarsToClone = new NullLogicVarBundle();
    	m_controlFuncToClone = new ControlFuncDoNothing(m_nDim);
    	
    	m_lfMinRad = 0.0;
    }
    
    /**
     * constructor specifying number of agents, the
     * dimensionality in which they live, and the size
     * of the containing box along each dimension.
     * @param nAgents number of agents
     * @param arrLfMin min corner of box in which agents live.
     * @param arrLfMax max corner of box in which agents live.
     * 
     */
    public RandomConnectedInit(int nAgents, double [] arrLfMin, 
			       double [] arrLfMax, double lfRad)
    {
    	setVerbosity();
    	m_randomGen = new Random();
    	m_nNumAgents = nAgents;
	
    	m_lfConnectRad = lfRad;
    	m_nDim = arrLfMin.length;
    	m_arrLfMin = new double[m_nDim];
    	m_arrLfMax = new double[m_nDim];
    	for(int i = 0; i < m_nDim; ++i) {
    		m_arrLfMin[i] = arrLfMin[i];
    		m_arrLfMax[i] = arrLfMax[i];
    	}
    	m_logicVarsToClone = new NullLogicVarBundle();
    	m_controlFuncToClone = new ControlFuncDoNothing(m_nDim);
    	
    	m_lfMinRad = 0.0;
    }
    
    public void setMinRad(double lfMinRad) {
    	m_lfMinRad = lfMinRad;
    }

    public void setDefaultIniVars(ILogicVarBundle copyMe)
    {
    	m_logicVarsToClone = copyMe.makeCopy();
    }

    public void setDefaultIniControlFunc(IControlFunc copyMe)
    {
    	m_controlFuncToClone = copyMe.makeCopy();
    }

    
    
    /**
     * Initial continuous state vector
     */
    public double [] getInitContState() 
    {
    	resetState();
    	return m_arrLfIniContState;
    }
    
    public void resetState() {
    	m_nNumAgentsActual = 0;
    	while(m_nNumAgentsActual != m_nNumAgents) {
    		
    		m_arrLfIniContState = createIniContState();
    		m_nNumAgentsActual = m_arrLfIniContState.length/m_nDim;
    		if(m_bVerbose) {
    			System.out.println("there are "+m_nNumAgentsActual+" agents, should be "+m_nNumAgents);
    		}
    	}
    }
    
    double [] createIniContState() 
    {
    	Vector<Vector<Double> > vecPoints;
    	vecPoints = getIniVec();
    	double [] arrLfOut = new double[vecPoints.size()*m_nDim];
    	for(int i = 0; i < vecPoints.size(); ++i) {
    		for(int j = 0; j < m_nDim; ++j) {
    			arrLfOut[i*m_nDim+j] = vecPoints.get(i).get(j);
    		}
    	}
		return arrLfOut;
    }
    
    Integer getDsuCount(Vector<Integer> vecCounts, Vector<Integer> vecReps, int i) {
    	Integer objRep = getDsuRep(vecReps, i);
    	return vecCounts.get(objRep.intValue());
    }
    
    Integer getDsuRep(Vector<Integer> vecReps, int i) {
    	Integer objJ = new Integer(i);
		while(vecReps.get(objJ).compareTo(objJ) != 0) {
	    	objJ = vecReps.get(objJ);
		}
		Integer objResult = objJ;
		Integer objNext = objJ;
		for(objJ = new Integer(i); objJ.compareTo(objResult) != 0; objJ = objNext) {
	    	objNext = vecReps.get(objJ);
	    	vecReps.set(objJ.intValue(), objResult);
		}
		return objResult;
    }

    void doDsuMerge(Vector<Integer> vecReps, Vector<Integer> vecCounts, int i, int j) {
    	Integer objIRep = getDsuRep(vecReps, i);
    	Integer objJRep = getDsuRep(vecReps, j);
    	if(objIRep.compareTo(objJRep) == 0) {
    		return;
    	}
    	vecCounts.set(objJRep, vecCounts.get(objJRep)+ vecCounts.get(objIRep));
    	vecReps.set(objIRep.intValue(), objJRep);
    }

    public class VecComp implements Comparator<Vector<Double>>
    {
    	public int compare(Vector<Double> v1, Vector<Double> v2) {
    		if(v1.get(0) < v2.get(0)) {
    			return -1;
    		}
    		if(v1.get(0) > v2.get(0) ) {
    			return 1;
    		}
    		return 0;
    	}
    	
    	public boolean equals(Vector<Double> v1, Vector<Double> v2) {
    		return !((v1.get(0) > v2.get(0)) || 
    				(v1.get(0) < v2.get(0)));
    	}
    }
    
    Vector<Vector<Double> > getIniVec() {
    	
    	Vector<Integer> vecIndices = new Vector<Integer>(m_nNumAgents, m_nNumAgents);
    	Vector<Integer> vecCounts = new Vector<Integer>(m_nNumAgents, m_nNumAgents);
    	
    	TreeMap<Vector<Double>, Integer> mapPos = 
    		new TreeMap<Vector<Double>, Integer>(new VecComp());
    	
    	
    	boolean bConnected = false;
    	
    	int i;
    	for(i = 0; !bConnected; ++i) {
    		vecIndices.add(i, i);
    		vecCounts.add(i,1);
    		Vector<Double> vecNextPos = new Vector<Double>(m_nDim);
    		int nAxis = 0;
    		for(nAxis = 0; nAxis < m_nDim; ++nAxis) {
    			Double objLfCurrPos = 
    				new Double(m_arrLfMin[nAxis] +
    						m_randomGen.nextDouble()*
    						(m_arrLfMax[nAxis]-m_arrLfMin[nAxis]));
    			vecNextPos.add(nAxis, objLfCurrPos);
    		}
    		
    		Iterator<Map.Entry<Vector<Double>, Integer> > iPosPairs = null;
    		
    		if(m_lfMinRad > 0.0 && i > 0) {

    			Vector<Double> vecMinCollideRange = new Vector<Double>(vecNextPos);
    			Vector<Double> vecMaxCollideRange = new Vector<Double>(vecNextPos);

    			vecMinCollideRange.set(0, vecMinCollideRange.get(0) - 2*m_lfMinRad);
    			vecMaxCollideRange.set(0, vecMaxCollideRange.get(0) + 2*m_lfMinRad);

    			boolean bCollides = false;

    			iPosPairs =  
    				mapPos.subMap(vecMinCollideRange, vecMaxCollideRange).entrySet().iterator();

    			while(iPosPairs.hasNext()) {
    				Map.Entry<Vector<Double>, Integer> currEntry = 
    					iPosPairs.next();

    				Integer nObjRep = currEntry.getValue();
    				Vector<Double> vecPosCmp = currEntry.getKey();
    				double lfDistSum = 0.0;
    				for(nAxis = 0; nAxis < m_nDim; ++nAxis) {
    					double lfAxisDist = vecNextPos.get(nAxis) -
    					vecPosCmp.get(nAxis);
    					lfDistSum += lfAxisDist*lfAxisDist;
    				}

    				if(lfDistSum < m_lfMinRad*m_lfMinRad && nObjRep.intValue() != i) {
    					bCollides = true;
    				}
    			}

    			if(bCollides) {
    				--i;
    				continue;
    			}

    		}
    		
    		mapPos.put(vecNextPos, i);
    		
    		Vector<Double> vecMinRange = new Vector<Double>(vecNextPos);
    		Vector<Double> vecMaxRange = new Vector<Double>(vecNextPos);
    		
    		vecMinRange.set(0, vecMinRange.get(0) - 1.2*m_lfConnectRad);
    		vecMaxRange.set(0, vecMaxRange.get(0) + 1.2*m_lfConnectRad);
    		
    		iPosPairs = 
    			mapPos.subMap(vecMinRange, vecMaxRange).entrySet().iterator();
//    		iPosPairs = mapPos.entrySet().iterator();
    		
    		while(iPosPairs.hasNext()) {
    			Map.Entry<Vector<Double>, Integer> currEntry = 
    				iPosPairs.next();
    			Integer nObjRep = currEntry.getValue();
    			Vector<Double> vecPosCmp = currEntry.getKey();
    			double lfDistSum = 0.0;
    			for(nAxis = 0; nAxis < m_nDim; ++nAxis) {
    				double lfAxisDist = vecNextPos.get(nAxis) -
    					vecPosCmp.get(nAxis);
    				lfDistSum += lfAxisDist*lfAxisDist;
    			}
    			
    			if(lfDistSum < m_lfConnectRad*m_lfConnectRad && nObjRep.intValue() != i) {
    				doDsuMerge(vecIndices, vecCounts, nObjRep, i);
    			}
    		}
    		Integer nObjNumAgentsThisComp = getDsuCount(vecCounts, vecIndices, i);
    		bConnected = (nObjNumAgentsThisComp >= m_nNumAgents);
    		if(m_bVerbose) {
    			System.out.println("UP to "+nObjNumAgentsThisComp +" agents, need "+m_nNumAgents);
    		}
    	}

    	Vector<Vector<Double> > vecResult = 
    		new Vector<Vector<Double> >(vecCounts.get(vecCounts.size()-1));
    	
    	Integer objSetRep = getDsuRep(vecIndices, vecIndices.size()-1);
    	int nAddPos = 0;
    	Vector<Double> vecOffset = null;
    	
    	Iterator<Map.Entry<Vector<Double>, Integer> > iPosPairs = 
			mapPos.entrySet().iterator();
    	
    	while(iPosPairs.hasNext()) {
    		Map.Entry<Vector<Double>, Integer> currEntry =
    			iPosPairs.next();
    		Integer nObjI = getDsuRep(vecIndices, currEntry.getValue());
    		if(objSetRep.compareTo(nObjI) == 0) {
    			vecResult.add(nAddPos, currEntry.getKey());
    			++nAddPos;
    		}
    	}
    	
    	int nSize = vecResult.size();
    	// now we permute the result vector for better results!!!
    	for(i = 0; i < vecResult.size(); ++i) {
    		int nSwap = this.m_randomGen.nextInt(nSize-i)+i;
    		Vector<Double> vecTmp = vecResult.elementAt(i);
    		vecResult.set(i, vecResult.elementAt(nSwap));
    		vecResult.set(nSwap, vecTmp);
    	}
    	
    	
    	for(i = 0; i < vecResult.size(); ++i) {
    		Vector<Double> vecAdjustMe = vecResult.get(i);
    		if(vecOffset == null) {
    			vecOffset = new Vector<Double>(vecResult.get(i).size());
    			for(int j = 0; j < vecResult.get(i).size(); ++j) {
    				vecOffset.add(j, vecResult.get(i).get(j));
    			}
    		}
    		for(int j = 0; j < vecAdjustMe.size(); ++j) {
    			vecAdjustMe.set(j, vecAdjustMe.get(j) - vecOffset.get(j));
    		}
    	}
    	
    	return vecResult;
    }

    public IEnvironment getEnv()
    {
	return new NdTrivialEnv(m_arrLfMin, m_arrLfMax);
    }
    


    public int getNumAgents() 
    {
    	return m_nNumAgentsActual;
    }
    

    /**
     * convenient to have this...
     */
    public StateBundle [] getInitDiscreteState()
    {
        StateBundle [] arrStates = new StateBundle[m_nNumAgentsActual];
    	if(m_bVerbose) {
    		System.out.println("Initializing to "+m_nNumAgentsActual+" agents\n");
    	}
    	for(int i = 0; i < m_nNumAgentsActual; ++i)
	    {
    		ILogicVarBundle currCopy = m_logicVarsToClone.makeCopy();
    		currCopy.setId(i);
    		currCopy.init();
    		arrStates[i] = new StateBundle(currCopy,
    				m_controlFuncToClone.makeCopy()
    		);
	    }
    	return arrStates;
    }

    public boolean isAgentCompatible(IAgent agent)
    {
    	return (agent != null);
    }
}