//// Authors: Stephen Chignell and Yeonuk Kim (RH Flux section).
//// Adapted from code written by: 
//// Poortinga et al. (2022). Chapter A2.5: Water Balance and Drought in Cloud-Based Remote Sensing with Google Earth Engine
//// https://docs.google.com/document/d/12cwzbNXtBQnm5switfL1vlv1stnp5D5gXcQeZMP_iJI/edit
 

// Import the study boundary.
// Note: to bound the analysis to a hand-drawn polygon,
//       comment out Option 1 and uncomment Option 2 below. Then draw study boundary within GEE map.

// Option 1: Select GEE asset
var study_boundary = ee.FeatureCollection('projects/ee-bale-wb/assets/Bale_ecoregionArea_WGS84');
  
// Option 2: Select hand-drawn polygon in map  
// var study_boundary = harenna;


// Import the Bale Mountains National Park (BMNP) boundary for reference
var BMNP = ee.FeatureCollection(
    'projects/ee-bale-wb/assets/BMNP');



// Center the map. Adjust scale level as necessary.
Map.centerObject(study_boundary, 8);

// Add the study boundary to the map.
Map.addLayer(study_boundary, {}, 'Study Boundary');


// Set start and end years.
var startYear = 2001;
var endYear = 2020;
 
// Create two date objects for start and end years.
var startDate = ee.Date.fromYMD(startYear, 1, 1);
var endDate = ee.Date.fromYMD(endYear + 1, 1, 1);
 
// Make a list with years.
var years = ee.List.sequence(startYear, endYear);

// Make a list with months. Edit and comment out the lines below 
// to extract the desired time period.
var months = ee.List.sequence(1, 12); //full year
//var months = ee.List([4,5]); //wet season (short)
//var months = ee.List([6,7]); //dry season (short)
//var months = ee.List([8,9,10]); //wet season (long)
//var months = ee.List([12,1,2]); //dry season (long)


// Import the CHIRPS precipitation dataset.
var CHIRPS = ee.ImageCollection('UCSB-CHG/CHIRPS/PENTAD');

// Filter for the relevant time period.
CHIRPS = CHIRPS.filterDate(startDate, endDate);


// We apply a nested loop where we first map over 
// the relevant years and then map over the relevant 
// months. The function returns an image with the total (sum)
// rainfall for each month. A flatten is applied to convert a
// feature collection of features into a single feature collection.
var monthlyPrecip = ee.ImageCollection.fromImages(
    years.map(function(y) {
        return months.map(function(m) {
            var w = CHIRPS.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum();
            return w.set('year', y)
                .set('month', m)
                .set('system:time_start', ee.Date
                    .fromYMD(y, m, 1));

        });
    }).flatten()
);

// Add the layer with monthly mean. Note that we clip by the study boundary.
// Note: the min/max values should be adjusted based on the range of the data in the study boundary.
var precipVis = {
    min: 0,
    max: 250,
    palette: 'red, orange, yellow, blue, darkblue'
};

Map.addLayer(monthlyPrecip.mean().clip(study_boundary),
    precipVis,
    'Mean monthly P');

// Set the title and axis labels for the chart.
var title = {
    title: 'Monthly precipitation',
    hAxis: {
        title: 'Date'    },
    vAxis: {
        title: 'Precipitation (mm)'
    },
};

// Plot the chart using the study boundary.
var chartMonthly = ui.Chart.image.seriesByRegion({
        imageCollection: monthlyPrecip,
        regions: study_boundary.geometry(),
        reducer: ee.Reducer.mean(),
        band: 'precipitation',
        scale: 5000,
        xProperty: 'system:time_start'
    }).setSeriesNames(['P'])
    .setOptions(title)
    .setChartType('ColumnChart');

// Print the chart.
print(chartMonthly);

// Import the PML V2 dataset.
var PML = ee.ImageCollection('CAS/IGSNRR/PML/V2_v017');
// Select the 3 different components of ET and sum them.
// Note, we exclude ET_water because it is known to be error-prone.
// Then multiply by 8 because the PML values are calculated over an 8-day period, but the units are per day.
var PML_ET = function(image){
  var image = image.addBands(
    image.expression(
    '(Ei+Es+Ec)*8', {
      'Ei': image.select('Ei'), // Interception from vegetation canopy
      'Es': image.select('Es'), // Soil evaporation
      'Ec': image.select('Ec') // Vegetation transpiration 
    }).rename('ET')
    );
    return image
}

var PML = PML.map(PML_ET).select('ET');

// Filter for the relevant time period.
PML = PML.filterDate(startDate, endDate);


// Import MOD16 dataset.
var mod16 = ee.ImageCollection('MODIS/006/MOD16A2').select('ET');

