#!/usr/bin/env python3
############################################################
# Program is part of MintPy                                #
# Copyright (c) 2013, Zhang Yunjun, Heresh Fattahi         #
# Author: Zhang Yunjun, Xiaohua Xu, Mar 2021               #
############################################################


import os
import sys
import glob
import numpy as np

try:
    from osgeo import gdal
except ImportError:
    raise ImportError('Can not import gdal!')

from mintpy.utils import (
    ptime,
    readfile,
    writefile,
    utils as ut,
)
from mintpy.utils.arg_utils import create_argument_parser


#########################################################################
EXAMPLE = """example:
  prep_gmtsar.py StHelensEnvDT156.txt
"""

def create_parser(subparsers=None):
    """Command line parser."""
    synopsis = 'Prepare GMTSAR metadata files.'
    epilog = EXAMPLE
    name = __name__.split('.')[-1]
    parser = create_argument_parser(
        name, synopsis=synopsis, description=synopsis, epilog=epilog, subparsers=subparsers)

    parser.add_argument('template_file', type=str, help='MintPy template file for GMTSAR products.')
    parser.add_argument('--mintpy-dir', dest='mintpy_dir', default='./',
                        help='MintPy directory (default: %(default)s).')
    parser.add_argument('--force', dest='update_mode', action='store_false',
                        help='Force to overwrite all .rsc metadata files.')
    return parser


def cmd_line_parse(iargs = None):
    parser = create_parser()
    inps = parser.parse_args(args=iargs)

    inps.template_file = os.path.abspath(inps.template_file)
    inps.mintpy_dir = os.path.expanduser(inps.mintpy_dir)
    inps.mintpy_dir = os.path.abspath(inps.mintpy_dir)
    return inps


#########################################################################
def get_prm_files(ifg_dir):
    """Get the date1/2.prm files in the interferogram directory."""
    prm_files = sorted(glob.glob(os.path.join(ifg_dir, '*.PRM')))
    prm_files = [i for i in prm_files if os.path.splitext(os.path.basename(i))[0].isdigit()]
    if len(prm_files) != 2:
        raise FileExistsError('NO two *.prm files found!')
    return prm_files


def get_multilook_number(ifg_dir, fbases=['corr', 'phase', 'phasefilt', 'unwrap']):
    """Get the number of multilook in range and azimuth direction."""
    # grab an arbitrary file in radar-coordiantes
    rdr_files = [os.path.join(ifg_dir, '{}.grd'.format(i)) for i in fbases]
    if len(rdr_files) == 0:
        raise ValueError('No radar-coord files found in {} with suffix: {}'.format(ifg_dir, fbases))

    # read step info into multilook number
    ds = gdal.Open(rdr_files[0], gdal.GA_ReadOnly)
    transform = ds.GetGeoTransform()
    rlooks = int(abs(transform[1]))
    alooks = int(abs(transform[5]))
    return alooks, rlooks


def get_lalo_ref(ifg_dir, prm_dict, fbases=['corr', 'phase', 'phasefilt', 'unwrap']):
    """Get the LAT/LON_REF1/2/3/4 from *_ll.grd file."""
    # grab an arbitrary file in geo-coordiantes
    geo_files = [os.path.join(ifg_dir, '{}_ll.grd'.format(i)) for i in fbases]
    if len(geo_files) == 0:
        print('WARNING: No geo-coord files found in {} with suffix: {}'.format(ifg_dir, fbases))
        return prm_dict

    # read corners lat/lon info
    ds = gdal.Open(geo_files[0], gdal.GA_ReadOnly)
    transform = ds.GetGeoTransform()
    x_step = abs(transform[1])
    y_step = abs(transform[5]) * -1.
    if not (1e-7 < x_step < 1.):
        raise ValueError('File {} is NOT geocoded!')

    W = transform[0] - x_step / 2.
    N = transform[3] - y_step / 2.
    E = W + x_step * ds.RasterXSize
    S = N + y_step * ds.RasterYSize

    if prm_dict['ORBIT_DIRECTION'].upper().startswith('ASC'):
        prm_dict['LAT_REF1'] = str(S)
        prm_dict['LAT_REF2'] = str(S)
        prm_dict['LAT_REF3'] = str(N)
        prm_dict['LAT_REF4'] = str(N)
        prm_dict['LON_REF1'] = str(W)
        prm_dict['LON_REF2'] = str(E)
        prm_dict['LON_REF3'] = str(W)
        prm_dict['LON_REF4'] = str(E)
    else:
        prm_dict['LAT_REF1'] = str(N)
        prm_dict['LAT_REF2'] = str(N)
        prm_dict['LAT_REF3'] = str(S)
        prm_dict['LAT_REF4'] = str(S)
        prm_dict['LON_REF1'] = str(E)
        prm_dict['LON_REF2'] = str(W)
        prm_dict['LON_REF3'] = str(E)
        prm_dict['LON_REF4'] = str(W)

    return prm_dict


