#! /usr/bin/env python2
############################################################
# Program is part of PySAR v1.0                            #
# Copyright(c) 2013, Heresh Fattahi                        #
# Author:  Heresh Fattahi                                  #
############################################################


import os
import sys
import getopt

import h5py
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FormatStrFormatter


def readGPSfile(gpsFile,gps_source):
    if gps_source in ['cmm4','CMM4']:
 
        gpsData = np.loadtxt(gpsFile,usecols = (1,2,3,4,5,6,7,8,9,10))
        Stations = np.loadtxt(gpsFile,dtype=str,usecols = (0,1))[:,0]
 
        St=[]; Lon=[];Lat=[];Ve=[];Vn=[];Se=[];Sn=[]
        for i in range(gpsData.shape[0]):
            if 'GPS' in Stations[i]:
                Lat.append(gpsData[i,0])
                Lon.append(gpsData[i,1])
                Ve.append(gpsData[i,2])
                Se.append(gpsData[i,3])
                Vn.append(gpsData[i,4])
                Sn.append(gpsData[i,5])
                St.append(Stations[i])
        Up=np.zeros(St.shape);Sup=np.zeros(St.shape)
    elif gps_source == 'pysar':
 
        #gpsData = np.loadtxt(gpsFile,usecols = (1,2,3,4,5,6,7,8,9))
        gpsData = np.loadtxt(gpsFile,usecols = (1,2,3,4,5,6))
        Stations = np.loadtxt(gpsFile,dtype=str,usecols = (0,1))[:,0]
 
        St=[];Lon=[];Lat=[];Ve=[];Vn=[];Se=[];Sn=[]
 
        for i in range(gpsData.shape[0]):
            Lon.append(gpsData[i,0]-360)
            Lat.append(gpsData[i,1])
            Ve.append(gpsData[i,2])
            Vn.append(gpsData[i,3])
            Se.append(gpsData[i,4])
            Sn.append(gpsData[i,5])
            St.append(Stations[i])
        Up=np.zeros(St.shape);Sup=np.zeros(St.shape)
    elif gps_source in ['usgs','USGS']:
        gpsData_Hz = np.loadtxt(gpsFile,usecols = (0,1,2,3,4,5,6))
        gpsData_up = np.loadtxt(gpsFile,usecols = (8,9))
        gpsData=np.hstack((gpsData_Hz,gpsData_up))
        Stations = np.loadtxt(gpsFile,dtype=str,usecols = (7,8))[:,0]
 
        St=[];Lon=[];Lat=[];Ve=[];Vn=[];Se=[];Sn=[];Up=[];Sup=[]

        for i in range(gpsData.shape[0]):
            Lat.append(gpsData[i,0])
            Lon.append(gpsData[i,1])
            Vn.append(gpsData[i,2])
            Ve.append(gpsData[i,3])
            Sn.append(gpsData[i,4])
            Se.append(gpsData[i,5])
            Up.append(gpsData[i,6])
            Sup.append(gpsData[i,7])
            St.append(Stations[i])

    return list(St),Lat,Lon,Ve,Se,Vn,Sn,Up,Sup 


def nearest(x, tbase,xstep):
    # """ find nearest neighbour """
    dist = np.sqrt((tbase -x)**2)
    if min(dist) <= np.abs(xstep):
        indx=dist==min(dist)
    else:
        indx=[]
    return indx

def find_row_column(Lon,Lat,lon,lat,lon_step,lat_step):
    ################################################
    # finding row and column numbers of the GPS point
  
    idx= nearest(Lon, lon, lon_step)
    idy= nearest(Lat, lat, lat_step)
    if idx !=[] and idy != []:
        IDX=np.where(idx==True)[0][0]
        IDY=np.where(idy==True)[0][0]
    else:
        IDX=np.nan
        IDY=np.nan
    return IDY, IDX


