package assignentthree;
import java.util.*;
public class KDTree implements Iterable<Datum>{

	KDNode 		rootNode;
	int    		k;
	int			numLeaves;

 	// constructor

	public KDTree(ArrayList<Datum> datalist) throws Exception {

		Datum[]  dataListArray  = new Datum[ datalist.size() ];

		if (datalist.size() == 0) {
			throw new Exception("Trying to create a KD tree with no data");
		}
		else
			this.k = datalist.get(0).x.length;

		int ct=0;
		for (Datum d :  datalist) {
			dataListArray[ct] = datalist.get(ct);
			ct++;
		}

	//   Construct a KDNode that is the root node of the KDTree.

		rootNode = new KDNode(dataListArray);
	}

	//   KDTree methods

	public Datum nearestPoint(Datum queryPoint) {
		return rootNode.nearestPointInNode(queryPoint);
	}


	public int height() {
		return this.rootNode.height();
	}

	public int countNodes() {
		return this.rootNode.countNodes();
	}

	public int size() {
		return this.numLeaves;
	}

	//-------------------  helper methods for KDTree   ------------------------------

	public static long distSquared(Datum d1, Datum d2) {

		long result = 0;
		for (int dim = 0; dim < d1.x.length; dim++) {
			result +=  (d1.x[dim] - d2.x[dim])*((long) (d1.x[dim] - d2.x[dim]));
		}
		// if the Datum coordinate values are large then we can easily exceed the limit of 'int'.
		return result;
	}

	public double meanDepth(){
		int[] sumdepths_numLeaves =  this.rootNode.sumDepths_numLeaves();
		return 1.0 * sumdepths_numLeaves[0] / sumdepths_numLeaves[1];
	}

	class KDNode {

		boolean leaf;
		Datum leafDatum;           //  only stores Datum if this is a leaf

		//  the next two variables are only defined if node is not a leaf

		int splitDim;      // the dimension we will split on
		int splitValue;    // datum is in low if value in splitDim <= splitValue, and high if value in splitDim > splitValue

		KDNode lowChild, highChild;   //  the low and high child of a particular node (null if leaf)
		  //  You may think of them as "left" and "right" instead of "low" and "high", respectively

                HashSet test=new HashSet();

		KDNode(Datum[] datalist) throws Exception
                {
                    ArrayList <Datum> highTemp=new ArrayList <> (0), lowTemp=new ArrayList <> (0);

                    if (datalist.length==1)
                    {
                        if (test.add(datalist[0]))
                        {
                            this.leaf=true;
                            this.leafDatum=datalist[0];
                        }
                        else
                        {
                            this.leaf=false;
                        }
                    }

                    else if (datalist.length>1)
                    {
                        this.leaf=false;
                        this.splitDim=getRangePos(datalist);

                        int highCheck=datalist[splitDim].x[0];
                        int lowCheck=datalist[splitDim].x[0];

                        for (int z=1; z<datalist[splitDim].x.length; z++)
                        {
                            if (datalist[splitDim].x[z]>highCheck)
                            {
                                highCheck=datalist[splitDim].x[z];
                            }

                            if (datalist[splitDim].x[z]<lowCheck)
                            {
                                lowCheck=datalist[splitDim].x[z];
                            }
                        }

                        splitValue=(highCheck+lowCheck)/2;

                        for (int y=0; y<datalist.length-1; y++)
                        {
                            if (datalist[y]!=null)
                            {
                                if (datalist[y].x[splitDim]<=splitValue)
                                {
                                    highTemp.ensureCapacity(highTemp.size()+1);
                                    highTemp.add(datalist[y]);
                                }
                                else
                                {
                                    lowTemp.ensureCapacity(lowTemp.size()+1);
                                    lowTemp.add(datalist[y]);
                                }
                            }
                        }

                        Datum [] high=highTemp.toArray(new Datum [highTemp.size()]);
                        Datum [] low=lowTemp.toArray(new Datum [lowTemp.size()]);

                        if(high.length>=1)
                        {
                            this.highChild=new KDNode(high);
                        }

                        if(low.length>=1)
                        {
                            this.lowChild=new KDNode(low);
                        }
                    }
		}

                public int getRangePos(Datum [] datalist)
                {
                    int range=0, rangePos=0, highCheck=0, lowCheck=0;

                    if (datalist.length<=0)
                    {
                        throw new NullPointerException("The list is empty.");
                    }

                    for (int y=0; y<datalist.length; y++)
                    {
                        if (datalist[y]!=null)
                        {
                            for (int z=0; z<datalist[y].x.length; z++)
                            {
                                if (datalist[y].x[z]>highCheck-lowCheck)
                                {
                                    highCheck=datalist[y].x[z];
                                }
                                else
                                {
                                    lowCheck=datalist[y].x[z];
                                }
                            }

                            if (highCheck-lowCheck>range)
                            {
                                range=highCheck-lowCheck;
                                rangePos=y;
                            }
                        }
                        else
                        {
                            break;
                        }
                    }
                    return rangePos;
                }

