Py.Cafe

from typing import Optional, Literal
import dash_bootstrap_components as dbc
import geopandas as gpd
import numpy as np
import pandas as pd
import vizro.models as vm
import vizro.plotly.express as px
from dash import html
from vizro import Vizro
from vizro.actions import filter_interaction
from vizro.models.types import capture
from vizro.tables import dash_ag_grid
from dash_ag_grid import AgGrid
from datetime import datetime
from datetime import timedelta


PLANNING_DATE = "2025-04-07"



class FlexContainer(vm.Container):
    """Custom flex `Container`."""

    type: Literal["flex_container"] = "flex_container"
    title: str = ""  # Title exists in vm.Container but we don't want to use it here.
    classname: str = "d-flex w-100 h-100"  # Ensure full width and height

    def build(self):
        """Returns a flex container."""
        return html.Div( 
            id=self.id,
            children=[component.build() for component in self.components],
            className=self.classname,
            style={
                "width": "100%",
                "height": "100%",
                "margin": "0",  # Remove any margin
                "padding": "0",  # Remove any padding
                "overflow": "hidden",  # Prevent scrollbars or extra space
            },
        )
    

vm.Page.add_type("components", FlexContainer)

@capture('figure')
def get_card(
        data_frame,
        value,
        value_format: str = "{value}",
        title: Optional[str] = None,
        icon: Optional[str] = None,
        color: str = None,
        date = False,
        variable = False
):
    if variable == "planning_date":
        value = len(data_frame[data_frame["optimal_harvest_date"]==datetime.strptime(date, "%Y-%m-%d")])
    elif variable == "next_field":
        value = data_frame[data_frame["optimal_harvest_date"]>=datetime.strptime(date, "%Y-%m-%d")].sort_values("optimal_harvest_date")["Field ID"].unique()[0]
    elif variable == "expected_sugar":
        value = data_frame[data_frame["optimal_harvest_date"]==datetime.strptime(date, "%Y-%m-%d")]["Expected TSH (t/ha)"].mean()
        if np.isnan(value):
            value_format = "-"
            value = "-"
    elif variable == "yield_change":
        sugar_today = data_frame[data_frame["optimal_harvest_date"]==datetime.strptime(date, "%Y-%m-%d")]["Expected TSH (t/ha)"].mean()
        sugar_tomorrow = data_frame[data_frame["optimal_harvest_date"]==(datetime.strptime(date, "%Y-%m-%d")+timedelta(days=1))]["Expected TSH (t/ha)"].mean()
        value = (sugar_tomorrow - sugar_today)/sugar_today*100
        if value < 0:
            color = "red"
        elif value == 0:
            color = "white"
        elif np.isnan(value):
            color = "white"
            value_format = "-"
            value = "-"
        else:
            color = "green"

    header = dbc.CardHeader(
        [
            html.P(icon, className="material-symbols-outlined") if icon else None,
            html.H4(title, className="card-kpi-title"),
        ]
    )
    body = dbc.CardBody(value_format.format(value=value))
    if color:
        return dbc.Card([header, body], class_name=f"card-kpi card-kpi-{color}-bar")
    else:
        return dbc.Card([header, body], class_name="card-kpi")
    

@capture('ag_grid')
def create_table(data_frame, planning_date, style, columnDefs, when_to_harvest):
    # data_frame = data_frame[data_frame["optimal_harvest_date"]==datetime.strptime(planning_date, "%Y-%m-%d")]
    data_frame["When to harvest"] = (data_frame["optimal_harvest_date"] - datetime.strptime(planning_date, "%Y-%m-%d")).dt.days
    data_frame["When to harvest"] = np.where(data_frame["When to harvest"]<0,
                                             "Harvest day already passed",
                                             np.where(data_frame["When to harvest"]==0,
                                                      "Today",
                                                      np.where(data_frame["When to harvest"]<7,
                                                               "This week",
                                                               np.where(data_frame["When to harvest"]<30,
                                                                        "Next weeks",
                                                                        "Next months"))))
    
    data_frame = data_frame.drop(columns=['geometry', 'trafficability_score'])

    if when_to_harvest:
        data_frame = data_frame[data_frame["When to harvest"].isin(when_to_harvest)]

    defaults = {
        "className": "ag-theme-quartz-dark ag-theme-vizro",
        "defaultColDef": {
            "resizable": True,
            "sortable": True,
            "filter": True,
            "filterParams": {
                "buttons": ["apply", "reset"],
                "closeOnApply": True,
            },
            "flex": 1,
            "minWidth": 70,
        },
        "style": style,
    }
    return AgGrid(
        columnDefs=columnDefs, rowData=data_frame.to_dict("records"), **defaults
    )
    
    return table

