//
//  RMQuadTree.m
//  MapView
//
//  Created by Thomas Rasch on 27.07.11.
//  Copyright 2011 Alpstein. All rights reserved.
//

#import "RMQuadTree.h"
#import "RMAnnotation.h"
#import "RMProjection.h"
#import "RMMapView.h"

#pragma mark -
#pragma mark RMQuadTreeNode implementation

#define kMinimumQuadTreeElementWidth 200.0 // projected meters
#define kMaxAnnotationsPerLeaf 4
#define kMinPixelDistanceForLeafClustering 100.0

@interface RMQuadTreeNode ()

- (id)initWithMapView:(RMMapView *)aMapView forParent:(RMQuadTreeNode *)aParentNode inBoundingBox:(RMProjectedRect)aBoundingBox;

- (void)addAnnotation:(RMAnnotation *)annotation;
- (void)removeAnnotation:(RMAnnotation *)annotation;

- (void)addAnnotationsInBoundingBox:(RMProjectedRect)aBoundingBox toMutableArray:(NSMutableArray *)someArray createClusterAnnotations:(BOOL)createClusterAnnotations withClusterSize:(RMProjectedSize)clusterSize findGravityCenter:(BOOL)findGravityCenter;

- (void)removeUpwardsAllCachedClusterAnnotations;

@end

@implementation RMQuadTreeNode

@synthesize nodeType;
@synthesize boundingBox, northWestBoundingBox, northEastBoundingBox, southWestBoundingBox, southEastBoundingBox;
@synthesize parentNode, northWest, northEast, southWest, southEast;

- (id)initWithMapView:(RMMapView *)aMapView forParent:(RMQuadTreeNode *)aParentNode inBoundingBox:(RMProjectedRect)aBoundingBox
{
    if (!(self = [super init]))
        return nil;

//    RMLog(@"New quadtree node at {(%.0f,%.0f),(%.0f,%.0f)}", aBoundingBox.origin.easting, aBoundingBox.origin.northing, aBoundingBox.size.width, aBoundingBox.size.height);

    mapView = aMapView;
    parentNode = [aParentNode retain];
    northWest = northEast = southWest = southEast = nil;
    annotations = [NSMutableArray new];
    boundingBox = aBoundingBox;
    cachedClusterAnnotation = nil;
    cachedClusterEnclosedAnnotations = nil;

    double halfWidth = boundingBox.size.width / 2.0, halfHeight = boundingBox.size.height / 2.0;
    northWestBoundingBox = RMProjectedRectMake(boundingBox.origin.x, boundingBox.origin.y + halfHeight, halfWidth, halfHeight);
    northEastBoundingBox = RMProjectedRectMake(boundingBox.origin.x + halfWidth, boundingBox.origin.y + halfHeight, halfWidth, halfHeight);
    southWestBoundingBox = RMProjectedRectMake(boundingBox.origin.x, boundingBox.origin.y, halfWidth, halfHeight);
    southEastBoundingBox = RMProjectedRectMake(boundingBox.origin.x + halfWidth, boundingBox.origin.y, halfWidth, halfHeight);

    nodeType = nodeTypeLeaf;

    return self;
}

- (void)dealloc
{
    mapView = nil;
    [cachedClusterAnnotation release]; cachedClusterAnnotation = nil;
    [cachedClusterEnclosedAnnotations release]; cachedClusterEnclosedAnnotations = nil;

    @synchronized (annotations)
    {
        for (RMAnnotation *annotation in annotations)
        {
            annotation.quadTreeNode = nil;
        }
    }

    [annotations release]; annotations = nil;

    [northWest release]; northWest = nil;
    [northEast release]; northEast = nil;
    [southWest release]; southWest = nil;
    [southEast release]; southEast = nil;
    [parentNode release]; parentNode = nil;

    [super dealloc];
}

- (NSArray *)annotations
{
    NSArray *immutableAnnotations = nil;

    @synchronized (annotations)
    {
        immutableAnnotations = [NSArray arrayWithArray:annotations];
    }

    return immutableAnnotations;
}

- (void)addAnnotationToChildNodes:(RMAnnotation *)annotation
{
    RMProjectedRect projectedRect = annotation.projectedBoundingBox;

    if (RMProjectedRectContainsProjectedRect(northWestBoundingBox, projectedRect))
    {
        if (!northWest)
            northWest = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:self inBoundingBox:northWestBoundingBox];

        [northWest addAnnotation:annotation];
    }
    else if (RMProjectedRectContainsProjectedRect(northEastBoundingBox, projectedRect))
    {
        if (!northEast)
            northEast = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:self inBoundingBox:northEastBoundingBox];

        [northEast addAnnotation:annotation];
    }
    else if (RMProjectedRectContainsProjectedRect(southWestBoundingBox, projectedRect))
    {
        if (!southWest)
            southWest = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:self inBoundingBox:southWestBoundingBox];

        [southWest addAnnotation:annotation];
    }
    else if (RMProjectedRectContainsProjectedRect(southEastBoundingBox, projectedRect))
    {
        if (!southEast)
            southEast = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:self inBoundingBox:southEastBoundingBox];

        [southEast addAnnotation:annotation];
    }
    else
    {
        @synchronized (annotations)
        {
            [annotations addObject:annotation];
        }

        annotation.quadTreeNode = self;
        [self removeUpwardsAllCachedClusterAnnotations];
    }
}

