from pprint import pprint
import matplotlib.pyplot as plt

"""
Explanation: <figure>
  <IMG SRC="../../logo/logo.png" WIDTH=250 ALIGN="right">
</figure>

IHE Python course, 2017
T.N.Olsthoorn, Feb 26, 2017
Working with lists, dicts and sets and their comprehensions
(Excercises, using 2017 world population data, reading data from Excel)
See Excel workbook in this folder
For this we list all files in the folder and check each if it has '.xls' in it and doesn't start with '~', because that means that the file is actually active in the Excel program.
End of explanation
"""


import os

os.listdir('.')

import os

xls_files = [f for f in os.listdir() if ".xls" in f and not f.startswith('~')]

print(xls_files)

"""
Explanation: Import the os module (operating system). It has the method listdir(), which gives a list of the files within a directory.
We apply os.listdir() to get a list of all files and immediately filter it to get only the files that have .xls and don't start with '~'
End of explanation
"""


import openpyxl # excel file functionality

wb = openpyxl.load_workbook(xls_files[0])

type(wb)

wb.get_sheet_names()

# Show the sheetnames inside the Excel workbook wb
wb.get_sheet_names()

"""
Explanation: Next we'll read the data from the Excel workbook
End of explanation
"""


wsPop = wb['Population']

type(wsPop)

"""
Explanation: We can now treat the workbook as a dictionary in which each key is the name of a worksheet
End of explanation
"""


for r in wsPop.rows:
    pass # do something

"""
Explanation: wsPop is an object. It has a number of methods and properties, which can be accessed by typing the name followed by a dot followed by a tab. To get more information it can be followed by a question mark.
On the mac wsPop.rows is a tuple of tuples. On some other systems it is a generator. It generates the tuples one after the other in a loop like
End of explanation
"""


data = [r for r in wsPop.rows]

"""
Explanation: but with a generator, one cannot say len(wsPop.rows) or index it like wsPop.rows[3].
However it's easy to first generate a list of tuples with a comprehension:
End of explanation
"""


cel = wsPop.rows[0][3]  # may not work if wsPop.rows is a generator on your system and not a tuple of tuple.
cel = data[0][3]  # this is a list of tuples, generated above.

cel.value

"""
Explanation: and then continue to work with data instead of wsPop.rows. Therefore, if you have trouble indexing wsPop.rows, first generate a list as shown above and use that wherever you see wsPop.rows below.
End of explanation
"""


print( type(wsPop.rows) )  # shows that the wsPop generates a tuple of tuples

print( len(wsPop.rows) ) # shows the number of rows

print( wsPop.rows[3][1] ) # shows some cell

wsPop.rows[3][1].value # shows the value of some cell

for r in wsPop.rows:
    print(r[1].value)

"""
Explanation: We gen get the contents of sheet as a series of rows using the attribute rows,
and then show an arbitrary value
End of explanation
"""


data = [[c.value for c in r ]   for r in wsPop.rows]

data[15][1]

# Show the first 5 lines
for i in range(5):
    print(data[i])

"""
Explanation: Turn the rows propery into a list of rows, with each row the values of the cells in them.
We can do that in a list comprehension, in this case a list comprehension around a list comprehension.
For each row we turn it into a list of the values of each cell in that row.
We do that for each row.
The result is a list of lists. This is done in one line:
End of explanation
"""


hdrs = data[0]
dims = data[1]

print()
print(hdrs)
print()
print(dims)
print()

"""
Explanation: To index the data directly, we don't want a list of lists but a dictionary with the country name as key, such that for country the data is kept in a dictionary with keys obtained form the headers in the first two rows of the excel file as is shown above.
End of explanation
"""


hdrs = [(h + str(d)).replace('None','')  for h, d in zip(hdrs, dims)]


pprint(hdrs)

h = data[13][:]


print(h)
name = h.pop(1)
print(name)
print(h)

"""
Explanation: Now glue together the hdrs and the dimes, and filter out the None texts:
All hdrs are strings, but that is not the case of the dims, where -2017 was turned in a value.
So when guening h and d together to a new string below, we have to use h + str(d).
The combined headers are obtained using a list comprehension, that also removes the text 'None' from the dims wherever it turns up:
End of explanation
"""


cntry = dict() # empty dict

for i in range(2, len(wsPop.rows)):
    row = [c.value for c in wsPop.rows[i]]
    cname = row.pop(1)
    cntry[cname] = row # entter the key  cname and the value row into the dict cntry

"""
Explanation: We could generate a dict with the country names as keys where each cntry[key] as a list of itmes. We cannot use a dict comprehension here because we need to pop the country value from each row in to get the key for that rows. Hence a for loop is used:
End of explanation
"""


