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

# In[ ]:


# processes historical precipitation data from CHIRPS


# In[1]:


import xarray as xr
import pandas as pd
import geopandas as gpd
import numpy as np
import datetime as dt
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import glob
import os


# In[2]:


# get the centroid of the catchment in order to determine the appropriate data source
centroid = gpd.read_file('AOI/luanda_catchment_level4.shp').to_crs(epsg = 4326).centroid


# In[3]:


# select data source

# use CHIRPS for 50°S - 50°N
# use WorldClim for others
if abs(centroid.y[0]) <= 50:
    data_folder = Path('D:/World Bank/CRP/data/CHIRPS')  # change file path as needed
    data_file = 'chirps-v2.0.monthly.nc'
else:
    data_folder = Path('D:/World Bank/CRP/data/WorldClim')  # change file path as needed
    data_file = ''
    # WorldClim is not yet downloaded. TODO


# In[4]:


output_folder = Path('output')


# In[5]:


# open data source
nc = xr.open_dataset(data_folder / data_file)


# In[6]:


# select data based on coordinate and time period (last 30 years)
df = nc.sel(time = slice(f'{str(dt.date.today().year-30)}-01-01', f'{str(dt.date.today().year-1)}-12-01')).sel(longitude = centroid.x[0], latitude = centroid.y[0], method = 'nearest')


# In[ ]:


# write data to csv
df = df.to_dataframe().drop(['latitude', 'longitude'], axis = 1)
df.to_csv(output_folder / 'precip.csv')


# In[71]:


# plot

# Create the line chart using Seaborn with connected lines
plt.figure(figsize=(10, 6))

sns.lineplot(data=df, x='time', y='precip', linewidth=2)

# Add labels and title
plt.xlabel('')
plt.ylabel('')
plt.title('Monthly precipitation 1993-2022 (mm)')

# plt.xticks(rotation=45)
plt.grid(True)
plt.show()