import csv
import pandas as pd
import os
import glob
import numpy as np

os.getcwd()
os.chdir(r'C:\Users\anson\Dropbox\Dissertation\Fracfocus\FracFocusCSV\STORET\1round\prism')


###merge PRISM files by year
path = r'C:\Users\anson\Dropbox\Dissertation\Fracfocus\FracFocusCSV\STORET\1round\prism\2012'  
all_files = glob.glob(os.path.join(path, "*.dta"))

df = pd.concat([pd.read_stata(fp).assign(fips=os.path.basename(fp).split('.')[0]) for fp in all_files])

#generate monthly level weather data for each grid
df.dateNum = pd.to_datetime(df.dateNum)
df['month'] = df['dateNum'].dt.month
df['year'] = df['dateNum'].dt.year
df['tMin_mean'] = df.groupby(["gridNumber","month"])["tMin"].transform("mean")
df['tMax_mean'] = df.groupby(["gridNumber","month"])["tMax"].transform("mean")
df['prec_sum'] = df.groupby(["gridNumber","month"])["prec"].transform("sum")

#generate geographic coordinate
df['gridNumber_1'] = df['gridNumber'] - 1
df['longitude'] = -125+(df['gridNumber_1'].mod(1405)/24)
df['latitude']  = 49.9375+1/48 - (df['gridNumber']/1405).apply(np.ceil)/24

df_keep = df[['gridNumber','fips','year','month','tMin_mean','tMax_mean','prec_sum','latitude','longitude']]
df_keep.drop_duplicates(inplace=True)
df_keep.to_csv('2012monthlevel.csv', index=False)

###repeat the above process for all years

#import the saved data
df_2012 = pd.read_csv('2012monthlevel.csv')
df_2013 = pd.read_csv('2013monthlevel.csv')
df_2014 = pd.read_csv('2014monthlevel.csv')
df_2015 = pd.read_csv('2015monthlevel.csv')
df_2016 = pd.read_csv('2016monthlevel.csv')
df_2017 = pd.read_csv('2017monthlevel.csv')
df_2018 = pd.read_csv('2018monthlevel.csv')
df_2019 = pd.read_csv('2019monthlevel.csv')
df_2020= pd.read_csv('2020monthlevel.csv')
df_2021= pd.read_csv('2021monthlevel.csv')
df_2020['year']=2020
df_2021['year']=2021

#extract geographic coordinate for each grid
df_drop = df_2012[['gridNumber','fips','latitude','longitude']]
df_drop.drop_duplicates(inplace=True)
df_drop.to_csv('grid.csv', index=False)


###merge PRISM with well data (executed in QGIS)
#import the grid data to QGIS
#import the well data to QGIS
#join the two layers by nearest attribute (set maximum neighbor as 1)
#export the joined layer as csv
df_gridlink = pd.read_csv('grid_joined.csv')
df_prism_raw = pd.concat([df_2012, df_2013, df_2014, df_2015, df_2016, df_2017,df_2018,df_2019,df_2020,df_2021], axis=0)
df_merged = df_gridlink.merge(df_prism_raw, on=['gridNumber','year','month'], how='left')
df_merged_keep = df_merged[['UploadKey','tMin_mean','tMax_mean','prec_sum']]

#generate weather variables
df_merged_keep['t_mean'] = (df_merged_keep['tMin_mean']+df_merged_keep['tMax_mean'])/2
df_merged_keep['wet_month'] = np.where(df_merged_keep['prec_sum']>100, 1, 0)
df_merged_keep['snow_month'] = np.where((df_merged_keep['prec_sum']>20)&(df_merged_keep['t_mean']<0), 1, 0)
df_merged_keep['rainsnow'] = np.where((df_merged_keep['wet_month']==1)|(df_merged_keep['snow_month']==1), 1, 0)

#merge with well data
df = pd.read_csv(r'C:\Users\anson\Dropbox\Dissertation\Fracfocus\FracFocusCSV\fracfocus_monitor.csv')
df_final = df.merge(df_merged_keep, on='UploadKey', how='left')
df_final.to_csv(r'C:\Users\anson\Dropbox\Dissertation\Fracfocus\FracFocusCSV\fracfocus_welllevel.csv', index=False)

