#!/usr/bin/env python
###############################################################################
# $Id: vrtwarp.py 4cd4d9dc01beab19ca0bef4735b5569b1b2fda07 2018-04-17 14:22:30 +0200 Even Rouault $
#
# Project:  GDAL/OGR Test Suite
# Purpose:  Test VRTWarpedDataset support.
# Author:   Frank Warmerdam <warmerdam@pobox.com>
#
###############################################################################
# Copyright (c) 2004, Frank Warmerdam <warmerdam@pobox.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
###############################################################################

import os
import shutil
import sys
from osgeo import gdal

sys.path.append('../pymod')

import gdaltest

###############################################################################
# Verify reading from simple existing warp definition.


def vrtwarp_1():

    tst = gdaltest.GDALTest('VRT', 'rgb_warp.vrt', 2, 21504)
    return tst.testOpen(check_filelist=False)

###############################################################################
# Create a new VRT warp in the temp directory.


def vrtwarp_2():

    try:
        os.remove('tmp/warp.vrt')
    except OSError:
        pass

    gcp_ds = gdal.OpenShared('data/rgb_gcp.vrt', gdal.GA_ReadOnly)

    gdaltest.vrtwarp_ds = gdal.AutoCreateWarpedVRT(gcp_ds)

    gcp_ds = None

    checksum = gdaltest.vrtwarp_ds.GetRasterBand(2).Checksum()
    expected = 21504
    if checksum != expected:
        gdaltest.post_reason('Got checksum of %d instead of expected %d.'
                             % (checksum, expected))
        return 'fail'

    return 'success'

###############################################################################
# Force the VRT warp file to be written to disk and close it.  Reopen, and
# verify checksum.


def vrtwarp_3():

    gdaltest.vrtwarp_ds.SetDescription('tmp/warp.vrt')
    gdaltest.vrtwarp_ds = None

    gdaltest.vrtwarp_ds = gdal.Open('tmp/warp.vrt', gdal.GA_ReadOnly)

    checksum = gdaltest.vrtwarp_ds.GetRasterBand(2).Checksum()
    expected = 21504

    gdaltest.vrtwarp_ds = None
    gdal.GetDriverByName('VRT').Delete('tmp/warp.vrt')

    if checksum != expected:
        gdaltest.post_reason('Got checksum of %d instead of expected %d.'
                             % (checksum, expected))
        return 'fail'

    return 'success'

###############################################################################
# Test implicit overviews with default source overview level strategy (AUTO)


def vrtwarp_4():

    src_ds = gdal.Open('../gcore/data/byte.tif')
    tmp_ds = gdal.GetDriverByName('GTiff').CreateCopy('tmp/vrtwarp_4.tif', src_ds)
    cs_main = tmp_ds.GetRasterBand(1).Checksum()
    tmp_ds.BuildOverviews('NONE', overviewlist=[2, 4])
    tmp_ds.GetRasterBand(1).GetOverview(0).Fill(127)
    cs_ov0 = tmp_ds.GetRasterBand(1).GetOverview(0).Checksum()
    tmp_ds.GetRasterBand(1).GetOverview(1).Fill(255)
    cs_ov1 = tmp_ds.GetRasterBand(1).GetOverview(1).Checksum()

    vrtwarp_ds = gdal.AutoCreateWarpedVRT(tmp_ds)
    tmp_ds = None

    for i in range(3):
        if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 2:
            gdaltest.post_reason('fail')
            return 'fail'
        if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
            print(i)
            gdaltest.post_reason('fail')
            return 'fail'
        if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != cs_ov0:
            gdaltest.post_reason('fail')
            return 'fail'
        if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != cs_ov1:
            gdaltest.post_reason('fail')
            return 'fail'
        if i == 0:
            vrtwarp_ds.SetDescription('tmp/vrtwarp_4.vrt')
            vrtwarp_ds = None
            vrtwarp_ds = gdal.Open('tmp/vrtwarp_4.vrt')
        elif i == 1:
            vrtwarp_ds = None
            tmp_ds = gdal.Open('tmp/vrtwarp_4.tif')
            vrtwarp_ds = gdal.AutoCreateWarpedVRT(tmp_ds)
            vrtwarp_ds.SetMetadataItem('SrcOvrLevel', 'AUTO')
            vrtwarp_ds.SetDescription('tmp/vrtwarp_4.vrt')
            tmp_ds = None

    # Add an explicit overview
    vrtwarp_ds.BuildOverviews('NEAR', overviewlist=[2, 4, 8])
    vrtwarp_ds = None

    ds = gdal.GetDriverByName('MEM').Create('', 3, 3, 1)
    ds.GetRasterBand(1).Fill(255)
    expected_cs_ov2 = ds.GetRasterBand(1).Checksum()
    ds = None

    vrtwarp_ds = gdal.Open('tmp/vrtwarp_4.vrt')
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 3:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != cs_ov0:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != cs_ov1:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(2).Checksum() != expected_cs_ov2:
        gdaltest.post_reason('fail')
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_4.vrt')
    gdal.Unlink('tmp/vrtwarp_4.tif')

    return 'success'

