#!/usr/bin/env ruby

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

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

require "getopts"
require "numru/netcdf"
module NumRu
  class NetCDFVar
    alias put scaled_put
    alias get scaled_get
  end
end

require "numru/ggraph"
require "colorbar"
require 'rational'
require 'complex'
require 'mathn'
include NumRu
include NMath


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

END{

#equivdepth_calc
#bunsan_calc
#bunsan_plot2

eqivdepth_plot

}

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

def GGraph.annotate(str_ary) # GGraph の annotate メソッド再定義

  charsize = 0.7 * DCL.uzpget('rsizec1')
  dvx = 0.01
  dvy = charsize*1.5
  raise TypeError,"Array expected" if ! str_ary.is_a?(Array)
  vxmin,vxmax,vymin,vymax = DCL.sgqvpt
#  vx = 0.66 
  vx = 0.50
#  vx = 0.40 
  vy = 0.12 - charsize/2
  str_ary.each{|str|
    DCL::sgtxzv(vx,vy,"[ #{str} ]",charsize,0,-1,1)
    vy -= dvy
  }
  nil
end


class GPhys

  # 座標, データ操作ログ記録関数の定義 (Grid.lost_axes 参照)
  def add_lost_axes( lost )
    GPhys.new( @grid.copy.add_lost_axes( lost.to_a ), @data)
  end
  def set_lost_axes( lost )
    GPhys.new( @grid.copy.set_lost_axes( lost.to_a ), @data)
  end

  # rename 再定義 (gphys class の値を返す)
  def rename(name)
    GPhys.new(@grid,@data.copy.rename(name))
  end

  # attribute 変更再定義 (gphys class の値を返す)
  def set_att(name,val)
    GPhys.new(@grid,@data.copy.set_att(name,val))
  end

end

def eqivdepth_plot

  eta1 = NArray.sfloat(76).indgen * 0.02
  x1, x2, x3  = equivdepth_calc(eta1)

  # 軸
  eta1 = VArray.new(eta1).rename("eta_1").
    put_att("units","1").put_att("ape_name","eta1")
  eta1 = Axis.new().set_pos(eta1)
  eta1 = Grid.new(eta1)

  # 値  
  x = Array.new

  x1r = VArray.new(x1.real.abs).rename("c_g_1_real").
    put_att("units","1").put_att("ape_name","c_g_1_real")
  x.push( GPhys.new( eta1, x1r) )

  x1i = VArray.new(x1.imag).rename("c_g_1_imag").
    put_att("units","1").put_att("ape_name","c_g_1_imag")
  x.push( GPhys.new( eta1, x1i) )

  x2r = VArray.new(x2.real.abs).rename("c_g_1_real").
    put_att("units","1").put_att("ape_name","c_g_1_real")
  x.push( GPhys.new( eta1, x2r) )

  x2i = VArray.new(x2.imag).rename("c_g_1_imag").
    put_att("units","1").put_att("ape_name","c_g_1_imag")
  x.push( GPhys.new( eta1, x2i) )


  x3r = VArray.new(x3.real.abs).rename("c_g_1_real").
    put_att("units","1").put_att("ape_name","c_g_1_real")
  x.push( GPhys.new( eta1, x3r) )

  x3i = VArray.new(x3.imag).rename("c_g_1_imag").
    put_att("units","1").put_att("ape_name","c_g_1_imag")
  x.push( GPhys.new( eta1, x3i) )


  # 描画

  DCL::gropn(1)
  GGraph.set_fig( 'viewport'=>[0.13,0.83,0.23,0.58]) 
  DCL.uzfact(0.6)                                   # 文字の大きさ倍
  DCL.sgpset('lcorner',false)    # コーナーマークはつけない
  DCL.slmgn(0.0, 0.0, 0.0, 0.0)  # 描画マージン指定
  DCL.sgpset('lfprop',true)      # プロポーショナルフォントを使う
  DCL.sgpset('lfull',true)       # 全画面表示
  DCL.sgpset('lcntl', false)     # アンダーバーは普通に表示
  DCL.uscset('cyspos', 'B' )      # y 軸単位を書く位置

