# +-======-+ 
#  Copyright (c) 2003-2018 United States Government as represented by 
#  the Admistrator of the National Aeronautics and Space Administration.  
#  All Rights Reserved.
#  
#  THIS OPEN  SOURCE  AGREEMENT  ("AGREEMENT") DEFINES  THE  RIGHTS  OF USE,
#  REPRODUCTION,  DISTRIBUTION,  MODIFICATION AND REDISTRIBUTION OF CERTAIN 
#  COMPUTER SOFTWARE ORIGINALLY RELEASED BY THE UNITED STATES GOVERNMENT AS 
#  REPRESENTED BY THE GOVERNMENT AGENCY LISTED BELOW ("GOVERNMENT AGENCY").  
#  THE UNITED STATES GOVERNMENT, AS REPRESENTED BY GOVERNMENT AGENCY, IS AN 
#  INTENDED  THIRD-PARTY  BENEFICIARY  OF  ALL  SUBSEQUENT DISTRIBUTIONS OR 
#  REDISTRIBUTIONS  OF THE  SUBJECT  SOFTWARE.  ANYONE WHO USES, REPRODUCES, 
#  DISTRIBUTES, MODIFIES  OR REDISTRIBUTES THE SUBJECT SOFTWARE, AS DEFINED 
#  HEREIN, OR ANY PART THEREOF,  IS,  BY THAT ACTION, ACCEPTING IN FULL THE 
#  RESPONSIBILITIES AND OBLIGATIONS CONTAINED IN THIS AGREEMENT.
#  
#  Government Agency: National Aeronautics and Space Administration
#  Government Agency Original Software Designation: GSC-15354-1
#  Government Agency Original Software Title:  GEOS-5 GCM Modeling Software
#  User Registration Requested.  Please Visit http://opensource.gsfc.nasa.gov
#  Government Agency Point of Contact for Original Software:  
#  			Dale Hithon, SRA Assistant, (301) 286-2691
#  
# +-======-+ 
"""
Python implementation of set_eta module under GMA_Shared/GMAO_hermes.

"""

from numpy import ones

ak = {}
bk = {}

# NCAR settings
# -------------

ak['18'] =    [ 291.70,  792.92,  2155.39,  4918.34,  8314.25,      
               7993.08, 7577.38,  7057.52,  6429.63,  5698.38,      
               4879.13, 3998.95,  3096.31,  2219.02,  1420.39,      
               754.13,  268.38,   0.0000,   0.0000 ]

bk['18'] =    [ 0.0000,    0.0000,    0.0000,   0.0000,   0.0000,   
                0.0380541, 0.0873088, 0.1489307, 0.2232996,         
                0.3099406, 0.4070096, 0.5112977, 0.6182465,         
                0.7221927, 0.8168173, 0.8957590, 0.9533137,         
                0.9851122, 1.0  ]
     
ak['26'] =    [ 219.4067,  489.5209,   988.2418,   1805.201,        
                2983.724,  4462.334,   6160.587,   7851.243,        
                7731.271,  7590.131,   7424.086,   7228.744,        
                6998.933,  6728.574,   6410.509,   6036.322,        
                5596.111,  5078.225,   4468.96,    3752.191,        
                2908.949,  2084.739,   1334.443,   708.499,         
                252.136,   0.,         0. ]

bk['26'] =    [ 0.,         0.,         0.,         0.,             
                0.,         0.,         0.,         0.,             
                0.01505309, 0.03276228, 0.05359622, 0.07810627,     
                0.1069411,  0.14086370, 0.180772,   0.227722,       
                0.2829562,  0.3479364,  0.4243822,  0.5143168,      
                0.6201202,  0.7235355,  0.8176768,  0.8962153,      
                0.9534761,  0.9851122,  1.        ]

