# I need
from glob import glob

# Pytroll/SatPy needs
from satpy import Scene, MultiScene
from satpy.writers import compute_writer_results

# For my latest PR#2394, I copied the enhanced multiscene.py to the wd
from satpy.modifiers.angles import get_satellite_zenith_angle
from functools import partial
import multiscene
from multiscene import stack

# Output of global_scene.available_composite_names() when all 16 GOES16/17 abi channels are loaded:
# *************************************************************************************************
# allcomposites = ['airmass', 'ash', 'cimss_cloud_type', 'cimss_cloud_type_raw', 'cimss_green',
#    'cimss_green_sunz', 'cimss_green_sunz_rayleigh', 'cimss_true_color', 'cimss_true_color_sunz',
#    'cimss_true_color_sunz_rayleigh', 'cira_day_convection', 'cira_fire_temperature', 'cloud_phase',
#    'cloud_phase_distinction', 'cloud_phase_distinction_raw', 'cloud_phase_raw', 'cloudtop',
#    'color_infrared', 'colorized_ir_clouds', 'convection', 'day_microphysics', 'day_microphysics_abi',
#    'day_microphysics_eum', 'dust', 'fire_temperature_awips', 'fog', 'green', 'green_crefl', 'green_nocorr',
#    'green_raw', 'green_snow', 'highlight_C14', 'ir108_3d', 'ir_cloud_day', 'land_cloud', 'land_cloud_fire',
#    'natural_color', 'natural_color_nocorr', 'natural_color_raw', 'night_fog', 'night_ir_alpha',
#    'night_ir_with_background', 'night_ir_with_background_hires', 'night_microphysics', 'night_microphysics_abi',
#    'overview', 'overview_raw', 'snow', 'snow_fog', 'so2', 'tropical_airmass', 'true_color', 'true_color_crefl',
#    'true_color_nocorr', 'true_color_raw', 'true_color_with_night_ir', 'true_color_with_night_ir_hires',
#    'water_vapors1', 'water_vapors2']

area = 'northamerica_10km'
composites = ['overview', 'C02']
generate = False

sat = 'GOESXX'
reader = 'abi_l1b'
outdir = './images'

# Use blend_type one of 'select_with_weights', 'blend_with_weights', 'stack_no_weights' or None
# blend_type = 'stack_no_weights'
# blend_type = 'select_with_weights'
blend_type = 'blend_with_weights'


filenames16 = glob('./slot16/*.nc')
filenames18 = glob('./slot18/*.nc')
global16 = Scene(filenames = filenames16, reader = reader)
global18 = Scene(filenames = filenames18, reader = reader)
mscn = MultiScene([global16, global18])
mscn.load(composites, generate=generate)

new_mscn = mscn.resample(area, resampler='nearest', radius_of_influence=20000, reduce_data=False)

if blend_type in ['select_with_weights', 'blend_with_weights']:
    weights=[]
    for scn in new_mscn.scenes:
        w = get_satellite_zenith_angle(scn[composites[0]]) / 90.
        w = (1. - w*w)
        weights.append(w * w)
    stack_with_weights = partial(stack, weights=weights, blend_type=blend_type)
    new_scn = new_mscn.blend(blend_function=stack_with_weights)
else:
    new_scn = new_mscn.blend()

if generate:
    for composite in composites:
        new_scn.save_dataset(composite, outdir+'/'+sat+'-'+composite+'-'+area+'-'+blend_type+'.png',
                             fill_value=0)
else:
    save_objects = []
    for composite in composites:
        save_objects.append(new_scn.save_dataset(composite, base_dir='.',
                            filename=outdir+'/'+sat+'-{name}-'+area+'-'+blend_type+'.png',
                            compute=False, fill_value=0))
    compute_writer_results(save_objects)