# coding: utf-8

# # APS - new snow

# ## Imports

from __future__ import print_function
import pylab as plt
import datetime
import numpy as np
plt.rcParams['figure.figsize'] = (14, 6)


# ## Parameters, categories and scores

# ## Hourly score- and decay-functions

# The score function for the new snow problem is of type $$ s = a \cdot \log_b{x} $$
# The decay function for the new snow problem is of type $$ d = a \cdot b^x $$

from matplotlib.patches import Rectangle
from scipy.optimize import curve_fit


def score_func(x, a, b):#, c):
    return (1 / (1+np.exp(-x+a))) * b
    #return a * (1. / np.log(b)) * np.log(x) + c


def decay_func(x, a, b):
    return a * b ** x


# ### Score function
# Hourly new snow amount in mm water equivalent
control_points = np.array([
        [-0.5, -2.],
        [0.0, 0.],
        [2., 3.],
        [10., 40.]
    ])
new_snow_1h_cat = control_points[:, 0]
new_snow_1h_score = control_points[:, 1] 


params = curve_fit(score_func, new_snow_1h_cat, new_snow_1h_score)

[sa, sb] = params[0]
print("Score function with parameters {0:.2f} and {1:.2f} results in daily increase of {2:.2f} points with 2 mm hourly precipitation.".format(sa, sb, score_func(2, sa, sb)*24))

x = np.arange(0, 20.0, 0.1)
res = score_func(x, sa, sb)

plt.scatter(new_snow_1h_cat, new_snow_1h_score)
plt.plot(x, res)

plt.xlabel('Hourly new snow amount (mm w.e.)')
plt.ylabel('Score')
plt.xlim(0, 20); plt.ylim(0, 100)
#plt.gca().add_patch(Rectangle((0, 0), 40, 100, edgecolor="lightgrey", facecolor="lightgrey"))


# ### Decay function

# In[18]:

# Hourly air temperature (ideally snow surface temperature, but that is not available at the moment)
control_points = np.array([
        [-40., 0.001],
        [-20.0, 0.01],
        [-5, 0.3],
        [0.0, 1.5],
        #[1., 0.],
        #[4., -10.],
        [5., 4.]
    ])
new_snow_1h_decay_cat = control_points[:, 0]
new_snow_1h_decay_score = control_points[:, 1] 

params = curve_fit(decay_func, new_snow_1h_decay_cat, new_snow_1h_decay_score)

[da, db] = params[0]
print("Decay function with parameters {0:.2f} and {1:.2f} results in daily reduction of {2:.2f} points at zero degrees Celsius.".format(da, db, decay_func(0, da, db)*24))
print("Decay function with parameters {0:.2f} and {1:.2f} results in daily reduction of {2:.2f} points at -10 degrees Celsius.".format(da, db, decay_func(-10, da, db)*24))
print("Decay function with parameters {0:.2f} and {1:.2f} results in daily reduction of {2:.2f} points at -20 degrees Celsius.".format(da, db, decay_func(-20, da, db)*24))

x = np.arange(-40, 10.0, 0.1)
res = decay_func(x, da, db)

plt.scatter(new_snow_1h_decay_cat, new_snow_1h_decay_score)
plt.plot(x, res)

plt.xlabel('Hourly air temperature (C)')
plt.ylabel('Decay')
plt.xlim(-42, 12); plt.ylim(-10, 10)
#plt.gca().add_patch(Rectangle((0, 0), 40, 100, edgecolor="lightgrey", facecolor="lightgrey"))


# ## Working with real data
# Load data from filefjell.db containing two weeks of met-data from the station.
# The database was generated by the notebook "xgeo_chartserver".

# In[19]:

import sqlite3
import pandas as pd


# In[20]:

db_name = 'filefjell.db'
conn = sqlite3.connect(db_name)
cur = conn.cursor()


# In[21]:

sql = "SELECT * from FILEFJELL"
df = pd.read_sql(sql, conn, index_col='index', parse_dates=['index']) #
#conn.close()
df.head()
df.columns


# In[22]:

df.plot(subplots='True', figsize=(14, 9))


# Derive required input for new snow problem:
# - new snow amount last 0-24 h
# - new snow amount last 24-72 h
# - temperature gradient last 6 h (relate temperature to settling rate of previous snow falls)

