#! /usr/bin/env python
# Time-stamp: <2015-02-17 17:01:57 pbrowne>
# MATLAB original by Javier Amezcua
# Python conversion by David Livings
# Python modifications by Philip Browne

###############################################################################
#    lorenz63_empire.py Implements Lorenz 1963 model with EMPIRE coupling
#
#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
###############################################################################

from mpi4py import MPI #import MPI stuff
import sys
import numpy as np

def initialise_mpi():
    #call mpi_init ( mpi_err )
    #not needed as called automatically when importing mpi4py
    global cpl_root,cpl_mpi_comm
    mdl_num_proc = 1
    da = np.array(0, dtype='i')
    mpi_comm_world = MPI.COMM_WORLD
    world_id = mpi_comm_world.Get_rank()
    world_size = mpi_comm_world.Get_size()

    models=mpi_comm_world.Split(da,world_id)
    models_size= models.Get_size()
    models_id  = models.Get_rank()
    mdlcolour = models_id / mdl_num_proc
    mdl_mpi_comm = models.Split(mdlcolour,models_id)
    mdl_id = mdl_mpi_comm.Get_rank()
    if ( mdl_id == 0):
        couple_colour = 9999
    else:
        couple_colour = MPI.UNDEFINED
    cpl_mpi_comm =  mpi_comm_world.Split(couple_colour,mdlcolour)
    if ( mdl_id == 0):
        nens = cpl_mpi_comm.Get_size()
        ptcl_id = cpl_mpi_comm.Get_rank()
        nda = world_size-models_size
        nens = nens-nda
        for da in range(1,nda+1):
            if ( ptcl_id < np.float64((nens*(da)))/np.float64(nda)):
                cpl_root = da-1 + nens
                break
            else:
                cpl_root = -1


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[:])
        #send state via EMPIRE as update has occured
        cpl_mpi_comm.Send([xt,len(xt),MPI.DOUBLE],dest=cpl_root,tag=1)
        cpl_mpi_comm.Recv([xt,len(xt),MPI.DOUBLE],source=cpl_root,tag=MPI.ANY_TAG)

        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


initialise_mpi() #call subroutine to initialise EMPIRE

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])
file.close()

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

#do the initial EMPIRE send and recieve as the initial state is now defined
cpl_mpi_comm.Send([x_0,len(x_0),MPI.DOUBLE_PRECISION],dest=cpl_root,  tag=1)
cpl_mpi_comm.Recv([x_0,len(x_0),MPI.DOUBLE_PRECISION],source=cpl_root,tag=MPI.ANY_TAG,status=None)

lorenz63(x_0)