- (void)addAnnotation:(RMAnnotation *)annotation
{
    if (nodeType == nodeTypeLeaf)
    {
        @synchronized (annotations)
        {
            [annotations addObject:annotation];
        }

        annotation.quadTreeNode = self;

        if ([annotations count] <= kMaxAnnotationsPerLeaf || boundingBox.size.width < (kMinimumQuadTreeElementWidth * 2.0))
        {
            [self removeUpwardsAllCachedClusterAnnotations];
            return;
        }

        nodeType = nodeTypeNode;

        // problem: all annotations that cross two quadrants will always be re-added here, which
        // might be a problem depending on kMaxAnnotationsPerLeaf

        NSArray *immutableAnnotations = nil;

        @synchronized (annotations)
        {
            immutableAnnotations = [NSArray arrayWithArray:annotations];
            [annotations removeAllObjects];
        }

        for (RMAnnotation *annotationToMove in immutableAnnotations)
        {
            [self addAnnotationToChildNodes:annotationToMove];
        }

        return;
    }

    [self addAnnotationToChildNodes:annotation];
}

- (void)removeAnnotation:(RMAnnotation *)annotation
{
    if (!annotation.quadTreeNode)
        return;

    annotation.quadTreeNode = nil;

    @synchronized (annotations)
    {
        [annotations removeObject:annotation];
    }

    [self removeUpwardsAllCachedClusterAnnotations];
}

- (void)annotationDidChangeBoundingBox:(RMAnnotation *)annotation
{
    if (RMProjectedRectContainsProjectedRect(boundingBox, annotation.projectedBoundingBox))
        return;

    [annotation retain];

    [self removeAnnotation:annotation];

    RMQuadTreeNode *nextParentNode = self;

    while ((nextParentNode = [nextParentNode parentNode]))
    {
        if (RMProjectedRectContainsProjectedRect(nextParentNode.boundingBox, annotation.projectedBoundingBox))
        {
            [nextParentNode addAnnotationToChildNodes:annotation];
            break;
        }
    }

    [annotation release];
}

- (NSUInteger)countEnclosedAnnotations
{
    NSUInteger count = [annotations count];
    count += [northWest countEnclosedAnnotations];
    count += [northEast countEnclosedAnnotations];
    count += [southWest countEnclosedAnnotations];
    count += [southEast countEnclosedAnnotations];

    return count;
}

- (NSArray *)enclosedAnnotations
{
    NSMutableArray *enclosedAnnotations = [NSMutableArray arrayWithArray:self.annotations];
    if (northWest) [enclosedAnnotations addObjectsFromArray:northWest.enclosedAnnotations];
    if (northEast) [enclosedAnnotations addObjectsFromArray:northEast.enclosedAnnotations];
    if (southWest) [enclosedAnnotations addObjectsFromArray:southWest.enclosedAnnotations];
    if (southEast) [enclosedAnnotations addObjectsFromArray:southEast.enclosedAnnotations];

    return enclosedAnnotations;
}

- (NSArray *)unclusteredAnnotations
{
    NSMutableArray *unclusteredAnnotations = [NSMutableArray array];

    @synchronized (annotations)
    {
        for (RMAnnotation *annotation in annotations)
        {
            if (!annotation.clusteringEnabled)
                [unclusteredAnnotations addObject:annotation];
        }
    }

    if (northWest) [unclusteredAnnotations addObjectsFromArray:[northWest unclusteredAnnotations]];
    if (northEast) [unclusteredAnnotations addObjectsFromArray:[northEast unclusteredAnnotations]];
    if (southWest) [unclusteredAnnotations addObjectsFromArray:[southWest unclusteredAnnotations]];
    if (southEast) [unclusteredAnnotations addObjectsFromArray:[southEast unclusteredAnnotations]];

    return unclusteredAnnotations;
}

- (RMAnnotation *)clusterAnnotation
{
    return cachedClusterAnnotation;
}

- (NSArray *)clusteredAnnotations
{
    return cachedClusterEnclosedAnnotations;
}