#  $fig_set_hash = { "window" => [ $x_size[0].real,$x_size[16].real,0.0,0.5 ] }
  $fig_set_hash = { "window" => [ 0,1.5,-0.4,0.4 ] }
  GGraph.set_fig($fig_set_hash)
  
  line_index_ary = Array.new; name_ary = Array.new
  
  name_ary = ["mode_1","mode_2","mode_3" ]
  line_index_ary = [92,22,12]

  sign = "YU.YAMADA #{Time.now.strftime("%Y/%m/%d %H:%M:%S JST")}"
  $lost_axis = [
    sign,
    "rezol: z=4", 
    "linear_heat, linear_equation", 
    "eta1 = -eta2"
  ]
  x[0] = x[0].set_lost_axes($lost_axis)

  # 描画
  GGraph.line( x[0].rename("C_g"), true, "index" => line_index_ary[2] )
  GGraph.line( x[1].rename("C_g"), false, "index" => line_index_ary[2],
	      'type' => 3 )
  GGraph.line( x[2].rename("C_g"), false, "index" => line_index_ary[1] )
  GGraph.line( x[3].rename("C_g"), false, "index" => line_index_ary[1],
	      'type' => 3 )
  GGraph.line( x[4].rename("C_g"), false, "index" => line_index_ary[0] )
  GGraph.line( x[5].rename("C_g"), false, "index" => line_index_ary[0],
	      'type' => 3 )
  # タイトル
  tani = "(1)"
  DCL::uxsttl("t", tani , -1.0 )
  DCL::uxmttl("t", "phase speed/growth rate", -1.0 ) 

  $file_label = ""
  
  charsize = 0.81 - $file_label.size * 0.0095 
  DCL::sgtxzv(charsize,0.6,$file_label,0.8 * DCL.uzpget('rsizec1'),0,-1,1) 

  line_index(name_ary,line_index_ary)
  
  DCL::grcls

end

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


def equivdepth_calc(eta1)

#  eta1 = 1.5
  eta1 = eta1
  eta2 = -eta1
  eta3 = 0.0
  
  eig = NArray.complex(3,5).fill!(0.0)
  gamma = NArray.complex(3).fill!(0.0)
  eta = NArray.complex(3,76).fill!(0.0)
  
  eta[0,true] = eta1
  eta[1,true] = eta2
  eta[2,true] = eta3
  
  # 固有関数 上下端で 0
  eig[true,0] = 0.0
  eig[true,4] = 0.0 
  3.times{ |num|
    
    # 固有値^0.5 の計算
    gamma[num] = 4.0 * ( 2.0 - 2.0 * cos( (num+1)*PI/(3+1) ) )**0.5
    
    # 固有関数 (規格化してない) 計算
    3.times{ |j|    
      eig[num,j+1] = sin( ((num+1)*(j+1)*PI)/(3+1) )
    }
  }

  3.times{ |num|
    eta[num] = eta[num] / eig[num,1] * gamma[num] / 4  * (2**0.5)  
  }

  h2 = eta[0,true] + eta[1,true] + eta[2,true]
#  h3 = 2**0.5 * (eta[0,true] - eta[2,true])
  h3 = (eta[0,true] - eta[2,true])
  h4 = eta[0,true] - eta[1,true] + eta[2,true]

  b = -(-10.0 + 3*h2 + 2*h3 + h4 )/4
  c = -(-6.0 + 4*h2 + h3 )/4
  d = -(-1.0 + h2)/4

  p = -1.0/3 * b**2 + c
  q = 2.0/27*(b**3) -b*c/3 + d

## 3 次方程式解法: カルダノの公式
## x^3 + 3px + q
## t^2 + qt - p =0, t = -q/2 +- (q**2/4 + p)**0.5

  t1 = NArray.complex(76).fill!(0.0)
  t2 = NArray.complex(76).fill!(0.0)
  
  t1  =  -q/2 + (q**2/4 + p**3/27)**0.5 
  t2  =  -q/2 - (q**2/4 + p**3/27)**0.5
  t1 = (t1)**(1.0/3)
  t2 = (t2)**(1.0/3)
  
  n1 = (t1 + t2)
  n2 = (t2 * Complex.polar(1, Math::PI*2/3) +
	t1 * Complex.polar(1, Math::PI*4/3) )
  n3 = (t1 * Complex.polar(1, Math::PI*2/3) +
	t2 * Complex.polar(1, Math::PI*4/3) )
  
  x1 = n1 - b/3
  x2 = n2 - b/3
  x3 = n3 - b/3
  
  x1 = (x1**0.5)*0.25*Complex(0,1)
  x2 = (x2**0.5)*0.25*Complex(0,1)
  x3 = -(x3**0.5)*0.25*Complex(0,1)