// Filter for the relevant time period.
mod16 = mod16.filterDate(startDate, endDate);

// We apply a nested loop where we first map over 
// the relevant years and then map over the relevant 
// months. The function returns an image with the total (sum)
// evapotranspiration for each month. A flatten is applied to convert a
// collection of collections into a single collection. 
// We multiply by 0.5 to take the mean (i.e. ensemble) the two products.
var monthlyEvap = ee.ImageCollection.fromImages(
    years.map(function(y) {
        return months.map(function(m) {

            var ET_PML = PML.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.5);

            var ET_mod16 = mod16.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.1)
                .multiply(0.5);

            var ET = ET_PML.add(ET_mod16).rename('ET_ensemble');

            return ET.set('year', y)
                .set('month', m)
                .set('system:time_start', ee.Date
                    .fromYMD(y, m, 1));

        });
    }).flatten()
);

// Add the layer with monthly mean. Note that we clip by the study boundary.
// Note: the min/max values should be adjusted based on the range of the data in the study boundary.
var evapVis = {
    min: 0,
    max: 140,
    palette: 'red, orange, yellow, blue, darkblue, purple'
};

Map.addLayer(monthlyEvap.mean().clip(study_boundary),
    evapVis,
    'Mean monthly ET');

// Set the title and axis labels for the chart.
var title = {
    title: 'Monthly evapotranspiration',
    hAxis: {
        title: 'Date'
    },
    vAxis: {
        title: 'Evapotranspiration (mm)'
    },
    colors: ['red']
};

// Plot the chart using the study boundary.
var chartMonthly = ui.Chart.image.seriesByRegion({
        imageCollection: monthlyEvap,
        regions: study_boundary.geometry(),
        reducer: ee.Reducer.mean(),
        band: 'ET_ensemble',
        scale: 500,
        xProperty: 'system:time_start'
    }).setSeriesNames(['ET_ensemble'])
    .setOptions(title)
    .setChartType('ColumnChart');

// Print the chart.
print(chartMonthly);


// We apply a nested loop where we first map over 
// the relevant years and then map over the relevant 
// months. The function returns an image with P - ET
// for each month. A flatten is applied to convert an
// collection of collections into a single collection.
// We multiply by 0.5 to take the mean (i.e. ensemble) the two products.
var waterBalance = ee.ImageCollection.fromImages(
    years.map(function(y) {
        return months.map(function(m) {

            var P = CHIRPS.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum();

            var ET_PML = PML.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.5);

            var ET_mod16 = mod16.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.1)
                .multiply(0.5);

            var ET = ET_PML.add(ET_mod16).rename('ET_ensemble');

            var wb = P.subtract(ET).rename('wb');

            return wb.set('year', y)
                .set('month', m)
                .set('system:time_start', ee.Date
                    .fromYMD(y, m, 1));

        });
    }).flatten()
);

// Add layer with monthly mean. Note that we clip by the study boundary.
// Note: the min/max values should be adjusted based on the range of the data in the study boundary.
var balanceVis = {
    min: -50,
    max: 200,
    palette: 'red, orange, yellow, blue, darkblue, purple'
};

Map.addLayer(waterBalance.mean().clip(study_boundary),
    balanceVis,
    'Mean monthly WB');

// Set the title and axis labels for the chart.
var title = {
    title: 'Monthly water balance',
    hAxis: {
        title: 'Date'
    },
    vAxis: {
        title: 'P-ET (mm)'
    },
    colors: ['green']
};

// Plot the chart using the study boundary.
var chartMonthly = ui.Chart.image.seriesByRegion({
        imageCollection: waterBalance,
        regions: study_boundary.geometry(),
        reducer: ee.Reducer.mean(),
        band: 'wb',
        scale: 500,
        xProperty: 'system:time_start'
    }).setSeriesNames(['WB'])
    .setOptions(title)
    .setChartType('ColumnChart');

// Print the chart.
print(chartMonthly);



// We apply a nested loop where we first map over 
// the relevant years and then map over the relevant 
// months. The function returns an image with bands for 
// water balance (wb), rainfall (P), and evapotranspiration (ET) 
// for each month. A flatten is applied to convert a collection 
// of collections into a single collection.
var ic = ee.ImageCollection.fromImages(
    years.map(function(y) {
        return months.map(function(m) {            // Calculate rainfall.
            var P = CHIRPS.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum();

            // Calculate evapotranspiration.
            var ET = mod16.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.1);

            // Calculate monthly water balance.
            var wb = P.subtract(ET).rename('wb');

            // Return an image with all images as bands.
            return ee.Image.cat([wb, P, ET])
                .set('year', y)
                .set('month', m)
                                .set('system:time_start', ee.Date
                    .fromYMD(y, m, 1));

        });
    }).flatten()
);


