# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Calculate location colours
# The colours for locations are calcuated and stored as a YAML file.

# ## Imports

from ag3 import release_data
ag3_data = release_data()

import yaml
import matplotlib.pyplot as plt
import pandas as pd
import cartopy as cart
import cartopy.crs as ccrs
import matplotlib.patches as patches
import numpy as np

# REQUIRES: https://pypi.org/project/colormath/
# pip install colormath
from colormath.color_objects import sRGBColor, LabColor
from colormath.color_conversions import convert_color 



# ## Settings

output_path = 'location_colours.yaml'

# Colour space
lightness = 50 # At 1, the midpoint appears black; at 100 it appears white. All points appear white around 700; black around -300.
colour_blend = 70 # At 1, all points appear equal grey, by `lightness`. Above 500, colours regions have harsh boundaries. Below 70, colours are greyish.
centre_at_equator = False
centre_at_prime_meridian = False
colour_space_orientation = {'NE': 'cyan', 'SE': 'green', 'SW': 'red', 'NW': 'magenta'}
use_location_min_max_as_colour_space = False

# Used these when use_location_min_max_as_colour_space = False:
colour_space_min_lat = -35 # southern tip of Africa
colour_space_max_lat = 20 # as north as the Sahara (roughly)
colour_space_min_long = -26 # west of Cape Verde
colour_space_max_long = 60 # east of Mauritius



# ### Eyeball colour space mapping

# +
ax = plt.axes(projection=ccrs.PlateCarree())
ax.stock_img()

xlim = ([colour_space_min_long - 10, colour_space_max_long + 10])
ylim = ([colour_space_min_lat - 10,  colour_space_max_lat + 10])
ax.set_xlim(xlim)
ax.set_ylim(ylim)

rect = patches.Rectangle(
    (colour_space_min_long, colour_space_min_lat), (colour_space_max_long - colour_space_min_long), (colour_space_max_lat - colour_space_min_lat),
    linewidth=1, edgecolor='black', facecolor='none', linestyle='--'
)
ax.add_patch(rect)

plt.show()
# -



# ## Functions

