# -*- coding: utf-8 -*-
# ku_map_daily.py
# 気象庁 解析雨量 (GRIB2形式) を読み、
# 日ごとに集計して
# Cartopy で地図上に表示する
# 2025-05-03, 2026-07-06 masudako
#
##### パッケージの予約
# - 日時の換算
from datetime import datetime, timedelta
# - 数値処理
import numpy as np
# - 異常終了
import sys
# - 作図関係
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
# - 作図関係 ここまで
# - GRIB2解読のために必要
from itertools import repeat
import struct
# - GRIB2解読のために必要
# - 都道府県境のために必要
import cartopy.io.shapereader as shapereader
from cartopy.feature import ShapelyFeature
# - 都道府県境のために必要 ここまで

### 解析雨量につかわれているGRIB2解読のための関数群 ###
def set_table(section5):
    max_level = struct.unpack_from('>H', section5, 15)[0]
    table = (
        -10, # define representative of level 0 (Missing Value)
        *struct.unpack_from('>'+str(max_level)+'H', section5, 18)
    )
    return np.array(table, dtype=np.int16)

def decode_runlength(code, hi_level):
    for raw in code:
        if raw <= hi_level:
            level = raw
            pwr = 0
            yield level
        else:
            length = (0xFF - hi_level)**pwr * (raw - (hi_level + 1))
            pwr += 1
            yield from repeat(level, length)

def load_jmara_grib2(file):
    with open(file, 'rb') as f:
        binary = f.read()

    len_ = {'sec0':16, 'sec1':21, 'sec3':72, 'sec4':82, 'sec6':6}

    end4 = len_['sec0'] + len_['sec1'] + len_['sec3'] + len_['sec4'] - 1
    len_['sec5'] = struct.unpack_from('>I', binary, end4+1)[0]
    section5 = binary[end4:(end4+len_['sec5']+1)]

    end6 = end4 + len_['sec5'] + len_['sec6']
    len_['sec7'] = struct.unpack_from('>I', binary, end6+1)[0]
    section7 = binary[end6:(end6+len_['sec7']+1)]

    highest_level = struct.unpack_from('>H', section5, 13)[0]
    level_table = set_table(section5)
    decoded = np.fromiter(
        decode_runlength(section7[6:], highest_level), dtype=np.int16
    ).reshape((3360, 2560))

    # convert level to representative
    return level_table[decoded]
### 解析雨量につかわれているGRIB2解読のための関数群 ここまで ###

### 解析雨量データを1時刻ぶん読む関数 (hamana, niigata 用) ###
def read_k_u(iyear, imon, iday, ihour, imin, datapath):
    syear = f'{iyear:04}'
    smon  = f'{imon:02}'
    sday  = f'{iday:02}'
    shour = f'{ihour:02}'
    smin  = f'{imin:02}'
# 読みこむファイル名
    filename = 'Z__C_RJTD_'+syear+smon+sday+shour+smin+'00' \
             +'_SRF_GPV_Ggis1km_Prr60lv_ANAL_grib2.bin'
    if   (iyear >= 2013) and (iyear <= 2019): # うちのアーカイブのディレクトリ構成対処
        filepath = datapath+syear+'/'+smon+'/'+sday+'/'+filename
    elif  iyear >= 2006: # 2020年以後はこちら
        filepath = datapath+syear+'/DATA/'+syear+'/'+smon+'/'+sday+'/'+filename
    elif (iyear == 2005) and (imon == 12) and (iday == 31) and (ihour >= 15):
        filepath = datapath+'2006/DATA/'+syear+'/'+smon+'/'+sday+'/'+filename
    else:
        print('Year '+syear+' is not supported.')
        sys.exit()
# データを読み、単位換算する
    print('Reading ', filepath) # 進行状況表示
    val = load_jmara_grib2(filepath) / 10
    return val
### 解析雨量データを1時刻ぶん読む関数 ここまで ###

### 作図をする関数 ###
def draw_a_contour_map(iyear, imon, iday, rain, lon, lat,\
                       lonmin, lonmax, latmin, latmax):

# 図の準備
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
    ax.set_extent([lonmin, lonmax, latmin, latmax], ccrs.PlateCarree())

# 数値の階級別に色でぬりわける
# (levels, colors の内容は需要に応じて変更する)

##  levels = [-1, 0, 0.5, 1.5, 2.5, 5, 10, 20, 40, 80, 160]
##  colors = ['white', 'khaki', 'gold',\
##            'yellowgreen', 'limegreen', 'springgreen',\
##            'cyan', 'royalblue', 'blue', 'midnightblue']
##  mappable = ax.contourf(lon, lat, rain, levels=levels, colors=colors,\
##                         extend='both', transform=ccrs.PlateCarree())
##  mappable.cmap.set_under('lightpink')
##  mappable.cmap.set_over('purple')

