#!/usr/bin/env ruby

# ----------------------------------------------
# local load path

$local_path = '/home/yukiko/work/ape/yukiko/lib' 
$: << $local_path

# ----------------------------------------------
# 必要なライブラリ, モジュールの読み込み

load "#{$local_path}/ape-view.rb"
load "#{$local_path}/lib-ape-view.rb"
load "#{$local_path}/lib-ape-base.rb"
load "#{$local_path}/lib-ape-stspctrum.rb"

# --------------------------------------------------------------------


END{

  $sstid = "control"
  $resol = "T319L48_non"
  set_directry
  mknc_stspctrum

}

def set_directry
  $groupid = $groupid_hash[$resol]
  $ncfile_path = "/home/yukiko/work/ape/NetCDF/#{$resol}/"
  $file_label = "#{$resol}_#{$expID}"
end

def mknc_stspctrum

  name_array = [
#    ["tr_tppn","precipitation_flux","kg m-2 s-1"] #, 
#    ["tr_lw_toa","toa_outgoing_longwave_flux","W m-2"] #, 
#    ["tr_om500","omega","Pa s-1"] #, 
#    ["tr_u250","eastward_wind","m s-1"] #, 
#    ["tr_v250","northward_wind","m s-1"] #, 
    ["tr_mslp","air_pressure_at_sea_level","Pa"]
  ]

  $name_array = name_array
    
  file = NetCDF.create("#{$resol}_#{$sstid}_wk_40smooth.nc")
#  file = NetCDF.create("hoge.nc")
  puts "#{$resol}_#{$sstid}_wk_40smooth.nc create start"

#  $t= ape_new "#{$ncfile_path}AGUforAPE-03a_TR_control_0031.nc"    
  $t= ape_new "hoge.nc"    

  name_array.size.times { |num| 

####    file = NetCDF.create("#{$resol}_#{$sstid}_wk_40smooth_#{name_array[num][2]}.nc")
####    puts "#{$resol}_#{$sstid}_wk_40smooth_#{name_array[num][2]}.nc create start"

#   g = $t.go(name_array[num][0]).cut(true,-15..15,true)
#    gphys = t319l48_non_opn(g,-15..15)
    
    gphys = $t.go(name_array[num][0])

    gphys_sym, gphys_asym, gphys_sym_wk, gphys_asym_wk, gphys_bg = 
      spct_wk1999(gphys)
    
    print "ncfile write #{Time.now} \n"
    # netCDF 初期化
    GPhys::NetCDF_IO.write(file, gphys_sym_wk )
    GPhys::NetCDF_IO.write(file, gphys_asym_wk )
    GPhys::NetCDF_IO.write(file, gphys_bg)
    GPhys::NetCDF_IO.write(file, gphys_sym)
    GPhys::NetCDF_IO.write(file, gphys_asym)

#    GPhys::NetCDF_IO.write(file, gphys )

   
  # 大域属性
  file.put_att("Conventions", "CF-1.0")
  file.put_att("title","Aqua Planet: #{$resol}_#{$sstid} Experiment, Wheeler and Kiladis, 1999 Plot")
  file.put_att("history", "Original data produced: #{Time.now} from GrADS by Yukiko YAMADA (AGU for APE)")
  file.put_att("institution", "#{$resol}")
  file.close


  }

  
end

# --------------------------------------------------------------------

def t319l48_non_opn(g,dimcut)

    grid_0    = g.grid_copy.axis(0)
    grid0size = grid_0.pos.val.size