def latlong_to_rgb_hex_via_lab(lat, long,
                               min_lat=-90, max_lat=90, min_long=-180, max_long=180,
                               centre_at_equator=False, centre_at_prime_meridian=False,
                               colour_space_orientation={'NE': 'cyan', 'SE': 'green', 'SW': 'red', 'NW': 'magenta'},
                               lightness=50, colour_blend=70
                              ):

    observer = '2' # Observer angle: default '2', alternative '10' degrees
    illuminant = 'd50' # default 'd50'
    # When observer '2', illuminant can have {'a', 'b', 'c', 'd50', 'd55', 'd65', 'd75', 'e', 'f2', 'f7', 'f11'}
    # When observer '10', illuminant can have {'d50', 'd55', 'd65', 'd75'}
    
    # Handle options to centre on the equator and/or the prime meridian
    if centre_at_equator:
        if -min_lat > max_lat:
            min_lat = -max_lat
    
    if centre_at_prime_meridian:
        if -min_long > max_long:
            min_Long = -max_long
     
    # Scale the lat-long 0 to 1, relative to the specified min and max
    lat_extent = max_lat - min_lat
    long_extent = max_long - min_long
    scaled_lat = (lat + abs(min_lat)) / lat_extent
    scaled_long = (long + abs(min_long)) / long_extent

    # Convert to range (-1 to 1) * f, according to the specified colour_space_orientation
    if colour_space_orientation == {'NE': 'red', 'SE': 'green', 'SW': 'cyan', 'NW': 'magenta'}:
        a = ((scaled_lat * 2) - 1) * colour_blend
        b = ((scaled_long * 2) - 1) * colour_blend
    elif colour_space_orientation == {'NE': 'red', 'SE': 'magenta', 'SW': 'cyan', 'NW': 'green'}:
        a = ((scaled_long * 2) - 1) * colour_blend
        b = ((scaled_lat * 2) - 1) * colour_blend
    elif colour_space_orientation == {'NE': 'cyan', 'SE': 'green', 'SW': 'red', 'NW': 'magenta'}:
        a = (1 - (scaled_long * 2)) * colour_blend
        b = (1 - (scaled_lat * 2)) * colour_blend
    elif colour_space_orientation == {'NE': 'cyan', 'SE': 'magenta', 'SW': 'red', 'NW': 'green'}:
        a = (1 - (scaled_lat * 2)) * colour_blend
        b = (1 - (scaled_long * 2)) * colour_blend
    else:
        raise ValueError(f'Unhandled colour_space_orientation: {colour_space_orientation}')
    
    # Without "clamping" or adjusting the scale factor (* 255 above), some colours will fall outside compatible space and raise an error during `get_rgb_hex()`
    # From https://python-colormath.readthedocs.io/en/latest/conversions.html
    # "RGB spaces tend to have a smaller gamut than some of the CIE color spaces. When converting to RGB, this can cause some of the coordinates to end up being out of the acceptable range (0.0-1.0 or 1-255, depending on whether your RGB color is upscaled)."
    # "Rather than clamp these for you, we leave them as-is. This allows for more accurate conversions back to the CIE color spaces. If you require the clamped (0.0-1.0 or 1-255) values, use the following properties on any RGB color: clamped_rgb_r, clamped_rgb_g, clamped_rgb_b"
    
    # Switch depending on supplied param type
    if isinstance(lat, int):
        lab = LabColor(lightness, a, b, observer, illuminant)
        rgb = convert_color(lab, sRGBColor)
        clamped_rgb = sRGBColor(rgb.clamped_rgb_r, rgb.clamped_rgb_g, rgb.clamped_rgb_b, is_upscaled=False)
        rgb_hex = clamped_rgb.get_rgb_hex()
        return rgb_hex
    elif isinstance(lat, pd.core.series.Series):
        # TODO: vectorize
        rgb_hex = []
        for i in range(len(lat)):
            lab = LabColor(lightness, a[i], b[i], observer, illuminant)
            rgb = convert_color(lab, sRGBColor)
            clamped_rgb = sRGBColor(rgb.clamped_rgb_r, rgb.clamped_rgb_g, rgb.clamped_rgb_b, is_upscaled=False)
            rgb_hex.append(clamped_rgb.get_rgb_hex())
        return rgb_hex
    else:
        raise ValueError('Unhandled parameter types')



# ## Spot check the lat-long colour conversion

spot_check_data = []
for lat in [-90, -45, 0, 45, 90]:
    for long in [-180, -90, 0, 90, 180]:
        spot = {'lat': lat, 'long': long, 'colour': latlong_to_rgb_hex_via_lab(lat, long, 
                                                                               centre_at_equator=centre_at_equator,
                                                                               centre_at_prime_meridian=centre_at_prime_meridian,
                                                                               colour_space_orientation=colour_space_orientation,
                                                                               lightness=lightness, colour_blend=colour_blend
                                                                              )
               }
        spot_check_data.append(spot)
spot_check_df = pd.DataFrame(spot_check_data)
spot_check_df.head()

# +
# Eyeball some lat-long to rgb_hex conversions
fig = plt.figure(figsize=(8, 8))
ax = plt.axes(projection=ccrs.PlateCarree())
ax.coastlines()
ax.add_feature(cart.feature.OCEAN, zorder=0, edgecolor='k', facecolor='lightblue')

plt.scatter(spot_check_df["long"], spot_check_df["lat"], c=spot_check_df["colour"], s=100_00)
plt.show()
# -



# ## Get the locations of all of the wild samples

wild_sample_meta = ag3_data.load_sample_set_metadata(ag3_data.all_wild_sample_sets)
wild_sample_meta.head()

# Check that there aren't different lat-longs for the same country-location
all_locations = wild_sample_meta[['country', 'location', 'latitude', 'longitude']]
unique_locations = all_locations.drop_duplicates(['country', 'location', 'latitude', 'longitude'])
duplicate_locations_by_lat_long = unique_locations.duplicated(['latitude', 'longitude'], keep=False)
unique_locations[duplicate_locations_by_lat_long]

