pages/01_lightning.py

import ee
import geemap
import ipyleaflet
import solara
import ipywidgets as widgets
from IPython.display import display

zoom = solara.reactive(3)
center = solara.reactive([20, 0])


def zonal_stats_chart(image, vector, **kwargs):
    if isinstance(vector, ee.Geometry):
        fc = ee.FeatureCollection(vector)
    elif isinstance(vector, ee.FeatureCollection):
        fc = vector
    else:
        raise ValueError(
            "The vector argument must be an ee.Geometry or ee.FeatureCollection."
        )
    result = geemap.zonal_stats(
        image, fc, statistics_type="SUM", return_fc=True, verbose=False, **kwargs
    )
    df = geemap.ee_to_df(result).T
    df.reset_index(inplace=True)
    df.columns = ["Type", "Area"]
    chart = geemap.bar_chart(df, "Type", "Area", x_label='', y_label="Area (m2)")
    chart.update_layout(
        margin=dict(l=0, r=0, t=10, b=0),
        height=280,
    )

    return chart


def add_analysis_gui(m=None, position='topright', opened=True):
    """Create a toolbar widget.
    Args:
        m (geemap.Map, optional): The geemap.Map instance. Defaults to None.
        opened (bool, optional): Whether to open the toolbar. Defaults to True.
    """

    fc = ee.FeatureCollection('users/giswqs/public/countries')
    countries = fc.aggregate_array('NAME').getInfo()
    countries.sort()
    gswe = ee.ImageCollection("users/h2i_lab/gswe/gswe_datasets")
    image = gswe.mosaic()
    # esa = gswe.select("esa").mosaic()
    # esri = gswe.select("esri").mosaic()
    # jrc = gswe.select("jrc").mosaic()
    # osm = gswe.select("osm").mosaic()
    # hydrolakes = gswe.select("hydrolakes").mosaic()

    widget_width = "270px"
    padding = "0px 0px 0px 5px"  # upper, right, bottom, left

    toolbar_button = widgets.ToggleButton(
        value=False,
        tooltip="Toolbar",
        icon="bar-chart",
        layout=widgets.Layout(width="28px", height="28px", padding="0px 0px 0px 4px"),
    )

    close_button = widgets.ToggleButton(
        value=False,
        tooltip="Close the tool",
        icon="times",
        button_style="primary",
        layout=widgets.Layout(height="28px", width="28px", padding="0px 0px 0px 4px"),
    )

    options = ["Draw an area", "Select a country"]

    radio = widgets.RadioButtons(
        options=options,
        layout=widgets.Layout(width=widget_width, padding=padding),
        style={"description_width": "initial"},
    )

    country = widgets.Dropdown(
        options=countries,
        value=None,
        layout=widgets.Layout(width=widget_width, padding=padding),
    )

    buttons = widgets.ToggleButtons(
        value=None,
        options=["Apply", "Reset", "Close"],
        tooltips=["Apply", "Reset", "Close"],
        button_style="primary",
        layout=widgets.Layout(padding="0px 2px 4px 2px"),
    )
    buttons.style.button_width = "88px"
    label = widgets.Label("Draw an area on the map first.")

    toolbar_widget = widgets.VBox()
    toolbar_widget.children = [toolbar_button]
    toolbar_header = widgets.HBox()
    toolbar_header.children = [close_button, toolbar_button]
    toolbar_footer = widgets.VBox()
    toolbar_footer.children = [
        radio,
        buttons,
    ]

    def change_radio(change):
        if change["new"] == "Select a country":
            toolbar_footer.children = [radio, country, buttons]
        else:
            toolbar_footer.children = [radio, buttons]

    radio.observe(change_radio, "value")

    m.selected_country = None

    def change_country(change):
        if change["new"]:
            country_name = country.value
            country_fc = fc.filter(ee.Filter.eq('NAME', country_name))
            vec_style = {'color': '000000ff', 'width': 3, 'fillColor': '00000000'}
            m.addLayer(country_fc.style(**vec_style), {}, 'Selected Country')
            m.centerObject(country_fc)
            m.selected_country = country_fc
            toolbar_footer.children = [radio, country, buttons]

    country.observe(change_country, "value")

    def toolbar_btn_click(change):
        if change["new"]:
            close_button.value = False
            toolbar_widget.children = [toolbar_header, toolbar_footer]
        else:
            if not close_button.value:
                toolbar_widget.children = [toolbar_button]

    toolbar_button.observe(toolbar_btn_click, "value")

    def close_btn_click(change):
        if change["new"]:
            toolbar_button.value = False
            if m is not None:
                if m.tool_control is not None and m.tool_control in m.controls:
                    m.remove_control(m.tool_control)
                    m.tool_control = None
            toolbar_widget.close()

    close_button.observe(close_btn_click, "value")

    def button_clicked(change):
        if change["new"] == "Apply":
            output = widgets.Output(
                layout=widgets.Layout(width=widget_width, padding=padding)
            )
            if radio.value == "Select a country":
                toolbar_footer.children = [radio, country, buttons, output]
            else:
                toolbar_footer.children = [radio, buttons, output]
            with output:
                output.clear_output()

                buttons.value = None

                if radio.value == "Draw an area":
                    if m.user_roi is None:
                        display(label)
                        buttons.value = None
                    else:
                        chart = zonal_stats_chart(image, m.user_roi, scale=100)
                        display(chart)
                elif radio.value == "Select a country":
                    if m.selected_country is not None:
                        chart = zonal_stats_chart(
                            image, m.selected_country.geometry(), scale=100
                        )
                        display(chart)

        elif change["new"] == "Reset":
            country.value = None
            radio.value = "Draw an area"
            toolbar_footer.children = [radio, buttons]

        elif change["new"] == "Close":
            if m is not None:
                if m.tool_control is not None and m.tool_control in m.controls:
                    m.remove_control(m.tool_control)
                    m.tool_control = None
            toolbar_widget.close()

        buttons.value = None

    buttons.observe(button_clicked, "value")

    toolbar_button.value = opened
    if m is not None:
        toolbar_control = ipyleaflet.WidgetControl(
            widget=toolbar_widget, position=position
        )

        if toolbar_control not in m.controls:
            m.add_control(toolbar_control)
            m.tool_control = toolbar_control
    else:
        return toolbar_widget


