import ee
import geemap
import solara

fireList = ["North Complex", "Dixie", "Cameron Peak", "August Complex"]
selected_fire = solara.reactive([fireList[0]])
dNBRvisParams = {'min': 0.0,'max': 0.8, 'palette': ['green', 'yellow','orange','red']}

class Map(geemap.Map):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.add_ee_data()
        self.add("layer_manager")
        self.add("inspector")

    def add_ee_data(self):
        # Add North Complex image
        north_complex_bb = ee.Geometry.BBox(-121.616097, 39.426723, -120.668526, 40.030845)
        self.north_complex = self.calc_nbr('2020-08-15', '2020-08-16', '2020-09-15', '2020-09-16', north_complex_bb)
        
        # Add Dixie image
        dixie_bb = ee.Geometry.BBox(-121.680467, 39.759303, -120.065477, 40.873387)
        self.dixie = self.calc_nbr('2021-07-12', '2021-07-13', '2021-09-15', '2021-09-16', dixie_bb)
        
        # Add Cameron Peak image
        cam_peak_bb = ee.Geometry.BBox(-106.014784, 40.377907, -105.116651, 40.822094)
        self.cam_peak = self.calc_nbr('2020-08-12', '2020-08-13', '2020-09-12', '2020-09-13', cam_peak_bb)
        
        # Add August Complex image
        aug_complex_bb = ee.Geometry.BBox(-123.668726, 39.337654, -122.355860, 40.498304)
        self.aug_complex = self.calc_nbr('2020-08-15', '2020-08-16', '2020-11-15', '2020-11-16', aug_complex_bb)

        # Visualization parameters
        self.dNBRvisParams = dNBRvisParams

    def calc_nbr(self, pre_start, pre_stop, post_start, post_stop, bbox):
        PREgoesCMI = ee.ImageCollection('NOAA/GOES/17/MCMIPF').filter(ee.Filter.date(pre_start, pre_stop)).mean()
        POSTgoesCMI = ee.ImageCollection('NOAA/GOES/17/MCMIPF').filter(ee.Filter.date(post_start, post_stop)).mean()
        preNBR = PREgoesCMI.select(['CMI_C03','CMI_C06']).normalizedDifference(['CMI_C03', 'CMI_C06']).toFloat().rename('NBR')
        postNBR = POSTgoesCMI.select(['CMI_C03','CMI_C06']).normalizedDifference(['CMI_C03', 'CMI_C06']).toFloat().rename('NBR')
        dNBR = preNBR.subtract(postNBR).select('NBR')
        dNBRclipped = dNBR.clip(bbox)
        return dNBRclipped

@solara.component
def Page():
    map_instance = Map(element_id="map")

    def update_map(selected_fire):
        map_instance.clear_layers()
        if selected_fire == "North Complex":
            map_instance.addLayer(map_instance.north_complex, map_instance.dNBRvisParams, 'North Complex GOES NBR', True)
        elif selected_fire == "Dixie":
            map_instance.addLayer(map_instance.dixie, map_instance.dNBRvisParams, 'Dixie Complex GOES NBR', True)
        elif selected_fire == "Cameron Peak":
            map_instance.addLayer(map_instance.cam_peak, map_instance.dNBRvisParams, 'Cameron Peak GOES NBR', True)
        elif selected_fire == "August Complex":
            map_instance.addLayer(map_instance.aug_complex, map_instance.dNBRvisParams, 'August Complex GOES NBR', True)

    with solara.Column(style={"min-width": "500px"}):
        map_widget = Map.element(
            center=[40, -100],
            zoom=4,
            height="600px",
        )

    solara.Select(label="Wildfire Case Study", value=selected_fire, values=fireList, on_value=update_map)
    solara.Markdown(f"**Selected**: {selected_fire.value}")