@capture('graph')
def create_map(data_frame, color_column, custom_data, planning_date, when_to_harvest):
    # Convert CRS
    data_frame = data_frame.to_crs(epsg=4326)

    data_frame["When to harvest"] = (data_frame["optimal_harvest_date"] - datetime.strptime(planning_date, "%Y-%m-%d")).dt.days
    data_frame["When to harvest"] = np.where(data_frame["When to harvest"]<0,
                                             "Harvest day already passed",
                                             np.where(data_frame["When to harvest"]==0,
                                                      "Today",
                                                      np.where(data_frame["When to harvest"]<7,
                                                               "This week",
                                                               np.where(data_frame["When to harvest"]<30,
                                                                        "Next weeks",
                                                                        "Next months"))))
    
    if when_to_harvest:
        data_frame = data_frame[data_frame["When to harvest"].isin(when_to_harvest)]

    map_zoom = 10
    bounds = data_frame.total_bounds
    map_center = dict(lat=(bounds[1] + bounds[3]) / 2,
                      lon=(bounds[0] + bounds[2]) / 2)

    data_frame["Harvest date"] = data_frame["optimal_harvest_date"].dt.strftime("%Y-%m-%d")

    figure = px.choropleth_map(
        data_frame=data_frame,
        geojson=data_frame.geometry,
        locations=data_frame.index,
        color=color_column,
        color_continuous_scale="Greens",
        map_style="satellite",
        center={"lat": -11.98, "lon": -46.21},
        zoom=map_zoom,
        opacity=0.7,
        custom_data=custom_data,
        hover_data=["Field ID", "Harvest date"]
    )

    figure.update_layout(
        margin={"l": 0, "r": 0, "t": 0, "b": 0},
        coloraxis_showscale=False,
    )

    # Update layout
    figure.update_geos(
        fitbounds="locations",
    )

    return figure


# Load parcels data
import os
import getpass
username = getpass.getuser()
# field_gdf = gpd.read_file(f"data/operations_model_results_preliminary.gpkg")
field_gdf = gpd.read_file("sanitized_data.geojson")
field_gdf["optimal_harvest_date"] = pd.to_datetime(field_gdf["optimal_harvest_date"])

# Subset to a few columns
field_gdf = field_gdf[["fieldId", "finca", "2025_sugar_yield", "geometry", "trafficability_score", "optimal_harvest_date"]]
# Rename columns
field_gdf.rename(columns={
    "fieldId": "Field ID",
    "finca": "Farm",
    "2025_sugar_yield": "Expected TSH (t/ha)",
}, inplace=True)

# Add a new column for trafficability (ramdom selection of Good, Fair, Poor)
field_gdf["Trafficability"] = np.where(field_gdf["trafficability_score"]<0.3,
                                       "Poor",
                                       np.where(field_gdf["trafficability_score"]<0.6,
                                                "Fair",
                                                "Good"))

field_gdf = field_gdf[["Farm", "Field ID", "Trafficability", "Expected TSH (t/ha)", "geometry", "optimal_harvest_date", "trafficability_score"]]

