#!/usr/bin/ruby

include Math
require "narray"
require './ebm_param'
require './ebm_fig'


class EBM

  # Constants for model mode
  IDHeatTransportBudyko  = 1
  IDHeatTransportSellers = 2

  # number of iterations
  NIteration = 100

  # convergence criterion for temperature
  ConvCriterionTemp = 1.0e-6
  # convergence criterion for solar flux (solar constant divided by 4)
  #ConvCriterionSolarFlux = 1.0e-3
  ConvCriterionSolarFlux = 1.0e-1

  # heat capacity devided by time step
  HeatCapDivByTimeStep = 0.0


  def initialize( jmax_in = 128, sTransferCoefC = TransferCoefC, sDiffCoefD = DiffCoefD, sAlbedoIce = AlbedoIce, sAlbedoIceFree = AlbedoIceFree )
    # number of grid
    @jmax = jmax_in

    # preparation of arrays
    y_Index = NArray.sfloat(@jmax  ).indgen!(1.0,1.0)
    q_Index = NArray.sfloat(@jmax+1).indgen!(0.0,1.0)

    # preparation of array
    @y_Temp  = NArray.sfloat(@jmax)
    @y_Lat   = 0.0 + 90.0 / @jmax / 2.0 + 90.0 / @jmax * ( y_Index - 1 )
    @q_Lat   = 0.0                      + 90.0 / @jmax * q_Index
    @sDelLat = PI / 2.0 / @jmax

    @y_DTDtRadSW = NArray.sfloat(@jmax)
    @y_DTDtRadLW = NArray.sfloat(@jmax)
    @y_DTDtTrans = NArray.sfloat(@jmax)

    # coefficient for heat transfer ( - TransferCoefC * ( T - \bar{T} ) )
    @sTransferCoefC  = sTransferCoefC

    # coefficient for heat diffusion
    @sDiffCoefD      = sDiffCoefD

    # albedo for ice free and ice covered surface
    @sAlbedoIce     = sAlbedoIce
    @sAlbedoIceFree = sAlbedoIceFree
  end


  def get_tempice
    return TempIce
  end
  def get_lat
    return @y_Lat
  end
  def get_temp
    return @y_Temp
  end
  def get_gmtemp
    return ( @y_Temp * NMath.cos(@y_Lat*PI/180.0) ).sum / NMath.cos(@y_Lat*PI/180.0).sum
  end
  def get_tend
    return @y_DTDtRadSW, @y_DTDtRadLW, @y_DTDtTrans
  end

  # Set solar radiatio as a function of latitude
  def solarrad( y_Lat )
    # Lindzen
    y_SolarDistr = 1.0 - 0.241 * ( 3.0 * NMath.sin(y_Lat*PI/180.0)**2 - 1.0 )

    return y_SolarDistr
  end


  # Fix solar radiation flux (solve simultaneous linear equation)
  def ebm( sIDHeatTransport = IDHeatTransportBudyko, sSolarRadTOA = 334.4, sTempInit = 300.0, sFlagDraw = false )

    # preparation of grids
    y_CosLat = NMath.cos( @y_Lat * PI / 180.0 )
    q_CosLat = NMath.cos( @q_Lat * PI / 180.0 )

    # solar radiation distribution function
    y_SolarDistr = solarrad( @y_Lat )


    # preparation of array for albedo
    y_Albedo = NArray.sfloat(@jmax)

    # temperature at a previous iteration
    y_TempB = NArray.sfloat(@jmax)


    yy_Mat = NMatrix.sfloat(@jmax,@jmax)
    y_Vec  = NVector.sfloat(@jmax)

    # Initial condition
    @y_Temp[true] = sTempInit

    # Iteration
    for n in 1..NIteration

      # save temperature
      y_TempB[true] = @y_Temp[true]

      # set albedo
      for j in 0..(@jmax-1)
        if @y_Temp[j] < TempIce then
          y_Albedo[j] = @sAlbedoIce
        else
          y_Albedo[j] = @sAlbedoIceFree
        end
      end


      case sIDHeatTransport
      when IDHeatTransportBudyko then
        # Budyko-type
        sSumWeight = y_CosLat.sum
        for k in 0..(@jmax-1)
          yy_Mat[true,k] = - @sTransferCoefC / sSumWeight * y_CosLat[true]
        end
        for k in 0..(@jmax-1)
          j = k
          yy_Mat[j,k] += PlanetRadCoefB + @sTransferCoefC
        end
        for k in 0..(@jmax-1)
          j = k
          yy_Mat[j,k] += HeatCapDivByTimeStep
        end
      when IDHeatTransportSellers then
        # Sellers-type
        yy_Mat[true,true] = 0.0
        for k in 0..0
          j = k
          yy_Mat[j  ,k] =   @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
          yy_Mat[j+1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
        end
        for k in (0+1)..((@jmax-1)-1)
          j = k
          yy_Mat[j-1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k]  = @sDiffCoefD / @sDelLat**2 * q_CosLat[j  ] / y_CosLat[j]
          yy_Mat[j  ,k] += @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
          yy_Mat[j+1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
        end
        for k in (@jmax-1)..(@jmax-1)
          j = k
          yy_Mat[j-1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k] =   @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
        end
      end

      y_Vec[true] = sSolarRadTOA * y_SolarDistr[true] * ( 1.0 - y_Albedo[true] ) - PlanetRadCoefA + HeatCapDivByTimeStep * y_TempB[true]


      y_Vec = y_Vec / yy_Mat
#      y_Vec = yy_Mat.lu.solve(y_Vec)
      @y_Temp[true] = y_Vec[true]

      draw_temp( TempIce, @y_Lat, @y_Temp, y_TempB ) if sFlagDraw

      # check covergence
      if ( @y_Temp - y_TempB ).abs.max < ConvCriterionTemp then
        break
      end

      if n >= NIteration then
        print "### Temperature has not been converged.\n"
      end
    end

    # calculate tendency
    @y_DTDtRadSW, @y_DTDtRadLW, @y_DTDtTrans = temp_tendency( sIDHeatTransport, sSolarRadTOA, y_SolarDistr, y_Albedo, y_CosLat, @y_Temp )
  end


  # Fix solar radiation flux
  def ebm_old( sSolarRadTOA = 334.4, sTempInit = 300.0, sFlagDraw = false )

    # preparation of grids
    y_CosLat = NMath.cos( @y_Lat * PI / 180.0 )

    # solar radiation distribution function
    y_SolarDistr = solarrad( @y_Lat )


    # preparation of array for albedo
    y_Albedo = NArray.sfloat(@jmax)

    # Initial condition
    @y_Temp[true] = sTempInit

    # temperature at a previous iteration
    y_TempB = NArray.sfloat(@jmax)

    # Iteration
    for n in 1..NIteration

      # save temperature
      y_TempB[true] = @y_Temp[true]

      # set albedo
      for j in 0..(@jmax-1)
        if @y_Temp[j] < TempIce then
          y_Albedo[j] = @sAlbedoIce
        else
          y_Albedo[j] = @sAlbedoIceFree
        end
      end

      sSumWeight = y_CosLat.sum
      sTempAve   = ( @y_Temp * y_CosLat ).sum
      sTempAve   = sTempAve / sSumWeight

      @y_Temp = ( sSolarRadTOA * y_SolarDistr * ( 1.0 - y_Albedo ) - PlanetRadCoefA + @sTransferCoefC * sTempAve ) / ( PlanetRadCoefB + @sTransferCoefC )

      draw_temp( TempIce, @y_Lat, @y_Temp, y_TempB ) if sFlagDraw

      # check covergence
      if ( @y_Temp - y_TempB ).abs.max < ConvCriterionTemp then
        break
      end

      if n >= NIteration then
        print "### Temperature has not been converged.\n"
      end
    end

    # calculate tendency
    @y_DTDtRadSW, @y_DTDtRadLW, @y_DTDtTrans = temp_tendency( sIDHeatTransport, sSolarRadTOA, y_SolarDistr, y_Albedo, y_CosLat, @y_Temp )
  end


  # Fix ice line latitude
  #   sIDHeatTransport: 1: Budyko-type,   2: diffusion-type
  #   sIDMinMax: 1: minimum temperature,   1: maximum temperature
  def ebm_fixlat( sIDHeatTransport = IDHeatTransportBudyko, sLatIce = 0.0, sIDMinMax = 1, sFlagDraw = false )

    # preparation of grids
    y_CosLat = NMath.cos( @y_Lat * PI / 180.0 )
    q_CosLat = NMath.cos( @q_Lat * PI / 180.0 )

    # solar radiation distribution function
    y_SolarDistr = solarrad( @y_Lat )


    # preparation of array for albedo
    y_Albedo = NArray.sfloat(@jmax)

    # temperature at a previous iteration
    y_TempB = NArray.sfloat(@jmax)


    yy_Mat = NMatrix.sfloat(@jmax,@jmax)
    y_Vec  = NVector.sfloat(@jmax)

    # find a grid which is closest to sLatIce
    sLatIceIndex = 0
    for j in 0..(@jmax-1)
      sLatIceIndex = j if ( @y_Lat[j] - sLatIce ).abs < ( @y_Lat[sLatIceIndex] - sLatIce ).abs
    end


    # set albedo
    case sIDMinMax
    when 1 then
      for j in 0..sLatIceIndex
        y_Albedo[j] = @sAlbedoIceFree
      end
      for j in (sLatIceIndex+1)..(@jmax-1)
        y_Albedo[j] = @sAlbedoIce
      end
    when 2 then
      for j in 0..(sLatIceIndex-1)
        y_Albedo[j] = @sAlbedoIceFree
      end
      for j in sLatIceIndex..(@jmax-1)
        y_Albedo[j] = @sAlbedoIce
      end
    else
      print "Unexpected mode\n"
      exit
    end


    # Initial condition
    @y_Temp[true] = TempIce
    sSolarRadTOA  = 0.0

    # Iteration
    for n in 1..NIteration do

      # save temperature
      y_TempB[true] = @y_Temp[true]
      # save solar radiation
      sSolarRadTOAB = sSolarRadTOA

      sSumWeight = y_CosLat.sum

#      # set albedo
#      #   This may not be needed. If low latitude is covered by ice, albedo
#      #   may be changed. 
#      for j in 0..(@jmax-1)
#        if @y_Temp[j] < TempIce then
#          y_Albedo[j] = AlbedoIce
#        else
#          y_Albedo[j] = AlbedoIceFree
#        end
#      end


      case sIDHeatTransport
      when IDHeatTransportBudyko then
        # Budyko-type
        for k in 0..(@jmax-1)
          yy_Mat[true,k] = - @sTransferCoefC / sSumWeight * y_CosLat[true]
        end
        for k in 0..(@jmax-1)
          j = k
          yy_Mat[j,k] += PlanetRadCoefB + @sTransferCoefC
        end

        j = sLatIceIndex
        yy_Mat[j,true] = - y_SolarDistr[true] * ( 1.0 - y_Albedo[true] )

        j = sLatIceIndex
        y_Vec[true] = @sTransferCoefC / sSumWeight * y_CosLat[j] * @y_Temp[j] - PlanetRadCoefA
        k = j
        y_Vec[k] += - ( PlanetRadCoefB + @sTransferCoefC ) * @y_Temp[j]

      when IDHeatTransportSellers then
        # Sellers-type
        yy_Mat[true,true] = 0.0
        for k in 0..0
          j = k
          yy_Mat[j  ,k] =   @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
          yy_Mat[j+1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
        end
        for k in (0+1)..((@jmax-1)-1)
          j = k
          yy_Mat[j-1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k]  = @sDiffCoefD / @sDelLat**2 * q_CosLat[j  ] / y_CosLat[j]
          yy_Mat[j  ,k] += @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
          yy_Mat[j+1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j+1] / y_CosLat[j]
        end
        for k in (@jmax-1)..(@jmax-1)
          j = k
          yy_Mat[j-1,k] = - @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k] =   @sDiffCoefD / @sDelLat**2 * q_CosLat[j] / y_CosLat[j]
          yy_Mat[j  ,k] += PlanetRadCoefB
          yy_Mat[j  ,k] += HeatCapDivByTimeStep
        end

        y_Vec[true] = - PlanetRadCoefA + HeatCapDivByTimeStep * y_TempB[true]

        j = sLatIceIndex
        for k in 0..(@jmax-1)
          y_Vec[k] -= yy_Mat[j,k] * @y_Temp[j]
          yy_Mat[j,k] = - y_SolarDistr[k] * ( 1.0 - y_Albedo[k] )
        end

      end



      y_Vec = y_Vec / yy_Mat
      @y_Temp[true]         = y_Vec[true]
      @y_Temp[sLatIceIndex] = TempIce
      sSolarRadTOA = y_Vec[sLatIceIndex]


      # check covergence
      if ( ( @y_Temp - y_TempB ).abs.max < ConvCriterionTemp ) && ( ( sSolarRadTOA - sSolarRadTOAB ).abs < ConvCriterionSolarFlux ) then
        break
      end

      # check albedo
      #   This may not be needed. If low latitude is covered by ice, albedo
      #   may be changed.
      for j in 0..(sLatIceIndex-1)
        if @y_Temp[j] < TempIce then
          print "Unexpected temperature at lat = ", @y_Lat[j], ", T = ", @y_Temp[j]
          exit
        end
      end
      j = sLatIceIndex
      @y_Temp[j] = TempIce
      for j in (sLatIceIndex+1)..(@jmax-1)
        if @y_Temp[j] > TempIce then
          print "Unexpected temperature at lat = ", @y_Lat[j], ", T = ", @y_Temp[j]
          exit
        end
      end

#      # modify temperature if needed for next iteration
#      for j in 0..(sLatIceIndex-1)
#        @y_Temp[j] = TempIce if @y_Temp[j] < TempIce
##        @y_Temp[j] = TempIce if @y_Temp[j] > TempIce
#      end
#      j = sLatIceIndex
#      @y_Temp[j] = TempIce
#      for j in (sLatIceIndex+1)..(@jmax-1)
#        @y_Temp[j] = TempIce if @y_Temp[j] > TempIce
##        @y_Temp[j] = TempIce if @y_Temp[j] < TempIce
#      end

      draw_temp( TempIce, @y_Lat, @y_Temp, y_TempB ) if sFlagDraw


      if n >= NIteration then
        print "### Temperature has not been converged.\n"
        print "### Delta Temp = ", ( @y_Temp - y_TempB ).abs.max, "\n"
        print "### ", ConvCriterionTemp, "\n"
        print "### Delta F0   = ", ( sSolarRadTOA - sSolarRadTOAB ), "\n"
        print "### ", ConvCriterionSolarFlux, "\n"
      end
    end

    # calculate tendency
    @y_DTDtRadSW, @y_DTDtRadLW, @y_DTDtTrans = temp_tendency( sIDHeatTransport, sSolarRadTOA, y_SolarDistr, y_Albedo, y_CosLat, @y_Temp )

    return sSolarRadTOA

  end


  def ice_line

    if @y_Temp[0] < TempIce then      # globally ice covered
      sIceLat =  0.0
    elsif @y_Temp[-1] > TempIce then  # globally ice free
      sIceLat = 90.0
    else
      for j in 1..(@y_Temp.size-1)
        if ( @y_Temp[j-1] - TempIce ) * ( @y_Temp[j] - TempIce ) <= 0.0 then
          sIceLat = ( @y_Lat[j] - @y_Lat[j-1] ) / ( @y_Temp[j] - @y_Temp[j-1] ) * ( TempIce - @y_Temp[j-1] ) + @y_Lat[j-1]
        end
      end
    end

    return sIceLat

  end

  def ice_line_j

    if @y_Temp[0] < TempIce then      # globally ice covered
      sIceLat =  0.0
      jIndex = 0
    elsif @y_Temp[-1] > TempIce then  # globally ice free
      sIceLat = 90.0
      jIndex = @jmax-1
    else
      for j in 1..(@y_Temp.size-1)
        if ( @y_Temp[j-1] - TempIce ) * ( @y_Temp[j] - TempIce ) <= 0.0 then
          sIceLat = ( @y_Lat[j] - @y_Lat[j-1] ) / ( @y_Temp[j] - @y_Temp[j-1] ) * ( TempIce - @y_Temp[j-1] ) + @y_Lat[j-1]
          jIndex = j-1
        end
      end
    end

    return jIndex

  end

  def temp_tendency( sIDHeatTransport, sSolarRadTOA, y_SolarDistr, y_Albedo, y_CosLat, y_Temp )

    y_DTDtRadSW = ( 1.0 - y_Albedo ) * sSolarRadTOA * y_SolarDistr
    y_DTDtRadLW = - ( PlanetRadCoefA + PlanetRadCoefB * y_Temp )
    y_DTDtTrans = - @sTransferCoefC * ( y_Temp - ( y_Temp * y_CosLat ).sum / y_CosLat.sum )

    if sIDHeatTransport == IDHeatTransportSellers then
      q_CosLat = NMath.cos( @q_Lat * PI / 180.0 )
      for j in 0..0
        y_DTDtTrans[j] = - ( - @sDiffCoefD * q_CosLat[j+1] * ( y_Temp[j+1] - y_Temp[j] ) / @sDelLat + 0.0 ) / ( y_CosLat[j] * @sDelLat )
      end
      for j in (0+1)..((@jmax-1)-1)
        y_DTDtTrans[j] = - ( - @sDiffCoefD * q_CosLat[j+1] * ( y_Temp[j+1] - y_Temp[j] ) / @sDelLat + @sDiffCoefD * q_CosLat[j] * ( y_Temp[j] - y_Temp[j-1] ) / @sDelLat ) / ( y_CosLat[j] * @sDelLat )
      end
      for j in (@jmax-1)..(@jmax-1)
        y_DTDtTrans[j] = - ( 0.0 + @sDiffCoefD * q_CosLat[j] * ( y_Temp[j] - y_Temp[j-1] ) / @sDelLat ) / ( y_CosLat[j] * @sDelLat )
      end
    end

    return y_DTDtRadSW, y_DTDtRadLW, y_DTDtTrans
  end

end