import ee
import geemap
import solara
import ipywidgets as widgets

ee.Authenticate()

# Initialize the library.
ee.Initialize(project='ovcrge-ssec-burn-scar-map-c116', opt_url='https://earthengine-highvolume.googleapis.com')


fireList = ["North Complex", "Dixie", "Cameron Peak", "August Complex"]
selected_fire = solara.reactive(fireList[0])
selected_days = solara.reactive(30)
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_basemap('OpenStreetMap')
        self.customize_ee_data(selected_fire.value, selected_days.value)
        self.add_selector()
        self.add_intSlider()
        self.add("layer_manager")
        self.remove("draw_control")

        
    def customize_ee_data(self, fire, elapDays):
        elapDayNum = ee.Number(elapDays)
        elapDay_plusOne = elapDayNum.add(ee.Number(1))
        
        north_startDate = ee.Date('2020-08-16')
        dixie_startDate = ee.Date('2021-07-13')
        cam_startDate = ee.Date('2020-08-13')
        aug_startDate = ee.Date('2020-08-15')
        
        north_complex_bb = ee.Geometry.BBox(-121.616097, 39.426723, -120.668526, 40.030845)
        dixie_bb = ee.Geometry.BBox(-121.680467, 39.759303, -120.065477, 40.873387)
        cam_peak_bb = ee.Geometry.BBox(-106.014784, 40.377907, -105.116651, 40.822094)
        aug_complex_bb = ee.Geometry.BBox(-123.668726, 39.337654, -122.355860, 40.498304)

        def calc_nbr(pre_start, pre_stop, post_start, post_stop, bbox):
                PREgoesCMI = ee.ImageCollection('NOAA/GOES/17/MCMIPC').filter(ee.Filter.date(pre_start, pre_stop))\
                                                                      .filter(ee.Filter.calendarRange(17, 21, 'hour')).mean()
                POSTgoesCMI = ee.ImageCollection('NOAA/GOES/17/MCMIPC').filter(ee.Filter.date(post_start, post_stop))\
                                                                      .filter(ee.Filter.calendarRange(17, 21, 'hour')).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
            
        self.clear_specific_layers()
        
        if fire == "North Complex": 
            north_complex = calc_nbr(north_startDate.advance(-7, 'day'), north_startDate, north_startDate.advance(elapDayNum, 'day'), north_startDate.advance(elapDay_plusOne,'day'), north_complex_bb) 
            self.addLayer(north_complex, dNBRvisParams, 'North Complex GOES NBR', True)
            self.centerObject(north_complex_bb, 9)
        elif fire == "Dixie":
            dixie = calc_nbr(dixie_startDate.advance(-7, 'day'), dixie_startDate, dixie_startDate.advance(elapDayNum, 'day'), dixie_startDate.advance(elapDay_plusOne,'day'), dixie_bb)
            self.addLayer(dixie, dNBRvisParams, 'Dixie Complex GOES NBR', True)
            self.centerObject(dixie_bb, 9)
        elif fire == "Cameron Peak":
            cam_peak = calc_nbr(cam_startDate.advance(-7, 'day'), cam_startDate, cam_startDate.advance(elapDayNum, 'day'), cam_startDate.advance(elapDay_plusOne,'day'),   cam_peak_bb)
            self.addLayer(cam_peak, dNBRvisParams, 'Cameron Peak GOES NBR', True)
            self.centerObject(cam_peak_bb, 9)
        elif fire == "August Complex":
            aug_complex = calc_nbr(aug_startDate.advance(-7, 'day'), aug_startDate, aug_startDate.advance(elapDayNum, 'day'), aug_startDate.advance(elapDay_plusOne,'day'),   aug_complex_bb)
            self.addLayer(aug_complex, dNBRvisParams, 'August Complex GOES NBR', True)
            self.centerObject(aug_complex_bb, 9)
    
    def clear_specific_layers(self):
        layers_to_keep = ['OpenStreetMap']
        layers = list(self.layers)
        for layer in layers:
            if layer.name not in layers_to_keep:
                self.remove_layer(layer)  
               
    def add_selector(self):
        selector = widgets.Dropdown(options=fireList, value="North Complex", description='Wildfire Case Study:') 

        def on_selector_change(change):
            if change['name'] == 'value': 
                selected_fire.value = change['new']
                self.customize_ee_data(selected_fire.value, selected_days.value)

        selector.observe(on_selector_change, names='value')
        self.add_widget(selector, position="topleft")

    def add_intSlider(self):
        slider = widgets.IntSlider(value=selected_days.value,min=1,max=30,step=1,description='Elapsed days:',)
        
        def on_slider_change(change):
            if change['name'] == 'value': 
                selected_days.value = change['new']
                self.customize_ee_data(selected_fire.value, selected_days.value)
            
        slider.observe(on_slider_change, names='value') 
        self.add_widget(slider, position="topleft")
        

@solara.component
def Page():
    
    with solara.Column(align="center"):
        markdown = """
        ## Burn Map Generator Case Studies """
        solara.Markdown(markdown)
    
    # Isolation is required to prevent the map from overlapping navigation (when screen width < 960px)
    with solara.Column(style={"isolation": "isolate"}):
        map_widget = Map.element(
            center=[39, -120.5],
            zoom=8,
            height="600px",
            toolbar_ctrl=False
        )