- (void)addAnnotationsInBoundingBox:(RMProjectedRect)aBoundingBox toMutableArray:(NSMutableArray *)someArray createClusterAnnotations:(BOOL)createClusterAnnotations withClusterSize:(RMProjectedSize)clusterSize findGravityCenter:(BOOL)findGravityCenter
{
    if (createClusterAnnotations)
    {
        double halfWidth = boundingBox.size.width / 2.0;
        BOOL forceClustering = (boundingBox.size.width >= clusterSize.width && halfWidth < clusterSize.width);
        NSArray *enclosedAnnotations = nil;

        // Leaf clustering
        if (!forceClustering && nodeType == nodeTypeLeaf && [annotations count] > 1)
        {
            NSMutableArray *annotationsToCheck = [NSMutableArray arrayWithArray:self.enclosedAnnotations];

            for (NSInteger i=[annotationsToCheck count]-1; i>0; --i)
            {
                BOOL mustBeClustered = NO;
                RMAnnotation *currentAnnotation = [annotationsToCheck objectAtIndex:i];

                for (NSInteger j=i-1; j>=0; --j)
                {
                    RMAnnotation *secondAnnotation = [annotationsToCheck objectAtIndex:j];

                    // This is of course not very accurate but is good enough for this use case
                    double distance = RMEuclideanDistanceBetweenProjectedPoints(currentAnnotation.projectedLocation, secondAnnotation.projectedLocation) / mapView.metersPerPixel;
                    if (distance < kMinPixelDistanceForLeafClustering)
                    {
                        mustBeClustered = YES;
                        break;
                    }
                }

                if (!mustBeClustered)
                {
                    [someArray addObject:currentAnnotation];
                    [annotationsToCheck removeObjectAtIndex:i];
                }
            }

            forceClustering = [annotationsToCheck count] > 0;

            if (forceClustering)
            {
                [cachedClusterAnnotation release]; cachedClusterAnnotation = nil;
                [cachedClusterEnclosedAnnotations release]; cachedClusterEnclosedAnnotations = nil;

                enclosedAnnotations = [NSArray arrayWithArray:annotationsToCheck];
            }
        }

        if (forceClustering)
        {
            if (!enclosedAnnotations)
                enclosedAnnotations = self.enclosedAnnotations;

            if (cachedClusterAnnotation && [enclosedAnnotations count] != [cachedClusterEnclosedAnnotations count])
            {
                [cachedClusterAnnotation release]; cachedClusterAnnotation = nil;
                [cachedClusterEnclosedAnnotations release]; cachedClusterEnclosedAnnotations = nil;
            }

            if (!cachedClusterAnnotation)
            {
                NSUInteger enclosedAnnotationsCount = [enclosedAnnotations count];

                if (enclosedAnnotationsCount < 2)
                {
                    [someArray addObjectsFromArray:enclosedAnnotations];
                    return;
                }

                RMProjectedPoint clusterMarkerPosition;

                if (findGravityCenter)
                {
                    double averageX = 0.0, averageY = 0.0;

                    for (RMAnnotation *annotation in enclosedAnnotations)
                    {
                        averageX += annotation.projectedLocation.x;
                        averageY += annotation.projectedLocation.y;
                    }

                    averageX /= (double)enclosedAnnotationsCount;
                    averageY /= (double) enclosedAnnotationsCount;

                    double halfClusterWidth = clusterSize.width / 2.0, halfClusterHeight = clusterSize.height / 2.0;

                    if (averageX - halfClusterWidth < boundingBox.origin.x)
                        averageX = boundingBox.origin.x + halfClusterWidth;
                    if (averageX + halfClusterWidth > boundingBox.origin.x + boundingBox.size.width)
                        averageX = boundingBox.origin.x + boundingBox.size.width - halfClusterWidth;
                    if (averageY - halfClusterHeight < boundingBox.origin.y)
                        averageY = boundingBox.origin.y + halfClusterHeight;
                    if (averageY + halfClusterHeight > boundingBox.origin.y + boundingBox.size.height)
                        averageY = boundingBox.origin.y + boundingBox.size.height - halfClusterHeight;

                    // TODO: anchorPoint
                    clusterMarkerPosition = RMProjectedPointMake(averageX, averageY);
                }
                else
                {
                    clusterMarkerPosition = RMProjectedPointMake(boundingBox.origin.x + halfWidth, boundingBox.origin.y + (boundingBox.size.height / 2.0));
                }

                CLLocationCoordinate2D clusterMarkerCoordinate = [[mapView projection] projectedPointToCoordinate:clusterMarkerPosition];

                cachedClusterAnnotation = [[RMAnnotation alloc] initWithMapView:mapView
                                                                     coordinate:clusterMarkerCoordinate
                                                                       andTitle:[NSString stringWithFormat:@"%d", enclosedAnnotationsCount]];
                cachedClusterAnnotation.annotationType = kRMClusterAnnotationTypeName;
                cachedClusterAnnotation.userInfo = self;

                cachedClusterEnclosedAnnotations = [[NSArray alloc] initWithArray:enclosedAnnotations];
            }

            [someArray addObject:cachedClusterAnnotation];
            [someArray addObjectsFromArray:[self unclusteredAnnotations]];

            return;
        }

        if (nodeType == nodeTypeLeaf)
        {
            @synchronized (annotations)
            {
                [someArray addObjectsFromArray:annotations];
            }

            return;
        }
    }
    else
    {
        if (nodeType == nodeTypeLeaf)
        {
            @synchronized (annotations)
            {
                [someArray addObjectsFromArray:annotations];
            }

            return;
        }
    }

    if (RMProjectedRectIntersectsProjectedRect(aBoundingBox, northWestBoundingBox))
        [northWest addAnnotationsInBoundingBox:aBoundingBox toMutableArray:someArray createClusterAnnotations:createClusterAnnotations withClusterSize:clusterSize findGravityCenter:findGravityCenter];
    if (RMProjectedRectIntersectsProjectedRect(aBoundingBox, northEastBoundingBox))
        [northEast addAnnotationsInBoundingBox:aBoundingBox toMutableArray:someArray createClusterAnnotations:createClusterAnnotations withClusterSize:clusterSize findGravityCenter:findGravityCenter];
    if (RMProjectedRectIntersectsProjectedRect(aBoundingBox, southWestBoundingBox))
        [southWest addAnnotationsInBoundingBox:aBoundingBox toMutableArray:someArray createClusterAnnotations:createClusterAnnotations withClusterSize:clusterSize findGravityCenter:findGravityCenter];
    if (RMProjectedRectIntersectsProjectedRect(aBoundingBox, southEastBoundingBox))
        [southEast addAnnotationsInBoundingBox:aBoundingBox toMutableArray:someArray createClusterAnnotations:createClusterAnnotations withClusterSize:clusterSize findGravityCenter:findGravityCenter];

    @synchronized (annotations)
    {
        for (RMAnnotation *annotation in annotations)
        {
            if (RMProjectedRectIntersectsProjectedRect(aBoundingBox, annotation.projectedBoundingBox))
                [someArray addObject:annotation];
        }
    }
}

