# -*- coding: utf-8 -*-

import math
import netCDF4
import numpy as np
import os

all_member_data = []

# *** ADD CODE TO READ IN ENSEMBLE MEMBER QPF FIELDS HERE ****
#
# Need to append the 2D QPF grid from each member to all_member_data
#
# Example code below (remove triple quotes around code to use/modify):
"""
member_names = ['hrrrv3.tl00', 'hrrrv3.tl01', 'hrwarw.tl00', 'hrwarw.tl01', 'hrwnmmb.tl00', 'hrwnmmb.tl01', 'hrwnssl.tl00', 'hrwnssl.tl01', 'namnest.tl00', 'namnest.tl01']

run = '2018060500'
qpf_hour_start = 12
qpf_hour_end = 24

dir_nc = '/www/ftp.nssl.noaa.gov/users/broberts/20200804_strauser_hrefqpf/'

for member in member_names:
    path_nc_start = os.path.join( dir_nc , run , member + '.f' + str(qpf_hour_start).zfill(3) + '.nc' )
    path_nc_end = os.path.join( dir_nc , run , member + '.f' + str(qpf_hour_end).zfill(3) + '.nc' )

    nc_start = netCDF4.Dataset( path_nc_start )
    nc_end = netCDF4.Dataset( path_nc_end )

    qpf_start = nc_start.variables['apcpsfc'][:]
    qpf_end = nc_end.variables['apcpsfc'][:]

    qpf_start[ qpf_start > 1e19 ] = 0.
    qpf_end[ qpf_end > 1e19 ] = 0.

    qpf_period = qpf_end - qpf_start

    all_member_data.append( qpf_period )
"""
# ************************************************************

# Get number of ensemble members
n_members = len(all_member_data)

# Create numpy array of all data
all_member_data = np.array(all_member_data)

# Mean of all members
mean_data = np.mean(all_member_data, axis=0)

# Shape of 2D grid
grid_shape = mean_data.shape

# Flatten and sort array of all member grid points
all_member_data = all_member_data.flatten()
all_member_data.sort()

# Flatten ensemble mean array
mean_data = mean_data.flatten()
# Indices of ensemble mean array, sorted by mean QPF value
sorted_mean_indices = np.argsort(mean_data)

# Match by percentile by choosing every N points from sorted full ensemble distribution
pmm_sorted = all_member_data[n_members-1::n_members]

# Construct PMM grid using ranked indices from ensemble mean grid
mean_data[sorted_mean_indices] = pmm_sorted[:]

# Final PMM array
pmm = np.array(mean_data, dtype="<f")
pmm = pmm.reshape( grid_shape )

# *** ADD CODE TO WRITE OUT PMM TO FILE HERE ****