#!/usr/bin/env python

# This file is part of chelsa_isimip3b_ba_1km.
#
# chelsa_isimip3b_ba_1km is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# chelsa_isimip3b_ba_1km is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with chelsa_isimip3b_ba_1km.  If not, see <https://www.gnu.org/licenses/>.


# import library ##############################################################
import os
import sys
from tqdm import tqdm

# *************************************************
# import functions and classes
# *************************************************
# add the directory of this script to allow module search
sys.path.append(os.path.dirname(os.path.abspath(os.path.join(sys.argv[0], '..'))))
from chelsa_isimip3b_ba_1km.functions import saga_functions
from chelsa_isimip3b_ba_1km.functions.chelsa_data_classes import Dem_data
from chelsa_isimip3b_ba_1km.functions.chelsa_data_classes import Aux_data

###############################################################################
# auxiliary functions #########################################################
def get_cmd_args():

    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

    parser = ArgumentParser(
        description='''# This python code is adapted for CHELSA_V2.1_ISIMIP3b
                         the code is adapted to the ISIMIP3b_BA data. The prepare_input code can be
                         used to clip all neccessary input
                         data to a given spatial extent to run CHELSA on a limited spatial extent.
                         Dependencies for ubuntu_18.04:
                         libwxgtk3.0-dev libtiff5-dev libgdal-dev libproj-dev 
                         libexpat-dev wx-common libogdi3.2-dev unixodbc-dev
                         g++ libpcre3 libpcre3-dev wget swig-4.0.1 python2.7-dev 
                         software-properties-common gdal-bin python-gdal 
                         python2.7-gdal libnetcdf-dev libgdal-dev
                         python-pip cdsapi saga_gis-7.6.0
                         All dependencies are resolved in the chelsa_V2.1.cont singularity container
                         Tested with: singularity version 3.3.0-809.g78ec427cc
                         ''',
        epilog='''author: Dirk N. Karger, dirk.karger@wsl.ch, Version 2.1''',
        formatter_class=ArgumentDefaultsHelpFormatter
        )
    parser._action_groups.pop()
    required = parser.add_argument_group('required arguments')
    optional = parser.add_argument_group('optional arguments')

    # collect the function arguments
    required.add_argument(
        '-i', '--input',
        dest='input',
        help='input directory containing the DEM related files etc., string',
        type=str,
        required=True,
    )
    required.add_argument(
        '-o', '--output',
        dest='output',
        help='output directory, string',
        type=str,
        required=True,
    )
    required.add_argument(
        '-au', '--aux',
        dest='aux',
        help='directory with auxilariy files, string',
        type=str,
        required=True,
    )
    required.add_argument(
        '-s', '--srad',
        dest='srad',
        help='srad input directory, string',
        type=str,
        required=True,
    )
    required.add_argument(
        '-xmn', '--xmin',
        dest='xmin',
        help='minimum longitude of the boundary box in geographic coordinates, float',
        type=float,
        required=True,
    )
    required.add_argument(
        '-xmx', '--xmax',
        dest='xmax',
        help='maximum longitude of the boundary box in geographic coordinates, float',
        type=float,
        required=True,
    )
    required.add_argument(
        '-ymn', '--ymin',
        dest='ymin',
        help='minimum latitude of the boundary box in geographic coordinates, float',
        type=float,
        required=True,
    )
    required.add_argument(
        '-ymx', '--ymax',
        dest='ymax',
        help='maximum latitude of the boundary box in geographic coordinates, float',
        type=float,
        required=True,
    )
    required.add_argument(
        '-bf', '--buffer',
        dest='buffer',
        help='buffer around the boundary box in geographic coordinates, float',
        type=float,
        required=True,
    )

    optional.add_argument(
        '-pg', '--show_progress_bar',
        dest='show_progress_bar',
        help='show progress bar (default: do not), flag',
        action='store_true',
        required=False,
        default=False,
    )

    args = parser.parse_args()

    return args

###############################################################################
# main code ###################################################################


# *************************************************
# Get arguments
# *************************************************
args = get_cmd_args()
print(args)