#    grid_1    = g.cut(true,dimcut,true).grid_copy.axis(1)
    grid_1    = g.grid_copy.axis(1)
    grid1size = grid_1.pos.val.size


    timesize = 360*4
    grid_time =
        Axis.new().
        set_pos(VArray.new(NArray.sfloat(timesize).indgen!/4+ 0.25).rename("time").
      put_att("units","days"))

    data =  NArray.sfloat(grid0size,grid1size,timesize).fill!(0.0)
    grid = Grid.new(grid_0, grid_1, grid_time)

      data[true,true,(-30*4*12)..(-30*4*11-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0031.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*11)..(-30*4*10-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0032.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*10)..(-30*4*9-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0033.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*9)..(-30*4*8-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0034.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*8)..(-30*4*7-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0035.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*7)..(-30*4*6-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0036.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*6)..(-30*4*5-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0037.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*5)..(-30*4*4-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0038.nc",
                            g.name).cut(true,dimcut,true).data.val

      data[true,true,(-30*4*4)..(-30*4*3-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0039.nc",
                            g.name).cut(true,dimcut,true).data.val
      data[true,true,(-30*4*3)..(-30*4*2-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0040.nc",
                            g.name).cut(true,dimcut,true).data.val
      data[true,true,(-30*4*2)..(-30*4-1)] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0041.nc",
                            g.name).cut(true,dimcut,true).data.val
      data[true,true,(-30*4)..-1] =
        GPhys::NetCDF_IO.open(
                            "#{$ncfile_path}AGUforAPE-03a_TR_control_0042.nc",
                            g.name).cut(true,dimcut,true).data.val


    data = VArray.new(data).rename(g.name)
    gphys = GPhys.new(grid, data)
#    gphys = gphys.set_att("ape_name",g.get_att("ape_name")). 
#                  set_att("units",g.get_att("units")). 
#                  set_lost_axes(g.lost_axes)

    return gphys

end
  

# --------------------------------------------------------------------

def spct_wk1999(gphys)

  gphys1 = gphys.cut(true,-15..0,true)
  gphys2 = gphys.cut(true,0..15,true)

#  gphys1 = gphys.cut(true,-5..0,true)
#  gphys2 = gphys.cut(true,0..5,true)

  work1 = gphys1[true,0,0..359].stspct_fft("spct").cut(-59..59,true) * 0.0
  work2 = gphys2[true,0,0..359].stspct_fft("spct").cut(-59..59,true) * 0.0

  dim0 = gphys1.coord(0).shape.to_s.to_i  # 緯度
  dim1 = gphys1.coord(1).shape.to_s.to_i  # 経度

  print "least-square + taper start #{Time.now} \n"


  g1 = gphys1[true,true,0..359].copy
  g2 = gphys2[true,true,0..359].copy

  # 最小二乗法 + hamming taper + 時空間スペクトル足しあわせ
  10.times { |num|

    print "#{num} #{Time.now}\n"

    g1[true,true,true] = gphys1[true,true,(num*120)..359+(num*120)].data.val
    g2[true,true,true] = gphys2[true,true,(num*120)..359+(num*120)].data.val

    x   = NArray.sfloat(dim0,dim1).fill!(0.0)
    xx  = NArray.sfloat(dim0,dim1).fill!(0.0)
    xy1 = NArray.sfloat(dim0,dim1).fill!(0.0)
    xy2 = NArray.sfloat(dim0,dim1).fill!(0.0)
    y1 =  NArray.sfloat(dim0,dim1).fill!(0.0)
    y2 =  NArray.sfloat(dim0,dim1).fill!(0.0)
    
    time =  NArray.sfloat(360).indgen!
    
    time.each{ |time_num|

      xy1 = time_num * g1.data.val[true,true,time_num] + xy1
      xy2 = time_num * g2.data.val[true,true,time_num] + xy2
      x  = time_num + x  
      y1  = g1.data.val[true,true,time_num] + y1
      y2  = g2.data.val[true,true,time_num] + y2
      xx = time_num + xx
      
    }

    a1 = (360 * xy1 - x * y1  ) / (360 * xx - x**2 )
    a2 = (360 * xy2 - x * y2  ) / (360 * xx - x**2 )
    b1 = ( xx * y1  - x * xy1 ) / (360 * xx - x**2 )
    b2 = ( xx * y2  - x * xy2 ) / (360 * xx - x**2 )

    time.each{ |time_num|
      g1[true,true,time_num] = ( g1.data.val[true,true,time_num] - ( a1 * time_num + b1 ) ) * 
	( 1.0 - cos(2*Math::PI*(time_num - 0.5)/360.0) ) / 2.0 
      g2[true,true,time_num] = ( g2.data.val[true,true,time_num] - ( a2 * time_num + b2 ) ) * 
	( 1.0 - cos(2*Math::PI*(time_num - 0.5)/360.0) ) / 2.0 
    }

    dim1.times { |dim1_num|
      work1 = g1[true,dim1_num,true].stspct_fft("spct").
                      cut(-59..59,true) + work1
      work2 = g2[true,dim1_num,true].stspct_fft("spct").
                      cut(-59..59,true) + work2
    }

  }

  gphys_asym = (( work1 - work2 )/2 ).
    set_att("ape_name", "#{$name_array[0][1]} S-T power spectrum").
    set_att("units", "1" ). 
    rename("#{gphys.name}_asym_org_spct")
  gphys_sym  = (( work1 + work2 )/2 ).
    set_att("ape_name", "#{$name_array[0][1]} S-T power spectrum").
    set_att("units", "1" ). 
    rename("#{gphys.name}_sym_org_spct")
  name = gphys.name

  gphys = ( gphys_sym + gphys_asym ).rename("#{gphys.name}")

  gphys_bg  = spct_121mean(gphys)

  gphys_sym_wk = spct_divide(gphys_sym,gphys_bg).
    set_att("ape_name", $name_array[0][1]).
    set_att("units","1"). 
    rename("#{name}_sym_spct")
  gphys_asym_wk = spct_divide(gphys_asym,gphys_bg).
    set_att("ape_name", $name_array[0][1] ).
    set_att("units","1"). 
    rename("#{name}_asym_spct")

  gphys_bg = gphys_bg.rename("#{gphys.name}_bg_spct").
    set_att("ape_name", $name_array[0][1] ).
    set_att("units", "1" )

  return gphys_sym, gphys_asym, gphys_sym_wk, gphys_asym_wk, gphys_bg
end


def spct_wk1999_ichimai(gphys)
  
  gphys = gphys.cut(true,0,true)[true,-401..-1]
  gphys = gphys.stspct_fft("spct")

  gphys_sym  = gphys.
    set_att("ape_name", "#{gphys.data.get_att("ape_name")} S-T power spectrum").
    set_att("units", "1" ). 
    rename("#{gphys.name}_asym_org_spct")
  gphys_bg  = spct_121mean(gphys_sym)

  gphys_sym_wk = spct_divide(gphys_sym,gphys_bg).
    set_att("ape_name", gphys.data.get_att("ape_name")).
    set_att("units","1"). 
    rename("#{gphys.name}_sym_spct")

  gphys_bg = gphys_bg.rename("#{gphys.name}_bg_spct").
    set_att("ape_name", gphys.data.get_att("ape_name")).
    set_att("units", "1" )

  gphys_asym = 0.0
  gphys_asym_wk = 0.0

  return gphys_sym, gphys_asym, gphys_sym_wk, gphys_asym_wk, gphys_bg
end

def spct_121mean(gphys)
  
  dim1 = gphys.coord(0).shape.to_s.to_i
  dim2 = gphys.coord(1).shape.to_s.to_i
  
  gphys_bg = gphys.copy
  work     = gphys.copy

  print "1-2-1 mean start #{Time.now} \n"

  40.times{ |num|

    print "#{num} #{Time.now} \n"
#  80.times{ |num|
    dim1.times { |dim1_num|
      if dim1_num == 0
	work[dim1_num,true] = 
	  ( gphys_bg[dim1_num,true] + gphys_bg[dim1_num+1,true] )/2 
      elsif dim1_num == (dim1-1)
	work[dim1_num,true] = 
	  ( gphys_bg[dim1_num-1,true] + gphys_bg[dim1_num,true] )/2 
      else
	work[dim1_num,true] = 
	  ( gphys_bg[dim1_num-1,true]+ gphys_bg[dim1_num,true] + gphys_bg[dim1_num,true] + gphys_bg[dim1_num+1,true] )/4 
      end
    }
    
    dim2.times { |dim2_num|
      if dim2_num == 0
	gphys_bg[true, dim2_num] = 
	  ( work[true,dim2_num] + work[true,dim2_num+1] )/2 
      elsif dim2_num == (dim2-1)
	gphys_bg[true, dim2_num] = 
	  ( work[true,dim2_num-1] + work[true,dim2_num] )/2 
      else
	gphys_bg[true, dim2_num] = 
	  ( work[true,dim2_num-1]+ work[true,dim2_num] + work[true,dim2_num] +  work[true,dim2_num+1] )/4 
      end
    }
  }   
  return gphys_bg

end

def spct_divide(gphys,gphys_bg)  

  dim1 = gphys.coord(0).shape.to_s.to_i
  dim2 = gphys.coord(1).shape.to_s.to_i

  dim1.times{ |dim1_num|
    dim2.times{ |dim2_num|
      unless gphys_bg[dim1_num,dim2_num].data.val == 0.0
	gphys[dim1_num,dim2_num]=
	  gphys[dim1_num,dim2_num].data.val/gphys_bg[dim1_num,dim2_num].data.val
      else 
	gphys[dim1_num,dim2_num] = 0.0 
      end
    }
  }

  return gphys
end