#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Probability Maps netCDF generation
"""
#import pdb; pdb.set_trace()

import os
import sys
from datetime import datetime
import json
import netCDF4
import numpy as np

DEBUG = True

PARAMS = ('SCONC_DUST', 'OD550_DUST')

INPUT_DIR = "/data/products/median/interpolated/"

DIR_PATH = os.path.abspath(os.path.dirname(__file__))
MODELS = json.load(open(os.path.join(DIR_PATH, '../conf/models.json')))
DATES = json.load(open(os.path.join(DIR_PATH, '../conf/dates.json')))
VARS = json.load(open(os.path.join(DIR_PATH, '../conf/vars.json')))
DELAY = DATES['delay']['delayed']
DELAY_DATE = DATES['delay']['start_date']
TSTEPS = 8

if DELAY:
    FIRST_STEP = 0
else:
    FIRST_STEP = 1
FIRST_STEPS_LEN = 3

LON0 = 5
LON1 = -10

# ProbabilityMaps class plots probability maps for D+1 and D+2 daily
# maximum of a given parameter (AOD and "SCONC_DUST") from a interpolated
# netcdf file and needs: input_dir, output_dir, parameter_name, threshold
# (The Domain is defined previously in the netcdf file)
# run_prob method needs: conversion_factor

class ProbabilityMaps:
    """Class handling creation of netCDF files with ensemble probability values
    """

    def __init__(self, curdate, input_dir, output_dir, parameter_name, threshold):
        self._input_dir = input_dir
        self._output_dir = output_dir
        self._parameter_name = parameter_name
        self._threshold = threshold
        self.curdate = curdate
    @property
    def input_dir(self):
        return self._input_dir
    @property
    def output_dir(self):
        return self._output_dir
    @property
    def parameter_name(self):
        return self._parameter_name
    @property
    def threshold(self):
        return self._threshold

    def run_prob(self, conversion_factor):
        """ Executes probability calculations """
        # Getting parameter from nc files
        param_maxprob_lists = {}
        weight_sum = 0.
        run_time = ['']
        lon = []
        lat = []
        first_steps = []
        for model in MODELS:
            if model == 'median':
                continue
            tpl = MODELS[model]['template']
            ifile = os.path.join(self.input_dir, f'{self.curdate}{tpl}.nc')
            if DEBUG:
                print ("input_file:", ifile)
            if not os.path.exists(ifile):
                if DEBUG:
                    print ("NOT FOUND!")
                continue
            weight = float(MODELS[model]['weight'])
            weight_sum += weight
            nc_file = netCDF4.Dataset(ifile)
            if self.parameter_name not in nc_file.variables:
                try:
                    param_values = nc_file.variables[self.parameter_name.lower()][:,:,LON0:LON1]
                except Exception as err:
                    if DEBUG:
                        print("ERROR", str(err))
                    continue
            else:
                param_values = nc_file.variables[self.parameter_name][:,:,LON0:LON1]
            first_steps = np.arange(FIRST_STEP, param_values[FIRST_STEP:,:,:].shape[0], TSTEPS) # [:-1]
            if len(first_steps) < FIRST_STEPS_LEN:
                continue
            if DEBUG:
                print("FIRST_STEPS", first_steps)
            for step in first_steps:
                if step not in param_maxprob_lists:
                    param_maxprob_lists[step] = []
            lon = nc_file.variables["lon"][LON0:LON1]
            lat = nc_file.variables["lat"][:]
            run_time =  nc_file.variables['time'].units.split()[2:]
            # Calculates parameter daily maximum for day1, from 24 to 48
            # Calculates parameter daily mean for day1, from 24 to 48
            # (steps are 12 15 18 21...24 27 30 33 36 39 42 45...)
            # Slicing an array:  a[start:end+1]
            # parameterArray has [step,lat,lon]

            for cur_first_step in first_steps:
                total_steps = cur_first_step + TSTEPS
                param_max = param_values[cur_first_step,:,:] #first step 24 hours
                param_mean = param_values[cur_first_step,:,:] #first step 24 hours
                for param in param_values[cur_first_step+1:total_steps,:,:]:
                    param_max = np.maximum(param_max, param)
                    param_mean = param_mean + param
                param_mean = param_mean/len(param_values[cur_first_step:total_steps,:,:])
                # Change matrix values to true/false when exceding the given threshold
                # and when *1 to 1/0 this way matrixs can be added.
                param_mean_prob = (param_mean > self.threshold*conversion_factor)*1*weight
                param_maxprob_lists[cur_first_step].append(np.array(param_mean_prob))

            nc_file.close()

        if DEBUG:
            print ("Done reading nc files")
        # We add all the matrix for each model
        # and divide by the number of models to get probability
        if DEBUG:
            print ("weight_sum", weight_sum)
        if DEBUG:
            print ("run_time0", run_time[0])
        # and then plotting the maps
        ######param_maxprob_lists = [param_max24_problist, param_max48_problist]
        day0  = datetime.strptime(run_time[0], "%Y-%m-%d")

        os.makedirs(os.path.join(self._output_dir, self.curdate), exist_ok=True)
        prob_file = "{}/{}/{}{}{}_{}_Prob{}.nc".format(self._output_dir, self.curdate,
                str(day0.year), str(day0.month).zfill(2), str(day0.day).zfill(2),
                self.parameter_name, str(self.threshold))

        tmp_data = [] # np.ones((2, len(lat), len(lon)))*-999.

        if DEBUG:
            print('*****', prob_file)
        for param_maxprob_key in param_maxprob_lists:
            param_maxprob = param_maxprob_lists[param_maxprob_key]
            if DEBUG:
                print('::::', param_maxprob_key, len(param_maxprob))
            result1 = np.array(param_maxprob[0])
            datalist_len = np.array(param_maxprob).shape[0]
            for data_idx in range(1, datalist_len):
                result1 = result1 + param_maxprob[data_idx]
            #data1=100*result1/float(datalist_len)
            data1 = 100*result1/weight_sum

            # dataframe to NETCDF #############################################
            if DEBUG:
                print (len(data1), len(data1[0]))
            if DEBUG:
                print ("lon:", len(lon))
            if DEBUG:
                print ("lat:", len(lat))
            if DEBUG:
                print (self.parameter_name)

            tmp_data.append(data1)

        if os.path.exists(prob_file):
            os.remove(prob_file)

        outfile = netCDF4.Dataset(prob_file, 'w')

        outfile.createDimension('lon', len(lon))
        outfile.createDimension('lat', len(lat))
        outfile.createDimension('time', FIRST_STEPS_LEN)

        var_time = outfile.createVariable('time', 'i', ('time',))
        var_time.units = "days since {} 00:00:00".format(datetime.strftime(day0, "%Y-%m-%d"))
        var_time[:] = np.arange(FIRST_STEPS_LEN)

        var_lon = outfile.createVariable('lon', 'f', ('lon',))
        var_lon.units = "degrees_east"
        var_lon[:] = lon[:]

        var_lat = outfile.createVariable('lat', 'f', ('lat',))
        var_lat.units = "degrees_north"
        var_lat[:]= lat[:]

        var_data = outfile.createVariable(self.parameter_name, 'f', ('time', 'lat', 'lon'))
        print(":::", var_data.shape, np.array(tmp_data).shape)
        var_data[:] = np.array(tmp_data)[:]

        outfile.close()


if __name__ == "__main__":
    # makes daily directories where the maps will be stored #before 20200715: Sahel_prova...
    today = datetime.today().strftime('%Y%m%d')
    print (today)
    cdate = sys.argv[1]

    # Probability Maps for Regional
    for param_name in PARAMS:
        threshold_list = VARS[param_name]['threshold_list']
        conv_factor = VARS[param_name]['mul']**-1
        for thresh in threshold_list:
            print(param_name, thresh)
            out_dir = f"/data/daily_dashboard/prob/{param_name.lower()}/{thresh}/netcdf/"
            ProbabilityMaps(cdate, INPUT_DIR,
                    out_dir, param_name,
                    thresh).run_prob(conv_factor)
