from overpasses import find_overpasses
import mutils as m
#from ProcessResults import gen_overpass_text_time_series
from ProcessResults import parse_sampleConfig, cdf_vals, get_rounded_cdf_percents
import numpy as np
from mutils.PosPacConst import *
import matplotlib.pyplot as plt
import json
import os
import sys
import datetime

from matplotlib import use
# Allow running headless from the command line
use("agg")

def writeKMLMarker(fid,lat,lon,hgt):
  fid.write('<Placemark>')
  fid.write('<styleUrl>#marker</styleUrl>')
  fid.write('<Point>')
  fid.write('<coordinates>')
  fid.write('%.8f,%.8f,%.3f' % (lon,lat,hgt))
  fid.write('</coordinates>')
  fid.write('</Point>')
  fid.write('</Placemark>\n')


def writeKMLHeader(fid,markerColor):
  fid.write('<?xml version="1.0" encoding="UTF-8"?>')
  fid.write('<kml xmlns="http://earth.google.com/kml/2.0"> <Document>')
  fid.write('<Style id="marker"><IconStyle>')
  fid.write('<color>%s</color>' % markerColor)
  fid.write('<scale>0.8</scale>')
  fid.write('<Icon><href>http://maps.google.com/mapfiles/kml/shapes/placemark_square.png</href></Icon>')
  fid.write('</IconStyle></Style>\n')

