#!/usr/bin/env python3

print ('\n################################################################################\n')
print ('     A probabilistic framework for Windows of Opportunity:\n')
print ('         The role of variance in critical transitions.\n')
print ('\n      Rescaling environmental signals.\n')
print ('         Coded by Jim van Belzen\n')
print ('################################################################################\n')

# Version 1.0 (19/10/2019) - Python implementation
# Version 1.0 (14/01/2021) - final version
# Coded and developed by Jim van Belzen
# published under the Creative Commons Attribution Non-Commercial license
# which allows users to read, copy, distribute and make derivative works
# for noncommercial purposes from the material, as long as the author of
# the original work is cited.
# This code comes with no warranties or support
# http://people.zeelandnet.nl/jbelzen/



###=== SETUP PYTHON ENVIRONMENT ===###
import numpy as np
import numpy.matlib
import matplotlib.pyplot as plt
from scipy import stats
from scipy import ndimage



###=== DEFINE FUNCTIONS ===###
def fRDS(E0, E1, k0, k1, Tend, dt):
    """
    Function to create Regular Dichotomous Signal (RDS)
    """
    N = int(Tend/dt + 1)
    t = np.linspace(0, N-1, N)*dt
    E = np.linspace(0, N-1, N)*0

    #- initial state
    if np.random.rand() <= k0/(k0+k1):
        q = 0;
    else:
        q = 1;

    #- create signal
    c = 0
    for i in range(N):
        # report current state
        if q == 0:
            E[i] = E0
        elif q == 1:
            E[i] = E1
        
        # determine new state
        if q == 0:
            if c >= 1./k0:
                q = 1
                c = 0
        elif q == 1:
            if c >= 1./k1:
                q = 0
                c = 0
        
        c += dt

    return E, t


def fDMN(E0, E1, k0, k1, Tend, dt):
    """
    Function to create Dichotomous Markov Noise (DMN) signal
    """
    #- temporal rescaling of transistion probabilities
    k0 = 1-(1-k0)**dt
    k1 = 1-(1-k1)**dt
    
    N = int(Tend/dt + 1)
    t = np.linspace(0, N-1, N)*dt
    E = np.linspace(0, N-1, N)*0
    
    #- initial state
    if np.random.rand() <= k0/(k0+k1):
        q = 0;
    else:
        q = 1;

    #- create signal
    for i in range(N):
        # report current state
        if q == 0:
            E[i] = E0
        elif q == 1:
            E[i] = E1
        
        # determine new state
        if q == 0:
            if np.random.rand() <= k0:
                q = 1
        elif q == 1:
            if np.random.rand() <= k1:
               q = 0

    return E, t


def fOUP(mu, Theta, sigma, Tend, dt):
    """
    Function to create correlated noise signal based on Ornstein-Uhlenbeck process (OUP)
    mu      average of stationary signal
    Theta   rate at which signal restores to average
    sigma   magnitude of Gaussian noise
    """
    #- setting up output variable
    N = int(Tend/dt + 1)
    t = np.linspace(0, N-1, N)*dt
    E = np.linspace(0, N-1, N)*0

    #- initial state (based on 6sigma rule for 99.9%)
    xt = mu + 3*sigma*np.random.normal()

    #- create signal
    for i in range(N):
        E[i] = xt
        dxt = Theta*(mu - xt)*dt + sigma*np.random.normal()
        xt += dxt

    return E, t
    
    
def fHRN(mu, tau, sigma, Tend, dt):
    """
    Function to create correlated noise signal based on Hasselmann (1976) variation of the
    Ornstein-Uhlenbeck process (OUP)
    mu      average of stationary signal
    tau     approximate period of the signal
    sigma   magnitude of Gaussian noise
    """
    
    #- setting up output variable
    N = int(Tend/dt + 1)
    t = np.linspace(0, N-1, N)*dt
    E = np.linspace(0, N-1, N)*0

    #- initial state (based on 6sigma rule for 99.9%)
    xt = mu + 3*sigma*np.random.normal()

    #- create signal
    for i in range(N):
        E[i] = xt
        dxt = (1/tau)*(mu - xt)*dt + sigma*np.random.normal()
        xt += dxt

    return E, t


def fpotwin(E, T0, dt):
    """
    Analyze environmental signal for potential windows
    """
    lbl, nlbl = ndimage.label(E<=T0)
    w = np.zeros((nlbl+1,1))
    wlbl = range(1, nlbl+1)
    for i in wlbl:
        w[i] = np.sum(lbl == i)*dt

    return w[1:]#w


