#!/usr/bin/env python
import os, sys, glob, subprocess, shutil
from datetime import datetime, timedelta
import re
import numpy as np
from satellite_info import env_paths, sat_ant_info

__authors__ = ["Ming Chen"]
__license__ = "Public Domain"
__maintainer__ = "Ming Chen"
__email__ = "ming.chen@noaa.gov"
__date__  = "Mar 24, 2023"

""" 
#####################################################################
# 
# ucar2umd_rnx.py
#
# Purpose:
# This Python library provides utilities for the COSMIC-2 GNSS data
# preprocessing workflow, including:
# 1) conversion of UCAR CRX data to RINEX format,
# 2) aggregation of sub-daily RINEX files into daily or hourly-session
#    datasets for processing with the BERNESE software package.
#
# History:  
#       Created by: Ming Chen 2023/03/24
#                   ming.chen@noaa.gov
# 
# Usage: from ucar2umd_rnx import UCAR2UMD_RNX
#
#
#####################################################################

"""

class UCAR2UMD_RNX(object):

   def __init__(self, mission='cosmic2'):
      self.mission = mission
      self.clean()

   def __del__(self):
      del self.header
      del self.dataset

   def clean(self):
      self.header = []
      self.dataset = {}
      self.add_header = True
      self.start_time = 'zzzz'
      self.end_time = '0000' 

   def write_rinex(self, outfile='Test.out', **kwargs):
       verbose = kwargs.get('verbose', False) 
       for i in range(len(self.header)):
          if self.header[i].find("FIRST") != -1: 
             self.header[i] = self.start_time
          elif self.header[i].find("LAST") != -1: 
             self.header[i] = self.end_time

       ds = self.dataset
       mykeys = list(ds.keys())
       mykeys.sort()
       self.dataset = {i: ds[i] for i in mykeys}
       
       if(verbose) : print('output:{}'.format(outfile))
       with open(outfile,'w') as f:
          for line in self.header:
             f.write(line)

          for k,v in self.dataset.items():
             #print(k,  v[0], len(v)
             for line in v:
                f.write(line)

   def add_rinex(self, filename, **kwargs):      
       verbose = kwargs.get('verbose', True) ; warning = False
       # ftype = os.path.splitext(filename)[1]
       if not os.path.isfile(filename):
         err = "Error: file not exis {} ... ".format(filename)
         print(err)
         return(False)
       
       btime = r'(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+\.\d+)\s+(\d+)\s+(\d+)(\w+)'
       htime = r'(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+\.\d+)\s+(\w+)\s+(\.+)'
       bdata = r'[^\+\-\.0-9\s]+'
       peh=re.compile(r'^\s*END OF HEADER\s*')
       pbt=re.compile(r'^\s*'+btime+r'\s*')
       pdd=re.compile(r'^\s*'+bdata)
 
       fp = open(filename, 'r' )
       nnn = 0 ; headerlines = 0
       while True :
           line = fp.readline() ; nnn += 1
           if len(line) == 0:
              break
           #line=line.replace("\n", "")
           if (headerlines == 0):
              while peh.match(line) == None:
                  headerlines += 1            
                  if self.add_header: self.header.append(line)
                  if line.find("FIRST") != -1:
                     if(verbose): print(line)
                     if line.strip() < self.start_time.strip(): self.start_time = line
                  if line.find("LAST") != -1:
                     if(verbose): print(line)
                     if line.strip() > self.end_time.strip(): self.end_time = line
                  line = fp.readline() ; nnn += 1
 
              if self.add_header:
                  self.header.append(line)
                  self.add_header = False
 	     
           if pbt.match(line):
              ts = pbt.search(line)
              dstr1="{:02d}{:02d}{:02d}".format(int(ts[1]),int(ts[2]),int(ts[3]))
              dstr2="T{:02d}:{:02d}:{:06.3f}".format(int(ts[4]),int(ts[5]),float(ts[6]))
              dkey=dstr1+dstr2

              sec    = float(ts[6])
              current_time = datetime(
                int(ts[1]), int(ts[2]), int(ts[3]), int(ts[4]), int(ts[5]),
                int(sec),
                int((sec % 1) * 1e6)
               )
              dkey = current_time
              datalines = [line] 
              n = 0 ; N = int(ts[8])*2
              while n < N:
                  line = fp.readline() ; nnn += 1
                  if len(line) == 0:
                     break
                  if len(line.strip()) == 0:
                     continue
                  if (pdd.match(line) != None) and (n==0):
                     if(warning):
                       print("additional non-data line ...." )
                       print(line.strip())
                       print("filename: {} lineNo: {} ".format(filename, nnn))
                     datalines.append(line)
                     continue
                  datalines.append(line)
                  n = n + 1
              if n == N:
                 if dkey in  self.dataset:
                    cn = len(self.dataset[dkey]) - 1
                    if(warning): print('existing ',dkey, N, cn)
                    if N > cn:
                       self.dataset[dkey] = datalines
                 else:
                    self.dataset[dkey] = datalines
              else:
                 if(verbose): print('Inconsistent record error ... {} != {} \n'.format(n,N))
 
       fp.close()

       return 
 

   def extract_rinex(self, outfile='Test.out', **kwargs):
       verbose = kwargs.get('verbose', False)
       gnss=['G','R','E','B','Z']
       btime = r'(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+\.\d+)\s+(\d+)\s+(\d+)(\w+)'
       htime = r'(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+\.\d+)\s+(\w+)\s+(\.+)'
       bdata = r'[^\+\-\.0-9\s]+'
  
       pbt=re.compile(r'^\s*'+btime+r'\s*')
       pdd=re.compile(r'^\s*'+bdata)
 
       for i in range(len(self.header)):
          if self.header[i].find("FIRST") != -1: 
             print ("FIRST: ",self.start_time)
          elif self.header[i].find("LAST") != -1: 
             print ("LST: ",self.end_time)

       ds = self.dataset
       mykeys = list(ds.keys())
       mykeys.sort()
       self.dataset = {i: ds[i] for i in mykeys}

       if verbose: print("output: {}".format(outfile))
       with open(outfile,'w') as f:
          for k,v in self.dataset.items():
             #print(k,  v[0], len(v))
             newline=''
             for line in v:
                 if pbt.match(line):
                    ts = pbt.search(line)
                    dstr1=ts[8] ; dstr2=ts[9]
                    sn="{:2d} ".format(int(ts[8]))
                    for s in gnss:
                        sn += "{}{:02d}".format(s, len(re.findall(s, dstr2)))
                    newline =  newline + " " + sn +"\n"
             f.write(k+ " " + newline)