################################################
def usage():
    print '''
   ************************************************************************************************
   
   Compares InSAR and GPS velocities. Option G can be used to specify the mode of Comparison.
   Out put is InSARvsGPS.png
 
   Usage:
   
   insar_vs_gps.py  -v InSARvelocity.h5 -g GPS velocity file -r Name of the reference GPS station -l ststion list to be compared
   
   -l :  if station list is not specified, all stations in GPS velocity file are compared with InSAR

   -m min value of the x and y axis of the plot
   -M max value of the x and y axis of the plot
   -r refernce GPS station
   -s second velocity map
   -S source of the GPS data: (usgs,cmm4,pysar)
      see documentation for more infromation
   -I incidence angle (if not given average look angle is used instead)
   -H Heading angle (if not given then the program reads it from the attributes of the velocity file)

   -G GPS components to be used to compare with InSAR: los_3D , los_Hz , los_Up , gps_Up. [default is los_3D]
       los_3D: to project three gps components to LOS
       los_Hz: to project horizontal gps components to LOS
       los_Up: to project only vertical gps component to LOS
       gps_Up: uses vertical GPS to compare with InSAR (InSAR LOS will be projected to Up)

   -A annotate the GPS station name on the plot [yes] or no
   -C annotation font color [default is green]
   -x annotation offset from the point in x direction [default=0]
   -y annotation offset from the point in y direction [default=0]
   -u to plot 1, 2 or 3 sigma uncertainty [default=1]
   -B marker size [default = 15]
   Example:
     
    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt -S usgs  -r BEMT -l 'HNPS,DHLG,SLMS,USGC'
    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt -S cmm4 -r BEMT -l 'HNPS,DHLG,SLMS,USGC,ROCH,MONP,SIO3,IVCO,TMAP,BMHL,BILL,OGHS'
    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt -S usgs -r BEMT
    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt -S usgs -r BEMT -c geo_temporal_coherence.h5 -t 0.95 
    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt -S usgs -r BEMT -c geo_temporal_coherence.h5 -t 0.95  -l 'HNPS,DHLG,SLMS,USGC'

    insar_vs_gps.py  -v geo_velocity_New_masked.h5 -g usgs_velocities_NAfixed.txt -r BEMT -S usgs -A yes -C green -x 0 -y 0.5 -H 193.0 -I 23.0

   +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

   When user wants to compare two InSAR velocities at GPS stations:

   Example: (to comaper geo_InSAR_velocity.h5 with simulated_velocity.h5 in GPS stations)

    insar_vs_gps.py  -v geo_InSAR_velocity.h5  -g gpsVelocity.txt  -r BEMT -s simulated_velocity.h5

   ************************************************************************************************
    '''