# *************************************************
# Set the directories from arguments
# *************************************************
INPUT = args.input
OUTPUT = args.output
AUX = args.aux
SRAD = args.srad
XMAX = args.xmax
XMIN = args.xmin
YMAX = args.ymax
YMIN = args.ymin
XYbuffer = args.buffer
show_progress_bar = bool(args.show_progress_bar)

dem_data = Dem_data(INPUT=INPUT)
aux_data = Aux_data(INPUT=INPUT, W5E5=AUX)

saga_functions.Load_Tool_Libraries(True)

# initailize the data classes
print('Start preparing input data for: xmin=' + str(XMIN-XYbuffer) + ' xmax='
                                              + str(XMAX+XYbuffer) + ' ymin='
                                              + str(YMIN-XYbuffer) + ' ymax='
                                              + str(YMAX+XYbuffer))

# create the extents
dem_data.set('dem_latlong3')

dem_clippend = saga_functions.Run_Clip_Grids(obj=dem_data.dem_latlong3,
                                             xmin=XMIN,
                                             xmax=XMAX,
                                             ymin=YMIN,
                                             ymax=YMAX,
                                             buffer=XYbuffer)

dem_clippend.Save(os.path.join(OUTPUT, 'dem_latlong3.sgrd'))

clipper = saga_functions.grid_calculator_simple(dem_clippend, 'a-a+1')
clipper_shp = saga_functions.Run_Vectorising_Grid_Classes(clipper)
clipper_shp.Save(os.path.join(OUTPUT, 'extent_wgs84_latlong.shp'))
clipper_prj = saga_functions.reproject_shape(clipper_shp)
clipper_prj.Save(os.path.join(OUTPUT, 'extent_merc.shp'))

# clip all the dem_data to the extent0
print('preparing primary input data ...')
print('clipping the projected raster ...')
rasterlist = ['demproj']
for raster in tqdm(rasterlist) if show_progress_bar else rasterlist:
    dem_data.set(raster)
    saga_functions.clip_with_polygon(dem_data.demproj, clipper_prj).Get_Grid(0).Save(os.path.join(OUTPUT, 'dem_merc3.sgrd'))
    dem_data.delete(raster)

# clip dem data
print('clipping the geographic raster ...')
rasterlist = ['dem_high', 'dem_low']
for raster in tqdm(rasterlist) if show_progress_bar else rasterlist:
    dem_data.set(raster)
    if raster == 'dem_high':
        saga_functions.clip_with_polygon(getattr(dem_data, raster), clipper_shp).Get_Grid(0).Save(
            os.path.join(OUTPUT, 'dem_latlong.sgrd'))
    if raster == 'dem_low':
        ds = getattr(dem_data, raster)
        # ds = change_latlong(ds)
        saga_functions.clip_with_polygon(ds, clipper_shp).Get_Grid(0).Save(os.path.join(OUTPUT, 'orography.sgrd'))
    dem_data.delete(raster)
print('preparing primary input data done.')

# clip all the aux data
print('preparing auxillary input data ...')
rasterlist = ['continents', 'dummy_W5E5', 'expocor', 'landseamask', 'oceans', 'patch', 'template_010',
              'template_025']
for raster in tqdm(rasterlist) if show_progress_bar else rasterlist:
    aux_data.set(raster)
    saga_functions.clip_with_polygon(getattr(aux_data, raster), clipper_shp).Get_Grid(0).Save(
        os.path.join(OUTPUT, raster + '.sgrd'))
    aux_data.delete(raster)
print('preparing auxillary input data done.')

# clip radiation data
print('preparing rsdscs input data ...')
dayrange = range(1, 367)
for day in tqdm(dayrange) if show_progress_bar else dayrange:
    file_prefix = 'CHELSA_stot_pj_'
    file_suffix = '{0}_V.2.1'.format(day)  # this is hard-coded in chelsa_data_classes...
    ds = saga_functions.import_gdal(os.path.join(SRAD, file_prefix + file_suffix + '.tif'))
    saga_functions.export_geotiff(
        saga_functions.clip_with_polygon(ds, clipper_shp).Get_Grid(0),
        os.path.join(OUTPUT, file_prefix + file_suffix + '.tif'))
    saga_functions.saga_api.SG_Get_Data_Manager().Delete(ds)
print('preparing rsdscs input data done.')