/*
** (c) 1996-2000 The Regents of the University of California (through
** E.O. Lawrence Berkeley National Laboratory), subject to approval by
** the U.S. Department of Energy.  Your use of this software is under
** license -- the license agreement is attached and included in the
** directory as license.txt or you may contact Berkeley Lab's Technology
** Transfer Department at TTD@lbl.gov.  NOTICE OF U.S. GOVERNMENT RIGHTS.
** The Software was developed under funding from the U.S. Government
** which consequently retains certain rights as follows: the
** U.S. Government has been granted for itself and others acting on its
** behalf a paid-up, nonexclusive, irrevocable, worldwide license in the
** Software to reproduce, prepare derivative works, and perform publicly
** and display publicly.  Beginning five (5) years after the date
** permission to assert copyright is obtained from the U.S. Department of
** Energy, and subject to any subsequent five (5) year renewals, the
** U.S. Government is granted for itself and others acting on its behalf
** a paid-up, nonexclusive, irrevocable, worldwide license in the
** Software to reproduce, prepare derivative works, distribute copies to
** the public, perform publicly and display publicly, and to permit
** others to do so.
*/

#include <winstd.H>

#include <algorithm>
#include <string>
#include <iostream>
#include <iomanip>
#include <fstream>

#include <REAL.H>
#include <IntVect.H>
#include <Box.H>
#include <Grid.H>
#include <Utility.H>
#include <VisMF.H>
#include <ParmParse.H>
#include <CoordSys.H>
#include <Geometry.H>
#include <BCTypes.H>

//
// This MUST be defined if don't have pubsetbuf() in I/O Streams Library.
//
#ifdef BL_USE_SETBUF
#define pubsetbuf setbuf
#endif

// declare global variables here
Real init_shrink    = 1.0;
Real     gravity    = 0.0;

int domain_bc[2*BL_SPACEDIM];
Array< Array<int> > bc;

Box prob_domain;
BoxArray grids;

static const std::string CheckPointVersion = "CheckPointVersion_1.0";

typedef char* charString;

#include <cstdio>
#include <new>
using std::setprecision;
#ifndef WIN32
using std::set_new_handler;
#endif

// functions called 
void print_usage(int,char**);
void parse_command_line(int, char**, charString&, charString&);
void writePlotFile(Grid*, Geometry&, int, Real, std::string&);
void writePlotFile(Grid*, Geometry&, std::string&, std::ostream&, int, Real);
void checkPoint(Grid*, Geometry, int, Real, Real, std::string&);

std::string thePlotFileType();

