#! /usr/bin/env python
# Time-stamp: <2014-05-16 14:37:04 pbrowne>
# MATLAB original by Javier Amezcua
# Python conversion by David Livings
# Python modifications by Philip Browne

###############################################################################
#    lorenz96.py Implements Lorenz 1996 model
#
#The MIT License (MIT)
#
#Copyright (c) 2014 Philip A. Browne
#
#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:
#
#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE.
#
#Email: p.browne@reading.ac.uk
###############################################################################

import sys
import numpy as np

def lorenz96(x_0):
    """This function computes the time evolution of the Lorenz 96 model.

    It is the general case for N variables; often N=40 is used.

    The Lorenz 1996 model is cyclical: dx[j]/dt=(x[j+1]-x[j-2])*x[j-1]-x[j]+F

    Inputs:   - x_0, original position.
    Outputs:  - x,   the nature run.   """

    # Initialize values for integration
    t = np.arange(0,tmax+tstep/2,tstep)

    x = x_0

    # The integration
    for i in range(len(t)-1): # for each time
        x[:] = rk4(x[:]) # solved via RK4
        print ' '.join(map(str,x[:]))

##____________________________________
## Functions for the integration
def rk4(Varsold):
    "This function contains the RK4 routine."
    k1 = f(Varsold                   )
    k2 = f(Varsold+(1.0/2.0)*tstep*k1)
    k3 = f(Varsold+(1.0/2.0)*tstep*k2)
    k4 = f(Varsold+          tstep*k3)
    Varsnew = Varsold + tstep*(k1+2*k2+2*k3+k4)/6.0
    return Varsnew

def f(x):
    "The actual Lorenz 1996 model."
    #global N, F
    k=np.empty_like(x)
    k.fill(np.nan)
    # Remember it is a cyclical model, hence we need modular algebra
    for j in range(N):
        k[j]=(x[(j+1)%N]-x[(j-2)%N])*x[(j-1)%N]-x[j]+F
    return k



global N,F,tstep,tmax

# Load the initial conditions, timestep, maximum time, and parameters
file = open('lorenz96.dat', 'r')
for line in file:
    if(line.split()[1] == 'tstep'):
        tstep  = eval(line.split()[0])
    elif(line.split()[1] == 'tmax'):
        tmax   = eval(line.split()[0])
    elif(line.split()[1] == 'F'):
        F      = eval(line.split()[0])
    elif(line.split()[1] == 'N'):
        N      = eval(line.split()[0])
    else:
        print 'Error in data file lorenz96.dat'
        sys.exit(1)
file.close()


x_0 = 0.0*np.array(range(0,N))
x_0[:] = F
pert = 0.05
pospert = np.ceil(N/2.0)-1
x_0[pospert] = F+pert

print ' '.join(map(str,x_0[:]))
lorenz96(x_0)