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

os.chdir(r'C:\Users\anson\Dropbox\Dissertation\Fracfocus\FracFocusCSV')

df = pd.read_csv("merged_raw_original.csv")


###clean CAS number
#NP refers to invalid chemical entries (labled as 'N/A'/'unknown'/'not assigned' or similar terms)
#TS refers to trade secret chemicals (labled as 'proprietary'/'trade secrets'/'confidential' or similar terms)
#SD refers to system disclosure entries without CAS numbers (labled as 'listed below'/'system disclosure' or similar terms)
df['CASNumber'] = df['CASNumber'].str.strip()
df['CASNumber'] = df['CASNumber'].str.lstrip('0')
df['CASNumber'] = df['CASNumber'].str.strip('*')
df['CASNumber'] = df['CASNumber'].str.strip('#')
df['CASNumber'] = df['CASNumber'].str.strip('(')
df['CASNumber'] = df['CASNumber'].str.strip(')')
df['CASNumber'] = df['CASNumber'].str.strip('"')
df['CASNumber'] = df['CASNumber'].str.strip("'")
df['CASNumber'] = df['CASNumber'].str.strip('?')
df['CASNumber'] = df['CASNumber'].str.replace('.','-')
CAS_ODD_dic = pd.read_csv('uniqueCAS.csv', header=None, index_col=0, squeeze=True).to_dict()
df['CAS_cleaned']= df['CASNumber'].map(CAS_ODD_dic)
df = df.drop(df[df.CAS_cleaned == 'SD'].index)


###construct trade secret variables
df['CAScount'] = df['CAS_cleaned'].groupby(df['UploadKey']).transform('count')
df['CASisTS'] = np.where(df['CAS_cleaned']== 'TS', True, False)
df['TScount'] = df['CASisTS'].groupby(df['UploadKey']).transform('sum')
df['CAScount_uq'] = df.groupby(['UploadKey'])['CAS_cleaned'].transform('nunique')

#TS_per
df['TS_perCASall'] = df['TScount']/df['CAScount']

#CAS_anony
df['TS_perCASuq'] = df['TScount']/(df['CAScount_uq']-1)

#TS_conc
df['HFper_TS'] = df['PercentHFJob']*df['CASisTS']
df_nowater = df[df['CAS_cleaned']!='7732185']
HFperTS_well_nowater = df_nowater.groupby(['UploadKey'])['HFper_TS','PercentHFJob'].sum().reset_index()
HFperTS_well_nowater['HFper_TS_snw'] = HFperTS_well_nowater['HFper_TS']/HFperTS_well_nowater['PercentHFJob']
HFperTS_well_nowater_keep = HFperTS_well_nowater[['UploadKey','HFper_TS_snw']]

#save TS variables
df_TS = df[['UploadKey', 'TS_perCASall', 'TS_perCASuq']]
df_TS.drop_duplicates(inplace=True)
TS_final = df_TS.merge(HFperTS_well_nowater_keep, on="UploadKey", how="left")
TS_final.replace([np.inf, -np.inf], np.nan, inplace=True)
TS_final.dropna(inplace=True)
TS_final.to_csv('TS_final.csv', index=False, encoding='utf-8-sig')


###prepare chemical-concentration data
#construct chemical-concentration-date dataset
df_noNP = df[df.CAS_cleaned!= 'NP']
df_noNP_agg = df_noNP.groupby(['UploadKey','CAS_cleaned'])['PercentHFJob'].sum().reset_index()
df_wellinfo = df[['UploadKey','APINumber', 'OperatorName',  'date', 'year', 'month']]
df_wellinfo.drop_duplicates(inplace=True)
df_dbscan = df_noNP_agg.merge(df_wellinfo, on='UploadKey', how='left')

#merge FracFocus data with WellDatabase using API number
df_WDB = pd.read_csv('header_prod.csv')

df_WDB_select = df_WDB[['API','Play','Latitude','Longitude','FirstSixMonthOil','FirstSixMonthBOE']]
df_dbscan_gis = df_dbscan.merge(df_WDB_select, left_on='APINumber', right_on='API', how='left')
df_dbscan_gis  = df_dbscan_gis[df_dbscan_gis['Play'].notna()]