// ###################################################################
// ##### MAIN PROGRAM
// ###################################################################
main(int argc, char *argv[])
{
  BoxLib::Initialize(argc,argv);

   charString   in_file = NULL;
   charString   res_file = NULL;
   parse_command_line(argc,argv,in_file,res_file);

   ParmParse pp;

   int max_step;
   int max_grid_size;
   Real stop_time;
   Real cfl;
   int  check_int;
   int  plot_int;
   std::string check_file_root = "chk";
   std::string plot_file_root = "plt";
   pp.get("max_step",max_step);
   pp.get("stop_time",stop_time);
   pp.get("cfl",cfl);
   pp.query("init_shrink",init_shrink);
   pp.query("gravity",gravity);
   pp.get("check_int",check_int);
   pp.query("check_file_root",check_file_root);
   pp.get("plot_int",plot_int);
   pp.query("plot_file_root",plot_file_root);
   std::string probin_file;
   if (pp.contains("probin_file"))
   {
       pp.get("probin_file",probin_file);
   }

   Array<int> n_cell(BL_SPACEDIM);
   pp.getarr("n_cell",n_cell,0,BL_SPACEDIM);
#if (BL_SPACEDIM == 2)
   IntVect pd_lo(0,0);
   IntVect pd_hi(n_cell[0]-1,n_cell[1]-1);
#elif (BL_SPACEDIM == 3)
   IntVect pd_lo(0,0,0);
   IntVect pd_hi(n_cell[0]-1,n_cell[1]-1,n_cell[2]-1);
#endif
   prob_domain = Box(pd_lo,pd_hi);
   pp.query("max_grid_size",max_grid_size);

   BoxList grids_list;
   grids_list.push_back(prob_domain);
   grids_list.maxSize(max_grid_size);
   grids.define(grids_list);

   // This reads problem size and coordinate system information.
   Geometry::Setup();
   Geometry geom;
   geom.define(prob_domain);

   pp.get("bcx_lo",domain_bc[0]);
   pp.get("bcx_hi",domain_bc[1]);
   pp.get("bcy_lo",domain_bc[2]);
   pp.get("bcy_hi",domain_bc[3]);

   if (domain_bc[0] == PERIODIC && domain_bc[0] != PERIODIC) {
     std::cout << "IF bcx_lo is periodic then bcx_hi must be periodic " << std::endl;  
     exit(0);
   }
   if (domain_bc[1] == PERIODIC && domain_bc[0] != PERIODIC) {
     std::cout << "IF bcx_hi is periodic then bcx_lo must be periodic " << std::endl;  
     exit(0);
   }
   if (domain_bc[2] == PERIODIC && domain_bc[3] != PERIODIC) {
     std::cout << "IF bcy_lo is periodic then bcy_hi must be periodic " << std::endl;  
     exit(0);
   }
   if (domain_bc[3] == PERIODIC && domain_bc[2] != PERIODIC) {
     std::cout << "IF bcy_hi is periodic then bcy_lo must be periodic " << std::endl;  
     exit(0);
   }

#if (BL_SPACEDIM == 3)
   pp.get("bcz_lo",domain_bc[4]);
   pp.get("bcz_hi",domain_bc[5]);
   if (domain_bc[4] == PERIODIC && domain_bc[5] != PERIODIC) {
     std::cout << "IF bcz_lo is periodic then bcz_hi must be periodic " << std::endl;  
     exit(0);
   }
   if (domain_bc[5] == PERIODIC && domain_bc[4] != PERIODIC) {
     std::cout << "IF bcz_hi is periodic then bcz_lo must be periodic " << std::endl;  
     exit(0);
   }
#endif

   bc.resize(grids.size());
   for (int i = 0; i < grids.size(); i++)
     bc[i].resize(2*BL_SPACEDIM,INTERIOR);

   const int* domlo = prob_domain.loVect();
   const int* domhi = prob_domain.hiVect();

   for (int i = 0; i < grids.size(); i++)
   {
     const int* lo = grids[i].loVect();
     const int* hi = grids[i].hiVect();
     for (int n = 0; n < BL_SPACEDIM; n++)
     {
       if (lo[n] == domlo[n]) bc[i][2*n  ] = domain_bc[2*n];
       if (hi[n] == domhi[n]) bc[i][2*n+1] = domain_bc[2*n+1];
     }
   }

   Grid* grid = new Grid(grids, geom, probin_file);

   Real time = 0.0;
   Real dt   = 0.0;
   int nstep = 0;

   grid->init(nstep,time,dt);

// First do the iterations to calculate an initial pressure
   if (nstep == 0) {
     dt = std::min(init_shrink * grid->estimateDt(dt,cfl), stop_time - time);
     grid->initialIter(time,dt);
     checkPoint(grid,geom,nstep,time,dt,check_file_root);
   }

   writePlotFile(grid,geom,nstep,time,plot_file_root);

   int last_checkpoint = nstep;
   int last_plotpoint  = nstep;

   while (nstep < max_step && time <= stop_time - 1.e-10) {

      if (nstep > 0) 
        dt = std::min(grid->estimateDt(dt,cfl), stop_time - time);

      nstep++;

      // do a timestep
      grid->advance(time, dt);
      time += dt;

      if (ParallelDescriptor::IOProcessor()) 
      {
        std::cout << "STEP = " << nstep << " TIME = " << time;
        std::cout << " DT = " << dt << std::endl;
        std::cout << std::endl; 
      }

      // checkpoint file?
      if (check_int > 0 && nstep%check_int == 0) {
         last_checkpoint = nstep;
         checkPoint(grid,geom,nstep,time,dt,check_file_root);
      };

      if (plot_int > 0 && nstep%plot_int == 0) {
         last_plotpoint = nstep;
         writePlotFile(grid,geom,nstep,time,plot_file_root);
      };

   };

   // dump final checkpoint file if needed
   if (nstep != last_checkpoint) {
      checkPoint(grid,geom,nstep,time,dt,check_file_root);
   };

   // dump final pltfile file if needed
   if (nstep != last_plotpoint) {
      writePlotFile(grid,geom,nstep,time,plot_file_root);
   };

   // clean up memory usage
   delete grid;

   BoxLib::Finalize();
}

// ###################################################################
// ##### print_usage
// ###################################################################
void print_usage(int /* argc */, char *argv[]) {
   std::cerr << "usage:\n";
   std::cerr << argv[0] << " infilename " << std::endl;
   std::cerr << "\n";
   exit(1);
}

// ###################################################################
// ##### parse_command_line
// ###################################################################
void parse_command_line(int argc, char *argv[],
                        charString &in_file, charString &res_file)
{

   // first input must be input file name
   if (argc < 2) print_usage(argc,argv);
   if (argv[1][0] == '-') {
      std::cerr << "first command line argument must be file name\n";
      print_usage(argc,argv);
   };
   in_file = argv[1];
}