def fnumwoo(E, T0, dt, lamb):
    """
    Numerically derive Windows of Opportunity from any environmental signal
    """
    #num = int(len(E)/2)
    num = int(50/dt)
    #num = 50
    #num = int((np.max(E) - T0)/(lamb*dt))
    #print('num:',num)
    #print('E:',np.max(E))
    
    #- run growth
    T = (E<=T0)*T0
    B = (T>0)
    for i in range(num):
        ii = i+1
        T += lamb*B*dt
        T[:-ii] = np.where(T[:-ii] > E[ii:], T[:-ii], 0)
        B = (T>0)
    
    #- get Windows of Opportunity
    lbl, nlbl = ndimage.label(B==1)
    w = np.zeros((nlbl+1,1))
    wlbl = range(1, nlbl+1)
    for i in wlbl:
        w[i] = np.sum(lbl == i)*dt
        
    return w[1:]#w




###=== SETTINGS & PARAMETERS ===###
#- set seed
seed = 2

#- Environmental condition, E
mu      = 1     # average of stationary signal
tau0    = 1     # approximate period of the signal
tau1    = 10    # approximate period of the signal
sigma   = 0.5   # magnitude of Gaussian noise

#- parameters to convert
a = 0.74
b = 0.493

#- Organism, T
T0 = mu/10

#- Temporal settings
dt      = .01
Tend    = 1000

#- estimate sigma of E, needed for rescaling
sigma_hat0 = sigma*a*(tau0/dt)**b
sigma_hat1 = sigma*a*(tau1/dt)**b

epsilon0 = sigma/(a*(tau0/dt)**b)
epsilon1 = sigma/(a*(tau1/dt)**b)


###=== SIM ===###

#- Environmental signal
np.random.seed(seed)
[E0, t] = fHRN(mu, tau0, sigma, Tend, dt)
np.random.seed(seed)
[E0_hat, t] = fHRN(mu, tau0, epsilon0, Tend, dt)

#- Environmental signal
np.random.seed(seed)
[E1, t] = fHRN(mu, tau1, sigma, Tend, dt)
np.random.seed(seed)
[E1_hat, t] = fHRN(mu, tau1, epsilon1, Tend, dt)


###=== STATS ===###
print('\n### Before rescaling of signal ###')
print('\nmu E0 = %1.3f, mu E1 = %1.3f' % (np.mean(E0),np.mean(E1)))
print('median of E0 = %1.3f and E1 = %1.3f' % (np.quantile(E0, 0.5),np.quantile(E1, 0.5)))
print('sigma E0 = %1.3f, sigma E1 = %1.3f' % (np.std(E0),np.std(E1)))
print('sigma hat E0 = %1.3f, sigma hat E1 = %1.3f' % (sigma_hat0,sigma_hat1))
print('\n### After rescaling of signal ###')
print('\nmu E0 = %1.3f, mu E1 = %1.3f' % (np.mean(E0_hat),np.mean(E1_hat)))
print('median of E0 = %1.3f and E1 = %1.3f' % (np.quantile(E0_hat, 0.5),np.quantile(E1_hat, 0.5)))
print('sigma E0 = %1.3f, sigma E1 = %1.3f' % (np.std(E0_hat),np.std(E1_hat)))

print('\n### Potential number of windows ###')
W00 = fpotwin(E0_hat, T0, dt)
W01 = fpotwin(E1_hat, T0, dt)
print('\nmu N0(W) = %d and N1(W) = %d' % (len(W00),len(W01)))


###=== GRAPHICAL OUTPUT ===###
# Figure panel properties
nrows = 2
ncols = 1
dpi = 300
tmp = plt.rcParams['figure.figsize']
plt.rcParams['figure.figsize'] = [tmp[0] * 1., tmp[1] * 1.]

plt.figure()

plt.subplot(nrows, ncols, 1)
plt.plot(t, E0, '-k')
plt.plot(t, E1, '-b')
plt.xlabel('Time, t')
plt.ylabel('Environment, E')

plt.subplot(nrows, ncols, 2)
plt.plot(t, E0_hat, '-k')
plt.plot(t, E1_hat, '-b')
plt.plot([0,Tend],[T0,T0],'--r')
plt.xlabel('Time, t')
plt.ylabel('Environment, E')
plt.show()


print ('\n\t###=== Rescaling environmental signals DONE ===###\n')