#  print "\n x1 = #{x1.real} +- #{x1.imag}i\n"
#  print " x2 = #{x2.real} +- #{x2.imag}i\n"
#  print " x3 = #{x3.real} +- #{x3.imag}i\n\n"


  return x1, x2, x3
end


def equivdepth_calc2


  b = (7-3*(2**0.5))/4
  c = (6-3*(2**0.5)/2)/4
  d = 1.0/4

  p = -1.0/3 * b**2 + c
  q = 2.0/27*(b**3) -b*c/3 + d

## 3 次方程式解法: カルダノの公式
## x^3 + 3px + q
## t^2 + qt - p =0, t = -q/2 +- (q**2/4 + p)**0.5

  p = p/3

  t1  =  -q/2 + (q**2/4 + p**3)**0.5
  t2  =  -q/2 - (q**2/4 + p**3)**0.5

#  t1 = Complex(t1,0)**(1.0/3)
  t1 = (-t1)**(1.0/3)*Complex.polar(1, Math::PI*2/3)
  t2 = (-t2)**(1.0/3)*Complex.polar(1, Math::PI*2/3)

  n1 = (t1 + t2)
  n2 = (t2 * Complex.polar(1, Math::PI*2/3) +
      t1 * Complex.polar(1, Math::PI*4/3) )
  n3 = (t1 *  Complex.polar(1, Math::PI*2/3) +
      t2 * Complex.polar(1, Math::PI*4/3))

  x1 = n1 - b/3
  x2 = n2 - b/3
  x3 = n3 - b/3
  
  print "x1 = #{(x1**0.5)*Complex(0,1)*0.25}\n"
  print "x2 = #{x2**0.5*Complex(0,1)*0.25}\n"
  print "x3 = #{x3**0.5*Complex(0,1)*0.25}\n"

end




def bunsan_plot

  DCL::gropn(1)
  GGraph.set_fig( 'viewport'=>[0.13,0.83,0.23,0.58]) 
  DCL.uzfact(0.6)                                   # 文字の大きさ倍
  DCL.sgpset('lcorner',false)    # コーナーマークはつけない
  DCL.slmgn(0.0, 0.0, 0.0, 0.0)  # 描画マージン指定
  DCL.sgpset('lfprop',true)      # プロポーショナルフォントを使う
  DCL.sgpset('lfull',true)       # 全画面表示
  DCL.sgpset('lcntl', false)     # アンダーバーは普通に表示
  DCL.uscset('cyspos', 'B' )      # y 軸単位を書く位置

#  $fig_set_hash = { "window" => [ $x_size[0].real,$x_size[16].real,0.0,0.5 ] }
  $fig_set_hash = { "window" => [ $x_size[0].real,$x_size[16].real,0.0,2.0 ] }
  GGraph.set_fig($fig_set_hash)
  
  line_index_ary = Array.new; name_ary = Array.new
  
  name_ary = ["kelvin", "mixed", "n=1", "n=2","n=3" ]
  line_index_ary = [12,22,92,32,42]

  sign = "YU.YAMADA #{Time.now.strftime("%Y/%m/%d %H:%M:%S JST")}"
  $lost_axis = [
    sign,
    "rezol: k=l=64, z=4", 
    "linear_heat, linear_equation", 
    "beta=1, eta1 = 1.5, eta2 = -1.5"
  ]
  $kelvin = $kelvin.set_lost_axes($lost_axis)

  # 描画
#  GGraph.line( $kelvin.rename("freq"), true, "index" => line_index_ary[0] )
  GGraph.line( $kelvin.rename("e-folding_time"),true,"index" => line_index_ary[0])
  GGraph.line( $mixed, false, "index" => line_index_ary[1] )
  3.times{ |num|
    GGraph.line( $grav_array[num], false, "index" => line_index_ary[num+2] )
    GGraph.line( $rossby_array[num], false, "index" => line_index_ary[num+2] )
  }
  GGraph.line( $axis0, false, "index" => line_index_ary[0] )

  # タイトル
  tani = "(#{$kelvin.data.get_att("units")})"
  DCL::uxsttl("t", tani , -1.0 )
  DCL::uxmttl("t", "growth rate", -1.0 ) 