class Map(geemap.Map):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.add_basemap('SATELLITE', show=False)
        self.add_ee_data()
        self.add_layer_manager()
        add_analysis_gui(self)
        # self.add_inspector()

    def add_ee_data(self):
        gswe = ee.ImageCollection("users/h2i_lab/gswe/gswe_datasets")
        self.addLayer(gswe.select("esa"), {'palette': ['red']}, "ESA")
        self.addLayer(gswe.select("esri"), {'palette': ['yellow']}, "ESRI")
        self.addLayer(gswe.select("jrc"), {'palette': ['blue']}, "JRC")
        self.addLayer(gswe.select("osm"), {'palette': ['green']}, "OSM")
        self.addLayer(gswe.select("hydrolakes"), {'palette': ['purple']}, "Hydrolakes")

        legend_dict = {
            'ESA': 'ff0000',
            'ESRI': 'ffff00',
            'JRC': '0000ff',
            'OSM': '00ff00',
            'Hydrolakes': '800080',
        }
        self.add_legend(legend_dict=legend_dict, position='bottomleft')
        fc = ee.FeatureCollection('users/giswqs/public/countries')
        style = {'color': '000000ff', 'width': 1, 'fillColor': '00000000'}
        self.addLayer(fc.style(**style), {}, 'Countries', False)


@solara.component
def Page():
    with solara.Column(style={"min-width": "500px"}):
        Map.element(  # type: ignore
            zoom=zoom.value,
            on_zoom=zoom.set,
            center=center.value,
            on_center=center.set,
            scroll_wheel_zoom=True,
            add_google_map=True,
            height="800px",
            data_ctrl=False,
        )