#! /usr/bin/env python
# Time-stamp: <2014-11-12 10:59:15 pbrowne>
# MATLAB original by Javier Amezcua
# Python conversion by David Livings
# Python modifications by Philip Browne

###############################################################################
#    lorenz63.py Implements Lorenz 1963 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 lorenz63(x_0):
    """Evolution of the Lorenz 1963 3-variable model.

    Inputs:  - x_0, an array containing the initial position
             - tmax, the maximum time.  Should be a multiple of the
               timestep 0.01.
    Outputs: - xt, the nature run """

    # The time array
    t = np.arange(0,tmax+tstep/2,tstep)

    xt = x_0

    # The cycle containing time integration
    for i in range(len(t)-1): # for each time step
        # We integrate using Runge-Kutta-4
        xt[:] = rk4(xt[:])

        print ' '.join(map(str,xt))


## Auxiliary functions

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):
    "This function contains the actual L63 model."
    # Initialize
    k = np.empty_like(x)
    k.fill(np.nan)
    # The Lorenz equations
    k[0] = sigma*(x[1]-x[0])
    k[1] = x[0]*(rho-x[2])-x[1]
    k[2] = x[0]*x[1]-beta*x[2]
    return k


global tmax,tstep,sigma,beta,rho
x_0 = np.array([1.0 , 1.0 , 1.0], dtype=np.float64)

# Load the initial conditions, timestep, maximum time, and parameters
file = open('lorenz63.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] == 'sigma'):
        sigma  = eval(line.split()[0])
    elif(line.split()[1] == 'beta'):
        beta   = eval(line.split()[0])     
    elif(line.split()[1] == 'rho'):
        rho    = eval(line.split()[0])
    elif(line.split()[1] == 'x1'):
        x_0[0] = eval(line.split()[0])
    elif(line.split()[1] == 'x2'):
        x_0[1] = eval(line.split()[0])
    elif(line.split()[1] == 'x3'):
        x_0[2] = eval(line.split()[0])
    else:
        print 'Error in data file lorenz63.dat'
        sys.exit(1)

file.close()

print ' '.join(map(str,x_0))

lorenz63(x_0)