#  DCL::uxmttl("t", "frequency", -1.0 ) 

  $file_label = "c_g =  0.167 + 0.082i"
  
  charsize = 0.81 - $file_label.size * 0.0095 
  DCL::sgtxzv(charsize,0.6,$file_label,0.8 * DCL.uzpget('rsizec1'),0,-1,1) 

  line_index(name_ary,line_index_ary)
  
  DCL::grcls
  
end


def bunsan_plot2


#  DCL.swpset('IWIDTH' ,800)
#  DCL.swpset('IHEIGHT' ,600)
  DCL::gropn(-1)
  GGraph.set_fig( 'viewport'=>[0.15,0.55,0.23,0.83]) 
  DCL.uzfact(0.6)                                   # 文字の大きさ倍
  DCL.sgpset('lcorner',false)    # コーナーマークはつけない
  DCL.slmgn(0.0, 0.0, 0.0, 0.0)  # 描画マージン指定
  DCL.sgpset('lfprop',true)      # プロポーショナルフォントを使う
  DCL.sgpset('lfull',true)       # 全画面表示
  DCL.sgpset('lcntl', false)     # アンダーバーは普通に表示
  DCL.uscset('cyspos', 'B' )      # y 軸単位を書く位置

  $fig_set_hash = { "window" => [ $x_size[8].real,$x_size[16].real,0.0,2.0 ] }
#  $fig_set_hash = { "window" => [ $x_size[8].real,$x_size[16].real,0.0,0.5 ] }
  GGraph.set_fig($fig_set_hash)
  
  line_index_ary = Array.new; name_ary = Array.new
  
  name_ary = ["kelvin", "mixed", "n=1", "n=2","n=3" ]
  line_index_ary = [12,22,92,32,42]


  # 描画
#  GGraph.line( $kelvin.rename("freq")[8..16], true, 
  GGraph.line( $kelvin.rename("e-folding_time")[8..16], true, 
	      "index" => line_index_ary[0] )
  GGraph.line( $mixed[8..16], false, "index" => line_index_ary[1] )
  GGraph.line( -$mixed2[8..16], false, 
	      "index" => line_index_ary[1],'type' => 3 )
  3.times{ |num|
    GGraph.line( ($grav_array[num])[8..16], false, 
		"index" => line_index_ary[num+2] )
    GGraph.line( -($grav2_array[num])[8..16], false, 
		"index" => line_index_ary[num+2],'type' => 3 )
    GGraph.line( -( $rossby2_array[num])[8..16], false, 
		"index" => line_index_ary[num+2], 'type' => 3 )
  }
#  GGraph.line( $axis0[9..16], false, "index" => line_index_ary[0] )

  # タイトル
  tani = "(#{$kelvin.data.get_att("units")})"
  DCL::uxsttl("t", tani , -1.0 )
  DCL::uxmttl("t", "growth rate", -1.0 ) 
#  DCL::uxmttl("t", "frequency", -1.0 ) 

  $file_label = "c_g =  0.167 + 0.082i"
  
  charsize = 0.53 - $file_label.size * 0.0095 
  DCL::sgtxzv(charsize,0.85,$file_label,0.8 * DCL.uzpget('rsizec1'),0,-1,1) 

  line_index(name_ary,line_index_ary)
  
  DCL::grcls
  
end

def line_index(str_ary,line_index_ary)
  charsize = 0.7 * DCL.uzpget('rsizec1')
  dvx = 0.01
  dvy = charsize*1.5
  raise TypeError,"Array expected" if ! str_ary.is_a?(Array)
  vxmin,vxmax,vymin,vymax = DCL.sgqvpt
  vx = 0.15
  vy = 0.12 - charsize/2
  str_ary.size.times{|num|
    DCL::sgtxzv(vx,vy,"--- #{str_ary[num]}", 
		charsize, 0, -1,line_index_ary[num])
    vy -= dvy
    if num == 2
      vx = 0.28
      vy = 0.12 - charsize/2
    end
  }
  nil
end


def bunsan_calc

#  $x_size = (NArray.sfloat(17).indgen - 8.0) *0.5
  $x_size = (NArray.sfloat(17).indgen - 8.0) 
  # omega = 2*3.1415/24/60/60
  # beta =  2*(omega)/(4e+7) * ( 60*60*24*(6370e+3) )  # (2*Omega/a)*cos(0)
  # c_g = ((200*9.8)**0.5)*60*60*24/(4e+7)

  beta = 1.0
#  c_g = 1.0

#  c_g = 0.167354855808284
#  c_g = Complex(0.167354855808284, 0.0823649746845925)

