require "narray"
require "ispack"

module SPML

  attr_reader :rn
  attr_reader :x_Lon, :y_Lat
  attr_reader :x_Lon_Weight, :y_Lat_Weight
  attr_reader :xy_Lon, :xy_Lat
  attr_reader :w_spectrum_VMiss

  def w_Initial(nm,im,jm)
    w_deriv_initial(nm,im,jm)
    @w_spectrum_VMiss = -999.000
  end
  def l_nm(arg0,arg1)
    if Numeric===arg0 && Numeric===arg1
      return l_nm_array00(arg0,arg1)
    elsif Numeric===arg0 && NArray===arg1
      return l_nm_array01(arg0,arg1)
    elsif NArray===arg0 && Numeric===arg1
      return l_nm_array10(arg0,arg1)
    elsif NArray===arg0 && NArray===arg1
       return l_nm_array11(arg0,arg1)
    else
      raise "arg must be Numeric or NArray"
    end
  end

  def w_xy(xy_data,ipow=0,iflag=0)
    return ISPACK::sntg2s(@nm,@im,@im,@jm,@jm,1,xy_data,
                          @it,@t,@y,@ip,@p,@r,@ia,@a,@q,@ws,@ww,
                          ipow,iflag)
  end
  def xy_w(w_data,ipow=0,iflag=0)
    return ISPACK::snts2g(@nm,@im,@im,@jm,@jm,1,w_data,
                          @it,@t,@y,@ip,@p,@r,@ia,@a,@q,@ws,@ww,
                          ipow,iflag).reshape!(@im,@jm)
  end

  def w_Lapla_w(w_data)
    return ISPACK::spclap(@nm,w_data,@rn[true,0])
  end
  def w_LaplaInv_w(w_data)
    return ISPACK::spclap(@nm,w_data,@rn[true,1])
  end

  def w_DLon_w(w_data)
    obj = NArray.float((@nm+1)*(@nm+1))
    ISPACK::spclam(@nm,w_data,obj,@irm)
    return obj
  end
  def xy_GradLon_w(w_data)
    return xy_w(w_data,1,-1)
  end
  def xy_GradLat_w(w_data)
    return xy_w(w_data,1,1)
  end
  def w_DivLon_xy(xy_data)
    return w_xy(xy_data,1,-1)
  end
  def w_DivLat_xy(xy_data)
    return w_xy(xy_data,1,1)
  end
  def w_Div_xy_xy(xy_u,xy_)
    return x_DivLon_xy(xy_u) + w_DivLat_xy(xy_v)
  end
  def w_Jacobian_w_w(w_a,w_b)
    obj = NArray.float((@nm+1)*(@nm+1))
    ISPACK::spnjcb(@nm,@im,@im,@jm,@jm,w_a,w_b,obj,@it,@t,@y,@ip2,@p2,@r2,@ip3,@p3,@r3,@ia,@q,@ws,@ww)
    return obj
  end

  def xy_GradLambda_w(w_data)
    return xy_w(w_data,0,-1)
  end
  def xy_GradMu_w(w_data)
    return xy_w(w_data,0,1)
  end
  def w_DivLambda_xy(xy_data)
    return w_xy(xy_data,2,-1)
  end
  def w_DivMu_xy(xy_data)
    return w_xy(xy_data,2,1)
  end

  def IntLon_x(x_data)
    return x_data.mul_add(@x_Lon_Weight,0)
  end
  def x_IntLat_xy(xy_data)
    ylw = @y_Lat_Weight.reshape(1,@jm)
    return (xy_data*ylw).sum(1)
  end
  def AvrLon_x(x_data)
    return IntLon_x(x_data)/@x_Lon_Weight.sum
  end
  def x_AvrLat_xy(xy_data)
    return x_IntLat_xy(xy_data)/@y_Lat_Weight.sum
  end
  def AvrLonLat_xy(xy_data)
    return AvrLon_x(x_AvrLat_xy(xy_data))
  end

  def nm_EnergyFromStreamFunc_w(w_StrFunc)
    obj = NArray.float(@nm+1,@nm*2+1).fill(w_spectrum_VMiss)
    for n in 0..@nm
      for m in -n..n
        obj[n,m+n] = 0.5*n*(n+1)*w_StrFunc[l_nm(n,m)]**2
      end
    end
    return obj
  end
  def n_EnergyFromStreamFunc_w(w_StrFunc)
    obj = NArray.float(@nm+1)
    for n in 0..@nm
      ind = NArray.sint(n*2+1).indgen(-n)
      obj[n] = 0.5*n*(n+1)*(w_StrFunc[l_nm(n,ind)]**2).sum
    end
    return obj
  end
  def nm_EnstrophyFromStreamFunc_w(w_StrFunc)
    obj = NArray.float(@nm+1,@nm*2+1).fill(w_spectrum_VMiss)
    for n in 0..@nm
      for m in -n..n
        obj[n,m+n] = 0.5*n**2*(n+1)**2*w_StrFunc[l_nm(n,m)]**2
      end
    end
    return obj
  end
  def n_EnstrophyFromStreamFunc_w(w_StrFunc)
    obj = NArray.float(@nm+1)
    for n in 0..@nm
      ind = NArray.sint(n*2+1).indgen(-n)
      obj[n] = 0.5*n**2*(n+1)**2*(w_StrFunc[l_nm(n,ind)]**2).sum
    end
    return obj
  end

  private

  def w_base_initial(n_in,i_in,j_in)
    @im = i_in
    @jm = j_in
    @nm = n_in

    @q = NArray.float(((@nm+1)/2+@nm+1)*@jm)
    iw = (@im+3*(@nm+1))*@jm
    @ws = NArray.float(iw) if @ws.nil? || @ws.length<iw
    @ww = NArray.float(iw) if @ww.nil? || @ww.length<iw

    @it,@t,@y,@ip,@p,@r,@ia,@a = ISPACK::sninit(@nm,@im,@jm)

    @x_Lon = 2.0*Math::PI*NArray.float(@im).indgen/@im
    @x_Lon_Weight = NArray.float(@im).fill(2.0*Math::PI/@im)

    @y = @y.reshape!(@jm/2,4)
    @y_Lat = NArray.float(@jm)
    @y_Lat[0...@jm/2] = NMath::asin(-@y[@jm/2-1..0,0])
    @y_Lat[@jm/2..-1] = NMath::asin(@y[0...@jm/2,0])
    @y_Lat_Weight = NArray.float(@jm)
    @y_Lat_Weight[0...@jm/2] = 2*@y[@jm/2-1..0,1]
    @y_Lat_Weight[@jm/2..-1] = 2*@y[0...@jm/2,1]
    @y = @y.reshape!(@jm*2)

    @xy_Lon = @x_Lon*NArray.float(1,@jm).fill(1)
    @xy_Lat = NArray.float(@im,1).fill(1)*@y_Lat.reshape(1,@jm)
  end

  def l_nm_array00(n,m)
    return ISPACK::snnm2l(n,m)-1
  end
  def l_nm_array01(n,marray)
    len = marray.length
    obj = NArray.float(len)
    len.times{|i|
      obj[i] = l_nm_array00(n,marray[i])
    }
    return obj
  end
  def l_nm_array10(narray,m)
    len = narray.length
    obj = NArray.float(len)
    len.times{|i|
      obj[i] = l_nm_array00(narray[i],m)
    }
    return obj
  end
  def l_nm_array11(narray,marray)
    len = narray.length
    unless marray.length==len
      raise "dimensions of input arrays n and m are different."
    end
    obj = NArray.float(len)
    len.times{|i|
      obj[i] = l_nm_array00(narray[i],marray[i])
    }
    return obj
  end

  def w_deriv_initial(n_in,i_in,j_in)
    w_base_initial(n_in,i_in,j_in)
    @rn = ISPACK::spnini(@nm).reshape!((@nm+1)*(@nm+1),2)
    @irm = ISPACK::spmini(@nm)

    @ip2,@p2,@r2 = ISPACK::snkini(@nm,@jm,2,@ip,@p,@r)
    @ip3,@p3,@r3 = ISPACK::snkini(@nm,@jm,3,@ip,@p,@r)

    iw = 3*[((@nm+1)/2*2+3)*(@nm/2+2)*2,@jm*((@nm+1)/2+@nm+1)*2,@jm*@jm].max
    @ws = NArray.float(iw) if @ws.nil? || @ws.length<iw
    @ww = NArray.float(iw) if @ww.nil? || @ww.length<iw
  end

end