###############################################################################
# Test implicit overviews with selection of the upper source overview level


def vrtwarp_5():

    src_ds = gdal.Open('../gcore/data/byte.tif')
    tmp_ds = gdal.GetDriverByName('GTiff').CreateCopy('tmp/vrtwarp_5.tif', src_ds)
    cs_main = tmp_ds.GetRasterBand(1).Checksum()
    tmp_ds.BuildOverviews('NONE', overviewlist=[2, 4])
    tmp_ds.GetRasterBand(1).GetOverview(0).Fill(127)
    tmp_ds.GetRasterBand(1).GetOverview(0).Checksum()
    tmp_ds.GetRasterBand(1).GetOverview(1).Fill(255)
    tmp_ds.GetRasterBand(1).GetOverview(1).Checksum()
    tmp_ds = None

    ds = gdal.Warp('', 'tmp/vrtwarp_5.tif', options='-of MEM -ovr NONE -overwrite -ts 10 10')
    expected_cs_ov0 = ds.GetRasterBand(1).Checksum()
    ds = None

    ds = gdal.GetDriverByName('MEM').Create('', 5, 5, 1)
    ds.GetRasterBand(1).Fill(127)
    expected_cs_ov1 = ds.GetRasterBand(1).Checksum()
    ds = None

    tmp_ds = gdal.Open('tmp/vrtwarp_5.tif')
    vrtwarp_ds = gdal.AutoCreateWarpedVRT(tmp_ds)
    vrtwarp_ds.SetMetadataItem('SrcOvrLevel', 'AUTO-1')
    tmp_ds = None
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 2:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != expected_cs_ov0:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != expected_cs_ov1:
        gdaltest.post_reason('fail')
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_5.tif')

    return 'success'

###############################################################################
# Test implicit overviews with GCP


def vrtwarp_6():

    src_ds = gdal.Open('../gcore/data/byte.tif')
    tmp_ds = gdal.GetDriverByName('GTiff').CreateCopy('tmp/vrtwarp_6.tif', src_ds)
    cs_main = tmp_ds.GetRasterBand(1).Checksum()
    tmp_ds.SetGeoTransform([0, 1, 0, 0, 0, 1])  # cancel geotransform
    gcp1 = gdal.GCP()
    gcp1.GCPPixel = 0
    gcp1.GCPLine = 0
    gcp1.GCPX = 440720.000
    gcp1.GCPY = 3751320.000
    gcp2 = gdal.GCP()
    gcp2.GCPPixel = 0
    gcp2.GCPLine = 20
    gcp2.GCPX = 440720.000
    gcp2.GCPY = 3750120.000
    gcp3 = gdal.GCP()
    gcp3.GCPPixel = 20
    gcp3.GCPLine = 0
    gcp3.GCPX = 441920.000
    gcp3.GCPY = 3751320.000
    src_gcps = (gcp1, gcp2, gcp3)
    tmp_ds.SetGCPs(src_gcps, src_ds.GetProjectionRef())
    tmp_ds.BuildOverviews('NEAR', overviewlist=[2, 4])
    cs_ov0 = tmp_ds.GetRasterBand(1).GetOverview(0).Checksum()
    cs_ov1 = tmp_ds.GetRasterBand(1).GetOverview(1).Checksum()

    vrtwarp_ds = gdal.AutoCreateWarpedVRT(tmp_ds)
    vrtwarp_ds.SetDescription('tmp/vrtwarp_6.vrt')
    vrtwarp_ds = None
    tmp_ds = None

    vrtwarp_ds = gdal.Open('tmp/vrtwarp_6.vrt')

    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 2:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != cs_ov0:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != cs_ov1:
        gdaltest.post_reason('fail')
        return 'fail'

    gdal.Unlink('tmp/vrtwarp_6.vrt')
    gdal.Unlink('tmp/vrtwarp_6.tif')

    return 'success'

###############################################################################
# Test implicit overviews with GCP (TPS)