def str2date(datestr):
    if len(datestr) == 8:
       mydate = datetime.strptime(datestr,'%Y%m%d')
    elif len(datestr) == 7:
       mydate = datetime.strptime(datestr,'%Y%j')
    else:
       print("Error str2date input, usage: yyyymndd or yyyyddd ")
       return None 
    return (mydate)


def ucar_crx2rnx(mission, satidx, antidx, datestr,  
         ucarpath=None, tmppath=None, utilpath=None, **kwargs):      

    verbose = kwargs.get('verbose', False) 
    satinfo = sat_ant_info(mission, satidx, antidx)
    if satinfo == None: return False
    filelist = []
    paths = env_paths()
    if utilpath == None: utilpath = paths['util']
    crx2rnx = os.path.join(utilpath, 'bin','CRX2RNX')
    if not os.path.isfile(crx2rnx):
       print('CRX2RNX {} utility is not available ....'.format(crx2rnx))
       return filelist
   
    ucardir = ucarpath ; tmpdir = tmppath
    if ucardir == None:
       ucardir =  paths['ucar']
    if tmpdir == None:
       tmpdir =  paths['tmp']

    satsvn =  satinfo['satsvn'] ; satname = satinfo['satname'] 
    antenna = satinfo['antsvn'] 
    if antenna == None: 
       print("error in antidx {}".format(antidx))
       return filelist
   
    sdate = str2date(datestr) ; dt = timedelta(days=1)
    year  = sdate.strftime('%Y')
    jday  = sdate.strftime('%j') 
   
    tmpin  = os.path.join(tmpdir, mission,'CRX',year+jday)
    tmpout = os.path.join(tmpdir, mission,'RNX',year+jday)
    if os.path.exists(tmpin): shutil.rmtree(tmpin)
    if os.path.exists(tmpout): shutil.rmtree(tmpout)
    os.makedirs(tmpin) ; os.makedirs(tmpout)

    # ucar zipped crx files -> unzipped
    ucarPodCrx =  os.path.join(ucardir, mission, 'podCrx', year, jday)
    fnames = 'podCrx_{}.{}.{}.*{}'.format(year, jday, satsvn, antenna)
    filelist = glob.glob(ucarPodCrx+'/'+fnames+'_crx.gz', recursive=False)        
    filelist= sorted(filelist)
  
    if len(filelist)< 1: 
      print( "no ucarPodCrx files found...")
      msg = '{} {} podCrx_*.{}.*{}  {}'.format(ucardir, mission, satsvn, antenna, datestr)
      print( msg) 
      return filelist
 
    for f in filelist:
       _, basename = os.path.split(os.path.abspath(f))
       #os.system('gunzip -f '+ indir+'/'+basename +' > '+ indir+'/'+basename +'.crx')
       os.system('gunzip -c -f '+ f +' > '+ tmpin+'/'+basename[:-7] +'.crx')

    # convert unzipped crx files to rnx files
    #fnames = 'podCrx_{}.{}.{}.*{}'.format(year, jday, satsvn, antenna)
    fnames = 'podCrx_*.{}.*{}'.format(satsvn, antenna)
    filelist = glob.glob(tmpin+'/'+fnames+'.crx', recursive=False)
    if len(filelist) < 1: return filelist
    fpattern = r'(\d{4}).(\d{3}).(\d{3}).(\d{2}).(\d{2})'
    pf = re.compile('_'+fpattern)
    for f in filelist:
       _, basename = os.path.split(os.path.abspath(f))
       #ts = pf.search(basename)
       outfile =   basename[:-4] +'.rnx'
       fileout = os.path.join(tmpout,  outfile) 
       #crx2rnx f ->outfile
       cmd_arg = [f]
       cmd_arg.append("-" )
       with  open(fileout,'w') as outfile:
          subprocess.run([crx2rnx]+cmd_arg, stdout=outfile)
       # cmd=crx2rnx +' '+ f + ' ->'+outfile
       # os.system( cmd )
    filelist = glob.glob(f'{tmpout}/*.rnx', recursive=False)
    return filelist




