# preview_snap.py
"""Create a PNG preview of the B3/B8 output from the supplied SNAP graph."""
import argparse
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import rasterio


def stretch(band):
    """Exclude nodata/non-finite values and stretch the valid 2–98% range."""
    band = np.ma.masked_invalid(np.ma.asarray(band, dtype=float))
    values = band.compressed()
    if not values.size:
        raise ValueError('A band contains no valid pixels.')
    low, high = np.percentile(values, [2, 98])
    if high <= low:
        raise ValueError('A band has no contrast between its 2nd and 98th percentiles.')
    return np.ma.clip((band - low) / (high - low), 0, 1), low, high


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('tiff', type=Path, help='Downloaded two-band SNAP GeoTIFF')
    parser.add_argument('--output', type=Path, default=Path('snap-subset-preview.png'))
    args = parser.parse_args()
    if args.output.suffix.lower() != '.png':
        parser.error('--output must name a PNG file.')
    if args.output.exists():
        parser.error('Output already exists. Choose a different --output filename.')

    with rasterio.open(args.tiff) as dataset:
        if dataset.count != 2:
            parser.error('Expected the two-band B3/B8 output from the supplied graph.')
        bands = dataset.read(masked=True)
    try:
        previews = [stretch(band) for band in bands]
    except ValueError as error:
        parser.error(str(error))

    # These band names and their order match nsis_l1c_subset.xml.
    names = ('B3', 'B8')
    cmap = plt.get_cmap('gray').copy()
    cmap.set_bad('#808080')
    figure, axes = plt.subplots(1, 2, figsize=(11, 6))
    figure.subplots_adjust(left=.06, right=.98, bottom=.15, top=.82, wspace=.18)
    for name, axis, (preview, low, high) in zip(names, axes, previews):
        axis.imshow(preview, cmap=cmap, vmin=0, vmax=1, interpolation='nearest')
        axis.set_title(f'{name}: 2–98% display stretch\nDN {low:.0f}–{high:.0f}')
        axis.set_xlabel('Column')
        axis.set_ylabel('Row')
        print(f'{name}: display range {low:.2f}–{high:.2f} DN')
    figure.suptitle('NSIS SNAP subset — separate band previews', fontsize=15, y=.97)
    figure.text(.5, .055, 'Display adjustment only; bands stretched independently.',
                ha='center', fontsize=10)
    figure.text(.5, .02, 'Nodata shown in grey. Original TIFF unchanged.',
                ha='center', fontsize=10)
    try:
        figure.savefig(args.output, dpi=140)
    finally:
        plt.close(figure)
    print('Preview saved:', args.output.resolve())


if __name__ == '__main__':
    main()