# In[23]:

#df['24h_precip'] = df['FILEFJELL - KYRKJESTØLANE (54710), Nedbør (mm)'].rolling(window=24).sum()
#df['72h_precip'] = df['FILEFJELL - KYRKJESTØLANE (54710), Nedbør (mm)'].rolling(window=72).sum()


# ## *TODO*:
# - Find real data with higher precip...
# 

# In[24]:

df['score'] = score_func(df['FILEFJELL - KYRKJESTØLANE (54710), Nedbør (mm)'], sa, sb)
df['decay'] = decay_func(df['FILEFJELL - KYRKJESTØLANE (54710), Lufttemperatur (°C)'], da, db)
df['new_snow_score'] = np.clip(df['score'] - df['decay'], 0, 120) # using 120 to see how often we exceed 100!


# In[27]:

df.plot(subplots='True', figsize=(14, 9))
plt.gcf().savefig('real.png', dpi=300)


# Select certain days using .loc. Works before or after the ['column_name']. See http://pandas.pydata.org/pandas-docs/stable/indexing.html#selection-by-label

# In[ ]:




# In[ ]:

df.loc['20160201']['new_snow_score'].plot()


# Or meteorological day

# In[ ]:

df['new_snow_score'].loc['20160201T06:00:00':'20160202T06:00:00'].plot()


# In[ ]:

sday = df.loc['20160201']
sday['new_snow_score'].describe()
sday['new_snow_score'].plot.box()


# In[ ]:




# ### Summing up the scores and decays

# In[ ]:

def score_sum(new_snow_score, new_snow_decay, wind_speed_score):
    _sum = np.zeros_like(new_snow_score)
    _sum[0] = np.clip((new_snow_score[0] * wind_speed_score[0] - new_snow_decay[0]), 0, 100)
    for i in np.arange(1, len(new_snow_score)):
        _sum[i] = np.clip(_sum[i-1] + (new_snow_score[i] * wind_speed_score[i] - new_snow_decay[i]), 0, 100)
                    
    return _sum


# In[ ]:

df['wind_score'] = score_wind_speed(df['FILEFJELL - KYRKJESTØLANE (54710), Vindhastighet 10m (m/s)'])

df['snow_score'] = score_new_snow_1h(df['FILEFJELL - KYRKJESTØLANE (54710), Nedbør (mm)'])
df['snow_decay'] = decay_func(df['FILEFJELL - KYRKJESTØLANE (54710), Nedbør (mm)'], a, b)
df['new_snow_score'] = score_sum(df['snow_score'], df['snow_decay'], df['wind_score'])

#TODO: add a wind_speed_decay function; should we need to regard wind_direction?


# In[ ]:

df.plot(subplots='True', figsize=(14, 23))


# # Outdated stuff

# ### Using three parameters a,b,c

# In[5]:

# Hourly air temperature (ideally snow surface temperature, but that is not available at the moment)
control_points = np.array([
        [-40., 0.01],
        [-20.0, 0.05],
        [-5, 1.],
        [0.0, 3.],
        #[1., 0.],
        #[4., -10.],
        [5., 4.]
    ])
new_snow_1h_decay_cat = control_points[:, 0]
new_snow_1h_decay_score = control_points[:, 1] 

params = curve_fit(decay_func, new_snow_1h_decay_cat, new_snow_1h_decay_score)

[a, b, c] = params[0]
print("Decay function with parameters {0:.2f}, {1:.2f} and {2:.2f} results in daily reduction of {3:.2f} points at zero degrees Celsius.".format(a, b, c, decay_func(0, a, b, c)*24))

x = np.arange(-40, 10.0, 0.1)
res = decay_func(x, a, b, c)

print(decay_func(0, a, b, c))

plt.scatter(new_snow_1h_decay_cat, new_snow_1h_decay_score)
plt.plot(x, res)

plt.xlabel('Hourly air temperature (C)')
plt.ylabel('Decay')
plt.xlim(-42, 12); plt.ylim(-10, 10)
#plt.gca().add_patch(Rectangle((0, 0), 40, 100, edgecolor="lightgrey", facecolor="lightgrey"))


# In[6]:

# Hourly new snow amount in mm water equivalent
control_points = np.array([
        [-2., 10.],
        [0.0, 2.0],
        [0.2, 0.5]#,
        #[0.5, 2.],
        #[1., 0.],
        #[4., -10.],
        #[10., -50.]
    ])
new_snow_1h_decay_cat = control_points[:, 0]
new_snow_1h_decay_score = control_points[:, 1] 


params = curve_fit(decay_func, new_snow_1h_decay_cat, new_snow_1h_decay_score)

[a, b] = params[0]
print("Decay function with parameters {0:.2f} and {1:.2f} results in daily reduction of {2:.2f} points with zero precipitation.".format(a, b, decay_func(0, a, b)*24))

x = np.arange(0, 20.0, 0.1)
res = decay_func(x, a, b)

plt.scatter(new_snow_1h_decay_cat, new_snow_1h_decay_score)
plt.plot(x, res)

plt.xlabel('Hourly new snow amount (mm w.e.)')
plt.ylabel('Decay')
plt.xlim(0, 20); plt.ylim(0, 10)
plt.gca().add_patch(Rectangle((0, 0), 40, 100, edgecolor="lightgrey", facecolor="lightgrey"))


# ### Main control factors

# In[3]:

# New snow amount last 24 h 0-60 cm [10 cm intervals]
new_snow_24h_cat = np.array([0, 20, 40, 60, 80, 100, 120])
new_snow_24h_score = np.array([0.5, 8.0, 15.0, 19., 21.0, 27.0, 33.3])

# Wind speed 0-100 km/h [0,10,20,30,40,50,60,80,100]
wind_speed_cat = np.array([-1.5, 0, 2.5, 5, 7.5, 10, 15, 20, 25, 30, 40]) # m/s
wind_speed_score = np.array([-1.0, 0.8, 2.0, 2.9, 3.2, 3.0, 1.1, 0.6, 0.4, 0.2, 0.0])


# ### Weighting
# Weights are added if they are independent of the value of the core factor or multiplied if they are related to the core factor.

# In[4]:

# New snow amount last 24-72h 0-100 cm [0,10,20,30,40,50,60,80,100]
new_snow_24_72h_cat = np.array([0, 10, 20, 30, 40, 50, 60, 80, 100])
new_snow_24_72h_weight = np.array([0.8, 0.83, 0.86, 0.89, 0.92, 0.95, 0.98, 0.99, 1.0]) # a weight for new_snow_24h

# Evolution of temperature
evolution_temperature_cat = ["constant very cold",
                             "constant cold",
                             "constant warm",
                             "rise towards 0 deg after snowfall",
                             "substantial cooling after snowfall"]

# Bonding to existing snowpack 
bonding_existing_snowpack_cat = ["favorable", "moderate", "poor"]

# Type of new snow 
type_new_snow_cat = ["loose-powder", "soft", "packed", "packed and moist"]


# The *new_snow_24_72h_weight* are used to weight the *new_snow_24h_scores* prior to multiplying it with *wind_speed_score*.
# 
# In order to achive a smooth fit within the range of interest I added some control points just right outside the normal range for the higher order polynomials.
# 
# The temperature evolution during a snowfall can be fitted to a curve which can then be compared to predefined curves/scenarios. The scenario with the best correlation is chosen to define the category.
# Temperature or change in snow depth will be used if the precipitation event is rain or snow when applied to data from a weather station. The AROME model generally supplies that separation.  
# 
# The *type_new_snow_cat* can be infered from *evolution_temperature* and *wind_speed*.
# 
# In the first place the categories can be set manually.

# ## Score functions

# ### New snow 24 h

# In[5]:

new_snow_24h_fit = np.polyfit(new_snow_24h_cat, new_snow_24h_score, 2)
score_new_snow_24h = np.poly1d(new_snow_24h_fit)

x = np.arange(0, 120.0)
res = score_new_snow_24h(x)

plt.scatter(new_snow_24h_cat, new_snow_24h_score)
plt.plot(x, res)
plt.xlabel('New snow amount (cm)')
plt.ylabel('Score')
plt.axhline(33.3, color='grey', ls='--')


# ### New snow 24-72 h

# In[6]:

new_snow_24_72h_fit = np.polyfit(new_snow_24_72h_cat, new_snow_24_72h_weight, 2)
score_new_snow_24_72h = np.poly1d(new_snow_24_72h_fit)