#  c_g = 0.0823649746845925**2 /4/0.02
#  c_g = 0.0823649746845925*$x_size - 0.02 * $x_size
  c_g = 0.0823649746845925

  # 軸
  knum = VArray.new($x_size).rename("knum").put_att("units","1")
  knum = Axis.new().set_pos(knum)
  grid = Grid.new(knum)

  # 軸
  $axis0 = $x_size * 0.0
  $axis0 = VArray.new($axis0).rename("axis0").
    put_att("units","1").put_att("ape_name","axis0")
  $axis0 = GPhys.new( grid, $axis0)

  ## 虚数を含む軸に変換
#  x_size = (NArray.complex(17).indgen - 8.0) *0.5
  x_size = (NArray.complex(17).indgen - 8.0) 

  # 初期化
  $rossby_array = Array.new
  $rossby2_array = Array.new
  $grav_array = Array.new
  $grav2_array = Array.new

  3.times{ |num|

    num = num +1
    
    ## 3 次方程式解法: カルダノの公式
    ## x^3 + 3px + q
    ## t^2 + qt - p =0, t = -q/2 +- (q**2/4 + p)**0.5
    p = - ( (c_g*x_size)**2 + c_g*beta*(2*num+1)) / 3
    q = - c_g**2*beta*x_size
    t1  =  -q/2 + (q**2/4 + p**3)**0.5
    t2  =  -q/2 - (q**2/4 + p**3)**0.5
    t1 = t1**(1.0/3)
    t2 = t2**(1.0/3)

    # 重力波
    grav1 = (t1 + t2).real
#    grav1 = (t1 + t2).imag
    grav1 = VArray.new(grav1).rename("grav1").
      put_att("units","1").put_att("ape_name","gravity")
    grav1 = GPhys.new( grid, grav1)
    
    # 反対重力波 (図示しない)
    grav2 = (t1 *  Complex.polar(1, Math::PI*2/3) + 
	     t2 * Complex.polar(1, Math::PI*4/3)).real

#    grav2 = (t1 *  Complex.polar(1, Math::PI*2/3) + 
#	     t2 * Complex.polar(1, Math::PI*4/3)).imag

    grav2 = (grav2 > 0) * grav2
    grav2 = VArray.new(grav2).rename("grav2").
      put_att("units","1").put_att("ape_name","gravity")
    grav2 = GPhys.new( grid, grav2)
    
    # ロスビー波
    rossby = (t2 * Complex.polar(1, Math::PI*2/3) + 
	      t1 * Complex.polar(1, Math::PI*4/3) ).real

#    rossby = (t2 * Complex.polar(1, Math::PI*2/3) + 
#	      t1 * Complex.polar(1, Math::PI*4/3) ).imag

    rossby2 = rossby
    rossby = (rossby > 0) * rossby
    rossby = VArray.new(rossby).rename("rossby").
      put_att("units","1").put_att("ape_name","rossby")
    rossby = GPhys.new( grid, rossby)

    rossby2 = VArray.new(rossby2).rename("rossby").
      put_att("units","1").put_att("ape_name","rossby")
    rossby2 = GPhys.new( grid, rossby2)

#    $rossby_array.push(rossby)

    $rossby_array.push(grav2)
    $rossby2_array.push(rossby2)
    $grav_array.push(grav1)
    $grav2_array.push(grav2)
  
  }

#  c_g = c_g.imag

  # 混合ロスビー
  $mixed =  c_g*x_size/2 + ((c_g*x_size/2)**2 + c_g*beta )**0.5 
  $mixed = VArray.new($mixed).rename("mixed").
    put_att("units","1").put_att("ape_name","mixed")
  $mixed = GPhys.new( grid, $mixed)

  # 反対混合ロスビー (図示しない)
  $mixed2 = c_g*x_size/2 - ((c_g*x_size/2)**2 + c_g*beta )**0.5
  $mixed2 = VArray.new($mixed2).rename("mixed").
    put_att("units","1").put_att("ape_name","mixed")
  $mixed2 = GPhys.new( grid, $mixed2)

  # ケルビン波
  $kelvin = (c_g*x_size).real
  $kelvin = ($kelvin > 0) * $kelvin
  $kelvin = VArray.new($kelvin).rename("kelvin").
    put_att("units","1").put_att("ape_name","kelvin")
  $kelvin = GPhys.new( grid, $kelvin)
  
end