#!/usr/bin/python

import sys
import os
import os.path
import glob
sys.path.insert(0, '/usr/lib/python2.7/site-packages')
from FreeFOAM.compat import *
import FreeFOAM.tutorial

class buoyantSimpleBuoyantCavityRunner(FreeFOAM.tutorial.CaseRunner):
   def __init__(self):
      FreeFOAM.tutorial.CaseRunner.__init__(self, 'buoyantSimple_buoyantCavity')
      self.add_app_step('blockMesh')
      self.add_app_step('buoyantSimple')
      self.add_app_step('sample', args=['-latestTime'])
      self.add_step('createPlots', self._createPlots, skip_test=True)

   def _createPlots(self, case_dir, stamp_file, test_mode):
      assert not test_mode
      # do we have numpy and matplotlib?
      try:
         import numpy as np
      except:
         stamp_file.write('Missing numpy, skipping graphs\nRESULT: FAILURE\n')
         return False
      try:
         import pylab as pl
      except:
         stamp_file.write(
               'Missing matplotlib/pylab, skipping graphs\nRESULT: FAILURE\n')
         return False
      # find latest time and the data directories
      timeDirs = [d for d in
            filter(os.path.isdir, glob.glob(os.path.join(case_dir, 'sets/*')))]
      imax = np.argmax(list(map(
         lambda d: float(os.path.basename(d)), timeDirs)))
      dataroot = timeDirs[imax]
      exptdataroot = os.path.join(case_dir, 'validation', 'exptData')
      # do the plotting
      try:
         fig = pl.figure()
         for i in map(int, '1 3 4 5 6 7 9'.split()):
            stamp_file.write(
                  'Processing temperature profile at y/yMax of 0.%d\n'%i)
            data = np.loadtxt(os.path.join(dataroot, 'y0.%d_T.xy'%i))
            exptdata = np.loadtxt(
                  os.path.join(exptdataroot, 'mt_z0_%d0_lo.dat'%i)
                  )/(1000, 1)+(0, 273.15)
            pl.cla()
            pl.xlim((0, 0.08))
            pl.ylim((285, 310))
            pl.xlabel('Channel width, x / [m]')
            pl.ylabel('Temperature / [K]')
            pl.grid(True)
            pl.plot(exptdata[:,0], exptdata[:,1], 'k-o', label='Expt 0.%d'%i)
            pl.plot(data[:,0], data[:,1], 'k-', lw=2, label='FreeFOAM 0.%d'%i)
            pl.legend(loc='upper left')
            pl.savefig(
                  os.path.join(case_dir, 'validation', 'OF_vs_EXPT_T%d.eps'%i),
                  format='eps')
         for i in map(int, '1 3 4 5 6 7 9'.split()):
            stamp_file.write('Processing velocity profile at y/yMax of 0.%d\n'%i)
            data = np.loadtxt(os.path.join(dataroot, 'y0.%d_U.xy'%i))
            exptdata = np.loadtxt(
                  os.path.join(exptdataroot, 'mv_z0_%d0_lo.dat'%i))/(1000, 1)
            pl.cla()
            pl.xlim((0, 0.08))
            pl.ylim((-0.2, 0.2))
            pl.xlabel('Channel width, x / [m]')
            pl.ylabel('Vertical velocity component, Uy / [m/s]')
            pl.grid(True)
            pl.plot(exptdata[:,0], exptdata[:,1], 'k-o', label='Expt 0.%d'%i)
            pl.plot(data[:,0], data[:,2], 'k-', lw=2, label='FreeFOAM 0.%d'%i)
            pl.legend(loc='upper left')
            pl.savefig(
                  os.path.join(case_dir, 'validation', 'OF_vs_EXPT_U%d.eps'%i),
                  format='eps')
      except Exception:
         e = sys.exc_info()[1]
         stamp_file.write('*** Error *** '+str(e)+'\nRESULT: FAILURE\n')
         return False

      stamp_file.write('RESULT: SUCCESS\n')
      return True

if __name__ == '__main__':
   os.chdir(os.path.abspath(os.path.dirname(sys.argv[0])))
   runner = FreeFOAM.tutorial.TutorialRunner()
   runner.add_case(buoyantSimpleBuoyantCavityRunner())
   sys.exit(runner.main())

# ------------------- vim: set sw=3 sts=3 ft=python et: ------------ end-of-file
