#!/usr/bin/env python
# coding: utf-8

# In[2]:


import xarray as xr
import pandas as pd
import os
import geopandas as gpd
from shapely.geometry import mapping
import numpy as np


# In[46]:


get_ipython().system('pip install rioxarray')


# In[51]:


# set paths
import regionmask
era_data_path="C:/Users/dboateng/Desktop/Datasets/ERA5/monthly_1950_2021/"
path_shapefile="C:/Users/dboateng/Desktop/Datasets/Station/Ghana/Ghana_ShapeFile/gh_wgs16dregions.shp"
afr_shape = "C:/Users/dboateng/Desktop/Datasets/Station/Ghana/Africa_shapefile/afr_g2014_2013_0.shp"

from1979to2012 = pd.date_range(start="1979-01-01", end="2012-12-31", freq="MS")
# read data

ghana_shape = gpd.read_file(path_shapefile)


# In[56]:


tp_monthly= xr.open_dataset(os.path.join(era_data_path, "tp_monthly.nc"))
tp_monthly = tp_monthly["tp"].sel(time=from1979to2012).mean(dim="time")
tp_monthly = tp_monthly.assign_coords({"longitude": (((tp_monthly.longitude + 180) % 360) - 180)})


# In[57]:


import rioxarray
tp_monthly.rio.set_spatial_dims(x_dim="longitude", y_dim="latitude", inplace=True)
tp_monthly.rio.write_crs("epsg:4326", inplace=True)


# In[58]:


clipped = tp_monthly.rio.clip(ghana_shape.geometry.apply(mapping), ghana_shape.crs, drop=True)


# In[60]:


clipped


# In[62]:


clipped.sortby("longitude").plot()


# In[5]:


africa_shape = gpd.read_file(afr_shape)


# In[6]:


africa_shape


# In[22]:


regions =ghana_shape["REGION"]
regions


# In[15]:


my_list


# In[23]:


countries_mask_poly = regionmask.Regions(outlines=ghana_shape.geometry)


# In[24]:


countries_mask_poly


# In[66]:


tp_monthly= xr.open_dataset(os.path.join(era_data_path, "tp_monthly.nc"))
tp_monthly = tp_monthly["tp"].sel(time=from1979to2012).mean(dim="time")


# In[67]:


mask = countries_mask_poly.mask(tp_monthly, lat_name="latitude", lon_name="longitude", wrap_lon=True)


# In[68]:


mask


# In[69]:


mask.plot()


# In[76]:


droped = mask.dropna(dim=("latitude"), how="any")


# In[77]:


droped


# In[41]:


countries = list(africa_shape["ISO3"])
my_list = set(list(countries))
indexes = [countries.index(x) for x in my_list]

countries_mask_poly = regionmask.Regions(name="ADM0_NAME", numbers=indexes, 
                                           names=africa_shape.ADM0_NAME[indexes],
                                           outlines=list(africa_shape.geometry.values[i] for i in range(0, africa_shape.shape[0])),
                                          abbrevs=africa_shape["ADM0_CODE"][indexes])


# In[42]:


countries_mask_poly


# In[45]:


countries_mask_poly[20]