def main(argv):

    annotation='yes'
    ann_x=0
    ann_y=0
    annotation_Color='green'
    disp_velocity='yes'
    GPS_InSAR_dif_thr=1
    gps_comp='los_3D'
    uncertainty_fac=1.0
    MarkerSize=5
    try:
        opts, args = getopt.getopt(argv,"v:r:g:G:l:c:t:m:M:s:S:A:B:C:x:y:I:H:u:")
 
    except getopt.GetoptError:
        usage() ; sys.exit(1)

    for opt,arg in opts:
  
        if opt == '-v':            velocityFile = arg
        elif opt == '-s':          velocityFile2 = arg
        elif opt == '-g':          gpsFile = arg
        elif opt == '-r':          refStation = arg
        elif opt == '-l':          stationsList = arg.split(',')
        elif opt == '-c':          coherenceFile = arg
        elif opt == '-t':          thr = float(arg)
        elif opt == '-m':          minV=float(arg)
        elif opt == '-M':          maxV=float(arg)
        elif opt == '-S':          gps_source = arg
        elif opt == '-A':          annotation = arg
        elif opt == '-C':          annotation_Color = arg
        elif opt == '-x':          ann_x = float(arg)
        elif opt == '-y':          ann_y = float(arg)
        elif opt == '-I':          theta = float(arg)
        elif opt == '-H':          heading = float(arg)
        elif opt == '-G':          gps_comp = arg
        elif opt == '-u':          uncertainty_fac = float(arg)
        elif opt == '-B':          MarkerSize = float(arg)

    try:
        velocityFile
        gpsFile 
        refStation
    except:
        usage();sys.exit(1)
 
    try:    thr
    except: thr=0.9

   
    h5file = h5py.File(velocityFile,'r')
    dset=h5file['velocity'].get('velocity')
    insarData=dset[0:dset.shape[0],0:dset.shape[1]]
    k=h5file.keys()
 
    try:
        h5file2 = h5py.File(velocityFile2,'r')
        dset2=h5file2['velocity'].get('velocity')
        insarData2=dset2[0:dset2.shape[0],0:dset2.shape[1]]
    except:
        print ''
  
    ullon=float(h5file[k[0]].attrs['X_FIRST'])
    ullat=float(h5file[k[0]].attrs['Y_FIRST'])
    lon_step=float(h5file[k[0]].attrs['X_STEP'])
    lat_step=float(h5file[k[0]].attrs['Y_STEP'])
    lon_unit=h5file[k[0]].attrs['Y_UNIT']
    lat_unit=h5file[k[0]].attrs['X_UNIT']
 
    Length,Width = np.shape(insarData)
 
    lllat=ullat+Length*lat_step
    urlon=ullon+Width*lon_step
    lat=np.arange(ullat,lllat,lat_step)
    lon=np.arange(ullon,urlon,lon_step)
    #################################################################################################
    # finding the raw an column of the reference gps station and referencing insar data to this pixel
    Stations,Lat,Lon,Ve,Se,Vn,Sn,Vu,Su=readGPSfile(gpsFile,gps_source)
    idxRef=Stations.index(refStation)
    IDYref,IDXref=find_row_column(Lon[idxRef],Lat[idxRef],lon,lat,lon_step,lat_step)

    #############################################
    #  Stations, gpsData = redGPSfile(gpsFile)
    #  idxRef=Stations.index(refStation)
    #  Lat,Lon,Vn,Ve,Sn,Se,Corr,Vu,Su = gpsData[idxRef,:]
    #  IDYref,IDXref=find_row_column(Lon,Lat,lon,lat,lon_step,lat_step)   
    ###################################################
 
    if (not np.isnan(IDYref)) and (not np.isnan(IDXref)):
        print ''
        print '-----------------------------------------------------------------------'
        print 'referencing InSAR data to the GPS station at : ' + str(IDYref) + ' , '+ str(IDXref)       
        if not np.isnan(insarData[IDYref][IDXref]):
             insarData=insarData - insarData[IDYref][IDXref]
        else:
            print ''' 
%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
       WARNING: nan value for InSAR data at the refernce pixel!
                reference station should be a pixel with valid value in InSAR data.
                                
                please select another GPS station as the reference station.
%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%                       
                  '''
            sys.exit(1)
    else:
        print 'WARNING:'
        print 'Reference GPS station is out of the area covered by InSAR data'
        print 'please select another GPS station as the reference station.'
        sys.exit(1)

    #######################################################################################
    
    try:
        stationsList
    except:
        stationsList = Stations 
        #stationsList.remove(refStation)
 
    try:
        print 'incidence angle = ' + str(theta)
    except:
        print 'using look angle from the velocity file. For more precise results input the incidence angle using option -I.'
        look_n=float(h5file['velocity'].attrs['LOOK_REF1'])
        look_f=float(h5file['velocity'].attrs['LOOK_REF2'])
        theta=(look_n+look_f)/2.
        print 'incidence angle = ' + str(theta) 
  
    try:
        print 'Heading angle = '+str(heading)
    except:
        heading = float(h5file['velocity'].attrs['HEADING'])
        if heading < 0:
            heading=heading+360 
 
    theta=theta*np.pi/180.0
    heading=heading*np.pi/180.0
 
    if gps_comp in ['los_3D','LOS_3D','los_3d']:
        unitVec=[np.cos(heading)*np.sin(theta),-np.sin(theta)*np.sin(heading),-np.cos(theta)]
        gps_comp_txt=' projecting three gps components to LOS' 
    elif gps_comp in ['los_Hz','LOS_HZ','los_hz','los_HZ','LOS_hz']:
        unitVec=[np.cos(heading)*np.sin(theta),-np.sin(theta)*np.sin(heading),0]
        gps_comp_txt=' projecting horizontal gps components to LOS'
    elif gps_comp in ['LOS_UP','los_Up','los_up','LOS_up']:
        unitVec=[0,0,-np.cos(theta)]
        gps_comp_txt=' projecting vertical gps components to LOS'
    elif gps_comp in ['gps_up','GPS_UP','GPS_Up','gps_Up']:
        unitVec=[0,0,1]
        gps_comp_txt=' comparing veryical gps with InSAR'
        print '-------------------------'
        print 'Projecting InSAR to vertical'
        insarData=-insarData/np.cos(theta)      
    print '-------------------------'
    print 'unit vector for :' + gps_comp_txt
    print unitVec
    print '-------------------------'
    gpsLOS_ref=unitVec[0]*Ve[idxRef]+unitVec[1]*Vn[idxRef]+unitVec[2]*Vu[idxRef]
    Sr= ((unitVec[0]**2)*Se[idxRef]**2+(unitVec[1]**2)*Sn[idxRef]**2+(unitVec[2]**2)*Su[idxRef]**2)**0.5
 
    print '######################################################################'
    try:
        h5coh = h5py.File(coherenceFile)
        kh5coh=h5coh.keys()
        dset=h5coh[kh5coh[0]].get(kh5coh[0])
        Coh=dset[0:dset.shape[0],0:dset.shape[1]]
    except:
        print 'No information about the coherence of the points'

    InSAR=[]
    GPS=[]
    InSAR1=[]
    GPS1=[]
    InSAR2=[]
    GPS2=[]
    coherence=[]
    GPSx=[]
    GPSy=[]
    GPSx1=[]
    GPSy1=[]
    GPSx2=[]
    GPSy2=[]
    GPS_station=[]
    GPS_std=[]
    for st in stationsList:
   
        try:
            idx=Stations.index(st)
            #Lat,Lon,Vn,Ve,Sn,Se,Corr,Vu,Su = gpsData[idx,:]
            gpsLOS=unitVec[0]*Ve[idx]+unitVec[1]*Vn[idx]+unitVec[2]*Vu[idx]
            Sg= ((unitVec[0]**2)*Se[idx]**2+(unitVec[1]**2)*Sn[idx]**2+(unitVec[2]**2)*Su[idx]**2)**0.5
            S=(Sg**2+Sr**2)**0.5
            gpsLOS=gpsLOS-gpsLOS_ref
            IDY,IDX=find_row_column(Lon[idx],Lat[idx],lon,lat,lon_step,lat_step)
            insar_velocity=-insarData[IDY][IDX]
            
            try: 
                gpsLOS=insarData2[IDY][IDX]-insarData2[IDYref][IDXref]
                gpsLOS=-1000.0*gpsLOS
            except:
                InSAR_GPS_Copmarison='yes'
    
            if not np.isnan(insarData[IDY][IDX]):
                print '%%%%%%%%%%%%%%%%%%%%'
                print st
                print 'GPS: ' + str(gpsLOS) + '  +/- '+str(S)  
                print 'INSAR: '+ str(-insarData[IDY][IDX]*1000.0)
                try:
                    print 'Coherence: ' + str(Coh[IDY][IDX])
                    coherence.append(Coh[IDY][IDX])
                    if Coh[IDY][IDX]>thr:
                        InSAR1.append(-insarData[IDY][IDX]*1000.0)
                        GPS1.append(gpsLOS)
                    else:
                        InSAR2.append(-insarData[IDY][IDX]*1000.0)
                        GPS2.append(gpsLOS)
                       
                except:
                    print 'No information about the coherence is available!'
        
                InSAR.append(-insarData[IDY][IDX]*1000.0)
                GPS.append(gpsLOS)
                GPS_station.append(st)
                GPSx.append(IDX)
                GPSy.append(IDY)
                GPS_std.append(S)
                if np.abs(gpsLOS+insarData[IDY][IDX]*1000.0) < GPS_InSAR_dif_thr:
                    GPSx1.append(IDX)
                    GPSy1.append(IDY)
                else:
                    GPSx2.append(IDX)
                    GPSy2.append(IDY)
        except:
            NoInSAR='yes'
      
    InSAR=np.array(InSAR)
    GPS=np.array(GPS)
    GPS_std=np.array(GPS_std)
    lt=len(InSAR)
    SAD=np.sum(np.abs(InSAR-GPS),0)/lt
    C1=np.zeros([2,len(InSAR)])
    C1[0][:]=InSAR
    C1[1][:]=GPS
    Cor = np.corrcoef(C1)[0][1]
    print '++++++++++++++++++++++++++++++++++++++++++++++'
    print 'Comparison summary:'
    print ''
    print 'AAD (average absolute difference)= '+str(SAD) + ' [mm/yr]'
    print 'Correlation = '+str(Cor)
    print ''
    print '++++++++++++++++++++++++++++++++++++++++++++++'
    ###############################################################
    try:
        minV
        maxV
    except:
        minV=np.min([InSAR,GPS])
        maxV=np.max([InSAR,GPS])


    fig = plt.figure()
    ax=fig.add_subplot(111)
    ax.errorbar(GPS,InSAR,yerr=0.0, xerr=uncertainty_fac*GPS_std, fmt='ko',ms=MarkerSize)
    ax.plot([minV-3,maxV+3],[minV-3,maxV+3],'k--')
    ax.set_ylabel('InSAR [mm/yr]',fontsize=26)
    ax.set_xlabel('GPS LOS [mm/yr]',fontsize=26)
    ax.set_ylim(minV-3,maxV+3)
    ax.set_xlim(minV-3,maxV+3)
 
    if annotation in ['yes','y','Y','Yes','YES']:
        for i in range(len(GPS)) :
            ax.annotate(GPS_station[i],xy=(GPS[i], InSAR[i]), xytext=(GPS[i]+ann_x, InSAR[i]+ann_y),color=annotation_Color)
 
    majorLocator = MultipleLocator(5)
    ax.yaxis.set_major_locator(majorLocator)
    minorLocator   = MultipleLocator(1)
    ax.yaxis.set_minor_locator(minorLocator)
    ax.xaxis.set_minor_locator(minorLocator)
    for tick in ax.xaxis.get_major_ticks():
                 tick.label.set_fontsize(26)
    for tick in ax.yaxis.get_major_ticks():
                 tick.label.set_fontsize(26)
    plt.tick_params(which='major', length=15,width=2)
    plt.tick_params(which='minor', length=6,width=2)
 
    figName = 'InSARvsGPS_errorbar.png'
    plt.savefig(figName) 
    ###############################################################
   
    fig = plt.figure()
    ax=fig.add_subplot(111)
    ax.plot(GPS,InSAR, 'ko',ms=MarkerSize)
    ax.plot([minV-3,maxV+3],[minV-3,maxV+3],'k--')
    # ax.plot([-10,20],[-10,20],'k--')
    ax.set_ylabel('InSAR [mm/yr]',fontsize=26)
    ax.set_xlabel('GPS LOS [mm/yr]',fontsize=26)
    ax.set_ylim(minV-3,maxV+3)
    ax.set_xlim(minV-3,maxV+3)
    # ax.set_ylim(-10,15)
    # ax.set_xlim(-10,15)
    if annotation in ['yes','y','Y','Yes','YES']:
        for i in range(len(GPS)) :
            ax.annotate(GPS_station[i],xy=(GPS[i], InSAR[i]), xytext=(GPS[i]+ann_x, InSAR[i]+ann_y),color=annotation_Color)
 
    majorLocator = MultipleLocator(5)
    ax.yaxis.set_major_locator(majorLocator)
    minorLocator   = MultipleLocator(1)
    ax.yaxis.set_minor_locator(minorLocator)
    ax.xaxis.set_minor_locator(minorLocator)
    for tick in ax.xaxis.get_major_ticks():
                 tick.label.set_fontsize(26)
    for tick in ax.yaxis.get_major_ticks():
                 tick.label.set_fontsize(26)
    plt.tick_params(which='major', length=15,width=2)
    plt.tick_params(which='minor', length=6,width=2)
 
    figName = 'InSARvsGPS.png'
    # plt.savefig(figName,pad_inches=0.0)  
    plt.savefig(figName)
  
    ######################################################
 
    try:
        Coh
        fig = plt.figure()
        ax=fig.add_subplot(111)
        ax.errorbar(GPS1,InSAR1,yerr=1.0, xerr=1.0, fmt='o')
        ax.errorbar(GPS2,InSAR2,yerr=1.0, xerr=1.0, fmt='^')
        ax.plot([minV-3,maxV+3],[minV-3,maxV+3],'--')
        ax.set_ylabel('InSAR [mm/yr]',fontsize=26)
        ax.set_xlabel('GPS LOS [mm/yr]',fontsize=26)
        ax.set_ylim(minV-3,maxV+3)
        ax.set_xlim(minV-3,maxV+3)
 
    except:
        print ''
    plt.show()

if __name__ == '__main__':
    main(sys.argv[1:])

