#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Apr 30 14:21:16 2020

@author: fsh
"""

from osgeo import gdal, ogr, osr
import matplotlib.pyplot as plt

gdal.UseExceptions()

#Utility function to load data
def loadData(infile, band=1):
    ds = gdal.Open(infile, gdal.GA_ReadOnly)
    #Data array
    data = ds.GetRasterBand(band).ReadAsArray
    #Map extent
    trans = ds.GetGeoTransform()
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    extent = [trans[0], trans[0], + xsize * trans[1], trans[3], 
              ysize * trans[5], trans[3]]
    
    ds = None
    return data, extent


#gdalwarp [--help-general] [--formats]
#    [-s_srs srs_def] [-t_srs srs_def] [-to "NAME=VALUE"]* [-novshiftgrid]
#    [-order n | -tps | -rpc | -geoloc] [-et err_threshold]
#    [-refine_gcps tolerance [minimum_gcps]]
#    [-te xmin ymin xmax ymax] [-te_srs srs_def]
#    [-tr xres yres] [-tap] [-ts width height]
#    [-ovr level|AUTO|AUTO-n|NONE] [-wo "NAME=VALUE"] [-ot Byte/Int16/...] [-wt Byte/Int16]
#    [-srcnodata "value [value...]"] [-dstnodata "value [value...]"]
#    [-srcalpha|-nosrcalpha] [-dstalpha]
#    [-r resampling_method] [-wm memory_in_mb] [-multi] [-q]
#    [-cutline datasource] [-cl layer] [-cwhere expression]
#    [-csql statement] [-cblend dist_in_pixels] [-crop_to_cutline]
#    [-of format] [-co "NAME=VALUE"]* [-overwrite]
#    [-nomd] [-cvmd meta_conflict_value] [-setci] [-oo NAME=VALUE]*
#    [-doo NAME=VALUE]*
#    srcfile* dstfile

def geocodeUsingGdalWarp(infile, latfile, lonfile, outfile,
                         insrs=4326, outsrs=None,
                         spacing=None, fmt='GTiff', bounds=None,
                         method='near'):
    '''
    Geocode a swath file using corresponding lat, lon files
    '''
    sourcexmltmpl = '''    <SimpleSource>
      <SourceFilename>{0}</SourceFilename>
      <SourceBand>{1}</SourceBand>
    </SimpleSource>'''
    
    driver = gdal.GetDriverByName('VRT')
    tempvrtname = '/media/sf_share/RTC_4-15-2020/mdd/Geocode_test1/geocode.vrt'
    inds = gdal.OpenShared(infile, gdal.GA_ReadOnly)
    
    tempds = driver.Create(tempvrtname, inds.RasterXSize, inds.RasterYSize, 0)
    
    for ii in range(inds.RasterCount):
        band = inds.GetRasterBand(1)
        tempds.AddBand(band.DataType)
        tempds.GetRasterBand(ii+1).SetMetadata({'source_0': sourcexmltmpl.format(infile, ii+1)}, 'vrt_sources')
  
    sref = osr.SpatialReference()
    sref.ImportFromEPSG(insrs)
    srswkt = sref.ExportToWkt()

    tempds.SetMetadata({'SRS' : srswkt,
                        'X_DATASET': lonfile,
                        'X_BAND' : '1',
                        'Y_DATASET': latfile,
                        'Y_BAND' : '1',
                        'PIXEL_OFFSET' : '0',
                        'LINE_OFFSET' : '0',
                        'PIXEL_STEP' : '1',
                        'LINE_STEP' : '1'}, 'GEOLOCATION')
    
    band = None
    tempds = None 
    inds = None
    
    if spacing is None:
        spacing = [None, None]
    
    warpOptions = gdal.WarpOptions(format=fmt,
                                   xRes=spacing[0], yRes=spacing[0],
                                   dstSRS=outsrs, outputBounds = bounds, 
                                   resampleAlg=method, geoloc=True)
    gdal.Warp(outfile, tempvrtname, options=warpOptions)

###Create and visualize geocoded output
geocodeUsingGdalWarp('/media/sf_share/RTC_4-15-2020/mdd/output/gamma_VH.10alks_10rlks.img.vrt',
                     '/media/sf_share/RTC_4-15-2020/mdd/geometry/lat.rdr.vrt',
                     '/media/sf_share/RTC_4-15-2020/mdd/geometry/lon.rdr.vrt',
                     '/media/sf_share/RTC_4-15-2020/mdd/Geocode_test1/geocoded_mdd-rtc.tif',
                     spacing=[0.0005, 0.0005])

#geodata, geoext = loadData('/media/sf_share/RTC_4-15-2020/mdd/Geocode_test1/geocoded_mdd-rtc.tif', band=2)

#plt.figure()
#plt.imshow(geodata, extent=geoext, cmap='gray')
#plt.show()