- (void)removeUpwardsAllCachedClusterAnnotations
{
    if (parentNode)
        [parentNode removeUpwardsAllCachedClusterAnnotations];

    [cachedClusterAnnotation release]; cachedClusterAnnotation = nil;
    [cachedClusterEnclosedAnnotations release]; cachedClusterEnclosedAnnotations = nil;
}

@end

#pragma mark -
#pragma mark RMQuadTree implementation

@implementation RMQuadTree

- (id)initWithMapView:(RMMapView *)aMapView
{
    if (!(self = [super init]))
        return nil;

    mapView = aMapView;
    rootNode = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:nil inBoundingBox:[[RMProjection googleProjection] planetBounds]];

    return self;
}

- (void)dealloc
{
    mapView = nil;
    [rootNode release]; rootNode = nil;
    [super dealloc];
}

- (void)addAnnotation:(RMAnnotation *)annotation
{
    @synchronized (self)
    {
        [rootNode addAnnotation:annotation];
    }
}

- (void)removeAnnotation:(RMAnnotation *)annotation
{
    @synchronized (self)
    {
        [annotation.quadTreeNode removeAnnotation:annotation];
    }
}

- (void)removeAllObjects
{
    @synchronized (self)
    {
        [rootNode release];
        rootNode = [[RMQuadTreeNode alloc] initWithMapView:mapView forParent:nil inBoundingBox:[[RMProjection googleProjection] planetBounds]];
    }
}

#pragma mark -

- (NSArray *)annotationsInProjectedRect:(RMProjectedRect)boundingBox
{
    return [self annotationsInProjectedRect:boundingBox createClusterAnnotations:NO withClusterSize:RMProjectedSizeMake(0.0, 0.0) findGravityCenter:NO];
}

- (NSArray *)annotationsInProjectedRect:(RMProjectedRect)boundingBox createClusterAnnotations:(BOOL)createClusterAnnotations withClusterSize:(RMProjectedSize)clusterSize findGravityCenter:(BOOL)findGravityCenter
{
    NSMutableArray *annotations = [NSMutableArray array];

    @synchronized (self)
    {
        [rootNode addAnnotationsInBoundingBox:boundingBox toMutableArray:annotations createClusterAnnotations:createClusterAnnotations withClusterSize:clusterSize findGravityCenter:findGravityCenter];
    }

    return annotations;
}

@end