#
# Splits data from the freeway in to "overpass" (within +/-200) and everything else grouped
# together in "clear"
#
if __name__ == "__main__":

  colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf', '#ffff00']

  dDataRate = 10 # Default rate is 10Hz

  if(len(sys.argv) == 2):
    with open(sys.argv[1],'r') as f: 
      configData = json.load(f)
  else:
    print("Usage:\n  python compareEnv.py config.json")
    sys.exit()

  now = datetime.datetime.now()
  YYYYMMDD = str(now.year) + '-' + str(now.month).zfill(2) + '-' + str(now.day).zfill(2)
  if not os.path.exists(YYYYMMDD):
    os.makedirs(YYYYMMDD)

  prefix = configData['prefix']
  prefix = YYYYMMDD + '/' + prefix

  if(os.path.exists(prefix + '-' + "FileLocations.txt")):
    os.remove(prefix + '-' + "FileLocations.txt")

  dataOut = {}
  allReceivers = []

  # Delete the concatenated debug data - we should only need to
  # look at the first dataSet. However, check them all in case there
  # is a different list of receivers. 
  for thisData in configData['dataSets']:
    for thisDevice in thisData['DUTs']: 
      thisLabel = thisDevice['label']
      # Remove the file where we'll save detailed information
      if os.path.exists(prefix + '-' + thisLabel+"-overpasses.txt"):
        os.remove(prefix + '-' + thisLabel+"-overpasses.txt")
      if os.path.exists(prefix + '-' + thisLabel+"-clear.txt"):
        os.remove(prefix + '-' + thisLabel+"-clear.txt")
      if os.path.exists(prefix + '-' + thisLabel+"-overpasses.kml"):
        os.remove(prefix + '-' + thisLabel+"-overpasses.kml")
      with open(prefix + '-' + thisLabel+"-overpasses.kml", "w") as fid:
        writeKMLHeader(fid,'ff0000ff')

      if os.path.exists(prefix + '-' + thisLabel+"-clear.kml"):
        os.remove(prefix + '-' + thisLabel+"-clear.kml")
      with open(prefix + '-' + thisLabel+"-clear.kml", "w") as fid:
        writeKMLHeader(fid,'ff00ff00')

  for thisData in configData['dataSets']:
    configFile = thisData['config']

    config = parse_sampleConfig(configFile)
    truthFile = config.truth_file
    pp = m.mutils.doload(truthFile)

    for thisDevice in thisData['DUTs']: 
      dataDir   = thisDevice['file']
      thisLabel = thisDevice['label']
        
      fullPath= os.path.realpath(dataDir)
      with open(prefix + '-' + "FileLocations.txt", "a") as fid:
        fid.write(fullPath + " " + thisLabel + "\n")

      if(thisLabel not in allReceivers):
        allReceivers.append(thisLabel)

      if(not thisLabel in dataOut):
        dataOut[thisLabel] = {}

      if('dataRate' in thisDevice):
        dataRate = thisDevice['dataRate']
      else:
        dataRate = dDataRate

      if('rec' in thisDevice):
        rec = thisDevice['rec']
      else:
        rec = '-d35:2'
      print(dataDir)
      d  = m.mutils.vd2cls(dataDir + '/*.T04',rec)

      start = d.TIME[0]
      stop  = d.TIME[-1]

      i = m.find( (pp[:,dPP_TIME] >= start) & (pp[:,dPP_TIME] <= stop) )
      thisPP = pp[i,:]
      overpassLoc, events = find_overpasses( thisPP, False ) # 200m is the default

      spans = np.zeros(len(d),dtype=bool)
      numOverpasses = len(events)

      for desc,start_stops in config.span:
        if(desc == 'Freeways'):
          desc = 'Freeway'
        
        if(desc != 'Freeway'):
          continue

        for start,stop in start_stops:
          i_d = m.find( (d.TIME>=start) & (d.TIME<=stop) )
          spans[i_d] = True

      # Now have an array "spans" epochs that are on the freeway are True
      i = m.find(spans)
      # The spans and the d arrays match - select only the data on the overpasses
      d = d[i]

      overpasses = np.zeros(len(d),dtype=bool)

      for thisEvent in events:
        _,i_event,i_d = m.intersect( thisEvent.TIME, d.TIME )
        overpasses[i_d] = True

      i = m.find(overpasses)
      d_overpasses = d[i]

      i = m.find(overpasses == False)
      d_clear = d[i]

      # Process the data in the clear
      _,i_pp,i_d = m.intersect( thisPP[:,dPP_TIME], d_clear.TIME )
      ref_pos = thisPP[i_pp,dPP_LAT:dPP_velU+1]
      dut_pos = np.array([d_clear.LAT[i_d],d_clear.LON[i_d],d_clear.HGT[i_d]]).T

      (dE, dN, dU) = m.llh2enu( dut_pos,ref_pos[:,0:3] )
      err_2d = np.sqrt(dE**2 + dN**2)
      err_3d = np.sqrt(dE**2 + dN**2 + dU**2)

      if(not 'clear' in dataOut[thisLabel]):
        dataOut[thisLabel]['clear'] = {'2d':err_2d,'3d':err_3d,'secs':len(err_2d)/dataRate}
      else:
        dataOut[thisLabel]['clear']['2d'] = np.concatenate((dataOut[thisLabel]['clear']['2d'],err_2d))
        dataOut[thisLabel]['clear']['3d'] = np.concatenate((dataOut[thisLabel]['clear']['3d'],err_3d))
        dataOut[thisLabel]['clear']['secs'] += len(err_2d)/dataRate

      with open(prefix + "-" + thisLabel+"-clear.txt","a") as fid:
        for i in range(len(dE)):
          fid.write("%.3f %d " % (d_clear.TIME[i_d[i]],d_clear.WEEK[i_d[i]]))
          fid.write("%.10f %.10f %.3f " % (dut_pos[i,0], dut_pos[i,1], dut_pos[i,2]))
          fid.write("%.10f %.10f %.3f " % (ref_pos[i,0], ref_pos[i,1], ref_pos[i,2]))
          fid.write("%.3f %.3f %.3f\n" % (dE[i],dN[i],dU[i]))

      # Output a 1Hz placemarker (at the DUT not truth) position
      with open(prefix + "-" + thisLabel+"-clear.kml","a") as fid:
        for i in range(len(dE)):
          timeStamp = d_clear.TIME[i_d[i]]
          if( (int(timeStamp * 100) % 100) == 0): 
            writeKMLMarker(fid,dut_pos[i,0],dut_pos[i,1],dut_pos[i,2]) 

      # Process the data under the overpass
      _,i_pp,i_d = m.intersect( thisPP[:,dPP_TIME], d_overpasses.TIME )
      ref_pos = thisPP[i_pp,dPP_LAT:dPP_velU+1]
      dut_pos = np.array([d_overpasses.LAT[i_d],d_overpasses.LON[i_d],d_overpasses.HGT[i_d]]).T

      (dE, dN, dU) = m.llh2enu( dut_pos,ref_pos[:,0:3] )
      err_2d = np.sqrt(dE**2 + dN**2)
      err_3d = np.sqrt(dE**2 + dN**2 + dU**2)

      if(not 'overpass' in dataOut[thisLabel]):
        dataOut[thisLabel]['overpass'] = {'2d':err_2d,'3d':err_3d,'secs':len(err_2d)/dataRate,'numOverpasses':numOverpasses}
      else:
        dataOut[thisLabel]['overpass']['2d'] = np.concatenate((dataOut[thisLabel]['overpass']['2d'],err_2d))
        dataOut[thisLabel]['overpass']['3d'] = np.concatenate((dataOut[thisLabel]['overpass']['3d'],err_3d))
        dataOut[thisLabel]['overpass']['secs'] += len(err_2d)/dataRate
        dataOut[thisLabel]['overpass']['numOverpasses'] += numOverpasses

      with open(prefix + "-" + thisLabel+"-overpasses.txt","a") as fid:
        for i in range(len(dE)):
          fid.write("%.3f %d " % (d_overpasses.TIME[i_d[i]],d_overpasses.WEEK[i_d[i]]))
          fid.write("%.10f %.10f %.3f " % (dut_pos[i,0], dut_pos[i,1], dut_pos[i,2]))
          fid.write("%.10f %.10f %.3f " % (ref_pos[i,0], ref_pos[i,1], ref_pos[i,2]))
          fid.write("%.3f %.3f %.3f\n" % (dE[i],dN[i],dU[i]))
      
      # Output a 1Hz placemarker (at the DUT not truth) position
      with open(prefix + "-" + thisLabel+"-overpasses.kml","a") as fid:
        for i in range(len(dE)):
          timeStamp = d_overpasses.TIME[i_d[i]]
          if( (int(timeStamp * 100) % 100) == 0): 
            writeKMLMarker(fid,dut_pos[i,0],dut_pos[i,1],dut_pos[i,2]) 


  # Close out the KML files
  for thisLabel in allReceivers:
    with open(prefix + '-' + thisLabel+"-overpasses.kml", "a") as fid:
      fid.write('</Document> </kml>')
    with open(prefix + '-' + thisLabel+"-clear.kml", "a") as fid:
      fid.write('</Document> </kml>')

  # Now plot the 2D and 3D data
  plots = ['2d','3d']
  for thisPlot in plots:
    colorIndex = 0
    plt.figure()
    for thisLabel in allReceivers:
      cx,cy = m.docdf(dataOut[thisLabel]['clear'][thisPlot])
      plotLabel = thisLabel + ": Clear | " +  "{:.1f}".format(dataOut[thisLabel]['clear']['secs']) + 's'
      plt.plot(cx,cy,color=colors[colorIndex], label=plotLabel)

      cx,cy = m.docdf(dataOut[thisLabel]['overpass'][thisPlot])
      plotLabel = thisLabel + ": Overpass | " +  "{:.1f}".format(dataOut[thisLabel]['overpass']['secs']) + 's' + ' | Num Overpasses ' + "{:d}".format(dataOut[thisLabel]['overpass']['numOverpasses'])
      plt.plot(cx,cy,linestyle='dashed', color=colors[colorIndex], label=plotLabel)
      colorIndex +=1
      colorIndex %= len(colors) # Recycle if we run out of colors

    plt.grid(True)
    plt.legend(fontsize=8)
    plt.xlim(0,1)
    plt.ylim(0,1)
    plt.xlabel('Error [m]')
    plt.ylabel(thisPlot.upper() + ' CDF')
    plt.title('Freeway data - clear vs +/-200m around overpasses')
    plt.tight_layout()
    m.save_compressed_png(prefix + '-' + thisPlot.upper() + "_cdf_1m.png", dpi=600, make_dir=False)

    plt.xlim(0,0.2)
    m.save_compressed_png(prefix + '-' + thisPlot.upper() + "_cdf_20cm.png", dpi=600, make_dir=False)

    # Now generate a Terrasat format graph with the x-axis in log format
    plt.xlim(0.01,10)
    plt.xscale('log')
    plt.tight_layout()
    m.save_compressed_png(prefix + '-' + thisPlot.upper() + "_cdf_LogX_10m.png", dpi=600, make_dir=False)
    plt.xlim(0.01,1)
    plt.xscale('log')
    plt.tight_layout()
    m.save_compressed_png(prefix + '-' + thisPlot.upper() + "_cdf_LogX_1m.png", dpi=600, make_dir=False)