# - cmap_under, cmap_over がきかない状況に対処 (hamana, niigata ではこちら)
##  levels = [-100, -1, 0, 0.5, 1.5, 2.5, 5, 10, 20, 40, 80, 160, 1280] # hourly
    levels = [-100, -1, 0, 1.5, 5, 10, 20, 40, 80, 160, 320, 640, 9600] # daily
    colors = ['lightpink', 'white', 'khaki', 'gold',\
              'yellowgreen', 'limegreen', 'springgreen', \
              'cyan', 'royalblue', 'blue', 'midnightblue', 'purple']
    mappable = ax.contourf(lon, lat, rain, levels=levels, colors=colors,\
                           transform=ccrs.PlateCarree())
# 凡例のカラーバーをかく
    divider = make_axes_locatable(ax)
    ax_cb = divider.new_horizontal(size=0.15, pad=0.1, axes_class=plt.Axes)
    fig.add_axes(ax_cb)
    cb = fig.colorbar(mappable, cax=ax_cb, orientation="vertical", ticks=levels, format="%.0f")
    cb.set_label('mm/day', labelpad=-25, y=1.10, rotation=0)
    for t in cb.ax.get_yticklabels():
        t.set_horizontalalignment('right')
        t.set_x(3.0)
# 海岸線をかく
    if (region == 'zenkoku') or (region == 'hondo'):
        sreso = '50m'
    else:
        sreso = '10m'
    ax.coastlines(resolution=sreso, lw=0.5)
# 都道府県境をかく
    if region != 'zenkoku':
        ax.add_feature(provinces)
# 緯度・経度目盛りをかく
    if   (region == 'zenkoku') or (region == 'hondo'):
        interv      = 5
        intervminor = 1
    elif (region == 'kanto') or (region == 'gkanto'):
        interv      = 1
        intervminor = 0.25
    else:
        interv      = 0.5
        intervminor = 0.1
    xticks      = np.arange(lonmin, lonmax+lolaeps, interv)
    xticksminor = np.arange(lonmin, lonmax+lolaeps, intervminor)
    yticks      = np.arange(latmin, latmax+lolaeps, interv)
    yticksminor = np.arange(latmin, latmax+lolaeps, intervminor)
    ax.tick_params(axis='x', which='both', bottom=True, top=True)
    ax.tick_params(axis='y', which='both', left=True, right=True)
    ax.set_xticks(xticks,      crs=ccrs.PlateCarree(), minor=False)
    ax.set_xticks(xticksminor, crs=ccrs.PlateCarree(), minor=True)
    ax.set_yticks(yticks,      crs=ccrs.PlateCarree(), minor=False)
    ax.set_yticks(yticksminor, crs=ccrs.PlateCarree(), minor=True)
    ax.set_xticklabels(xticks, fontsize=12)
    ax.set_yticklabels(yticks, fontsize=12)
    lon_formatter = LongitudeFormatter(zero_direction_label=True)
    lat_formatter = LatitudeFormatter()
    ax.xaxis.set_major_formatter(lon_formatter)
    ax.yaxis.set_major_formatter(lat_formatter)
# 緯線・経線をかく
    gl = ax.gridlines(crs=ccrs.PlateCarree(), linestyle=':', color='grey', linewidth=1)
    gl.xlocator = mticker.FixedLocator(xticksminor)
    gl.ylocator = mticker.FixedLocator(yticksminor)
# タイトルをかく
    syear = f'{iyear:04}'
    smon  = f'{imon:02}'
    sday  = f'{iday:02}'
    ax.set_title('JMA Radar Raingauge Analyzed Precipitation\n' \
        +syear+' '+smon+' '+sday+' 00-24 JST' )
# 画像出力
    if savefig:
# - 画像をファイルに保存
        imagefilepathname = imagefilepath+'rr'+syear+smon+sday+'jst'
        if region == '':
            imagefilepathname = imagefilepathname + '.png'
        else:
            imagefilepathname = imagefilepathname + '_' + region + '.png'
        print('Writing: ',imagefilepathname) # 進行状況表示
        plt.savefig(imagefilepathname)
        plt.clf()
        plt.close()
    else:
# - 画像を画面に表示
        plt.show() #画面表示
### 作図をする関数 ここまで ###

##### Main Program #####
# パラメータ設定: 実行ごとに変更する可能性が高い変数値の設定をまとめておく
# - 作図範囲名
region  = 'gkanto'
# - 入力データの所在
datapath = "/raid/kaiseki_uryo/" # 解析雨量データが置かれているディレクトリ
# - 作図さき
savefig = False # True:画像ファイル、False:画面
imagefilepath = "image/" # savefig==True のときだけ意味がある
# - range の端をこえるための緯度経度への加算 (単位:度)
lolaeps = 0.01
# パラメータ設定: ここまで