ak['30'] =    [ 225.523952394724, 503.169186413288, 1015.79474285245, 
               1855.53170740604, 3066.91229343414,  4586.74766123295, 
               6332.34828710556, 8070.14182209969,  9494.10423636436, 
              11169.321089983,  13140.1270627975,  15458.6806893349,  
              18186.3352656364, 17459.799349308,   16605.0657629967,  
              15599.5160341263, 14416.541159153,   13024.8308181763,  
              11387.5567913055,  9461.38575673103,  7534.44507718086, 
               5765.89405536652, 4273.46378564835,  3164.26791250706, 
               2522.12174236774, 1919.67375576496,  1361.80268600583, 
                853.108894079924, 397.881818935275,    0.,            
                  0.  ]

bk['30'] =    [ 0.,                 0.,                                  
                0.,                 0.,                0.,               
                0.,                 0.,                0.,               
                0.,                 0.,                0.,               
                0.,                 0.,                0.03935482725501, 
                0.085653759539127,  0.140122056007385, 0.20420117676258, 
                0.279586911201477,  0.368274360895157, 0.47261056303978, 
                0.576988518238068,  0.672786951065063, 0.75362843275070, 
                0.813710987567902,  0.848494648933411, 0.88112789392471, 
                0.911346435546875,  0.938901245594025, 0.96355980634689, 
                0.985112190246582,  1.   ]

# NASA DAO settings
# -----------------

ak['32'] =     [0.00000,     106.00000,     224.00000,     
              411.00000,     685.00000,    1065.00000,     
             1565.00000,    2179.80000,    2900.00000,     
             3680.00000,    4550.00000,    5515.00000,     
             6607.00000,    7844.00000,    9236.56616,     
            10866.34280,   12783.70000,   15039.29900,     
            17693.00000,   20815.20900,   24487.49020,     
            28808.28710,   32368.63870,   33739.96480,     
            32958.54300,   30003.29880,   24930.12700,     
            18568.89060,   12249.20510,    6636.21191,     
             2391.51416,       0.00000,       0.00000 ]

bk['32'] =     [0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.00000,       0.00000,     
                0.00000,       0.01523,       0.06132,     
                0.13948,       0.25181,       0.39770,     
                0.55869,       0.70853,       0.83693,     
                0.93208,       0.98511,       1.00000 ]

ak['48'] =    [40.00000,     100.00000,     200.00000,     
              350.00000,     550.00000,     800.00000,     
             1085.00000,    1390.00000,    1720.00000,     
             2080.00000,    2470.00000,    2895.00000,     
             3365.00000,    3890.00000,    4475.00000,     
             5120.00000,    5830.00000,    6608.00000,     
             7461.00000,    8395.00000,    9424.46289,     
            10574.46900,   11864.80330,   13312.58850,     
            14937.03770,   16759.70760,   18804.78670,     
            21099.41250,   23674.03720,   26562.82650,     
            29804.11680,   32627.31601,   34245.89759,     
            34722.29104,   34155.20062,   32636.50533,     
            30241.08406,   27101.45052,   23362.20912,     
            19317.04955,   15446.17194,   12197.45091,     
             9496.39912,    7205.66920,    5144.64339,     
             3240.79521,    1518.62245,       0.00000,     
                0.00000  ]

bk['48'] =    [0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00000,       0.00000,      
               0.00000,       0.00813,       0.03224,      
               0.07128,       0.12445,       0.19063,      
               0.26929,       0.35799,       0.45438,      
               0.55263,       0.64304,       0.71703,      
               0.77754,       0.82827,       0.87352,      
               0.91502,       0.95235,       0.98511,      
               1.00000      ]

ak['55'] =    [ 1.00000,       2.00000,       3.27000,                   
              4.75850,       6.60000,       8.93450,                     
             11.97030,      15.94950,      21.13490,                     
             27.85260,      36.50410,      47.58060,                     
             61.67790,      79.51340,     101.94420,                     
            130.05080,     165.07920,     208.49720,                     
            262.02120,     327.64330,     407.65670,                     
            504.68050,     621.68000,     761.98390,                     
            929.29430,    1127.68880,    1364.33920,                     
           1645.70720,    1979.15540,    2373.03610,                     
           2836.78160,    3380.99550,    4017.54170,                     
           4764.39320,    5638.79380,    6660.33770,                     
           7851.22980,    9236.56610,   10866.34270,                     
          12783.70000,   15039.30000,   17693.00000,                     
          20119.20876,   21686.49129,   22436.28749,                     
          22388.46844,   21541.75227,   19873.78342,                     
          17340.31831,   13874.44006,   10167.16551,                     
           6609.84274,    3546.59643,    1270.49390,                     
              0.00000,       0.00000   ]

