# first stage of bringing MODIS radiances into terms of AIRS spatial
# footprints (after collocation)

import numpy
import sys

from flags import *
from pyhdf.SD import SD, SDC
from util import HdfWriter

# helper function for reading in MODIS radiances. returns a numpy
# array with dimensions:
#   - 16 (MODIS band, 20-25,27-36)
#   - 2030 or 2040 (MODIS along track index)
#   - 1354 (MODIS across track index)
def read_modis_radiances(filename):

    sd = SD(filename)
    sds = sd.select('EV_1KM_Emissive')
    scales = numpy.array(sds.radiance_scales).reshape((16, 1, 1))
    offsets = numpy.array(sds.radiance_offsets).reshape((16, 1, 1))
    raw = sds.get()
    bad_idxs = (raw >= 32768)
    rads = scales * (raw - offsets)
    rads[bad_idxs] = -9999.0
    return rads
    
# verify command line arguments
if len(sys.argv) != 4:
    print ('usage: %s <modis_file> <collocation_file> <output_file>' %
           sys.argv[0])
    sys.exit(1)
modis_file = sys.argv[1]
coll_file = sys.argv[2]
output_file = sys.argv[3]

# read in the MODIS radiances
in_rads = read_modis_radiances(modis_file)

# and the collocation data
coll_sd = SD(coll_file)
num_colls = coll_sd.select('Number_Of_MODIS_FOV').get()
modis_along_idxs = coll_sd.select('MODIS_Along_Track_IDX').get()
modis_across_idxs = coll_sd.select('MODIS_Across_Track_IDX').get()

# set up numpy arrays for output
n_arr = numpy.empty((16, 135, 90), numpy.int32)
rad_sum_arr = numpy.empty((16, 135, 90), numpy.float64)
rad_squared_sum_arr = numpy.empty((16, 135, 90), numpy.float64)
flags_arr = numpy.zeros((16, 135, 90), numpy.int8)

# loop over each AIRS pixel
for airs_row in range(135):
    for airs_col in range(90):

        # avoid excessive indexing in the code below
        n = n_arr[:,airs_row,airs_col]
        rads_sum = rad_sum_arr[:,airs_row,airs_col]
        rads_squared_sum = rad_squared_sum_arr[:,airs_row,airs_col]
        flags = flags_arr[:,airs_row,airs_col]

        # use the collocation data to pull out overlapping MODIS
        # radiances (skip this AIRS FOV if there are none)
        num = num_colls[airs_row,airs_col]
        if num == 0:
            n[:] = 0
            rads_sum[:] = 0.0
            rads_squared_sum[:] = 0.0
            continue
        along_idxs = modis_along_idxs[airs_row,airs_col,:num] - 1
        across_idxs = modis_across_idxs[airs_row,airs_col,:num] - 1
        rads = in_rads[:,along_idxs,across_idxs]

        # if there's any invalid data amongst the collocated MODIS
        # pixels, filter it out and set the INVALID flag for the
        # affected bands
        bad_idxs = (rads == -9999.0)
        flags[bad_idxs.any(axis=1)] |= INVALID

        # determine how many pixels are being used per band
        n[:] = num - bad_idxs.sum(axis=1)

        # get radiance sum and squared sum (bad measurements are set
        # to zero so they have no effect on the result)
        # FIXME: use masked arrays for this instead?
        rads[bad_idxs] = 0.0
        rads_sum[:] = rads.sum(axis=1)
        rads_squared_sum[:] = (rads * rads).sum(axis=1)

# write results to HDF
writer = HdfWriter(output_file)
writer.write('MODIS_N', n_arr)
writer.write('MODIS_Radiance_Sum', rad_sum_arr)
writer.write('MODIS_Radiance_Squared_Sum', rad_squared_sum_arr)
writer.write('MODIS_Flags', flags_arr)