# 解析雨量データの座標情報 (西端の経度、経度間隔、北端の緯度、緯度間隔)
# (ここは作図範囲をかえたときも変更しない)
origlon  = 118.0
deltalon =   1.0/80.0
origlat  =  48.0
deltalat =  -1.0/120.0
nx = 2560
ny = 3360
lon = np.linspace(118, 150, nx, endpoint=False) + 1/80 / 2
lat = np.linspace( 48,  20, ny, endpoint=False) - 1/80/1.5 / 2
# 解析雨量データの座標情報 ここまで

# 作図範囲
if   region == 'zenkoku':
# - 全国 (データのある全域)
    lonmin = 120.0
    lonmax = 150.0
    latmin =  20.0
    latmax =  50.0
elif region == 'hondo':
# - 本土 (四大島をふくむ)
    lonmin = 128.0
    lonmax = 146.0
    latmin =  30.0
    latmax =  46.0
elif region == 'gkanto':
# - 関東を主として中部地方・東北南部に拡大 ("gkanto" ← Greater Kanto)
    lonmin = 136.5
    lonmax = 141.5
    latmin =  34.5
    latmax =  38.5
elif region == 'kanto':
# - 関東
    lonmin = 138.5
    lonmax = 141.5
    latmin =  34.5
    latmax =  37.5
else:
# - 小地域 (緯度経度の値を需要に応じて変更する)
    lonmin = 137.0
    lonmax = 138.0
    latmin =  35.0
    latmax =  36.0

# - 指定した経度・緯度でかこまれた内側にある画素番号
if region == 'zenkoku':
    ix0 = 0
    ix1 = nx-1
    iy0 = 0
    iy1 = ny-1
else:
    ix0 = int(np.ceil(  (lonmin - (origlon + deltalon/2.0) ) / deltalon ) )
    ix1 = int(np.floor( (lonmax - (origlon + deltalon/2.0) ) / deltalon ) )
    iy0 = int(np.ceil(  (latmax - (origlat + deltalat/2.0) ) / deltalat ) )
    iy1 = int(np.floor( (latmin - (origlat + deltalat/2.0) ) / deltalat ) )
# 作図範囲 ここまで
    
# 都道府県境のデータの準備
if region != 'zenkoku':
# - Natural Earthの行政界から日本を抽出したもの
    shapefilepath = '/home/masudako/mapdata/ne/'
    shapefilename = 'ne_10m_admin_1_japan.shp'
# - 国土数値情報の行政界のうち都道府県境だけ
##  shapefilepath = '/home/masudako/mapdata/kokudo/'
##  shapefilename = 'N03-20240101_prefecture.shp'
# - 国土数値情報の行政界全部
##  shapefilepath = '/home/masudako/mapdata/kokudo/'
##  shapefilename = 'N03-20240101.shp'
###
    shapefilepathname = shapefilepath + shapefilename
    province_geos = list(shapereader.Reader(shapefilepathname).geometries())
    provinces = ShapelyFeature(province_geos,
                           crs=ccrs.PlateCarree(),
                           facecolor='none', edgecolor='black', lw=0.25)
# 都道府県境のデータの準備 ここまで

# 対象日時 (日本標準時で指定する)
iyearjst       = 2022
imonjst        =    9
# - 開始日
idayjststart   =    1
# - 終了日
idayjstend     =   30

# 日のループ
for idayjst in range(idayjststart, idayjstend+1):
# 時のループ
    for ihourjst in range(1, 24+1):
# 分のループ (現状では正時だけ読む)
        for iminjst in [0]:
# 日時を世界時に換算する (読みこむファイル名は世界時でつけられている)
            if (ihourjst == 24):
                datetimeutc = datetime(iyearjst, imonjst, idayjst, ihourjst-1, iminjst)\
                            - timedelta(hours=8)
            else:
                datetimeutc = datetime(iyearjst, imonjst, idayjst, ihourjst, iminjst)\
                            - timedelta(hours=9)
            iyear = datetimeutc.year
            imon  = datetimeutc.month
            iday  = datetimeutc.day
            ihour = datetimeutc.hour
            imin  = datetimeutc.minute
#
# 解析雨量データを1時刻ぶん読む
            rain = read_k_u(iyear, imon, iday, ihour, imin, datapath)

# たしこむ (欠損値をたしこまないようにする)
        if(ihourjst == 1):
            rainsum = rain
        else:
            for iy in range(iy0, iy1+1):
                for ix in range(ix0, ix1+1):
                    if(rain[iy,ix] != -1):
                        if( rainsum[iy,ix] == -1):
                            rainsum[iy,ix] = rain[iy,ix]
                        else:
                            rainsum[iy,ix] += rain[iy,ix]
# 24時刻のループはここまで
#
# 作図する
    draw_a_contour_map(iyearjst, imonjst, idayjst,\
            rainsum[iy0:iy1+1, ix0:ix1+1], lon[ix0:ix1+1], lat[iy0:iy1+1],\
            lonmin, lonmax, latmin, latmax)
# ##### End of program #####