bk['55'] =     [0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,            
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00000,       0.00000,       0.00000,           
                0.00696,       0.02801,       0.06372,           
                0.11503,       0.18330,       0.27033,           
                0.37844,       0.51046,       0.64271,           
                0.76492,       0.86783,       0.94329,           
                0.98511,       1.00000  ]


# NCEP's 64 sigma layers
# ----------------------

ak['64'] =     [1.00000,       3.90000,       8.70000,      
               15.42000,      24.00000,      34.50000,      
               47.00000,      61.50000,      78.60000,      
               99.13500,     124.12789,     154.63770,      
              191.69700,     236.49300,     290.38000,      
              354.91000,     431.82303,     523.09300,      
              630.92800,     757.79000,     906.45000,      
             1079.85000,    1281.00000,    1515.00000,      
             1788.00000,    2105.00000,    2470.00000,      
             2889.00000,    3362.00000,    3890.00000,      
             4475.00000,    5120.00000,    5830.00000,      
             6608.00000,    7461.00000,    8395.00000,      
             9424.46289,   10574.46880,   11864.80270,      
            13312.58890,   14937.03710,   16759.70700,      
            18804.78710,   21099.41210,   23674.03710,      
            26562.82810,   29804.11720,   32627.31640,      
            34245.89840,   34722.28910,   34155.19920,      
            32636.50390,   30241.08200,   27101.44920,      
            23362.20700,   19317.05270,   15446.17090,      
            12197.45210,    9496.39941,    7205.66992,      
             5144.64307,    3240.79346,    1518.62134,      
                0.00000,       0.00000 ]

bk['64'] =     [0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00000,      
                0.00000,       0.00000,       0.00813,      
                0.03224,       0.07128,       0.12445,      
                0.19063,       0.26929,       0.35799,      
                0.45438,       0.55263,       0.64304,      
                0.71703,       0.77754,       0.82827,      
                0.87352,       0.91502,       0.95235,      
                0.98511,       1.00000 ]


ak['72'] = [ 1.0000000,   2.0000002,       3.2700005,       4.7585009,       6.6000011, 
         8.9345014,       11.970302,       15.949503,       21.134903,       27.852606, 
         36.504108,       47.580610,       61.677911,       79.513413,       101.94402, 
         130.05102,       165.07903,       208.49704,       262.02105,       327.64307, 
         407.65710,       504.68010,       621.68012,       761.98417,       929.29420, 
         1127.6902,       1364.3402,       1645.7103,       1979.1604,       2373.0405, 
         2836.7806,       3381.0007,       4017.5409,       4764.3911,       5638.7912, 
         6660.3412,       7851.2316,       9236.5722,       10866.302,       12783.703, 
         15039.303,       17693.003,       20119.201,       21686.501,       22436.301, 
         22389.800,       21877.598,       21214.998,       20325.898,       19309.696, 
         18161.897,       16960.896,       15625.996,       14290.995,       12869.594, 
         11895.862,       10918.171,       9936.5219,       8909.9925,       7883.4220, 
         7062.1982,       6436.2637,       5805.3211,       5169.6110,       4533.9010, 
         3898.2009,       3257.0809,       2609.2006,       1961.3106,       1313.4804, 
         659.37527,       4.8048257,       0.0000000 ]