// Define the water balance chart and print it to the console.
var chartWB =
    ui.Chart.image.series({
        imageCollection: ic.select(['wb', 'precipitation', 'ET']),
        region: study_boundary,
        reducer: ee.Reducer.mean(),
        scale: 5000,
                xProperty: 'system:time_start'
    })
    .setSeriesNames(['ET', 'P', 'WB'])
    .setOptions({
        title: 'Water Balance',
        hAxis: {
            title: 'Date',
            titleTextStyle: {
                italic: false,
                bold: true
            }
        },
        vAxis: {
            title: 'Water (mm)',
            titleTextStyle: {
                italic: false,
                bold: true
            }
        },
        lineWidth: 1,
        colors: ['green', 'blue', 'red'],
        curveType: 'function'
    });

// Print the water balance chart.
print(chartWB);




/// Written by Yeonuk Kim
// Net radiation and SFE ET driven by ERA5-Land
// Import ERA5-Land dataset
var ERA5L = ee.ImageCollection('ECMWF/ERA5_LAND/MONTHLY_AGGR').filterDate(startDate, endDate);

// calculate net radiation and SFE evapotranspiration (add band to the imagecollection)
var ERA5_function = function(image) {
  // Net radiation (unit J/m2/month)
   var image = image.addBands(
    image.expression(
    '-(LE + H)', {
      'LE': image.select('surface_latent_heat_flux_sum'),
      'H': image.select('surface_sensible_heat_flux_sum')
    }).rename('Rnet_G'));
  // vapour pressure (ea)
  var image = image.addBands(
  image.expression(
    '611.2 * 2.718282 ** ((17.62 * (Tdew-273.15))/(243.12 + (Tdew-273.15)))', {
      'Tdew': image.select('dewpoint_temperature_2m')
    }).rename('ea'));
  // specific humidity (q)  
  var image = image.addBands(
      image.expression(
    '0.622 * ea/(PA - (1 - 0.622) * ea)', {
      'ea': image.select('ea'),
      'PA': image.select('surface_pressure')
    }).rename('q'));
  // reversed Bowen ratio under SFE state  
  var image = image.addBands(
      image.expression(
    'q * (((2.501 - 0.00237 * (Tair-273.15)) * 1000000)**2)/(1005 * 461.5 * (Tair ** 2))', {
      'q': image.select('q'),
      'Tair': image.select('temperature_2m')
    }).rename('SFE_rev_B'));  
  // SFE evapotranspiration (mm/month)
  var image = image.addBands(
      image.expression(
    'SFE_rev_B/(SFE_rev_B+1)*Rnet_G/((2.501 - 0.00237 * (Tair-273.15)) * 1000000)', {
      'SFE_rev_B': image.select('SFE_rev_B'),
      'Rnet_G': image.select('Rnet_G'),
      'Tair': image.select('temperature_2m')
    }).rename('SFE_ET'));

  return image
};

var ERA5L = ERA5L.map(ERA5_function);
var ERA5L_Rnet = ERA5L.select('Rnet_G'); // unit: J/m2/month
var ERA5L_SFE_ET = ERA5L.select('SFE_ET'); // unit: mm/month

Map.addLayer(ERA5L_SFE_ET.mean().clip(study_boundary),
    evapVis,
    'Mean monthly SFE ET');

// Set the title and axis labels for the chart.
var title = {
    title: 'Monthly evapotranspiration',
    hAxis: {
        title: 'Date'
    },
    vAxis: {
        title: 'SFE ET (mm)'
    },
    colors: ['red']
};

// Plot the chart using the study_boundary boundary.
var chartMonthly = ui.Chart.image.seriesByRegion({
        imageCollection: ERA5L_SFE_ET,
        regions: study_boundary.geometry(),
        reducer: ee.Reducer.mean(),
        band: 'SFE_ET',
        scale: 500,
        xProperty: 'system:time_start'
    }).setSeriesNames(['SFE_ET'])
    .setOptions(title)
    .setChartType('ColumnChart');

// Print the chart.
print(chartMonthly);



// We apply a nested loop where we first map over 
// the relevant years and then map over the relevant 
// months. The function returns an image with ET - SFE ET = relative humidity flux
// for each month. A flatten is applied to convert an
// collection of collections into a single collection.
var Flux_rh = ee.ImageCollection.fromImages(
    years.map(function(y) {
        return months.map(function(m) {

            var ET_PML = PML.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.5);

            var ET_mod16 = mod16.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum()
                .multiply(0.1)
                .multiply(0.5);

            var ET = ET_PML.add(ET_mod16).rename('ET_ensemble');
                
            var SFE_ET = ERA5L_SFE_ET.filter(ee.Filter
                    .calendarRange(y, y, 'year'))
                .filter(ee.Filter.calendarRange(m, m,
                    'month'))
                .sum();

            var Flux_rh = ET.subtract(SFE_ET).rename('Flux_rh');

            return ee.Image.cat([Flux_rh, ET, SFE_ET])
                .set('year', y)
                .set('month', m)
                                .set('system:time_start', ee.Date
                    .fromYMD(y, m, 1));

        });
    }).flatten()
);

