import numpy as np
import pygrib
import sys

if len(sys.argv) == 7:
    inputFile = sys.argv[1]
    varName = sys.argv[2]
    model = sys.argv[6]
    if '.' in inputFile:
        extension = inputFile.split('.')[-1]
    else: 
        extension = ''
    if extension == 'gr2':
        if model == 'WRF-12km' or model == 'WRF-3km':
            with pygrib.open(inputFile) as grbfile:
                grb = grbfile.select(parameterName='Total precipitation', stepRange=f'{int(sys.argv[5])-24}-{sys.argv[5]}')[0]
                met_data, lats, lons = grb.data()
                msg = grb.tostring()
                
                outFile = '/'.join(inputFile.split('/')[0:5])+'/processed/'+'/'.join(inputFile.split('/')[6:9])+'/tp_24h/'+inputFile.split('/')[-1]
                grbout = open(outFile, 'wb')
                grbout.write(msg)
                grbout.close()
                
                attrs = {}
                attrs['valid'] = sys.argv[3]
                attrs['init'] = sys.argv[4]
                attrs['lead'] = sys.argv[5]+'0000'
                attrs['accum'] = '240000'
                attrs['name'] = varName
                attrs['long_name'] = ' '.join(str(grb).split(':')[1].split(' ')[0:2])
                attrs['level'] = 'A24' 
                attrs['units'] = ' '.join(str(grb).split(':')[2].split(' ')[0:2])
                attrs['grid'] = outFile
        elif model == 'NCEP-GFS':
            with pygrib.open(inputFile) as grbfile:
                grb = grbfile.select(shortName='tp')[0]
                met_data, lats, lons = grb.data()
                msg = grb.tostring()

                outFile = '/'.join(inputFile.split('/')[0:5])+'/processed/'+'/'.join(inputFile.split('/')[6:9])+'/tp_24h/'+inputFile.split('/')[-1]
                grbout = open(outFile, 'wb')
                grbout.write(msg)
                grbout.close()
                
                attrs = {}
                attrs['valid'] = sys.argv[3]
                attrs['init'] = sys.argv[4]
                attrs['lead'] = sys.argv[5]+'0000'
                attrs['accum'] = '240000'
                attrs['name'] = varName
                attrs['long_name'] = ' '.join(str(grb).split(':')[1].split(' ')[0:2])
                attrs['level'] = 'A24' 
                attrs['units'] = str(grb).split(':')[2].split(' ')[0]
                attrs['grid'] = outFile
    elif extension == '':
        if model == 'ECMWF-HRES':
            if int(sys.argv[5]) <= 24:
                with pygrib.open(inputFile) as grbfile:
                    grb = grbfile.select(shortName='tp')[0]
                    met_data, lats, lons = grb.data()
                    msg = grb.tostring()
    
                    outFile = '/'.join(inputFile.split('/')[0:5])+'/processed/'+'/'.join(inputFile.split('/')[6:9])+'/tp_24h/'+inputFile.split('/')[-1]
                    grbout = open(outFile, 'wb')
                    grbout.write(msg)
                    grbout.close()
                    
                    attrs = {}
                    attrs['valid'] = sys.argv[3]
                    attrs['init'] = sys.argv[4]
                    attrs['lead'] = sys.argv[5]+'0000'
                    attrs['accum'] = '240000'
                    attrs['name'] = varName
                    attrs['long_name'] = ' '.join(str(grb).split(':')[1].split(' ')[0:2])
                    attrs['level'] = 'A24' 
                    attrs['units'] = str(grb).split(':')[2].split(' ')[0]
                    attrs['grid'] = outFile
            else:
                with pygrib.open(inputFile) as grbfile:
                    grb = grbfile.select(shortName='tp')[0]
                    met_data, lats, lons = grb.data()
                inputFileS = '/'.join(inputFile.split('/')[0:9])+'/'+inputFile.split('/')[-1][0:13]+str('{:02d}'.format(int(inputFile.split('/')[-1][13:15])-1))+inputFile.split('/')[-1][-5:]
                with pygrib.open(inputFileS) as grbfileS:
                    grbS = grbfileS.select(shortName='tp')[0]
                    met_dataS, lats, lons = grbS.data()
                met_data = np.subtract(met_data, met_dataS)
                met_data = np.multiply(met_data, 1000)
                #Can't change grb.values directly, we need to remove existing values first
                grb.values = grb.values * 0
                grb.values = met_data
                msg = grb.tostring()
                
                outFile = '/'.join(inputFile.split('/')[0:5])+'/processed/'+'/'.join(inputFile.split('/')[6:9])+'/tp_24h/'+inputFile.split('/')[-1]
                grbout = open(outFile, 'wb')
                grbout.write(msg)
                grbout.close()
                    
                attrs = {}
                attrs['valid'] = sys.argv[3]
                attrs['init'] = sys.argv[4]
                attrs['lead'] = sys.argv[5]+'0000'
                attrs['accum'] = '240000'
                attrs['name'] = varName
                attrs['long_name'] = ' '.join(str(grb).split(':')[1].split(' ')[0:2])
                attrs['level'] = 'A24' 
                attrs['units'] = str(grb).split(':')[2].split(' ')[0]
                attrs['grid'] = outFile