// ###################################################################
// ##### writePlotFile
// ###################################################################

// -------------------------------------------------------------
void
writePlotFile(Grid* grid, Geometry& geom, int nstep, Real time, 
              std::string& plot_file_root)
{
  std::string pltfile = plot_file_root;
  char buf[sizeof(int)+1];
  sprintf(buf,"%04d",nstep);
  pltfile += buf;

  if (ParallelDescriptor::IOProcessor()) 
    std::cout << "Writing plotfile: " << pltfile << std::endl;

  // create directory
  if (!(BoxLib::UtilCreateDirectory(pltfile,0755))) {
   BoxLib::CreateDirectoryFailed(pltfile);
  }  

  std::string HeaderFileName = pltfile + "/Header";
  
  VisMF::IO_Buffer io_buffer(VisMF::IO_Buffer_Size);
  
  std::ofstream HeaderFile;
  
  HeaderFile.rdbuf()->pubsetbuf(io_buffer.dataPtr(), io_buffer.size());
  
  int old_prec;
  
  HeaderFile.open(HeaderFileName.c_str(), std::ios::out|std::ios::trunc);
  
  if (!HeaderFile.good())
    BoxLib::FileOpenFailed(HeaderFileName);

  old_prec = HeaderFile.precision(15);
  
  static const std::string RunstatString("write_pltfile");

  writePlotFile(grid,geom,pltfile,HeaderFile,nstep,time);
  
  HeaderFile.precision(old_prec);
  
  if (ParallelDescriptor::IOProcessor() && !HeaderFile.good())
    BoxLib::Error("main::writePlotFile() failed");
}

// ###################################################################
// ##### writePlotFile
// ###################################################################

void writePlotFile(Grid* grid, Geometry& geom, std::string& dir, 
                   std::ostream& os, int nstep, Real time)
{
  int numPlotComp;

  const Real* dx = geom.CellSize();

  int i;
  int finest_level = 0;
  int n_state  = grid->NumState();
  int n_derive = grid->NumDerive();

  numPlotComp = n_state + 1 + BL_SPACEDIM + n_derive;

  if (ParallelDescriptor::IOProcessor()) 
  {
    // plotfile type
    os << thePlotFileType() << '\n';
  
    // number of components 
    os << numPlotComp << '\n';
 
    // names
    os << "x_velocity"   << '\n';
    os << "y_velocity"   << '\n';
#if (BL_SPACEDIM == 3)
    os << "z_velocity"   << '\n';
#endif
    os << "density" << '\n';
    os << "tracer"  << '\n';
    //  IF YOU WISH TO ADD A NEW SCALAR VARIABLE, UNCOMMENT THIS LINE 
    //  os << "new_scalar"  << '\n';
  
    os << "pressure"<< '\n';
    os << "px"      << '\n';
    os << "py"      << '\n';
#if (BL_SPACEDIM == 3)
    os << "pz"      << '\n';
#endif
    os << "vort"    << '\n';

    static const std::string MultiFabBaseName("/MultiFab"); 

    // dimensions
    os << BL_SPACEDIM << '\n';
    // current time
    os << time << '\n';
    // finest_level
    os << finest_level << '\n';
    // Problem domain in real coordinates
    for (i = 0; i < BL_SPACEDIM; i++) os << geom.ProbLo(i) << ' ';
    os << '\n';
    for (i = 0; i < BL_SPACEDIM; i++) os << geom.ProbHi(i) << ' '; 
    os << '\n';
    // Refinement ratio
    os << '\n';
    // Problem domain in index space
    os << geom.Domain() << ' ';
    os << '\n';
  
    os << nstep << '\n';
    for (int k = 0; k < BL_SPACEDIM; k++) 
         os << dx[k] << ' ';
    os << '\n';
    // Coordinate system
    os << (int) CoordSys::Coord() << '\n';
    // Write bndry data
    os << "0\n";
  }
  
  int thisLevel = 0;
  int numGridsOnLevel = grids.size();
  int IAMRSteps = 0;

  // 
  // Build the directory to hold the MultiFab at this level.
  // The name is relative to the directory containing the Header file.
  // 
  static const std::string BaseName = "/Cell";
  char buf[64];
  sprintf(buf, "Level_%d",thisLevel);
  std::string Level = buf;

  // 
  // Now for the full pathname of that directory.
  // 
  std::string FullPath = dir;
  if (!FullPath.empty() && FullPath[FullPath.length()-1] != '/')
    FullPath += '/';
  FullPath += Level;
  
  // 
  // Only the I/O processor makes the directory if it doesnt already exist.
  // 
  if (ParallelDescriptor::IOProcessor())
    if (!BoxLib::UtilCreateDirectory(FullPath, 0755))
      BoxLib::CreateDirectoryFailed(FullPath);
  // 
  // Force other processors to wait until the directory is built.
  // 
  ParallelDescriptor::Barrier();

  if (ParallelDescriptor::IOProcessor())
  {
    os << thisLevel << ' ' << numGridsOnLevel << ' ' << time << '\n';
    os << IAMRSteps << '\n';
    for (int i = 0; i < grids.size(); ++i)
    {
      for (int n = 0; n < BL_SPACEDIM; n++) 
      { 
        Real grid_lo = (grids[i].loVect()[n]  ) * dx[n];
        Real grid_hi = (grids[i].hiVect()[n]+1) * dx[n];
        os << grid_lo << ' ' << grid_hi << '\n';
      }
    }
    //
    // The full relative pathname of the MultiFabs at this level.
    // The name is relative to the Header file containing this name.
    // It's the name that gets written into the Header.
    //
    if (numPlotComp > 0)
    {
        std::string PathNameInHeader = Level;
        PathNameInHeader += BaseName;
        os << PathNameInHeader << '\n';
    }
  }

  int nGrow = 0;
  MultiFab plotMF(grids,numPlotComp,nGrow);

  grid->writePlotFile(plotMF,time);

  //
  // Use the Full pathname when naming the MultiFab.
  //
  std::string TheFullPath = FullPath;
  TheFullPath += BaseName;
  VisMF::Write(plotMF,TheFullPath,VisMF::OneFilePerCPU);
}


