#!/usr/bin/env python3

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

REF_LATITUDE = 49.58666
REF_LONGITUDE = 11.01250
# REF_LATITUDE = 1
# REF_LONGITUDE = 0


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}"


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}")
"""

print("Loading Geodata…")

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

print("Finding boundaries…")

# 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})…")

    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_LATITUDE * np.pi / 180
ref_lon = REF_LONGITUDE * np.pi / 180

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

#print(simplegeodata['AW']['proj_coordinates'])
###print(simplegeodata['DE']['proj_coordinates'])

#debug plot
# 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.01;
        }
"""))

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}…")

    color = random_country_color()

    group = doc.g()

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

        pgon = doc.polygon(points, **{
                           'class': 'country',
                           'fill': color})

        group.add(pgon)

    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, fill='none',
                       stroke_width=0.1, stroke='black'))

print(f"Saving {doc.filename}…")
doc.save(pretty=True)

exit(0)

# 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()