#!/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{

#$rezol = ["T39L48_eml",
#  "T39L24_eml",
#  "T39L96_eml",
#  "T79L48_eml" ]

# $rezol = ["T79L48_non"]
#$rezol = ["T39L48_eml"]
#$resol = "T39L48_eml"
#$resol = "T159L48_non"
#$resol = "T319L48_non"

  
# ["T39L48_ias","T39L48_ksc","T39L48_kuo","T39L48_mca","T39L48_non",
#    "T39L24_eml","T39L96_eml","T39L24_non","T39L96_non",
#    "T79L48_eml"].each{ |resol|
#
#    $sstid = "control"
#    $resol = resol
#    mknc_stspctrum
#  }

    $sstid = "control"
    $resol = "T39L48_ias"
    mknc_stspctrum

}


def mknc_stspctrum

  name_array = [
    ["T",0.254,"tr_t250","air_temperature","K"], 
    ["T",0.166,"tr_t170","air_temperature","K"], 
    ["T",0.566,"tr_t570","air_temperature","K"], 
    ["Q",0.947,"tr_q950","specific_humidity","1"], 
    ["Q",0.852,"tr_q850","specific_humidity","1"], 
    ["Q",0.755,"tr_q760","specific_humidity","1"], 
    ["U",0.254,"tr_u250","eastward_wind","m s-1"], 
    ["U",0.166,"tr_u170","eastward_wind","m s-1"], 
    ["V",0.254,"tr_v250","northward_wind","m s-1"], 
    ["V",0.166,"tr_v170","northward_wind","m s-1"], 
    ["Z",0.254,"tr_phi250","geopotential_height","m"], 
    ["Z",0.166,"tr_phi170","geopotential_height","m"]
  ]

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

  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"
    gphys = __grads_opn(name_array[num])
    
    #  puts name
    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)

  }

  
  # 大域属性
  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 __grads_opn(name_array)

  name = name_array[0]
  layer = name_array[1]
  rename = name_array[2]
  apename = name_array[3]
  units = name_array[4]
  
  print "#{name}, #{layer}\n"

  file_path = "/home/yukiko/work/ape/GrADS/#{$resol}_expID01/"
  
  if $resol =~ /T159L48/
 
    file1 = Ape.new("#{file_path}/0011/#{name}.ctl")
    file2 = Ape.new("#{file_path}/0012/#{name}.ctl")
    file3 = Ape.new("#{file_path}/0013/#{name}.ctl")
    file4 = Ape.new("#{file_path}/0014/#{name}.ctl")
    
    #  tr_tppn
    gphys = file1.gphys_open(name).cut(true,-15..15,layer,true)
    
    # 軸
    t_size = 
      NArray.sfloat(gphys.grid_copy.axis(2).pos.val.size*4).indgen*1.0/4.0 + 0.25
    y_size = gphys.grid_copy.axis(1).pos.val
    x_size = gphys.grid_copy.axis(0).pos.val

    # data
    data = NArray.sfloat(x_size.size,y_size.size,t_size.size).indgen * 0.0

    gphys = file1.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,0..(t_size.size/4-1)] = gphys.data.val.to_na

    gphys = file2.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,(t_size.size/4)..(t_size.size/2-1)] = gphys.data.val.to_na

    gphys = file3.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,(t_size.size/2)..(t_size.size*3/4-1)] = gphys.data.val.to_na

    gphys = file4.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,(t_size.size*3/4)..-1] = gphys.data.val.to_na

    data = VArray.new(data).rename(rename).
      put_att("units",units).
      put_att("ape_name",apename)

    x_name = gphys.grid_copy.axis(0).name
    y_name = gphys.grid_copy.axis(1).name
    x_units = gphys.grid_copy.axis(0).pos.get_att("units")
    y_units = gphys.grid_copy.axis(1).pos.get_att("units")

    x = VArray.new(x_size).rename(x_name).put_att("units",x_units)
    y = VArray.new(y_size).rename(y_name).put_att("units",y_units)
    t = VArray.new(t_size).rename("time").put_att("units","days")
    x = Axis.new().set_pos(x)
    y = Axis.new().set_pos(y)
    t = Axis.new().set_pos(t)
    grid = Grid.new(x,y,t)

    data = GPhys.new( grid, data)
  
  elsif $resol =~ /T39L/ || $resol =~ /T79L/

    file1 = Ape.new("#{file_path}/0006/#{name}.ctl")
    file2 = Ape.new("#{file_path}/0007/#{name}.ctl")
    
    #  tr_tppn
    gphys = file1.gphys_open(name).cut(true,-15..15,layer,true)
    
    # 軸
    t_size = 
      NArray.sfloat(gphys.grid_copy.axis(2).pos.val.size*2).indgen*1.0/4.0 + 0.25
    y_size = gphys.grid_copy.axis(1).pos.val
    x_size = gphys.grid_copy.axis(0).pos.val

    # data
    data = NArray.sfloat(x_size.size,y_size.size,t_size.size).indgen * 0.0

    gphys = file1.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,0..(t_size.size/2-1)] = gphys.data.val.to_na

    gphys = file2.gphys_open(name).cut(true,-15..15,layer,true)
    data[true,true,(t_size.size/2)..-1] = gphys.data.val.to_na


    data = VArray.new(data).rename(rename).
      put_att("units",units).
      put_att("ape_name",apename)

    x_name = gphys.grid_copy.axis(0).name
    y_name = gphys.grid_copy.axis(1).name
    x_units = gphys.grid_copy.axis(0).pos.get_att("units")
    y_units = gphys.grid_copy.axis(1).pos.get_att("units")

    x = VArray.new(x_size).rename(x_name).put_att("units",x_units)
    y = VArray.new(y_size).rename(y_name).put_att("units",y_units)
    t = VArray.new(t_size).rename("time").put_att("units","days")
    x = Axis.new().set_pos(x)
    y = Axis.new().set_pos(y)
    t = Axis.new().set_pos(t)
    grid = Grid.new(x,y,t)

    data = GPhys.new( grid, data)

  end

  return data
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)

  dim = gphys1.coord(1).shape.to_s.to_i  
  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"

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

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

    g1 = gphys1[true,true,(num*120)..359+(num*120)].copy
    g2 = gphys2[true,true,(num*120)..359+(num*120)].copy
    	
    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", "#{gphys.data.get_att("ape_name")} S-T power spectrum").
    set_att("units", "1" ). 
    rename("#{gphys.name}_asym_org_spct")
  gphys_sym  = (( work1 + work2 )/2 ).
    set_att("ape_name", "#{gphys.data.get_att("ape_name")} S-T power spectrum").
    set_att("units", "1" ). 
    rename("#{gphys.name}_sym_org_spct")
  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", gphys.data.get_att("ape_name")).
    set_att("units","1"). 
    rename("#{gphys.name}_sym_spct")
  gphys_asym_wk = spct_divide(gphys_asym,gphys_bg).
    set_att("ape_name", gphys.data.get_att("ape_name") ).
    set_att("units","1"). 
    rename("#{gphys.name}_asym_spct")

  gphys_bg = gphys_bg.rename("#{gphys.name}_bg_spct").
    set_att("ape_name", gphys.data.get_att("ape_name")).
    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