def ucar_crx2rnx(crxlist, pdate, crx2rnx_cmd_file,  tmppath=None, **kwargs):      

    filelist = []  
    # ucar zipped crx files -> unzipped
    if len(crxlist)< 1: 
      print( "no PodCrx files found...")
      return filelist

    #crx2rnx = os.path.join(utilpath, 'bin','CRX2RNX')
    crx2rnx = crx2rnx_cmd_file
    if not os.path.isfile(crx2rnx):
       print('CRX2RNX {} utility is not available ....'.format(crx2rnx))
       return filelist
   
    tmpdir = './TMP' if tmppath is None else tmppath

    crxtmp = os.path.join(tmpdir, 'CRX')
    rnxtmp = os.path.join(tmpdir, 'RNX')
    if os.path.exists(crxtmp): shutil.rmtree(crxtmp)
    if os.path.exists(rnxtmp): shutil.rmtree(rnxtmp)
    os.makedirs(crxtmp) ; os.makedirs(rnxtmp)

    crxlist = sorted(crxlist)

    for f in crxlist:
       _, basename = os.path.split(os.path.abspath(f))
       #os.system('gunzip -f '+ indir+'/'+basename +' > '+ indir+'/'+basename +'.crx')
       os.system('gunzip -c -f '+ f +' > '+ crxtmp+'/'+basename[:-7] +'.crx')

    # convert unzipped crx files to rnx files
    crxlist = glob.glob(f'{crxtmp}/podCrx*.crx', recursive=False)
    if len(crxlist) < 1: return filelist
    for f in crxlist:
       _, basename = os.path.split(os.path.abspath(f))
       outfile =   basename[:-4] +'.rnx'
       fileout = os.path.join(rnxtmp,  outfile) 
       #crx2rnx f ->outfile
       cmd_arg = [f]
       cmd_arg.append("-" )
       with  open(fileout,'w') as outfile:
          subprocess.run([crx2rnx]+cmd_arg, stdout=outfile)
       # cmd=crx2rnx +' '+ f + ' ->'+outfile
       # os.system( cmd )
    filelist = glob.glob(f'{rnxtmp}/*.rnx', recursive=False)
    return filelist
		
if __name__ == "__main__":
  print("UCAR2UMD_RNX")