bk['72'] = [0.0000000,    0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,       0.0000000,       0.0000000,       0.0000000,       0.0000000, 
         0.0000000,   8.1754130e-09,    0.0069600246,     0.028010041,     0.063720063, 
        0.11360208,      0.15622409,      0.20035011,      0.24674112,      0.29440312, 
        0.34338113,      0.39289115,      0.44374018,      0.49459020,      0.54630418, 
        0.58104151,      0.61581843,      0.65063492,      0.68589990,      0.72116594, 
        0.74937819,      0.77063753,      0.79194696,      0.81330397,      0.83466097, 
        0.85601798,      0.87742898,      0.89890800,      0.92038701,      0.94186501, 
        0.96340602,      0.98495195,       1.0000000 ]


#  ECMWF Nature
#  ------------
ak['91'] =    [    0.000000000000,    2.000040054321,      3.980832099915,     7.387186050415, 
                  12.908319473267,   21.413604736328,     33.952865600586,    51.746597290039, 
                  76.167663574219,  108.715560913086,    150.986022949219,   204.637451171875, 
                 271.356445312500,  352.824462890625,    450.685791015625,   566.519287109375,  
                 701.813232421875,  857.945800781250,                                          
                1036.166503906250, 1237.585449218750,   1463.163818359375,  1713.709716796875, 
                1989.874511718750, 2292.155517578125,   2620.898437500000,  2976.302246093750, 
                3358.425781250000, 3767.196044921875,   4202.417968750000,  4663.777343750000, 
                5150.859375000000, 5663.156250000000,   6199.839843750000,  6759.726562500000, 
                7341.468750000000, 7942.925781250000,   8564.625000000000,  9208.304687500000, 
                9873.562500000000, 10558.882812500000,  11262.484375000000, 11982.660156250000, 
               12713.898437500000, 13453.226562500000,  14192.011718750000, 14922.687500000000, 
               15638.054687500000, 16329.562500000000,  16990.625000000000, 17613.281250000000, 
               18191.031250000000, 18716.968750000000,  19184.546875000000, 19587.515625000000, 
               19919.796875000000, 20175.394531250000,  20348.917968750000, 20434.156250000000, 
               20426.218750000000, 20319.011718750000,  20107.031250000000, 19785.359375000000, 
               19348.777343750000, 18798.824218750000,  18141.296875000000, 17385.593750000000, 
               16544.585937500000, 15633.566406250000,  14665.644531250000, 13653.218750000000, 
               12608.382812500000, 11543.167968750000,  10471.312500000000,  9405.222656250000, 
                8356.253906250000,  7335.164062500000,   6353.921875000000,  5422.800781250000, 
                4550.214843750000,  3743.464355468750,   3010.146972656250,  2356.202636718750, 
                1784.854492187500,  1297.656250000000,    895.193603515625,   576.314208984375, 
                 336.772460937500,   162.043426513672,     54.208343505859 ,    6.575628280640, 
                   0.003160000080,     0.000000000000]

bk['91'] =     [  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000000000,     0.000000000000, 
                  0.000000000000,     0.000000000000,      0.000000272400,     0.000013911600, 
                  0.000054667194,     0.000131364097,      0.000278884778,     0.000548384152, 
                  0.001000134507,     0.001701075351,      0.002764719306,     0.004267048091, 
                  0.006322167814,     0.009034991264,      0.012508261949,     0.016859579831, 
                  0.022188644856,     0.028610348701,      0.036226909608,     0.045146133751, 
                  0.055474229157,     0.067316174507,      0.080777287483,     0.095964074135, 
                  0.112978994846,     0.131934821606,      0.152933537960,     0.176091074944, 
                  0.201520144939,     0.229314863682,      0.259554445744,     0.291993439198, 
                  0.326329410076,     0.362202584743,      0.399204790592,     0.436906337738, 
                  0.475016415119,     0.513279736042,      0.551458477974,     0.589317142963, 
                  0.626558899879,     0.662933588028,      0.698223590851,     0.732223808765, 
                  0.764679491520,     0.795384764671,      0.824185431004,     0.850950419903, 
                  0.875518381596,     0.897767245770,      0.917650938034,     0.935157060623, 
                  0.950273811817,     0.963007092476,      0.973466038704,     0.982238113880, 
                  0.989152967930,     0.994204163551,      0.997630119324,     1.000000000000]


