#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# I need
from glob import glob

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

# This is stack from satpy v0.40 (Adams selecting pixels version only)
from satpy.multiscene import stack

# Below is my version from PR#2394 where I copied multiscene.py to wd
# import multiscene
# from multiscene import stack

import xarray as xr
import numpy as np

sat = 'MetopB'
reader = 'nwcsaf-pps_nc'
area = 'eurol'
generate = False

# Available NWCSAF composites in this script   (see .../satpy/etc/composites/visir.yaml):
# ***************************************************************************************
# All these pseudo/fake 'composites' are available for DAY and NIG passes:  ['cloud_top_height',
# 'cloud_top_pressure', 'cloud_top_temperature', 'cloudmask', 'cloudmask_extended', 'cloudtype']
composites = ['cloudmask', 'cloudtype']

outdir = './images'

# Set blend_type  either 'stack_no_weights' or 'select_with_weights'
blend_type = 'stack_no_weights'
#blend_type = 'select_with_weights'

def get_pixw(scn, dsn):
    j0, i0 = scn[dsn].shape
    lon, lat = scn[dsn].attrs['area'].get_lonlats()
    warr = np.zeros((i0))
    for i in range(i0):
        w = (2 * i + 1 - i0) / (i0 - 1)
        warr[i] = 1 - w * w
    warr = np.ones((j0, i0)) * warr
    pixw = xr.DataArray(data=warr, dims=["y", "x"],
                        coords=dict(
                            lon=(["y", "x"], lon),
                            lat=(["y", "x"], lat)),
                        attrs=dict(
                            description="Pixel weights from index.",
                            units="[1]",
    ),)
    return pixw


scenes = []
for passnum in ['1', '2', '3']:
    passfiles = glob('./pass'+passnum+'/*.nc')
    scenes.append(Scene(filenames=passfiles, reader=reader))
print(scenes[0].available_dataset_names())
mscn = MultiScene(scenes)
mscn.load(composites, generate=generate)

if  blend_type == 'select_with_weights':
    for scn in mscn.scenes:
        scn['pixw'] = get_pixw(scn, 'ct')
        scn['pixw'].attrs['area'] = scn['ct'].attrs['area']
    weights = True
else:
    weights = False

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

if weights:
    weights = []
    for scn in new_mscn.scenes:
        weights.append(scn['pixw'])
    stack_with_weights = partial(stack, weights=weights) # PR#2394: , 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)