cntry['Italy']

"""
Explanation: Now we can get the data of any cntry like so:
End of explanation
"""


print(hdrs)
hdrs.pop(1)
print()
print(hdrs)

"""
Explanation: However, this is not smart enough. We can't see what the figure in the list mean. Therefore, we should use a dictionary for the contents of each cntry with the fields as keys.
These fields are now simply obtained by popping the second item from the hdrs list:
End of explanation
"""


cntry = dict()
for i in range(2, len(wsPop.rows)):  # skip the first two lines, which are headers
    
    row = [c.value for c in wsPop.rows[i]] # turn the Excel row in a list
  
    cname = row.pop(1)      # pop of the country name, which becomes the key
    
    cntry[cname] = {fld: v for fld, v in zip(hdrs, row)}  # use dict comprehension

"""
Explanation: Now generate the dict again, with the contens of each cntry itself a dict of its fields:
End of explanation
"""


cntry['Netherlands']['Population-2017']

"""
Explanation: Now the contents of an arbitrary country looks like this, it's a dict with fields and values.
End of explanation
"""


totPop2017 = 0.  # start out with zero

for c in cntry.keys():
    totPop2017 += cntry[c]['Population-2017'] # get the value directly from indexing the field

print('Total population in the world is {:.2f} billion'.format(totPop2017/1e9))

# Now compute the toal populatin in 30 years, using the current growth rates

popWorld = [0 for i in range(1, 31)]
for k in cntry.keys():
    pc  = cntry[k]['Population-2017']   # country population 2017
    try:
        frc = cntry[k]['Fert.Rate'] / 100.  # growth rate fraction
    except:
        # needed in case frc contains `None`
        frc = 0 # simply use 0 in those cases

    # for the country
    popCntry = [pc * (1 + frc)**i for i in range(1, 31) ]
    
    # for the entire world
    popWorld = [pc + pw for pc, pw in zip(popCntry, popWorld)]

for i, pw in enumerate(popWorld):
    yr = 2017 + i + 1
    print('popWorld [{}] = {:5.2f} billion'.format(yr, pw/1e9))

"""
Explanation: Let's now compute the total population of the world.
We do this by summing for each country.
End of explanation
"""


wscont = wb['CountriesByContinent'] # read it

print(wscont)  # shows it's a worksheet object

print(wscont.rows[0])  # shows it's a tuple of two cell objects

wscont.rows[0][0].value, wscont.rows[0][1].value  # turned into a tuple of tow strings

"""
Explanation: Population per continent.
We don't have the Continent field associated with the countries. But we do have a list of countries and their continents in a second worksheet:
End of explanation
"""


cont = { v.value : k.value for k, v in wb['CountriesByContinent'] }

pprint(cont)  # notic the key : value pairs separated by commas

"""
Explanation: We can immediately, in a single line, turn this worksheet into a dictionary with name cont that has the country as key and the continent as field (a string).
End of explanation
"""


for c in ['Bahamas', 'Spain', 'Morocco', 'Honduras', 'Cambodia']:
    print('{:20} lies in {}'.format(c, cont[c]))

"""
Explanation: Show for a few countries in which continent they are:
End of explanation
"""


df_cont_cntry = set(cont.keys()) - set(cntry.keys())
pprint(df_cont_cntry)

print("\n{} countries are in dict `cont` that are not in dict `cntry`".format(len(df_cont_cntry)))

"""
Explanation: Adding the field continent to the cntry dict
We like to add the field Continent to our cntry dict using the cont dict.
This would be easy if the country names in both dicts would be exactly the same.
Let's see if this is the case.
We can check this using set logic.
Convert the keys of he cntry dict to a set and do the same with those of the cont dict
First step: What countries are in th cont dict that are not in the cntry dict?
End of explanation
"""


df_cntry_cont = set(cntry.keys()) - set(cont.keys())
pprint(df_cntry_cont)

print("\n{} countries are in dict `cntry` that are not in dict `cont`".format(len(df_cntry_cont)))

"""
Explanation: And likewise: which are in the cntry dict that are not in the cont dict?
End of explanation
"""


# notice that the columsns in this excel sheet are in columns 2 and 3 and not in 1 and 2 !!
# We construct the dict in one line, using a single dict comprehension

missing = {rw[2].value : rw[1].value for rw in wb['missing']}  # don't need .rows

pprint(missing)

"""
Explanation: The set of contries that are in cont but not in cntry obviously have differently spelled names.
Probably the easiest way is to take the names from cntry that are not in cont and look up the continent in which each of these contries lies, and use that to supplement the missing countries. We don't then have to worry about the misspelled names.
Using a list of missing names with their continent to complete the cntry dict
This list of missing contries with their continent is in sheet missed of our workbook.
End of explanation
"""