ak['96'] =    [  1.00000,       2.32782,       3.34990, 
               4.49484,       5.62336,       6.93048, 
               8.41428,      10.06365,      11.97630, 
              14.18138,      16.70870,      19.58824, 
              22.84950,      26.52080,      30.62845, 
              35.19588,      40.24273,      45.78375, 
              51.82793,      58.43583,      65.62319, 
              73.40038,      81.77154,      90.73373, 
             100.27628,     110.82243,     122.47773, 
             135.35883,     149.59464,     165.32764, 
             182.71530,     201.93164,     223.16899, 
             246.63988,     272.57922,     301.24661, 
             332.92902,     367.94348,     406.64044, 
             449.40720,     496.67181,     548.90723, 
             606.63629,     670.43683,     740.94727, 
             818.87329,     904.99493,    1000.17395, 
            1105.36304,    1221.61499,    1350.09326, 
            1492.08362,    1649.00745,    1822.43469, 
            2014.10168,    2225.92627,    2460.02905, 
            2718.75195,    3004.68530,    3320.69092, 
            3669.93066,    4055.90015,    4482.46240, 
            4953.88672,    5474.89111,    6050.68994, 
            6687.04492,    7390.32715,    8167.57373, 
            9026.56445,    9975.89648,   11025.06934, 
           12184.58398,   13466.04785,   14882.28320, 
           16447.46289,   18177.25781,   20088.97461, 
           21886.89453,   23274.16602,   24264.66602, 
           24868.31641,   25091.15430,   24935.41016, 
           24399.52148,   23478.13281,   22162.01758, 
           20438.00586,   18288.83984,   15693.01172, 
           12624.54199,    9584.35352,    6736.55713, 
            4231.34326,    2199.57910,     747.11890, 
              0.00000 ]

bk['96'] =    [0.00000,       0.00000,       0.00000,
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00000,       0.00000,       0.00000, 
              0.00315,       0.01263,       0.02853, 
              0.05101,       0.08030,       0.11669, 
              0.16055,       0.21231,       0.27249, 
              0.34169,       0.42062,       0.51005, 
              0.61088,       0.70748,       0.79593, 
              0.87253,       0.93400,       0.97764,  
              1.00000 ]

def getEdge(km):
    """Return tuple with edge values of (ak,bk) for the input number of 
       vertical layers *km*. Notice that the ak's are in units of Pascal (Pa). """
    ae = ones(km+1)
    be = ones(km+1)
    ae[:] = ak[str(km)]
    be[:] = bk[str(km)] 
    return (ae,be)

def getMid(km):
    """Return tuple with mid-layer values of (ak,bk) for the input number of 
       vertical layers *km*. Notice that the ak's are in units of Pascal (Pa). """
    ae, be = (ak[str(km)], bk[str(km)]) 
    am = ones(km)
    bm = ones(km)
    for k in range(km):
        am[k] = (ae[k+1] + ae[k]) / 2.
        bm[k] = (be[k+1] + be[k]) / 2.
    return (am,bm)

def getDelta(km):
    """Return tuple with delta values of (ak,bk) for the input number of 
       vertical layers *km*. Notice that the ak's are in units of Pascal (Pa). """
    ae, be = (ak[str(km)], bk[str(km)]) 
    dak = ones(km)
    dbk = ones(km)
    for k in range(km):
        dak[k] = ae[k+1] - ae[k]
        dbk[k] = be[k+1] - be[k]
    return (dak,dbk)

def getPe(km,p_ref=100000.):
    """Return pressure at edges given a reference pressure."""
    ae, be = getEdge(km)
    return (ae + p_ref * be)

def getPm(km,p_ref=100000.):
    """Return pressure at mid-layer given a reference pressure."""
    am, bm = getMid(km)
    return (am + p_ref * bm)

def getDelp(km,p_ref=100000.):
    """Return pressure thickness at mid-layer given a reference pressure."""
    dak, dbk = getDelta(km)
    return (dak + p_ref * dbk)