# Define table styles
# Define cell styles for the different columns
trafficability_style = {
    "styleConditions": [
        {
            "condition": "params.value === 'Good'",
            "style": {"backgroundColor": "#c6efcd", "color": "#006100"},
        },
        {
            "condition": "params.value === 'Fair'",
            "style": {"backgroundColor": "#fff2cc", "color": "#9c6500"},
        },
        {
            "condition": "params.value === 'Poor'",
            "style": {"backgroundColor": "#ffd9cc", "color": "#9c0006"},
        },
    ]
}

harvest_style = {
    "styleConditions": [
        {
            "condition": "params.value === 'Today'",
            "style": {"backgroundColor": "#c6efcd", "color": "#006100"},
        },
        {
            "condition": "params.value === 'This week'",
            "style": {"backgroundColor": "#fff2cc", "color": "#9c6500"},
        },
        {
            "condition": "params.value === 'Next weeks'",
            "style": {"backgroundColor": "#ffe6cc", "color": "#be5505"},
        },
        {
            "condition": "params.value === 'Next months'",
            "style": {"backgroundColor": "#ffd9cc", "color": "#9c0006"},
        },
    ]
}

yield_style = {
    "styleConditions": [
        {
            "condition": "params.value >= 9.7",
            "style": {"backgroundColor": "#c6efcd", "color": "#006100"},
        },
        {
            "condition": "params.value < 9.3",
            "style": {"backgroundColor": "#ffd9cc", "color": "#9c0006"},
        },
        {
            "condition": "params.value >= 9.3 && params.value < 9.7",
            "style": {"backgroundColor": "#fff2cc", "color": "#9c6500"},
        },
    ]
}

# Define column definitions
columnDefs = [
    {
        "field": "Farm",
        "flex": 3
    },
    {
        "field": "Field ID",
        "flex": 3
    },
    {
        "field": "When to harvest",
        "cellStyle": harvest_style,
        "flex": 5
    },
    {
        "field": "Trafficability",
        "cellStyle": trafficability_style,
        "flex": 4
    },
    {
        "field": "Expected TSH (t/ha)",
        "valueFormatter": {"function": "d3.format('.2f')(params.value)"},
        "cellStyle": yield_style,
        "flex": 6
    },
]


kpi_banner = FlexContainer(
    id='kpi_banner',
    components=[
        vm.Figure(
            id="next_field",
            figure=get_card(
                data_frame=field_gdf,
                value=42,
                value_format="{value}",
                title="Next field to harvest",
                color=None,
                date=PLANNING_DATE,
                variable="next_field"
            ),
        ),
        vm.Figure(
            id="planning_date",
            figure=get_card(
                data_frame=field_gdf,
                value=0,
                value_format="{value}",
                title="Fields to harvest today",
                color=None,
                date=PLANNING_DATE,
                variable="planning_date"
            ),
        ),
        vm.Figure(
            id="soil_trafficability",
            figure=get_card(
                data_frame=field_gdf,
                value=42,
                value_format="Good",
                title="Soil Trafficability",
                color='green',
            ),
        ),
        vm.Figure(
            id="expected_sugar",
            figure=get_card(
                data_frame=field_gdf,
                value=10.6,
                value_format="{value:.2f} t/ha",
                title="Expected sugar yield (t/ha)",
                color=None,
                date=PLANNING_DATE,
                variable="expected_sugar"
            ),
        ),
        vm.Figure(
            id="yield_change",
            figure=get_card(
                data_frame=field_gdf,
                value=-3.5,
                value_format="{value:.2f} %",
                title="Yield change tomorrow (%)",
                color='red',
                date=PLANNING_DATE,
                variable="yield_change"
            ),
        ),
    ],
    classname="kpi-banner",
)


