import os.path as p

import numpy as np
import matplotlib.pyplot as plt

self_path = p.dirname(p.abspath(__file__))
root_path = p.dirname(self_path)

import sys
sys.path.append(p.join(root_path, 'test'))
sys.path.append(p.join(root_path, 'reproductions/Davis_et_al_2014'))

import output_preservation_test as opt

import subprocess as sub

import MIMutils as mim

import time

mim_exec = p.join(root_path, "MIM")


def run_davis_et_al_2014(nx,ny,layers,nTimeSteps,dt):
    with opt.working_directory(p.join(self_path, "Davis_et_al_2014")):
        run_experiment(write_input_davis_et_al_2014, nx, ny, layers, nTimeSteps, dt)


# Generic function to run an experiment
def run_experiment(write_input, nx, ny, layers, nTimeSteps,dt):
    """ Run MIM"""
    sub.check_call(["rm", "-rf", "input/"])
    sub.check_call(["rm", "-rf", "output/"])
    sub.check_call(["mkdir", "-p", "output/"])
    sub.check_call(["rm", "-rf", "run_finished.txt"])

    with opt.working_directory(root_path):
        sub.check_call(["make", "MIM"])
    with opt.working_directory("input"):
        write_input(nx, ny, layers,nTimeSteps,dt)
    tweak_parameters(nx, ny, layers,nTimeSteps, dt)

    then = time.time()
    sub.check_call([mim_exec])
    print "MIM execution took", time.time() - then


def tweak_parameters(nx, ny, layers, nTimeSteps,dt):
    """Alter the parameters.in file"""
    sub.check_call(
        "cat parameters.in " +
        "| sed 's/^ nx =.*,$/ nx = %d,/'" % (nx,) +
        "| sed 's/^ ny =.*,$/ ny = %d,/'" % (ny,) +
        "| sed 's/^ layers =.*,$/ layers = %d,/'" % (layers,) +
        "| sed 's/^ nTimeSteps =.*,$/ nTimeSteps = %d,/'" % (nTimeSteps,) +
        "| sed 's/^ dt =.*,$/ dt = %d,/'" % (dt,) +
        "> parameters.new", shell=True)
    sub.check_call(["mv", "parameters.new", "parameters.in"])


# Creat the inputs
def write_input_davis_et_al_2014(nx, ny, layers, nTimeSteps,dt):
    assert layers == 1
    xlen = 1.5e6
    ylen = 2.7e6
    grid = mim.Grid(nx, ny, xlen / nx, ylen / ny)

    opt.write_f_plane(nx,ny, 14.5842318e-5) # Coriolis at North Pole
    
    write_davis_wetmask(grid)
    write_davis_wind(grid)
    write_davis_sponge(grid)
    write_wind_time_series(nTimeSteps,dt)


    with opt.fortran_file('initH.bin', 'w') as f:
        X,Y = np.meshgrid(grid.x,grid.y)
        initH = 400.*np.ones(X.shape)
        f.write_record(initH.astype(np.float64))


@@ -2382,12 +2382,725 @@
 .float64))%0A%0A
+%0Adef write_davis_wetmask(grid):%0A    %22%22%22Write the wet mask file for a recreation of Davis et al. (2014).%22%22%22%0A%0A    with opt.fortran_file('wetmask.bin', 'w') as f:%0A        X,Y = np.meshgrid(grid.x,grid.y)%0A        # start with land everywhere and carve out space for water%0A        wetmask = np.zeros(X.shape, dtype=np.float64)%0A        # circular gyre region%0A        wetmask%5B((Y-1950e3)**2 + (X-750e3)**2) %3C 750e3**2%5D = 1%0A        # 150 km wide channel%0A        wetmask%5B(X-750e3)**2 %3C 75e3**2%5D = 1%0A        # sponge region%0A        wetmask%5BY%3C780e3%5D = 1%0A        # clean up the edges%0A        wetmask%5B 0, :%5D = 0%0A        wetmask%5B-1, :%5D = 0%0A        wetmask%5B :, 0%5D = 0%0A        wetmask%5B :,-1%5D = 0%0A        f.write_record(wetmask)%0A%0A

create wetmask