"""JAXA Earth API 0.1.6: two monthly mean precipitation-rate rasters.
Data is fetched by the reader. Read collection license before commercial use.
No API key, GPU, LLM API, or TLS-verification bypass is used.
"""
import argparse
import json
from datetime import datetime, timezone
from pathlib import Path

COLLECTION = 'JAXA.EORC_GSMaP_standard.Gauge.00Z-23Z.v6_monthly'
DATES = ['2024-06-01T00:00:00', '2024-07-31T23:59:59']
MONTHS = ['2024-06', '2024-07']


def validate_bbox(bbox):
    import math
    if len(bbox) != 4 or not all(math.isfinite(x) for x in bbox):
        raise ValueError('bbox must contain four finite numbers')
    west, south, east, north = bbox
    if not (-180 <= west < east <= 180 and -60 <= south < north <= 60):
        raise ValueError('Use west < east, south < north; GSMaP latitude range is -60..60')
    if east - west > 20 or north - south > 20:
        raise ValueError('This small-area tutorial allows at most 20 degrees per side')


def export(collection, out):
    import numpy as np
    import rasterio
    from rasterio.transform import from_bounds
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    a = collection.raster.img
    ids = [x.rstrip('/') for x in collection.stac_date.id]
    if ids != MONTHS or a.ndim != 4 or a.shape[0] != 2 or a.shape[3] != 1:
        raise RuntimeError(f'Unexpected response: dates={ids}, shape={a.shape}')
    if collection.proj_params.epsg != 4326:
        raise RuntimeError('Expected EPSG:4326')
    units = [x.unit for x in collection.cinfo]
    if units != ['mm/hr', 'mm/hr']:
        raise RuntimeError(f'Unit changed: {units}; inspect metadata before proceeding')
    south, north = map(float, collection.raster.latlim[0])
    west, east = map(float, collection.raster.lonlim[0])
    h, w = a.shape[1:3]
    transform = from_bounds(west, south, east, north, w, h)
    arrays = []
    out.mkdir(parents=True, exist_ok=True)
    for i, month in enumerate(ids):
        data = np.asarray(a[i, :, :, 0], dtype='float32').copy()
        # All negative values are invalid for this precipitation-rate product.
        data[(data < 0) | ~np.isfinite(data)] = np.nan
        if not np.isfinite(data).any():
            raise RuntimeError(f'No valid data for {month}')
        path = out / f'gsmap-{month}.tif'
        with rasterio.open(path, 'w', driver='GTiff', height=h, width=w,
                           count=1, dtype='float32', crs='EPSG:4326',
                           transform=transform, nodata=float('nan'), compress='deflate') as dst:
            dst.write(data, 1)
            dst.update_tags(collection=COLLECTION, month=month, unit='mm/hr',
                            attribution='JAXA / GSMaP; acquired with JAXA Earth API')
        # Check exported orientation, pixel-center coordinates, and all values.
        with rasterio.open(path) as src:
            np.testing.assert_allclose(src.read(1), data, equal_nan=True)
            assert src.crs.to_epsg() == 4326 and src.transform.e < 0
            np.testing.assert_allclose(src.xy(0, 0),
                                       (west + (east-west)/w/2, north - (north-south)/h/2))
        arrays.append(data)
    vmax = float(max(np.nanmax(x) for x in arrays))
    fig, axes = plt.subplots(1, 2, figsize=(11, 5), layout='constrained')
    for ax, data, month in zip(axes, arrays, ids):
        im = ax.imshow(data, origin='upper', extent=(west, east, south, north),
                       vmin=0, vmax=vmax, cmap='viridis', interpolation='nearest')
        ax.set(title=month, xlabel='Longitude (deg E)', ylabel='Latitude (deg N)')
    fig.colorbar(im, ax=axes, label='Monthly mean precipitation rate (mm/hr)', shrink=0.8)
    fig.suptitle('JAXA / GSMaP via JAXA Earth API | same color scale; no coastline overlay')
    fig.savefig(out/'rain-comparison.png', dpi=160)
    plt.close(fig)
    metadata = {
        'retrievedAt': datetime.now(timezone.utc).isoformat(),
        'collection': COLLECTION, 'datesUTC': DATES, 'months': ids,
        'bbox': [west, south, east, north], 'shape': list(a.shape),
        'ppu': float(collection.raster.ppu), 'epsg': 4326, 'unit': 'mm/hr',
        'meaning': 'Monthly mean precipitation rate, NOT monthly accumulated rainfall',
        'source': collection.stac_collection.url[0][0],
        'license': collection.stac_collection.json[0][0].get('license'),
        'licenseLinks': [x['href'] for x in collection.stac_collection.json[0][0]['links']
                         if x['rel'] == 'license'],
        'attribution': 'JAXA / GSMaP; acquired with JAXA Earth API for Python 0.1.6',
    }
    (out/'metadata.json').write_text(json.dumps(metadata, indent=2), encoding='utf-8')
    print(json.dumps({'months': ids, 'shape': list(a.shape), 'output': str(out)}))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--bbox', type=float, nargs=4, default=[128, 30, 146, 46],
                   metavar=('WEST', 'SOUTH', 'EAST', 'NORTH'))
    p.add_argument('--out', type=Path, default=Path('jaxa-output'))
    args = p.parse_args()
    try:
        validate_bbox(args.bbox)
    except ValueError as exc:
        p.error(str(exc))
    from jaxa.earth import je
    c = (je.ImageCollection(COLLECTION, ssl_verify=True)
         .filter_date(DATES).filter_resolution(10)
         .filter_bounds(args.bbox).select('PRECIP').get_images())
    export(c, args.out)


if __name__ == '__main__':
    main()