df_dbscan_gis.to_csv('frac_DBSCAN.csv', index=False, encoding='utf-8-sig')


###generate control variables
df_dbscan_gis_well = df_dbscan_gis[['UploadKey','OperatorName','Play','date','year','month','FirstSixMonthOil','FirstSixMonthBOE']]
df_dbscan_gis_well.drop_duplicates(inplace=True)

#WELLNUM_new
df_dbscan_gis_well_byplay = df_dbscan_gis_well.groupby(['OperatorName','year','Play'])['UploadKey'].count().reset_index()
df_dbscan_gis_well_byplay.rename(columns={'UploadKey':'wellcount_byplay'}, inplace=True)
df_dbscan_gis_well_byplay['log_wellcount_byplay'] = np.log2(df_dbscan_gis_well_byplay.wellcount_byplay+1) 

#note: WELLNUM_new and WELLNUM_all are log-transformed with base 2, while OGPROD is log-transformed with base e.
#These variables were constructed at different times in different files, and this is an 'omission' I found after
#the paper was online. As shown in Table 2 and Figure 3, the main findings remain robust without adding any 
#control variables. The results remain robust regardless of how these controls are transformed (with base 2 or e).

#WELLNUM_all
df_dbscan_gis_well_byplay['cumsum_byplay'] = df_dbscan_gis_well_byplay.groupby(['OperatorName','Play']).wellcount_byplay.cumsum()
df_dbscan_gis_well_byplay['log_cumsum_byplay'] = np.log2(df_dbscan_gis_well_byplay.cumsum_byplay+1) 

#OPEXP
df_dbscan_gis_well_byplay['firstyear_byplay'] = df_dbscan_gis_well_byplay.groupby(['OperatorName','Play'])['year'].transform('first')
df_dbscan_gis_well_byplay['opexp_byplay']=df_dbscan_gis_well_byplay['year']-df_dbscan_gis_well_byplay['firstyear_byplay']+1 

#OGPROD
df_dbscan_gis_well['OGprod'] = (df_dbscan_gis_well['FirstSixMonthOil'] + df_dbscan_gis_well['FirstSixMonthBOE'])/1000
df_dbscan_gis_well['OGprod_byplay'] = df_dbscan_gis_well.groupby(['OperatorName','Play','year','month'])['OGprod'].transform('mean')
df_dbscan_gis_well['OGprod_byplay_y'] = df_dbscan_gis_well.groupby(['OperatorName','Play','year'])['OGprod'].transform('mean')
df_dbscan_gis_well['ln_OGprod_byplay'] = np.log(df_dbscan_gis_well.OGprod_byplay+1) 
df_dbscan_gis_well['ln_OGprod_byplay_y'] = np.log(df_dbscan_gis_well.OGprod_byplay_y+1) 
df_dbscan_gis_well_og = df_dbscan_gis_well[['OperatorName','Play','year','month','ln_OGprod_byplay','ln_OGprod_byplay_y']]
df_dbscan_gis_well_og.drop_duplicates(inplace=True)

#merge all controls
df_dbscan_gis_well_c = df_dbscan_gis_well.merge(df_dbscan_gis_well_byplay, on=['OperatorName','Play','year'], how='left')
df_dbscan_gis_well_c = df_dbscan_gis_well_c.drop(['FirstSixMonthOil', 'FirstSixMonthBOE','OGprod','OGprod_byplay','firstyear_byplay','wellcount_byplay','cumsum_byplay'], axis=1)
df_dbscan_gis_well_c.to_csv('controls.csv', index=False, encoding='utf-8-sig' )


###split chemical-concentration data by shale play (for DBSCAN)
df_dbscan_gis = df_dbscan_gis.drop(['FirstSixMonthOil', 'FirstSixMonthBOE','API'], axis=1)

for i, x in df_dbscan_gis.groupby('Play'):
    p = os.path.join(os.getcwd(), "playDBSCAN_{}.csv".format(i.lower()))
    x.to_csv(p, index=False, encoding='utf-8-sig')

