import os
import numpy as np
import h5py

#
# 2026-03-19
# Land, Vegetation, and Ice (LVIS) Team
# http://lvis.gsfc.nasa.gov/
#
# EXAMPLE USAGE:
#   from read_lvis1b import read_lvis1b
#   data = read_lvis1b(lvis1b_file)
#   data = read_lvis1b(lvis1b_file, quiet=True)
#

def read_lvis1b(file, quiet=False):
    """
    Python code to read in all LVIS Level-1B release files
       v1.04, v1.05, v2.0.1, v2.0.2, v2.0.3

    Parameters
    ----------
    file : str
        LVIS Level 1B .h5 file to read in.
    quiet : bool, optional
        If False (default), prints status messages.

    Returns
    -------
    data : dict
        Data returned as a dictionary of numpy arrays, or None on failure.

    """

    # -------------------------------------------------------------------------
    # File check
    # -------------------------------------------------------------------------
    if not os.path.isfile(file):
        print(f'File not found: {file}')
        return None

    infile = file

    if not quiet:
        print(f'Reading in LVIS LEVEL-1B file: {infile}')

    # -------------------------------------------------------------------------
    # Open HDF5 file and read all datasets
    # -------------------------------------------------------------------------
    with h5py.File(infile, 'r') as hdf:

        # Check number of shots
        numshots = hdf['LFID'].shape[0]
        if numshots <= 0:
            print(f'No shots found: LFID has {numshots} elements')
            return None

        # ---------------------------------------------------------------------
        # Read rxwave and txwave to determine format version 
        # ---------------------------------------------------------------------

        rxwave = hdf['RXWAVE'][:]   # shape: (numshots, rx_samples)
        txwave = hdf['TXWAVE'][:]   # shape: (numshots, tx_samples)
        rx_samples = rxwave.shape[1]
        tx_samples = txwave.shape[1]

        if not quiet:
            print(f'rx samples: {rx_samples},  tx samples: {tx_samples}')

        # v2 LVISF: rx=1216, tx=128 or tx=256
        if rx_samples == 1216:
            if not quiet:
                print('LDS v2 LVIS-Facility')
            lon_low_key = 'LON1215'
            lat_low_key = 'LAT1215'
            z_low_key   = 'Z1215'
            low_sample  = rx_samples - 1   # 1215
            has_date    = False

        # v2 LVISC: rx=1024
        elif rx_samples == 1024:
            if not quiet:
                print('LDS v2 LVIS-Classic')
            lon_low_key = 'LON1023'
            lat_low_key = 'LAT1023'
            z_low_key   = 'Z1023'
            low_sample  = rx_samples - 1   # 1023
            has_date    = False

        # v1.04 IceBridge: rx=528
        elif rx_samples == 528:
            if not quiet:
                print('LDS v1.04: IceBridge 2009-2015')
            lon_low_key = 'LON527'
            lat_low_key = 'LAT527'
            z_low_key   = 'Z527'
            low_sample  = rx_samples - 1   # 527
            has_date    = False

        # v1.05 2003-2009: rx=432
        elif rx_samples == 432:
            if not quiet:
                print('LDS v1.05: Rereleases of 2003-2009 data')
            lon_low_key = 'LON431'
            lat_low_key = 'LAT431'
            z_low_key   = 'Z431'
            low_sample  = rx_samples - 1   # 431
            has_date    = True

        # v1.05 1998-1999: rx=352
        elif rx_samples == 352:
            if not quiet:
                print('LDS v1.05: Rereleases of 1998-1999 data')
            lon_low_key = 'LON351'
            lat_low_key = 'LAT351'
            z_low_key   = 'Z351'
            low_sample  = rx_samples - 1   # 351
            has_date    = True

        else:
            print('This file does not have the standard dimensions '
                  'for an LVIS Level-1B release file: {infile}')
            print(f'rx waveform samples: {rx_samples}')
            print(f'tx waveform samples: {tx_samples}')
            return None

        # ---------------------------------------------------------------------
        # Read all common fields
        # ---------------------------------------------------------------------
        data = {
            'lfid':          hdf['LFID'][:].astype(np.uint32),
            'shotnumber':    hdf['SHOTNUMBER'][:].astype(np.uint32),
            'azimuth':       hdf['AZIMUTH'][:].astype(np.float32),
            'incidentangle': hdf['INCIDENTANGLE'][:].astype(np.float32),
            'range':         hdf['RANGE'][:].astype(np.float32),
            'time':          hdf['TIME'][:].astype(np.float64),
            'lon0':          hdf['LON0'][:].astype(np.float64),
            'lat0':          hdf['LAT0'][:].astype(np.float64),
            'z0':            hdf['Z0'][:].astype(np.float32),
            'sigmean':       hdf['SIGMEAN'][:].astype(np.float32),
            'txwave':        txwave.astype(np.uint16),
            'rxwave':        rxwave.astype(np.uint16),
        }

        # Version-specific low-end waveform position fields
        data[f'lon{low_sample}'] = hdf[lon_low_key][:].astype(np.float64)
        data[f'lat{low_sample}'] = hdf[lat_low_key][:].astype(np.float64)
        data[f'z{low_sample}']   = hdf[z_low_key][:].astype(np.float32)

        # v1.05 only: date field
        if has_date:
            data['date'] = hdf['DATE'][:].astype(np.int32)

    return data


# =============================================================================
# Example usage
# =============================================================================
if __name__ == '__main__':

    lvis_file = 'LVIS_file.h5'

    # Retrieve data
    data = read_lvis1b(lvis_file)

    # Retrieve data silently
    data = read_lvis1b(lvis_file, quiet=True)

    if data is not None:
        print(f'Number of shots:  {len(data["lfid"])}')
        print(f'Fields available: {list(data.keys())}')
        print(f'First lfid:       {data["lfid"][0]}')
        print(f'First time:       {data["time"][0]}')
        print(f'rxwave shape:     {data["rxwave"].shape}')
        print(f'txwave shape:     {data["txwave"].shape}')