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

# # Assignment 2: Mono Lake Water Balance with Jupyter Notebook

# *4 pts, due 2/2/17*

# This tutorial-style assignment will give you a chance to apply your newfound Python and Jupyter Notebook skills, and to hone your experience with modeling water balances. As usual, feel free to work with others to complete this assignment, but each person should turn in his or her own notebook. Please submit your assignment as an .ipnb file.

# ## Getting started

# Working in this Jupyter Notebook file, start with the basics: insert a new markdown cell to put your name on the file! Then, execute the block of code below, which imports the packages that we will use in this exercise.

# In[ ]:


# Import numerical tools
import numpy as np

# Import pyplot for plotting
import matplotlib.pyplot as plt

#Import pandas for reading in and managing data
import pandas as pd

# Magic function to make matplotlib inline; other style specs must come AFTER
get_ipython().run_line_magic('matplotlib', 'inline')

get_ipython().run_line_magic('config', "InlineBackend.figure_formats = {'svg',}")
#Comment the above line and uncomment the line below if svg graphics are not working in your browser.
#%config InlineBackend.figure_formats = {'png', 'retina'}


# ## Importing data

# The next step is to import the data files that served as inputs for the Mono Lake water balance. We will input them as arrays using the "pandas" package. First you want a datafile to import! We will import the data as a CSV (comma-separated value) file, which is one of the most portable types of data files. Navigate to the now-familiar [Mono Lake water basin spreadsheet](https://docs.google.com/spreadsheets/d/1oTJ5i9jse58mUIq5pC6Dl049WM-j8gCWJ188d1Vq5cc/edit?usp=sharing) from the past assignment, and go to File>Download As>Comma-separated values. 
# 
# Then copy and paste the path and filename of this file into the definition of *fname* below. 

# In[ ]:


# Use pd.read_csv() to read in the data and store in a DataFrame
fname = '/Users/lglarsen/Downloads/Mono Lake Water Balance Model.csv'
df = pd.read_csv(fname)


# In[ ]:


# Look at the contents (first 10 indices)
df.head()


# We imported the whole model (i.e., the whole spreadsheet) out of convenience, but we're going to pretend we don't already know the answer, and that all we know is the starting lake volume, the record of diversions, and the precipitation at nearby stations, for the 1941-2011 period. For convenience, we will define separate arrays for each of these variables. 

# ## Defining arrays of input variables

# One of the nice things about using data frames is that it is easy to refer to specific columns. We just do so by name!

# In[ ]:


#Define an array for year
year = df['Year']
#Take a look at how this new variable, *year*, is formatted. Notice that the index of the first
#element in the array is 0. This is different from Matlab, which indexes arrays beginning with 
#1!
year


# In this next cell we will show one way for pulling out the input diversion data up to and including 2011 without having to go back and look at the original spreadsheet to figure out the row that coincides with 2011.

# In[ ]:


#First, pull out the diversion column
div_data = df['Diversions']
div_data = div_data[year<=2011] #This takes only the rows of div_data for which the statement
#'year is less than or equal to 2011' is true.


# Now you should do the same for rainfall from the nearby gauge station! Insert a code cell below and create an array (call it *nearby_rain*) that contains just the rainfall from the nearby gauge up to and including the year 2011.

# ## Define initial conditions and user-defined parameters

# Next we need to define the initial conditions of the model. In this case, the initial condition is the starting lake volume in 1941. This is also where we define the parameters (the things that we want to tweak in each version of the model).

# In[ ]:


V_init = 6300000000 #Initial lake volume, meters^3
input_precip = 0.531 #Estimated precipitation, m/year (this is at the nearby station)
input_ET = 1.25 #Estimated ET, m/year


# ## Model calculations

# Now we get to the heart of the model. The heart of the model consists of a for-loop, which calculates all of the variables that we are interested in, on a year-by-year basis. Each new calculation of volume takes the volume solved for in the previous year and adds to it the difference between inflows and outflows, multiplied by the one-year time step. (Whenever I am coding, I find it helpful to keep track of the units in the comments.)
# 
# Before we get to this for loop, there are a few other things that we need to do to set it up. For example, we need to tell it that the "old" volume we should use for the first calculation is the initial condition, 'V_init.' Secondly, we need to establish that we want to save the whole history of lake levels and salinities (in other words, after we run this code, we don't just want the LAST timestep; we want the whole time series so that we can plot it!). This is called 'initialization.'

