#Import modules (including xlwt: http://xlwt.readthedocs.io/en/latest/)
import sys, os
import pandas as pd
from openpyxl import load_workbook

#Get the file names
dataDir = '../../Data'
inDataFN = dataDir + os.sep + 'WaterBalanceData.csv'
inXlsxFN = dataDir + os.sep + 'USWaterBalanceSheet.xlsx'

#Load the data into a pandas dataframe
dfData = pd.read_csv(inDataFN)

#Row and column indices
Aq = 2
Do = 3
In = 4
Ic = 5
Ig = 6
Li = 7
Mi = 8
PS = 9
TC = 10
TR = 11
Supply = 12
Sf = 21 #Surface
Gw = 22 #Groundwater

#Create the dictionary of GroupNames and cell locations
celLocs = {'Aquaculture_Surface':(Sf,Aq),
           'Aquaculture_Groundwater':(Gw,Aq),
           'Domestic_Surface':(Sf,Do),
           'Domestic_Groundwater':(Gw,Do),
           'Industrial_Surface':(Sf,In),
           'Industrial_Groundwater':(Gw,In),
           'Irrigation_Crop_Surface':(Sf,Ic),
           'Irrigation_Crop_Groundwater':(Gw,Ic),
           'Irrigation_Golf_Surface':(Sf,Ig),
           'Irrigation_Golf_Groundwater':(Gw,Ig),
           'Livestock_Surface':(Sf,Li),
           'Livestock_Groundwater':(Gw,Li),
           'Mining_Surface':(Sf,Mi),
           'Mining_Groundwater':(Gw,Mi),
           'PublicSupply_Surface':(Sf,PS),
           'PublicSupply_Groundwater':(Gw,PS),
           'ThermoElec_OnceThru_Surface':(Sf,TC),
           'ThermoElec_OnceThru_Groundwater':(Gw,TC),
           'ThermoElec_Recirc_Surface':(Sf,TR),
           'ThermoElec_Recirc_Groundwater':(Gw,TR),
           'Supply':(4,12)
          } 

#Create the workbook object
wb = load_workbook(filename = inXlsxFN)

for year in (2000,2005, 2010):
    #Get the year worksheet in the workbook
    ws = wb.get_sheet_by_name(str(year))
    
    #Label the sheet
    ws.cell(column=1,row=1,value="US Water Balance: {}. Values in MGal/Year".format(year))
    
    #use the dictionary to insert values
    for name, cellLoc in celLocs.items():
        #Get the value for selected year
        val = dfData[(dfData.Group == name) & (dfData.YEAR == year)]['MGal'].iloc[0]
        #insert it into the Excel file
        ws.cell(column = cellLoc[1],row = cellLoc[0],value = val)

wb.save(dataDir+os.sep+'BalanceSheet.xlsx')