from rasterio import open
from rasterio.transform import rowcol, xy


def img2coord(tl_x, x_res, tl_y, y_res, row, col):
    ''' Straightforward version with both resolution values (+vae x, -ve y) '''
    return ((col * x_res) + tl_x, (row * y_res) + tl_y)

def coord2img(tl_x, x_res, tl_y, y_res, x, y):
    ''' Straightforward version with both resolution values (+ve x, -ve y) '''
    return (int((y - tl_y) / y_res), int((x - tl_x) / x_res)) 


def img2coord_gt(gt, row, col):
    ''' GDAL GeoTransform version '''
    return ((col * gt[1]) + gt[0], (row * gt[5]) + gt[3])

def coord2img_gt(gt, x, y):
    ''' GDAL GeoTransform version '''
    return (int((y - gt[3]) / gt[5]), int((x - gt[0]) / gt[1])) 


def img2coord_single_r(tl_x, tl_y, res, row, col):
    ''' Straightforward version with only one resolution value (+ve) '''
    return ((col * res) + tl_x, tl_y - (row * res))

def coord2img_single_r(tl_x, tl_y, res, x, y):
    ''' Straightforward version with only one resolution value (+ve) '''
    return (int((tl_y - y) / res), int((x - tl_x) / res)) 


def img2coord_affine(a, r, c):
    ''' version using the notation facilitated by the Affine library'''
    return a * (c, r)

def coord2img_affine(a, e, n):
    ''' version using the notation facilitated by the Affine library'''
    c, r = ~a * (e, n)
    return (r, c)	# note this returns (col, row), here I reverse it so it matches the others
        

# open a raster dataset (as DatasetReader object)
with open("data/helvellyn/Helvellyn-50.tif") as ds:

    # get affine transformation matrix
    a = ds.transform

    # construct GDAL style GeoTransform object
    '''
    GT(0) x-coordinate of the upper-left corner of the upper-left pixel.
    GT(1) w-e pixel resolution / pixel width.
    GT(2) row rotation (typically zero).
    GT(3) y-coordinate of the upper-left corner of the upper-left pixel.
    GT(4) column rotation (typically zero).
    GT(5) n-s pixel resolution / pixel height (negative value for a north-up image).
    '''
    gt = a.to_gdal()

    # this is how to expand GDAL GeoTransform into into named variables
    tl_x, x_res, _, tl_y, _, y_res = gt
    
    # this is how to expand Affine matrix into into named variables
    x_res, _, tl_x, _, y_res, tl_y, _, _, _ = a

    # get data from the raster
    band = ds.read(1)
    bounds = ds.bounds

    # make a set of locations for testing test (just use the corners and centre of the raster)
    locations = [
        ("Bottom Left:", bounds.left, bounds.bottom),
        ("Top Left:", bounds.left, bounds.top),
        ("Top Right:", bounds.right, bounds.top),
        ("Bottom Right:", bounds.right, bounds.bottom),
        ("Centre:", bounds.left + (bounds.right - bounds.left) / 2, bounds.bottom + (bounds.top - bounds.bottom) / 2)
    ]

    # test all functions at all locations
    for name, e, n in locations:
        print(f"\n{name}")

        print("\ngdal geotransform (returns cell origin)")
        print(e, n)
        r, c = coord2img_gt(gt, e, n)
        print(r, c)
        x, y = img2coord_gt(gt, r, c)
        print (x, y)

        print("\nboth resolution values (returns cell origin)")
        print(e, n)
        r, c = coord2img(tl_x, x_res, tl_y, y_res, e, n)
        print(r, c)
        x, y = img2coord(tl_x, x_res, tl_y, y_res, r, c)
        print (x, y)

        print("\nsingle (+ve) resolution value (returns cell origin)")
        print(e, n)
        r, c = coord2img_single_r(tl_x, tl_y, x_res, e, n)
        print(r, c)
        x, y = img2coord_single_r(tl_x, tl_y, x_res, r, c)
        print (x, y)

        print("\naffine object (returns cell origin)")
        print(e, n)
        r, c = coord2img_affine(a, e, n)
        print(r, c)
        x, y = img2coord_affine(a, r, c)
        print (x, y)

        print("\nrasterio.transform (returns cell centre)")
        print(e, n)
        r, c = rowcol(a, e, n)
        print(r, c)
        x, y = xy(a, r, c)
        print (x, y)

        print("\nrasterio via dataset object (returns cell centre)")
        print(e, n)
        r, c = ds.index(e, n)
        print(r, c)
        x, y = ds.xy(r, c)
        print (x, y)
        print("\n------\n")