def get_slant_range_distance(ifg_dir, prm_dict, fbases=['corr', 'phase', 'phasefilt', 'unwrap']):
    """Get a constant slant range distance in the image center, for dataset in geo-coord."""
    # grab an arbitrary file in radar-coordiantes
    rdr_files = [os.path.join(ifg_dir, '{}.grd'.format(i)) for i in fbases]
    if len(rdr_files) == 0:
        raise ValueError('No radar-coord files found in {} with suffix: {}'.format(ifg_dir, fbases))

    # read width from rdr_file
    ds = gdal.Open(rdr_files[0], gdal.GA_ReadOnly)
    width = ds.RasterXSize

    near_range = float(prm_dict['STARTING_RANGE'])
    range_pixel_size = float(prm_dict['RANGE_PIXEL_SIZE'])
    slant_range_dist = near_range + range_pixel_size * width / 2.
    prm_dict['SLANT_RANGE_DISTANCE'] = slant_range_dist

    return prm_dict


def extract_gmtsar_metadata(unw_file, template_file, rsc_file=None, update_mode=True):
    """Extract metadata from GMTSAR interferogram stack."""

    # update_mode: check existing rsc_file
    if update_mode and ut.run_or_skip(rsc_file, in_file=unw_file, check_readable=False) == 'skip':
        return readfile.read_roipac_rsc(rsc_file)

    ifg_dir = os.path.dirname(unw_file)

    # 1. read *.PRM file
    prm_file = get_prm_files(ifg_dir)[0]
    meta = readfile.read_gmtsar_prm(prm_file)
    meta['PROCESSOR'] = 'gmtsar'

    # 2. read template file: HEADING, ORBIT_DIRECTION
    template = readfile.read_template(template_file)
    for key in ['HEADING', 'ORBIT_DIRECTION']:
        if key in template.keys():
            meta[key] = template[key].lower()
        else:
            raise ValueError('Attribute {} is missing! Please manually specify it in the template file.')

    # 3. grab A/RLOOKS from radar-coord data file
    meta['ALOOKS'], meta['RLOOKS'] = get_multilook_number(ifg_dir)
    meta['AZIMUTH_PIXEL_SIZE'] *= meta['ALOOKS']
    meta['RANGE_PIXEL_SIZE'] *= meta['RLOOKS']

    # 4. grab LAT/LON_REF1/2/3/4 from geo-coord data file
    meta = get_lalo_ref(ifg_dir, meta)

    # 5. grab X/Y_FIRST/STEP from unw_file if in geo-coord
    ds = gdal.Open(unw_file, gdal.GA_ReadOnly)
    transform = ds.GetGeoTransform()
    x_step = abs(transform[1])
    y_step = abs(transform[5]) * -1.
    if 1e-7 < x_step < 1.:
        meta['X_STEP'] = x_step
        meta['Y_STEP'] = y_step
        meta['X_FIRST'] = transform[0] - x_step / 2.
        meta['Y_FIRST'] = transform[3] - y_step / 2.
        # constrain longitude within (-180, 180]
        if meta['X_FIRST'] > 180.:
            meta['X_FIRST'] -= 360.

    # 6. extra metadata for the missing geometry dataset: SLANT_RANGE_DISTANCE / INCIDENCE_ANGLE
    # for dataset in geo-coordinates
    if 'Y_FIRST' in meta.keys():
        meta = get_slant_range_distance(ifg_dir, meta)
        Re = float(meta['EARTH_RADIUS'])
        H = float(meta['HEIGHT'])
        Rg = float(meta['SLANT_RANGE_DISTANCE'])
        Inc = (np.pi - np.arccos((Re**2 + Rg**2 - (Re+H)**2) / (2*Re*Rg))) * 180./np.pi
        meta['INCIDENCE_ANGLE'] = Inc

    # convert all value to string format
    for key, value in meta.items():
        meta[key] = str(value)

    # write to .rsc file
    meta = readfile.standardize_metadata(meta)
    if rsc_file:
        print('writing ', rsc_file)
        os.makedirs(os.path.dirname(rsc_file), exist_ok=True)
        writefile.write_roipac_rsc(meta, rsc_file)

    return meta


