import h3

import pandas as pd

from shapely import geometry

from shapely.geometry import Polygon, MultiPolygon

import plotly.express as px

import pyspark.sql.functions as func

​

​

# function to plot a single field of hexes, and use a column as the coloration to visualize the fields distribution

def plot_field_op_labelled_hexes(df, operation_guid, operation_col='FieldOperationGuid', hex_col='hex_index_L12',

                                 col_to_color=None, title=None, color_scale='rdylgn', mapbox='open-street-map'):

    # filter pyspark df to just desired op

    df = df.where(func.col(operation_col).eqNullSafe(operation_guid))

    # convert from pyspark to pandas

    df = df.toPandas()

    # get list of all hex ids

    hex_list = df[hex_col]

    # use h3 to geo boundary to get a geojson polygon list

    polygon_list = [[hex_index, Polygon(h3.h3_to_geo_boundary(hex_index, geo_json=True))] for hex_index in hex_list]

    # make dict

    geo_dict = {}

    # fill dictionary with polygon list vals

    geo_dict["type"] = "FeatureCollection"

    geo_dict["features"] = [{"type": "Feature", 'properties': {'hex_index_L12': a[0]}, "geometry": a[1]} for a in

                            [[b[0], geometry.mapping(b[1])] for b in polygon_list]]

    for i, x in enumerate(geo_dict['features']):

        x['id'] = i

    # set center of field

    center = polygon_list[0][1].centroid.wkt.replace('(', '').replace(')', '').split(' ')[1:]

    # plot op using choropleth mapbox

    fig = px.choropleth_mapbox(df, geojson=geo_dict, locations=hex_col, color=col_to_color, hover_data=[col_to_color],

                               hover_name=hex_col, featureidkey="properties.hex_index_L12",

                               mapbox_style=mapbox,

                               zoom=15, center={"lat": float(center[1]), "lon": float(center[0])},

                               opacity=0.8,

                               title=title,

                               color_continuous_scale=color_scale,

                               height=1000,

                               width=1000)

    # show plot

    fig.show()