#!/usr/bin/env python3

import numpy as np 
import argparse
import os
import isce
import isceobj
import shelve
import datetime
from isceobj.Location.Offset import OffsetField
from iscesys.StdOEL.StdOELPy import create_writer
from mroipac.ampcor.DenseAmpcor import DenseAmpcor
#from isceobj.Utils.denseoffsets import denseoffsets
#import pickle
from isceobj.Util.decorators import use_api

def createParser():
    '''
    Command line parser.
    '''

    parser = argparse.ArgumentParser( description='Generate offset field between two Sentinel swaths')
    parser.add_argument('-m','--reference', type=str, dest='reference', required=True,
            help='Reference image')
    parser.add_argument('-s', '--secondary',type=str, dest='secondary', required=True,
            help='Secondary image')

    parser.add_argument('--ww', type=int, dest='winwidth', default=64,
            help='Window Width')
    parser.add_argument('--wh', type=int, dest='winhgt', default=64,
            help='Window height')
    parser.add_argument('--sw', type=int, dest='srcwidth', default=20,
            help='Search window width')
    parser.add_argument('--sh', type=int, dest='srchgt', default=20,
            help='Search window height')
    parser.add_argument('--mm', type=int, dest='margin', default=50,
            help='Margin')
    parser.add_argument('--kw', type=int, dest='skipwidth', default=64,
            help='Skip across')
    parser.add_argument('--kh', type=int, dest='skiphgt', default=64,
            help='Skip down')

    parser.add_argument('-o','--outPrefix', type=str, dest='outprefix', default='dense_ampcor',
            help='Output prefix')

    parser.add_argument('--aa', type=int, dest='azshift', default=0,
            help='Gross azimuth offset')

    parser.add_argument('--rr', type=int, dest='rgshift', default=0,
            help='Gross range offset')
    parser.add_argument('--oo', type=int, dest='oversample', default=32,
            help = 'Oversampling factor')

    return parser

def cmdLineParse(iargs = None):
    parser = createParser()
    inps =  parser.parse_args(args=iargs)

    return inps

@use_api
def estimateOffsetField(reference, secondary, inps=None):
    '''
    Estimate offset field between burst and simamp.
    '''

    ###Loading the secondary image object
    sim = isceobj.createSlcImage()
    sim.load(secondary+'.xml')
    sim.setAccessMode('READ')
    sim.createImage()

    ###Loading the reference image object
    sar = isceobj.createSlcImage()
    sar.load(reference + '.xml')
    sar.setAccessMode('READ')
    sar.createImage()

    width = sar.getWidth()
    length = sar.getLength()

    objOffset = DenseAmpcor(name='dense')
    objOffset.configure()

#   objOffset.numberThreads = 6
    objOffset.setWindowSizeWidth(inps.winwidth)
    objOffset.setWindowSizeHeight(inps.winhgt)
    objOffset.setSearchWindowSizeWidth(inps.srcwidth)
    objOffset.setSearchWindowSizeHeight(inps.srchgt)
    objOffset.skipSampleAcross = inps.skipwidth
    objOffset.skipSampleDown = inps.skiphgt
    objOffset.margin = inps.margin 
    objOffset.oversamplingFactor = inps.oversample

    objOffset.setAcrossGrossOffset(inps.rgshift)
    objOffset.setDownGrossOffset(inps.azshift)

    objOffset.setFirstPRF(1.0)
    objOffset.setSecondPRF(1.0)
    if sar.dataType.startswith('C'):
        objOffset.setImageDataType1('mag')
    else:
        objOffset.setImageDataType1('real')

    if sim.dataType.startswith('C'):
        objOffset.setImageDataType2('mag') 
    else:
        objOffset.setImageDataType2('real')

    objOffset.offsetImageName = inps.outprefix + '.bil'
    objOffset.snrImageName = inps.outprefix +'_snr.bil'

    
    objOffset.denseampcor(sar, sim)

    sar.finalizeImage()
    sim.finalizeImage()
    return objOffset



def main(iargs=None):        
    '''
    Generate offset fields burst by burst.
    '''

    inps = cmdLineParse(iargs)
    outDir = os.path.dirname(inps.outprefix)
    os.makedirs(outDir, exist_ok=True)

    objOffset = estimateOffsetField(inps.reference, inps.secondary, inps)

    
    print('Top left corner of offset image: ', objOffset.locationDown[0][0],objOffset.locationAcross[0][0])

if __name__ == '__main__':
    
    main()
