import numpy as np
import pandas as pd
import pickle
import json
import os
from sklearn.decomposition import PCA
from copy import deepcopy
from glob import glob
from multiprocessing import Pool
from sklearn.linear_model import LinearRegression
from hagelslag.evaluation.ProbabilityMetrics import DistributedROC
from os.path import join
from hagelslag.util.make_proj_grids import read_arps_map_file, read_ncar_map_file

try:
    from sklearn.model_selection import KFold
except ImportError:
    from sklearn.cross_validation import KFold


class TrackModeler(object):
    """
    TrackModeler is designed to load and process data generated by TrackProcessing and then use that data to fit
    machine learning models to predict whether or not hail will occur, hail size, and translation errors in time and
    space.
    """

    def __init__(self,
                 ensemble_name,
                 train_data_path,
                 forecast_data_path,
                 member_files,
                 start_dates,
                 end_dates,
                 sector,
                 map_file,
                 group_col="Microphysics"):
        self.train_data_path = train_data_path
        self.forecast_data_path = forecast_data_path
        self.ensemble_name = ensemble_name
        self.member_files = member_files
        self.start_dates = start_dates
        self.end_dates = end_dates
        self.data = {"train": {}, "forecast": {}}
        self.condition_models = {}
        self.size_models = {}
        self.size_distribution_models = {}
        self.track_models = {"translation-x": {},
                             "translation-y": {},
                             "start-time": {}}
        self.sector = sector
        self.map_file = map_file
        self.group_col = group_col
            
        if self.map_file[-3:] == "map":                                  
            self.proj_dict, self.grid_dict = read_arps_map_file(self.map_file)
        else:                                               
            self.proj_dict, self.grid_dict = read_ncar_map_file(self.map_file)   
            
        
        return

    def load_data(self, mode="train", format="csv"):
        """
        Load data from flat data files containing total track information and information about each timestep.
        The two sets are combined using merge operations on the Track IDs. Additional member information is gathered
        from the appropriate member file.
        Args:
            mode: "train" or "forecast"
            format:  file format being used. Default is "csv"
        """
        if mode in self.data.keys():
            run_dates = pd.DatetimeIndex(start=self.start_dates[mode],
                                        end=self.end_dates[mode],freq="1D")
            run_date_str = [d.strftime("%Y%m%d-%H%M") for d in run_dates.date]
            print(run_date_str)
            all_total_track_files = sorted(glob(getattr(self, mode + "_data_path") +
                                                "*total_" + self.ensemble_name + "*." + format))
            all_step_track_files = sorted(glob(getattr(self, mode + "_data_path") +
                                               "*step_" + self.ensemble_name + "*." + format))
            total_track_files = []
            for track_file in all_total_track_files:
                file_date = track_file.split("_")[-1][:-4]
                if file_date in run_date_str:
                    total_track_files.append(track_file)
            step_track_files = []
            for step_file in all_step_track_files:
                file_date = step_file.split("_")[-1][:-4]
                if file_date in run_date_str:
                    step_track_files.append(step_file)            
	  	   
            self.data[mode]["total"] = pd.concat(map(pd.read_csv, total_track_files),
                                                 ignore_index=True)
            self.data[mode]["total"] = self.data[mode]["total"].fillna(value=0)
            self.data[mode]["total"] = self.data[mode]["total"].replace([np.inf, -np.inf], 0)
            self.data[mode]["step"] = pd.concat(map(pd.read_csv, step_track_files),
                                                ignore_index=True)
            self.data[mode]["step"] = self.data[mode]["step"].fillna(value=0)
            self.data[mode]["step"] = self.data[mode]["step"].replace([np.inf, -np.inf], 0)
            if mode == "forecast":
                self.data[mode]["step"] = self.data[mode]["step"].drop_duplicates("Step_ID")
            self.data[mode]["member"] = pd.read_csv(self.member_files[mode])
            self.data[mode]["combo"] = pd.merge(self.data[mode]["step"],
                                                self.data[mode]["total"],
                                                on=["Track_ID", "Ensemble_Name", "Ensemble_Member", "Run_Date"])
            self.data[mode]["combo"] = pd.merge(self.data[mode]["combo"],
                                                self.data[mode]["member"],
                                                on="Ensemble_Member") 
            self.data[mode]["total_group"] = pd.merge(self.data[mode]["total"],
                                                      self.data[mode]["member"],
                                                      on="Ensemble_Member")

    def calc_copulas(self,
                     output_file,
                     model_names=("start-time", "translation-x", "translation-y"),
                     label_columns=("Start_Time_Error", "Translation_Error_X", "Translation_Error_Y")):
        """
        Calculate a copula multivariate normal distribution from the training data for each group of ensemble members.
        Distributions are written to a pickle file for later use.
        Args:
            output_file: Pickle file
            model_names: Names of the tracking models
            label_columns: Names of the data columns used for labeling
        Returns:
        """
        if len(self.data['train']) == 0:
            self.load_data()
        groups = self.data["train"]["member"][self.group_col].unique()
        copulas = {}
        label_columns = list(label_columns)
        for group in groups:
            print(group)
            group_data = self.data["train"]["total_group"].loc[
                self.data["train"]["total_group"][self.group_col] == group]
            group_data = group_data.dropna()
            group_data.reset_index(drop=True, inplace=True)
            copulas[group] = {}
            copulas[group]["mean"] = group_data[label_columns].mean(axis=0).values
            copulas[group]["cov"] = np.cov(group_data[label_columns].values.T)
            copulas[group]["model_names"] = list(model_names)
            del group_data
        pickle.dump(copulas, open(output_file, "w"), pickle.HIGHEST_PROTOCOL)

    def fit_condition_models(self, model_names,
                             model_objs,
                             input_columns,
                             output_column="Matched",
                             output_threshold=0.0):
        """
        Fit machine learning models to predict whether or not hail will occur.
        Args:
            model_names: List of strings with the names for the particular machine learning models
            model_objs: scikit-learn style machine learning model objects.
            input_columns: list of the names of the columns used as input for the machine learning model
            output_column: name of the column used for labeling whether or not the event occurs
            output_threshold: splitting threshold to determine if event has occurred. Default 0.0
        """
        print("Fitting condition models")
        groups = self.data["train"]["member"][self.group_col].unique()
        
        weights = None

        for group in groups:
            print(group)
            group_data = self.data["train"]["combo"].loc[self.data["train"]["combo"][self.group_col] == group] 
            if self.sector:
        
                lon_obj = data.loc[:,'Centroid_Lon']
                lat_obj = data.loc[:,'Centroid_Lat']

                left_lon,right_lon = self.grid_dict["sw_lon"],self.grid_dict["ne_lon"]
                lower_lat,upper_lat = self.grid_dict["sw_lat"],self.grid_dict["ne_lat"]

                weights = np.where((left_lon<=lon_obj)&(right_lon>=lon_obj) &\
                    (lower_lat<=lat_obj)&(upper_lat>=lat_obj),1,0.3)
            
            output_data = np.where(group_data[output_column] > output_threshold, 1, 0)
            print("Ones: ", np.count_nonzero(output_data > 0), "Zeros: ", np.count_nonzero(output_data == 0))
            self.condition_models[group] = {}
            for m, model_name in enumerate(model_names):
                print(model_name)
                self.condition_models[group][model_name] = deepcopy(model_objs[m])
                try:
                    self.condition_models[group][model_name].fit(group_data[input_columns], 
                                output_data,sample_weight=weights)
                except:
                    self.condition_models[group][model_name].fit(group_data[input_columns], output_data)

                if hasattr(self.condition_models[group][model_name], "best_estimator_"):
                    print(self.condition_models[group][model_name].best_estimator_)
                    print(self.condition_models[group][model_name].best_score_)

    def fit_condition_threshold_models(self, model_names, model_objs, input_columns, output_column="Matched",
                                       output_threshold=0.5, num_folds=5, threshold_score="ets"):
        """
        Fit models to predict hail/no-hail and use cross-validation to determine the probaility threshold that
        maximizes a skill score.

        Args:
            model_names: List of machine learning model names
            model_objs: List of Scikit-learn ML models
            input_columns: List of input variables in the training data
            output_column: Column used for prediction
            output_threshold: Values exceeding this threshold are considered positive events; below are nulls
            num_folds: Number of folds in the cross-validation procedure
            threshold_score: Score available in ContingencyTable used for determining the best probability threshold
        Returns:
            None
        """
        print("Fitting condition models")
        groups = self.data["train"]["member"][self.group_col].unique()
        
        weights=None

        for group in groups:
            print(group)
            group_data = self.data["train"]["combo"].iloc[
                np.where(self.data["train"]["combo"][self.group_col] == group)[0]]
            
            if self.sector:
                lon_obj = group_data.loc[:,'Centroid_Lon']
                lat_obj = group_data.loc[:,'Centroid_Lat']
                
                conus_lat_lon_points = zip(lon_obj.values.ravel(),lat_obj.values.ravel())
                center_lon, center_lat = self.proj_dict["lon_0"],self.proj_dict["lat_0"] 
            
                distances = np.array([np.sqrt((x-center_lon)**2+\
                        (y-center_lat)**2) for (x, y) in conus_lat_lon_points])
            
                min_dist, max_minus_min = min(distances),max(distances)-min(distances)

                distance_0_1 = [1.0-((d - min_dist)/(max_minus_min)) for d in distances]
                weights = np.array(distance_0_1)
        
            output_data = np.where(group_data.loc[:, output_column] > output_threshold, 1, 0)
            ones = np.count_nonzero(output_data > 0)
            print("Ones: ", ones, "Zeros: ", np.count_nonzero(output_data == 0))
            self.condition_models[group] = {}
            num_elements = group_data[input_columns].shape[0]
            
            for m, model_name in enumerate(model_names):
                print(model_name)    
                roc = DistributedROC(thresholds=np.arange(0, 1.1, 0.01))
                self.condition_models[group][model_name] = deepcopy(model_objs[m])

                try:
                    kf = KFold(n_splits=num_folds)
                    for train_index, test_index in kf.split(group_data[input_columns].values):
                        if np.count_nonzero(output_data[train_index]) > 0:
                            try:
                                self.condition_models[group][model_name].fit(
                                        group_data.iloc[train_index][input_columns],
                                        output_data[train_index],sample_weight=weights[train_index])
                            except:
                                self.condition_models[group][model_name].fit(
                                        group_data.iloc[train_index][input_columns],
                                        output_data[train_index])
                            
                            cv_preds = self.condition_models[group][model_name].predict_proba(
                                group_data.iloc[test_index][input_columns])[:,1]
                            
                            roc.update(cv_preds, output_data[test_index])
                        
                        else:
                            continue

                except TypeError:
                    kf = KFold(num_elements,n_folds=num_folds)
                    for train_index, test_index in kf:

                        if np.count_nonzero(output_data[train_index]) > 0:
                            try:
                                self.condition_models[group][model_name].fit(
                                        group_data.iloc[train_index][input_columns],
                                        output_data[train_index],sample_weight=weights[train_index])
                            except:
                                self.condition_models[group][model_name].fit(
                                        group_data.iloc[train_index][input_columns],
                                        output_data[train_index])
                            cv_preds = self.condition_models[group][model_name].predict_proba(
                                group_data.iloc[test_index][input_columns])[:, 1]
                            
                            roc.update(cv_preds, output_data[test_index])
                        
                        else:
                            continue

                self.condition_models[group][
                    model_name + "_condition_threshold"], _ = roc.max_threshold_score(threshold_score)
                print(model_name + " condition threshold: {0:0.3f}".format(
                self.condition_models[group][model_name + "_condition_threshold"]))
                self.condition_models[group][model_name].fit(group_data[input_columns], output_data)

    def predict_condition_models(self, model_names,
                                 input_columns,
                                 metadata_cols,
                                 data_mode="forecast",
                                 ):
        """
        Apply condition modelsto forecast data.
        Args:
            model_names: List of names associated with each condition model used for prediction
            input_columns: List of columns in data used as input into the model
            metadata_cols: Columns from input data that should be included in the data frame with the predictions.
            data_mode: Which data subset to pull from for the predictions, "forecast" by default
        Returns:
            A dictionary of data frames containing probabilities of the event and specified metadata
        """
        groups = self.condition_models.keys()
        predictions = pd.DataFrame(self.data[data_mode]["combo"][metadata_cols])
        for group in groups:
            print(group)
            print(self.condition_models[group])
            g_idxs = self.data[data_mode]["combo"][self.group_col] == group
            group_count = np.count_nonzero(g_idxs)
            if group_count > 0:
                for m, model_name in enumerate(model_names):
                    mn = model_name.replace(" ", "-")
                    predictions.loc[g_idxs, mn + "_conditionprob"] = self.condition_models[group][
                                                                         model_name].predict_proba(
                        self.data[data_mode]["combo"].loc[g_idxs, input_columns])[:, 1]
                    predictions.loc[g_idxs,
                                    mn + "_conditionthresh"] = np.where(predictions.loc[g_idxs, mn + "_conditionprob"]
                                                                        >= self.condition_models[group][
                                                                            model_name + "_condition_threshold"], 1, 0)

        return predictions

    def fit_size_distribution_models(self, model_names, model_objs, input_columns,
                                     output_columns=None, calibrate=False):
        """
        Fits multitask machine learning models to predict the parameters of a size distribution
        Args:
            model_names: List of machine learning model names
            model_objs: scikit-learn style machine learning model objects
            input_columns: Training data columns used as input for ML model
            output_columns: Training data columns used for prediction
            calibrate: Whether or not to fit a log-linear regression to predictions from ML model
        """
        if output_columns is None:
            output_columns = ["Shape", "Location", "Scale"]
        groups = np.unique(self.data["train"]["member"][self.group_col])
        
        weights=None
        
        for group in groups:
            group_data = self.data["train"]["combo"].loc[self.data["train"]["combo"][self.group_col] == group]
            group_data = group_data.dropna()
            group_data = group_data[group_data[output_columns[-1]] > 0]
            if self.sector:
        
                lon_obj = group_data.loc[:,'Centroid_Lon']
                lat_obj = group_data.loc[:,'Centroid_Lat']
                
                conus_lat_lon_points = zip(lon_obj.values.ravel(),lat_obj.values.ravel())
                center_lon, center_lat = self.proj_dict["lon_0"],self.proj_dict["lat_0"] 
            
                distances = np.array([np.sqrt((x-center_lon)**2+\
                        (y-center_lat)**2) for (x, y) in conus_lat_lon_points])
            
                min_dist, max_minus_min = min(distances),max(distances)-min(distances)

                distance_0_1 = [1.0-((d - min_dist)/(max_minus_min)) for d in distances]
                weights = np.array(distance_0_1)

            self.size_distribution_models[group] = {"multi": {}, "lognorm": {}}
            if calibrate:
                self.size_distribution_models[group]["calshape"] = {}
                self.size_distribution_models[group]["calscale"] = {}
            log_labels = np.log(group_data[output_columns].values)
            log_means = log_labels.mean(axis=0)
            log_sds = log_labels.std(axis=0)
            self.size_distribution_models[group]['lognorm']['mean'] = log_means
            self.size_distribution_models[group]['lognorm']['sd'] = log_sds
            for m, model_name in enumerate(model_names):
                print(group, model_name)
                self.size_distribution_models[group]["multi"][model_name] = deepcopy(model_objs[m])
                try:
                    self.size_distribution_models[group]["multi"][model_name].fit(group_data[input_columns],
                                                                              (log_labels - log_means) / log_sds,
                                                                        sample_weight=weights)
                except:
                    self.size_distribution_models[group]["multi"][model_name].fit(group_data[input_columns],
                                                                              (log_labels - log_means) / log_sds)
                if calibrate:
                    training_predictions = self.size_distribution_models[
                        group]["multi"][model_name].predict(group_data[input_columns])
                    self.size_distribution_models[group]["calshape"][model_name] = LinearRegression()
                    self.size_distribution_models[group]["calshape"][model_name].fit(training_predictions[:, 0:1],
                                                                                     (log_labels[:, 0] - log_means[0]) /
                                                                                     log_sds[
                                                                                         0],
                                                                                    sample_weight=weights)
                    self.size_distribution_models[group]["calscale"][model_name] = LinearRegression()
                    self.size_distribution_models[group]["calscale"][model_name].fit(training_predictions[:, 1:],
                                                                                     (log_labels[:, 1] - log_means[1]) /
                                                                                     log_sds[
                                                                                         1],
                                                                            sample_weight=weights)


    def fit_size_distribution_component_models(self, model_names, model_objs, input_columns, output_columns):
        """
        This calculates 2 principal components for the hail size distribution between the shape and scale parameters.
        Separate machine learning models are fit to predict each component.

        Args:
            model_names: List of machine learning model names
            model_objs: List of machine learning model objects.
            input_columns: List of input variables
            output_columns: Output columns, should contain Shape and Scale.

        Returns:

        """
        groups = np.unique(self.data["train"]["member"][self.group_col])
        
        weights=None

        for group in groups:
            print(group)
            group_data = self.data["train"]["combo"].loc[self.data["train"]["combo"][self.group_col] == group]
            group_data = group_data.dropna()
            group_data = group_data.loc[group_data[output_columns[-1]] > 0]
            
            if self.sector:
        
                lon_obj = group_data.loc[:,'Centroid_Lon']
                lat_obj = group_data.loc[:,'Centroid_Lat']
                
                conus_lat_lon_points = zip(lon_obj.values.ravel(),lat_obj.values.ravel())
                center_lon, center_lat = self.proj_dict["lon_0"],self.proj_dict["lat_0"] 
            
                distances = np.array([np.sqrt((x-center_lon)**2+\
                        (y-center_lat)**2) for (x, y) in conus_lat_lon_points])
            
                min_dist, max_minus_min = min(distances),max(distances)-min(distances)

                distance_0_1 = [1.0-((d - min_dist)/(max_minus_min)) for d in distances]
                weights = np.array(distance_0_1)

            
            self.size_distribution_models[group] = {"lognorm": {}}
            self.size_distribution_models[group]["lognorm"]["pca"] = PCA(n_components=len(output_columns))
            log_labels = np.log(group_data[output_columns].values)
            log_labels[:, np.where(output_columns == "Shape")[0]] *= -1
            log_means = log_labels.mean(axis=0)
            log_sds = log_labels.std(axis=0)
            log_norm_labels = (log_labels - log_means) / log_sds
            out_pc_labels = self.size_distribution_models[group]["lognorm"]["pca"].fit_transform(log_norm_labels)
            self.size_distribution_models[group]['lognorm']['mean'] = log_means
            self.size_distribution_models[group]['lognorm']['sd'] = log_sds
            for comp in range(len(output_columns)):
                self.size_distribution_models[group]["pc_{0:d}".format(comp)] = dict()
                for m, model_name in enumerate(model_names):
                    print(model_name, comp)
                    self.size_distribution_models[group][
                        "pc_{0:d}".format(comp)][model_name] = deepcopy(model_objs[m])
                    try:
                        self.size_distribution_models[group][
                            "pc_{0:d}".format(comp)][model_name].fit(group_data[input_columns],
                                                                 out_pc_labels[:, comp],
                                                            sample_weight=weights)
                    except:
                        self.size_distribution_models[group][
                            "pc_{0:d}".format(comp)][model_name].fit(group_data[input_columns],
                                                                 out_pc_labels[:, comp])
        return

    def predict_size_distribution_models(self, model_names, input_columns, metadata_cols,
                                         data_mode="forecast", location=6, calibrate=False):
        """
        Make predictions using fitted size distribution models.
        Args:
            model_names: Name of the models for predictions
            input_columns: Data columns used for input into ML models
            metadata_cols: Columns from input data that should be included in the data frame with the predictions.
            data_mode: Set of data used as input for prediction models
            location: Value of fixed location parameter
            calibrate: Whether or not to apply calibration model
        Returns:
            Predictions in dictionary of data frames grouped by group type
        """
        groups = self.size_distribution_models.keys()
        predictions = {}
        for group in groups:
            group_data = self.data[data_mode]["combo"].loc[self.data[data_mode]["combo"][self.group_col] == group]
            predictions[group] = group_data[metadata_cols]
            if group_data.shape[0] > 0:
                log_mean = self.size_distribution_models[group]["lognorm"]["mean"]
                log_sd = self.size_distribution_models[group]["lognorm"]["sd"]
                for m, model_name in enumerate(model_names):
                    multi_predictions = self.size_distribution_models[group]["multi"][model_name].predict(
                        group_data[input_columns])
                    if calibrate:
                        multi_predictions[:, 0] = self.size_distribution_models[group]["calshape"][model_name].predict(
                            multi_predictions[:, 0:1])
                        multi_predictions[:, 1] = self.size_distribution_models[group]["calscale"][model_name].predict(
                            multi_predictions[:, 1:])
                    multi_predictions = np.exp(multi_predictions * log_sd + log_mean)
                    if multi_predictions.shape[1] == 2:
                        multi_predictions_temp = np.zeros((multi_predictions.shape[0], 3))
                        multi_predictions_temp[:, 0] = multi_predictions[:, 0]
                        multi_predictions_temp[:, 1] = location
                        multi_predictions_temp[:, 2] = multi_predictions[:, 1]
                        multi_predictions = multi_predictions_temp
                    for p, pred_col in enumerate(["shape", "location", "scale"]):
                        predictions[group][model_name].loc[:, model_name.replace(" ", "-") + "_" + pred_col] = \
                            multi_predictions[:, p]
        return predictions

    def predict_size_distribution_component_models(self, model_names, input_columns, output_columns, metadata_cols,
                                                   data_mode="forecast", location=6):
        """
        Make predictions using fitted size distribution models.
        Args:
            model_names: Name of the models for predictions
            input_columns: Data columns used for input into ML models
            output_columns: Names of output columns
            metadata_cols: Columns from input data that should be included in the data frame with the predictions.
            data_mode: Set of data used as input for prediction models
            location: Value of fixed location parameter
        Returns:
            Predictions in dictionary of data frames grouped by group type
        """
        groups = self.size_distribution_models.keys()
        predictions = pd.DataFrame(self.data[data_mode]["combo"][metadata_cols])
        for group in groups:
            group_idxs = self.data[data_mode]["combo"][self.group_col] == group
            group_count = np.count_nonzero(group_idxs)
            print(self.size_distribution_models[group])
            if group_count > 0:
                log_mean = self.size_distribution_models[group]["lognorm"]["mean"]
                log_sd = self.size_distribution_models[group]["lognorm"]["sd"]
                for m, model_name in enumerate(model_names):
                    raw_preds = np.zeros((group_count, len(output_columns)))
                    for c in range(len(output_columns)):
                        raw_preds[:, c] = self.size_distribution_models[group][
                            "pc_{0:d}".format(c)][model_name].predict(self.data[data_mode]["combo"].loc[group_idxs,
                                                                                                        input_columns])
                    log_norm_preds = self.size_distribution_models[group]["lognorm"]["pca"].inverse_transform(raw_preds)
                    log_norm_preds[:, 0] *= -1
                    multi_predictions = np.exp(log_norm_preds * log_sd + log_mean)
                    if multi_predictions.shape[1] == 2:
                        multi_predictions_temp = np.zeros((multi_predictions.shape[0], 3))
                        multi_predictions_temp[:, 0] = multi_predictions[:, 0]
                        multi_predictions_temp[:, 1] = location
                        multi_predictions_temp[:, 2] = multi_predictions[:, 1]
                        multi_predictions = multi_predictions_temp
                    for p, pred_col in enumerate(["shape", "location", "scale"]):
                        predictions.loc[group_idxs, model_name.replace(" ", "-") + "_" + pred_col] = \
                            multi_predictions[:, p]
        return predictions

    def fit_size_models(self, model_names,
                        model_objs,
                        input_columns,
                        output_column="Hail_Size",
                        output_start=5,
                        output_step=5,
                        output_stop=100):
        """
        Fit size models to produce discrete pdfs of forecast hail sizes.
        Args:
            model_names: List of model names
            model_objs: List of model objects
            input_columns: List of input variables
            output_column: Output variable name
            output_start: Hail size bin start
            output_step: hail size bin step
            output_stop: hail size bin stop
        """
        print("Fitting size models")
        groups = self.data["train"]["member"][self.group_col].unique()
        output_start = int(output_start)
        output_step = int(output_step)
        output_stop = int(output_stop)
        for group in groups:
            group_data = self.data["train"]["combo"].loc[self.data["train"]["combo"][self.group_col] == group]
            group_data.dropna(inplace=True)
            group_data = group_data[group_data[output_column] >= output_start]
            output_data = group_data[output_column].values.astype(int)
            output_data[output_data > output_stop] = output_stop
            discrete_data = ((output_data - output_start) // output_step) * output_step + output_start
            self.size_models[group] = {}
            self.size_models[group]["outputvalues"] = np.arange(output_start, output_stop + output_step, output_step,
                                                                dtype=int)
            for m, model_name in enumerate(model_names):
                print("{0} {1}".format(group, model_name))
                self.size_models[group][model_name] = deepcopy(model_objs[m])
                self.size_models[group][model_name].fit(group_data[input_columns], discrete_data)

    def predict_size_models(self, model_names,
                            input_columns,
                            metadata_cols,
                            data_mode="forecast"):
        """
        Apply size models to forecast data.
        Args:
            model_names:
            input_columns:
            metadata_cols:
            data_mode:
        """
        groups = self.size_models.keys()
        predictions = {}
        for group in groups:
            group_data = self.data[data_mode]["combo"].loc[self.data[data_mode]["combo"][self.group_col] == group]
            if group_data.shape[0] > 0:
                predictions[group] = {}
                output_values = self.size_models[group]["outputvalues"].astype(int)
                for m, model_name in enumerate(model_names):
                    print("{0} {1}".format(group, model_name))
                    pred_col_names = [model_name.replace(" ", "-") + "_{0:02d}".format(p) for p in output_values]
                    predictions[group][model_name] = group_data[metadata_cols]
                    pred_vals = self.size_models[group][model_name].predict_proba(group_data[input_columns])
                    pred_classes = self.size_models[group][model_name].classes_
                    pred_pdf = np.zeros((pred_vals.shape[0], output_values.size))
                    for pcv, pc in enumerate(pred_classes):
                        idx = np.where(output_values == pc)[0][0]
                        pred_pdf[:, idx] = pred_vals[:, pcv]
                    for pcn, pred_col_name in enumerate(pred_col_names):
                        predictions[group][model_name].loc[:, pred_col_name] = pred_pdf[:, pcn]
        return predictions

    def fit_track_models(self,
                         model_names,
                         model_objs,
                         input_columns,
                         output_columns,
                         output_ranges,
                         ):
        """
        Fit machine learning models to predict track error offsets.
            model_names:
            model_objs:
            input_columns:
            output_columns:
            output_ranges:
        """
        print("Fitting track models")
        groups = self.data["train"]["member"][self.group_col].unique()
        for group in groups:
            group_data = self.data["train"]["combo"].loc[self.data["train"]["combo"][self.group_col] == group]
            group_data = group_data.dropna()
            group_data = group_data.loc[group_data["Duration_Step"] == 1]
            for model_type, model_dict in self.track_models.items():
                model_dict[group] = {}
                output_data = group_data[output_columns[model_type]].values.astype(int)
                output_data[output_data < output_ranges[model_type][0]] = output_ranges[model_type][0]
                output_data[output_data > output_ranges[model_type][1]] = output_ranges[model_type][1]
                discrete_data = (output_data - output_ranges[model_type][0]) // output_ranges[model_type][2] * \
                                output_ranges[model_type][2] + output_ranges[model_type][0]
                model_dict[group]["outputvalues"] = np.arange(output_ranges[model_type][0],
                                                              output_ranges[model_type][1] +
                                                              output_ranges[model_type][2],
                                                              output_ranges[model_type][2])
                for m, model_name in enumerate(model_names):
                    print("{0} {1} {2}".format(group, model_type, model_name))
                    model_dict[group][model_name] = deepcopy(model_objs[m])
                    model_dict[group][model_name].fit(group_data[input_columns], discrete_data)

    def predict_track_models(self, model_names,
                             input_columns,
                             metadata_cols,
                             data_mode="forecast",
                             ):
        """
        Predict track offsets on forecast data.
        Args:
            model_names:
            input_columns:
            metadata_cols:
            data_mode:
        Returns:
        """
        predictions = {}
        for model_type, track_model_set in self.track_models.items():
            predictions[model_type] = {}
            groups = track_model_set.keys()
            for group in groups:
                group_data = self.data[data_mode]["combo"].loc[self.data[data_mode]["combo"][self.group_col] == group]
                if group_data.shape[0] > 0:
                    predictions[model_type][group] = {}
                    output_values = track_model_set[group]["outputvalues"].astype(int)
                    for m, model_name in enumerate(model_names):
                        print("{0} {1} {2}".format(group, model_type, model_name))
                        pred_col_names = [model_name.replace(" ", "-") + "_{0:02d}".format(p) for p in output_values]
                        predictions[model_type][group][model_name] = group_data[metadata_cols]
                        pred_vals = track_model_set[group][model_name].predict_proba(group_data[input_columns])
                        pred_classes = track_model_set[group][model_name].classes_
                        pred_pdf = np.zeros((pred_vals.shape[0], output_values.size))
                        for pcv, pc in enumerate(pred_classes):
                            idx = np.where(pc == output_values)[0][0]
                            pred_pdf[:, idx] = pred_vals[:, pcv]
                        for pcn, pred_col_name in enumerate(pred_col_names):
                            predictions[model_type][group][model_name].loc[:, pred_col_name] = pred_pdf[:, pcn]
        return predictions

    def save_models(self, model_path):
        """
        Save machine learning models to pickle files.
        """
        for group, condition_model_set in self.condition_models.items():
            for model_name, model_obj in condition_model_set.items():
                out_filename = model_path + \
                               "{0}_{1}_condition.pkl".format(group,
                                                              model_name.replace(" ", "-"))
                with open(out_filename, "wb") as pickle_file:
                    pickle.dump(model_obj,
                                pickle_file,
                                pickle.HIGHEST_PROTOCOL)
        for group, size_model_set in self.size_models.items():
            for model_name, model_obj in size_model_set.items():
                out_filename = model_path + \
                               "{0}_{1}_size.pkl".format(group,
                                                         model_name.replace(" ", "-"))
                with open(out_filename, "wb") as pickle_file:
                    pickle.dump(model_obj,
                                pickle_file,
                                pickle.HIGHEST_PROTOCOL)
        for group, dist_model_set in self.size_distribution_models.items():
            for model_type, model_objs in dist_model_set.items():
                for model_name, model_obj in model_objs.items():
                    out_filename = model_path + \
                                   "{0}_{1}_{2}_sizedist.pkl".format(group,
                                                                     model_name.replace(" ", "-"),
                                                                     model_type)
                    with open(out_filename, "wb") as pickle_file:
                        pickle.dump(model_obj,
                                    pickle_file,
                                    pickle.HIGHEST_PROTOCOL)
        for model_type, track_type_models in self.track_models.items():
            for group, track_model_set in track_type_models.items():
                for model_name, model_obj in track_model_set.items():
                    out_filename = model_path + \
                                   "{0}_{1}_{2}_track.pkl".format(group,
                                                                  model_name.replace(" ", "-"),
                                                                  model_type)
                    with open(out_filename, "wb") as pickle_file:
                        pickle.dump(model_obj,
                                    pickle_file,
                                    pickle.HIGHEST_PROTOCOL)

        return

    def load_models(self, model_path):
        """
        Load models from pickle files.
        """
        condition_model_files = sorted(glob(model_path + "*_condition.pkl"))
        if len(condition_model_files) > 0:
            for condition_model_file in condition_model_files:
                model_comps = condition_model_file.split("/")[-1][:-4].split("_")
                if model_comps[0] not in self.condition_models.keys():
                    self.condition_models[model_comps[0]] = {}
                model_name = model_comps[1].replace("-", " ")
                with open(condition_model_file, "rb") as cmf:
                    if "condition_threshold" in condition_model_file:
                        self.condition_models[model_comps[0]][model_name + "_condition_threshold"] = pickle.load(cmf)
                    else:
                        self.condition_models[model_comps[0]][model_name] = pickle.load(cmf)

        size_model_files = sorted(glob(model_path + "*_size.pkl"))
        if len(size_model_files) > 0:
            for size_model_file in size_model_files:
                model_comps = size_model_file.split("/")[-1][:-4].split("_")
                if model_comps[0] not in self.size_models.keys():
                    self.size_models[model_comps[0]] = {}
                model_name = model_comps[1].replace("-", " ")
                with open(size_model_file, "rb") as smf:
                    self.size_models[model_comps[0]][model_name] = pickle.load(smf)

        size_dist_model_files = sorted(glob(model_path + "*_sizedist.pkl"))
        if len(size_dist_model_files) > 0:
            for dist_model_file in size_dist_model_files:
                model_comps = dist_model_file.split("/")[-1][:-4].split("_")
                if model_comps[0] not in self.size_distribution_models.keys():
                    self.size_distribution_models[model_comps[0]] = {}
                if "_".join(model_comps[2:-1]) not in self.size_distribution_models[model_comps[0]].keys():
                    self.size_distribution_models[model_comps[0]]["_".join(model_comps[2:-1])] = {}
                model_name = model_comps[1].replace("-", " ")
                with open(dist_model_file, "rb") as dmf:
                    self.size_distribution_models[model_comps[0]]["_".join(model_comps[2:-1])][
                        model_name] = pickle.load(dmf)

        track_model_files = sorted(glob(model_path + "*_track.pkl"))
        if len(track_model_files) > 0:
            for track_model_file in track_model_files:
                model_comps = track_model_file.split("/")[-1][:-4].split("_")
                group = model_comps[0]
                model_name = model_comps[1].replace("-", " ")
                model_type = model_comps[2]
                if model_type not in self.track_models.keys():
                    self.track_models[model_type] = {}
                if group not in self.track_models[model_type].keys():
                    self.track_models[model_type][group] = {}
                with open(track_model_file, "rb") as tmf:
                    self.track_models[model_type][group][model_name] = pickle.load(tmf)

    def output_forecasts_json(self, forecasts,
                              condition_model_names,
                              size_model_names,
                              dist_model_names,
                              track_model_names,
                              json_data_path,
                              out_path):
        """
        Output forecast values to geoJSON file format.
        :param forecasts:
        :param condition_model_names:
        :param size_model_names:
        :param track_model_names:
        :param json_data_path:
        :param out_path:
        :return:
        """
        total_tracks = self.data["forecast"]["total"]
        for r in np.arange(total_tracks.shape[0]):
            track_id = total_tracks.loc[r, "Track_ID"]
            print(track_id)
            track_num = track_id.split("_")[-1]
            ensemble_name = total_tracks.loc[r, "Ensemble_Name"]
            member = total_tracks.loc[r, "Ensemble_Member"]
            group = self.data["forecast"]["member"].loc[self.data["forecast"]["member"]["Ensemble_Member"] == member,
                                                        self.group_col].values[0]
            run_date = track_id.split("_")[-4][:8]
            step_forecasts = {}
            for ml_model in condition_model_names:
                step_forecasts["condition_" + ml_model.replace(" ", "-")] = forecasts["condition"][group].loc[
                    forecasts["condition"][group]["Track_ID"] == track_id, ml_model]
            for ml_model in size_model_names:
                step_forecasts["size_" + ml_model.replace(" ", "-")] = forecasts["size"][group][ml_model].loc[
                    forecasts["size"][group][ml_model]["Track_ID"] == track_id]
            for ml_model in dist_model_names:
                step_forecasts["dist_" + ml_model.replace(" ", "-")] = forecasts["dist"][group][ml_model].loc[
                    forecasts["dist"][group][ml_model]["Track_ID"] == track_id]
            for model_type in forecasts["track"].keys():
                for ml_model in track_model_names:
                    mframe = forecasts["track"][model_type][group][ml_model]
                    step_forecasts[model_type + "_" + ml_model.replace(" ", "-")] = mframe.loc[
                        mframe["Track_ID"] == track_id]
            json_file_name = "{0}_{1}_{2}_model_track_{3}.json".format(ensemble_name,
                                                                       run_date,
                                                                       member,
                                                                       track_num)
            full_json_path = json_data_path + "/".join([run_date, member]) + "/" + json_file_name
            with open(full_json_path) as json_file_obj:
                try:
                    track_obj = json.load(json_file_obj)
                except FileNotFoundError:
                    print(full_json_path + " not found")
                    continue
            for f, feature in enumerate(track_obj['features']):
                del feature['properties']['attributes']
                for model_name, fdata in step_forecasts.items():
                    ml_model_name = model_name.split("_")[1]
                    if "condition" in model_name:
                        feature['properties'][model_name] = fdata.values[f]
                    else:
                        predcols = []
                        for col in fdata.columns:
                            if ml_model_name in col:
                                predcols.append(col)
                        feature['properties'][model_name] = fdata.loc[:, predcols].values[f].tolist()
            full_path = []
            for part in [run_date, member]:
                full_path.append(part)
                if not os.access(out_path + "/".join(full_path), os.R_OK):
                    try:
                        os.mkdir(out_path + "/".join(full_path))
                    except OSError:
                        print("directory already created")
            out_json_filename = out_path + "/".join(full_path) + "/" + json_file_name
            with open(out_json_filename, "w") as out_json_obj:
                json.dump(track_obj, out_json_obj, indent=1, sort_keys=True)
        return

    def output_forecasts_csv(self, forecasts, mode, csv_path, run_date_format="%Y%m%d-%H%M"):
        """
        Output hail forecast values to csv files by run date and ensemble member.

        Args:
            forecasts:
            mode:
            csv_path:
        Returns:
        """
        merged_forecasts = pd.merge(forecasts["condition"],
                                    forecasts["dist"],
                                    on=["Step_ID","Track_ID","Ensemble_Member","Forecast_Hour"])
        all_members = self.data[mode]["combo"]["Ensemble_Member"]
        members = np.unique(all_members)
        all_run_dates = pd.DatetimeIndex(self.data[mode]["combo"]["Run_Date"])
        run_dates = pd.DatetimeIndex(np.unique(all_run_dates))
        print(run_dates)
        for member in members:
            for run_date in run_dates:
                mem_run_index = (all_run_dates == run_date) & (all_members == member)
                member_forecast = merged_forecasts.loc[mem_run_index]
                member_forecast.to_csv(join(csv_path, "hail_forecasts_{0}_{1}_{2}.csv".format(self.ensemble_name,
                                                                                              member,
                                                                                              run_date.strftime
																							  (run_date_format))))
        return

    def output_forecasts_json_parallel(self, forecasts,
                                       condition_model_names,
                                       size_model_names,
                                       dist_model_names,
                                       track_model_names,
                                       json_data_path,
                                       out_path,
                                       num_procs):
        pool = Pool(num_procs)
        total_tracks = self.data["forecast"]["total_group"]
        for r in total_tracks.index:
            track_id = "{0}".format(total_tracks.loc[r, "Track_ID"])
            print(track_id)
            track_num = track_id.split("_")[-1]
            ensemble_name = "{0}".format(total_tracks.loc[r, "Ensemble_Name"])
            member = "{0}".format(total_tracks.loc[r, "Ensemble_Member"])
            group = total_tracks.loc[r, self.group_col]
            run_date = track_id.split("_")[-4][:8]
            step_forecasts = {}
            for ml_model in condition_model_names:
                step_forecasts["condition_" + ml_model.replace(" ", "-")] = forecasts["condition"][group].loc[
                    forecasts["condition"][group]["Track_ID"] == track_id, ml_model].copy()
            # for ml_model in size_model_names:
            #    step_forecasts["size_" + ml_model.replace(" ", "-")] = forecasts["size"][group][ml_model].loc[
            #        forecasts["size"][group][ml_model]["Track_ID"] == track_id].copy()
            for ml_model in dist_model_names:
                step_forecasts["dist_" + ml_model.replace(" ", "-")] = forecasts["dist"][group][ml_model].loc[
                    forecasts["dist"][group][ml_model]["Track_ID"] == track_id].copy()
            # for model_type in forecasts["track"].keys():
            #    for ml_model in track_model_names:
            #        mframe = forecasts["track"][model_type][group][ml_model]
            #        step_forecasts[model_type + "_" + ml_model.replace(" ", "-")] = mframe.loc[
            #            mframe["Track_ID"] == track_id].copy()
            pool.apply_async(output_forecast, (step_forecasts, run_date, ensemble_name, member, track_num,
                                               json_data_path,
                                               out_path))
        pool.close()
        pool.join()
        return


def output_forecast(step_forecasts, run_date, ensemble_name, member, track_num, json_data_path, out_path):
    json_file_name = "{0}_{1}_{2}_model_track_{3}.json".format(ensemble_name,
                                                               run_date,
                                                               member,
                                                               track_num)
    full_json_path = json_data_path + "/".join([run_date, member]) + "/" + json_file_name
    try:
        json_file_obj = open(full_json_path)
        track_obj = json.load(json_file_obj)
        json_file_obj.close()
    except IOError:
        print(full_json_path + " not found")
        return
    for f, feature in enumerate(track_obj['features']):
        del feature['properties']['attributes']
        for model_name, fdata in step_forecasts.items():
            ml_model_name = model_name.split("_")[1]
            if "condition" in model_name:
                feature['properties'][model_name] = fdata.values[f]
            else:
                predcols = []
                for col in fdata.columns:
                    if ml_model_name in col:
                        predcols.append(col)
                feature['properties'][model_name] = fdata.loc[:, predcols].values[f].tolist()
    full_path = []
    for part in [run_date, member]:
        full_path.append(part)
        if not os.access(out_path + "/".join(full_path), os.R_OK):
            try:
                os.mkdir(out_path + "/".join(full_path))
            except OSError:
                print("directory already created")
    out_json_filename = out_path + "/".join(full_path) + "/" + json_file_name
    try:
        out_json_obj = open(out_json_filename, "w")
        json.dump(track_obj, out_json_obj, indent=1, sort_keys=True)
        out_json_obj.close()
    except IOError:
        print(out_json_filename + " not found")
        return
    return