		public Datum nearestPointInNode(Datum queryPoint)
                {
                        KDTreeIterator iterator=new KDTreeIterator(this);
                        int i = 0, j = iterator.list.size()-1, mid = 0;

                        if (KDTree.distSquared(iterator.pos(0), queryPoint)<=0)
                        {
                            return iterator.pos(0);
                        }

                        if (KDTree.distSquared(iterator.pos(0), queryPoint)>=0)
                        {
                            return iterator.pos(iterator.list.size()-1);
                        }

                        while (i < j)
                        {
                            mid = (i + j) / 2;

                            if (KDTree.distSquared(iterator.pos(mid), queryPoint)==0)
                            {
                                return iterator.pos(mid);
                            }

                                    /* If target is less than array element,
                                       then search in left */
                            if (KDTree.distSquared(iterator.pos(mid), queryPoint)<0)
                            {

                                        // If target is greater than previous
                                        // to mid, return closest of two
                                if (mid>0 && KDTree.distSquared(iterator.pos(mid-1), queryPoint)>0)
                                {
                                    if (KDTree.distSquared(queryPoint, iterator.pos(mid-1)) >= KDTree.distSquared(iterator.pos(mid), queryPoint))
                                        return iterator.pos(mid);
                                    else
                                        return iterator.pos(mid-1);
                                }
                                        /* Repeat for left half */
                                j = mid;
                            }

                                    // If target is greater than mid
                            else
                            {
                                if (mid < iterator.list.size()-1 && (KDTree.distSquared(iterator.pos(mid+1), queryPoint)>0))
                                {
                                    if (KDTree.distSquared(queryPoint, iterator.pos(mid-1)) >= KDTree.distSquared(iterator.pos(mid), queryPoint))
                                        return iterator.pos(mid+1);
                                    else
                                        return iterator.pos(mid);
                                }
                                i = mid + 1;
                            }
                        }
                        // Only single element left after search
                        return iterator.pos(mid);
                }

		// -----------------  KDNode helper methods (might be useful for debugging) -------------------

		public int height() {
                    if (this==null)
                        {
                            throw new NullPointerException();
                        }
			if (this.leaf)
				return 0;
			else
                            if (this.lowChild!=null && this.highChild!=null)
                            {
                                return 1 + Math.max( this.lowChild.height(), this.highChild.height());
                            }
                            else if (this.lowChild!=null)
                            {
                                return 1 + this.lowChild.height();
                            }
                            else if (this.highChild!=null)
                            {
                                return 1 + this.highChild.height();
                            }
                            else
                            {
                                return 0;
                            }
		}

                /*if (this.leaf)
				return 0;
			else {
				return 1 + Math.max( this.lowChild.height(), this.highChild.height());
			}*/

		public int countNodes() {
                        if (this==null)
                        {
                            throw new NullPointerException();
                        }
			if (this.leaf)
				return 1;
			else
                            if (this.lowChild!=null && this.highChild!=null)
                            {
                                return 1 + this.lowChild.countNodes() + this.highChild.countNodes();
                            }
                            else if (this.lowChild!=null)
                            {
                                return 1 + this.lowChild.countNodes();
                            }
                            else if (this.highChild!=null)
                            {
                                return 1 + this.highChild.countNodes();
                            }
                            else
                            {
                                return 0;
                            }
		}

		/*
		 * Returns a 2D array of ints.  The first element is the sum of the depths of leaves
		 * of the subtree rooted at this KDNode.   The second element is the number of leaves
		 * this subtree.    Hence,  I call the variables  sumDepth_size_*  where sumDepth refers
		 * to element 0 and size refers to element 1.
		 */

		public int[] sumDepths_numLeaves(){
			int[] sumDepths_numLeaves_low, sumDepths_numLeaves_high;
			int[] return_sumDepths_numLeaves = new int[2];

			/*
			 *  The sum of the depths of the leaves is the sum of the depth of the leaves of the subtrees,
			 *  plus the number of leaves (size) since each leaf defines a path and the depth of each leaf
			 *  is one greater than the depth of each leaf in the subtree.
			 */

			if (this.leaf) {  // base case
				return_sumDepths_numLeaves[0] = 0;
				return_sumDepths_numLeaves[1] = 1;
			}
			else {
				sumDepths_numLeaves_low  = this.lowChild.sumDepths_numLeaves();
				sumDepths_numLeaves_high = this.highChild.sumDepths_numLeaves();
				return_sumDepths_numLeaves[0] = sumDepths_numLeaves_low[0] + sumDepths_numLeaves_high[0] + sumDepths_numLeaves_low[1] + sumDepths_numLeaves_high[1];
				return_sumDepths_numLeaves[1] = sumDepths_numLeaves_low[1] + sumDepths_numLeaves_high[1];
			}
			return return_sumDepths_numLeaves;
		}

        }

        @Override
	public Iterator<Datum> iterator(){
		return new KDTreeIterator(this.rootNode);
	}

	private class KDTreeIterator implements Iterator<Datum>
        {

            ArrayList<Datum> list=new ArrayList();
            int x;

            KDTreeIterator(KDNode curr)
            {
                while (!curr.leaf)
                {
                    new KDTreeIterator(curr.lowChild);
                    new KDTreeIterator(curr.highChild);

                    if (curr.leaf)
                    {
                        list.add(curr.leafDatum);
                        break;
                    }
                }
                x=0;
            }

            @Override
            public boolean hasNext()
            {
                return !(list.get(x+1)==null);
            }

            @Override
            public Datum next()
            {
                if (hasNext())
                {
                    x++;
                    return list.get(x+1);
                }
                else
                {
                    throw new NullPointerException();
                }
            }

            public Datum pos(int x)
            {
                return list.get(x);
            }
        }
}