#!/usr/bin/ruby -Ke
# -*- coding: euc-jp -*-
#
# MPI でコア毎に出力されたファイルをひとまとめにするためのスクリプト. 
# 設定ファイルより, MPI のコア数とファイル名の接頭詞を読み込む. 
# なお, 出力先は, 時刻名のディレクトリである. 
#
# USAGE:
#
#   $ ruby arare_unite.rb <設定ファイル> <分割数>
#
# 想定するディレクトリ構造
#   ./arare.conf         設定ファイル
#   ./data/              生データ置き場
#          VelX_rank000000.nc
#          VelX_rank000001.nc
#          VelX_rank000002.nc
#          ....
#   ./TIME_XXXX-XXXXX/   結合したファイル置き場. 
#          VelX.nc


require "numru/ggraph"
include NumRu

conf    = ARGV[0]
dir     = ""
maxsize = 512 * 512 * 250 * 8

###
### 設定ファイルを読み込む. ファイル名の接頭詞と並列数を取り出す.
###
xprocs = 1
yprocs = 1

conf = open( ARGV[0] )
while line = conf.gets
  line.chomp!
  if (/FilePrefix/ =~ line)
    prefix = line.split("=")[1]
    prefix = prefix.gsub("\"","").gsub(" ","")
#    print prefix
  elsif (/xsub/ =~ line)
    xprocs = line.split("=")[1].to_i
#    print xprocs
  elsif (/ysub/ =~ line) 
    yprocs = line.split("=")[1].to_i
#    print yprocs
  end
end
conf.close


###
### 変数毎にファイルを取り出してまとめる. 1 変数 1 ファイルを仮定. 
###
Dir.glob("data/#{prefix}*rank000000.nc").sort.each{|nc|
  
  # ファイル名から変数名を判断
  tmp = nc.gsub("data/#{prefix}", "")
  /^([A-Z].+)_rank000000.nc$/ =~ tmp 

  # 変数に対する処理.
  if $1
    var = $1.to_s
    p var

    ####
    #### rank 0 のファイルをオープンして座標の情報を取り出した後に, 
    #### ファイル名の配列を作り, オープンし, 出力する. 
    ####
    gphys0 = GPhys::IO.open( nc, var )  
    tmp2 = gphys0.coord(0).name

    if tmp2 == "x"

      # ファイル名の配列を作成. 
      files = Array.new
      yprocs.times{|j|
        files[j] = Array.new      
        xprocs.times{|i|
          k = xprocs * yprocs - (j + 1) * xprocs + i
          files[j].push( "data/#{prefix}#{var}_rank#{sprintf("%06d", k)}.nc")
        }
      }
      
      # ファイルのオープン
      if yprocs == 1 
        gphys = GPhys::IO.open( files[0], var )
      else
        gphys = GPhys::IO.open( files, var )
      end

      # 軸情報
      x  = gphys.coord(0).val
      y  = gphys.coord(1).val
      z  = gphys.coord(2).val
      t  = gphys.coord(3).val

      # 前回の計算に当該時刻が含まれていないかのチェック
      edir = Dir.glob("TIME_*").sort.pop
      if /#{sprintf('%08d', t.min.to_i)}/ =~ edir
        t = t[1..t.size-1]
      end

      # 1 ファイルに入れる時間方向のグリッド数を決める
      xyzsize = x.size * y.size * z.size
      tsize = maxsize / xyzsize 
      taxis = t.size
      div   = taxis / tsize  + 1
      p "div: #{div}; n: #{taxis / div}"

      # 2 GB の壁問題があるので, 分割出力する
      #
      if div > 1
        
        n = taxis / div  # 1 ファイルあたりのデータ数
        
        # 割り切れないときのための処理. 
        if ( n % div > 0 )           
          n = n + 1 
        end

        div.times{|pp|
          n1 = n * pp
          n2 = [n * ( pp + 1 ) - 1, t.size - 1].min
#          p "#{pp}, #{n1} -> #{n2}"

          # ディレクトリ作成
          dir = "TIME_#{sprintf('%08d',t[n1])}-#{sprintf('%08d',t[n2])}"
          Dir::mkdir( dir ) unless FileTest.exist?( dir )

          # ファイル書き出し
          newfile = "#{dir}/#{prefix}#{var}.nc"
          unless FileTest.exist?( newfile )
            outfile = NetCDF.create(newfile)
            GPhys::NetCDF_IO.write( outfile, gphys[true,true,true,n1..n2] ) 
            outfile.close
            p "time : #{t[n1]}--#{t[n2]}"
            p "WRITE: #{newfile}"
          end
        }

      else

        # ディレクトリ作成
        dir = "TIME_#{sprintf('%08d',t.min)}-#{sprintf('%08d',t.max)}"
        Dir::mkdir( dir ) unless FileTest.exist?( dir )

        # ファイルの書き出し. 
        newfile = "#{dir}/#{prefix}#{var}.nc"
        unless FileTest.exist?( newfile )
          outfile = NetCDF.create(newfile)
          GPhys::NetCDF_IO.write( outfile, gphys ) unless FileTest.exist?( outfile )
          outfile.close
          p "WRITE: #{newfile}"
        end
      end

#      files.flatten.each{|file|
#        dir1 = "#{dir}/mpi"
#        Dir::mkdir( dir1 ) unless FileTest.exist?( dir1 )
#        p "mv #{file} #{dir1}" 
#        system( "mv #{file} #{dir1}" )
#      }


    else
      p "#{nc} does not have x-direction, SKIP! "
    end
  else
    p "SKIP: #{tmp}"
  end  
}