def prepare_geometry(geom_files, meta, update_mode=True):
    """Prepare .rsc file for all geometry files."""
    num_file = len(geom_files)
    if num_file == 0:
        raise FileNotFoundError('NO geometry file found!')

    # write .rsc file for each geometry file
    for geom_file in geom_files:
        # copy over the common metadata
        geom_meta = {}
        for key, value in meta.items():
            geom_meta[key] = value

        # update from .grd file
        geom_meta.update(readfile.read_gdal_vrt(geom_file))

        # write .rsc file
        rsc_file = geom_file+'.rsc'
        writefile.write_roipac_rsc(geom_meta, rsc_file,
                                   update_mode=update_mode,
                                   print_msg=True)

    return


def prepare_stack(unw_files, meta, update_mode=True):
    """Prepare .rsc file for all unwrapped interferogram files."""
    num_file = len(unw_files)
    if num_file == 0:
        raise FileNotFoundError('NO unwrapped interferogram file found!')

    # write .rsc file for each interferogram file
    prog_bar = ptime.progressBar(maxValue=num_file)
    for i, unw_file in enumerate(unw_files):
        ifg_dir = os.path.dirname(unw_file)
        ifg_meta = {}

        # copy over the common metadata
        for key, value in meta.items():
            ifg_meta[key] = value

        # update from .grd file
        ifg_meta.update(readfile.read_gdal_vrt(unw_file))

        # add DATE12
        prm_files = get_prm_files(ifg_dir)
        date1, date2 = [os.path.splitext(os.path.basename(i))[0] for i in prm_files]
        ifg_meta['DATE12'] = '{}-{}'.format(ptime.yymmdd(date1), ptime.yymmdd(date2))

        # and P_BASELINE_TOP/BOTTOM_HDR
        baseline_file = os.path.join(ifg_dir, 'baseline.txt')
        if os.path.isfile(baseline_file):
            bDict = readfile.read_template(baseline_file)
            ifg_meta['P_BASELINE_TOP_HDR'] = bDict['B_perpendicular']
            ifg_meta['P_BASELINE_BOTTOM_HDR'] = bDict['B_perpendicular']
        else:
            ifg_meta['P_BASELINE_TOP_HDR'] = '0'
            ifg_meta['P_BASELINE_BOTTOM_HDR'] = '0'
            msg = 'WARNING: No baseline file found in: {}. '.format(baseline_file)
            msg += 'Set P_BASELINE* to 0 and continue.'
            print(msg)

        # write .rsc file
        rsc_file = unw_file+'.rsc'
        writefile.write_roipac_rsc(ifg_meta, rsc_file,
                                   update_mode=update_mode,
                                   print_msg=False)

        prog_bar.update(i+1, suffix='{}_{}'.format(date1, date2))
    prog_bar.close()
    return


#########################################################################
def main(iargs=None):
    inps = cmd_line_parse(iargs)

    # read file path from template file
    template = readfile.read_template(inps.template_file)
    inps.unw_files = sorted(glob.glob(template['mintpy.load.unwFile']))
    inps.cor_files = sorted(glob.glob(template['mintpy.load.corFile']))
    inps.dem_file = glob.glob(template['mintpy.load.demFile'])[0]

    # extract common metadata
    rsc_file = os.path.join(inps.mintpy_dir, 'inputs/data.rsc')
    meta = extract_gmtsar_metadata(unw_file=inps.unw_files[0],
                                   template_file=inps.template_file,
                                   rsc_file=rsc_file,
                                   update_mode=inps.update_mode)

    # prepare metadata for geometry files
    prepare_geometry([inps.dem_file], meta=meta, update_mode=inps.update_mode)

    # prepare metadata for interferogram files
    prepare_stack(inps.unw_files, meta=meta, update_mode=inps.update_mode)

    print('Done.')
    return


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