#!/usr/bin/env python3

import sys
import svgwrite
import numpy as np
import matplotlib.pyplot as pp
from matplotlib.colors import hsv_to_rgb
import json
import random
import argparse

REF_LATITUDE = 49.58666
REF_LONGITUDE = 11.01250
# REF_LATITUDE = -30
# REF_LONGITUDE = 90


def map_azimuthal_equidistant(lat, lon, ref_lat, ref_lon, R=1):
    """ Azimuthal equidistant projection.

    This function takes a point to map in latitude/longitude format as well as
    a reference point which becomes the "center" of the map.

    It then projects the point into the 2D plane such that the distance from
    the center is proportional to the distance on the Great Circle through the
    projected and the reference point. The angle represents the azimuthal
    direction of the projected point.

    Args:
        lat(numpy.array):   Latitudes of the point to project.
        lon(numpy.array):   Longitudes of the point to project.
        ref_lat(float):     Latitude of the reference point.
        ref_lon(float):     Longitude of the reference point.
        R(float):           Radius (scale) of the map.

    Returns:
        x(numpy.array):   The calculated x coordinates.
        y(numpy.array):   The calculated y coordinates.

    """
    dlon = lon - ref_lon

    rho_linear_norm = np.arccos(np.sin(ref_lat) * np.sin(lat)
                                + np.cos(ref_lat)
                                * np.cos(lat) * np.cos(dlon)) / np.pi

    rho = R * rho_linear_norm

    theta = np.arctan2(np.cos(lat) * np.sin(dlon),
                       (np.cos(ref_lat) * np.sin(lat)
                        - np.sin(ref_lat) * np.cos(lat)
                        * np.cos(dlon)))

    x = rho * np.sin(theta)
    y = -rho * np.cos(theta)

    return x, y


def random_country_color():
    h = random.random()
    s = 0.7
    v = 0.8
    r, g, b = [int(255.99*x) for x in hsv_to_rgb([h, s, v])]
    return f"#{r:02x}{g:02x}{b:02x}"


def is_point_in_polygon(point, polygon):
    # Idea: draw an infinite line from the test point along the x axis to the
    # right. Then check how many polygon edges this line intersects. If the
    # number is even, the point is outside the polygon.

    edges = []  # list of lists, containing two points each

    for i in range(len(polygon)-1):
        edges.append([polygon[i], polygon[i+1]])

    # the closing edge
    edges.append([polygon[-1], polygon[0]])

    num_intersects = 0

    test_x, test_y = point

    for edge in edges:
        start_x = edge[0][0]
        start_y = edge[0][1]
        end_x = edge[1][0]
        end_y = edge[1][1]

        # quick exclusion tests
        if start_x < test_x and end_x < test_x:
            continue  # edge is completely left of the test point

        if start_y < test_y and end_y < test_y:
            continue  # edge is completely below the test point

        if start_y > test_y and end_y > test_y:
            continue  # edge is completely above the test point

        # calculate the x coordinate where the edge intersects the whole
        # horizontal line
        intersect_x = start_x + (end_x - start_x) \
            * (test_y - start_y) \
            / (end_y - start_y)

        if intersect_x > test_x:
            # we found an intersection!
            num_intersects += 1

    if num_intersects % 2 == 0:
        return False  # even number of intersects -> outside polygon
    else:
        return True  # odd number of intersects -> inside polygon


def svg_make_inverse_country_path(doc, map_radius, polygon, **kwargs):
    # build a closed circle path covering the whole map
    commands = [f"M 0, {map_radius}",
                f"a {map_radius},{map_radius} 0 1,0 {map_radius*2},0",
                f"a {map_radius},{map_radius} 0 1,0 {-map_radius*2},0",
                "z"]

    # "subtract" the country polygon
    commands.append(f"M {polygon[0][0]} {polygon[0][1]}")

    # add lines for each polygon point
    for point in polygon[1:]:
        commands.append(f"L {point[0]} {point[1]}")

    # ensure straight closing line
    commands.append(f"L {polygon[0][0]} {polygon[0][1]}")

    # close the inner path
    commands.append("z")

    return doc.path(commands, **kwargs)