# ### Initilization

# In[ ]:


V_old = V_init #The first value of V_old that we'll use in the calculations is V_init
max_lake_elev = 0 #Maximum lake elevation attained through the whole modeled time series.At the
#outset, we don't know, so we set it equal to zero.
elev = np.zeros(len(year)) #This initializes our array of lake elevations that we want to save with 
#as many entries as there are years.
salinity = np.zeros(len(year)) #Just another array of zeros, where we will save salinity.


# ### The For Loop!

# Note the syntax being used here in the for loop. In python, range() generates integers between zero and 1- the number specified as the argument, which here is the length of the *year* array. The indexing statement at the beginning of the for loop ends with a colon. For every value of *i* that it runs through, it will execute all of the lines of code that are indented below it. Take note that in python, indentation is important!

# In[ ]:


for i in range(len(year)): #Execute this trivial for loop to see how for loops and the range()
#function work in python.
    print(i)
    


# Now let's put the real calculations into the for loop. I will get it started with the calculation of lake depth, area and a few other quantities. You should complete the for loop with the other important calculations that are columns on the spreadsheet model. I have created 'pseudocode' for you (in other words, outlining what you need to do in the comments). Your task is to complete the code and produce plots of lake level and salinity over time. Remember that you can find all of the original mathematical formulas in the Mono Lake water balance spreadsheet. 
# 
# Note the syntax of the "if" statement. You will need to use similar syntax when writing a formula for diversions (the hardest part of this assignment!). Diversions, precipitation, and ET all are estimated differently before and after 2011. You can incorporate the calculation of ET into the if statement I have started below. Note that you will need multiple nested if statements for the diversions calculation, and you that you need to solve for diversions AFTER solving for runoff (so that you know how much water is available to take in the sustainable period).
# 
# *Another important note*: It is rare that you will write code correctly on the first time. Fortunately, Python lets you know where the error is with an arrow and gives you an inkling of what the error is. When I am debugging a code, I find it useful to insert the print() command into the code so that I can make sure my variables look as I expect them to. **Note that when you re-execute the for loop after running it once, you will need to rerun the initialization cell, otherwise you will get incorrect values!**

# In[ ]:


for i in range(len(year)):
    depth = (V_old/9496.7)**0.3333 #meters. Note here the syntax for the exponent, **
    elev[i] = (depth+1866.8)*3.28 #meters.
    max_lake_elev = max(elev) #Figure out the maximum lake elevation we've seen
    area = np.pi*(depth*95.24)**2 #meters squared. The function that calculates pi 
    #comes from the numpy package (which we imported as np), so we use np to call it.
    
    #Solve for precipitation rate in the Mono Lake basin and ET (you do the ET part)
    if year[i]<=2011: #If the current year is less than or equal to 2011...
        rain_mono = nearby_rain[i]*0.37 #Scale the precipitation measured at a nearby gauge
        
    else: #Otherwise, if this is the "future..." (well, if it's beyond 2011)
        rain_mono = input_precip*0.37 #Use the pre-specified input precipitation, scaled
    
    #Solve for the precipitation flux.
    
    #Solve for the groundwater flux.
    
    #Solve for runoff.
    
    #Solve for evapotranspiration flux in m^3/yr
                        
    #Now solve for diversions
    
    
    #Solve for the 'net' inflow (see formula in column Q of spreadsheet)
    
    #Solve for salinity
    
    #Now solve for the new lake volume!
    V_new = V_old + (net_inflow-ET_flux)*1 #m^3. The 1 isn't really necessary, but it is included for completeness. It represents 1 year.
    V_old = V_new #Not to sound philosophical, but all that is new shall become old.


# ## Graph the results

# Now you should make two plots of the variables whose time series you have saved. We set an example with the first one.

# In[ ]:


#Make a plot for lake elevation
plt.plot(year, elev, '-')
plt.xlabel('Year')
plt.ylabel('Elevation, m')
plt.title('Historical and projected lake level')

#Make a plot for salinity


# ## Finishing up

# You are done! If you want, feel free to play around with the user-set parameters (post-2011 average evaporation rate and precipitation rate) to see how it changes the plots. Remember you will have to go through and execute all of the cells again, in order, if you do this. Save and turn in your final notebook, showing the execution of all cells. 

# In[ ]: