#!/usr/bin/env python
# Be ready for python 3, supported since python 2.6
from __future__ import print_function, division, unicode_literals # gdastool requires at least python version 2.7
import sys
if sys.version_info < (2,7):
    print("You are using python version: ")
    print(" ", sys.version)
    print("gdastool supports python versions 2.7 or later.")
    sys.exit(-1)

from argparse import ArgumentParser, RawTextHelpFormatter
import subprocess
import logging
import os.path
import struct

try:
    import numpy as np
    from scipy.optimize import curve_fit
    from scipy.interpolate import InterpolatedUnivariateSpline
except ImportError as E:
    print("\n!!! ERROR: gdastool requires the scipy and nump packages.\n")
    raise E

try:
    import matplotlib.pyplot as plt
    matplotlibAvailable = True
except:
    matplotlibAvailable = False

# predefined observatory locations
observatories = {'lofar': {'latitude': 52.9, 'longitude': 6.9},
                 'tunka-rex': {'latitude': 51.8, 'longitude': 103.1},
                 'ska-low': {'latitude': -26.7, 'longitude': 116.6},
                 'grand': {'latitude': 42.9, 'longitude': 86.7},
                 'aera': {'latitude': -35.2, 'longitude': -69.3}}

# parsing inputs
parser = ArgumentParser(description='''Creates an atmosphere profile for CORSIKA/CoREAS from GDAS data.

Downloads GDAS model data for the defined location and time and fits
a 5-layer model of the atmosphere to the data. Based on the fit, a
table for the refractive index is created for usage in CoREAS.''',
        epilog='''Please contact
  Pragati Mitra <pragati9163@gmail.com>,
  Arthur Corstanje <a.corstanje@astro.ru.nl> or
  Tobias Winchen <tobias.winchen@rwth-aachen.de>
in case of questions or bugs.''',
        formatter_class=RawTextHelpFormatter)
parser.add_argument("-t", "--utctimestamp", help="UTC time stamp of the event " )
parser.add_argument("-o", "--output", default="ATMOSPHERE.DAT", help="Name of the outputfile.")
cogroup = parser.add_mutually_exclusive_group(required=True)
cogroup.add_argument('--observatory', type=str, choices=observatories.keys(),
        help="Preset of observatory coordinates.")
cogroup.add_argument('-c', '--coordinates', nargs=2, type=float, help='Coordinates of the observatory lat=-90..90 lon=0..360 in deg, e.g. --coordinates 50.85 4.25 for Brussels.')
parser.add_argument('-m', '--minheight', type=float, default=-1E3, help='Mimimum hight for the interpolation. Default is -1.0 km.')
parser.add_argument('-s', '--interpolationSteps', default=1.0, type=float, help='Step length for interpolation. Default is 1 m. ')

parser.add_argument("-p", "--gdaspath", default='.', help="path to local gdas file directory. If required file is not there, it will be downloaded.")
parser.add_argument('-v', '--verbose', default=0, action='count', help='Set log level, -vv for debug.')
if matplotlibAvailable:
    parser.add_argument('-g', '--createplot', default=False, action='store_true', help='plot density profile.')


options = parser.parse_args()
from datetime import datetime
utcstring = str(datetime.utcfromtimestamp(int(options.utctimestamp)))
import dateutil.parser
try:
    date_time = dateutil.parser.parse(utcstring)
except:
    print('ERROR: Cannot parse date: ' + ' '.join(options.datetime))
    sys.exit(-1)

if options.observatory in observatories:
    options.coordinates = [observatories[options.observatory]['latitude'],
            observatories[options.observatory]['longitude']]
print('')
print('    Coordinates: lat = {:+.2f} deg, lon = {:+.2f} deg'.format(options.coordinates[0], options.coordinates[1]))
print('     Time [UTC]: {}'.format(date_time.ctime()))

# Round-off of time to nearest 3-hour time 
seconds = date_time.hour * 3600 + date_time.minute * 60 + date_time.second
delta_seconds = round(seconds / (3.0*3600)) * 3*3600 - seconds
delta_seconds = int(delta_seconds) # Number of seconds to add to get onto a 3-hour grid
from datetime import timedelta
date_time = date_time + timedelta(seconds=delta_seconds) # Using this date&time in what follows
print(' ')
print('     Time on 3-hour grid [UTC]: {}'.format(date_time.ctime()))


month_name = ["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"]
part = 5
if (date_time.day < 29):
    part -= 1
if (date_time.day < 22):
    part -= 1
if (date_time.day < 15):
    part -= 1
if (date_time.day < 8):
    part -= 1

year_gdas = str(date_time.year)
gdasname = "gdas1.{0}{1}.w{2}".format(month_name[int(date_time.month) - 1], year_gdas[2:], part)
print('Using GDAS file:', gdasname)

gdaspath = os.path.abspath(options.gdaspath)
if not os.path.isdir(gdaspath):
    logging.info("Creating nonexisting local gdas directory {}.".format(gdaspath))
    os.mkdir(gdaspath)

if os.path.isfile(os.path.join(gdaspath, gdasname)):
    print("Found {} in {}, no download.".format(gdasname, gdaspath))
else:
    gdasurl = 'ftp://arlftp.arlhq.noaa.gov/pub/archives/gdas1'
    print("File not found in {}, download from {}".format(gdaspath, gdasurl))
    cmd = ['wget', '--directory-prefix={}'.format(gdaspath), '{}/{}'.format(gdasurl, gdasname)]
    #if options.verbose == 0:
    #    cmd.append('--no-verbose')
    #    cmd.append('--show-progress')
    v = subprocess.call(cmd)
    if not v == 0:
        print('ERROR DOWNLOADING {}/{} -- ABORTING!'.format(gdasurl, gdasname))
        sys.exit(-1)

# converting to gdas coordinates
time_gdas = int(date_time.hour)
day_gdas = int(date_time.day)
month_gdas = int(date_time.month)
lat = 90 + int(round(options.coordinates[0]))
lon = int(round(options.coordinates[1]))
altitude = np.zeros([24])
temp = np.zeros([24])
relh = np.zeros([24])
pressure = np.array([0, 1000, 975, 950, 925, 900, 850, 800, 750, 700,
                    650, 600, 550, 500, 450, 400, 350, 300, 250, 200, 150, 100, 50, 20])
RI = np.zeros([24])
h1 = int(np.floor(int(time_gdas) / 3.0) * 3)
h2 = (h1 % 2) * 3
string_id = str(int(year_gdas[2:])).rjust(2) + str(month_gdas).rjust(2) + str(day_gdas).rjust(2) + str(h1).rjust(2) + str(h2).rjust(2)
logging.debug("string id = {}".format(string_id))
f = open(os.path.join(gdaspath, gdasname), 'rb')
s = f.read().decode('latin_1')
f.close()
skip = 0
start = skip + s[skip:].find("INDX")
block = s[start:start + 2000]
nx, ny, nz = int(block[129:132]), int(block[132:135]), int(block[136:138])
nxy = nx * ny

tmp_data = np.zeros([ny, nx])

skip = 0
found = True


for j in range(250):
    if j % 25 == 0:
        print('\rParsing gdas file: [{:10}]'.format(''.join(['.' for i in range(j // 25 + 1)])), end='') # At least python 2.7 is required
        sys.stdout.flush()
    next = s[skip:].find(str(string_id))
    if (next < 0):
        found = False
    start = skip + next
    block = s[start:start + 1000]
    if (found):
        lvl, keyword, nexp, precision, value = int(block[10:12]), block[
            14:18], int(block[20:23]), float(block[23:36]), float(block[36:50])
        scale = 2.0 ** (7 - nexp)
        if (keyword == 'PRSS' or keyword == 'RH2M' or keyword == 'SHGT' or keyword == 'T02M'
           or keyword == 'HGTS' or keyword == 'TEMP' or keyword == 'RELH'):
            datablock = bytes(s[start + 50:start + nxy + 50].encode('latin_1'))
            diffs = (
                (np.array([struct.unpack('65160B', datablock)]) - 127) / scale)[0]
            vold = value
            indx = 0
            for k in np.arange(0, ny):
                for l in np.arange(0, nx):
                    tmp_data[k, l] = vold + diffs[indx]
                    indx += 1
                    vold = tmp_data[k, l]
                vold = tmp_data[k, 0]
        if (keyword == 'PRSS'):
            logging.debug("Found PRSS block at ".format(start))
            pressure[0] = tmp_data[lat, lon]
        elif (keyword == 'SHGT'):
            logging.debug("Found SHGT block at".format(start))
            altitude[0] = tmp_data[lat, lon]
        elif (keyword == 'RH2M'):
            logging.debug("Found RH2M block at".format(start))
            relh[0] = tmp_data[lat, lon]
        elif (keyword == 'T02M'):
            logging.debug("Found TO2M block at".format(start))
            prss_data = tmp_data
            temp[0] = tmp_data[lat, lon]  # temperature
        elif (keyword == 'HGTS'):
            logging.debug("Found HGTS block for level {} at {}".format(lvl, start))
            altitude[lvl] = tmp_data[lat, lon]  # altitude
        elif (keyword == 'TEMP'):
            logging.debug("Found TEMP block for level {} at {}".format(lvl, start))
            temp[lvl] = tmp_data[lat, lon]
        elif (keyword == 'RELH'):
            logging.debug("Found RELH block for level {} at {}".format(lvl, start))
            relh[lvl] = tmp_data[lat, lon]  # relative humidity
        skip = start + 100
print('')

# calculating atmospheric variables

def geopot_to_geometric(lat,h):
    # Formula from Abreu et al., Description of Atmospheric Conditions at the Pierre Auger Observatory using the Global Data Assimilation System (GDAS)
    # Orig. link Mahoney et al., 2008
    z = (1+0.002644*np.cos(2*lat*np.pi/ 180. )) * h + (1+0.0089*np.cos(2*lat*np.pi/180.)) * (h * h / 6245000.)
    return z

alt_geopot = altitude
alt_geometric = geopot_to_geometric(lat-90,alt_geopot)
altitude = alt_geometric  # setting altitude to geometric
pressure = 100 * pressure  # hPa to Pa
tempC = temp - 273.15
part_press = np.zeros([24])
# Warning: Magnus formula for partial pressure of water vapor is designed to return hectopascals
# So, the additional factor 100 converts to pascals which are used here; (relh / 100) is the relative humidity fraction
for j in np.arange(24):
    if (tempC[j] < 0):
        part_press[j] = (relh[j] / 100.) * 100 * 6.1064 * \
            np.exp(21.88 * tempC[j] / (265.5 + tempC[j]))  # Pa
    else:
        part_press[j] = (relh[j] / 100.) * 100 * 6.1070 * \
            np.exp(17.15 * tempC[j] / (234.9 + tempC[j]))  # Pa

M_dry = 0.02897     # Molar mass of dry air, in kg/mol
M_water = 0.01802   # Molar mass of water, idem
M_CO2 = 0.04401     # Molar mass of CO_2, idem

phi_CO2 = 385.0 * 1e-6 # CO2 fraction (from ppm-volume ~ 385)
phi_water = part_press / pressure
phi_dry = 1 - phi_water - phi_CO2
M_air = phi_dry * M_dry + phi_water * M_water + phi_CO2 * M_CO2
density = (pressure * M_air / temp / 8.31451)  # kg m^-3; Gas constant R = 8.31451 J/(mol K)
density = density / 1000  # g / cm^3
RI = np.zeros([24])
pressure_dry = pressure - part_press
# Refractivity formula from Rueger, Refractive index formulae for radio waves, Proceedings of FIG XXII International Congress, 2002
# Designed to take pressure in hPa as input, hence divide by 100 below.
RI = 77.689 * (pressure_dry / temp) + 71.2952 * (part_press / temp) + 375463 * (part_press / (temp * temp))
RI = RI / 100  #  for converting p from hpa to pa

# interpolation of refractivity

alt_ground = altitude[0]
alt_max = altitude[23]
alt_values = np.arange(alt_ground, alt_max, options.interpolationSteps)

RI_interp = np.zeros([len(alt_values)])
interpolation = InterpolatedUnivariateSpline(altitude, np.log(RI), k=1)
RI_interp = interpolation(alt_values)
RI_interp = np.exp(RI_interp)
refractiveIndex = (RI_interp * 1e-6 + 1)

# fitting to corsika atmosthere model


def fn_atm_depth(x, par1, par2, par3):
    return par1 + par2 * np.exp(-1e5 * x / par3)


def a_from_atmdep(atmdep, x, b, c):
    return atmdep - b * np.exp(-1e5 * x / c)


def Density(x, b, c):
    return b / c * np.exp(-1e5 * x / c)


def find_par_b(rho, c, x0):
    return rho * c * np.exp(1e5 * x0 / c)


def fit_lay(x, rho, x0, c):
    b = find_par_b(rho, c, x0)
    return Density(x, b, c)

def rms(x):
    return np.sqrt(np.dot(x,x) / len(x))


boun1 = 10
boun2 = 17
altitude = altitude / 1000  # to km
alt_bc1=altitude[boun1]
alt_bc2=altitude[boun2]
alt_bc3=altitude[23]
x_lay1 = altitude[ :boun1 + 1]
x_lay2 = altitude[boun1:boun2 + 1]
x_lay3 = altitude[boun2:]
x_lay4=[alt_bc3, 100]
den_lay1 = density[ :boun1 + 1]
den_lay2 = density[boun1:boun2 + 1]
den_lay3 = density[boun2:]
den_lay4 = [density[23], 1e-09]

# fitting the linear equation for density profile
ans1 = np.zeros([2])
ans2 = np.zeros([2])
ans3 = np.zeros([2])
mat1 = np.zeros([2, 2])
mat2 = np.zeros([2, 2])
mat3 = np.zeros([2, 2])
mat1 = np.array([[altitude[1], 1], [altitude[boun1], 1]])
cons1 = [np.log(density[1]), np.log(density[boun1])]
ans1 = np.linalg.solve(np.asarray(mat1), np.asarray(cons1))
mat2 = np.array([[altitude[boun1], 1], [altitude[boun2], 1]])
cons2 = [np.log(density[boun1]), np.log(density[boun2])]
ans2 = np.linalg.solve(np.asarray(mat2), np.asarray(cons2))
mat3 = np.array([[altitude[boun2], 1], [altitude[23], 1]])
cons3 = [np.log(density[boun2]), np.log(density[23])]
ans3 = np.linalg.solve(np.asarray(mat3), np.asarray(cons3))
mat4 = np.array([[altitude[23], 1], [100, 1]])
cons4 = [np.log(density[23]), np.log(1e-09)]
ans4 = np.linalg.solve(np.asarray(mat4), np.asarray(cons4))
A1 = ans1[0]
A2 = ans2[0]
A3 = ans3[0]
A4 = ans4[0]
B1 = np.log(density[boun1]) - A1 * alt_bc1
C1 = -1e5 / A1
C2 = -1e5 / A2
C3 = -1e5 / A3
C4 = -1e5 / A4
b1 = C1 * np.exp(B1)
coeff1, cov1 = curve_fit(Density, np.asarray(x_lay1), np.asarray(den_lay1), np.array([b1, C1]))
b1_new = coeff1[0]
c1_new = coeff1[1]
den_bc1 = Density(alt_bc1, b1_new, c1_new)



def fit_lay2(x, c):
    return fit_lay(x, den_bc1, alt_bc1, c)

coeff2, cov2 = curve_fit(fit_lay2, np.asarray(x_lay2), np.asarray(den_lay2), C2)
c2_new = coeff2[0]
b2_new = find_par_b(den_bc1, c2_new, alt_bc1)
den_bc2 = Density(alt_bc2, b2_new, c2_new)


def fit_lay3(x, c):
    return fit_lay(x, den_bc2, alt_bc2, c)

coeff3, cov3 = curve_fit(fit_lay3, np.asarray(x_lay3), np.asarray(den_lay3), C3)
c3_new = coeff3[0]
b3_new = find_par_b(den_bc2, c3_new, alt_bc2)
den_bc3 = Density(alt_bc3, b3_new, c3_new)


def fit_lay4(x, c):
    return fit_lay(x, den_bc3, alt_bc3, c)

coeff4, cov4 = curve_fit(fit_lay4, np.asarray(x_lay4), np.asarray(den_lay4), C4)
c4_new = coeff4[0]
b4_new = find_par_b(den_bc3, c4_new, alt_bc3)
den2_bc3 = Density(alt_bc3, b4_new, c4_new)
xmax = 0.01128292 - 1e-09 * 100. * 100000.  # mass overburden at the boundary of 4th layer
a4 = a_from_atmdep(xmax, 100, b4_new, c4_new)
atmdep_bc3 = fn_atm_depth(alt_bc3, a4, b4_new, c4_new)
a3 = a_from_atmdep(atmdep_bc3, alt_bc3, b3_new, c3_new)
atmdep_bc2 = fn_atm_depth(alt_bc2, a3, b3_new, c3_new)
a2 = a_from_atmdep(atmdep_bc2, alt_bc2, b2_new, c2_new)
atmdep_bc1 = fn_atm_depth(alt_bc1, a2, b2_new, c2_new)
a1 = a_from_atmdep(atmdep_bc1, alt_bc1, b1_new, c1_new)
a = [a1, a2, a3, a4, 0.01128292]
b = [b1_new, b2_new, b3_new, b4_new, 1]
c = [c1_new, c2_new, c3_new, c4_new, 1e9]

# fit quality and plotting

den1_simu = Density(x_lay1 , b1_new , c1_new)
den2_simu = Density(x_lay2[1:] , b2_new , c2_new)
den3_simu = Density(x_lay3[1:], b3_new , c3_new)
den_simu = den1_simu.tolist() + den2_simu.tolist() + den3_simu.tolist()
den_data = density.tolist()
diff_density_rel = ( np.asarray(den_simu) - np.asarray(den_data)) / np.asarray(den_simu)
density_rms = rms(diff_density_rel)
if matplotlibAvailable:
    if(options.createplot):
        alt_simu_gdas=np.arange(0,altitude[23],0.2)
        rho_gdas = np.zeros([len(alt_simu_gdas )])
        for k in np.arange(len(alt_simu_gdas)):
            par=np.zeros([3])
            if(alt_simu_gdas[k]<altitude[boun1]):
                par=[a1,b1_new,c1_new]
            if(altitude[boun1]<=alt_simu_gdas[k]<altitude[boun2]):
                par=[a2,b2_new,c2_new]
            if(altitude[boun2]<=alt_simu_gdas[k]<altitude[23]):
                par=[a3,b3_new,c3_new]
            if(altitude[23]<=alt_simu_gdas[k]<=100):
                par=[a4,b4_new,c4_new]
            rho_gdas[k]=Density(alt_simu_gdas[k],par[1],par[2])

        plt.plot(altitude*1000,density,'bo',label='data')
        plt.plot(alt_simu_gdas*1000,rho_gdas,'r-',label='fit')
        plt.title("GDAS density profile with 5-layer fit")
        plt.xlabel("altitude (m)" , fontsize=20)
        plt.ylabel("air density (g/cm$^3$)" ,fontsize=20)
        plt.legend(loc="upper right")
        plt.tight_layout()
        plt.show()

# interpolation below ground
lowest_altitude = alt_ground - options.interpolationSteps * np.round((alt_ground - options.minheight) / options.interpolationSteps) # Make sure the altitudes below and above ground match seamlessly.

alt_values_extended = np.arange(lowest_altitude, alt_ground, options.interpolationSteps)
density_underground = Density(alt_values_extended/1000,b1_new,c1_new)
density_ground = den_lay1[0]
RI_ground = RI[0]
RI_interp_underground = RI_ground *density_underground / density_ground
refractiveIndex_underground = (RI_interp_underground * 1e-6 + 1)
refractiveIndex_total = refractiveIndex_underground.tolist() + refractiveIndex.tolist()
alt_values_total = alt_values_extended.tolist() + alt_values.tolist()

print('\nModel parameters:')
for i in range(5):
    print("{}: a = {: .4E}   b = {: .4E}   c = {: .4E}".format(i, a[i], b[i], c[i]))
print("Layers: L5 = {}, L4 = {}, L3 = {}, L2 = {}".format(100, altitude[23], altitude[boun2], altitude[boun1]))
print("Measure of fit quality: rms density value of relative difference between data and fit= {}"  .format(density_rms))
print(" ")
# writing in to file
f = open(options.output, 'w')
f.write("# atmospheric parameters ATMLAY, A, B, C respectively\n")
# write 0 layer here that is not changeable in corsika file
f.write('{: .8E} {: .8E} {: .8E} {: .8E} {: .8E}\n'.format(0,
    altitude[boun1]*1.0e5, altitude[boun2]*1.0e5, altitude[23]*1.0e5,
    100.0*1.0e5 ))
f.write('{: .8E} {: .8E} {: .8E} {: .8E} {: .8E}\n'.format(*a))
f.write('{: .8E} {: .8E} {: .8E} {: .8E} {: .8E}\n'.format(*b))
f.write('{: .8E} {: .8E} {: .8E} {: .8E} {: .8E}\n'.format(*c))
f.write("# atmospheric height [m] and refractive index columns \n")
for i in np.arange(len(refractiveIndex_total)):
    f.write("{: .8E} {: .8E}\n".format(alt_values_total[i], refractiveIndex_total[i]))
f.close()
print('Output written to: {}'.format(options.output))