std::string
thePlotFileType()
{
    static const std::string the_plot_file_type("HyperCLaw-V1.1");
    return the_plot_file_type;
}

void checkPoint(Grid* grid, Geometry geom, int nstep, Real time, Real dt,
                std::string& check_file_root)
{
    //
    // In checkpoint files always write out FABs in NATIVE format.
    //
    FABio::Format thePrevFormat = FArrayBox::getFormat();

    FArrayBox::setFormat(FABio::FAB_NATIVE);

    Real dCheckPointTime0 = ParallelDescriptor::second();

    std::string chkfile = check_file_root;
    const std::string ckfile = BoxLib::Concatenate(chkfile,nstep);
    //
    // Only the I/O processor makes the directory if it doesn't already exist.
    //
    if (ParallelDescriptor::IOProcessor())
        if (!BoxLib::UtilCreateDirectory(ckfile, 0755))
            BoxLib::CreateDirectoryFailed(ckfile);
    //
    // Force other processors to wait till directory is built.
    //
    ParallelDescriptor::Barrier();

    std::string HeaderFileName = ckfile + "/Header";

    VisMF::IO_Buffer io_buffer(VisMF::IO_Buffer_Size);

    std::ofstream HeaderFile;

    HeaderFile.rdbuf()->pubsetbuf(io_buffer.dataPtr(), io_buffer.size());

    int old_prec;

    if (ParallelDescriptor::IOProcessor())
    {
        std::cout << "Writing checkpoint file at step " << nstep << std::endl;
        //
        // Only the IOProcessor() writes to the header file.
        //
        HeaderFile.open(HeaderFileName.c_str(), std::ios::out|std::ios::trunc);

        if (!HeaderFile.good())
            BoxLib::FileOpenFailed(HeaderFileName);

        old_prec = HeaderFile.precision(30);

        HeaderFile << CheckPointVersion << '\n'
                   << BL_SPACEDIM       << '\n'
                   << time      <<        '\n';
        HeaderFile << geom      <<        '\n';
        HeaderFile << dt        << ' ' << '\n';
        HeaderFile << nstep     << ' ' << '\n';
    }

    grid->checkPoint(ckfile,HeaderFile);

    if (ParallelDescriptor::IOProcessor())
    {
        HeaderFile.precision(old_prec);

        if (!HeaderFile.good())
            BoxLib::Error("Amr::checkpoint() failed");
    }
    //
    // Don't forget to reset FAB format.
    //
    FArrayBox::setFormat(thePrevFormat);

    const int IOProc     = ParallelDescriptor::IOProcessorNumber();
    Real dCheckPointTime = ParallelDescriptor::second() - dCheckPointTime0;

    ParallelDescriptor::ReduceRealMax(dCheckPointTime,IOProc);

    if (ParallelDescriptor::IOProcessor())
        std::cout << "checkPoint() time = " << dCheckPointTime << " secs." << std::endl;
}