for k in missing.keys():
    cont[k] = missing[k]

"""
Explanation: Now we can complete our cntry dict with a Continent attribute for every country in it by using the dict cont and the dict missing.
Let's just join them, like so:
End of explanation
"""


for k in cntry.keys():
    cntry[k]['Continent'] = cont[k]

"""
Explanation: And then add the continent to the cntry dict
End of explanation
"""


for k in cntry.keys():
    print("{:30} lies in {}".format(k, cntry[k]['Continent']))

"""
Explanation: Now we can print the country name with its continent next to it:
End of explanation
"""


continents = { cntry[k]['Continent'] for k in cntry.keys()} # set comprehension

print(continents)

"""
Explanation: Population growth per continent
It is now possible to compute and show the population and its growth per continent
The first step is extract the continents from the dict using set logic, i.e. by set comprehension. The result is a set with the unique values from the Continent field.
End of explanation
"""


continent='Europe'

# population and fertility rate for the countries of this continent:
popCont =  [(cntry[k]['Population-2017'], cntry[k]['Fert.Rate']) for k in cntry.keys() if cntry[k]['Continent']==continent]

print("The number of countries in {} is {}".format(continent, len(popCont)))

"""
Explanation: It's then possible to compute the future population be selecting the countries for this continent to do the computation on:
End of explanation
"""


pTot = dict()  # empty dict
Nyr = 50  # year to predict
print("The predicted population inover the next {} years is:".format(Nyr))

for c in continents:
    pTot[c] = [0 for i in range(Nyr)] # start with empty total for each continent
    
    # generate a list of [Pop, fertility rate] for each country in this continent
    popCont =  [(cntry[k]['Population-2017'], cntry[k]['Fert.Rate'])
                            for k in cntry.keys() if cntry[k]['Continent'] == c]

    # compute the country's future population and add to continentn total
    for p, fr in popCont:
        try:
            p  = float(p)
            fr = float(fr)
            # population of country in coming years
            pcntry = [p * (1 + fr/100.)**i for i in range(1, Nyr + 1)]
            
            # add to continent total
            pTot[c]= [pt + p for pt, p in zip(pTot[c], pcntry)]        
        except:
            # it crashes when fertility rate is 'None', we ignore these frew contries
            pass
        
    print("{:10s}".format(c), end="")  # print continent
    for p in pTot[c]:
           print("{:6.2f}".format(p / 1.0e9), end="") # print Pop values
    print()

"""
Explanation: The last but one step is to compute the population for the coming years in each continent. We make a dictionary pTot with the continent as key and which has the list of future population values as a list:
End of explanation
"""


print("The predicted population in billions:")

continents = pTot.keys()

styear = 2018 # starting year
print("{:4s}".format("Year"), end="")
for c in continents:
    print("{:>12}".format(c), end="")

print()
for i in range(Nyr):
    print("{:4d}".format(styear + i), end="")
    for c in continents:
        print("{:12.2f}".format(pTot[c][i]/1e9), end="")
    print()
        

"""
Explanation: A better overview of the results would be to have columns, one per continent, and the numbers vertical, with the year as first column. Although the numbers are now not naturally ordered to do this, it can be realized with little effort as follows:
End of explanation
"""


years = [2018 + i for i in range(Nyr)]
for c in continents:
    plt.plot(years, pTot[c], label=c)
plt.xlabel('year')
plt.ylabel('Population [billions]')
plt.title('Population development')
plt.legend(loc='best', fontsize='x-small')
plt.show()

#plt.pie?

"""
Explanation: Finally, make a plot of the growth curves:
End of explanation
"""


contPop = dict()
for yr in range(Nyr):
    contPop[2018 + yr] = { c : v for c, v in zip(continents, [pTot[c][yr] for c in continents])}

# show it
pprint(contPop[2025])

"""
Explanation: We can also put the year results in a dict that has as key the year and as values a dict with the continents.
End of explanation
"""


import numpy as np

fig, axs = plt.subplots(2,2, sharex=True, sharey=True)
axs = axs.ravel()

for ax, yr in zip(axs.ravel(), [2020, 2030, 2040,2050]):
    ax.set_title(str(yr))  
    x = np.array(list(contPop[yr].values()))/1.0e9
    y = list(contPop[yr].keys())
    #print(x)
    #print(y)
    r = np.sqrt(np.sum(x))/4
    ax.xlim = []
    ax.pie(x, labels=y, radius=r)
plt.show()

"""
Explanation: Then finally we could make pie charts for the population at say 4 points in time.
End of explanation
"""