def render(ref_lat, ref_lon, output_stream):
    random.seed(0)

    """ Test code
    test_lat = [np.pi/2, np.pi/4, 0, -np.pi/4, -np.pi/2]
    test_lon = np.arange(-np.pi, np.pi, np.pi/4)

    for lat in test_lat:
        for lon in test_lon:
            x, y = map_azimuthal_equidistant(np.array([lat]), np.array([lon]), np.pi/2, 0)

            print(f"{lat*180/np.pi:6.3f}, {lon*180/np.pi:6.3f} => {x[0]:6.3f}, {y[0]:6.3f}", file=sys.stderr)
    """

    print("Loading Geodata…", file=sys.stderr)

    with open('geo-countries/data/countries.geojson', 'r') as jfile:
        geojson = json.load(jfile)

    print("Finding boundaries…", file=sys.stderr)

    # key: 3-letter country identifier
    # data: {full_name, numpy.array(coordinates), numpy.array(proj_coordinates)}.
    # coordinates is a list of 2xN arrays, where N is the number of points. Row 0
    # contains the longitude, Row 1 the latitude.
    # proj_coordinates is a list of 2xN arrays, where N is the number of points.
    # Row 0 contains the projected x, Row 1 the projected y.
    simplegeodata = {}

    features = geojson['features']

    for feature in features:
        name = feature['properties']['ADMIN']
        key = feature['properties']['ISO_A2']

        # handle duplicate keys (can happen for small countries)
        if key in simplegeodata.keys():
            key = name

        print(f"Preparing {key} ({name})…", file=sys.stderr)

        multipoly = feature['geometry']['coordinates']

        conv_polys = []

        for poly in multipoly:
            for subpoly in poly:
                coords_list = []  # list of lists

                assert(len(subpoly[0]) == 2)
                coords_list += subpoly

                # convert coordinates to numpy array and radians
                coords = np.array(coords_list).T * np.pi / 180

                conv_polys.append(coords)

        simplegeodata[key] = {"name": name, "coordinates": conv_polys}

    ref_lat = ref_lat * np.pi / 180
    ref_lon = ref_lon * np.pi / 180

    antipodal_lat = -ref_lat
    antipodal_lon = ref_lon + np.pi

    if antipodal_lon > np.pi:
        antipodal_lon -= 2*np.pi

    R = 500

    """
    # Override data with test coordinate system
    coords = []

    N = 128

    # constant-latitude circles
    coords.append(np.array([np.linspace(-np.pi, np.pi, N),
                            np.ones(N) * np.pi/4]))
    coords.append(np.array([np.linspace(-np.pi, np.pi, N),
                            np.ones(N) * 0]))
    coords.append(np.array([np.linspace(-np.pi, np.pi, N),
                            np.ones(N) * -np.pi/4]))

    # constant-longitude half-circles
    coords.append(np.array([np.ones(N) * -4*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * -3*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * -2*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * -1*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * 0*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * 1*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * 2*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))
    coords.append(np.array([np.ones(N) * 3*np.pi/4,
                            np.linspace(-np.pi/2, np.pi/2, N)]))

    simplegeodata = {"XY": {'name': 'test', 'coordinates': coords}}
    """

    # apply azimuthal equidistant projection
    for k, v in simplegeodata.items():
        proj_polys = []

        for poly in v['coordinates']:
            lat = poly[1, :]
            lon = poly[0, :]

            x, y = map_azimuthal_equidistant(lat, lon, ref_lat, ref_lon, R)

            coords = np.array([x, y])

            # remove any points that contain a NaN coordinate
            coords = coords[:, np.any(np.invert(np.isnan(coords)), axis=0)]

            proj_polys.append(coords)

        v['proj_coordinates'] = proj_polys

    # generate the SVG

    doc = svgwrite.Drawing("/tmp/test.svg", size=(2*R, 2*R))

    doc.defs.add(doc.style("""
            .country {
                stroke: black;
                stroke-width: 0.01px;
            }

            .dist_circle_label, .azimuth_line_label {
                font-size: 3px;
                font: sans-serif;
                text-anchor: middle;
            }

            .dist_circle, .azimuth_line {
                fill: none;
                stroke: black;
                stroke-width: 0.1px;
            }

            .maidenhead_line {
                fill: none;
                stroke: red;
                stroke-width: 0.1px;
                opacity: 0.5;
            }
    """))

    doc.add(doc.circle(center=(R, R), r=R, fill='#ddeeff',
                       stroke_width=1, stroke='black'))

    for k, v in simplegeodata.items():
        print(f"Exporting {k}…", file=sys.stderr)

        color = random_country_color()

        group = doc.g()

        for i in range(len(v['proj_coordinates'])):
            poly = v['proj_coordinates'][i]
            points = poly.T + R  # shift to the center of the drawing

            # check if the antipodal point is inside this polygon. If so, it
            # needs to be "inverted", i.e. subtracted from the surrounding map
            # circle.

            if is_point_in_polygon((antipodal_lon, antipodal_lat),
                                   v['coordinates'][i].T):
                print("!!! Found polygon containing the antipodal point!")
                obj = svg_make_inverse_country_path(doc, R, np.flipud(points), **{
                                  'class': 'country',
                                  'fill': color})
            else:
                obj = doc.polygon(points, **{
                                  'class': 'country',
                                  'fill': color})

            group.add(obj)

        group.set_desc(title=v['name'])
        doc.add(group)

    # generate equidistant circles

    d_max = 40075/2
    for distance in [500, 1000, 2000, 3000, 4000, 5000, 6000, 8000, 10000, 12000,
                     14000, 16000, 18000, 20000]:
        r = R * distance / d_max
        doc.add(doc.circle(center=(R, R), r=r,
                           **{'class': 'dist_circle'}))

        doc.add(doc.text(f"{distance} km", (R, R-r+5),
                **{'class': 'dist_circle_label'}))

    # generate azimuth lines in 30° steps

    for azimuth in np.arange(0, np.pi, np.pi/6):
        start_x = R + R * np.cos(azimuth-np.pi/2)
        start_y = R + R * np.sin(azimuth-np.pi/2)
        end_x = R - R * np.cos(azimuth-np.pi/2)
        end_y = R - R * np.sin(azimuth-np.pi/2)

        doc.add(doc.line((start_x, start_y), (end_x, end_y),
                **{'class': 'azimuth_line'}))

        azimuth_deg = int(np.round(azimuth * 180 / np.pi))
        textpos = (2*R - 10, R - 2)

        txt = doc.text(f"{azimuth_deg:d} °", textpos,
                       **{'class': 'azimuth_line_label'})
        txt.rotate(azimuth_deg - 90, center=(R, R))
        doc.add(txt)

        txt = doc.text(f"{azimuth_deg+180:d} °", textpos,
                       **{'class': 'azimuth_line_label'})
        txt.rotate(azimuth_deg - 90 + 180, center=(R, R))
        doc.add(txt)

    # generate Maidenhead locator grid (first two letters only)

    group = doc.g()

    N = 18  # subdivisions of Earth
    resolution = 4096

    for x in range(0, N):
        lon = x * 2 * np.pi / N
        lat = np.linspace(-np.pi/2, np.pi/2, resolution)

        x, y = map_azimuthal_equidistant(lat, lon, ref_lat, ref_lon, R)
        points = np.array([x, y]).T + R

        group.add(doc.polyline(points, **{'class': 'maidenhead_line'}))

    for y in range(0, N):
        lon = np.linspace(-np.pi, np.pi, resolution)
        lat = y * np.pi / N - np.pi/2

        x, y = map_azimuthal_equidistant(lat, lon, ref_lat, ref_lon, R)
        points = np.array([x, y]).T + R

        group.add(doc.polyline(points, **{'class': 'maidenhead_line'}))

    doc.add(group)

    """
    for x in range(0, 26):
        for y in range(0, 26):
            sectorname = chr(ord('A')+x) + chr(ord('A')+y)
    """

    print(f"Writing output…", file=sys.stderr)
    doc.write(output_stream, pretty=True)

    return

    # Debug Plot

    for k, v in simplegeodata.items():
        for poly in v['proj_coordinates']:
            pp.plot(poly[0, :], poly[1, :])

    pp.plot([-1, 1], [0, 0], 'k', linewidth=0.5)
    pp.plot([0, 0], [-1, 1], 'k', linewidth=0.5)

    t = np.linspace(-np.pi, np.pi, 256)
    ct, st = np.cos(t), np.sin(t)
    pp.plot(ct, st, 'k', linewidth=0.5)

    U = 40075
    for distance in np.arange(0, U/2, 2000):
        f = distance / (U/2)
        pp.plot(f*ct, f*st, 'k', linewidth=0.2)

    pp.axis('equal')

    pp.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
            description="Render an azimuthal equidistant map of the world " +
                        "centered on the given point")

    parser.add_argument(metavar='ref-lat', type=float, dest='ref_lat',
                        help='Reference Latitude')
    parser.add_argument(metavar='ref-lon', type=float, dest='ref_lon',
                        help='Reference Longitude')
    parser.add_argument('-o', '--output-file', type=argparse.FileType('w'),
                        help='The output SVG file (default: print to stdout)',
                        default=sys.stdout)

    args = parser.parse_args()

    render(args.ref_lat, args.ref_lon, args.output_file)