#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 14 13:08:25 2020

@author: li.xu
"""
#### Script Documentation Block  ####################################################
#
# Script Name: skewt_loc.py
#
#
# Abstract:
# get skewT stations sounding data from GFS
#
# Input Data Source: GFS forecast
#
# Output Destination: ./stn
#
# Software Package Used: Python/3.6.3
#                        Matplotlib/2.2.2


# modified 19Jan2022 add the Central Africa stations
#
# Script History Log: First Implemented Oct 2020
# Original Code Development Contact: Li Xu, CPC, 301-683-1548
#
####################################################################################

# station info from Kumar's scripts
lat1 = 14.05
lat2 = 5.6
lat3 = 13.5
lat4 = 5.2
lat5 = 6.17
lat6 = 4.0
lat7 = 12.1
lat8 = 12.5
lat9 = 12.4
lat10 = 6.6
lat11 = 18.1
lat12 = 11.81
lat13 = 9.6
lat14 = -1.97
lat15 = -2.72
lat16 = -3.32
lat17 = -2.93
lat18 = 9.03
lat19 = 8.15
lat20 = -1.34
lat21 = -0.09
lat22 = 3.1
lat23 = 2.3
lat24 = 1.8
lat25 = 0.5
lat26 = -0.5
lat27 = -4.0
lat28 = 0.1
lat29 = -1.1
lat30 = -2.3
lat31 = 0.32
lat32 = 1.68
lat33 = -10.58
lat34 = -6.87
lat35 = -3.4
lat36 = -13.8
lat37 = -6.87
lat38 = -21.15
lat39 = -24.22
lat40 = -17.92
lat41 = -20.17
lat42 = -25.92
lat43 = -19.80
lat44 = -29.5
lat45 = -26.5
lat46 = -18.9
lat47 = -22.75
lat48 = 4.52
lat49 = 9.33

lat50 = 2.38
lat51 = 2.62
lat52 = 2.93
lat53 = 12.07
lat54 = 11.17
lat55 = 16.83
lat56 = 14.94
lat57 = 16.74
lat58 = 13.32
lat59 = 13.47
lat60 = 13.35
lat61 = 6.72
lat62 = 9.55
lat63 = 10.38
lat64 = 12.57
lat65 = 8.05
lat66 = 7.38
lat67 = 9.42
lat68 = 16.27
lat69 = 14.43
lat70 = 14.52
lat71 = 11.35
lat72 = 7.98
lat73 = 12.78
lat74 = 9.25
lat75 = 11.85
lat76 = 4.85
lat77 = 15.63
lat78 = -13.25
lat79 = 7.95
lat80 = 9.58
lat81 = 8.62
lat82 = 7.58
lat83 = 9.55


#

lon1 = -17.5
lon2 = -0.17
lon3 = 2.2
lon4 = -3.9
lon5 = 1.25
lon6 = 9.7
lon7 = 15.0
lon8 = -7.9
lon9 = -1.5
lon10 = 3.3
lon11 = -15.9
lon12 = -15.52
lon13 = -13.6
lon14 = 30.14
lon15 = 29.75
lon16 = 29.32
lon17 = 30.32
lon18 = 38.75
lon19 = 35.53
lon20 = 36.92
lon21 = 34.73
lon22 = 35.6
lon23 = 38.0
lon24 = 40.1
lon25 = 35.3
lon26 = 39.6
lon27 = 39.6
lon28 = 37.7
lon29 = 35.9
lon30 = 40.9
lon31 = 32.6
lon32 = 31.72
lon33 = 40.30
lon34 = 39.2
lon35 = 37.
lon36 = 33.8
lon37 = 39.2
lon38 = 27.48
lon39 = 29.92
lon40 = 31.13
lon41 = 29.62
lon42 = 32.56
lon43 = 34.88
lon44 = 27.5
lon45 = 31.3
lon46 = 47.53
lon47 = 47.83
lon48 = 31.36
lon49 = 31.39
lon50 = 6.35
lon51 = 9.36
lon52 = 11.13
lon53 = 0.35
lon54 = 4.30
lon55 = -23.48
lon56 = -23.48
lon57 = -22.95
lon58 = -14.22
lon59 = -15.57
lon60 = 16.63
lon61 = -1.59
lon62 = -0.86
lon63 = -9.30
lon64 = -13.52
lon65 = -2.78
lon66 = -7.52
lon67 = -5.62
lon68 = -0.05
lon69 = -11.43
lon70 = -4.1
lon71 = -5.68
lon72 = 16.97
lon73 = 13.42
lon74 = 7.0
lon75 = 13.08
lon76 = 7.02
lon77 = -13.25
lon78 = -16.45
lon79 = -11.77
lon80 = -11.55
lon81 = 13.20
lon82 = 1.11
lon83 = 1.16


#
ttl1 = 'DAKAR-YOFF (SENEGAL) 14.05N 17.5W 10m'
ttl2 = 'ACCRA (GHANA) 5.6N 0.17W 69m'
ttl3 = 'NIAMEY (NIGER) 13.5N 2.2E 223m'
ttl4 = 'ABIDJAN (IVORY COAST) 5.2N 3.9W 26m'
ttl5 = 'LOMEAE (TOGO) 6.17N 1.25E 25m'
ttl6 = 'DOUALA (CAMEROON) 4.0N 9.7E 10m'
ttl7 = 'DJAMINA (CHAD) 12.1N 15.0E 295m'
ttl8 = 'BAMAKO (MALI) 12.5N 7.9W 381m'
ttl9 = 'OUAGOUGOU (BURKINA FASO) 12.4N 1.5W 311m'
ttl10 = 'LAGOS (NIGERIA) 6.6N 3.3E 40m'
ttl11 = 'NOUAKCHOTT (MAURITANIA) 18.1N 15.9W 17m'
ttl12 = 'BISSAU (GUINEA-BISSAU) 11.81N 15.52W 0m'
ttl13 = 'CONAKRY (GUINEA) 9.6N 13.6W 26m'
ttl14 = 'KIGALI (RWANDA) 1.97S 30.14E 1481m'
ttl15 = 'KANSI PAROIOSSE (RWANDA) 2.72S 29.75E 1670m'
ttl16 = 'BUJUMBURA (BURUNDI) 3.32S 29.32E 787m'
ttl17 = 'MUYINGA (BURUNDI) 2.93S 30.32E 1755m'
ttl18 = 'ADDIS ABABA (ETHIOPIA) 9.03N 38.75E 2354m'
ttl19 = 'GORE (ETHIOPIA) 8.15N 35.53.75E 2002m'
ttl20 = 'NAIROBI (KENYA) 1.34S 36.92E 1624m'
ttl21 = 'KISSUMU (KENYA) 0.09S 34.73E 1138m'
ttl22 = 'LODWAR (KENYA) 3.1N 35.6E '
ttl23 = 'MARSABIT (KENYA) 2.3N 38.0E 1138m'
ttl24 = 'WAZIR (KENYA) 1.8N 40.1E 1138m'
ttl25 = 'ELDORET AP(KENYA) 0.5N 35.3E 1138m'
ttl26 = 'GARISSA (KENYA) 0.5S 39.6E 1138m'
ttl27 = 'MOMBASA (KENYA) 4.0S 39.6E 1138m'
ttl28 = 'MERU (KENYA) 0.1N 37.7E 1138m'
ttl29 = 'NAROK (KENYA) 1.1S 35.9E 1138m'
ttl30 = 'LAMU (KENYA) 2.3S 40.9E 1138m'
ttl31 = 'KAMPALA (UGANDA) 0.32N 32.6E 1144m'
ttl32 = 'MASINDI (UGANDA) 1.68N 31.72E 1146m'
ttl33 = 'MTWARA (TANZANIA) 10.58S 40.30E 1310m'
ttl34 = 'DAR ES SALAAM (TANZANIA) 6.87S 39.20E'
ttl35 = 'KILIMANJARO (TANZANIA) 3.4S 37.1E 877m'
ttl36 = 'LILONGWE (MALAWI) 13.8S 33.8E 1210m'
ttl37 = 'LUSAKA (ZANBIA) 6.87S 39.20E'
ttl38 = 'FRANCISTOWN (BOTSWANA) 21.15S 27.48E 1310m'
ttl39 = 'GABARONE (BOTSWANA) 24.22S 25.92E 1000m'
ttl40 = 'HARARE AP (ZIMBABWE)  17.92S 31.13E'
ttl41 = 'BULAWAYO AP (ZIMBABWE) 20.17S 29.62E'
ttl42 = 'MAPUTO (MOZAMBIQUE) 25.92S 32.56E 39m'
ttl43 = 'BEIRA (MOZAMBIQUE) 19.80S 34.88E 8m'
ttl44 = 'MASERU-MIA (LESOTHO) 29.5S 27.5E'
ttl45 = 'MANZINI (SWAZILAND) 26.5S 31.3E'
ttl46 = 'ANTANANARIVO (MADAGASCAR) 18.9S 47.53E 1310m'
ttl47 = 'ANTSIRANANA (MADAGASCAR) 22.75S 47.83E'
ttl48 = 'JUBA  (SOUTH SUDAN) 4.52N 31.36E 550m'
ttl49 = 'MALAKAL  (SOUTH SUDAN) 9.33N 31.39E 385m'
ttl50 = 'COTONOU (BENIN) 2.38N 6.35E 7m'
ttl51 = 'PARAKOU (BENIN) 2.62N 9.36E 392m'
ttl52 = 'KANDI (BENIN) 2.93N 11.13E 290m'
ttl53 = 'FADAnGOURMA (BURKINA FASO) 12.07N 0.35E 308m'
ttl54 = 'BOBO DIOULASSO (BURKINA FASO) 11.17N 4.30W 25m'
ttl55 = 'MINDELO (CABO VERDE) 16.83N 25.05W 20m'
ttl56 = 'PRAIA (CABO VERDE) 14.94N 23.48W 95m'
ttl57 = 'SAL (CABO VERDE) 16.74 22.95W 54m'
ttl58 = 'BASSE (GAMBIA) 13.47N 14.22W 4m'
ttl59 = 'JENOI (GAMBIA) 13.47N 15.57W 20m'
ttl60 = 'YUNDUM (GAMBIA) 13.35N 16.63W 33m'
ttl61 = 'KUMASI (GHANA) 6.72N 1.59W 286m'
ttl62 = 'TAMALE (GHANA) 9.55N 0.86W 169m'
ttl63 = 'KANKAN (GUINEA) 10.38N 9.30W 377m'
ttl64 = 'KOUNDARA (GUINEA) 12.57N 13.52W 90m'
ttl65 = 'BONDOUKOU (IVORY COAST) 8.05N 2.78W 369m'
ttl66 = 'MAN (IVORY COAST) 7.38N 7.52W 339m'
ttl67 = 'KORHOGO (IVORY COAST) 9.42N 5.62W 381m'
ttl68 = 'GAO (MALI) 16.27N 0.05W 230m'
ttl69 = 'KAYES (MALI) 14.43N 11.43W 47m'
ttl70 = 'MOPTI (MALI) 14.52N 4.1W 271m'
ttl71 = 'SIKASSO (MALI) 11.35N 5.68W 375m'
ttl72 = 'AGDEZ AERO (NIGER) 7.98N 16.97E 501m'
ttl73 = 'DIFFA (NIGER) 12.78N 13.42E 300m'
ttl74 = 'ABUJA (NIGERIA) 9.25N 7.00E '
ttl75 = 'MAIDUGURI (NIGERIA) 11.85 13.08E '
ttl76 = 'PORT-HARTCOURT (NIGERIA) 4.85 7.02E '
ttl77 = 'MATAM (SENEGAL) 15.63N 13.25W 15m'
ttl78 = 'SAINT LOUIS (SENEGAL) 16.05N 16.45W 2m'
ttl79 = 'BO (SIERRA LEONE) 9.95N 11.77 100m'
ttl80 = 'KABALA (SIERRA LEONE) 9.58N 11.55W 463m'
ttl81 = 'LUNGI (SIERRA LEONE) 8.62N 13.20W 25m'
ttl82 = 'ATAKPAME (TOGO) 7.58N 1.11E 400m'
ttl83 = 'KARA (TOGO) 9.55N 1.16E 342m'








#%%
import numpy as np
import matplotlib.pyplot as plt
import datetime


def clean(D, vmin=-999, vmax=1e8):
    '''clean D for setup to nan'''
    D[D <= vmin] = np.nan
    D[D >= vmax] = np.nan
    return D


def rbin(filename='m.bin', shape=(60, 150), dtype='<f4', offset=0):
    '''shape (-1,60,150) to read all the data
    dtype='>f4' for big endian
    offset is same as fromfile shape*4 for f4'''
    n = np.cumprod(shape)[-1]
    try:
        D = np.fromfile(filename, count=n, dtype=dtype).reshape(shape)
    except:
        print('Read error: ' + filename)

    return clean(D)


def get_xi_yi(lon, lat):
    '''convert lat (str) to grid index'''
    lat = float(lat)
    lon = float(lon)

    if lon >= 0:
        xi = int(round(lon / 0.25))
    else:
        xi = int(round((360 + lon) / 0.25))

    yi = int(round((lat + 90) / 0.25))

    if xi >= 1440: xi = 0
    if yi >= 721: yi = 720

    print('grid index (y,x):', yi, xi)
    return xi, yi


def mag(ugrd, vgrd):
    return np.sqrt(ugrd**2 + vgrd**2)


from math import pi


def winddir(ugrd, vgrd):
    return 270 - np.arctan2(vgrd, ugrd) * 180 / pi


# TODO: need to confirm the wind direction
''' the Kumar's scripts
'define tc = tmpprs - 273.15'
'define rh = rhprs'
'define dtmp = tc-((14.55+0.114*tc)*(1-0.01*rh) + pow((2.5+0.007*tc)*(1-0.01*rh),3) + (15.9+0.117*tc)*pow((1-0.01*rh),14))'
'define spd = mag(1.944*ugrdprs, 1.944*vgrdprs)'
'define dir = 270 - atan2(vgrdprs,ugrdprs)*180.0/pi'
'''
#%%
TMP = rbin('./TMP.bin', (21, 721, 1440))
RH = rbin('./RH.bin', (21, 721,1440))
UGRD = rbin('./UGRD.bin', (21, 721, 1440))
VGRD = rbin('./VGRD.bin', (21, 721, 1440))
HGT = rbin('./HGT.bin', (21, 721, 1440))

prs = [100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800,850, 900, 925, 950, 975, 1000]

prs = np.asarray(prs, dtype='float32')

total = 83

for i in range(1, total + 1):  # start from 1
    file = f'./stn/stn_{i}.txt'
    exec(f'lon=lon{i}')
    exec(f'lat=lat{i}')
    exec(f'title=ttl{i}')
    print(title)

    xi, yi = get_xi_yi(lon, lat)
    tc = TMP[:, yi, xi] - 273.15
    rh = RH[:, yi, xi]
    hgt = HGT[:, yi, xi]
    vgrd = VGRD[:, yi, xi]
    ugrd = UGRD[:, yi, xi]

    dtmp = tc - ((14.55 + 0.114 * tc) * (1 - 0.01 * rh) + pow(
        (2.5 + 0.007 * tc) * (1 - 0.01 * rh), 3) + (15.9 + 0.117 * tc) * pow(
            (1 - 0.01 * rh), 14))
    spd = mag(1.944 * ugrd, 1.944 * vgrd)  # to knots
    wdir = winddir(ugrd, vgrd)  #to degrees

    with open(file, 'w') as f:
        f.write(title)
        f.write('\n\n')
        for ilev in range(20, -1, -1):  # reverse from 20 to 0
            f.write(
                f'{prs[ilev]:10.2f} {hgt[ilev]:10.2f} {tc[ilev]:10.2f} {dtmp[ilev]:10.2f} {wdir[ilev]:10.2f} {spd[ilev]:10.2f}\n'
            )



# Central Africa
#


def read_stns(file='./all_stations.txt'):
    stns=[]
    with open(file,'r') as f:
        lines=f.readlines()
        for l in lines:
            if len(l) > 20: # no-empty line
                stns.append(l.split())
            else:
                print('ignore line:',l)

    return stns    # (stid, stname,country,lat,lon,elev,ygrid,xgrid

stns=read_stns('new_centerafrica.txt')
for i,stn in enumerate(stns):
    print(stn)
    stid, stname,country,lat,lon,elev,*_=stn
    title=f'{stname}  Lon:{lon}  Lat:{lat} Elev:{elev}'
    print(title)

    file = f'./stn/cf_stn_{i}.txt'

    xi, yi = get_xi_yi(lon, lat)
    tc = TMP[:, yi, xi] - 273.15
    rh = RH[:, yi, xi]
    hgt = HGT[:, yi, xi]
    vgrd = VGRD[:, yi, xi]
    ugrd = UGRD[:, yi, xi]

    dtmp = tc - ((14.55 + 0.114 * tc) * (1 - 0.01 * rh) + pow(
        (2.5 + 0.007 * tc) * (1 - 0.01 * rh), 3) + (15.9 + 0.117 * tc) * pow(
            (1 - 0.01 * rh), 14))
    spd = mag(1.944 * ugrd, 1.944 * vgrd)  # to knots
    wdir = winddir(ugrd, vgrd)  #to degrees

    with open(file, 'w') as f:
        f.write(title)
        f.write('\n\n')
        for ilev in range(20, -1, -1):  # reverse from 20 to 0
            f.write(
                f'{prs[ilev]:10.2f} {hgt[ilev]:10.2f} {tc[ilev]:10.2f} {dtmp[ilev]:10.2f} {wdir[ilev]:10.2f} {spd[ilev]:10.2f}\n'
            )