#!/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      Numerical simulations to obtain number of Windows of Opportunity\n ')
print ('         Coded by Jim van Belzen\n')
print ('################################################################################\n')

# Version 1.0 (19/10/2019) - Python implementation
# Version 1.1 (23/10/2019)
# Version 1.2 (05/11/2019)
# Version 2.0 (08/11/2019) - Numerical experiments
# Version 2.1 (30/06/2020)
# Version 2.2 (01/07/2020)
# Version 2.3 (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 ndimage
import pandas as pd
import time


###=== 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


def fnumwoo(E, T0, dt, lamb):
    """
    Numerically derive Windows of Opportunity from any environmental signal
    """
    if T0>=np.min(E) and lamb>0:
       
        num = np.array((np.max(E) - T0)/(lamb*dt), dtype=int)
        
        #- 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
    else:
        w = 0
        
    return w



###=== NUMERICAL EXPERIMENTS ===###
# Start the timer:
start_time = time.time()

##- Parameters & settings
seed = 0

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

#- parameters to convert
a = 0.74    # 5/6
b = 0.493   # 19/40

#- Organism, T
T0 = mu/10

#- Temporal settings
dt      = .01
Tend    = 1000

#- number of replicate's
N = 1000

#- lambda's to run
lambs   = np.arange(0, 2.6, 0.1)

nt      = len(lambs)

# Setup output matrices
Estats  = np.zeros((N,5))  # number of potential Windows
woos    = np.zeros((N,nt)) # number of actual Windows of Opportunity
fwoos   = np.zeros((N,nt)) # Fraction of Windows succefully used

#- estimate epsilon to rescale sigma of E
epsilon = sigma/(a*(tau/dt)**b)

c = 0
for i in range(N):
    #- set seed
    np.random.seed(seed + i)
    
    #- Environmental signal
    [E, t] = fHRN(mu, tau, epsilon, Tend, dt)
    
    #- Store E stats
    w0 = fpotwin(E, T0, dt)
    potwoos = np.sum(w0>0)
    Estats[i,0] = np.mean(E)
    Estats[i,1] = np.min(E)
    Estats[i,2] = np.max(E)
    Estats[i,3] = np.std(E)
    Estats[i,4] = potwoos
    
    #- loop
    for j in range(nt):
        w  = fnumwoo(E, T0, dt, lambs[j])
        
        #- Store stats
        nwoos     = np.sum(w>0)
        woos[i,j] = nwoos
        if nwoos>0 and potwoos>0:
            fwoos[i,j] = nwoos/potwoos
        else:
            fwoos[i,j] = 0
            
        c += 1
        print("Simulation at: %1.2f perc." % (c/(N*nt)*100))

end_time = time.time()
elapsed_time = end_time - start_time
print("Simulation took: %1.1f s" % (elapsed_time))
print("Finished!\n")


df = pd.DataFrame(fwoos)
df.to_csv('results/fwoos_tau_%1.2f_sigma_%1.2f.csv' % (tau,sigma), header=lambs)

df = pd.DataFrame(woos)
df.to_csv('results/woos_tau_%1.2f_sigma_%1.2f.csv' % (tau,sigma), header=lambs)

df = pd.DataFrame(Estats)
df.to_csv('results/Estats_tau_%1.2f_sigma_%1.2f.csv' % (tau,sigma))

plt.figure()
plt.plot(lambs, np.quantile(fwoos, 0.5, axis=0),'.-r')
plt.xlabel(r'Growth rate, $\lambda$')
plt.ylabel('Opportunities used/# windows')
plt.savefig('results/Fig_WoOs_tau_%1.2f_sigma_%1.2f.png' % (tau,sigma), dpi=300, bbox_inches = 'tight')

plt.close()

print ('\n\t###=== THE END ===###\n')
