# tifffile/examples/issue125.py

"""Create a Fsspec ReferenceFileSystem for a sequence of TIFF files on S3.

This Python script uses the Tifffile and Fsspec libraries to create a
multiscale ReferenceFileSystem JSON file for a sequence of cloud optimized
GeoTIFF (COG) files stored on S3. The tiles of the COG files are used as
chunks. No additional Numcodecs codec needs to be registered since the COG
files use Zlib compression. An Xarray dataset is created from the
ReferenceFileSystem file and a subset of the dataset is plotted.

See https://github.com/cgohlke/tifffile/issues/125

"""

import fsspec
import kerchunk.utils
import xarray
from matplotlib import pyplot

import tifffile

# get a list of cloud optimized GeoTIFF files stored on S3
remote_options = {
    'anon': True,
    'client_kwargs': {'endpoint_url': 'https://mghp.osn.xsede.org'},
}
fs = fsspec.filesystem('s3', **remote_options)
files = [f's3://{f}' for f in fs.ls('/rsignellbucket1/lcmap/cog')]

# write the ReferenceFileSystem of each file to a JSON file
with open('issue125.json', 'w', encoding='utf-8', newline='\n') as jsonfile:
    for i, filename in enumerate(tifffile.natural_sorted(files)):
        url, name = filename.rsplit('/', 1)
        with fs.open(filename) as fh, tifffile.TiffFile(fh, name=name) as tif:
            print(tif.geotiff_metadata)
            with tif.series[0].aszarr() as store:
                store.write_fsspec(
                    jsonfile,
                    url=url,
                    zarr_format=2,
                    # using an experimental API:
                    _shape=[len(files)],  # shape of file sequence
                    _axes=['T'],  # axes of file sequence
                    _index=[i],  # index of this file in sequence
                    _append=i > 0,  # if True, only write index keys+value
                    _close=i == len(files) - 1,  # if True, no more appends
                    # groupname='0',  # required for non-pyramidal series
                )


# create an Xarray dataset from the ReferenceFileSystem file

store = kerchunk.utils.refs_as_store('issue125.json')
dataset = xarray.open_zarr(store, consolidated=False, mask_and_scale=False)
print(dataset)

# plot a slice of the 5th pyramidal level of the dataset
xarray.plot.imshow(dataset['5'][0, 32:-32, 32:-32])
pyplot.show()