# +
# TODO: Resolve/Explain these duplicate lat-longs
# -

# We only need the unique locations
locations = wild_sample_meta[['country', 'location', 'latitude', 'longitude']].drop_duplicates().reset_index(drop=True)
locations



# ## Calculate the location colours

# +
if use_location_min_max_as_colour_space:
    min_lat = locations['latitude'].min()
    max_lat = locations['latitude'].max()
    min_long = locations['longitude'].min()
    max_long = locations['longitude'].max()
else:
    min_lat = colour_space_min_lat
    max_lat = colour_space_max_lat
    min_long = colour_space_min_long
    max_long = colour_space_max_long

locations['rgb_hex'] = latlong_to_rgb_hex_via_lab(locations['latitude'],
                                                  locations['longitude'],
                                                  min_lat=min_lat, max_lat=max_lat,
                                                  min_long=min_long, max_long=max_long,
                                                  centre_at_equator=centre_at_equator,
                                                  centre_at_prime_meridian=centre_at_prime_meridian,
                                                  colour_space_orientation=colour_space_orientation,
                                                  lightness=lightness, colour_blend=colour_blend
                                                 )
locations
# -



# ## Check the location colours on a map

# +
ax = plt.axes(projection=ccrs.PlateCarree())
ax.stock_img()
#ax.coastlines(color='grey')

xlim = ([locations["longitude"].min() - 10,  locations["longitude"].max() + 10])
ylim = ([locations["latitude"].min() - 10,  locations["latitude"].max() + 10])
ax.set_xlim(xlim)
ax.set_ylim(ylim)

plt.scatter(locations["longitude"], locations["latitude"], edgecolors=locations["rgb_hex"], alpha=1, s=100, color='None', linewidth=2)

plt.show()
# -



# ## Make a reference colour map

# Generate lat-longs limited to the same boundaries
# TODO: use a less brutal method
ref_lat_longs = []
lat_step = 1
long_step = 1
for lat in np.arange(min_lat, max_lat, lat_step):
    for long in np.arange(min_long, max_long, long_step):
        ref_lat_longs.append({'lat': lat, 'long': long})
ref_lat_longs[0]

# Add the generated lat-longs to a DataFrame, for convenience
ref_locations_df = pd.DataFrame(data=ref_lat_longs, columns=['lat', 'long'])
ref_locations_df.head()

# Get the colours for the lat-longs
ref_locations_df['rgb_hex'] = latlong_to_rgb_hex_via_lab(ref_locations_df['lat'],
                                                            ref_locations_df['long'],
                                                            min_lat=min_lat, max_lat=max_lat,
                                                            min_long=min_long, max_long=max_long,
                                                            centre_at_equator=centre_at_equator,
                                                            colour_space_orientation=colour_space_orientation,
                                                            lightness=lightness, colour_blend=colour_blend
                                                           )
ref_locations_df.head()

# ### Plot the coloured lat-longs on a map

# +
fig = plt.figure(figsize=(10, 10))
ax = plt.axes(projection=ccrs.PlateCarree())

xlim = ([-20, 60]) # west, east (long)
ylim = ([-40,  40]) # south, north (lat)
ax.set_xlim(xlim)
ax.set_ylim(ylim)

ax.add_feature(cart.feature.OCEAN, zorder=1, edgecolor='k', facecolor='white')

s = 45 # Size of square found manually
plt.scatter(ref_locations_df["long"], ref_locations_df["lat"], c=ref_locations_df["rgb_hex"], marker='s', s=s)

plt.show()
# -

# ### Check the locations on the reference map

# +
fig = plt.figure(figsize=(10, 10))
ax = plt.axes(projection=ccrs.PlateCarree())

xlim = ([min_long - 3, max_long + 7]) # west, east (long)
ylim = ([min_lat - 15,  max_lat + 25]) # south, north (lat)
ax.set_xlim(xlim)
ax.set_ylim(ylim)

