import math
import numpy as np
import matplotlib.pyplot as plt
from geopy.distance import geodesic
from shapely.geometry import Polygon

# min and max latitudes, longitudes of Washington DC
WA = \
    {
        "min_lat": 38.79164435,
        "max_lat": 39.031386,
        "min_lng": -77.11979522,
        "max_lng": -76.867218
    }

SQRT3 = 1.73205080756887729352744634150587236


def _calculate_polygons(startx, starty, endx, endy, radius):
    """
    Calculate a grid of hexagon coordinates of the given radius
    given lower-left and upper-right coordinates
    Returns a list of lists containing 6 tuples of x, y point coordinates
    These can be used to construct valid regular hexagonal polygons

    You will probably want to use projected coordinates for this
    """
    # calculate side length given radius
    sl = (2 * radius) * math.tan(math.pi / 6)
    # calculate radius for a given side-length
    # (a * (math.cos(math.pi / 6) / math.sin(math.pi / 6)) / 2)
    # see http://www.calculatorsoup.com/calculators/geometry-plane/polygon.php

    # calculate coordinates of the hexagon points
    # sin(30)
    p = sl * 0.5
    b = sl * math.cos(math.radians(30))
    w = b * 2
    h = 2 * sl

    # offset start and end coordinates by hex widths and heights to guarantee coverage
    startx = startx - w
    starty = starty - h
    endx = endx + w
    endy = endy + h

    origx = startx
    origy = starty

    # offsets for moving along and up rows
    xoffset = b
    yoffset = 3 * p

    polygons = []
    row = 1
    counter = 0

    while starty < endy:
        if row % 2 == 0:
            startx = origx + xoffset
        else:
            startx = origx
        while startx < endx:
            p1x = startx
            p1y = starty + p
            p2x = startx
            p2y = starty + (3 * p)
            p3x = startx + b
            p3y = starty + h
            p4x = startx + w
            p4y = starty + (3 * p)
            p5x = startx + w
            p5y = starty + p
            p6x = startx + b
            p6y = starty
            poly = [
                (p1x, p1y),
                (p2x, p2y),
                (p3x, p3y),
                (p4x, p4y),
                (p5x, p5y),
                (p6x, p6y),
                (p1x, p1y)]
            polygons.append(poly)
            counter += 1
            startx += w
        starty += yoffset
        row += 1
    return polygons


def _set_radius(radius: float):
    """
    :param radius: Radius must be selected based on distance between centroids of polygons.
    Radius must make distance 75 meters for each polygon centroids.
    """
    # Radius selected based on distance between centroids of polygons.
    # This radius makes distance 75 meters for each polygon centroids.
    polygons = _calculate_polygons(WA['min_lat'], WA['min_lng'],
                                  WA['max_lat'], WA['max_lng'], radius)
    pol_1 = Polygon(polygons[1])
    pol_2 = Polygon(polygons[2])
    pol_1_centre = pol_1.centroid
    pol_2_centre = pol_2.centroid
    polygon_1_centroid_coords = pol_1_centre.bounds[0], pol_1_centre.bounds[1]
    polygon_2_centroid_coords = pol_2_centre.bounds[0], pol_2_centre.bounds[1]

    # Distance between centroid of first polygon to centroid of second polygon
    # This is equal to 75.04 meters,
    distance = geodesic(polygon_1_centroid_coords, polygon_2_centroid_coords).meters
    assert int(distance) == 75, 'Distance must be equal to 75 meters'

    return radius


try:
    # Radius selected based on distance between centroids of polygons.
    # This radius makes distance 75 meters for each polygon centroids.
    radius = _set_radius(0.000338)
    polygons = _calculate_polygons(WA['min_lat'], WA['min_lng'],
                                  WA['max_lat'], WA['max_lng'], radius)
except AssertionError as error:
    raise error

# Number of polygons that covers Washington DC. Number of polygons = 155295
num_polygons = len(polygons)


def gen_bins_with_polycollection(data: list, gridsize=366):
    """
    From the given array of data points ([[x, y], ...]), generate hexagon bins by using matplotlib PolyCollection object
    Returns hexagon centroids and corresponding vehicle counts in that hexagon.
    :param data: list of latitude and longitude pairs.
    :param gridsize: The number of hexagons in the *x*-direction.
    """
    data = np.array(data)
    x = data[:, 0]
    y = data[:, 1]
    # initialize matplotlib PolyCollection object
    # gridsize is set to a value, which makes the number of polygons
    # generated by plt.hexbin approximately equal to  "num_polygons"
    polycollection = plt.hexbin(x, y, gridsize=gridsize, extent=[WA['min_lat'],WA['max_lat'],
                                                            WA['min_lng'], WA['max_lng']])
    plt.close()

    # centroids of hexagon bins.
    hexagon_centroids = polycollection.get_offsets()

    # corresponding count of vehicles in each hexagon bin.
    count_vehicles = polycollection.get_array()

    # indices where count of vehciles in bins is not 0.
    indices = np.where(count_vehicles != 0)

    non_zero_vehicle_counts = count_vehicles[indices]
    hexagon_centroids = hexagon_centroids[indices]

    sorted_count_indices = np.argsort(non_zero_vehicle_counts)[::-1]

    sorted_hexagon_centers = hexagon_centroids[sorted_count_indices]
    sorted_vehicle_counts = non_zero_vehicle_counts[sorted_count_indices]

    sorted_vehicle_counts_gt_1 = sorted_vehicle_counts[np.where(sorted_vehicle_counts > 1)[0]]
    sorted_hexagon_centers_gt_1 = sorted_hexagon_centers[np.where(sorted_vehicle_counts > 1)[0]]

    return sorted_hexagon_centers_gt_1, sorted_vehicle_counts_gt_1


def _nearest_bins_centerpoint(value: float, scale: float):
    """
    Returns the nearest two centerpoint positions (along the same axis) for the
    closest bins to the specified value.
    :param value: Either the x or y value of the point
    :param scale: The scale of the hexagons
    :returns: tuple of closest bin centerpoints
    """
    div = math.floor(value/(scale / 2))
    rounded = scale / 2 * (div + (1 if div % 2 == 1 else 0))
    rounded_scaled = scale / 2 * (div + (1 if div % 2 == 0 else 0))
    return [rounded, rounded_scaled]


def _distance(x: float, y: float, x1: float, y1: float):
    """
    Returns the Euclidean distance from [x,y] to [x1,y1]
    """
    return math.sqrt((x - x1) * (x - x1) + (y - y1) * (y - y1))


def gen_bins_with_bin_width(data: list, bin_width: float):
    """
    From the given array of data points ([[x, y], ...]), generate an array of
    hexagonal bins of given bin_width
    :param data: list of latitude and longitude pairs.
    :param bin_width: the size of the bins to create
    """
    bins = {}
    data = np.array(data)
    for p in data:
        x, y = p[0], p[1]

        px_nearest = _nearest_bins_centerpoint(x, bin_width)
        py_nearest = _nearest_bins_centerpoint(y, bin_width * SQRT3)

        z1 = _distance(x, y, px_nearest[0], py_nearest[0])
        z2 = _distance(x, y, px_nearest[1], py_nearest[1])

        if z1 < z2:
            bin = str([px_nearest[0], py_nearest[0]])
        else:
            bin = str([px_nearest[1], py_nearest[1]])

        bins[bin] = 1 + bins.get(bin, 0)

    hexagon_centers = list()
    vehicle_counts = list()
    for k, v in bins.items():
        hexagon_centers.append(eval(k))
        vehicle_counts.append(v)
    return hexagon_centers, vehicle_counts