// Add layer with monthly mean. Note that we clip by the study_boundary.
// Note: the min/max values should be adjusted based on the range of the data in the study boundary.
var balanceVis = {
    min: -100,
    max: 100,
    palette: 'red, orange, yellow, blue, darkblue, purple'
};

Map.addLayer(Flux_rh.select('Flux_rh').mean().clip(study_boundary),
    balanceVis,
    'Mean monthly ET - SFE_ET');

// Set the title and axis labels for the chart.
var title = {
    title: 'Monthly ET - SFE_ET',
    hAxis: {
        title: 'Date'
    },
    vAxis: {
        title: 'ET - SFE_ET (mm)'
    },
    colors: ['green']
};

// Plot the chart using the study_boundary boundary.
var chartMonthly = ui.Chart.image.seriesByRegion({
        imageCollection: Flux_rh,
        regions: study_boundary.geometry(),
        reducer: ee.Reducer.mean(),
        band: 'Flux_rh',
        scale: 5000,
        xProperty: 'system:time_start'
    }).setSeriesNames(['Flux_rh'])
    .setOptions(title)
    .setChartType('ColumnChart');

// Print the chart.
print(chartMonthly);


// Define ET and SFE_ET chart and print it to the console.
var chartET_SFE_ET =
    ui.Chart.image.series({
        imageCollection: Flux_rh.select([ 'ET_ensemble','Flux_rh', 'SFE_ET']),
        region: study_boundary,
        reducer: ee.Reducer.mean(),
        scale: 5000,
                xProperty: 'system:time_start'
    })
    .setSeriesNames([ 'ET_ensemble', 'Flux_rh','SFE_ET'])
    .setOptions({
        title: 'ET components',
        hAxis: {
            title: 'Date',
            titleTextStyle: {
                italic: false,
                bold: true
            }
        },
        vAxis: {
            title: 'ET (mm)',
            titleTextStyle: {
                italic: false,
                bold: true
            }
        },
        lineWidth: 1,
        colors: ['green', 'blue', 'red'],
        curveType: 'function'
    });

// Print the water balance chart.
print(chartET_SFE_ET);


// Add the BMNP boundary to the map for reference.
// Set styling
var styling = {color: 'black', fillColor: '00000000'};
Map.addLayer(BMNP.style(styling), {}, 'BMNP');



/// Define color bar
var palette = ["red", "orange", "yellow", "blue", "darkblue", "purple"]
function createColorBar(titleText, palette, min, max) {
  // Legend Title
  var title = ui.Label({
    value: titleText,
    style: {fontWeight: 'bold', textAlign: 'center', stretch: 'horizontal'}});
  // Colorbar
  var legend = ui.Thumbnail({
    image: ee.Image.pixelLonLat().select(0),
    params: {
      bbox: [0, 0, 1, 0.1],
      dimensions: '200x20',
      format: 'png',
      min: 0, max: 1,
      palette: palette},
    style: {stretch: 'horizontal', margin: '8px 8px', maxHeight: '40px'},
  });
  // Legend Labels
  var labels = ui.Panel({
    widgets: [
      ui.Label(min, {margin: '4px 10px',textAlign: 'left', stretch: 'horizontal'}),
      ui.Label((min+max)/2, {margin: '4px 20px', textAlign: 'center', stretch: 'horizontal'}),
      ui.Label(max, {margin: '4px 10px',textAlign: 'right', stretch: 'horizontal'})],
    layout: ui.Panel.Layout.flow('horizontal')});
  // Create a panel with all 3 widgets
  var legendPanel = ui.Panel({
    widgets: [title, legend, labels],
    style: {position: 'bottom-center', padding: '8px 15px'}
  })
  return legendPanel
}
// Call the function to create a colorbar legend. 
// Edit and comment all but one of the lines below to
// create a legend with the desired variable and value range. 
// Change default values depending on range of each variable.

//var colorBar = createColorBar('P-ET (mm)', palette, -50, 50) 
//var colorBar = createColorBar('P-ET (mm)', palette, -40, 40)
//var colorBar = createColorBar('RH flux (mm)', palette, -35, 35)
//var colorBar = createColorBar('Precipitation (mm)', palette, 50, 120)
var colorBar = createColorBar('Evapotranspiration (mm)', palette, 50, 120)

Map.add(colorBar);