ax.add_feature(cart.feature.OCEAN, zorder=1, edgecolor='k', facecolor='white')
plt.scatter(ref_locations_df["long"], ref_locations_df["lat"], c=ref_locations_df["rgb_hex"], marker='s', s=s)
rect = patches.Rectangle(
    (min_long, min_lat), (max_long - min_long), (max_lat - min_lat),
    linewidth=1, edgecolor='black', facecolor='none', linestyle='--', zorder=2
)
ax.add_patch(rect)
plt.plot([min_long - 3, max_long + 7], [0, 0], color='black', linewidth=3, marker='None', linestyle='--')
plt.scatter(locations["longitude"], locations["latitude"], c=locations["rgb_hex"], edgecolor='k', alpha=1, s=100, linewidth=1.5, zorder=3)

plt.show()
# -



# ## Store the location colours as a file

# Get the location_labels and their corresponding rgb_hex colours
location_colours = locations[['country', 'location', 'rgb_hex']].set_index(['country', 'location'])
location_colours

# ### Convert the DataFrame values into a dictionary

# +
# Note: Panda's to_dict(orient='index') converts multi-column index into a tuple,
# which isn't what we want for the YAML
#location_colours_as_dict = location_colours.to_dict(orient='index')
#list(location_colours_as_dict.keys())[0]
#ignore_index=True
# -

# Compose the dict manually
location_colours_as_dict = {}
for country, country_group in location_colours.groupby('country'):
    location_colours_as_dict[country] = {}
    for location, location_group in country_group.groupby('location'):
        location_colours_as_dict[country][location] = location_group['rgb_hex'][0]

location_colours_as_dict.keys()

location_colours_as_dict['Burkina Faso']



# ### Write the dictionary to YAML

with open(output_path, 'w') as ymlfile:
    yaml.dump(location_colours_as_dict, ymlfile, sort_keys=True, allow_unicode=True)



# ## Check the location colours file

# Load the data from the YAML file
with open(output_path) as file:
    reloaded_data_dict = yaml.load(file, Loader=yaml.Loader)

reloaded_data_dict.keys()

reloaded_data_dict['Burkina Faso']

# ### Convert the dict into a DataFrame

# +
# Note: Panda's from_dict() and from_records() don't interpret the dict in the way we want
#reloaded_df = pd.DataFrame.from_dict(reloaded_data_dict)
#reloaded_df.head()
# -

# Compose the DataFrame manually
reloaded_df = pd.DataFrame(columns=['country', 'location', 'rgb_hex'])
for country in reloaded_data_dict.keys():
    for location in reloaded_data_dict[country].keys():
        reloaded_df = reloaded_df.append(
            {'country': country, 'location': location, 'rgb_hex': reloaded_data_dict[country][location]},
            ignore_index=True
        )
reloaded_df

# Match the location labels with the lat longs we have
original_locations = locations[['country', 'location', 'longitude', 'latitude']]
recovered_locations = reloaded_df.merge(original_locations, on=['country', 'location'])
recovered_locations

# ### Eyeball the colours from the file on a map

# +
fig = plt.figure(figsize=(8, 8))
ax = plt.axes(projection=ccrs.PlateCarree())

xlim = ([recovered_locations["longitude"].min() - 10,  recovered_locations["longitude"].max() + 10])
ylim = ([recovered_locations["latitude"].min() - 10,  recovered_locations["latitude"].max() + 10])
ax.set_xlim(xlim)
ax.set_ylim(ylim)

ax.gridlines(color='blue')
ax.coastlines()

plt.scatter(recovered_locations["longitude"], recovered_locations["latitude"], color=recovered_locations["rgb_hex"], s=100, zorder=3)

plt.show()
# -



# ## Test `ag3.py` usage

import ag3
import importlib
importlib.reload(ag3)

imported_location_colours = ag3.release_data().location_colours
imported_location_colours['Cameroon'].keys()

imported_location_colours['Cameroon']['Séboré']



StarcoderdataJupyterScriptsDedupFiltered