weather_banner = FlexContainer(
    id='weather_banner',
    components=[
        vm.Figure(
            figure=get_card(
                data_frame=field_gdf,
                value=32,
                value_format="{value}°C - Sunny",
                title="Today",
                color=None,
                icon="Sunny",
            ),
        ),
        vm.Figure(
            figure=get_card(
                data_frame=field_gdf,
                value=28,
                value_format="{value}°C - Clouds",
                title="Tomorrow",
                color=None,
                icon="Cloud",
            ),
        ),
        vm.Figure(
            figure=get_card(
                data_frame=field_gdf,
                value=21,
                value_format="{value}°C - Rainy",
                title="Wednesday",
                color=None,
                icon="Rainy",
            ),
        ),
        vm.Figure(
            figure=get_card(
                data_frame=field_gdf,
                value=18,
                value_format="{value}°C - Rainy",
                title="Thursday",
                color=None,
                icon="Rainy",
            ),
        ),
        vm.Figure(
            figure=get_card(
                data_frame=field_gdf,
                value=30,
                value_format="{value}°C - Sunny",
                title="Friday",
                color=None,
                icon="Sunny",
            ),
        ),
    ],
    classname="kpi-banner",
)


@capture("action")
def update_page_title(click_action, page_title, append_text=""):
    if click_action and append_text not in page_title:
        page_title = page_title + [append_text]
    return page_title


page = vm.Page(
    title="Welcome to your Weather Forecast & Harvest Planning Dashboard",
    layout=vm.Grid(grid=[
        [0, 0],
        [1, 1],
        [2, 3],
        [2, 3],
        [2, 3],
        [2, 3],
        [2, 3],
    ]),
    # Define components with initial figures (using default/initial data)
    components=[
        weather_banner,
        kpi_banner,
        vm.Graph(
            id="field_map",
            figure=create_map(
                data_frame=field_gdf,
                color_column="Expected TSH (t/ha)",
                custom_data=['Field ID'],
                planning_date=PLANNING_DATE,
                when_to_harvest=None
            ),
            actions=[
                vm.Action(
                    function=filter_interaction(
                        targets=["field_map", "field_table"],
                    ),
                ),
                vm.Action(
                    function=update_page_title(append_text=" - Refresh the page to reset filter"),
                    inputs=["field_map.clickData", "page-title.children"],
                    outputs=["page-title.children"],
                ),
            ],
        ),
        vm.AgGrid(
            id="field_table",
            figure=create_table(
                data_frame=field_gdf,
                planning_date=PLANNING_DATE,
                style={"font-size": "18px"},
                columnDefs=columnDefs,
                when_to_harvest=None
            ),
        ),
    ],
    controls=[
        vm.Filter(column="Farm"),
        vm.Filter(column="Field ID"),
        vm.Filter(
            column="Expected TSH (t/ha)",
            selector=vm.RangeSlider(
                min=0,
                max=15,
                step=0.5,
                marks={i: str(i) for i in range(0, 16, 5)},
                value=[5, 15],
            ),
        ),
        vm.Filter(
            column="Trafficability",
            selector=vm.Checklist(
                options=['Good', 'Fair', 'Poor'],
                value=['Good', 'Fair', 'Poor'],
            ),
        ),
        vm.Parameter(
            targets=["planning_date.date", 
                     "field_table.planning_date", 
                     "next_field.date", 
                     "expected_sugar.date", 
                     "yield_change.date", 
                     "field_map.planning_date"],
            selector=vm.DatePicker(
                title="Planning date",
                min="2025-04-07",
                max="2025-06-16",
                range=False,
                value=PLANNING_DATE
                ),
                ),
        vm.Parameter(
            targets=["field_table.when_to_harvest", "field_map.when_to_harvest"],
            selector=vm.Checklist(
                title="When to harvest (from the planning date)",
                options=["Today", "This week", "Next weeks", "Next months"],
                value=["Today", "This week", "Next weeks", "Next months"]
            )
        )
    ],
)

dashboard = vm.Dashboard(
    title="Agriculture Analytics & Insights Control Tower",
    pages=[page],
)

app = Vizro().build(dashboard)

# Add footer in the bottom right corner
app.dash.layout.children.append(
    dbc.NavLink(
        ["From McKinsey - Powered by ACRE"],
        href="https://www.mckinsey.com/industries/agriculture/how-we-help-clients/acre",
        target="_blank",
        className="anchor-container",
    )
)

server = app.dash.server

if __name__ == "__main__":
    app.run()