def vrtwarp_7():

    src_ds = gdal.Open('../gcore/data/byte.tif')
    tmp_ds = gdal.GetDriverByName('GTiff').CreateCopy('tmp/vrtwarp_7.tif', src_ds)
    cs_main = tmp_ds.GetRasterBand(1).Checksum()
    tmp_ds.SetGeoTransform([0, 1, 0, 0, 0, 1])  # cancel geotransform
    gcp1 = gdal.GCP()
    gcp1.GCPPixel = 0
    gcp1.GCPLine = 0
    gcp1.GCPX = 440720.000
    gcp1.GCPY = 3751320.000
    gcp2 = gdal.GCP()
    gcp2.GCPPixel = 0
    gcp2.GCPLine = 20
    gcp2.GCPX = 440720.000
    gcp2.GCPY = 3750120.000
    gcp3 = gdal.GCP()
    gcp3.GCPPixel = 20
    gcp3.GCPLine = 0
    gcp3.GCPX = 441920.000
    gcp3.GCPY = 3751320.000
    src_gcps = (gcp1, gcp2, gcp3)
    tmp_ds.SetGCPs(src_gcps, src_ds.GetProjectionRef())
    tmp_ds.BuildOverviews('NEAR', overviewlist=[2, 4])
    cs_ov0 = tmp_ds.GetRasterBand(1).GetOverview(0).Checksum()
    cs_ov1 = tmp_ds.GetRasterBand(1).GetOverview(1).Checksum()
    tmp_ds = None

    vrtwarp_ds = gdal.Warp('tmp/vrtwarp_7.vrt', 'tmp/vrtwarp_7.tif', options='-overwrite -of VRT -tps')
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 2:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
        print(cs_main)
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != cs_ov0:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != cs_ov1:
        gdaltest.post_reason('fail')
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_7.vrt')
    gdal.Unlink('tmp/vrtwarp_7.tif')

    return 'success'

###############################################################################
# Test implicit overviews with RPC


def vrtwarp_8():

    shutil.copy('../gcore/data/byte.tif', 'tmp/vrtwarp_8.tif')
    shutil.copy('../gcore/data/test_rpc.txt', 'tmp/vrtwarp_8_rpc.txt')
    ds = gdal.Open('tmp/vrtwarp_8.tif', gdal.GA_Update)
    ds.BuildOverviews('NEAR', overviewlist=[2])
    ds = None

    ds = gdal.Warp('', 'tmp/vrtwarp_8.tif', options='-of MEM -rpc')
    expected_cs_main = ds.GetRasterBand(1).Checksum()
    ds = None

    vrtwarp_ds = gdal.Warp('tmp/vrtwarp_8.vrt', 'tmp/vrtwarp_8.tif', options='-overwrite -of VRT -rpc')
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 1:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != expected_cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != 1214:
        gdaltest.post_reason('fail')
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).XSize)
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).YSize)
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum())
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_8.vrt')
    gdal.Unlink('tmp/vrtwarp_8.tif')
    gdal.Unlink('tmp/vrtwarp_8_rpc.txt')

    return 'success'

###############################################################################
# Test implicit overviews with GEOLOCATION


def vrtwarp_9():

    shutil.copy('../gcore/data/sstgeo.tif', 'tmp/sstgeo.tif')

    f = open('tmp/sstgeo.vrt', 'wb')
    f.write('''<VRTDataset rasterXSize="60" rasterYSize="39">
  <Metadata domain="GEOLOCATION">
    <MDI key="SRS">GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],TOWGS84[0,0,0,0,0,0,0],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0,AUTHORITY["EPSG","8901"]],UNIT["degree",0.0174532925199433,AUTHORITY["EPSG","9108"]],AXIS["Lat",NORTH],AXIS["Long",EAST],AUTHORITY["EPSG","4326"]]</MDI>
    <MDI key="X_DATASET">tmp/sstgeo.tif</MDI>
    <MDI key="X_BAND">1</MDI>
    <MDI key="PIXEL_OFFSET">0</MDI>
    <MDI key="PIXEL_STEP">1</MDI>
    <MDI key="Y_DATASET">tmp/sstgeo.tif</MDI>
    <MDI key="Y_BAND">2</MDI>
    <MDI key="LINE_OFFSET">0</MDI>
    <MDI key="LINE_STEP">1</MDI>
  </Metadata>
  <VRTRasterBand dataType="Int16" band="1">
    <ColorInterp>Gray</ColorInterp>
    <NoDataValue>-32767</NoDataValue>
    <SimpleSource>
      <SourceFilename relativeToVRT="1">sstgeo.tif</SourceFilename>
      <SourceBand>3</SourceBand>
      <SrcRect xOff="0" yOff="0" xSize="60" ySize="39"/>
      <DstRect xOff="0" yOff="0" xSize="60" ySize="39"/>
    </SimpleSource>
  </VRTRasterBand>
</VRTDataset>
'''.encode('ascii'))
    f.close()
    ds = gdal.Open('tmp/sstgeo.vrt', gdal.GA_Update)
    ds.BuildOverviews('NEAR', overviewlist=[2])
    ds = None

    ds = gdal.Warp('', 'tmp/sstgeo.vrt', options='-of MEM -geoloc')
    expected_cs_main = ds.GetRasterBand(1).Checksum()
    ds = None

    vrtwarp_ds = gdal.Warp('tmp/vrtwarp_9.vrt', 'tmp/sstgeo.vrt', options='-overwrite -of VRT -geoloc')
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 1:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != expected_cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != 63859:
        gdaltest.post_reason('fail')
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).XSize)
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).YSize)
        print(vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum())
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_9.vrt')
    gdal.Unlink('tmp/sstgeo.vrt')
    gdal.Unlink('tmp/sstgeo.vrt.ovr')
    gdal.Unlink('tmp/sstgeo.tif')

    return 'success'

###############################################################################
# Test implicit overviews with selection of the full resolution level


def vrtwarp_10():

    src_ds = gdal.Open('../gcore/data/byte.tif')
    tmp_ds = gdal.GetDriverByName('GTiff').CreateCopy('tmp/vrtwarp_10.tif', src_ds)
    cs_main = tmp_ds.GetRasterBand(1).Checksum()
    tmp_ds.BuildOverviews('NONE', overviewlist=[2, 4])
    tmp_ds.GetRasterBand(1).GetOverview(0).Fill(127)
    tmp_ds.GetRasterBand(1).GetOverview(0).Checksum()
    tmp_ds.GetRasterBand(1).GetOverview(1).Fill(255)
    tmp_ds.GetRasterBand(1).GetOverview(1).Checksum()
    tmp_ds = None

    ds = gdal.Warp('', 'tmp/vrtwarp_10.tif', options='-of MEM -ovr NONE -ts 10 10')
    expected_cs_ov0 = ds.GetRasterBand(1).Checksum()
    ds = None

    ds = gdal.Warp('', 'tmp/vrtwarp_10.tif', options='-of MEM -ovr NONE -ts 5 5')
    expected_cs_ov1 = ds.GetRasterBand(1).Checksum()
    ds = None

    tmp_ds = gdal.Open('tmp/vrtwarp_10.tif')
    vrtwarp_ds = gdal.AutoCreateWarpedVRT(tmp_ds)
    vrtwarp_ds.SetMetadataItem('SrcOvrLevel', 'NONE')
    tmp_ds = None
    if vrtwarp_ds.GetRasterBand(1).GetOverviewCount() != 2:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).Checksum() != cs_main:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(0).Checksum() != expected_cs_ov0:
        gdaltest.post_reason('fail')
        return 'fail'
    if vrtwarp_ds.GetRasterBand(1).GetOverview(1).Checksum() != expected_cs_ov1:
        gdaltest.post_reason('fail')
        return 'fail'
    vrtwarp_ds = None

    gdal.Unlink('tmp/vrtwarp_10.tif')

    return 'success'

###############################################################################
# Test implicit overviews with dest alpha band (#6081)


def vrtwarp_11():

    ds = gdal.Open('data/bug6581.vrt')
    cs1 = ds.GetRasterBand(1).Checksum()
    cs2 = ds.GetRasterBand(2).Checksum()
    cs3 = ds.GetRasterBand(3).Checksum()
    ds = None

    if cs1 != 22122 or cs2 != 56685 or cs3 != 22122:
        gdaltest.post_reason('fail')
        print(cs1)
        print(cs2)
        print(cs3)
        return 'fail'

    return 'success'

###############################################################################
# Test reading a regular VRT whose source is a warped VRT inlined


def vrtwarp_read_vrt_of_warped_vrt():

    ds = gdal.Open('data/vrt_of_warped_vrt.vrt')
    cs = ds.GetRasterBand(1).Checksum()
    if cs != 4672:
        print(cs)
        return 'fail'

    return 'success'

###############################################################################
# Test different nodata values on bands and partial blocks (#6581)


gdaltest_list = [
    vrtwarp_1,
    vrtwarp_2,
    vrtwarp_3,
    vrtwarp_4,
    vrtwarp_5,
    vrtwarp_6,
    vrtwarp_7,
    vrtwarp_8,
    vrtwarp_9,
    vrtwarp_10,
    vrtwarp_11,
    vrtwarp_read_vrt_of_warped_vrt
]


if __name__ == '__main__':

    gdaltest.setup_run('vrtwarp')

    gdaltest.run_tests(gdaltest_list)

    gdaltest.summarize()
