import snake
from snake.cuibm.simulation import CuIBMSimulation


if snake.__version__ != '0.1.2':
  warnings.warn('The figures were originally created with snake-0.1.2, '+
                'you are using snake-{}'.format(snake.__version__))

# Computes the mean force coefficients from the cuIBM simulation
# with grid-spacing h=0.006 in the uniform region and atol=1.0E-06 for the 
# Poisson solver.
# The force coefficients are averaged between 32 and 64 time-units.
simulation_directory = os.path.join(os.path.dirname(__file__), 
                                    'h0.006_vatol16_patol6_dt0.0002')
simulation = CuIBMSimulation(description='atol=1.0E-06',
                             directory=simulation_directory)
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])

# Computes the mean force coefficients from the cuIBM simulation
# with grid-spacing h=0.006 in the uniform region and atol=1.0E-08 for the 
# Poisson solver.
# The force coefficients are averaged between 32 and 64 time-units.
simulation_directory = os.path.join(os.path.dirname(__file__), 
                                    'h0.006_vatol16_patol8_dt0.0002')
simulation2 = CuIBMSimulation(description='atol=1.0E-08',
                              directory=simulation_directory)
simulation2.read_forces()
simulation2.get_mean_forces(limits=[32.0, 64.0])

# Creates a table with the time-averaged force coefficients.

