# plot_asrolrzonann.py
# 2026-06-22 masudako
# Reading CERES monthly data
# Plotting annual average
# Absorbed shortwave radiation and Outgoing longwave radiation

import numpy  as np
import xarray as xr
import matplotlib.pyplot as plt
import os
import sys
#
savefig = True
#
# データがふくむ格子点の経度、緯度 (かならずしも作図範囲にかぎらない)
nlo = 360
nla = 180
lat = np.linspace(-89.5,  89.5,  nla,   endpoint=True)
##print(lon)
##print(lat)
##sys.exit()

### 太陽放射吸収
svar = 'asr'
filepath = '/raid/ceres/binmon/'
nmonth = 12
# 作業配列
sum   = np.zeros((nla,nmonth), dtype=np.float32)
count = np.zeros((nla,nmonth), dtype=np.int32)
countthresh = nlo/2 # 便宜上きめる
#
for month in range(1, nmonth+1):
    smonth = f'{month:02}'
# isrとnsrのファイルを読み、ひきざんする
    filename = 'isr'+'clim'+smonth+'.r4l'
    filepathname = filepath + filename
    print('Reading: ', filepathname)
    f = open(filepathname, mode='rb')
    val1 = np.fromfile(f, dtype='<f',sep='').reshape(nla,nlo)
    filename = 'nsr'+'clim'+smonth+'.r4l'
    filepathname = filepath + filename
    print('Reading: ', filepathname)
    f = open(filepathname, mode='rb')
    val2 = np.fromfile(f, dtype='<f',sep='').reshape(nla,nlo)
    val  = val1 - val2
# たしこむ
    for la in range(nla):
       for lo in range(nlo):
           if not np.isnan(val[la,lo]):
               sum[  la,month-1] = sum[  la,month-1] + val[la,lo]
               count[la,month-1] = count[la,month-1] + 1
# 東西全経度平均
       if count[la,month-1] >= countthresh:
           sum[ la,month-1] =  sum[la,month-1] / count[la, month-1]
       else:
           sum[ la,month-1] = np.nan
# 年平均
daymo = [31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
#         1   2      3   4   5   6   7   8   9  10  11  12
swzonann= np.zeros(nla, dtype=np.float32)
weight  = np.zeros(nla, dtype=np.float32)
for la in range(nla):
    for month in range(1, 12+1):
        if not np.isnan(sum[la, month-1]):
            swzonann[la] = swzonann[la] + sum[la, month-1] * daymo[month-1]
            weight[ la]  = weight[  la] + daymo[month-1]
    swzonann[la] = swzonann[la] / weight[la]

### 地球放射射出
# 作業配列
sum   = np.zeros((nla,nmonth), dtype=np.float32)
count = np.zeros((nla,nmonth), dtype=np.int32)
countthresh = nlo/2 # 便宜上きめる
#
for month in range(1, nmonth+1):
    smonth = f'{month:02}'
# olrのファイルを読む
    filename = 'olr'+'clim'+smonth+'.r4l'
    filepathname = filepath + filename
    print('Reading: ', filepathname)
    f = open(filepathname, mode='rb')
    val = np.fromfile(f, dtype='<f',sep='').reshape(nla,nlo)
# たしこむ
    for la in range(nla):
       for lo in range(nlo):
           if not np.isnan(val[la,lo]):
               sum[  la,month-1] = sum[  la,month-1] + val[la,lo]
               count[la,month-1] = count[la,month-1] + 1
# 東西全経度平均
       if count[la,month-1] >= countthresh:
           sum[ la,month-1] =  sum[la,month-1] / count[la, month-1]
       else:
           sum[ la,month-1] = np.nan
# 年平均
daymo = [31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
#         1   2      3   4   5   6   7   8   9  10  11  12
lwzonann= np.zeros(nla, dtype=np.float32)
weight  = np.zeros(nla, dtype=np.float32)
for la in range(nla):
    for month in range(1, 12+1):
        if not np.isnan(sum[la, month-1]):
            lwzonann[la] = lwzonann[la] + sum[la, month-1] * daymo[month-1]
            weight[ la]  = weight[  la] + daymo[month-1]
    lwzonann[la] = lwzonann[la] / weight[la]
    
plt.title('Radiation at TOA [W/m${}^2$] (CERES 2000.03-2025.12 "climatology")')
plt.xlabel('Latitude (deg. N)', fontsize=12)
plt.xlim(-95,   95)
plt.ylim(-25,  425)
plt.xticks(np.arange(-90,  90+1,  15), fontsize=12)
plt.yticks(np.arange(  0, 400+1,  50), fontsize=12)
plt.grid(linewidth=0.25, color='black', linestyle='dashed')
plt.plot(lat, swzonann, '-', color='blue', label='Absorbed SW rad.')
plt.plot(lat, lwzonann, '-', color='red',  label='Outgonig LW rad.')
plt.legend(fontsize=12)
#
if savefig:
# 画像ファイル出力
    outimagefile = 'asrolrclimzonann.png'
    print('Writing ', outimagefile)
    plt.savefig(outimagefile, bbox_inches='tight', pad_inches=0.0)
    plt.clf()
    plt.close()
else:
# 画面表示
    plt.show()
##### END OF PROGRAM #####