x = np.arange(0, 100.0)
res = score_new_snow_24_72h(x)

plt.scatter(new_snow_24_72h_cat, new_snow_24_72h_weight)
plt.plot(x, res)
plt.xlabel('New snow last 24-72 h (cm)')
plt.ylabel('Weight')
plt.axhline(1.0, color='grey', ls='--')


# ### Wind speed

# In[7]:

wind_speed_fit = np.polyfit(wind_speed_cat, wind_speed_score, 5)
score_wind_speed = np.poly1d(wind_speed_fit)

x = np.arange(-5, 45.0)
res = score_wind_speed(x)

plt.scatter(wind_speed_cat, wind_speed_score)
plt.plot(x, res)
plt.xlabel('Wind speed (m/s)')
plt.ylabel('Score')
plt.xlim(0, 30)
plt.ylim(0, 4)
#plt.axvspan(0.0, 36.0, facecolor='grey', alpha=0.5) # model validity range
plt.axhline(3.0, color='grey', ls='--')


# ## New snow vs. wind speed

# In[8]:

new_snow = np.matrix(np.arange(0, 125.0))
sns = score_new_snow_24h(new_snow)

# weighted by new snow amount of the previous two days
new_snow_72 = 40
ns_weight = score_new_snow_24_72h(new_snow_72)

sns *= ns_weight

wind_speed = np.matrix(np.arange(0, 40.0))
swp = score_wind_speed(wind_speed)

M = np.multiply(sns, swp.T)
#print(M)
plt.contourf(M)#np.flipud(M.T))
print("Min {0}; Max {1}".format(np.amin(M), np.amax(M)))
plt.colorbar()
plt.xlabel("New snow last 24h (cm)")
plt.ylabel("Wind speed (m/s)")


# # ToDo

# - calculate new_snow_score for some weeks
# - compare to chosen AP in regional forecast
# - maybe extent to a larger grid
# 
# 
# ...continue with hemsedal_jan2016.py in Test

# # Random scripting testing

# In[ ]:

new_snow_cat = ["0-5", "5-10", "10-15", "15-20"]
new_snow_thres = {(0, 5): 0.2, (5, 10): 0.5, (10, 15): 1, (15, 20): 3}

wind_cat = ["0-3", "4-7", "8-10", "10-15", "16-30"]
wind_thres = {(0, 3): 0.2, (3, 7): 1, (7, 10): 2, (10, 15): 0.2, (15, 30): 0.01}


new_snow_region = np.array([[0, 4, 6, 18],
                             [0, 4, 6, 18],
                             [0, 4, 6, 18]])

wind_region = np.array([[0, 4, 12, 18],
                         [4, 0, 18, 6],
                         [18, 12, 6, 0]])


# In[ ]:

def get_score(a, score_dict):
    for key, value in score_dict.items():
        if key[0] <= a < key[1]:
    #    if a < key:
            return value
            break
    return None


# the dict is not sorted and the comparison *less than* is random...

# In[ ]:

new_snow_region_score = [get_score(a, new_snow_thres) for a in new_snow_region.flatten()]
new_snow_region_score = np.array(new_snow_region_score).reshape(new_snow_region.shape)
print(new_snow_region_score)


# In[ ]:

wind_region_score = [get_score(a, wind_thres) for a in wind_region.flatten()]
wind_region_score = np.array(wind_region_score).reshape(wind_region.shape)
print(wind_region_score)


# In[ ]:

print(wind_region_score * new_snow_region_score)


# In[ ]:

X = np.matrix(np.arange(0, 11.0))
Y = np.matrix(np.arange(10.0, 21.0))

Z = np.multiply(X, Y.T)
print(X)
print(Y.T)
print(Z)
plt.imshow(Z)
print("Min {0}; Max {1}".format(np.amin(Z), np.amax(Z)))
plt.colorbar()


# In[ ]:

arr = np.random.geometric(0.3, (20, 20))
plt.pcolor(arr)


# In[ ]:

window = arr[1:11,1:11]
arr[1:11,1:11] = 1
plt.pcolor(arr)


# In[ ]:

print(np.median(window))
print(window.flatten())
print(np.bincount(window.flatten()))
print(np.sort(window, axis=None))


# In[ ]:



