#!/usr/bin/env python

import sys
from itertools import chain

import numpy as np

from osgeo import gdal
from osgeo import gdalconst

from matplotlib import pyplot as plt

from shapely.geometry import Polygon
from shapely.geometry.polygon import LinearRing

bounds = Polygon(
    LinearRing([
        (137.8, -10.6),
        (153.2, -10.6),
        (153.2, -28.2),
        (137.8, -28.2),
        (137.8, -10.6)
    ])
).bounds

ds = gdal.Open(sys.argv[1], gdalconst.GA_ReadOnly)

inv_matrix = gdal.InvGeoTransform(ds.GetGeoTransform())[1]


bounsds_xy = list(chain(*[gdal.ApplyGeoTransform(inv_matrix, x, y)
                          for x, y in zip(bounds[::2], bounds[1::2])]))

minX, minY, maxX, maxY = bounds
ulX, ulY = gdal.ApplyGeoTransform(inv_matrix, minX, maxY)
lrX, lrY = gdal.ApplyGeoTransform(inv_matrix, maxX, minY)

xOffset = int(ulX)
yOffset = int(ulY)
ySize = int(lrY - ulY)
xSize = int(lrX - ulX)
size = min(xSize, ySize)

xs = ds.ReadAsArray(xOffset, yOffset, xSize, ySize)
xs[xs == -9999] = np.nan

plt.clf()
plt.imshow(xs)
plt.colorbar()
plt.savefig("test.png")

driver = gdal.GetDriverByName("GTiff")
dst_ds = driver.Create(
    "test.tif", xs.shape[1], xs.shape[0], 1, gdal.GDT_Float32
)
dst_ds.GetRasterBand(1).WriteArray(xs)