saves the figure, and prints a data-frame that contains the mean values.
"""

import os

from snake.petibm.simulation import PetIBMSimulation
from snake.cuibm.simulation import CuIBMSimulation


simulation = PetIBMSimulation(description='PetIBM')
simulation.read_forces()
time_limits = (32.0, 64.0)
simulation.get_mean_forces(limits=time_limits)
simulation.get_strouhal(limits=time_limits, order=200)

krishnan = CuIBMSimulation(description='Krishnan et al. (2014)')
filepath = os.path.join(os.environ['SNAKE'],
                        'resources',
                        'flyingSnake2d_cuibm_anush',
                        'flyingSnake2dRe2000AoA35',
                        'forces')
krishnan.read_forces(file_path=filepath)
krishnan.get_mean_forces(limits=time_limits)
krishnan.get_strouhal(limits=time_limits, order=200)

simulation.plot_forces(display_coefficients=True,
                       coefficient=2.0,
                       display_extrema=True, order=200,
                       limits=(0.0, 80.0, 0.0, 3.0),
                       other_simulations=krishnan,
                       other_coefficients=2.0,

"""
Computes, plots, and saves the 2D vorticity field from a cuIBM simulation at
saved time-steps.
"""

from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation()
simulation.read_grid()

for time_step in simulation.get_time_steps():
  simulation.read_fields('vorticity', time_step)
  simulation.plot_contour('vorticity',
                          field_range=(-5.0, 5.0, 40),
                          filled_contour=False,
                          view=(-1.0, -3.0, 15.0, 3.0),
                          colorbar=False,
                          cmap=None,
                          colors='k',
                          style='seaborn-dark',
                          width=8.0)

simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
cd_original = 2.0*simulation.forces[0].mean['value']
cl_original = 2.0*simulation.forces[1].mean['value']

# Computes the mean coefficients from the cuIBM simulations
# reported in Krishnan et al. (2014).
# The force coefficients are averaged between 32 and 64 time-units.
cd_krishnan, cl_krishnan = [], []
for re in ['Re1000', 'Re2000']:
  for aoa in ['AoA25', 'AoA30', 'AoA35', 'AoA40']:
    simulation_directory = os.path.join(os.environ['SNAKE'],
                                        'resources',
                                        'flyingSnake2d_cuibm_anush',
                                        'flyingSnake2d'+re+aoa)
    krishnan = CuIBMSimulation(directory=simulation_directory)
    krishnan.read_forces()
    krishnan.get_mean_forces(limits=[32.0, 64.0])
    cd_krishnan.append(2.0*krishnan.forces[0].mean['value'])
    cl_krishnan.append(2.0*krishnan.forces[1].mean['value'])


# plot figure
aoa  = [25, 30, 35, 40]
pyplot.style.use(os.path.join(os.environ['SNAKE'], 
                              'snake', 
                              'styles', 
                              'snakeReproducibility.mplstyle'))
fig = pyplot.figure(figsize=(6, 8))
gs = gridspec.GridSpec(3, 2, 
                       height_ratios=[1, 1, 0.5])

                    help='file containing the list of simulation directories')
parser.add_argument('--save-dir', 
                    dest='save_directory', 
                    type=str, 
                    default=script_directory, 
                    help='directory where to save the figures')
args = parser.parse_args()
with open(args.map, 'r') as infile:
  dirs = yaml.load(infile)


# Plots the instantaneous force coefficients at Re=1000 and AoA=35deg
# using a current version of cuIBM with CUSP-0.5.1
# and compares to the results reported in Krishnan et al. (2014).
simulation_directory = dirs['cuibm-current-cusp051']['Re1000AoA35']
simulation = CuIBMSimulation(description='cuIBM (current) - cusp-0.5.1',
                             directory=simulation_directory)
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

krishnan = CuIBMSimulation(description='Krishnan et al. (2014)')
krishnan.read_forces(file_path=os.path.join(os.environ['SNAKE'],
                                            'resources',
                                            'flyingSnake2d_cuibm_anush',
                                            'flyingSnake2dRe1000AoA35',
                                            'forces'))
krishnan.get_mean_forces(limits=[32.0, 64.0])
krishnan.get_strouhal(limits=[32.0, 64.0], order=200)

dataframe = simulation.create_dataframe_forces(display_strouhal=True,
                                               display_coefficients=True,

print('\nPython version:\n{}'.format(sys.version))
print('\nsnake version: {}\n'.format(snake.__version__))

simulation = IBAMRSimulation(description='IBAMR',
                             directory=os.path.join(os.environ['HOME'],
                                                    'snakeReproducibilityPackages',
                                                    'ibamr',
                                                    'bodyMarkers',
                                                    'Re2000AoA30'))
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

krishnan = CuIBMSimulation(description='Krishnan et al. (2014)')
krishnan.read_forces(file_path=os.path.join(os.environ['SNAKE'],
                                            'resources',
                                            'flyingSnake2d_cuibm_anush',
                                            'flyingSnake2dRe2000AoA30',
                                            'forces'))
krishnan.get_mean_forces(limits=[32.0, 64.0])
krishnan.get_strouhal(limits=[32.0, 64.0], order=200)

dataframe = simulation.create_dataframe_forces(display_strouhal=True,
                                               display_coefficients=True,
                                               coefficient=-2.0)
dataframe2 = krishnan.create_dataframe_forces(display_strouhal=True,
                                              display_coefficients=True,
                                              coefficient=2.0)
dataframe = dataframe.append(dataframe2)

"""
Computes, plots, and saves the 2D vorticity field from a cuIBM simulation at
saved time-steps.
"""

from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation()
simulation.read_grid()

for time_step in simulation.get_time_steps():
  simulation.read_fields('vorticity', time_step)
  simulation.plot_contour('vorticity',
                          field_range=(-5.0, 5.0, 101),
                          filled_contour=True,
                          view=[-2.0, -5.0, 15.0, 5.0],
                          style='mesnardo',
                          width=8.0)

"""
Plots the vorticity field of a 2D cuIBM simulation.
"""

import os
from snake.cuibm.simulation import CuIBMSimulation
from snake.body import Body


simulation = CuIBMSimulation()
simulation.read_grid()

for time_step in simulation.get_time_steps():
  all_bodies = Body(file_path=os.path.join('{:0>7}'.format(time_step),
                                           'bodies'))
  n_total = all_bodies.x.size
  bodies = [Body(), Body()]
  bodies[0].x = all_bodies.x[:n_total // 2]
  bodies[0].y = all_bodies.y[:n_total // 2]
  bodies[1].x = all_bodies.x[n_total // 2:]
  bodies[1].y = all_bodies.y[n_total // 2:]
  simulation.read_fields('vorticity', time_step)
  simulation.plot_contour('vorticity',
                          field_range=(-2.0, 2.0, 40),
                          filled_contour=True,
                          bodies=bodies,
                          view=(-3.0, -5.0, 15.0, 5.0),
                          style='seaborn-dark',
                          cmap='viridis',
                          width=8.0)

# file: plotVorticitySnake.py
# author: Olivier Mesnard ([email protected])
# description: Plots the 2D vorticity field near the snake.
# Run this script from the simulation directory.


from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation()
simulation.read_grid()

for time_step in simulation.get_time_steps():
    simulation.read_fields("vorticity", time_step)
    simulation.plot_contour(
        "vorticity", field_range=[-5.0, 5.0, 101], filled_contour=True, view=[-0.75, -1.0, 1.50, 1.0], width=8.0
    )

# file: plotliftCoefficientCompareKrishnanEtAl2014.py
# author: Olivier Mesnard ([email protected])
# description: Plots the instantaneous lift coefficient
#              and compare to results from Krishnan et al. (2014).
# Run this script from the simulation directory.


import os

from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation(description='cuIBM')
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

krishnan = CuIBMSimulation(description='Krishnan et al. (2014)')
krishnan.read_forces(file_path='{}/resources/flyingSnake2d_cuibm_anush/'
                               'flyingSnake2dRe2000AoA35/forces'
                               ''.format(os.environ['SNAKE']))
krishnan.get_mean_forces(limits=[32.0, 64.0])
krishnan.get_strouhal(limits=[32.0, 64.0], order=200)

simulation.plot_forces(indices=[1],
                       display_coefficients=True,
                       coefficient=2.0,
                       display_extrema=True, order=200,
                       limits=(0.0, 80.0, 0.0, 3.0),
                       other_simulations=krishnan,
                       other_coefficients=2.0,

                                            'linear'),
                  'folders': ['30', '90', '270', '810']},
        'uquad': {'directory': os.path.join('uquad_T00.25_20_0.00050',
                                            'linear'),
                  'folders': ['20', '60', '180', '540']},
        'vquad': {'directory': os.path.join('vquad_T00.25_20_0.00050',
                                            'linear'),
                  'folders': ['20', '60', '180', '540']}}

time_step = 500
field_names = ['x-velocity', 'y-velocity']

for key, series in data.items():
  series['cases'] = []
  for folder in series['folders']:
    case = CuIBMSimulation(directory=os.path.join(series['directory'], folder),
                           description=folder)
    case.read_grid()
    case.read_fields(field_names, time_step)
    nx, ny = case.grid[0].size - 1, case.grid[1].size - 1
    masks = read_mask(os.path.join(case.directory, 'mask.txt'), nx, ny)
    for i, name in enumerate(field_names):
      case.fields[name].values = numpy.multiply(case.fields[name].values,
                                                masks[i])

    series['cases'].append(case)
  series['first'] = get_observed_orders(series['cases'][:-1],
                                        field_names,
                                        series['cases'][0],
                                        save_name=None)
  series['last'] = get_observed_orders(series['cases'][1:],
                                       field_names,

    ) / numpy.log(ratio)


time_increments = ["dt=5.0E-05", "dt=1.0E-04", "dt=2.0E-04"]
final_time_steps = [2000, 1500, 1250]
refinement_ratio = 2.0

fx, fy = [], []
p = []
qx, qy = [], []

main_directory = os.path.dirname(__file__)

for time_increment, final_time_step in zip(time_increments, final_time_steps):
    directory = os.path.join(main_directory, time_increment.replace("=", ""))
    simu = CuIBMSimulation(directory=directory, description=time_increment)
    simu.read_forces()
    fx.append(simu.forces[0].values[-1])
    fy.append(simu.forces[1].values[-1])
    simu.read_grid()
    p_simu = simu.read_pressure(final_time_step)
    p.append(p_simu.values)
    qx_simu, qy_simu = simu.read_fluxes(final_time_step)
    qx.append(qx_simu.values)
    qy.append(qy_simu.values)


file_path = os.path.join(os.path.dirname(__file__), "temporalConvergence.txt")
with open(file_path, "w") as outfile:
    outfile.write("\n* Drag force:\n")
    outfile.write("Value and relative difference with the reference value (smallest dt)\n")

"""
Plots the vorticity field of a 2D cuIBM simulation.
"""

import os
from snake.cuibm.simulation import CuIBMSimulation
from snake.body import Body


simulation = CuIBMSimulation()
simulation.read_grid()

for time_step in simulation.get_time_steps():
  body = Body(file_path=os.path.join('{:0>7}'.format(time_step), 'bodies'))
  simulation.read_fields('vorticity', time_step)
  simulation.plot_contour('vorticity',
                          field_range=(-10.0, 10.0, 10),
                          filled_contour=False,
                          bodies=body,
                          view=(-1.0, -1.0, 3.5, 1.0),
                          style='seaborn-dark',
                          colors='grey',
                          colorbar=False,
                          width=8.0)

from snake.cuibm.simulation import CuIBMSimulation


if snake.__version__ != '0.1.2':
  warnings.warn('The figures were originally created with snake-0.1.2, '+
                'you are using snake-{}'.format(snake.__version__))


cases = []
fx, fy = [], []
resolutions = ['h=0.00267', 'h=0.004', 'h=0.006']
time_limits = (20.0, 28.0)

for resolution in resolutions:
  case = CuIBMSimulation(directory=os.path.join(os.path.dirname(__file__),
                                                resolution.replace('=', '')),
                         description=resolution)
  case.read_forces()
  case.get_mean_forces(limits=time_limits)
  fx.append(case.forces[0].mean['value'])
  fy.append(case.forces[1].mean['value'])
  cases.append(case)

# Calculates the observed order of convergence for the time-averaged force
# coefficients.
ratio = 1.5
order = numpy.log((fx[2]-fx[1])/(fx[1]-fx[0]))/numpy.log(ratio)
print(order)
order = numpy.log((fy[2]-fy[1])/(fy[1]-fy[0]))/numpy.log(ratio)
print(order)

from snake.cuibm.simulation import CuIBMSimulation
from snake.solutions.koumoutsakosLeonard1995 import KoumoutsakosLeonard1995


# Parse from the command-line the directory of the runs.
parser = argparse.ArgumentParser()
parser.add_argument('--directory',
                    dest='directory',
                    default=os.getcwd(),
                    type=str,
                    help='directory of the runs')
args = parser.parse_args()
directory = args.directory

simulation = CuIBMSimulation(directory=directory,
                             description='cuIBM')
simulation.read_forces()

# Reads drag coefficient of Koumoutsakos and Leonard (1995) for Re=40.
file_name = 'koumoutsakos_leonard_1995_cylinder_dragCoefficientRe40.dat'
file_path = os.path.join(os.environ['CUIBM_DIR'], 'data', file_name)
kl1995 = KoumoutsakosLeonard1995(file_path=file_path, Re=40)

# Plots the instantaneous drag coefficients.
images_directory = os.path.join(directory, 'images')
if not os.path.isdir(images_directory):
  os.makedirs(images_directory)
pyplot.style.use('seaborn-dark')
kwargs_data = {'label': simulation.description,
               'color': '#336699',
               'linestyle': '-',

# file: plotForceCoefficients.py
# author: Olivier Mesnard ([email protected])
# description: Plots the instantaneous force coefficients.
# Run this script from the simulation directory.


from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation()
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

simulation.plot_forces(display_coefficients=True,
                       coefficient=2.0,
                       display_extrema=True, order=200,
                       limits=(0.0, 80.0, 0.0, 3.0),
                       save_name='forceCoefficients')
dataframe = simulation.create_dataframe_forces(display_strouhal=True,
                                               display_coefficients=True,
                                               coefficient=2.0)
print(dataframe)

#              and compare to results from Krishnan et al. (2014).
# Run this script from the simulation directory.


import os

from snake.openfoam.simulation import OpenFOAMSimulation
from snake.cuibm.simulation import CuIBMSimulation


simulation = OpenFOAMSimulation(description='IcoFOAM')
simulation.read_forces(display_coefficients=True)
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

krishnan = CuIBMSimulation(description='Krishnan et al. (2014)')
krishnan.read_forces(file_path='{}/resources/flyingSnake2d_cuibm_anush/'
                               'flyingSnake2dRe2000AoA35/forces'
                               ''.format(os.environ['SNAKE']))
krishnan.get_mean_forces(limits=[32.0, 64.0])
krishnan.get_strouhal(limits=[32.0, 64.0], order=200)

simulation.plot_forces(indices=[1],
                       display_coefficients=True,
                       display_extrema=True, order=200,
                       limits=(0.0, 80.0, 0.0, 3.0),
                       other_simulations=krishnan,
                       other_coefficients=2.0,
                       save_name='liftCoefficientCompareKrishnanEtAl2014')
dataframe = simulation.create_dataframe_forces(indices=[1],
                                               display_strouhal=True,

# description: Plots the instantaneous force coefficients of three runs:
#              1- 2k35 / current cuIBM with CUSP-0.4.0;
#              2- 2k35 / current cuIBM with CUSP-0.5.1;
#              3- 2k35 / older cuIBM with CUSP-0.4.0;


import os

from matplotlib import pyplot

from snake.cuibm.simulation import CuIBMSimulation


simulation = CuIBMSimulation(description='cuIBM (current) - CUSP-0.4.0',
                             directory=os.path.join(os.environ['HOME'],
                                                    'simulations_cuIBM',
                                                    'production-cusp-0.4.0',
                                                    'flyingSnake2dRe2000AoA35_20160502'))
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

other = CuIBMSimulation(description='cuIBM (current) - CUSP-0.5.1',
                        directory=os.path.join(os.environ['HOME'],
                                               'simulations_cuIBM',
                                               'production-cusp-0.5.1',
                                               'flyingSnake2dRe2000AoA35_20160502'))
other.read_forces()
other.get_mean_forces(limits=[32.0, 64.0])
other.get_strouhal(limits=[32.0, 64.0], order=200)

import snake
from snake.cuibm.simulation import CuIBMSimulation


if snake.__version__ != '0.1.2':
  warnings.warn('The figures were originally created with snake-0.1.2, '+
                'you are using snake-{}'.format(snake.__version__))

atol_min, atol_max = 5, 16
time_step = 1500

main_directory = os.path.dirname(__file__)

reference_directory = os.path.join(main_directory,
                                   'atol16')
reference = CuIBMSimulation(directory=reference_directory,
                            description='atol=1.0E-16')
reference.read_forces()
fx_reference = reference.forces[0].values[-1]
fy_reference = reference.forces[1].values[-1]
reference.read_grid()
p_reference = reference.read_pressure(time_step)
qx_reference, qy_reference = reference.read_fluxes(time_step)

fx_errors, fy_errors = [], []
p_errors = []
qx_errors, qy_errors = [], []

for atol in range(atol_min, atol_max):
  directory = os.path.join(main_directory,
                           'atol{}'.format(atol))
  simu = CuIBMSimulation(directory=directory,

import os
import sys

from matplotlib import pyplot

import snake
from snake.cuibm.simulation import CuIBMSimulation


print('\nPython version:\n{}'.format(sys.version))
print('\nsnake version: {}\n'.format(snake.__version__))

simulation = CuIBMSimulation(description='cuIBM (current) - CUSP-0.4.0',
                             directory=os.path.join(os.environ['HOME'],
                                                    'snakeReproducibilityPackages',
                                                    'cuibm',
                                                    'current',
                                                    'cusp040',
                                                    'Re2000AoA35'))
simulation.read_forces()
simulation.get_mean_forces(limits=[32.0, 64.0])
simulation.get_strouhal(limits=[32.0, 64.0], order=200)

other = CuIBMSimulation(description='cuIBM (current) - CUSP-0.5.1',
                        directory=os.path.join(os.environ['HOME'],
                                                    'snakeReproducibilityPackages',
                                                    'cuibm',
                                                    'current',
                                                    'cusp051',
                                                    'Re2000AoA35'))
other.read_forces()