mlbaseline.py

# -*- coding: utf-8 -*-
"""mlbaseline.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1T7c2kh_IH8EJcQlEs3h9YBRnKUDEp9V2

# Crop Yield Prediction - ML Baseline

We use WOFOST crop growth indicators, weather variables, geographic information, soil data and remote sensing indicators to predict the yield.

## Google Colab Notes

**To run the script in Google Colab environment**
1. Download the data directory and save it somewhere convenient.
2. Open the notebook using Google Colaboratory.
3. Create a copy of the notebook for yourself.
4. Click connect on the right hand side of the bar below menu items. When you are connected to a machine, you will see a green tick mark and bars showing RAM and disk.
5. Click the folder icon on the left sidebar and click upload. Upload the data files you downloaded. Click *Ok* when you see a warning saying the files will be deleted after the session is disconnected.
6. Use *Runtime* -> *Run before* option to run all cells before **Set Configuration**.
7. Run the remaining cells except **Python Script Main**. The configuration subsection allows you to change configuration and rerun experiments.

## Global Variables and Spark Installation/Initialization

Initialize Spark session and global variables. Package installation is required only in Google Colab.
"""

#%%writefile globals.py
# test_env = 'notebook'
test_env = 'cluster'
# test_env = 'pkg'

# change to False to skip tests
run_tests = False

# NUTS levels
nuts_levels = ['NUTS' + str(i) for i in range(4)]

# country codes
countries = ['BG', 'DE', 'ES', 'FR', 'HU', 'IT', 'NL', 'PL', 'RO']

# debug levels
debug_levels = [i for i in range(5)]

# Keeping these two mappings inside CYPConfiguration leads to SPARK-5063 error
# when lambda functions use them. Therefore, they are defined as globals now.

# crop name to id mapping
crop_id_dict = {
    'grain maize': 2,
    'sugar beet' : 6,
    'sugarbeet' : 6,
    'sugarbeets' : 6,
    'sugar beets' : 6,
    'total potatoes' : 7,
    'potatoes' : 7,
    'potato' : 7,
    'winter wheat' : 90,
    'soft wheat' : 90,
    'sunflower' : 93,
    'spring barley' : 95,
}

# crop id to name mapping
crop_name_dict = {
    2 : 'grain maize',
    6 : 'sugarbeet',
    7 : 'potatoes',
    90 : 'soft wheat',
    93 : 'sunflower',
    95 : 'spring barley',
}

import pyspark

from pyspark.sql import functions as SparkF
from pyspark.sql import types as SparkT

from pyspark import SparkContext
from pyspark import SparkConf
from pyspark.sql import SparkSession
from pyspark.sql import SQLContext

SparkContext.setSystemProperty('spark.executor.memory', '12g')
SparkContext.setSystemProperty('spark.driver.memory', '6g')
spark = SparkSession.builder.master("local[*]").getOrCreate()
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

sc = SparkContext.getOrCreate()
sqlContext = SQLContext(sc)

"""## Utility Functions"""

#%%writefile util.py
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# crop name and id mappings
def cropNameToID(crop_id_dict, crop):
  """
  Return id of given crop. Relies on crop_id_dict.
  Return 0 if crop name is not in the dictionary.
  """
  crop_lcase = crop.lower()
  try:
    crop_id = crop_id_dict[crop_lcase]
  except KeyError as e:
    crop_id = 0

  return crop_id

def cropIDToName(crop_name_dict, crop_id):
  """
  Return crop name for given crop ID. Relies on crop_name_dict.
  Return 'NA' if crop id is not found in the dictionary.
  """
  try:
    crop_name = crop_name_dict[crop_id]
  except KeyError as e:
    crop_name = 'NA'

  return crop_name

def getYear(date_str):
  """Extract year from date in yyyyMMdd or dd/MM/yyyy format."""
  return SparkF.when(SparkF.length(date_str) == 8,
                     SparkF.year(SparkF.to_date(date_str, 'yyyyMMdd')))\
                     .otherwise(SparkF.year(SparkF.to_date(date_str, 'dd/MM/yyyy')))

def getMonth(date_str):
  """Extract month from date in yyyyMMdd or dd/MM/yyyy format."""
  return SparkF.when(SparkF.length(date_str) == 8,
                     SparkF.month(SparkF.to_date(date_str, 'yyyyMMdd')))\
                     .otherwise(SparkF.month(SparkF.to_date(date_str, 'dd/MM/yyyy')))

def getDay(date_str):
  """Extract day from date in yyyyMMdd or dd/MM/yyyy format."""
  return SparkF.when(SparkF.length(date_str) == 8,
                     SparkF.dayofmonth(SparkF.to_date(date_str, 'yyyyMMdd')))\
                     .otherwise(SparkF.dayofmonth(SparkF.to_date(date_str, 'dd/MM/yyyy')))

# 1-10: Dekad 1
# 11-20: Dekad 2
# > 20 : Dekad 3
def getDekad(date_str):
  """Extract dekad from date in YYYYMMDD format."""
  month = getMonth(date_str)
  day = getDay(date_str)
  return SparkF.when(day < 30, (month - 1)* 3 +
                     SparkF.ceil(day/10)).otherwise((month - 1) * 3 + 3)

# Machine Learning Utility Functions

# This definition is from the suggested answer to:
# https://stats.stackexchange.com/questions/58391/mean-absolute-percentage-error-mape-in-scikit-learn/294069#294069
def mean_absolute_percentage_error(Y_true, Y_pred):
  """Mean Absolute Percentage Error"""
  Y_true, Y_pred = np.array(Y_true), np.array(Y_pred)
  return np.mean(np.abs((Y_true - Y_pred) / Y_true)) * 100

def printFeatures(features, indices, log_fh=None):
  """Print names of features in groups of 5"""
  num_features = len(indices)
  groups = int(num_features/5) + 1

  feature_str = '\n'
  for g in range(groups):
    group_start = g * 5
    group_end = (g + 1) * 5
    if (group_end > num_features):
      group_end = num_features

    group_indices = indices[group_start:group_end]
    for idx in group_indices:
      if (idx == group_indices[-1]):
        feature_str += str(idx+1) + ': ' + features[idx]
      else:
        feature_str += str(idx+1) + ': ' + features[idx] + ', '

    feature_str += '\n'

  print(feature_str)
  if (log_fh is not None):
    log_fh.write(feature_str)

def getPredictionScores(Y_true, Y_predicted, metrics):
  """Get values of metrics for given Y_predicted and Y_true"""
  pred_scores = {}

  for met in metrics:
    score_function = metrics[met]
    met_score = score_function(Y_true, Y_predicted)
    # for RMSE, score_function is mean_squared_error, take square root
    # normalize RMSE
    if (met == 'RMSE'):
      met_score = np.round(100*np.sqrt(met_score)/np.mean(Y_true), 2)
      pred_scores['NRMSE'] = met_score
    # normalize mean absolute errors except MAPE which is already a percentage
    elif ((met == 'MAE') or (met == 'MdAE')):
      met_score = np.round(100*met_score/np.mean(Y_true), 2)
      pred_scores['N' + met] = met_score
    # MAPE, R2, ... : no postprocessing
    else:
      met_score = np.round(met_score, 2)
      pred_scores[met] = met_score

  return pred_scores

def getFilename(crop, country, yield_trend,
                early_season, early_season_end, nuts_level=None):
  """Get filename based on input arguments"""
  suffix = crop.replace(' ', '_')
  suffix += '_' + country

  if (nuts_level is not None):
    suffix += '_' + nuts_level

  if (yield_trend):
    suffix += '_trend'
  else:
    suffix += '_notrend'

  if (early_season):
    suffix += '_early' + str(early_season_end)

  return suffix

def getLogFilename(crop, country, yield_trend,
                   early_season, early_season_end):
  """Get filename for experiment log"""
  log_file = getFilename(crop, country, yield_trend,
                         early_season, early_season_end)
  return log_file + '.log'

def getFeatureFilename(crop, country, yield_trend,
                       early_season, early_season_end):
  """Get unique filename for features"""
  feature_file = 'ft_'
  suffix = getFilename(crop, country, yield_trend, early_season, early_season_end)
  feature_file += suffix
  return feature_file

def getPredictionFilename(crop, country, nuts_level, yield_trend,
                          early_season, early_season_end):
  """Get unique filename for predictions"""
  pred_file = 'pred_'
  suffix = getFilename(crop, country, yield_trend,
                       early_season, early_season_end, nuts_level)
  pred_file += suffix
  return pred_file

def plotTrend(years, actual_values, trend_values, trend_label):
  """Plot a linear trend and scatter plot of actual values"""
  plt.scatter(years, actual_values, color="blue", marker="o")
  plt.plot(years, trend_values, '--')
  plt.xticks(np.arange(years[0], years[-1] + 1, step=len(years)/5))
  ax = plt.axes()
  plt.xlabel("YEAR")
  plt.ylabel(trend_label)
  plt.title(trend_label + ' Trend by YEAR')
  plt.show()

def plotTrueVSPredicted(actual, predicted):
  """Plot actual and predicted values"""
  fig, ax = plt.subplots()
  ax.scatter(np.asarray(actual), predicted)
  ax.plot([actual.min(), actual.max()], [actual.min(), actual.max()], 'k--', lw=4)
  ax.set_xlabel('Actual')
  ax.set_ylabel('Predicted')
  plt.show()

"""## Configuration"""

#%%writefile config.py
from sklearn.model_selection import GridSearchCV
from sklearn.linear_model import Lasso
from sklearn.linear_model import Ridge
from sklearn.svm import SVR
from sklearn.ensemble import ExtraTreesRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.neural_network import MLPRegressor

from sklearn.pipeline import Pipeline
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import mutual_info_regression
from sklearn.feature_selection import SelectFromModel
from sklearn.feature_selection import RFE

from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import RobustScaler

from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import r2_score
from sklearn.metrics import explained_variance_score
from sklearn.metrics import median_absolute_error

class CYPConfiguration:
  def __init__(self, crop_name='potatoes', country_code='NL', season_cross='N'):
    self.config = {
        'crop_name' : crop_name,
        'crop_id' : cropNameToID(crop_id_dict, crop_name),
        'season_crosses_calendar_year' : season_cross,
        'country_code' : country_code,
        'nuts_level' : 'NUTS2',
        'data_sources' : [ 'WOFOST', 'METEO_DAILY', 'SOIL', 'YIELD' ],
        'use_yield_trend' : 'N',
        'predict_yield_residuals' : 'N',
        'find_optimal_trend_window' : 'N',
        # set it to a list with one entry for fixed window
        'trend_windows' : [5, 7, 10],
        'use_centroids' : 'N',
        'use_remote_sensing' : 'Y',
        'early_season_prediction' : 'N',
        'early_season_end_dekad' : 0,
        'data_path' : '.',
        'output_path' : '.',
        'save_features' : 'N',
        'use_saved_features' : 'N',
        'save_predictions' : 'Y',
        'use_saved_predictions' : 'N',
        'compare_with_mcyfs' : 'N',
        'debug_level' : 0,
    }

    # Description of configuration parameters
    # This should be in sync with config above
    self.config_desc = {
        'crop_name' : 'Crop name',
        'crop_id' : 'Crop ID',
        'season_crosses_calendar_year' : 'Crop growing season crosses calendar year boundary',
        'country_code' : 'Country code (e.g. NL)',
        'nuts_level' : 'NUTS level for yield prediction',
        'data_sources' : 'Input data sources',
        'use_yield_trend' : 'Estimate and use yield trend',
        'predict_yield_residuals' : 'Predict yield residuals instead of full yield',
        'find_optimal_trend_window' : 'Find optimal trend window',
        'trend_windows' : 'List of trend window lengths (number of years)',
        'use_centroids' : 'Use centroid coordinates and distance to coast',
        'use_remote_sensing' : 'Use remote sensing data (FAPAR)',
        'early_season_prediction' : 'Predict yield early in the season',
        'early_season_end_dekad' : 'Early season end dekad relative to harvest',
        'data_path' : 'Path to all input data. Default is current directory.',
        'output_path' : 'Path to all output files. Default is current directory.',
        'save_features' : 'Save features to a CSV file',
        'use_saved_features' : 'Use features from a CSV file',
        'save_predictions' : 'Save predictions to a CSV file',
        'use_saved_predictions' : 'Use predictions from a CSV file',
        'compare_with_mcyfs' : 'Compare predictions with MARS Crop Yield Forecasting System',
        'debug_level' : 'Debug level to control amount of debug information',
    }

    ########### Machine learning configuration ###########
    # test fraction
    self.test_fraction = 0.3

    # scaler
    self.scaler = StandardScaler()

    # Feature selection algorithms. Initialized in getFeatureSelectors().
    self.feature_selectors = {}

    # prediction algorithms
    self.estimators = {
        # linear model
        'Ridge' : {
            'estimator' : Ridge(alpha=1, random_state=42, max_iter=1000,
                                copy_X=True, fit_intercept=True),
            'fs_param_grid' : dict(estimator__alpha=[1e-1]),
            'param_grid' : dict(estimator__alpha=[1e-5, 1e-2, 1e-1, 0.5, 1, 5, 10])
        },
        'KNN' : {
            'estimator' : KNeighborsRegressor(weights='distance'),
            'fs_param_grid' : dict(estimator__n_neighbors=[5]),
            'param_grid' : dict(estimator__n_neighbors=[3, 5, 7, 9])
        },
        # SVM regression
        'SVR' : {
            'estimator' : SVR(kernel='rbf', gamma='scale', max_iter=-1,
                              shrinking=True, tol=0.001),
            'fs_param_grid' : dict(estimator__C=[10.0],
                                   estimator__epsilon=[0.5]),
            'param_grid' : dict(estimator__C=[1e-1, 5e-1, 1.0, 5.0, 10.0, 50.0, 100.0, 200.0],
                                estimator__epsilon=[1e-2, 1e-1, 0.5, 1.0, 5.0]),
        },
        # random forest
        #'RF' : {
        #    'estimator' : RandomForestRegressor(bootstrap=True, random_state=42,
        #                                        oob_score=True, min_samples_leaf=5),
        #    'fs_param_grid' : dict(estimator__max_depth=[7],
        #                           estimator__n_estimators=[100]),
        #    'param_grid' : dict(estimator__max_depth=[5, 7],
        #                        estimator__n_estimators=[100, 500])
        #},
        # extra randomized trees
        #'ERT' : {
        #    'estimator' : ExtraTreesRegressor(bootstrap=True, random_state=42,
        #                                      oob_score=True, min_samples_leaf=5),
        #    'fs_param_grid' : dict(estimator__max_depth=[7],
        #                           estimator__n_estimators=[100]),
        #    'param_grid' : dict(estimator__max_depth=[5, 7],
        #                        estimator__n_estimators=[100, 500])
        #},
        # gradient boosted decision trees
        'GBDT' : {
            'estimator' : GradientBoostingRegressor(learning_rate=0.01,
                                                    subsample=0.8, loss='huber',
                                                    min_samples_leaf=5,
                                                    random_state=42),
            'fs_param_grid' : dict(estimator__max_depth=[5],
                                   estimator__n_estimators=[100]),
            'param_grid' : dict(estimator__max_depth=[5, 10, 15],
                                estimator__n_estimators=[100, 500])
        },
        #'MLP' : {
        #    'estimator' : MLPRegressor(batch_size='auto', learning_rate='adaptive',
        #                               solver='sgd', activation='relu',
        #                               learning_rate_init=0.01, power_t=0.5,
        #                               max_iter=1000, shuffle=True,
        #                               random_state=42, tol=0.001,
        #                               verbose=False, warm_start=False,
        #                               momentum=0.9, nesterovs_momentum=True,
        #                               early_stopping=True,
        #                               validation_fraction=0.4, beta_1=0.9,
        #                               beta_2=0.999, epsilon=1e-08),
        #    'fs_param_grid' : dict(estimator__hidden_layer_sizes=[(10, 10), (15,15)],
        #                           estimator__alpha=[0.2, 0.3]),
        #    'param_grid' : dict(estimator__hidden_layer_sizes=[(10, 10), (15, 15), (20, 20)],
        #                        estimator__alpha=[0.1, 0.2, 0.3]),
        #},
   }

    # k-fold validation metric for feature selection
    self.fs_cv_metric = 'neg_mean_squared_error'
    # k-fold validation metric for training
    self.est_cv_metric = 'neg_mean_squared_error'

    # Performance evaluation metrics:
    # sklearn supports these metrics:
    # 'explained_variance', 'max_error', 'neg_mean_absolute_error
    # 'neg_mean_squared_error', 'neg_root_mean_squared_error'
    # 'neg_mean_squared_log_error', 'neg_median_absolute_error', 'r2'
    self.eval_metrics = {
        # EXP_VAR (y_true, y_obs) = 1 - ( var(y_true - y_obs) / var (y_true) )
        #'EXP_VAR' : explained_variance_score,
        # MAE (y_true, y_obs) = ( 1 / n ) * sum_i-n ( | y_true_i - y_obs_i | )
        'MAE' : mean_absolute_error,
        # MdAE (y_true, y_obs) = median ( | y_true_1 - y_obs_1 |, | y_true_2 - y_obs_2 |, ... )
        #'MdAE' : median_absolute_error,
        # MAPE (y_true, y_obs) = ( 1 / n ) * sum_i-n ( ( y_true_i - y_obs_i ) / y_true_i )
        'MAPE' : mean_absolute_percentage_error,
        # MSE (y_true, y_obs) = ( 1 / n ) * sum_i-n ( y_true_i - y_obs_i )^2
        'RMSE' : mean_squared_error,
        # R2 (y_true, y_obs) = 1 - ( ( sum_i-n ( y_true_i - y_obs_i )^2 )
        #                           / sum_i-n ( y_true_i - mean(y_true) )^2)
        'R2' : r2_score,
    }

  ########### Setters and getters ###########
  def setCropName(self, crop_name):
    """Set the crop name"""
    crop = crop_name.lower()
    assert crop in crop_id_dict
    self.config['crop_name'] = crop
    self.config['crop_id'] = cropNameToID(crop_id_dict, crop)

  def getCropName(self):
    """Return the crop name"""
    return self.config['crop_name']

  def setCropID(self, crop_id):
    """Set the crop ID"""
    assert crop_id in crop_name_dict
    self.config['crop_id'] = crop_id
    self.config['crop_name'] = cropIDToName(crop_name_dict, crop_id)

  def getCropID(self):
    """Return the crop ID"""
    return self.config['crop_id']

  def setSeasonCrossesCalendarYear(self, season_crosses):
    """Set whether the season crosses calendar year boundary"""
    scross = season_crosses.upper()
    assert scross in ['Y', 'N']
    self.config['season_crosses_calendar_year'] = scross

  def seasonCrossesCalendarYear(self):
    """Return whether the season crosses calendar year boundary"""
    return (self.config['season_crosses_calendar_year'] == 'Y')

  def setCountryCode(self, country_code):
    """Set the country code"""
    if (country_code is None):
      self.config['country_code'] = None
    else:
      ccode = country_code.upper()
      assert len(ccode) == 2
      assert ccode in countries
      self.config['country_code'] = ccode

  def getCountryCode(self):
    """Return the country code"""
    return self.config['country_code']

  def setNUTSLevel(self, nuts_level):
    """Set the NUTS level"""
    nuts = nuts_level.upper()
    assert nuts in nuts_levels
    self.config['nuts_level'] = nuts

  def getNUTSLevel(self):
    """Return the NUTS level"""
    return self.config['nuts_level']

  def setDataSources(self, data_sources):
    """Get the data sources"""
    # TODO: some validation
    self.config['data_sources'] = data_sources

  def updateDataSources(self, data_src, include_src, nuts_level=None):
    """add or remove data_src from data sources"""
    src_nuts = self.getNUTSLevel()
    if (nuts_level is not None):
      src_nuts = nuts_level

    data_sources = self.config['data_sources']
    # no update required
    if (((include_src == 'Y') and (data_src in data_sources)) or
        ((include_src == 'N') and (data_src not in data_sources))):
      return

    if (include_src == 'Y'):
      if (isinstance(data_sources, dict)):
        data_sources[data_src] = src_nuts
      else:
        data_sources.append(data_src)
    else:
      if (isinstance(data_sources, dict)):
        del data_sources[data_src]
      else:
        data_sources.remove(data_src)

    self.config['data_sources'] = data_sources

  def getDataSources(self):
    """Return the data sources"""
    return self.config['data_sources']

  def setUseYieldTrend(self, use_trend):
    """Set whether to use yield trend"""
    use_yt = use_trend.upper()
    assert use_yt in ['Y', 'N']
    self.config['use_yield_trend'] = use_yt

  def useYieldTrend(self):
    """Return whether to use yield trend"""
    return (self.config['use_yield_trend'] == 'Y')

  def setPredictYieldResiduals(self, pred_res):
    """Set whether to use predict yield residuals"""
    pred_yres = pred_res.upper()
    assert pred_yres in ['Y', 'N']
    self.config['predict_yield_residuals'] = pred_yres

  def predictYieldResiduals(self):
    """Return whether to use predict yield residuals"""
    return (self.config['predict_yield_residuals'] == 'Y')

  def setFindOptimalTrendWindow(self, find_opt):
    """Set whether to find optimal trend window for each year"""
    find_otw = find_opt.upper()
    assert find_otw in ['Y', 'N']
    self.config['find_optimal_trend_window'] = find_otw

  def findOptimalTrendWindow(self):
    """Return whether to find optimal trend window for each year"""
    return (self.config['find_optimal_trend_window'] == 'Y')

  def setTrendWindows(self, trend_windows):
    """Set trend window lengths (years)"""
    assert isinstance(trend_windows, list)
    assert len(trend_windows) > 0

    # trend windows less than 2 years do not make sense
    for tw in trend_windows:
      assert tw > 2

    self.config['trend_windows'] = trend_windows

  def getTrendWindows(self):
    """Return trend window lengths (years)"""
    return self.config['trend_windows']

  def setUseCentroids(self, use_centroids):
    """Set whether to use centroid coordinates and distance to coast"""
    use_ct = use_centroids.upper()
    assert use_ct in ['Y', 'N']
    self.config['use_centroids'] = use_ct
    self.updateDataSources('CENTROIDS', use_ct)

  def useCentroids(self):
    """Return whether to use centroid coordinates and distance to coast"""
    return (self.config['use_centroids'] == 'Y')

  def setUseRemoteSensing(self, use_remote_sensing):
    """Set whether to use remote sensing data"""
    use_rs = use_remote_sensing.upper()
    assert use_rs in ['Y', 'N']
    self.config['use_remote_sensing'] = use_rs
    self.updateDataSources('REMOTE_SENSING', use_rs, 'NUTS2')

  def useRemoteSensing(self):
    """Return whether to use remote sensing data"""
    return (self.config['use_remote_sensing'] == 'Y')

  def setEarlySeasonPrediction(self, early_season):
    """Set whether to do early season prediction"""
    ep = early_season.upper()
    assert ep in ['Y', 'N']
    self.config['early_season_prediction'] = ep

  def earlySeasonPrediction(self):
    """Return whether to do early season prediction"""
    return (self.config['early_season_prediction'] == 'Y')

  def setEarlySeasonEndDekad(self, end_dekad):
    """Set early season prediction dekad"""
    dekads_range = [dek for dek in range(1, 37)]
    assert end_dekad in dekads_range
    self.config['early_season_end_dekad'] = end_dekad

  def getEarlySeasonEndDekad(self):
    """Return early season prediction dekad"""
    return self.config['early_season_end_dekad']

  def setDataPath(self, data_path):
    """Set the data path"""
    # TODO: some validation
    self.config['data_path'] = data_path

  def getDataPath(self):
    """Return the data path"""
    return self.config['data_path']

  def setOutputPath(self, out_path):
    """Set the path to output files. TODO: some validation."""
    self.config['output_path'] = out_path

  def getOutputPath(self):
    """Return the path to output files."""
    return self.config['output_path']

  def setSaveFeatures(self, save_ft):
    """Set whether to save features in a CSV file"""
    sft = save_ft.upper()
    assert sft in ['Y', 'N']
    self.config['save_features'] = sft

  def saveFeatures(self):
    """Return whether to save features in a CSV file"""
    return (self.config['save_features'] == 'Y')

  def setUseSavedFeatures(self, use_saved):
    """Set whether to use features from CSV file"""
    saved = use_saved.upper()
    assert saved in ['Y', 'N']
    self.config['use_saved_features'] = saved

  def useSavedFeatures(self):
    """Return whether to use to use features from CSV file"""
    return (self.config['use_saved_features'] == 'Y')

  def setSavePredictions(self, save_pred):
    """Set whether to save predictions in a CSV file"""
    spd = save_pred.upper()
    assert spd in ['Y', 'N']
    self.config['save_predictions'] = spd

  def savePredictions(self):
    """Return whether to save predictions in a CSV file"""
    return (self.config['save_predictions'] == 'Y')

  def setUseSavedPredictions(self, use_saved):
    """Set whether to use predictions from CSV file"""
    saved = use_saved.upper()
    assert saved in ['Y', 'N']
    self.config['use_saved_predictions'] = saved

  def useSavedPredictions(self):
    """Return whether to use to use predictions from CSV file"""
    return (self.config['use_saved_predictions'] == 'Y')

  def setCompareWithMCYFS(self, compare_mcyfs):
    """Set whether to compare predictions with MCYFS"""
    comp_mcyfs = compare_mcyfs.upper()
    assert comp_mcyfs in ['Y', 'N']
    self.config['compare_with_mcyfs'] = comp_mcyfs

  def compareWithMCYFS(self):
    """Return whether to compare predictions with MCYFS"""
    return (self.config['compare_with_mcyfs'] == 'Y')

  def setDebugLevel(self, debug_level):
    """Set the debug level"""
    assert debug_level in debug_levels
    self.config['debug_level'] = debug_level

  def getDebugLevel(self):
    """Return the debug level"""
    return self.config['debug_level']

  def updateConfiguration(self, config_update):
    """Update configuration"""
    assert isinstance(config_update, dict)
    for k in config_update:
      assert k in self.config

      # keys that need special handling
      special_cases = {
          'crop_name' : self.setCropName,
          'crop_id' : self.setCropID,
          'use_centroids' : self.setUseCentroids,
          'use_remote_sensing' : self.setUseRemoteSensing,
      }

      if (k not in special_cases):
        self.config[k] = config_update[k]
        continue

      # special case
      special_cases[k](config_update[k])

  def printConfig(self, log_fh):
    """Print current configuration and write configuration to log file."""
    config_str = '\nCurrent ML Baseline Configuration'
    config_str += '\n--------------------------------'
    for k in self.config:
      if (isinstance(self.config[k], dict)):
        conf_keys = list(self.config[k].keys())
        if (not isinstance(conf_keys[0], str)):
          conf_keys = [str(k) for k in conf_keys]

        config_str += '\n' + self.config_desc[k] + ': ' + ', '.join(conf_keys)
      elif (isinstance(self.config[k], list)):
        conf_vals = self.config[k]
        if (not isinstance(conf_vals[0], str)):
          conf_vals = [str(k) for k in conf_vals]

        config_str += '\n' + self.config_desc[k] + ': ' + ', '.join(conf_vals)
      else:
        conf_val = self.config[k]
        if (not isinstance(conf_val, str)):
          conf_val = str(conf_val)

        config_str += '\n' + self.config_desc[k] + ': ' + conf_val

    config_str += '\n'
    log_fh.write(config_str + '\n')
    print(config_str)

  # Machine learning configuration
  def getTestFraction(self):
    """Return test set fraction (of full dataset)"""
    return self.test_fraction

  def setTestFraction(self, test_fraction):
    """Set test set fraction (of full dataset)"""
    assert (test_fraction > 0.0 and test_fraction < 1.0)
    self.test_fraction = test_fraction

  def getFeatureScaler(self):
    """Return feature scaling method"""
    return self.scaler

  def setFeatureScaler(self, scaler):
    """Set feature scaling method"""
    assert (isinstance(scaler, MinMaxScaler) or isinstance(scaler, StandardScaler))
    self.scaler = scaler

  def getFeatureSelectionCVMetric(self):
    """Return metric for feature selection using K-fold validation"""
    return self.fs_cv_metric

  def setFeatureSelectionCVMetric(self, fs_metric):
    """Return metric for feature selection using K-fold validation"""
    assert fs_metric in self.eval_metrics
    self.fs_cv_metric = fs_metric

  def getAlgorithmTrainingCVMetric(self):
    """Return metric for hyperparameter optimization using K-fold validation"""
    return self.est_cv_metric

  def setFeatureSelectionCVMetric(self, est_metric):
    """Return metric for hyperparameter optimization using K-fold validation"""
    assert est_metric in self.eval_metrics
    self.est_cv_metric = est_metric

  def getFeatureSelectors(self, X_train, Y_train, num_features,
                          custom_cv):
    """Feature selection methods"""
    # already defined?
    if (len(self.feature_selectors) > 0):
      return self.feature_selectors

    # NOTE: X_train, Y_train, custom_cv
    # are for optimizing hyperparamters of rf and lasso used
    # to define feature selectors. At the moment, we don't
    # optimize hyperparameters.

    # Early season prediction can have less than 10 features
    min_features = 10 if num_features > 10 else num_features
    max_features = [min_features]

    if (num_features > 15):
      max_features.append(15)
    if (num_features > 20):
      max_features.append(20)

    use_yield_trend = self.useYieldTrend()
    if ((num_features > 25) and (use_yield_trend)):
      max_features.append(25)

    rf = RandomForestRegressor(n_estimators=100, max_depth=5,
                               bootstrap=True, random_state=42,
                               oob_score=True, min_samples_leaf=5)

    lasso = Lasso(alpha=0.1, copy_X=True, fit_intercept=True,
                  random_state=42,selection='cyclic', tol=0.01)

    self.feature_selectors = {
      # random forest
      'random_forest' : {
          'selector' : SelectFromModel(rf, threshold='median'),
          'param_grid' : dict(selector__max_features=max_features)
      },
      # recursive feature elimination using Lasso
      'RFE_Lasso' : {
          'selector' : RFE(lasso),
          'param_grid' : dict(selector__n_features_to_select=max_features)
      },
      # NOTE: Mutual info raises an error when used with spark parallel backend.
      # univariate feature selection
      # 'mutual_info' : {
      #     'selector' : SelectKBest(mutual_info_regression),
      #     'param_grid' : dict(selector__k=max_features)
      # },
    }

    return self.feature_selectors

  def setFeatureSelectors(self, ft_sel):
    """Set feature selection algorithms"""
    assert isinstance(ft_sel, dict)
    assert len(ft_sel) > 0
    for sel in ft_sel:
      assert isinstance(sel, dict)
      assert 'selector' in sel
      assert 'param_grid' in sel
      # add cases if other feature selection methods are used
      assert (isinstance(sel['selector'], SelectKBest) or
              isinstance(sel['selector'], SelectFromModel) or
              isinstance(sel['selector'], RFE))
      assert isinstance(sel['param_grid'], dict)

    self.feature_selectors = ft_sel

  def getEstimators(self):
    """Return machine learning algorithms for prediction"""
    return self.estimators
  
  def setEstimators(self, estimators):
    """Set machine learning algorithms for prediction"""
    assert isinstance(estimators, dict)
    assert len(estimators) > 0
    for est in estimators:
      assert isinstance(est, dict)
      assert 'estimator' in est
      assert 'param_grid' in est
      assert 'fs_param_grid' in est
      assert isinstance(est['param_grid'], dict)
      assert isinstance(est['fs_param_grid'], dict)

    self.estimators = estimators
  
  def getEvaluationMetrics(self):
    """Return metrics to evaluate predictions of algorithms"""
    return self.eval_metrics
  
  def setEvaluationMetrics(self, metrics):
    assert isinstance(metrics, dict)
    self.eval_metrics = metrics

"""## Workflow

### Data Loading and Preprocessing

#### Data Loader Class
"""

#%%writefile data_loading.py
class CYPDataLoader:
  def __init__(self, spark, cyp_config):
    self.spark = spark
    self.data_path = cyp_config.getDataPath()
    self.country_code = cyp_config.getCountryCode()
    self.nuts_level = cyp_config.getNUTSLevel()
    self.data_sources = cyp_config.getDataSources()
    self.verbose = cyp_config.getDebugLevel()
    self.data_dfs = {}

  def loadFromCSVFile(self, data_path, src, nuts, country_code):
    """
    The implied filename for each source is:
    <data_source>_<nuts_level>_<country_code>.csv
    Examples: CENTROIDS_NUTS2_NL.csv, WOFOST_NUTS2_NL.csv.
    Schema is inferred from the file. We might want to specify the schema at some point.
    """
    if (country_code is not None):
      datafile = data_path + '/' + src  + '_' + nuts + '_' + country_code + '.csv'
    else:
      datafile = data_path + '/' + src  + '_' + nuts + '.csv'

    if (self.verbose > 1):
      print('Data file name', '"' + datafile + '"')

    df = self.spark.read.csv(datafile, header = True, inferSchema = True)
    return df

  def loadData(self, src, nuts_level):
    """
    Load data for a specific data source.
    nuts_level may one level or a list of levels.
    """
    data_path = self.data_path
    country_code = self.country_code
    assert src in self.data_sources

    if (isinstance(nuts_level, list)):
      src_dfs = []
      for nuts in nuts_level:
        df = self.loadFromCSVFile(data_path, src, nuts, country_code)
        src_dfs.append(df)

    elif (isinstance(nuts_level, str)):
      src_dfs = self.loadFromCSVFile(data_path, src, nuts_level, country_code)
    else:
      src_dfs = None

    return src_dfs

  def loadAllData(self):
    """
    NOT SUPPORTED:
    1. Schema is not defined.
    2. Loading data for multiple countries.
    3. Loading data from folders.
    Ioannis: Spark has a nice way of loading several files from a folder,
    and associating the file name on each record, using the function
    input_file_name. This allows to extract medatada from the path
    into the dataframe. In your case it could be the country name, etc.
    """
    data_dfs = {}
    for src in self.data_sources:
      nuts_level = self.nuts_level
      if (isinstance(self.data_sources, dict)):
        nuts_level = self.data_sources[src]
      # REMOTE_SENSING data is at NUTS2. If nuts_level is None, leave as is.
      elif ((src == 'REMOTE_SENSING') and (nuts_level is not None)):
        nuts_level = 'NUTS2'

      if ('METEO' in src):
        data_dfs['METEO'] = self.loadData(src, nuts_level)
      else:
        data_dfs[src] = self.loadData(src, nuts_level)

    if (self.verbose > 1):
      data_sources_str = ''
      for src in data_dfs:
        data_sources_str = data_sources_str + src + ', '

      # remove the comma and space from the end
      print('Loaded data:', data_sources_str[:-2])
      print('\n')

    return data_dfs

"""#### Data Preprocessor Class"""

#%%writefile data_preprocessing.py
from pyspark.sql import Window

class CYPDataPreprocessor:
  def __init__(self, spark, cyp_config):
    self.spark = spark
    self.verbose = cyp_config.getDebugLevel()

  def extractYearDekad(self, df):
    """Extract year and dekad from date_col in yyyyMMdd format."""
    # Conversion to string type is required to make getYear(), getMonth() etc. work correctly.
    # They use to_date() function to verify valid dates and to_date() expects the date column to be string.
    df = df.withColumn('DATE', df['DATE'].cast("string"))
    df = df.select('*',
                   getYear('DATE').alias('FYEAR'),
                   getDekad('DATE').alias('DEKAD'))

    # Bring FYEAR, DEKAD to the front
    col_order = df.columns[:2] + df.columns[-2:] + df.columns[2:-2]
    df = df.select(col_order).drop('DATE')
    return df

  def getCropSeasonInformation(self, wofost_df, season_crosses_calyear):
    """Crop season information based on WOFOST DVS"""
    join_cols = ['IDREGION', 'FYEAR']
    if (('DATE' in wofost_df.columns) and ('FYEAR' not in wofost_df.columns)):
      wofost_df = self.extractYearDekad(wofost_df)

    crop_season = wofost_df.select(join_cols).distinct()
    diff_window = Window.partitionBy(join_cols).orderBy('DEKAD')
    cs_window = Window.partitionBy('IDREGION').orderBy('FYEAR')

    wofost_df = wofost_df.withColumn('VALUE', wofost_df['DVS'])
    wofost_df = wofost_df.withColumn('PREV', SparkF.lag(wofost_df['VALUE']).over(diff_window))
    wofost_df = wofost_df.withColumn('DIFF', SparkF.when(SparkF.isnull(wofost_df['PREV']), 0)\
                                     .otherwise(wofost_df['VALUE'] - wofost_df['PREV']))
    # calculate end of season dekad
    dvs_nochange_filter = ((wofost_df['VALUE'] >= 200) & (wofost_df['DIFF'] == 0.0))
    year_end_filter = (wofost_df['DEKAD'] == 36)
    if (season_crosses_calyear):
      value_zero_filter =  (wofost_df['VALUE'] == 0)
    else:
      value_zero_filter =  ((wofost_df['PREV'] >= 200) & (wofost_df['VALUE'] == 0))

    end_season_filter = (dvs_nochange_filter | value_zero_filter | year_end_filter)
    crop_season = crop_season.join(wofost_df.filter(end_season_filter).groupBy(join_cols)\
                                   .agg(SparkF.min('DEKAD').alias('SEASON_END_DEKAD')), join_cols)
    wofost_df = wofost_df.drop('VALUE', 'PREV', 'DIFF')

    # We take the max of SEASON_END_DEKAD for current campaign and next campaign
    # to determine which dekads go to next campaign year.
    max_year = crop_season.agg(SparkF.max('FYEAR')).collect()[0][0]
    min_year = crop_season.agg(SparkF.min('FYEAR')).collect()[0][0]
    crop_season = crop_season.withColumn('NEXT_SEASON_END', SparkF.when(crop_season['FYEAR'] == max_year,
                                                                        crop_season['SEASON_END_DEKAD'])\
                                         .otherwise(SparkF.lead(crop_season['SEASON_END_DEKAD']).over(cs_window)))
    crop_season = crop_season.withColumn('SEASON_END',
                                         SparkF.when(crop_season['SEASON_END_DEKAD'] > crop_season['NEXT_SEASON_END'],
                                                     crop_season['SEASON_END_DEKAD'])\
                                         .otherwise(crop_season['NEXT_SEASON_END']))
    crop_season = crop_season.withColumn('PREV_SEASON_END', SparkF.when(crop_season['FYEAR'] == min_year, 0)\
                                         .otherwise(SparkF.lag(crop_season['SEASON_END']).over(cs_window)))
    crop_season = crop_season.select(join_cols + ['PREV_SEASON_END', 'SEASON_END'])

    return crop_season

  def alignDataToCropSeason(self, df, crop_season, season_crosses_calyear):
    """Calculate CAMPAIGN_YEAR, CAMPAIGN_DEKAD based on crop_season"""
    join_cols = ['IDREGION', 'FYEAR']
    max_year = crop_season.agg(SparkF.max('FYEAR')).collect()[0][0]
    min_year = crop_season.agg(SparkF.min('FYEAR')).collect()[0][0]
    df = df.join(crop_season, join_cols)

    # Dekads > SEASON_END belong to next campaign year
    df = df.withColumn('CAMPAIGN_YEAR',
                       SparkF.when(df['DEKAD'] > df['SEASON_END'], df['FYEAR'] + 1)\
                       .otherwise(df['FYEAR']))
    # min_year has no previous season information. We align CAMPAIGN_DEKAD to end in 36.
    # For other years, dekads < SEASON_END are adjusted based on PREV_SEASON_END.
    # Dekads > SEASON_END get renumbered from 1 (for next campaign).
    df = df.withColumn('CAMPAIGN_DEKAD',
                       SparkF.when(df['CAMPAIGN_YEAR'] == min_year, df['DEKAD'] + 36 - df['SEASON_END'])\
                       .otherwise(SparkF.when(df['DEKAD'] > df['SEASON_END'], df['DEKAD'] - df['SEASON_END'])\
                                  .otherwise(df['DEKAD'] + 36 - df['PREV_SEASON_END'])))

    # Columns should be IDREGION, FYEAR, DEKAD, ..., CAMPAIGN_YEAR, CAMPAIGN_DEKAD.
    # Bring CAMPAIGN_YEAR and CAMPAIGN_DEKAD to the front.
    col_order = df.columns[:3] + df.columns[-2:] + df.columns[3:-2]
    df = df.select(col_order)
    if (season_crosses_calyear):
      # For crop with two seasons, remove the first year. Data from the first year
      # only contributes to the second year and we have already moved useful data
      # to the second year (or first campaign year).
      df = df.filter(df['CAMPAIGN_YEAR'] > min_year)

    # In both cases, remove extra rows beyond max campaign year
    df = df.filter(df['CAMPAIGN_YEAR'] <= max_year)
    return df

  def preprocessWofost(self, wofost_df, crop_season, season_crosses_calyear):
    """
    Extract year and dekad from date. Use crop_season to compute
    CAMPAIGN_YEAR and CAMPAIGN_DEKAD.
    """
    drop_cols = crop_season.columns[2:]
    if (('DATE' in wofost_df.columns) and ('FYEAR' not in wofost_df.columns)):
      wofost_df = self.extractYearDekad(wofost_df)

    join_cols = ['IDREGION', 'FYEAR']
    wofost_df = self.alignDataToCropSeason(wofost_df, crop_season, season_crosses_calyear)

    # WOFOST indicators come after IDREGION, FYEAR, DEKAD, CAMPAIGN_YEAR, CAMPAIGN_DEKAD.
    wofost_inds = wofost_df.columns[5:]
    # set indicators values for dekads after end of season to zero.
    # TODO - Dilli: Find a way to avoid the for loop.
    for ind in wofost_inds:
      wofost_df = wofost_df.withColumn(ind,
                                       SparkF.when(wofost_df['DEKAD'] < wofost_df['SEASON_END'],
                                                   wofost_df[ind])\
                                       .otherwise(0))

    wofost_df = wofost_df.drop(*drop_cols)
    return wofost_df

  def preprocessMeteo(self, meteo_df, crop_season, season_crosses_calyear):
    """
    Extract year and dekad from date, calculate CWB.
    Use crop_season to compute CAMPAIGN_YEAR and CAMPAIGN_DEKAD.
    """
    join_cols = ['IDREGION', 'FYEAR']
    drop_cols = crop_season.columns[2:]
    meteo_df = meteo_df.drop('IDCOVER')
    meteo_df = meteo_df.withColumn('CWB',
                                   SparkF.bround(meteo_df['PREC'] - meteo_df['ET0'], 2))
    if (('DATE' in meteo_df.columns) and ('FYEAR' not in meteo_df.columns)):
      meteo_df = self.extractYearDekad(meteo_df)

    meteo_df = self.alignDataToCropSeason(meteo_df, crop_season, season_crosses_calyear)
    meteo_df = meteo_df.drop(*drop_cols)
    return meteo_df

  def preprocessMeteoDaily(self, meteo_df):
    """
    Convert daily meteo data to dekadal. Takes avg for all indicators
    except TMAX (take max), TMIN (take min), PREC (take sum), ET0 (take sum), CWB (take sum).
    """
    self.spark.catalog.dropTempView('meteo_daily')
    meteo_df.createOrReplaceTempView('meteo_daily')
    join_cols = ['IDREGION', 'CAMPAIGN_YEAR', 'CAMPAIGN_DEKAD']
    join_df = meteo_df.select(join_cols + ['FYEAR', 'DEKAD']).distinct()

    # We are ignoring VPRES, WSPD and RELH at the moment
    # avg(VPRES) as VPRES1, avg(WSPD) as WSPD1, avg(RELH) as RELH1,
    # TMAX| TMIN| TAVG| VPRES| WSPD| PREC| ET0| RAD| RELH| CWB
    #
    # It seems keeping same name after aggregation is fine. We are using a
    # different name just to be sure nothing untoward happens.
    query = 'select IDREGION, CAMPAIGN_YEAR, CAMPAIGN_DEKAD, '
    query = query + ' max(TMAX) as TMAX1, min(TMIN) as TMIN1, '
    query = query + ' bround(avg(TAVG), 2) as TAVG1, bround(sum(PREC), 2) as PREC1, '
    query = query + ' bround(sum(ET0), 2) as ET01, bround(avg(RAD), 2) as RAD1, '
    query = query + ' bround(sum(CWB), 2) as CWB1 '
    query = query + ' from meteo_daily group by IDREGION, CAMPAIGN_YEAR, CAMPAIGN_DEKAD '
    query = query + ' order by IDREGION, CAMPAIGN_YEAR, CAMPAIGN_DEKAD'
    meteo_df = self.spark.sql(query).cache()

    # rename the columns
    selected_cols = ['TMAX', 'TMIN', 'TAVG', 'PREC', 'ET0', 'RAD', 'CWB']
    for scol in selected_cols:
      meteo_df = meteo_df.withColumnRenamed(scol + '1', scol)

    meteo_df = meteo_df.join(join_df, join_cols)
    # Bring FYEAR, DEKAD to the front
    col_order = meteo_df.columns[:1] + meteo_df.columns[-2:] + meteo_df.columns[1:-2]
    meteo_df = meteo_df.select(col_order)

    return meteo_df

  def remoteSensingNUTS2ToNUTS3(self, rs_df, nuts3_regions):
    """
    Convert NUTS2 remote sensing data to NUTS3.
    Remote sensing values for NUTS3 regions are inherited from parent regions.
    NOTE this function is called before preprocessRemoteSensing.
    preprocessRemoteSensing expects crop_season and rs_df to be at the same NUTS level.
    """
    NUTS3_dict = {}

    for nuts3 in nuts3_regions:
      nuts2 = nuts3[:4]
      try:
        existing = NUTS3_dict[nuts2]
      except KeyError as e:
        existing = []

      NUTS3_dict[nuts2] = existing + [nuts3]

    rs_NUTS3 = rs_df.rdd.map(lambda r: (NUTS3_dict[r[0]] if r[0] in NUTS3_dict else [], r[1], r[2]))
    rs_NUTS3_df = rs_NUTS3.toDF(['NUTS3_REG', 'DATE', 'FAPAR'])
    rs_NUTS3_df = rs_NUTS3_df.withColumn('IDREGION', SparkF.explode('NUTS3_REG')).drop('NUTS3_REG')
    rs_NUTS3_df = rs_NUTS3_df.select('IDREGION', 'DATE', 'FAPAR')

    return rs_NUTS3_df

  def preprocessRemoteSensing(self, rs_df, crop_season, season_crosses_calyear):
    """
    Extract year and dekad from date.
    Use crop_season to compute CAMPAIGN_YEAR and CAMPAIGN_DEKAD.
    NOTE crop_season and rs_df must be at the same NUTS level.
    """
    join_cols = ['IDREGION', 'FYEAR']
    drop_cols = crop_season.columns[2:]
    if (('DATE' in rs_df.columns) and ('FYEAR' not in rs_df.columns)):
      rs_df = self.extractYearDekad(rs_df)

    rs_df = self.alignDataToCropSeason(rs_df, crop_season, season_crosses_calyear)
    rs_df = rs_df.drop(*drop_cols)
    return rs_df

  def preprocessCentroids(self, centroids_df):
    df_cols = centroids_df.columns
    centroids_df = centroids_df.withColumn('CENTROID_X', SparkF.bround('CENTROID_X', 2))
    centroids_df = centroids_df.withColumn('CENTROID_Y', SparkF.bround('CENTROID_Y', 2))

    return centroids_df

  def preprocessSoil(self, soil_df):
    # SM_WC = water holding capacity
    soil_df = soil_df.withColumn('SM_WHC', SparkF.bround(soil_df.SM_FC - soil_df.SM_WP, 2))
    soil_df = soil_df.select(['IDREGION', 'SM_WHC'])

    return soil_df

  def preprocessAreaFractions(self, af_df, crop_id):
    """Filter area fractions data by crop id"""
    af_df = af_df.withColumn("FYEAR", af_df["FYEAR"].cast(SparkT.IntegerType()))
    af_df = af_df.filter(af_df.CROP_ID == crop_id).drop('CROP_ID')

    return af_df

  def preprocessYield(self, yield_df, crop_id):
    """
    Yield preprocessing depends on the data format.
    Here we cover preprocessing for France (NUTS3), Germany (NUTS3) and the Netherlands (NUTS2).
    """
    # Delete trailing empty columns
    empty_cols = [ c for c in yield_df.columns if c.startswith('_c') ]
    for c in empty_cols:
      yield_df = yield_df.drop(c)

    # Special case for Netherlands and Germany: convert yield columns into rows
    years = [int(c) for c in yield_df.columns if c[0].isdigit()]
    if (len(years) > 0):
      yield_by_year = yield_df.rdd.map(lambda x: (x[0], cropNameToID(crop_id_dict, x[0]), x[1],
                                                  [(years[i], x[i+2]) for i in range(len(years))]))

      yield_df = yield_by_year.toDF(['CROP', 'CROP_ID', 'IDREGION', 'YIELD'])
      yield_df = yield_df.withColumn('YR_YIELD', SparkF.explode('YIELD')).drop('YIELD')
      yield_by_year = yield_df.rdd.map(lambda x: (x[0], x[1], x[2], x[3][0], x[3][1]))
      yield_df = yield_by_year.toDF(['CROP', 'CROP_ID', 'IDREGION', 'FYEAR', 'YIELD'])
    else:
      yield_by_year = yield_df.rdd.map(lambda x: (x[0], cropNameToID(crop_id_dict, x[0]), x[1], x[2], x[3]))
      yield_df = yield_by_year.toDF(['CROP', 'CROP_ID', 'IDREGION', 'FYEAR', 'YIELD'])

    yield_df = yield_df.filter(yield_df.CROP_ID == crop_id).drop('CROP', 'CROP_ID')
    if (yield_df.count() == 0):
      return None

    yield_df = yield_df.filter(yield_df.YIELD.isNotNull())
    yield_df = yield_df.withColumn("YIELD", yield_df["YIELD"].cast(SparkT.FloatType()))
    yield_df = yield_df.filter(yield_df['YIELD'] > 0.0)

    return yield_df

  def preprocessYieldMCYFS(self, mcyfs_df, crop_id):
    """Preprocess MCYFS NUTS0 level yield predictions"""
    # the input columns are IDREGION, CROP, PREDICTION_DATE, FYEAR, YIELD_PRED
    mcyfs_df = mcyfs_df.withColumn('PRED_DEKAD', getDekad('PREDICTION_DATE'))
    # the columns should now be IDREGION, CROP, PREDICTION_DATE, FYEAR, YIELD_PRED, PRED_DEKAD
    yield_by_year = mcyfs_df.rdd.map(lambda x: (x[1], cropNameToID(crop_id_dict, x[1]),
                                                x[0], x[3], x[2], x[4], x[5]))
    mcyfs_df = yield_by_year.toDF(['CROP', 'CROP_ID', 'IDREGION', 'FYEAR',
                                         'PRED_DATE', 'YIELD_PRED', 'PRED_DEKAD'])
    mcyfs_df = mcyfs_df.filter(mcyfs_df.CROP_ID == crop_id).drop('CROP', 'CROP_ID')
    if (mcyfs_df.count() == 0):
      return None

    mcyfs_df = mcyfs_df.withColumn("YIELD_PRED", mcyfs_df["YIELD_PRED"].cast(SparkT.FloatType()))

    return mcyfs_df

"""#### Run Data Preprocessing"""

#%%writefile run_data_preprocessing.py
def printPreprocessingInformation(df, data_source, order_cols, crop_season=None):
  """Print preprocessed data and additional debug information"""
  df_regions = [reg[0] for reg in df.select('IDREGION').distinct().collect()]
  print(data_source , 'data available for', len(df_regions), 'region(s)')
  if (crop_season is not None):
    print('Season end information')
    crop_season.orderBy(['IDREGION', 'FYEAR']).show(10)

  print(data_source, 'data')
  df.orderBy(order_cols).show(10)

def preprocessData(cyp_config, cyp_preprocessor, data_dfs):
  crop_id = cyp_config.getCropID()
  nuts_level = cyp_config.getNUTSLevel()
  season_crosses_calyear = cyp_config.seasonCrossesCalendarYear()
  use_centroids = cyp_config.useCentroids()
  use_remote_sensing = cyp_config.useRemoteSensing()
  debug_level = cyp_config.getDebugLevel()

  order_cols = ['IDREGION', 'CAMPAIGN_YEAR', 'CAMPAIGN_DEKAD']
  # wofost data
  wofost_df = data_dfs['WOFOST']
  wofost_df = wofost_df.filter(wofost_df['CROP_ID'] == crop_id).drop('CROP_ID')
  crop_season = cyp_preprocessor.getCropSeasonInformation(wofost_df, season_crosses_calyear)
  wofost_df = cyp_preprocessor.preprocessWofost(wofost_df, crop_season, season_crosses_calyear)
  wofost_regions = [reg[0] for reg in wofost_df.select('IDREGION').distinct().collect()]
  data_dfs['WOFOST'] = wofost_df
  if (debug_level > 1):
    printPreprocessingInformation(wofost_df, 'WOFOST', order_cols, crop_season)

  # meteo data
  meteo_df = data_dfs['METEO']
  meteo_df = cyp_preprocessor.preprocessMeteo(meteo_df, crop_season, season_crosses_calyear)
  assert (meteo_df is not None)
  data_dfs['METEO'] = meteo_df
  if (debug_level > 1):
    printPreprocessingInformation(meteo_df, 'METEO', order_cols)

  # remote sensing data
  rs_df = None
  if (use_remote_sensing):
    rs_df = data_dfs['REMOTE_SENSING']
    rs_df = rs_df.drop('IDCOVER')

    # if other data is at NUTS3, convert rs_df to NUTS3 using parent region data
    if (nuts_level == 'NUTS3'):
      rs_df = cyp_preprocessor.remoteSensingNUTS2ToNUTS3(rs_df, wofost_regions)

    rs_df = cyp_preprocessor.preprocessRemoteSensing(rs_df, crop_season, season_crosses_calyear)
    assert (rs_df is not None)
    data_dfs['REMOTE_SENSING'] = rs_df
    if (debug_level > 1):
      printPreprocessingInformation(rs_df, 'REMOTE_SENSING', order_cols)

  order_cols = ['IDREGION']
  # centroids and distance to coast
  centroids_df = None
  if (use_centroids):
    centroids_df = data_dfs['CENTROIDS']
    centroids_df = cyp_preprocessor.preprocessCentroids(centroids_df)
    data_dfs['CENTROIDS'] = centroids_df
    if (debug_level > 1):
      printPreprocessingInformation(centroids_df, 'CENTROIDS', order_cols)

  # soil data
  soil_df = data_dfs['SOIL']
  soil_df = cyp_preprocessor.preprocessSoil(soil_df)
  data_dfs['SOIL'] = soil_df
  if (debug_level > 1):
    printPreprocessingInformation(soil_df, 'SOIL', order_cols)

  order_cols = ['IDREGION', 'FYEAR']
  # yield_data
  yield_df = data_dfs['YIELD']
  if (debug_level > 1):
    print('Yield before preprocessing')
    yield_df.show(10)

  yield_df = cyp_preprocessor.preprocessYield(yield_df, crop_id)
  assert (yield_df is not None)
  data_dfs['YIELD'] = yield_df
  if (debug_level > 1):
    print('Yield after preprocessing')
    yield_df.show(10)

  return data_dfs

"""### Training and Test Split

**Custom validation splits**

When yield trend is used, custom sliding validation split is used.

#### Training and Test Splitter Class
"""

#%%writefile train_test_sets.py
import numpy as np

class CYPTrainTestSplitter:
  def __init__(self, cyp_config):
    self.use_yield_trend = cyp_config.useYieldTrend()
    self.test_fraction = cyp_config.getTestFraction()
    self.verbose = cyp_config.getDebugLevel()

  def getTestYears(self, all_years, test_fraction=None, use_yield_trend=None):
    num_years = len(all_years)
    test_years = []
    if (test_fraction is None):
      test_fraction = self.test_fraction

    if (use_yield_trend is None):
      use_yield_trend = self.use_yield_trend

    if (use_yield_trend):
      # If test_year_start 15, years with index >= 15 are added to the test set
      test_year_start = num_years - np.floor(num_years * test_fraction).astype('int')
      test_years = all_years[test_year_start:]
    else:
      # If test_year_pos = 5, every 5th year is added to test set.
      # indices start with 0, so test_year_pos'th year has index (test_year_pos - 1)
      test_year_pos = np.floor(1/test_fraction).astype('int')
      test_years = all_years[test_year_pos - 1::test_year_pos]

    return test_years

  def trainTestSplit(self, yield_df, test_fraction=None, use_yield_trend=None):
    all_years = sorted([yr[0] for yr in yield_df.select('FYEAR').distinct().collect()])
    test_years = self.getTestYears(all_years, test_fraction, use_yield_trend)

    return test_years

  # Returns an array containings tuples (train_idxs, test_idxs) for each fold
  # NOTE Y_train should include IDREGION, FYEAR as first two columns.
  def customKFoldValidationSplit(self, Y_train_full, num_folds, log_fh=None):
    """
    Custom K-fold Validation Splits:
    When using yield trend, we cannot do k-fold cross-validation. The custom
    K-Fold validation splits data in time-ordered fashion. The test data
    always comes after the training data.
    """
    all_years = sorted(np.unique(Y_train_full[:, 1]))
    num_years = len(all_years)
    num_test_years = 1
    num_train_years = num_years - num_test_years * num_folds

    custom_cv = []
    custom_split_info = '\nCustom sliding validation train, test splits'
    custom_split_info += '\n----------------------------------------------'

    for k in range(num_folds):
      train_years_start = k * num_test_years
      test_years_start = train_years_start + num_train_years
      train_years = all_years[train_years_start:test_years_start]
      test_years = all_years[test_years_start:test_years_start + num_test_years]
      test_indexes = np.ravel(np.nonzero(np.isin(Y_train_full[:, 1], test_years)))
      train_indexes = np.ravel(np.nonzero(np.isin(Y_train_full[:, 1], train_years)))
      custom_cv.append(tuple((train_indexes, test_indexes)))

      train_years = [str(y) for y in train_years]
      test_years = [str(y) for y in test_years]
      custom_split_info += '\nValidation set ' + str(k + 1) + ' training years: ' + ', '.join(train_years)
      custom_split_info += '\nValidation set ' + str(k + 1) + ' test years: ' + ', '.join(test_years)

    custom_split_info += '\n'
    if (log_fh is not None):
      log_fh.write(custom_split_info)

    if (self.verbose > 1):
      print(custom_split_info)

    return custom_cv

"""#### Run Training and Test Split"""

#%%writefile run_train_test_split.py
def printTrainTestSplits(train_df, test_df, src, order_cols):
  """Print Training and Test Splits"""
  print('\n', src, 'training data')
  train_df.orderBy(order_cols).show(5)
  print('\n', src, 'test data')
  test_df.orderBy(order_cols).show(5)

# Training, Test Split
# --------------------
def splitTrainingTest(df, year_col, test_years):
  """Splitting given df into training and test dataframes."""
  train_df = df.filter(~df[year_col].isin(test_years))
  test_df = df.filter(df[year_col].isin(test_years))

  return [train_df, test_df]

def splitDataIntoTrainingTestSets(cyp_config, preprocessed_dfs, log_fh):
  """
  Split preprocessed data into training and test sets based on
  availability of yield data.
  """
  nuts_level = cyp_config.getNUTSLevel()
  use_centroids = cyp_config.useCentroids()
  use_remote_sensing = cyp_config.useRemoteSensing()
  debug_level = cyp_config.getDebugLevel()

  yield_df = preprocessed_dfs['YIELD']
  train_test_splitter = CYPTrainTestSplitter(cyp_config)
  test_years = train_test_splitter.trainTestSplit(yield_df)
  test_years_info = '\nTest years: ' + ', '.join([str(y) for y in sorted(test_years)]) + '\n'
  log_fh.write(test_years_info + '\n')
  print(test_years_info)

  # Times series data used for feature design.
  ts_data_sources = {
      'WOFOST' : preprocessed_dfs['WOFOST'],
      'METEO' : preprocessed_dfs['METEO'],
  }

  if (use_remote_sensing):
    ts_data_sources['REMOTE_SENSING'] = preprocessed_dfs['REMOTE_SENSING']

  train_test_dfs = {}
  for ts_src in ts_data_sources:
    train_test_dfs[ts_src] = splitTrainingTest(ts_data_sources[ts_src], 'CAMPAIGN_YEAR', test_years)

  # SOIL, CENTROIDS data are static.
  train_test_dfs['SOIL'] = [preprocessed_dfs['SOIL'], preprocessed_dfs['SOIL']]
  if (use_centroids):
    train_test_dfs['CENTROIDS'] = [preprocessed_dfs['CENTROIDS'],
                                   preprocessed_dfs['CENTROIDS']]

  # yield data
  train_test_dfs['YIELD'] = splitTrainingTest(yield_df, 'FYEAR', test_years)

  if (debug_level > 2):
    for src in train_test_dfs:
      if (src in ts_data_sources):
        order_cols = ['IDREGION', 'CAMPAIGN_YEAR', 'CAMPAIGN_DEKAD']
      elif ((src == 'YIELD') or (src == 'AREA_FRACTIONS')):
        order_cols = ['IDREGION', 'FYEAR']
      else:
        order_cols = ['IDREGION']

      train_df = train_test_dfs[src][0]
      test_df = train_test_dfs[src][1]
      printTrainTestSplits(train_df, test_df, src, order_cols)

  return train_test_dfs, test_years

"""### Data Summary

#### Data Summarizer Class
"""

#%%writefile data_summary.py
from pyspark.sql import Window
import functools

class CYPDataSummarizer:
  def __init__(self, cyp_config):
    self.verbose = cyp_config.getDebugLevel()

  def wofostDVSSummary(self, wofost_df, early_season_end=None):
    """
    Summary of crop calendar based on DVS.
    Early season end is relative to end of the season, hence a negative number.
    """
    join_cols = ['IDREGION', 'CAMPAIGN_YEAR']
    dvs_summary = wofost_df.select(join_cols).distinct()

    # We find the start and end dekads for DVS ranges
    my_window = Window.partitionBy(join_cols).orderBy('CAMPAIGN_DEKAD')

    wofost_df = wofost_df.withColumn('VALUE', wofost_df['DVS'])
    wofost_df = wofost_df.withColumn('PREV', SparkF.lag(wofost_df['VALUE']).over(my_window))
    wofost_df = wofost_df.withColumn('DIFF', SparkF.when(SparkF.isnull(wofost_df['PREV']), 0)\
                                     .otherwise(wofost_df['VALUE'] - wofost_df['PREV']))
    wofost_df = wofost_df.withColumn('SEASON_ALIGN', wofost_df['CAMPAIGN_DEKAD'] - wofost_df['DEKAD'])

    del_cols = ['VALUE', 'PREV', 'DIFF']
    dvs_summary = dvs_summary.join(wofost_df.filter(wofost_df['VALUE'] > 0.0).groupBy(join_cols)\
                                   .agg(SparkF.min('CAMPAIGN_DEKAD').alias('START_DVS')), join_cols)
    dvs_summary = dvs_summary.join(wofost_df.filter(wofost_df['DVS'] >= 100).groupBy(join_cols)\
                                   .agg(SparkF.min('CAMPAIGN_DEKAD').alias('START_DVS1')), join_cols)
    dvs_summary = dvs_summary.join(wofost_df.filter(wofost_df['DVS'] >= 200).groupBy(join_cols)\
                                   .agg(SparkF.min('CAMPAIGN_DEKAD').alias('START_DVS2')), join_cols)
    dvs_summary = dvs_summary.join(wofost_df.filter(wofost_df['DVS'] >= 200).groupBy(join_cols)\
                                   .agg(SparkF.min('DEKAD').alias('HARVEST')), join_cols)
    dvs_summary = dvs_summary.join(wofost_df.filter(wofost_df['DVS'] >= 200).groupBy(join_cols)\
                                   .agg(SparkF.max('SEASON_ALIGN').alias('SEASON_ALIGN')), join_cols)

    # Calendar year end season and early season dekads for comparing with MCYFS
    # Campaign year early season dekad to filter data during feature design
    dvs_summary = dvs_summary.withColumn('CALENDAR_END_SEASON', dvs_summary['HARVEST'] + 1)
    dvs_summary = dvs_summary.withColumn('CAMPAIGN_EARLY_SEASON',
                                         dvs_summary['CALENDAR_END_SEASON'] + dvs_summary['SEASON_ALIGN'])
    if (early_season_end is not None):
      dvs_summary = dvs_summary.withColumn('CALENDAR_EARLY_SEASON',
                                         dvs_summary['CALENDAR_END_SEASON'] + early_season_end)
      dvs_summary = dvs_summary.withColumn('CAMPAIGN_EARLY_SEASON',
                                           dvs_summary['CAMPAIGN_EARLY_SEASON'] + early_season_end)

    wofost_df = wofost_df.drop(*del_cols)
    dvs_summary = dvs_summary.drop('HARVEST', 'SEASON_ALIGN')

    return dvs_summary

  def indicatorsSummary(self, df, min_cols, max_cols, avg_cols):
    """long term min, max and avg values of selected indicators by region"""
    avgs = []
    if (avg_cols[1:]):
      avgs = [SparkF.bround(SparkF.avg(x), 2).alias('avg(' + x + ')') for x in avg_cols[1:]]
    
    if (min_cols[:1]):
      summary = df.select(min_cols).groupBy('IDREGION').min()
    else:
      summary = df.select(min_cols).groupBy('IDREGION')

    if (max_cols[1:]):
      summary = summary.join(df.select(max_cols).groupBy('IDREGION').max(), 'IDREGION')

    if (avgs):
      summary = summary.join(df.select(avg_cols).groupBy('IDREGION').agg(*avgs), 'IDREGION')
    return summary

  def yieldSummary(self, yield_df):
    """long term min, max and avg values of yield by region"""
    select_cols = ['IDREGION', 'YIELD']
    yield_summary = yield_df.select(select_cols).groupBy('IDREGION').min('YIELD')
    yield_summary = yield_summary.join(yield_df.select(select_cols).groupBy('IDREGION')\
                                       .agg(SparkF.max('YIELD')), 'IDREGION')
    yield_summary = yield_summary.join(yield_df.select(select_cols).groupBy('IDREGION')\
                                       .agg(SparkF.bround(SparkF.avg('YIELD'), 2)\
                                            .alias('avg(YIELD)')), 'IDREGION')
    return yield_summary

"""#### Run Data Summary"""

#%%writefile run_data_summary.py
def printDataSummary(df, data_source):
  """Print summary information"""
  if (data_source == 'WOFOST_DVS'):
    print('Crop calender information based on WOFOST data')
    max_year = df.select('CAMPAIGN_YEAR').agg(SparkF.max('CAMPAIGN_YEAR')).collect()[0][0]
    df.filter(df.CAMPAIGN_YEAR == max_year).orderBy('IDREGION').show(10)
  else:
    print(data_source, 'indicators summary')
    df.orderBy('IDREGION').show()

def getWOFOSTSummaryCols():
  """WOFOST columns used for data summary"""
  # only RSM has non-zero min values
  min_cols = ['IDREGION', 'RSM']
  max_cols = ['IDREGION'] + ['WLIM_YB', 'WLIM_YS', 'DVS',
                             'WLAI', 'RSM', 'TWC', 'TWR']
  # biomass and DVS values grow over time
  avg_cols = ['IDREGION', 'WLAI', 'RSM', 'TWC', 'TWR']

  return [min_cols, max_cols, avg_cols]

def getMeteoSummaryCols():
  """Meteo columns used for data summary"""
  col_names = ['TMAX', 'TMIN', 'TAVG', 'PREC', 'ET0', 'CWB', 'RAD']
  min_cols = ['IDREGION'] + col_names
  max_cols = ['IDREGION'] + col_names
  avg_cols = ['IDREGION'] + col_names

  return [min_cols, max_cols, avg_cols]

def getRemoteSensingSummaryCols():
  """Remote Sensing columns used for data summary"""
  col_names = ['FAPAR']
  min_cols = ['IDREGION'] + col_names
  max_cols = ['IDREGION'] + col_names
  avg_cols = ['IDREGION'] + col_names

  return [min_cols, max_cols, avg_cols]

def summarizeData(cyp_config, cyp_summarizer, train_test_dfs):
  """
  Summarize data. Create DVS summary to infer crop calendar.
  Summarize selected indicators for each data source.
  """
  wofost_train_df = train_test_dfs['WOFOST'][0]
  meteo_train_df = train_test_dfs['METEO'][0]
  yield_train_df = train_test_dfs['YIELD'][0]

  use_remote_sensing = cyp_config.useRemoteSensing()
  debug_level = cyp_config.getDebugLevel()
  early_season = cyp_config.earlySeasonPrediction()
  early_season_end = None
  if (early_season):
    early_season_end = cyp_config.getEarlySeasonEndDekad()

  # DVS summary (crop calendar)
  # NOTE this summary of crops based on wofost data should be used with caution
  # 1. The summary is per region per year.
  # 2. The summary is based on wofost simulations not real sowing and harvest dates
  dvs_summary = cyp_summarizer.wofostDVSSummary(wofost_train_df, early_season_end)
  dvs_summary = dvs_summary.drop('CALENDAR_END_SEASON', 'CALENDAR_EARLY_SEASON')
  if (debug_level > 1):
    printDataSummary(dvs_summary, 'WOFOST_DVS')

  summary_cols = {
      'WOFOST' : getWOFOSTSummaryCols(),
      'METEO' : getMeteoSummaryCols(),
  }

  summary_sources_dfs = {
      'WOFOST' : wofost_train_df,
      'METEO' : meteo_train_df,
  }

  if (use_remote_sensing):
    rs_train_df = train_test_dfs['REMOTE_SENSING'][0]
    summary_cols['REMOTE_SENSING'] = getRemoteSensingSummaryCols()
    summary_sources_dfs['REMOTE_SENSING'] = rs_train_df

  summary_dfs = {}
  for sum_src in summary_sources_dfs:
    summary_dfs[sum_src] = cyp_summarizer.indicatorsSummary(summary_sources_dfs[sum_src],
                                                            summary_cols[sum_src][0],
                                                            summary_cols[sum_src][1],
                                                            summary_cols[sum_src][2])

  for src in summary_dfs:
    if (debug_level > 2):
      printDataSummary(summary_dfs[src], src)

  yield_summary = cyp_summarizer.yieldSummary(yield_train_df)
  if (debug_level > 2):
    printDataSummary(yield_summary, 'YIELD')

  summary_dfs['WOFOST_DVS'] = dvs_summary
  summary_dfs['YIELD'] = yield_summary

  return summary_dfs

"""### Crop Calendar

We infer crop calendar using WOFOST DVS.
"""

#%%writefile crop_calendar.py
import numpy as np

def getCropCalendarPeriods(df):
  """Periods for per year crop calendar"""
  # (maximum of 4 months = 12 dekads).
  # Subtracting 11 because both ends of the period are included.
  # p0 : if CAMPAIGN_EARLY_SEASON > df.START_DVS
  #        START_DVS - 11 to START_DVS
  #      else
  #        START_DVS - 11 to CAMPAIGN_EARLY_SEASON
  p0_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > df.START_DVS,
                          (df.CAMPAIGN_DEKAD >= (df.START_DVS - 11)) &
                          (df.CAMPAIGN_DEKAD <= df.START_DVS))\
                          .otherwise((df.CAMPAIGN_DEKAD >= (df.START_DVS - 11)) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))
  # p1 : if CAMPAIGN_EARLY_SEASON > (df.START_DVS + 1)
  #        (START_DVS - 1) to (START_DVS + 1)
  #      else
  #        (START_DVS - 1) to CAMPAIGN_EARLY_SEASON
  p1_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > (df.START_DVS + 1),
                          (df.CAMPAIGN_DEKAD >= (df.START_DVS - 1)) &
                          (df.CAMPAIGN_DEKAD <= (df.START_DVS + 1)))\
                          .otherwise((df.CAMPAIGN_DEKAD >= (df.START_DVS - 1)) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))
  # p2 : if CAMPAIGN_EARLY_SEASON > df.START_DVS1
  #        START_DVS to START_DVS1
  #      else
  #        START_DVS to CAMPAIGN_EARLY_SEASON
  p2_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > df.START_DVS1,
                          (df.CAMPAIGN_DEKAD >= df.START_DVS) &
                          (df.CAMPAIGN_DEKAD <= df.START_DVS1))\
                          .otherwise((df.CAMPAIGN_DEKAD >= df.START_DVS) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))
  # p3 : if CAMPAIGN_EARLY_SEASON > (df.START_DVS1 + 1)
  #        (START_DVS1 - 1) to (START_DVS1 + 1)
  #      else
  #        (START_DVS1 - 1) to CAMPAIGN_EARLY_SEASON
  p3_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > (df.START_DVS1 + 1),
                          (df.CAMPAIGN_DEKAD >= (df.START_DVS1 - 1)) &
                          (df.CAMPAIGN_DEKAD <= (df.START_DVS1 + 1)))\
                          .otherwise((df.CAMPAIGN_DEKAD >= (df.START_DVS1 - 1)) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))
  # p4 : if CAMPAIGN_EARLY_SEASON > df.START_DVS2
  #        START_DVS1 to START_DVS2
  #      else
  #        START_DVS1 to CAMPAIGN_EARLY_SEASON
  p4_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > df.START_DVS2,
                          (df.CAMPAIGN_DEKAD >= df.START_DVS1) &
                          (df.CAMPAIGN_DEKAD <= df.START_DVS2))\
                          .otherwise((df.CAMPAIGN_DEKAD >= df.START_DVS1) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))
  # p5 : if CAMPAIGN_EARLY_SEASON > (df.START_DVS2 + 1)
  #        (START_DVS2 - 1) to (START_DVS2 + 1)
  #      else
  #        (START_DVS2 - 1) to CAMPAIGN_EARLY_SEASON
  p5_filter = SparkF.when(df.CAMPAIGN_EARLY_SEASON > (df.START_DVS2 + 1),
                          (df.CAMPAIGN_DEKAD >= (df.START_DVS2 - 1)) &
                          (df.CAMPAIGN_DEKAD <= (df.START_DVS2 + 1)))\
                          .otherwise((df.CAMPAIGN_DEKAD >= (df.START_DVS2 - 1)) &
                                     (df.CAMPAIGN_DEKAD <= df.CAMPAIGN_EARLY_SEASON))

  cc_periods = {
      'p0' : p0_filter,
      'p1' : p1_filter,
      'p2' : p2_filter,
      'p3' : p3_filter,
      'p4' : p4_filter,
      'p5' : p5_filter,
  }

  return cc_periods

def getCountryCropCalendar(crop_cal):
  """Take averages to make the crop calendar per country"""
  crop_cal = crop_cal.withColumn('COUNTRY', SparkF.substring('IDREGION', 1, 2))
  aggrs = [ SparkF.bround(SparkF.avg(crop_cal['START_DVS'])).alias('START_DVS'),
            SparkF.bround(SparkF.avg(crop_cal['START_DVS1'])).alias('START_DVS1'),
            SparkF.bround(SparkF.avg(crop_cal['START_DVS2'])).alias('START_DVS2'),
            SparkF.bround(SparkF.avg(crop_cal['CAMPAIGN_EARLY_SEASON'])).alias('CAMPAIGN_EARLY_SEASON') ]

  crop_cal = crop_cal.groupBy('COUNTRY').agg(*aggrs)
  return crop_cal

def getCropCalendar(cyp_config, dvs_summary, log_fh):
  """Use DVS summary to infer the crop calendar"""
  pd_dvs_summary = dvs_summary.toPandas()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  debug_level = cyp_config.getDebugLevel()

  avg_dvs_start = np.round(pd_dvs_summary['START_DVS'].mean(), 0)
  avg_dvs1_start = np.round(pd_dvs_summary['START_DVS1'].mean(), 0)
  avg_dvs2_start = np.round(pd_dvs_summary['START_DVS2'].mean(), 0)

  # We look at 6 windows
  # 0. Preplanting window (maximum of 4 months = 12 dekads).
  # Subtracting 11 because both ends of the period are included.
  p0_start = 1 if (avg_dvs_start - 11) < 1 else (avg_dvs_start - 11)
  p0_end = avg_dvs_start

  # 1. Planting window
  p1_start = avg_dvs_start - 1
  p1_end = avg_dvs_start + 1

  # 2. Vegetative phase
  p2_start = avg_dvs_start
  p2_end = avg_dvs1_start

  # 3. Flowering phase
  p3_start = avg_dvs1_start - 1
  p3_end = avg_dvs1_start + 1

  # 4. Yield formation phase
  p4_start = avg_dvs1_start
  p4_end = avg_dvs2_start

  # 5. Harvest window
  p5_start = avg_dvs2_start - 1
  p5_end = avg_dvs2_start + 1

  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = 36
  if (early_season_prediction):
    early_season_end = np.round(pd_dvs_summary['CAMPAIGN_EARLY_SEASON'].mean(), 0)
    p0_end = early_season_end if (p0_end > early_season_end) else p0_end
    p1_end = early_season_end if (p1_end > early_season_end) else p1_end
    p2_end = early_season_end if (p2_end > early_season_end) else p2_end
    p3_end = early_season_end if (p3_end > early_season_end) else p3_end
    p4_end = early_season_end if (p4_end > early_season_end) else p4_end
    p5_end = early_season_end if (p5_end > early_season_end) else p5_end

  crop_cal = {}
  if (p0_end > p0_start):
    crop_cal['p0'] = { 'desc' : 'pre-planting window', 'start' : p0_start, 'end' : p0_end }
  if (p1_end > p1_start):
    crop_cal['p1'] = { 'desc' : 'planting window', 'start' : p1_start, 'end' : p1_end }
  if (p2_end > p2_start):
    crop_cal['p2'] = { 'desc' : 'vegetative phase', 'start' : p2_start, 'end' : p2_end }
  if (p3_end > p3_start):
    crop_cal['p3'] = { 'desc' : 'flowering phase', 'start' : p3_start, 'end' : p3_end }
  if (p4_end > p4_start):
    crop_cal['p4'] = { 'desc' : 'yield formation phase', 'start' : p4_start, 'end' : p4_end }
  if (p5_end > p5_start):
    crop_cal['p5'] = { 'desc' : 'harvest window', 'start' : p5_start, 'end' : p5_end }

  if (early_season_prediction):
    early_season_rel_harvest = cyp_config.getEarlySeasonEndDekad()
    early_season_info = '\nEarly Season Prediction Dekad: ' + str(early_season_rel_harvest)
    early_season_info += ', Campaign Dekad: ' + str(early_season_end)
    log_fh.write(early_season_info + '\n')
    if (debug_level > 1):
      print(early_season_info)

  crop_cal_info = '\nCrop Calendar'
  crop_cal_info += '\n-------------'
  for p in crop_cal:
    crop_cal_info += '\nPeriod ' + p + ' (' + crop_cal[p]['desc'] + '): '
    crop_cal_info += 'Campaign Dekads ' + str(crop_cal[p]['start']) + '-' + str(crop_cal[p]['end'])

  log_fh.write(crop_cal_info + '\n')
  if (debug_level > 1):
    print(crop_cal_info)

  return crop_cal

"""### Feature Design

For WOFOST, Meteo and Remote Sensing features, we aggregate indicators or count days/dekads with indicators above or below some thresholds. We use 4 thresholds: +/- 1 STD and +/- 2STD above or below the average.

We determine the start and end dekads using crop calendar inferred from WOFOST DVS summary. In the case of early season prediction, end dekad is set to the prediction dekad.

#### Featurizer Class
"""

#%%writefile feature_design.py
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import functools

class CYPFeaturizer:
  def __init__(self, cyp_config):
    self.verbose = cyp_config.getDebugLevel()
    self.lt_stats = {}

  def extractFeatures(self, df, data_source, crop_cal,
                      max_cols, avg_cols, extreme_cols,
                      join_cols, fit=False):
    """
    Extract aggregate and extreme features.
    If fit=True, compute and save long-term stats.
    """
    df = df.withColumn('COUNTRY', SparkF.substring('IDREGION', 1, 2))
    crop_cal = getCountryCropCalendar(crop_cal)
    df = df.join(SparkF.broadcast(crop_cal), 'COUNTRY')

    cc_periods = getCropCalendarPeriods(df)
    aggrs = []
    # max aggregation
    for p in max_cols:
      if (max_cols[p]):
        aggrs += [SparkF.bround(SparkF.max(SparkF.when(cc_periods[p], df[x])), 2)\
                  .alias('max' + x + p) for x in max_cols[p] ]

    # avg aggregation
    for p in avg_cols:
      if (avg_cols[p]):
        aggrs += [SparkF.bround(SparkF.avg(SparkF.when(cc_periods[p], df[x])), 2)\
                  .alias('avg' + x + p) for x in avg_cols[p] ]

    # if not computing extreme features, we can return
    if (not extreme_cols):
      ft_df = df.groupBy(join_cols).agg(*aggrs)
      return ft_df

    # compute long-term stats and save them
    if (fit):
      stat_aggrs = []
      for p in extreme_cols:
        if (extreme_cols[p]):
          stat_aggrs += [ SparkF.bround(SparkF.avg(SparkF.when(cc_periods[p], df[x])), 2)\
                         .alias('avg' + x + p) for x in extreme_cols[p] ]
          stat_aggrs += [ SparkF.bround(SparkF.stddev(SparkF.when(cc_periods[p], df[x])), 2)\
                         .alias('std' + x + p) for x in extreme_cols[p] ]

      if (stat_aggrs):
        lt_stats = df.groupBy('COUNTRY').agg(*stat_aggrs)
        self.lt_stats[data_source] = lt_stats

    df = df.join(SparkF.broadcast(self.lt_stats[data_source]), 'COUNTRY')

    # features for extreme conditions
    for p in extreme_cols:
      if (extreme_cols[p]):
        # Count of days or dekads with values crossing threshold
        for i in range(1, 3):
          aggrs += [ SparkF.sum(SparkF.when((df[x] > (df['avg' + x + p] + i * df['std' + x + p])) &
                                            cc_periods[p], 1))\
                    .alias(x + p + 'gt' + str(i) + 'STD') for x in extreme_cols[p] ]
          aggrs += [ SparkF.sum(SparkF.when((df[x] < (df['avg' + x + p] - i * df['std' + x + p])) &
                                            cc_periods[p], 1))\
                    .alias(x + p + 'lt' + str(i) + 'STD') for x in extreme_cols[p] ]

    ft_df = df.groupBy(join_cols).agg(*aggrs)
    ft_df = ft_df.na.fill(0.0)
    return ft_df

"""#### Yield Trend Estimator Class"""

#%%writefile yield_trend.py
import numpy as np
import functools
from pyspark.sql import Window
from sklearn.metrics import mean_squared_error

class CYPYieldTrendEstimator:
  def __init__(self, cyp_config):
    self.verbose = cyp_config.getDebugLevel()
    self.trend_windows = cyp_config.getTrendWindows()

  def setTrendWindows(self, trend_windows):
    """Set trend window lengths"""
    assert isinstance(trend_windows, list)
    assert len(trend_windows) > 0
    assert isinstance(trend_windows[0], int)
    self.trend_windows = trend_windows

  def getTrendWindowYields(self, df, trend_window, reg_id=None):
    """Extract previous years' yield values to separate columns"""
    sel_cols = ['IDREGION', 'FYEAR', 'YIELD']
    my_window = Window.partitionBy('IDREGION').orderBy('FYEAR')

    yield_fts = df.select(sel_cols)
    if (reg_id is not None):
      yield_fts = yield_fts.filter(yield_fts.IDREGION == reg_id)

    for i in range(trend_window):
      yield_fts = yield_fts.withColumn('YIELD-' + str(i+1),
                                       SparkF.lag(yield_fts.YIELD, i+1).over(my_window))
      yield_fts = yield_fts.withColumn('YEAR-' + str(i+1),
                                       SparkF.lag(yield_fts.FYEAR, i+1).over(my_window))

    # drop columns withs null values
    for i in range(trend_window):
      yield_fts = yield_fts.filter(SparkF.col('YIELD-' + str(i+1)).isNotNull())

    prev_yields = [ 'YIELD-' + str(i) for i in range(trend_window, 0, -1)]
    prev_years = [ 'YEAR-' + str(i) for i in range(trend_window, 0, -1)]
    sel_cols = ['IDREGION', 'FYEAR'] + prev_years + prev_yields
    yield_fts = yield_fts.select(sel_cols)

    return yield_fts

  # Christos, Ante's suggestion
  # - To avoid overfitting, trend estimation could use a window which skips a year in between
  # So a window of 3 will use 6 years of data
  def printYieldTrendRounds(self, df, reg_id, trend_windows=None):
    """Print the sequence of years used for yield trend estimation"""
    reg_years = sorted([yr[0] for yr in df.filter(df['IDREGION'] == reg_id).select('FYEAR').distinct().collect()])
    num_years = len(reg_years)
    if (trend_windows is None):
      trend_windows = self.trend_windows

    for trend_window in trend_windows:
      rounds = (num_years - trend_window)
      if ((self.verbose > 2) and (trend_window == trend_windows[0])):
        print('Trend window', trend_window)
    
      for rd in range(rounds):
        test_year = reg_years[-(rd + 1)]
        start_year = reg_years[-(rd + trend_window + 1)]
        end_year = reg_years[-(rd + 2)]

        if ((self.verbose > 2) and (trend_window == trend_windows[0])):
          print('Round:', rd, 'Test year:', test_year,
                'Trend Window:', start_year, '-', end_year)

  def getLinearYieldTrend(self, window_years, window_yields, pred_year):
    """Linear yield trend prediction"""
    coefs = np.polyfit(window_years, window_yields, 1)
    return float(np.round(coefs[0] * pred_year + coefs[1], 2))

  def getFixedWindowTrendFeatures(self, df, trend_window=None, pred_year=None):
    """Predict the yield trend for each IDREGION and FYEAR using a fixed window"""
    join_cols = ['IDREGION', 'FYEAR']
    if (trend_window is None):
      trend_window = self.trend_windows[0]

    yield_ft_df = self.getTrendWindowYields(df, trend_window)
    if (pred_year is not None):
      yield_ft_df = yield_ft_df.filter(yield_ft_df['FYEAR'] == pred_year)

    pd_yield_ft_df = yield_ft_df.toPandas()
    region_years = pd_yield_ft_df[join_cols].values
    prev_year_cols = ['YEAR-' + str(i) for i in range(1, trend_window + 1)]
    prev_yield_cols = ['YIELD-' + str(i) for i in range(1, trend_window + 1)]
    window_years = pd_yield_ft_df[prev_year_cols].values
    window_yields = pd_yield_ft_df[prev_yield_cols].values

    yield_trend = []
    for i in range(region_years.shape[0]):
      yield_trend.append(self.getLinearYieldTrend(window_years[i, :],
                                                  window_yields[i, :],
                                                  region_years[i, 1]))

    pd_yield_ft_df['YIELD_TREND'] = yield_trend
    return pd_yield_ft_df

  def getFixedWindowTrend(self, df, reg_id, pred_year, trend_window=None):
    """
    Return linear trend prediction for given region and year
    using fixed trend window.
    """
    if (trend_window is None):
      trend_window = self.trend_windows[0]

    reg_df = df.filter((df['IDREGION'] == reg_id) & (df['FYEAR'] <= pred_year))
    pd_yield_ft_df = self.getFixedWindowTrendFeatures(reg_df, trend_window, pred_year)
    if (len(pd_yield_ft_df.index) == 0):
      print('No data to estimate yield trend')
      return None

    if (len(pd_yield_ft_df.index) == 0):
      return None

    reg_year_filter = (df['IDREGION'] == reg_id) & (df['FYEAR'] == pred_year)
    pd_yield_ft_df['ACTUAL'] = df.filter(reg_year_filter).select('YIELD').collect()[0][0]
    pd_yield_ft_df = pd_yield_ft_df.rename(columns={'YIELD_TREND' : 'PREDICTED'})

    return pd_yield_ft_df

  def getL1OutCVPredictions(self, pd_yield_ft_df, trend_window, join_cols, iter):
    """1 iteration of leave-one-out cross-validation"""
    iter_year_cols = ['YEAR-' + str(i) for i in range(1, trend_window + 1) if i != iter]
    iter_yield_cols = ['YIELD-' + str(i) for i in range(1, trend_window + 1) if i != iter]
    window_years = pd_yield_ft_df[iter_year_cols].values
    window_yields = pd_yield_ft_df[iter_yield_cols].values

    # We are going to predict yield value for YEAR-<iter>.
    pred_years = pd_yield_ft_df['YEAR-' + str(iter)].values
    predicted_trend = []
    for i in range(pred_years.shape[0]):
      predicted_trend.append(self.getLinearYieldTrend(window_years[i, :],
                                                      window_yields[i, :],
                                                      pred_years[i]))

    pd_iter_preds = pd_yield_ft_df[join_cols]
    pd_iter_preds['YTRUE' + str(iter)] = pd_yield_ft_df['YIELD-' + str(iter)]
    pd_iter_preds['YPRED' + str(iter)] = predicted_trend

    if (self.verbose > 2):
      print('Leave-one-out cross-validation: iteration', iter)
      print(pd_iter_preds.head(5))

    return pd_iter_preds

  def getL1outRMSE(self, cv_actual, cv_predicted):
    """Compute RMSE for leave-one-out predictions"""
    return float(np.round(np.sqrt(mean_squared_error(cv_actual, cv_predicted)), 2))

  def getMinRMSEIndex(self, cv_rmses):
    """Index of min rmse values"""
    return np.nanargmin(cv_rmses)

  def getL1OutCVRMSE(self, df, trend_window, join_cols, pred_year=None):
    """Run leave-one-out cross-validation and compute RMSE"""
    join_cols = ['IDREGION', 'FYEAR']
    pd_yield_ft_df = self.getFixedWindowTrendFeatures(df, trend_window, pred_year)
    pd_l1out_preds = None
    for i in range(1, trend_window + 1):
      pd_iter_preds = self.getL1OutCVPredictions(pd_yield_ft_df, trend_window,
                                                 join_cols, i)
      if (pd_l1out_preds is None):
        pd_l1out_preds = pd_iter_preds
      else:
        pd_l1out_preds = pd_l1out_preds.merge(pd_iter_preds, on=join_cols)

    region_years = pd_l1out_preds[join_cols].values
    ytrue_cols = ['YTRUE' + str(i) for i in range(1, trend_window + 1)]
    ypred_cols = ['YPRED' + str(i) for i in range(1, trend_window + 1)]
    l1out_ytrue = pd_l1out_preds[ytrue_cols].values
    l1out_ypred = pd_l1out_preds[ypred_cols].values
    cv_rmse = []
    for i in range(region_years.shape[0]):
      cv_rmse.append(self.getL1outRMSE(l1out_ytrue[i, :],
                                       l1out_ypred[i, :]))

    pd_l1out_rmse = pd_yield_ft_df[join_cols]
    pd_l1out_rmse['YIELD_TREND' + str(trend_window)] = pd_yield_ft_df['YIELD_TREND']
    pd_l1out_rmse['CV_RMSE' + str(trend_window)] = cv_rmse

    return pd_l1out_rmse

  def getOptimalTrendWindows(self, df, pred_year=None, trend_windows=None):
    """
    Compute optimal yield trend values based on leave-one-out
    cross validation errors for different trend windows.
    """
    join_cols = ['IDREGION', 'FYEAR']
    if (trend_windows is None):
      trend_windows = self.trend_windows

    pd_tw_rmses = None
    for tw in trend_windows:
      pd_l1out_rmse = self.getL1OutCVRMSE(df, tw, join_cols, pred_year)
      if (pd_tw_rmses is None):
        pd_tw_rmses = pd_l1out_rmse
      else:
        pd_tw_rmses = pd_tw_rmses.merge(pd_l1out_rmse, on=join_cols, how='left')

    if (self.verbose > 2):
      print('Leave-one-out cross-validation: RMSE')
      print(pd_tw_rmses.sort_values(by=join_cols).head(5))

    region_years = pd_tw_rmses[join_cols].values
    tw_rmse_cols = ['CV_RMSE' + str(tw) for tw in trend_windows]
    tw_trend_cols = ['YIELD_TREND' + str(tw) for tw in trend_windows]
    tw_cv_rmses = pd_tw_rmses[tw_rmse_cols].values
    tw_yield_trend = pd_tw_rmses[tw_trend_cols].values

    opt_windows = []
    yield_trend_preds = []
    for i in range(region_years.shape[0]):
      min_rmse_index = self.getMinRMSEIndex(tw_cv_rmses[i, :])
      opt_windows.append(trend_windows[min_rmse_index])
      yield_trend_preds.append(tw_yield_trend[i, min_rmse_index])

    pd_opt_win_df = pd_tw_rmses[join_cols]
    pd_opt_win_df['OPT_TW'] = opt_windows
    pd_opt_win_df['YIELD_TREND'] = yield_trend_preds
    if (self.verbose > 2):
      print('Optimal trend windows')
      print(pd_opt_win_df.sort_values(by=join_cols).head(5))

    return pd_opt_win_df

  def getOptimalWindowTrendFeatures(self, df, trend_windows=None):
    """
    Get previous year yield values and predicted yield trend
    by determining optimal trend window for each region and year.
    NOTE: We have to select the same number of features, so we
    select previous trend_windows[0] yield values.
    """
    join_cols = ['IDREGION', 'FYEAR']
    if (trend_windows is None):
      trend_windows = self.trend_windows

    pd_yield_ft_df = self.getTrendWindowYields(df, trend_windows[0]).toPandas()
    pd_opt_win_df = self.getOptimalTrendWindows(df, trend_windows=trend_windows)
    pd_opt_win_df = pd_opt_win_df.drop(columns=['OPT_TW'])
    pd_yield_ft_df = pd_yield_ft_df.merge(pd_opt_win_df, on=join_cols)

    return pd_yield_ft_df

  def getOptimalWindowTrend(self, df, reg_id, pred_year, trend_windows=None):
    """
    Compute the optimal trend window for given region and year based on
    leave-one-out cross validation errors for different trend windows.
    """
    df = df.filter(df['IDREGION'] == reg_id)
    pd_opt_win_df = self.getOptimalTrendWindows(df, pred_year, trend_windows)
    if (len(pd_opt_win_df.index) == 0):
      return None

    reg_year_filter = (df['IDREGION'] == reg_id) & (df['FYEAR'] == pred_year)
    pd_opt_win_df['ACTUAL'] = df.filter(reg_year_filter).select('YIELD').collect()[0][0]
    pd_opt_win_df = pd_opt_win_df.rename(columns={'YIELD_TREND' : 'PREDICTED'})

    return pd_opt_win_df

"""#### Create WOFOST, Meteo and Remote Sensing Features"""

#%%writefile run_feature_design.py
def wofostMaxFeatureCols():
  """columns or indicators using max aggregation"""
  # must be in sync with crop calendar periods
  max_cols = {
      'p0' : [],
      'p1' : [],
      'p2' : ['WLIM_YB', 'TWC', 'WLAI'],
      'p3' : [],
      'p4' : ['WLIM_YB', 'WLIM_YS', 'TWC', 'WLAI'],
      'p5' : [],
  }

  return max_cols

def wofostAvgFeatureCols():
  """columns or indicators using avg aggregation"""
  # must be in sync with crop calendar periods
  avg_cols = {
      'p0' : [],
      'p1' : [],
      'p2' : ['RSM'],
      'p3' : [],
      'p4' : ['RSM'],
      'p5' : [],
  }

  return avg_cols

def wofostCountFeatureCols():
  """columns or indicators using count aggregation"""
  # must be in sync with crop calendar periods
  count_cols = {
      'p0' : [],
      'p1' : ['RSM'],
      'p2' : ['RSM'],
      'p3' : ['RSM'],
      'p4' : ['RSM'],
      'p5' : [],
  }

  return count_cols

# Meteo Feature ideas:
# Two dry summers caused drop in ground water level:
#   rainfall sums going back to second of half of previous year
# Previous year: high production, prices low, invest less in crop
#   Focus on another crop
def meteoMaxFeatureCols():
  """columns or indicators using max aggregation"""
  # must be in sync with crop calendar periods
  max_cols = { 'p0' : [], 'p1' : [], 'p2' : [], 'p3' : [], 'p4' : [], 'p5' : [] }

  return max_cols

def meteoAvgFeatureCols():
  """columns or indicators using avg aggregation"""
  # must be in sync with crop calendar periods
  avg_cols = {
      'p0' : ['TAVG', 'PREC', 'CWB'],
      'p1' : ['TAVG', 'PREC'],
      'p2' : ['TAVG', 'CWB'],
      'p3' : ['PREC'],
      'p4' : ['CWB'],
      'p5' : ['PREC'],
  }

  return avg_cols

def meteoCountFeatureCols():
  """columns or indicators using count aggregation"""
  # must be in sync with crop calendar periods
  count_cols = {
      'p0' : [],
      'p1' : ['TMIN', 'PREC'],
      'p2' : [],
      'p3' : ['PREC', 'TMAX'],
      'p4' : [],
      'p5' : ['PREC'],
  }

  return count_cols

def rsMaxFeatureCols():
  """columns or indicators using max aggregation"""
  # must be in sync with crop calendar periods
  max_cols = { 'p0' : [], 'p1' : [], 'p2' : [], 'p3' : [], 'p4' : [], 'p5' : [] }

  return max_cols

def rsAvgFeatureCols():
  """columns or indicators using avg aggregation"""
  # must be in sync with crop calendar periods
  avg_cols = {
      'p0' : [],
      'p1' : [],
      'p2' : ['FAPAR'],
      'p3' : [],
      'p4' : ['FAPAR'],
      'p5' : [],
  }

  return avg_cols

def convertFeaturesToPandas(ft_dfs, join_cols):
  """Convert features to pandas and merge"""
  train_ft_df = ft_dfs[0]
  test_ft_df = ft_dfs[1]
  train_ft_df = train_ft_df.withColumnRenamed('CAMPAIGN_YEAR', 'FYEAR')
  test_ft_df = test_ft_df.withColumnRenamed('CAMPAIGN_YEAR', 'FYEAR')
  pd_train_df = train_ft_df.toPandas()
  pd_test_df = test_ft_df.toPandas()

  return [pd_train_df, pd_test_df]

def dropZeroColumns(pd_ft_dfs):
  """Drop columns which have all zeros in training data"""
  pd_train_df = pd_ft_dfs[0]
  pd_test_df = pd_ft_dfs[1]

  pd_train_df = pd_train_df.loc[:, (pd_train_df != 0.0).any(axis=0)]
  pd_train_df = pd_train_df.dropna(axis=1)
  pd_test_df = pd_test_df[pd_train_df.columns]

  return [pd_train_df, pd_test_df]

def printFeatureData(pd_feature_dfs, join_cols):
  for src in pd_feature_dfs:
    pd_train_fts = pd_feature_dfs[src][0]
    if (pd_train_fts is None):
      continue

    pd_test_fts = pd_feature_dfs[src][1]
    all_cols = list(pd_train_fts.columns)
    aggr_cols = [ c for c in all_cols if (('avg' in c) or ('max' in c))]
    if (len(aggr_cols) > 0):
      print('\n', src, 'Aggregate Features: Training')
      print(pd_train_fts[join_cols + aggr_cols].head(5))
      print('\n', src, 'Aggregate Features: Test')
      print(pd_test_fts[join_cols + aggr_cols].head(5))

    ext_cols = [ c for c in all_cols if (('Z+' in c) or ('Z-' in c) or
                                         ('lt' in c) or ('gt' in c))]
    if (len(ext_cols) > 0):
      print('\n', src, 'Features for Extreme Conditions: Training')
      print(pd_train_fts[join_cols + ext_cols].head(5))
      print('\n', src, 'Features for Extreme Conditions: Test')
      print(pd_test_fts[join_cols + ext_cols].head(5))

def createFeatures(cyp_config, cyp_featurizer, train_test_dfs,
                   summary_dfs, log_fh):
  """Create WOFOST, Meteo and Remote Sensing features"""
  nuts_level = cyp_config.getNUTSLevel()
  use_remote_sensing = cyp_config.useRemoteSensing()
  debug_level = cyp_config.getDebugLevel()

  wofost_train_df = train_test_dfs['WOFOST'][0]
  wofost_test_df = train_test_dfs['WOFOST'][1]
  meteo_train_df = train_test_dfs['METEO'][0]
  meteo_test_df = train_test_dfs['METEO'][1]
  yield_train_df = train_test_dfs['YIELD'][0]

  dvs_summary = summary_dfs['WOFOST_DVS']

  # Filter out years before min yield year
  yield_min_year = yield_train_df.agg({"FYEAR": "MIN"}).collect()[0][0]
  print('Yield min year', yield_min_year)

  wofost_train_df = wofost_train_df.filter(wofost_train_df.CAMPAIGN_YEAR >= yield_min_year)
  meteo_train_df = meteo_train_df.filter(meteo_train_df.CAMPAIGN_YEAR >= yield_min_year)
  dvs_summary = dvs_summary.filter(dvs_summary.CAMPAIGN_YEAR >= yield_min_year)

  rs_train_df = None
  rs_test_df = None
  if (use_remote_sensing):
    rs_train_df = train_test_dfs['REMOTE_SENSING'][0]
    rs_test_df = train_test_dfs['REMOTE_SENSING'][1]
    rs_train_df = rs_train_df.filter(rs_train_df.CAMPAIGN_YEAR >= yield_min_year)
    rs_test_df = rs_test_df.filter(rs_test_df.CAMPAIGN_YEAR >= yield_min_year)

  join_cols = ['IDREGION', 'CAMPAIGN_YEAR']
  aggr_ft_cols = {
      'WOFOST' : [wofostMaxFeatureCols(), wofostAvgFeatureCols()],
      'METEO' : [meteoMaxFeatureCols(), meteoAvgFeatureCols()],
  }

  count_ft_cols = {
      'WOFOST' : wofostCountFeatureCols(),
      'METEO' : meteoCountFeatureCols(),
  }

  train_ft_src_dfs = {
      'WOFOST' : wofost_train_df,
      'METEO' : meteo_train_df,
  }

  test_ft_src_dfs = {
      'WOFOST' : wofost_test_df,
      'METEO' : meteo_test_df,
  }

  if (use_remote_sensing):
    train_ft_src_dfs['REMOTE_SENSING'] = rs_train_df
    test_ft_src_dfs['REMOTE_SENSING'] = rs_test_df
    aggr_ft_cols['REMOTE_SENSING'] = [rsMaxFeatureCols(), rsAvgFeatureCols()]
    count_ft_cols['REMOTE_SENSING'] = {}

  crop_cal_train = dvs_summary
  crop_cal_test = dvs_summary

  train_ft_dfs = {}
  test_ft_dfs = {}
  for ft_src in train_ft_src_dfs:
    train_ft_dfs[ft_src] = cyp_featurizer.extractFeatures(train_ft_src_dfs[ft_src],
                                                          ft_src,
                                                          crop_cal_train,
                                                          aggr_ft_cols[ft_src][0],
                                                          aggr_ft_cols[ft_src][1],
                                                          count_ft_cols[ft_src],
                                                          join_cols,
                                                          True)
    test_ft_dfs[ft_src] = cyp_featurizer.extractFeatures(test_ft_src_dfs[ft_src],
                                                         ft_src,
                                                         crop_cal_test,
                                                         aggr_ft_cols[ft_src][0],
                                                         aggr_ft_cols[ft_src][1],
                                                         count_ft_cols[ft_src],
                                                         join_cols)

  pd_conversion_dict = {
      'WOFOST' : [ train_ft_dfs['WOFOST'], test_ft_dfs['WOFOST'] ],
      'METEO' : [ train_ft_dfs['METEO'], test_ft_dfs['METEO'] ],
  }

  if (use_remote_sensing):
      pd_conversion_dict['REMOTE_SENSING'] = [ train_ft_dfs['REMOTE_SENSING'], test_ft_dfs['REMOTE_SENSING'] ]

  pd_feature_dfs = {}
  for ft_src in pd_conversion_dict:
    pd_feature_dfs[ft_src] = convertFeaturesToPandas(pd_conversion_dict[ft_src], join_cols)

  # Check and drop features with all zeros (possible in early season prediction).
  for ft_src in pd_feature_dfs:
    pd_feature_dfs[ft_src] = dropZeroColumns(pd_feature_dfs[ft_src])

  if (debug_level > 1):
    join_cols = ['IDREGION', 'FYEAR']
    printFeatureData(pd_feature_dfs, join_cols)

  return pd_feature_dfs

"""#### Create Yield Trend Features"""

#%%writefile run_trend_feature_design.py
def createYieldTrendFeatures(cyp_config, cyp_trend_est,
                             yield_train_df, yield_test_df, test_years):
  """Create yield trend features"""
  join_cols = ['IDREGION', 'FYEAR']
  find_optimal = cyp_config.findOptimalTrendWindow()
  trend_window = cyp_config.getTrendWindows()[0]
  debug_level = cyp_config.getDebugLevel()
  yield_df = yield_train_df.union(yield_test_df.select(yield_train_df.columns))

  if (find_optimal):
    pd_train_features = cyp_trend_est.getOptimalWindowTrendFeatures(yield_train_df)
    pd_test_features = cyp_trend_est.getOptimalWindowTrendFeatures(yield_df)
  else:
    pd_train_features = cyp_trend_est.getFixedWindowTrendFeatures(yield_train_df)
    pd_test_features = cyp_trend_est.getFixedWindowTrendFeatures(yield_df)

  pd_test_features = pd_test_features[pd_test_features['FYEAR'].isin(test_years)]
  prev_year_cols = ['YEAR-' + str(i) for i in range(1, trend_window + 1)]
  pd_train_features = pd_train_features.drop(columns=prev_year_cols)
  pd_test_features = pd_test_features.drop(columns=prev_year_cols)

  if (debug_level > 1):
    print('\nYield Trend Features: Train')
    join_cols = ['IDREGION', 'FYEAR']
    print(pd_train_features.sort_values(by=join_cols).head(5))
    print('Total', len(pd_train_features.index), 'rows')
    print('\nYield Trend Features: Test')
    print(pd_test_features.sort_values(by=join_cols).head(5))
    print('Total', len(pd_test_features.index), 'rows')

  return pd_train_features, pd_test_features

"""### Combine features and labels

Combine wofost, meteo and soil with remote sensing. Combine with centroids or yield trend both if configured. Combine with yield data in the end.
"""

#%%writefile combine_features.py
def combineFeaturesLabels(cyp_config, sqlCtx,
                          prep_train_test_dfs, pd_feature_dfs,
                          join_cols, log_fh):
  """
  Combine wofost, meteo and soil with remote sensing. Combine centroids
  and yield trend if configured. Combine with yield data in the end.
  If configured, save features to a CSV file.
  """
  pd_soil_df = prep_train_test_dfs['SOIL'][0].toPandas()
  pd_yield_train_df = prep_train_test_dfs['YIELD'][0].toPandas()
  pd_yield_test_df = prep_train_test_dfs['YIELD'][1].toPandas()

  # Feature dataframes have already been converted to pandas
  pd_wofost_train_ft = pd_feature_dfs['WOFOST'][0]
  pd_wofost_test_ft = pd_feature_dfs['WOFOST'][1]
  pd_meteo_train_ft = pd_feature_dfs['METEO'][0]
  pd_meteo_test_ft = pd_feature_dfs['METEO'][1]

  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  use_yield_trend = cyp_config.useYieldTrend()
  use_centroids = cyp_config.useCentroids()
  use_remote_sensing = cyp_config.useRemoteSensing()
  save_features = cyp_config.saveFeatures()
  debug_level = cyp_config.getDebugLevel()
  
  combine_info = '\nCombine Features and Labels'
  combine_info += '\n---------------------------'
  yield_min_year = pd_yield_train_df['FYEAR'].min()
  combine_info += '\nYield min year ' + str(yield_min_year) + '\n'

  # start with static SOIL data
  pd_train_df = pd_soil_df.copy(deep=True)
  pd_test_df = pd_soil_df.copy(deep=True)
  combine_info += '\nData size after including SOIL data: '
  combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
  combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  if (use_centroids):
    # combine with region centroids
    pd_centroids_df = prep_train_test_dfs['CENTROIDS'][0].toPandas()
    pd_train_df = pd_train_df.merge(pd_centroids_df, on=['IDREGION'])
    pd_test_df = pd_test_df.merge(pd_centroids_df, on='IDREGION')
    combine_info += '\nData size after including CENTROIDS data: '
    combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
    combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  # combine with WOFOST features
  static_cols = list(pd_train_df.columns)
  pd_train_df = pd_train_df.merge(pd_wofost_train_ft, on=['IDREGION'])
  pd_test_df = pd_test_df.merge(pd_wofost_test_ft, on=['IDREGION'])
  wofost_cols = list(pd_wofost_train_ft.columns)
  col_order = ['IDREGION', 'FYEAR'] + static_cols[1:] + wofost_cols[2:]
  pd_train_df = pd_train_df[col_order]
  pd_test_df = pd_test_df[col_order]
  combine_info += '\nData size after including WOFOST features: '
  combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
  combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  # combine with METEO features
  pd_train_df = pd_train_df.merge(pd_meteo_train_ft, on=join_cols)
  pd_test_df = pd_test_df.merge(pd_meteo_test_ft, on=join_cols)
  combine_info += '\nData size after including METEO features: '
  combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
  combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  # combine with remote sensing features
  if (use_remote_sensing):
    pd_rs_train_ft = pd_feature_dfs['REMOTE_SENSING'][0]
    pd_rs_test_ft = pd_feature_dfs['REMOTE_SENSING'][1]

    pd_train_df = pd_train_df.merge(pd_rs_train_ft, on=join_cols)
    pd_test_df = pd_test_df.merge(pd_rs_test_ft, on=join_cols)
    combine_info += '\nData size after including REMOTE_SENSING features: '
    combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
    combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  if (use_yield_trend):
    # combine with yield trend features
    pd_yield_trend_train_ft = pd_feature_dfs['YIELD_TREND'][0]
    pd_yield_trend_test_ft = pd_feature_dfs['YIELD_TREND'][1]
    pd_train_df = pd_train_df.merge(pd_yield_trend_train_ft, on=join_cols)
    pd_test_df = pd_test_df.merge(pd_yield_trend_test_ft, on=join_cols)
    combine_info += '\nData size after including yield trend features: '
    combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
    combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  # combine with yield data
  pd_train_df = pd_train_df.merge(pd_yield_train_df, on=join_cols)
  pd_test_df = pd_test_df.merge(pd_yield_test_df, on=join_cols)
  pd_train_df = pd_train_df.sort_values(by=join_cols)
  pd_test_df = pd_test_df.sort_values(by=join_cols)
  combine_info += '\nData size after including yield (label) data: '
  combine_info += '\nTrain ' + str(len(pd_train_df.index)) + ' rows.'
  combine_info += '\nTest ' + str(len(pd_test_df.index)) + ' rows.\n'

  log_fh.write(combine_info + '\n')
  if (debug_level > 1):
    print(combine_info)
    print('\nAll Features and labels: Training')
    print(pd_train_df.head(5))
    print('\nAll Features and labels: Test')
    print(pd_test_df.head(5))

  if (save_features):
    early_season_prediction = cyp_config.earlySeasonPrediction()
    early_season_end = cyp_config.getEarlySeasonEndDekad()
    feature_file_path = cyp_config.getOutputPath()
    features_file = getFeatureFilename(crop, country, use_yield_trend,
                                       early_season_prediction, early_season_end)
    save_ft_path = feature_file_path + '/' + features_file
    save_ft_info = '\nSaving features to: ' + save_ft_path + '[train, test].csv'
    log_fh.write(save_ft_info + '\n')
    if (debug_level > 1):
      print(save_ft_info)

    pd_train_df.to_csv(save_ft_path + '_train.csv', index=False, header=True)
    pd_test_df.to_csv(save_ft_path + '_test.csv', index=False, header=True)

    # NOTE: In some environments, Spark can write, but pandas cannot.
    # In such cases, use the following code.
    # spark_train_df = sqlCtx.createDataFrame(pd_train_df)
    # spark_train_df.coalesce(1).write.option('header','true').mode('overwrite').csv(save_ft_path + '_train')
    # spark_test_df = sqlCtx.createDataFrame(pd_test_df)
    # spark_test_df.coalesce(1).write.option('header','true').mode('overwrite').csv(save_ft_path + '_test')

  return pd_train_df, pd_test_df

"""### Load Saved Features and Labels"""

#%%writefile load_saved_features.py
import pandas as pd

def loadSavedFeaturesLabels(cyp_config, spark):
  """Load saved features from a CSV file"""
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()
  debug_level = cyp_config.getDebugLevel()

  feature_file_path = cyp_config.getOutputPath()
  feature_file = getFeatureFilename(crop, country, use_yield_trend,
                                    early_season_prediction, early_season_end)

  load_ft_path = feature_file_path + '/' + feature_file
  pd_train_df = pd.read_csv(load_ft_path + '_train.csv', header=0)
  pd_test_df = pd.read_csv(load_ft_path + '_test.csv', header=0)

  # NOTE: In some environments, Spark can read, but pandas cannot.
  # In such cases, use the following code.
  # spark_train_df = spark.read.csv(load_ft_path + '_train.csv', header=True, inferSchema=True)
  # spark_test_df = spark.read.csv(load_ft_path + '_test.csv', header=True, inferSchema=True)
  # pd_train_df = spark_train_df.toPandas()
  # pd_test_df = spark_test_df.toPandas()

  if (debug_level > 1):
    print('\nAll Features and labels')
    print(pd_train_df.head(5))
    print(pd_test_df.head(5))

  return pd_train_df, pd_test_df

"""### Machine Learning using scikit learn

#### Feature Selector Class
"""

#%%writefile feature_selection.py
import numpy as np
import pandas as pd

from sklearn.pipeline import Pipeline
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import SelectFromModel
from sklearn.feature_selection import RFE
from sklearn.model_selection import GridSearchCV
from sklearn.utils import parallel_backend

class CYPFeatureSelector:
  def __init__(self, cyp_config, X_train, Y_train, custom_cv, all_features):
    self.X_train = X_train
    self.Y_train = Y_train
    self.custom_cv = custom_cv
    self.all_features = all_features
    self.scaler = cyp_config.getFeatureScaler()
    self.cv_metric = cyp_config.getFeatureSelectionCVMetric()
    self.verbose = cyp_config.getDebugLevel()

  def setCustomCV(self, custom_cv):
    """Set custom K-Fold validation splits"""
    self.custom_cv = custom_cv

  def setFeatures(self, features):
    """Set the list of all features"""
    self.all_features = all_features

  # K-fold validation to find the optimal number of features
  # and optimal hyperparameters for estimator.
  def featureSelectionGridSearch(self, selector, est, param_grid):
    """Use grid search with k-fold validation to optimize the number of features"""
    X_train_copy = np.copy(self.X_train)
    pipeline = Pipeline([("scaler", self.scaler),
                         ("selector", selector),
                         ("estimator", est)])

    grid_search = GridSearchCV(estimator=pipeline, param_grid=param_grid,
                               scoring=self.cv_metric, cv=self.custom_cv)
    with parallel_backend('spark', n_jobs=-1):
      grid_search.fit(X_train_copy, np.ravel(self.Y_train))

    best_score = grid_search.best_score_
    best_estimator = grid_search.best_estimator_
    indices = []
    # feature selectors should have 'get_support' function
    if ((isinstance(selector, SelectFromModel)) or (isinstance(selector, SelectKBest)) or
        (isinstance(selector, RFE))):
      sel = grid_search.best_estimator_.named_steps['selector']
      indices = sel.get_support(indices=True)

    result = {
        'indices' : indices,
        'best_score' : best_score,
    }

    return result

  # Compare different feature selectors using cross validation score.
  # Also compare combined features with the best individual feature selector.
  def compareFeatureSelectors(self, feature_selectors,
                              est, est_param_grid):
    """Compare feature selectors based on K-fold validation scores"""
    combined_indices = []
    # set it to a large negative value
    best_score = -1000
    best_indices = []
    best_selector = ''
    fs_scores = {}
    for sel_name in feature_selectors:
      selector = feature_selectors[sel_name]['selector']
      sel_param_grid = feature_selectors[sel_name]['param_grid']
      param_grid = sel_param_grid.copy()
      param_grid.update(est_param_grid)
      result = self.featureSelectionGridSearch(selector, est, param_grid)
      param_grid.clear()

      if (self.verbose > 2):
        print('\nFeature selection using', sel_name)
        print('Best cross-validation', self.cv_metric + ':',
              np.round(result['best_score'], 3))

        print('\nSelected Features:')
        print('-------------------')
        printFeatures(self.all_features, result['indices'])

      combined_indices = list(set(combined_indices) | set(result['indices']))

      fs_scores[sel_name] = result['best_score']

      if (result['best_score'] > best_score):
        best_indices = result['indices']
        best_score = result['best_score']
        best_selector = sel_name

    result = {
        'best_selector' : best_selector,
        'best_score' : best_score,
        'best_indices' : best_indices,
        'combined_indices' : combined_indices,
        'scores' : fs_scores
    }

    return result

  # Between the optimal sets for each estimator select the set with the higher score.
  def selectOptimalFeatures(self, feature_selectors,
                            est, est_name, est_param_grid, log_fh):
    """
    Select optimal features by comparing individual feature selectors
    and combined features
    """
    X_train_selected = None
    # set it to a large negative value
    best_score = -1000
    fs_summary = {}
    row_count = 1

    est_info = '\nEstimator: ' + est_name
    est_info += '\n---------------------------'
    log_fh.write(est_info)
    print(est_info)

    result = self.compareFeatureSelectors(feature_selectors,
                                          est, est_param_grid)

    # result includes
    # - 'best_selector' : name of the best selector
    # - 'best_score' : cv score of features selected with estimator
    # - 'best_indices' : indices of features selected
    # - 'best_estimator' : estimator with settings that gave best score
    # - 'combined_indices' : indices of combined features
    # - 'scores' : dict with scores of all feature selectors

    for sel_name in result['scores']:
      est_sel_row = [est_name, sel_name, np.round(result['scores'][sel_name], 2)]
      fs_summary['row' + str(row_count)] = est_sel_row
      row_count += 1

    # calculate cross-validation score for combined features
    X_train_selected = self.X_train[:, result['combined_indices']]
    pipeline = Pipeline([("scaler", self.scaler), ("estimator", est)])
    grid_search = GridSearchCV(estimator=pipeline, param_grid=est_param_grid,
                               scoring=self.cv_metric, cv=self.custom_cv)
    X_train_selected_copy = np.copy(X_train_selected)
    with parallel_backend('spark', n_jobs=-1):
      grid_search.fit(X_train_selected_copy, np.ravel(self.Y_train))

    combined_score = grid_search.best_score_
    combo_sel_row = [est_name, 'combined', np.round(combined_score, 2)]

    fs_summary['row' + str(row_count)] = combo_sel_row
    row_count += 1

    # We check if combined features give us a better score
    # than the best feature selection method
    if (combined_score < result['best_score']):
      selector = result['best_selector']
      selected_indices = result['best_indices']
    else:
      selected_indices = result['combined_indices']

    fs_df_columns = ['estimator', 'selector', self.cv_metric]
    fs_df = pd.DataFrame.from_dict(fs_summary, orient='index', columns=fs_df_columns)

    ftsel_summary_info = '\nFeature Selection Summary'
    ftsel_summary_info += '\n---------------------------'
    ftsel_summary_info += '\n' + fs_df.to_string(index=False) + '\n'
    log_fh.write(ftsel_summary_info)
    print(ftsel_summary_info)

    return selected_indices

"""#### Algorithm Evaluator Class"""

#%%writefile algorithm_evaluation.py
import numpy as np
import pandas as pd 

from sklearn.pipeline import Pipeline
from sklearn.linear_model import Ridge
from sklearn.model_selection import GridSearchCV
from sklearn.utils import parallel_backend

class CYPAlgorithmEvaluator:
  def __init__(self, cyp_config, custom_cv):
    self.scaler = cyp_config.getFeatureScaler()
    self.estimators = cyp_config.getEstimators()
    self.custom_cv = custom_cv
    self.cv_metric = cyp_config.getAlgorithmTrainingCVMetric()
    self.metrics = cyp_config.getEvaluationMetrics()
    self.verbose = cyp_config.getDebugLevel()
    self.use_yield_trend = cyp_config.useYieldTrend()
    self.trend_windows = cyp_config.getTrendWindows()
    self.predict_residuals = cyp_config.predictYieldResiduals()

  def setCustomCV(self, custom_cv):
    """Set custom K-Fold validation splits"""
    self.custom_cv = custom_cv

  # Nash-Sutcliffe Model Efficiency
  def nse(self, Y_true, Y_pred):
    """
        Nash Sutcliffe efficiency coefficient
        input:
          Y_pred: predicted
          Y_true: observed
        output:
          nse: Nash Sutcliffe efficient coefficient
        """
    return (1 - np.sum(np.square(Y_pred - Y_true))/np.sum(np.square(Y_true - np.mean(Y_true))))

  def updateAlgorithmsSummary(self, alg_summary, alg_name,
                              train_scores, test_scores):
    """Update algorithms summary with scores for given algorithm"""
    alg_row = [alg_name]
    alg_index = len(alg_summary)
    for met in train_scores:
      alg_row += [train_scores[met], test_scores[met]]

    alg_summary['row' + str(alg_index)] = alg_row

  def createPredictionDataFrames(self, Y_train_pred, Y_test_pred, data_cols):
    """"Create pandas data frames from true and predicted values"""
    pd_train_df = pd.DataFrame(data=Y_train_pred, columns=data_cols)
    pd_test_df = pd.DataFrame(data=Y_test_pred, columns=data_cols)

    return pd_train_df, pd_test_df

  def printPredictionDataFrames(self, pd_train_df, pd_test_df, log_fh):
    """"Print true and predicted values from pandas data frames"""
    train_info = '\nYield Predictions Training Set'
    train_info += '\n--------------------------------'
    train_info += '\n' + pd_train_df.head(6).to_string(index=False)
    log_fh.write(train_info + '\n')
    print(train_info)

    test_info = '\nYield Predictions Test Set'
    test_info += '\n--------------------------------'
    test_info += '\n' + pd_test_df.head(6).to_string(index=False)
    log_fh.write(test_info + '\n')
    print(test_info)

  def getNullMethodPredictions(self, Y_train_full, Y_test_full, log_fh):
    """
    The Null method or poor man's prediction. Y_*_full includes IDREGION, FYEAR.
    If using yield trend, Y_*_full also include YIELD_TREND.
    The null method predicts the YIELD_TREND or the average of the training set.
    """
    Y_train = Y_train_full[:, 2]
    if (self.use_yield_trend):
      Y_train = Y_train_full[:, 3]

    min_yield = np.round(np.min(Y_train), 2)
    max_yield = np.round(np.max(Y_train), 2)
    avg_yield = np.round(np.mean(Y_train), 2)
    median_yield = np.round(np.median(np.ravel(Y_train)), 2)

    null_method_label = 'Null Method: '
    if (self.use_yield_trend):
      null_method_label += 'Predicting linear yield trend:'
      data_cols = ['IDREGION', 'FYEAR', 'YIELD_TREND', 'YIELD']
      pd_train_df, pd_test_df = self.createPredictionDataFrames(Y_train_full, Y_test_full,
                                                                data_cols)
    else:
      Y_train_full_n = np.insert(Y_train_full, 2, avg_yield, axis=1)
      Y_test_full_n = np.insert(Y_test_full, 2, avg_yield, axis=1)
      null_method_label += 'Predicting average of the training set:'
      data_cols = ['IDREGION', 'FYEAR', 'YIELD_PRED', 'YIELD']
      pd_train_df, pd_test_df = self.createPredictionDataFrames(Y_train_full_n, Y_test_full_n,
                                                                data_cols)

    null_method_info = '\n' + null_method_label
    null_method_info += '\nMin Yield: ' + str(min_yield) + ', Max Yield: ' + str(max_yield)
    null_method_info += '\nMedian Yield: ' + str(median_yield) + ', Mean Yield: ' + str(avg_yield)
    log_fh.write(null_method_info + '\n')
    print(null_method_info)
    self.printPredictionDataFrames(pd_train_df, pd_test_df, log_fh)

    result = {
        'train' : pd_train_df,
        'test' : pd_test_df,
    }

    return result

  def evaluateNullMethodPredictions(self, pd_train_df, pd_test_df, alg_summary):
    """Evaluate the predictions of the null method and add an entry to alg_summary"""
    Y_train = pd_train_df['YIELD'].values
    Y_test = pd_test_df['YIELD'].values

    if (self.use_yield_trend):  
      alg_name = 'trend'
      Y_pred_train = pd_train_df['YIELD_TREND'].values
      Y_pred_test = pd_test_df['YIELD_TREND'].values
    else:
      alg_name = 'average'
      Y_pred_train = pd_train_df['YIELD_PRED'].values
      Y_pred_test = pd_test_df['YIELD_PRED'].values

    train_scores = getPredictionScores(Y_train, Y_pred_train, self.metrics)
    test_scores = getPredictionScores(Y_test, Y_pred_test, self.metrics)
    self.updateAlgorithmsSummary(alg_summary, alg_name, train_scores, test_scores)

  def getYieldTrendML(self, X_train, Y_train, X_test):
    """
    Predict yield trend using a linear model.
    No need to scale features. They are all yield values.
    """
    est = Ridge(alpha=1, random_state=42, max_iter=1000,
                copy_X=True, fit_intercept=True)
    est_param_grid = dict(alpha=[1e-3, 1e-2, 1e-1, 1, 10])

    grid_search = GridSearchCV(estimator=est, param_grid=est_param_grid,
                               scoring=self.cv_metric, cv=self.custom_cv)
    X_train_copy = np.copy(X_train)
    with parallel_backend('spark', n_jobs=-1):
      grid_search.fit(X_train_copy, np.ravel(Y_train))

    best_params = grid_search.best_params_
    best_estimator = grid_search.best_estimator_

    if (self.verbose > 1):
      for param in est_param_grid:
        print(param + '=', best_params[param])

    Y_pred_train = np.reshape(best_estimator.predict(X_train), (X_train.shape[0], 1))
    Y_pred_test = np.reshape(best_estimator.predict(X_test), (X_test.shape[0], 1))

    return Y_pred_train, Y_pred_test

  def estimateYieldTrendAndDetrend(self, X_train, Y_train, X_test, Y_test, features):
    """Estimate yield trend using machine learning and detrend"""
    trend_window = self.trend_windows[0]
    # NOTE assuming previous years' yield values are at the end
    X_train_trend = X_train[:, -trend_window:]
    X_test_trend = X_test[:, -trend_window:]

    Y_train_trend, Y_test_trend = self.getYieldTrendML(X_train_trend, Y_train, X_test_trend)
    # New features exclude previous years' yield and include YIELD_TREND
    features_n = features[:-trend_window] + ['YIELD_TREND']
    X_train_n = np.append(X_train[:, :-trend_window], Y_train_trend, axis=1)
    X_test_n = np.append(X_test[:, :-trend_window], Y_test_trend, axis=1)
    Y_train_res = np.reshape(Y_train, (X_train.shape[0], 1)) - Y_train_trend
    Y_test_res = np.reshape(Y_test, (X_test.shape[0], 1)) - Y_test_trend

    result =  {
        'X_train' : X_train_n,
        'Y_train' : Y_train_res,
        'Y_train_trend' : Y_train_trend,
        'X_test' : X_test_n,
        'Y_test' : Y_test_res,
        'Y_test_trend' : Y_test_trend,
        'features' : features_n,
    }

    return result

  def yieldPredictionsFromResiduals(self, pd_train_df, Y_train, pd_test_df, Y_test):
    """Predictions are residuals. Add trend back to get yield predictions."""
    pd_train_df['YIELD_RES'] = pd_train_df['YIELD']
    pd_train_df['YIELD'] = Y_train
    pd_test_df['YIELD_RES'] = pd_test_df['YIELD']
    pd_test_df['YIELD'] = Y_test
    
    for alg in self.estimators:
      pd_train_df['YIELD_RES_PRED_' + alg] = pd_train_df['YIELD_PRED_' + alg]
      pd_train_df['YIELD_PRED_' + alg] = pd_train_df['YIELD_RES_PRED_' + alg] + pd_train_df['YIELD_TREND']
      pd_test_df['YIELD_RES_PRED_' + alg] = pd_test_df['YIELD_PRED_' + alg]
      pd_test_df['YIELD_PRED_' + alg] = pd_test_df['YIELD_RES_PRED_' + alg] + pd_test_df['YIELD_TREND']

    sel_cols = ['IDREGION', 'FYEAR', 'YIELD_TREND', 'YIELD_RES']
    for alg in self.estimators:
      sel_cols += ['YIELD_RES_PRED_' + alg, 'YIELD_PRED_' + alg]

    sel_cols.append('YIELD')
    result = {
        'train' : pd_train_df[sel_cols],
        'test' : pd_test_df[sel_cols],
    }

    return result

  def trainAndTest(self, X_train, Y_train, X_test,
                   est, est_name, est_param_grid):
    """
    Use k-fold validation to tune hyperparameters and evaluate performance.
    """
    if (self.verbose > 1):
      print('\nEstimator', est_name)
      print('---------------------------')
      print(est)

    pipeline = Pipeline([("scaler", self.scaler), ("estimator", est)])
    grid_search = GridSearchCV(estimator=pipeline, param_grid=est_param_grid,
                               scoring=self.cv_metric, cv=self.custom_cv)
    X_train_copy = np.copy(X_train)
    with parallel_backend('spark', n_jobs=-1):
      grid_search.fit(X_train_copy, np.ravel(Y_train))

    best_params = grid_search.best_params_
    best_estimator = grid_search.best_estimator_

    if (self.verbose > 1):
      for param in est_param_grid:
        print(param + '=', best_params[param])

    Y_pred_train = np.reshape(best_estimator.predict(X_train), (X_train.shape[0], 1))
    Y_pred_test = np.reshape(best_estimator.predict(X_test), (X_test.shape[0], 1))

    return Y_pred_train, Y_pred_test

  def combineAlgorithmPredictions(self, pd_ml_predictions, pd_alg_predictions, alg):
    """Combine predictions of ML algorithms."""
    join_cols = ['IDREGION', 'FYEAR']

    if (pd_ml_predictions is None):
      pd_ml_predictions = pd_alg_predictions
      pd_ml_predictions = pd_ml_predictions.rename(columns={'YIELD_PRED': 'YIELD_PRED_' + alg })
    else:
      pd_alg_predictions = pd_alg_predictions[join_cols + ['YIELD_PRED']]
      pd_ml_predictions = pd_ml_predictions.merge(pd_alg_predictions, on=join_cols)
      pd_ml_predictions = pd_ml_predictions.rename(columns={'YIELD_PRED': 'YIELD_PRED_' + alg })
      # Put YIELD at the end
      all_cols = list(pd_ml_predictions.columns)
      col_order = all_cols[:-2] + ['YIELD_PRED_' + alg, 'YIELD']
      pd_ml_predictions = pd_ml_predictions[col_order]

    return pd_ml_predictions

  def getMLPredictions(self, X_train, Y_train_full, X_test, Y_test_full,
                       cyp_ftsel, ft_selectors, features, log_fh):
    """Train and evaluate crop yield prediction algorithms"""
    Y_train = Y_train_full[:, -1]
    Y_test = Y_test_full[:, -1]
    pd_test_predictions = None
    pd_train_predictions = None

    # feature selection frequency
    # NOTE must be in sync with crop calendar periods
    feature_selection_counts = {
        'static' : {},
        'p0' : {},
        'p1' : {},
        'p2' : {},
        'p3' : {},
        'p4' : {},
        'p5' : {},
    }

    for est_name in self.estimators:
      # feature selection
      est = self.estimators[est_name]['estimator']
      param_grid = self.estimators[est_name]['fs_param_grid']
      selected_indices = cyp_ftsel.selectOptimalFeatures(ft_selectors,
                                                         est, est_name, param_grid,
                                                         log_fh)
      sel_fts_info = '\nSelected Features:'
      sel_fts_info += '\n-------------------'
      log_fh.write(sel_fts_info)
      print(sel_fts_info)
      printFeatures(features, selected_indices, log_fh)

      # update feature selection counts
      for idx in selected_indices:
        ft_count = 0
        ft = features[idx]
        ft_period = None
        for p in feature_selection_counts:
          if p in ft:
            ft_period = p

        if (ft_period is None):
          ft_period = 'static'

        if (ft in feature_selection_counts[ft_period]):
          ft_count = feature_selection_counts[ft_period][ft]

        feature_selection_counts[ft_period][ft] = ft_count + 1

      X_train_sel = X_train[:, selected_indices]
      X_test_sel = X_test[:, selected_indices]

      # Training and testing
      param_grid = self.estimators[est_name]['param_grid']
      # yield/yield residual predictions
      Y_pred_train, Y_pred_test = self.trainAndTest(X_train_sel, Y_train, X_test_sel,
                                                    est, est_name, param_grid)
      data_cols = ['IDREGION', 'FYEAR']
      if (self.use_yield_trend):
        data_cols.append('YIELD_TREND')
        Y_train_full_n = np.insert(Y_train_full, 3, Y_pred_train[:, 0], axis=1)
        Y_test_full_n = np.insert(Y_test_full, 3, Y_pred_test[:, 0], axis=1)
      else:
        Y_train_full_n = np.insert(Y_train_full, 2, Y_pred_train[:, 0], axis=1)
        Y_test_full_n = np.insert(Y_test_full, 2, Y_pred_test[:, 0], axis=1)

      data_cols += ['YIELD_PRED', 'YIELD']
      pd_train_df, pd_test_df = self.createPredictionDataFrames(Y_train_full_n, Y_test_full_n, data_cols)
      pd_train_predictions = self.combineAlgorithmPredictions(pd_train_predictions, pd_train_df, est_name)
      pd_test_predictions = self.combineAlgorithmPredictions(pd_test_predictions, pd_test_df, est_name)

    ft_counts_info = '\nFeature Selection Frequencies'
    ft_counts_info += '\n-------------------------------'
    for ft_period in feature_selection_counts:
      ft_count_str = ft_period + ': '
      for ft in sorted(feature_selection_counts[ft_period],
                       key=feature_selection_counts[ft_period].get, reverse=True):
        ft_count_str += ft + '(' + str(feature_selection_counts[ft_period][ft]) + '), '

      if (len(feature_selection_counts[ft_period]) > 0):
        # drop ', ' from the end
        ft_count_str = ft_count_str[:-2]

      ft_counts_info += '\n' + ft_count_str

    ft_counts_info += '\n'
    log_fh.write(ft_counts_info)
    if (self.verbose > 1):
      print(ft_counts_info)

    result = {
        'train' : pd_train_predictions,
        'test' : pd_test_predictions,
    }

    return result

  def evaluateMLPredictions(self, pd_train_predictions, pd_test_predictions, alg_summary):
    """Evaluate predictions of ML algorithms and add entries to alg_summary."""
    Y_train = pd_train_predictions['YIELD'].values
    Y_test = pd_test_predictions['YIELD'].values

    for alg in self.estimators:
      alg_col = 'YIELD_PRED_' + alg
      Y_pred_train = pd_train_predictions[alg_col].values
      Y_pred_test = pd_test_predictions[alg_col].values
      train_scores = getPredictionScores(Y_train, Y_pred_train, self.metrics)
      test_scores = getPredictionScores(Y_test, Y_pred_test, self.metrics)

      self.updateAlgorithmsSummary(alg_summary, alg, train_scores, test_scores)

"""#### Run Machine Learning"""

#%%writefile run_machine_learning.py
import pandas as pd
import numpy as np
from joblibspark import register_spark

def warn(*args, **kwargs):
    pass

import warnings
warnings.warn = warn

def getValidationSplits(cyp_config, pd_train_df, pd_test_df, log_fh):
  """Split features and label into training and test sets"""
  use_yield_trend = cyp_config.useYieldTrend()
  debug_level = cyp_config.getDebugLevel()

  regions = [reg for reg in pd_train_df['IDREGION'].unique()]
  num_regions = len(regions)

  original_headers = list(pd_train_df.columns.values)
  features = []
  labels = []
  if (use_yield_trend):
    features = original_headers[2:-2]
    labels = original_headers[:2] + original_headers[-2:]
  else:
    features = original_headers[2:-1]
    labels = original_headers[:2] + original_headers[-1:]

  X_train = pd_train_df[features].values
  Y_train = pd_train_df[labels].values

  train_info = '\nTraining Data Size: ' + str(len(pd_train_df.index)) + ' rows'
  train_info += '\nX cols: ' + str(X_train.shape[1]) + ', Y cols: ' + str(Y_train.shape[1])
  train_info += '\n' + pd_train_df.head(5).to_string(index=False)
  log_fh.write(train_info + '\n')
  if (debug_level > 1):
    print(train_info)

  X_test = pd_test_df[features].values
  Y_test = pd_test_df[labels].values

  test_info = '\nTest Data Size: ' + str(len(pd_test_df.index)) + ' rows'
  test_info += '\nX cols: ' + str(X_test.shape[1]) + ', Y cols: ' + str(Y_test.shape[1])
  test_info += '\n' + pd_test_df.head(5).to_string(index=False)
  log_fh.write(test_info + '\n')
  if (debug_level > 1):
    print(test_info)

  # print feature names
  num_features = len(features)
  indices = [idx for idx in range(num_features)]
  feature_info = '\nAll features'
  feature_info += '\n-------------'
  log_fh.write(feature_info)
  print(feature_info)
  printFeatures(features, indices, log_fh)

  # num_folds for k-fold cv
  num_folds = 5
  custom_cv = num_folds
  if (use_yield_trend):
    cyp_cv_splitter = CYPTrainTestSplitter(cyp_config)
    custom_cv = cyp_cv_splitter.customKFoldValidationSplit(Y_train, num_folds, log_fh)

  # L1 penalty = error (y_pred - y_obs)^2 + alpha * sum (|w_i|) = sparsity regularization
  # L2 penalty = error (y_pred - y_obs)^2 + alpha * sqrt ( sum (w_i^2) ) = weight decay regularization

  result = {
      'X_train' : X_train,
      'Y_train_full' : Y_train,
      'X_test' : X_test,
      'Y_test_full' : Y_test,
      'custom_cv' : custom_cv,
      'features' : features,
  }

  return result

def getMachineLearningPredictions(cyp_config, pd_train_df, pd_test_df, log_fh):
  """Train and evaluate algorithms"""
  metrics = cyp_config.getEvaluationMetrics()
  use_yield_trend = cyp_config.useYieldTrend()
  predict_residuals = cyp_config.predictYieldResiduals()
  debug_level = cyp_config.getDebugLevel()

  # register spark parallel backend
  register_spark()

  eval_info = '\nTraining and Evaluation'
  eval_info += '\n-------------------------'
  log_fh.write(eval_info)
  if (debug_level > 1):
    print(eval_info)

  data_splits = getValidationSplits(cyp_config, pd_train_df, pd_test_df, log_fh)
  X_train = data_splits['X_train']
  Y_train_full = np.copy(data_splits['Y_train_full'])
  X_test = data_splits['X_test']
  Y_test_full = np.copy(data_splits['Y_test_full'])
  features = data_splits['features']
  custom_cv = data_splits['custom_cv']

  alg_summary = {}
  cyp_algeval = CYPAlgorithmEvaluator(cyp_config, custom_cv)
  null_preds = cyp_algeval.getNullMethodPredictions(Y_train_full, Y_test_full, log_fh)

  if (use_yield_trend and predict_residuals):
    result = cyp_algeval.estimateYieldTrendAndDetrend(X_train, Y_train_full[:, -1],
                                                      X_test, Y_test_full[:, -1], features)
    X_train = result['X_train']
    Y_train_full[:, -1] = result['Y_train'][:, 0]
    Y_train_full[:, -2] = result['Y_train_trend'][:, 0]
    X_test = result['X_test']
    Y_test_full[:, -1] = result['Y_test'][:,0]
    Y_test_full[:, -2] = result['Y_test_trend'][:, 0]
    features = result['features']

  # feature selection methods
  num_features = len(features)
  ft_selectors = cyp_config.getFeatureSelectors(X_train, Y_train_full[:, -1],
                                                num_features, custom_cv)
  cyp_ftsel = CYPFeatureSelector(cyp_config, X_train, Y_train_full[:, -1], custom_cv, features)
  ml_preds = cyp_algeval.getMLPredictions(X_train, Y_train_full, X_test, Y_test_full,
                                          cyp_ftsel, ft_selectors, features, log_fh)
  if (use_yield_trend and predict_residuals):
    # NOTE Y_train_full, Y_test_full can get modified above
    ml_preds = cyp_algeval.yieldPredictionsFromResiduals(ml_preds['train'],
                                                         data_splits['Y_train_full'][:, -1],
                                                         ml_preds['test'],
                                                         data_splits['Y_test_full'][:, -1])

  cyp_algeval.evaluateNullMethodPredictions(null_preds['train'], null_preds['test'], alg_summary)
  cyp_algeval.evaluateMLPredictions(ml_preds['train'], ml_preds['test'], alg_summary)
  cyp_algeval.printPredictionDataFrames(ml_preds['train'], ml_preds['test'], log_fh)

  alg_df_columns = ['algorithm']
  for met in metrics:
    alg_df_columns += ['train_' + met, 'test_' + met]

  alg_df = pd.DataFrame.from_dict(alg_summary, orient='index', columns=alg_df_columns)

  eval_summary_info = '\nAlgorithm Evaluation Summary'
  eval_summary_info += '\n-----------------------------'
  eval_summary_info += '\n' + alg_df.to_string(index=False) + '\n'
  log_fh.write(eval_summary_info)
  print(eval_summary_info)

  return ml_preds['test']

def saveMLPredictions(cyp_config, sqlCtx, pd_ml_predictions):
  """Save ML predictions to a CSV file"""
  debug_level = cyp_config.getDebugLevel()
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  nuts_level = cyp_config.getNUTSLevel()
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()
  ml_algs = cyp_config.getEstimators()

  output_path = cyp_config.getOutputPath()
  output_file = getPredictionFilename(crop, country, nuts_level, use_yield_trend,
                                      early_season_prediction, early_season_end)

  save_pred_path = output_path + '/' + output_file
  if (debug_level > 1):
    print('\nSaving predictions to', save_pred_path + '.csv')
    print(pd_ml_predictions.head(5))

  pd_ml_predictions.to_csv(save_pred_path + '.csv', index=False, header=True)

  # NOTE: In some environments, Spark can write, but pandas cannot.
  # In such cases, use the following code.
  # spark_predictions_df = sqlCtx.createDataFrame(pd_ml_predictions)
  # spark_predictions_df.coalesce(1)\
  #                     .write.option('header','true')\
  #                     .mode("overwrite").csv(save_pred_path)

"""### Load Saved Predictions"""

#%%writefile load_saved_predictions.py
import pandas as pd

def loadSavedPredictions(cyp_config, spark):
  """Load machine learning predictions from saved CSV file"""
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  nuts_level = cyp_config.getNUTSLevel()
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()
  debug_level = cyp_config.getDebugLevel()

  pred_file_path = cyp_config.getOutputPath()
  pred_file = getPredictionFilename(crop, country, nuts_level, use_yield_trend,
                                    early_season_prediction, early_season_end)
  
  pred_file += '.csv'
  pd_ml_predictions = pd.read_csv(pred_file_path + '/' + pred_file, header=0)

  # NOTE: In some environments, Spark can read, but pandas cannot.
  # In such cases, use the following code.
  # all_pred_df = spark.read.csv(pred_file_path + '/' + pred_file, header=True, inferSchema=True)
  # pd_ml_predictions = all_pred_df.toPandas()

  if (debug_level > 1):
    print(pd_ml_predictions.head(5))

  return pd_ml_predictions

"""### Compare Predictions with MCYFS"""

#%%writefile compare_with_mcyfs.py
import numpy as np
import pandas as pd

def saveNUTS0Predictions(cyp_config, sqlCtx, nuts0_ml_predictions):
  """Save predictions aggregated to NUTS0"""
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  nuts_level = 'NUTS0'
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()
  debug_level = cyp_config.getDebugLevel()

  output_path = cyp_config.getOutputPath()
  output_file = getPredictionFilename(crop, country, nuts_level, use_yield_trend,
                                      early_season_prediction, early_season_end)

  save_pred_path = output_path + '/' + output_file
  if (debug_level > 1):
    print('\nNUTS0 Predictions of ML algorithms')
    print(nuts0_ml_predictions.head(5))
    print('\nSaving predictions to', save_pred_path + '.csv')

  nuts0_ml_predictions.to_csv(save_pred_path + '.csv', index=False, header=True)

  # NOTE: In some environments, Spark can write, but pandas cannot.
  # In such cases, use the following code.
  # spark_predictions_df = sqlCtx.createDataFrame(nuts0_ml_predictions)
  # spark_predictions_df.coalesce(1)\
  #                     .write.option('header','true')\
  #                     .mode("overwrite").csv(save_pred_path)

def getDataForMCYFSComparison(spark, cyp_config, test_years):
  """Load and preprocess data for MCYFS comparison"""
  data_path = cyp_config.getDataPath()
  crop_id = cyp_config.getCropID()
  nuts_level = cyp_config.getNUTSLevel()
  season_crosses_calyear = cyp_config.seasonCrossesCalendarYear()
  early_season_end = cyp_config.getEarlySeasonEndDekad()
  debug_level = cyp_config.getDebugLevel()
  area_nuts = ['NUTS' + str(i) for i in range(int(nuts_level[-1]), 0, -1)]
  data_sources = {
      'WOFOST' : nuts_level,
      'AREA_FRACTIONS' : area_nuts,
      'YIELD' : 'NUTS0',
      'YIELD_PRED_MCYFS' : 'NUTS0',
  }

  if (run_tests):
    test_util = TestUtil(spark)
    test_util.runAllTests()

  print('##############')
  print('# Load Data  #')
  print('##############')

  if (run_tests):
    test_loader = TestDataLoader(spark)
    test_loader.runAllTests()

  cyp_config.setDataSources(data_sources)
  cyp_loader = CYPDataLoader(spark, cyp_config)
  data_dfs = cyp_loader.loadAllData()

  wofost_df = data_dfs['WOFOST']
  area_dfs = data_dfs['AREA_FRACTIONS']
  nuts0_yield_df = data_dfs['YIELD']
  mcyfs_yield_df = data_dfs['YIELD_PRED_MCYFS']

  print('####################')
  print('# Preprocess Data  #')
  print('####################')

  if (run_tests):
    test_preprocessor = TestDataPreprocessor(spark)
    test_preprocessor.runAllTests()

  cyp_preprocessor = CYPDataPreprocessor(spark, cyp_config)
  wofost_df = wofost_df.filter(wofost_df['CROP_ID'] == crop_id).drop('CROP_ID')
  crop_season = cyp_preprocessor.getCropSeasonInformation(wofost_df, season_crosses_calyear)
  wofost_df = cyp_preprocessor.preprocessWofost(wofost_df, crop_season, season_crosses_calyear)

  for i in range(len(area_dfs)):
    af_df = area_dfs[i]
    af_df = cyp_preprocessor.preprocessAreaFractions(af_df, crop_id)
    af_df = af_df.filter(af_df['FYEAR'].isin(test_years))
    area_dfs[i] = af_df

  if (debug_level > 1):
    print('NUTS0 Yield before preprocessing')
    nuts0_yield_df.show(10)

  nuts0_yield_df = cyp_preprocessor.preprocessYield(nuts0_yield_df, crop_id)
  nuts0_yield_df = nuts0_yield_df.filter(nuts0_yield_df['FYEAR'].isin(test_years))
  if (debug_level > 1):
    print('NUTS0 Yield after preprocessing')
    nuts0_yield_df.show(10)

  if (debug_level > 1):
    print('MCYFS yield predictions before preprocessing')
    mcyfs_yield_df.show(10)

  mcyfs_yield_df = cyp_preprocessor.preprocessYieldMCYFS(mcyfs_yield_df, crop_id)
  mcyfs_yield_df = mcyfs_yield_df.filter(mcyfs_yield_df['FYEAR'].isin(test_years))
  if (debug_level > 1):
    print('MCYFS yield predictions after preprocessing')
    mcyfs_yield_df.show(10)

  # Check we have yield data for crop
  assert (nuts0_yield_df is not None)
  assert (mcyfs_yield_df is not None)

  if (run_tests):
    test_summarizer = TestDataSummarizer(spark)
    test_summarizer.runAllTests()

  cyp_summarizer = CYPDataSummarizer(cyp_config)
  dvs_summary = cyp_summarizer.wofostDVSSummary(wofost_df, early_season_end)
  dvs_summary = dvs_summary.filter(dvs_summary['CAMPAIGN_YEAR'].isin(test_years))

  data_dfs = {
      'WOFOST_DVS' : dvs_summary,
      'AREA_FRACTIONS' : area_dfs,
      'YIELD_NUTS0' : nuts0_yield_df,
      'YIELD_PRED_MCYFS' : mcyfs_yield_df
  }

  return data_dfs

def fillMissingDataWithAverage(pd_pred_df, print_debug):
  """Fill missing data with regional average or zero"""
  regions = pd_pred_df['IDREGION'].unique()

  for reg_id in regions:
    reg_filter = (pd_pred_df['IDREGION'] == reg_id)
    pd_reg_pred_df = pd_pred_df[reg_filter]

    if (len(pd_reg_pred_df[pd_reg_pred_df['YIELD_PRED'].notnull()].index) == 0):
      if (print_debug):
        print('No data for', reg_id)

      pd_pred_df.loc[reg_filter, 'FRACTION'] = 0.0
      pd_pred_df.loc[reg_filter, 'YIELD_PRED'] = 0.0
    else:
      reg_avg_yield_pred = pd_pred_df.loc[reg_filter, 'YIELD_PRED'].mean()
      pd_pred_df.loc[reg_filter, 'YIELD_PRED'] = pd_pred_df.loc[reg_filter, 'YIELD_PRED']\
                                                           .fillna(reg_avg_yield_pred)  

  return pd_pred_df

def recalculateAreaFractions(pd_pred_df, print_debug):
  """Recalculate area fractions by excluding regions with missing data"""
  join_cols = ['IDREG_PARENT', 'FYEAR']
  pd_af_sum = pd_pred_df.groupby(join_cols).agg(FRACTION_SUM=('FRACTION', 'sum')).reset_index()
  pd_pred_df = pd_pred_df.merge(pd_af_sum, on=join_cols, how='left')
  pd_pred_df['FRACTION'] = pd_pred_df['FRACTION'] / pd_pred_df['FRACTION_SUM']
  pd_pred_df = pd_pred_df.drop(columns=['FRACTION_SUM'])

  return pd_pred_df

def aggregatePredictionsToNUTS0(cyp_config, pd_ml_predictions,
                                area_dfs, test_years, join_cols):
  """Aggregate regional predictions to national level"""
  pd_area_dfs = []
  nuts_level = cyp_config.getNUTSLevel()
  use_yield_trend = cyp_config.useYieldTrend()
  crop_id = cyp_config.getCropID()
  alg_names = list(cyp_config.getEstimators().keys())
  debug_level = cyp_config.getDebugLevel()

  for af_df in area_dfs:
    pd_af_df = af_df.toPandas()
    pd_af_df = pd_af_df[pd_af_df['FYEAR'].isin(test_years)]
    pd_area_dfs.append(pd_af_df)

  nuts0_pred_df = None
  for alg in alg_names:
    sel_cols = ['IDREGION', 'FYEAR', 'YIELD_PRED_' + alg]
    pd_alg_pred_df = pd_ml_predictions[sel_cols]
    pd_alg_pred_df = pd_alg_pred_df.rename(columns={'YIELD_PRED_' + alg : 'YIELD_PRED'})

    for idx in range(len(pd_area_dfs)):
      pd_af_df = pd_area_dfs[idx]
      # merge with area fractions to get all regions and years
      pd_alg_pred_df = pd_af_df.merge(pd_alg_pred_df, on=join_cols)
      print_debug = (debug_level > 2) and (alg == alg_names[0])
      pd_alg_pred_df = fillMissingDataWithAverage(pd_alg_pred_df, print_debug)
      pd_alg_pred_df['IDREG_PARENT'] = pd_alg_pred_df['IDREGION'].str[:-1]
      pd_alg_pred_df = recalculateAreaFractions(pd_alg_pred_df, print_debug)
      if (print_debug):
        print('\nAggregation to NUTS' + str(len(pd_area_dfs) - (idx + 1)))
        print(pd_alg_pred_df[pd_alg_pred_df['FYEAR'] == test_years[0]].head(10))

      pd_alg_pred_df['YPRED_WEIGHTED'] = pd_alg_pred_df['YIELD_PRED'] * pd_alg_pred_df['FRACTION']
      pd_alg_pred_df = pd_alg_pred_df.groupby(by=['IDREG_PARENT', 'FYEAR'])\
                                     .agg(YPRED_WEIGHTED=('YPRED_WEIGHTED', 'sum')).reset_index()
      pd_alg_pred_df = pd_alg_pred_df.rename(columns={'IDREG_PARENT': 'IDREGION',
                                                      'YPRED_WEIGHTED': 'YIELD_PRED' })

    pd_alg_pred_df = pd_alg_pred_df.rename(columns={ 'YIELD_PRED': 'YIELD_PRED_' + alg })
    if (nuts0_pred_df is None):
      nuts0_pred_df = pd_alg_pred_df
    else:
      nuts0_pred_df = nuts0_pred_df.merge(pd_alg_pred_df, on=join_cols)

  return nuts0_pred_df

def getMCYFSPrediction(pd_mcyfs_pred_df, pred_year, pred_dekad, print_debug):
  """Get MCYFS prediction for given year with prediction date close to pred_dekad"""
  pd_pred_year = pd_mcyfs_pred_df[pd_mcyfs_pred_df['FYEAR'] == pred_year]
  mcyfs_pred_dekads = pd_pred_year['PRED_DEKAD'].unique()
  if (len(mcyfs_pred_dekads) == 0):
    return 0.0

  mcyfs_pred_dekads = sorted(mcyfs_pred_dekads)
  mcyfs_pred_dekad = mcyfs_pred_dekads[-1]
  if (pred_dekad < mcyfs_pred_dekad):
    for dek in mcyfs_pred_dekads:
      if dek >= pred_dekad:
        mcyfs_pred_dekad = dek
        break

  pd_pred_dek = pd_pred_year[pd_pred_year['PRED_DEKAD'] == mcyfs_pred_dekad]
  yield_pred_list = pd_pred_dek['YIELD_PRED'].values

  if (print_debug):
    print('\nAll MCYFS dekads for', pred_year, ':', mcyfs_pred_dekads)
    print('MCYFS prediction dekad', mcyfs_pred_dekad)
    print('ML Baseline prediction dekad', pred_dekad)
    print('MCYFS prediction:', yield_pred_list[0], '\n')

  return yield_pred_list[0]

def getNUTS0Yield(pd_nuts0_yield_df, pred_year, print_debug):
  """Get the true (reported) Eurostat yield value"""
  nuts0_yield_year = pd_nuts0_yield_df[pd_nuts0_yield_df['FYEAR'] == pred_year]
  pred_year_yield = nuts0_yield_year['YIELD'].values
  if (len(pred_year_yield) == 0):
    return 0.0

  if (print_debug):
    print(pred_year, 'Eurostat yield', pred_year_yield[0])

  return pred_year_yield[0]

def comparePredictionsWithMCYFS(sqlCtx, cyp_config, pd_ml_predictions, log_fh):
  """Compare ML Baseline predictions with MCYFS predictions"""
  # We need AREA_FRACTIONS, MCYFS yield predictions and NUTS0 Eurostat YIELD
  # for comparison with MCYFS
  country_code = cyp_config.getCountryCode()
  debug_level = cyp_config.getDebugLevel()
  alg_names = list(cyp_config.getEstimators().keys())
  test_years = list(pd_ml_predictions['FYEAR'].unique())
  early_season_prediction = cyp_config.earlySeasonPrediction()

  spark = sqlCtx.sparkSession
  data_dfs = getDataForMCYFSComparison(spark, cyp_config, test_years)
  pd_dvs_summary = data_dfs['WOFOST_DVS'].toPandas()
  pd_nuts0_yield_df = data_dfs['YIELD_NUTS0'].toPandas()
  pd_mcyfs_pred_df = data_dfs['YIELD_PRED_MCYFS'].toPandas()
  area_dfs = data_dfs['AREA_FRACTIONS']
  join_cols = ['IDREGION', 'FYEAR']
  test_years = pd_ml_predictions['FYEAR'].unique()
  metrics = cyp_config.getEvaluationMetrics()
  nuts0_pred_df = aggregatePredictionsToNUTS0(cyp_config, pd_ml_predictions,
                                              area_dfs, test_years, join_cols)

  crop_season_cols = ['IDREGION', 'CAMPAIGN_YEAR', 'CALENDAR_END_SEASON', 'CALENDAR_EARLY_SEASON']
  pd_dvs_summary = pd_dvs_summary[crop_season_cols].rename(columns={ 'CAMPAIGN_YEAR' : 'FYEAR' })
  pd_dvs_summary = pd_dvs_summary.groupby('FYEAR').agg(END_SEASON=('CALENDAR_END_SEASON', 'mean'),
                                                       EARLY_SEASON=('CALENDAR_EARLY_SEASON', 'mean'))\
                                                       .round(0).reset_index()
  if (debug_level > 1):
    print(pd_dvs_summary.head(5).to_string(index=False))

  alg_summary = {}
  Y_pred_mcyfs = []
  Y_true = []
  nuts0_pred_df['YIELD_PRED_MCYFS'] = 0.0
  nuts0_pred_df['YIELD'] = 0.0
  nuts0_pred_df = nuts0_pred_df.sort_values(by=join_cols)
  ml_pred_years = nuts0_pred_df['FYEAR'].unique()
  mcyfs_pred_years = []
  print_debug = (debug_level > 2)
  if (print_debug):
    print('\nPredictions and true values for', country_code)

  for yr in ml_pred_years:
    pred_dekad = pd_dvs_summary[pd_dvs_summary['FYEAR'] == yr]['END_SEASON'].values[0]
    if (early_season_prediction):
      pred_dekad = pd_dvs_summary[pd_dvs_summary['FYEAR'] == yr]['EARLY_SEASON'].values[0]

    mcyfs_pred = getMCYFSPrediction(pd_mcyfs_pred_df, yr, pred_dekad, print_debug)
    nuts0_yield = getNUTS0Yield(pd_nuts0_yield_df, yr, print_debug)
    if ((mcyfs_pred > 0.0) and (nuts0_yield > 0.0)):
      nuts0_pred_df.loc[nuts0_pred_df['FYEAR'] == yr, 'YIELD'] = nuts0_yield
      nuts0_pred_df.loc[nuts0_pred_df['FYEAR'] == yr, 'YIELD_PRED_MCYFS'] = mcyfs_pred
      mcyfs_pred_years.append(yr)

  nuts0_pred_df = nuts0_pred_df[nuts0_pred_df['FYEAR'].isin(mcyfs_pred_years)]
  Y_true = nuts0_pred_df['YIELD'].values

  if (print_debug):
    print(nuts0_pred_df.head(5))

  if (len(mcyfs_pred_years) > 0):
    for alg in alg_names:
      Y_pred_alg = nuts0_pred_df['YIELD_PRED_' + alg].values
      alg_nuts0_scores = getPredictionScores(Y_true, Y_pred_alg, metrics)

      alg_row = [alg]
      for met in alg_nuts0_scores:
        alg_row.append(alg_nuts0_scores[met])

      alg_index = len(alg_summary)
      alg_summary['row' + str(alg_index)] = alg_row

    Y_pred_mcyfs = nuts0_pred_df['YIELD_PRED_MCYFS'].values
    mcyfs_nuts0_scores = getPredictionScores(Y_true, Y_pred_mcyfs, metrics)
    alg_row = ['MCYFS_Predictions']
    for met in mcyfs_nuts0_scores:
      alg_row.append(mcyfs_nuts0_scores[met])

    alg_index = len(alg_summary)
    alg_summary['row' + str(alg_index)] = alg_row

    alg_df_columns = ['algorithm']
    for met in metrics:
      alg_df_columns += ['test_' + met]

    alg_df = pd.DataFrame.from_dict(alg_summary, orient='index',
                                    columns=alg_df_columns)
    eval_summary_info = '\nAlgorithm Evaluation Summary (NUTS0) for ' + country_code
    eval_summary_info += '\n-------------------------------------------'
    eval_summary_info += '\n' + alg_df.to_string(index=False) + '\n'
    log_fh.write(eval_summary_info)
    print(eval_summary_info)

  save_predictions = cyp_config.savePredictions()
  if (save_predictions):
    saveNUTS0Predictions(cyp_config, sqlCtx, nuts0_pred_df)

"""## Tests

### Test Utility Functions
"""

#%%writefile test_util.py
import numpy as np

class TestUtil():
  def __init__(self, spark):
    self.good_date = spark.createDataFrame([(1, '19940102'),
                                          (2, '15831224')],
                                         ['ID', 'DATE'])
    self.bad_date = spark.createDataFrame([(1, '14341224'),
                                           (2, '12345678'),
                                          (3, '123-12-24')],
                                         ['ID', 'DATE'])

  def testDateFormat(self):
    print('\n Test Date Format')
    self.good_date = self.good_date.withColumn('FYEAR', getYear('DATE'))
    self.good_date.show()
    self.bad_date = self.bad_date.withColumn('FYEAR', getYear('DATE'))
    self.bad_date.show()
    assert (self.bad_date.filter(self.bad_date.FYEAR.isNull()).count() == 2)
    self.bad_date = self.bad_date.withColumn('MONTH', getMonth('DATE'))
    self.bad_date.show()
    assert (self.bad_date.filter(self.bad_date.MONTH.isNull()).count() == 2)
    self.bad_date = self.bad_date.withColumn('DAY', getDay('DATE'))
    # check the day here for first date, it's incorrect
    # seems to be a Spark issue
    self.bad_date.show()
    assert (self.bad_date.filter(self.bad_date.DAY.isNull()).count() == 2)
    self.bad_date = self.bad_date.withColumn('DEKAD', getDekad('DATE'))
    self.bad_date.show()
    assert (self.bad_date.filter(self.bad_date.DEKAD.isNull()).count() == 2)

  def testGetYear(self):
    print('\n Test getYear')
    self.good_date = self.good_date.withColumn('FYEAR', getYear('DATE'))
    self.good_date.show()
    year1 = self.good_date.filter(self.good_date.ID == 1).select('FYEAR').collect()[0][0]
    assert year1 == 1994
    year2 = self.good_date.filter(self.good_date.ID == 2).select('FYEAR').collect()[0][0]
    assert year2 == 1583

  def testGetMonth(self):
    print('\n Test getMonth')
    self.good_date = self.good_date.withColumn('MONTH', getMonth('DATE'))
    self.good_date.show()
    month1 = self.good_date.filter(self.good_date.ID == 1).select('MONTH').collect()[0][0]
    assert month1 == 1
    month2 = self.good_date.filter(self.good_date.ID == 2).select('MONTH').collect()[0][0]
    assert month2 == 12

  def testGetDay(self):
    print('\n Test getDay')
    self.good_date = self.good_date.withColumn('DAY', getDay('DATE'))
    self.good_date.show()
    day1 = self.good_date.filter(self.good_date.ID == 1).select('DAY').collect()[0][0]
    assert day1 == 2
    day2 = self.good_date.filter(self.good_date.ID == 2).select('DAY').collect()[0][0]
    assert day2 == 24

  def testGetDekad(self):
    print('\n Test getDekad')
    self.good_date = self.good_date.withColumn('DEKAD', getDekad('DATE'))
    self.good_date.show()
    dekad1 = self.good_date.filter(self.good_date.ID == 1).select('DEKAD').collect()[0][0]
    assert dekad1 == 1
    dekad2 = self.good_date.filter(self.good_date.ID == 2).select('DEKAD').collect()[0][0]
    assert dekad2 == 36

  def testCropIDToName(self):
    print('\n Test cropIDToName')
    crop_name = cropIDToName(crop_name_dict, 6)
    print(6, ':' + crop_name)
    assert crop_name == 'sugarbeet'
    crop_name = cropIDToName(crop_name_dict, 8)
    print(8, ':' + crop_name)
    assert crop_name == 'NA'

  def testCropNameToID(self):
    print('\n Test cropNameToID')
    crop_id = cropNameToID(crop_id_dict, 'Potatoes')
    print('Potatoes:', crop_id)
    assert crop_id == 7
    crop_id = cropNameToID(crop_id_dict, 'Soybean')
    print('Soybean:', crop_id)
    assert crop_id == 0

  def testPrintFeatures(self):
    print('\n Test printFeatures')
    features = ['feat' + str(i+1) for i in range(15)]
    num_features = len(features)
    num_half = np.cast['int64'](np.floor(num_features/2))
    indices1 = [ i for i in range(num_features)]
    indices2 = [ 2*i for i in range(num_half)]
    indices3 = [ (2*i + 1) for i in range(num_half)]

    printFeatures(features, indices1)
    printFeatures(features, indices2)
    printFeatures(features, indices3)

  def testPlotTrend(self):
    print('\n Test plotTrend')
    years = [yr for yr in range(2000, 2010)]
    trend_values = [ (i + 1) for i in range(50, 60)]
    actual_values = []
    for tval in trend_values:
      if (tval % 2) == 0:
        actual_values.append(tval + 0.5)
      else:
        actual_values.append(tval - 0.5)

    plotTrend(years, actual_values, trend_values, 'YIELD')

  def testPlotTrueVSPredicted(self):
    print('\n Test plotTrueVSPredicted')
    Y_true = [ (i + 1) for i in range(50, 60)]
    Y_predicted = []
    for tval in Y_true:
      if (tval % 2) == 0:
        Y_predicted.append(tval + 0.5)
      else:
        Y_predicted.append(tval - 0.5)

    Y_true = np.asarray(Y_true)
    Y_predicted = np.asarray(Y_predicted)

    plotTrueVSPredicted(Y_true, Y_predicted)

  def runAllTests(self):
    print('\nTest Utility Functions BEGIN\n')
    self.testDateFormat()
    self.testGetYear()
    self.testGetMonth()
    self.testGetDay()
    self.testGetDekad()
    self.testCropIDToName()
    self.testCropNameToID()
    self.testPrintFeatures()
    self.testPlotTrend()
    self.testPlotTrueVSPredicted()
    print('\nTest Utility Functions END\n')

"""### Test Data Loading"""

#%%writefile test_data_loading.py
class TestDataLoader():
  def __init__(self, spark):
    cyp_config = CYPConfiguration()
    self.nuts_level = cyp_config.getNUTSLevel()
    data_sources = { 'SOIL' : self.nuts_level }
    cyp_config.setDataSources(data_sources)
    cyp_config.setDebugLevel(2)

    self.data_loader = CYPDataLoader(spark, cyp_config)

  def testDataLoad(self):
    print('\nTest loadData, loadAllData')
    soil_df = self.data_loader.loadData('SOIL', self.nuts_level)
    assert soil_df is not None
    soil_df.show(5)

    all_dfs = self.data_loader.loadAllData()
    soil_df = all_dfs['SOIL']
    assert soil_df is not None
    soil_df.show(5)

  def runAllTests(self):
    print('\nTest Data Loader BEGIN\n')
    self.testDataLoad()
    print('\nTest Data Loader END\n')

"""### Test Data Preprocessing"""

#%%writefile test_data_preprocessing.py
class TestDataPreprocessor():
  def __init__(self, spark):
    cyp_config = CYPConfiguration()
    cyp_config.setDebugLevel(2)
    self.preprocessor = CYPDataPreprocessor(spark, cyp_config)

    # create a small wofost data set
    # preprocessing currently extracts the year and dekad only
    self.wofost_df = spark.createDataFrame([(6, 'NL11', '19790110', 0.0, 0),
                                            (6, 'NL11', '19790121', 0.0, 0),
                                            (6, 'NL11', '19790331', 0.0, 50),
                                            (6, 'NL11', '19790510', 0.0, 100),
                                            (6, 'NL11', '19790821', 0.0, 150),
                                            (6, 'NL11', '19790831', 0.0, 201),
                                            (6, 'NL11', '19790910', 0.0, 201),
                                            (6, 'NL11', '19800110', 0.0, 0),
                                            (6, 'NL11', '19800121', 0.0, 0),
                                            (6, 'NL11', '19800331', 0.0, 50),
                                            (6, 'NL11', '19800610', 0.0, 100),
                                            (6, 'NL11', '19800721', 0.0, 150),
                                            (6, 'NL11', '19800831', 0.0, 200),
                                            (6, 'NL11', '19800910', 0.0, 201),
                                            (6, 'NL11', '19800921', 0.0, 201)],
                                           ['CROP_ID', 'IDREGION', 'DATE', 'POT_YB', 'DVS'])

    self.crop_season = None

    # Create a small meteo dekadal data set
    # Preprocessing currently extracts the year and dekad, and computes climate
    # water balance.
    self.meteo_dekdf = spark.createDataFrame([('NL11', '19790110', 1.2, 0.2),
                                              ('NL11', '19790121', 2.1, 0.2),
                                              ('NL11', '19790131', 0.2, 1.2),
                                              ('NL11', '19790210', 0.1, 2.0),
                                              ('NL11', '19790221', 0.0, 2.1),
                                              ('NL11', '19790228', 1.0, 0.8),
                                              ('NL11', '19790310', 1.1, 1.0),
                                              ('NL11', '19800110', 1.2, 0.2),
                                              ('NL11', '19800121', 2.1, 0.2),
                                              ('NL11', '19800131', 0.2, 1.2),
                                              ('NL11', '19800210', 0.1, 2.0),
                                              ('NL11', '19800221', 0.0, 2.1),
                                              ('NL11', '19800228', 1.0, 0.8),
                                              ('NL11', '19800310', 1.1, 1.0)],
                                             ['IDREGION', 'DATE', 'PREC', 'ET0'])

    # Create a small meteo daily data set
    # Preprocessing currently converts daily data to dekadal data by taking AVG
    # for all indicators except TMAX (MAX is used instead) and TMIN (MIN is used instead).
    query = 'select IDREGION, FYEAR, DEKAD, max(TMAX) as TMAX, min(TMIN) as TMIN, '
    query = query + ' bround(avg(TAVG), 2) as TAVG, bround(sum(PREC), 2) as PREC, '
    query = query + ' bround(sum(ET0), 2) as ET0, bround(avg(RAD), 2) as RAD, '
    query = query + ' bround(sum(CWB), 2) as CWB '
    self.meteo_daydf = spark.createDataFrame([('NL11', '19790101', 1.2, 0.2, 8.5, -1.2, 5.5, 10000.0),
                                              ('NL11', '19790102', 2.1, 0.2, 9.1, 0.3, 6.1, 12000.0),
                                              ('NL11', '19790103', 0.2, 1.2, 10.4, 1.2, 7.2, 14000.0),
                                              ('NL11', '19790104', 0.1, 2.0, 8.1, -1.5, 5.2, 10000.0),
                                              ('NL11', '19790105', 0.0, 2.1, 10.2, 1.0, 7.5, 13000.0),
                                              ('NL11', '19790106', 1.2, 0.2, 11.2, 2.5, 8.2, 16000.0),
                                              ('NL11', '19790112', 2.1, 0.5, 9.2, 0.5, 5.5, 12000.0),
                                              ('NL11', '19790113', 0.2, 1.1, 10.2, 1.4, 7.1, 14000.0),
                                              ('NL11', '19790114', 0.1, 2.0, 12.0, 3.2, 8.3, 15000.0),
                                              ('NL11', '19790115', 0.0, 1.5, 13.1, 4.5, 9.2, 17000.0),
                                              ('NL11', '19790122', 2.1, 0.5, 9.2, 0.5, 5.5, 12000.0),
                                              ('NL11', '19790123', 0.2, 1.1, 10.2, 1.4, 7.1, 14000.0),
                                              ('NL11', '19790124', 0.1, 2.0, 12.0, 3.2, 8.3, 15000.0),
                                              ('NL11', '19790125', 0.0, 1.5, 13.1, 4.5, 9.2, 17000.0),
                                              ('NL11', '19800101', 1.2, 0.2, 8.5, -1.2, 5.5, 10000.0),
                                              ('NL11', '19800102', 2.1, 0.2, 9.1, 0.3, 6.1, 12000.0),
                                              ('NL11', '19800103', 0.2, 1.2, 10.4, 1.2, 7.2, 14000.0),
                                              ('NL11', '19800104', 0.1, 2.0, 8.1, -1.5, 5.2, 10000.0),
                                              ('NL11', '19800105', 0.0, 2.1, 10.2, 1.0, 7.5, 13000.0),
                                              ('NL11', '19800106', 1.2, 0.2, 11.2, 2.5, 8.2, 16000.0),
                                              ('NL11', '19800112', 2.1, 0.5, 9.2, 0.5, 5.5, 12000.0),
                                              ('NL11', '19800113', 0.2, 1.1, 10.2, 1.4, 7.1, 14000.0),
                                              ('NL11', '19800114', 0.1, 2.0, 12.0, 3.2, 8.3, 15000.0),
                                              ('NL11', '19800115', 0.0, 1.5, 13.1, 4.5, 9.2, 17000.0),
                                              ('NL11', '19800122', 2.1, 0.5, 9.2, 0.5, 5.5, 12000.0),
                                              ('NL11', '19800123', 0.2, 1.1, 10.2, 1.4, 7.1, 14000.0),
                                              ('NL11', '19800124', 0.1, 2.0, 12.0, 3.2, 8.3, 15000.0),
                                              ('NL11', '19800125', 0.0, 1.5, 13.1, 4.5, 9.2, 17000.0)],
                                             ['IDREGION', 'DATE', 'PREC', 'ET0', 'TMAX', 'TMIN', 'TAVG', 'RAD'])

    # Create a small remote sensing data set
    # Preprocessing currently extracts the year and dekad
    self.rs_df1 = spark.createDataFrame([('NL11', '19790321', 0.47),
                                         ('NL11', '19790331', 0.49),
                                         ('NL11', '19790410', 0.55),
                                         ('NL11', '19790421', 0.49),
                                         ('NL11', '19790430', 0.64),
                                         ('NL11', '19800110', 0.42),
                                         ('NL11', '19800121', 2.43),
                                         ('NL11', '19800131', 0.41),
                                         ('NL11', '19800210', 0.42),
                                         ('NL11', '19800221', 0.44),
                                         ('NL11', '19800228', 0.45),
                                         ('NL11', '19800310', 2.43)],
                                        ['IDREGION', 'DATE', 'FAPAR'])

    self.rs_df2 = spark.createDataFrame([('FR10', '19790110', 0.42),
                                         ('FR10', '19790121', 0.43),
                                         ('FR10', '19790131', 0.41),
                                         ('FR10', '19790210', 0.42),
                                         ('FR10', '19790221', 0.44),
                                         ('FR10', '19790228', 0.45),
                                         ('FR10', '19790310', 0.47),
                                         ('FR10', '19790321', 0.49),
                                         ('FR10', '19790331', 0.55),
                                         ('FR10', '19790410', 0.62),
                                         ('FR10', '19790421', 0.66),
                                         ('FR10', '19800110', 0.42),
                                         ('FR10', '19800121', 2.43),
                                         ('FR10', '19800131', 0.41),
                                         ('FR10', '19800210', 0.42),
                                         ('FR10', '19800221', 0.44),
                                         ('FR10', '19800228', 0.45),
                                         ('FR10', '19800310', 2.43)],
                                        ['IDREGION', 'DATE', 'FAPAR'])
  
    self.crop_season_nuts3 = spark.createDataFrame([('FR101', '1979', 0, 27),
                                                    ('FR101', '1980', 27, 28),
                                                    ('FR102', '1979', 0, 27),
                                                    ('FR102', '1980', 27, 29)],
                                                   ['IDREGION', 'FYEAR', 'PREV_SEASON_END', 'SEASON_END'])

    # Create small yield data sets
    # Two formats are preprocessed: (1) year and yield are columns,
    # (2) years are columns with yield values in rows
    # Preprocessing currently converts (2) into 1
    self.yield_df1 = spark.createDataFrame([('potatoes', 'FR102', '1989', 29.75),
                                            ('potatoes', 'FR102', '1990', 25.44),
                                            ('potatoes', 'FR103', '1989', 30.2),
                                            ('potatoes', 'FR103', '1990', 29.9),
                                            ('sugarbeet', 'FR102', '1989', 66.0),
                                            ('sugarbeet', 'FR102', '1990', 55.0),
                                            ('sugarbeet', 'FR103', '1989', 69.3),
                                            ('sugarbeet', 'FR103', '1990', 59.1)],
                                           ['Crop', 'IDREGION', 'FYEAR', 'YIELD'])

    self.yield_df2 = spark.createDataFrame([('Total potatoes', 'NL11', 38.0, 40.5, 40.0),
                                            ('Total potatoes', 'NL12', 49.0, 44.0, 46.8),
                                            ('Spring barley', 'NL13', 4.6, 5.5, 6.6),
                                            ('Spring barley', 'NL12', 5.6, 6.1, 7.0)],
                                           ['Crop', 'IDREGION', '1994', '1995', '1996'])

  def testExtractYearDekad(self):
    print('WOFOST data after extracting year and dekad')
    print('-------------------------------------------')
    self.preprocessor.extractYearDekad(self.wofost_df).show(10)

  def testPreprocessWofost(self):
    print('WOFOST data after preprocessing')
    print('--------------------------------')
    self.wofost_df = self.wofost_df.filter(self.wofost_df['CROP_ID'] == 6).drop('CROP_ID')
    self.crop_season = self.preprocessor.getCropSeasonInformation(self.wofost_df,
                                                                  False)
    self.wofost_df = self.preprocessor.preprocessWofost(self.wofost_df,
                                                        self.crop_season,
                                                        False)
    self.wofost_df.show(5)
    self.crop_season.show(5)

  def testPreprocessMeteo(self):
    print('Meteo dekadal data after preprocessing')
    print('--------------------------------------')
    self.meteo_dekdf = self.preprocessor.preprocessMeteo(self.meteo_dekdf,
                                                         self.crop_season,
                                                         False)
    self.meteo_dekdf.show(5)

  def testPreprocessMeteoDaily(self):
    self.meteo_daydf = self.preprocessor.preprocessMeteo(self.meteo_daydf,
                                                         self.crop_season,
                                                         False)
    self.meteo_daydf = self.preprocessor.preprocessMeteoDaily(self.meteo_daydf)
    print('Meteo daily data after preprocessing')
    print('------------------------------------')
    self.meteo_daydf.show(5)

  def testPreprocessRemoteSensing(self):
    self.rs_df1 = self.preprocessor.preprocessRemoteSensing(self.rs_df1,
                                                            self.crop_season,
                                                            False)
    print('Remote sensing data after preprocessing')
    print('---------------------------------------')
    self.rs_df1.show(5)

  def testRemoteSensingNUTS2ToNUTS3(self):
    print('Remote sensing data before preprocessing')
    print('---------------------------------------')
    self.rs_df2.show()
    nuts3_regions = [reg[0] for reg in self.yield_df1.select('IDREGION').distinct().collect()]
    self.rs_df2 = self.preprocessor.remoteSensingNUTS2ToNUTS3(self.rs_df2, nuts3_regions)
    print('Remote sensing data at NUTS3')
    print('-----------------------------')
    self.rs_df2.show(5)

    self.rs_df2 = self.preprocessor.preprocessRemoteSensing(self.rs_df2,
                                                            self.crop_season_nuts3,
                                                            False)
    print('Remote sensing data after preprocessing')
    print('---------------------------------------')
    self.rs_df2.show(5)

  def testPreprocessYield(self):
    self.yield_df1 = self.preprocessor.preprocessYield(self.yield_df1, 7)
    print('Yield data format 1 after preprocessing')
    print('--------------------------------------')
    self.yield_df1.show(5)

    print('Yield data format 2 before preprocessing')
    print('----------------------------------------')
    self.yield_df2.show(5)

    self.yield_df2 = self.preprocessor.preprocessYield(self.yield_df2, 7)
    print('Yield data format 2 after preprocessing')
    print('----------------------------------------')
    self.yield_df2.show(5)

  def runAllTests(self):
    print('\nTest Data Preprocessor BEGIN\n')
    self.testExtractYearDekad()
    self.testPreprocessWofost()
    self.testPreprocessMeteo()
    self.testPreprocessMeteoDaily()
    self.testPreprocessRemoteSensing()
    self.testRemoteSensingNUTS2ToNUTS3()
    self.testPreprocessYield()
    print('\nTest Data Preprocessor END\n')

"""### Test Data Summary"""

#%%writefile test_data_summary.py
class TestDataSummarizer():
  def __init__(self, spark):
    cyp_config = CYPConfiguration()
    cyp_config.setDebugLevel(2)
    self.data_summarizer = CYPDataSummarizer(cyp_config)

    # create a small wofost data set
    self.wofost_df = spark.createDataFrame([('NL11', '1979', 14, 0.0, 0.0, 5.0),
                                            ('NL11', '1979', 15, 2.0, 1.0, 10.0),
                                            ('NL11', '1979', 16, 5.0, 4.0, 7.0),
                                            ('NL11', '1979', 17, 15.0, 12.0, 4.0),
                                            ('NL11', '1979', 18, 40.0, 35.0, 6.0),
                                            ('NL11', '1979', 19, 100.0, 80.0, 5.0),
                                            ('NL11', '1979', 20, 150.0, 120.0, 3.0),
                                            ('NL11', '1980', 13, 0.0, 0.0, 15.0),
                                            ('NL11', '1980', 14, 5.0, 2.0, 12.0),
                                            ('NL11', '1980', 15, 15.0, 12.0, 10.0),
                                            ('NL11', '1980', 16, 50.0, 40.0, 8.0),
                                            ('NL11', '1980', 17, 100.0, 80.0, 4.0),
                                            ('NL11', '1980', 18, 200.0, 140.0, 5.0),
                                            ('NL11', '1980', 19, 200.0, 150.0, 12.0),
                                            ('NL11', '1980', 20, 200.0, 150.0, 12.0)],
                                           ['IDREGION', 'FYEAR', 'DEKAD', 'POT_YB', 'WLIM_YB', 'RSM'])

    self.wofost_df2 = spark.createDataFrame([('NL11', '1979', 12, '1979', 22, 0),
                                             ('NL11', '1979', 13, '1979', 23, 1),
                                             ('NL11', '1979', 14, '1979', 24, 4),
                                             ('NL11', '1979', 15, '1979', 25, 70),
                                             ('NL11', '1979', 16, '1979', 26, 101),
                                             ('NL11', '1979', 19, '1979', 29, 150),
                                             ('NL11', '1979', 21, '1979', 31, 180),
                                             ('NL11', '1979', 23, '1979', 33, 200),
                                             ('NL11', '1979', 24, '1979', 34, 201),
                                             ('NL11', '1979', 25, '1979', 35, 201),
                                             ('NL11', '1980', 12, '1980', 22, 0),
                                             ('NL11', '1980', 13, '1980', 23, 2),
                                             ('NL11', '1980', 14, '1980', 24, 15),
                                             ('NL11', '1980', 15, '1980', 25, 80),
                                             ('NL11', '1980', 16, '1980', 26, 99),
                                             ('NL11', '1980', 19, '1980', 29, 140),
                                             ('NL11', '1980', 21, '1980', 31, 170),
                                             ('NL11', '1980', 23, '1980', 33, 195),
                                             ('NL11', '1980', 24, '1980', 34, 201),
                                             ('NL11', '1980', 25, '1980', 35, 201)],
                                           ['IDREGION', 'FYEAR', 'DEKAD', 'CAMPAIGN_YEAR', 'CAMPAIGN_DEKAD', 'DVS'])

    # Create a small meteo dekadal data set
    self.meteo_df = spark.createDataFrame([('NL11', '1979', 1, '1979', 11, 1.2, 8.5, -1.2, 5.5, 10000.0),
                                           ('NL11', '1979', 2, '1979', 12, 2.1, 9.1, 0.3, 6.1, 12000.0),
                                           ('NL11', '1979', 3, '1979', 13, 0.2, 10.4, 1.2, 7.2, 14000.0),
                                           ('NL11', '1979', 4, '1979', 14, 0.1, 8.1, -1.5, 5.2, 10000.0),
                                           ('NL11', '1979', 5, '1979', 15, 0.0, 10.2, 1.0, 7.5, 13000.0),
                                           ('NL12', '1979', 1, '1979', 12, 1.2, 11.2, 2.5, 8.2, 16000.0),
                                           ('NL12', '1979', 2, '1979', 13, 2.1, 9.2, 0.5, 5.5, 12000.0),
                                           ('NL12', '1979', 3, '1979', 14, 0.2, 10.2, 1.4, 7.1, 14000.0),
                                           ('NL12', '1979', 4, '1979', 15, 0.1, 12.0, 3.2, 8.3, 15000.0),
                                           ('NL12', '1979', 5, '1979', 16, 0.0, 13.1, 4.5, 9.2, 17000.0)],
                                          ['IDREGION', 'FYEAR', 'DEKAD', 'CAMPAIGN_YEAR', 'CAMPAIGN_DEKAD', 'PREC', 'TMAX', 'TMIN', 'TAVG', 'RAD'])

    # Create a small remote sensing data set
    # Preprocessing currently extracts the year and dekad
    self.rs_df = spark.createDataFrame([('NL11', '1979', 1, 0.42),
                                         ('NL11', '1979', 2, 0.41),
                                         ('NL11', '1979', 3, 0.42),
                                         ('NL11', '1980', 1, 0.44),
                                         ('NL11', '1980', 2, 0.45),
                                        ('NL11', '1980', 3, 0.43)],
                                        ['IDREGION', 'FYEAR', 'DEKAD', 'FAPAR'])

    # Create a small yield data set
    self.yield_df = spark.createDataFrame([(7, 'FR102', '1989', 29.75),
                                           (7, 'FR102', '1990', 25.44),
                                           (7, 'FR103', '1989', 30.2),
                                           (7, 'FR103', '1990', 29.9),
                                           (6, 'FR102', '1989', 66.0),
                                           (6, 'FR102', '1990', 55.0),
                                           (6, 'FR103', '1989', 69.3),
                                           (6, 'FR103', '1990', 59.1)],
                                          ['CROP_ID', 'IDREGION', 'FYEAR', 'YIELD'])

  def testWofostDVSSummary(self):
    print('WOFOST Crop Calendar Summary using DVS')
    print('-----------------------------')
    self.data_summarizer.wofostDVSSummary(self.wofost_df2).show()

  def testWofostIndicatorsSummary(self):
    print('WOFOST indicators summary')
    print('--------------------------')
    min_cols = ['IDREGION']
    max_cols = ['IDREGION', 'POT_YB', 'WLIM_YB']
    avg_cols = ['IDREGION', 'RSM']
    self.data_summarizer.indicatorsSummary(self.wofost_df, min_cols, max_cols, avg_cols).show()

  def testMeteoIndicatorsSummary(self):
    print('Meteo indicators summary')
    print('-------------------------')
    meteo_cols = self.meteo_df.columns[3:]
    min_cols = ['IDREGION'] + meteo_cols
    max_cols = ['IDREGION'] + meteo_cols
    avg_cols = ['IDREGION'] + meteo_cols
    self.data_summarizer.indicatorsSummary(self.meteo_df, min_cols, max_cols, avg_cols).show()

  def testRemoteSensingSummary(self):
    print('Remote sensing indicators summary')
    print('----------------------------------')
    rs_cols = ['FAPAR']
    min_cols = ['IDREGION'] + rs_cols
    max_cols = ['IDREGION'] + rs_cols
    avg_cols = ['IDREGION'] + rs_cols
    self.data_summarizer.indicatorsSummary(self.rs_df, min_cols, max_cols, avg_cols).show()

  def testYieldSummary(self):
    crop = 'potatoes'
    print('Yield summary for', crop)
    print('-----------------------------')
    crop_id = cropNameToID(crop_id_dict, crop)
    self.yield_df = self.yield_df.filter(self.yield_df.CROP_ID == crop_id)
    self.data_summarizer.yieldSummary(self.yield_df).show()

  def runAllTests(self):
    print('\nTest Data Summarizer BEGIN\n')
    self.testWofostDVSSummary()
    self.testWofostIndicatorsSummary()
    self.testMeteoIndicatorsSummary()
    self.testRemoteSensingSummary()
    self.testYieldSummary()
    print('\nTest Data Summarizer END\n')

"""### Test Yield Trend Estimation"""

#%%writefile test_yield_trend.py
class TestYieldTrendEstimator():
  def __init__(self, yield_df):
    # TODO: Create a small yield data set
    self.yield_df = yield_df
    cyp_config = CYPConfiguration()
    self.verbose = 2
    cyp_config.setDebugLevel(self.verbose)
    self.trend_est = CYPYieldTrendEstimator(cyp_config)

  def testYieldTrendTwoRegions(self):
    print('\nFind the optimal trend window and estimate trend for first 2 regions')
    pd_yield_df = self.yield_df.toPandas()
    regions = sorted(pd_yield_df['IDREGION'].unique())
    reg1 = regions[0]
    pd_reg1_df = pd_yield_df[pd_yield_df['IDREGION'] == reg1]
    reg1_num_years = len(pd_reg1_df.index)
    reg1_max_year = pd_reg1_df['FYEAR'].max()
    reg1_min_year = pd_reg1_df['FYEAR'].min()

    if (self.verbose > 2):
      print('\nPrint Yield Trend Rounds')
      print('------------------------')

    trend_windows = [5]
    self.trend_est.printYieldTrendRounds(self.yield_df, reg1, trend_windows)

    if (self.verbose > 1):
      print('\n Fixed Trend Window prediction for region 1')
      print('---------------------------------------------------')
    trend_window = 5
    pd_fixed_win_df = self.trend_est.getFixedWindowTrend(self.yield_df, reg1, reg1_max_year,
                                                         trend_window)
    if (self.verbose > 1):
      print(pd_fixed_win_df.head(1))

    if (self.verbose > 1):
      print('\n Optimal Trend Window and prediction for region 1')
      print('---------------------------------------------------')
  
    trend_windows = [5, 7]
    pd_opt_win_df = self.trend_est.getOptimalWindowTrend(self.yield_df, reg1, reg1_max_year,
                                                         trend_windows)
    if (self.verbose > 1):
      print(pd_opt_win_df.head(1))

    reg2 = regions[1]
    pd_reg2_df = pd_yield_df[pd_yield_df['IDREGION'] == reg2]
    reg2_num_years = len(pd_reg2_df.index)
    reg2_max_year = pd_reg2_df['FYEAR'].max()
    reg2_min_year = pd_reg2_df['FYEAR'].min()

    if (self.verbose > 1):
      print('\n Fixed Trend Window prediction for region 2')
      print('---------------------------------------------------')
    trend_window = 5
    pd_fixed_win_df = self.trend_est.getFixedWindowTrend(self.yield_df, reg2, reg2_max_year,
                                                         trend_window)
    if (self.verbose > 1):
      print(pd_fixed_win_df.head(1))

    if (self.verbose > 1):
      print('\n Optimal Trend Window and prediction for region 2')
      print('---------------------------------------------------')

    pd_opt_win_df = self.trend_est.getOptimalWindowTrend(self.yield_df, reg2, reg2_max_year,
                                                         trend_windows)
    if (self.verbose > 1):
      print(pd_opt_win_df.head(1))

  def testYieldTrendAllRegions(self):
    print('\nYield trend estimation for all regions')

    print('\nOptimal Trend Windows')
    pd_trend_df = self.trend_est.getOptimalWindowTrendFeatures(self.yield_df)
    print(pd_trend_df.head(5))

    print('\nFixed Trend Window')
    pd_trend_df = self.trend_est.getFixedWindowTrendFeatures(self.yield_df)
    print(pd_trend_df.head(5))

  def runAllTests(self):
    print('\nTest Yield Trend Estimator BEGIN\n')
    self.testYieldTrendTwoRegions()
    self.testYieldTrendAllRegions()
    print('\nTest Yield Trend Estimator END\n')

"""### Test custom train, test split"""

#%%writefile test_train_test_split.py
class TestCustomTrainTestSplit:
  def __init__(self, yield_df):
    cyp_config = CYPConfiguration()
    self.verbose = 2
    cyp_config.setDebugLevel(self.verbose)
    self.yield_df = yield_df
    self.trTsSplitter = CYPTrainTestSplitter(cyp_config)

  def testCustomTrainTestSplit(self):
    print('\nTest customTrainTestSplit')
    test_fraction = 0.2
    regions = [reg[0] for reg in self.yield_df.select('IDREGION').distinct().collect()]
    num_regions = len(regions)
    test_years = self.trTsSplitter.trainTestSplit(self.yield_df, test_fraction, True)
    all_years = [yr[0] for yr in self.yield_df.select('FYEAR').distinct().collect()]
    yield_train_df = self.yield_df.filter(~self.yield_df['FYEAR'].isin(test_years))
    yield_test_df = self.yield_df.filter(self.yield_df['FYEAR'].isin(test_years))

    if(self.verbose > 1):
      print('\nCustom training, test split using yield trend')
      print('---------------------------------------------')
      print('Estimated size of test data', num_regions * np.floor(len(all_years) * test_fraction))
      print('Data Size:', yield_train_df.count(), yield_test_df.count())
      print('Test years:', test_years)

    test_years = self.trTsSplitter.trainTestSplit(self.yield_df, test_fraction, False)

    if(self.verbose > 1):
      print('\ncustom training, test split without yield trend')
      print('------------------------------------------------')
      print('Estimated size of test data', num_regions * np.floor(len(all_years) * test_fraction))
      print('Data Size:', yield_train_df.count(), yield_test_df.count())
      print('Test years:', test_years)

  def testCustomKFoldValidationSplit(self):
    print('\nTest customKFoldValidationSplit')
    test_fraction = 0.2
    num_folds = 5
    test_years = self.trTsSplitter.trainTestSplit(self.yield_df, test_fraction, 'Y')
    yield_train_df = self.yield_df.filter(~self.yield_df['FYEAR'].isin(test_years))
    yield_test_df = self.yield_df.filter(self.yield_df['FYEAR'].isin(test_years))
    yield_cols = yield_train_df.columns
    pd_yield_train_df = yield_train_df.toPandas()
    Y_train_full = pd_yield_train_df[yield_cols].values

    custom_cv = self.trTsSplitter.customKFoldValidationSplit(Y_train_full, num_folds)

  def runAllTests(self):
    print('\nTest Custom Train, Test Splitter BEGIN\n')
    self.testCustomTrainTestSplit()
    self.testCustomKFoldValidationSplit()
    print('\nTest Custom Train, Test Splitter END\n')

"""## Run Workflow

### Set Configuration
"""

if (test_env == 'notebook'):
  cyp_config = CYPConfiguration()

  if (run_tests):
    test_util = TestUtil(spark)
    test_util.runAllTests()

  my_config = {
      'crop_name' : 'sugarbeet',
      'season_crosses_calendar_year' : 'N',
      'country_code' : 'NL',
      'data_sources' : [ 'WOFOST', 'METEO_DAILY', 'SOIL', 'YIELD'],
      'data_path' : '.',
      'output_path' : '.',
      'nuts_level' : 'NUTS2',
      'use_yield_trend' : 'Y',
      'predict_yield_residuals' : 'N',
      'trend_windows' : [5, 7, 10],
      'use_centroids' : 'N',
      'use_remote_sensing' : 'Y',
      'early_season_prediction' : 'Y',
      'early_season_end_dekad' : -6,
      'save_features' : 'N',
      'use_saved_features' : 'N',
      'save_predictions' : 'N',
      'use_saved_predictions' : 'N',
      'compare_with_mcyfs' : 'Y',
      'debug_level' : 2,
  }

  cyp_config.updateConfiguration(my_config)
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  nuts_level = cyp_config.getNUTSLevel()
  debug_level = cyp_config.getDebugLevel()
  use_saved_predictions = cyp_config.useSavedPredictions()
  use_saved_features = cyp_config.useSavedFeatures()
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()

  print('##################')
  print('# Configuration  #')
  print('##################')
  output_path = cyp_config.getOutputPath()
  log_file = getLogFilename(crop, country, use_yield_trend,
                            early_season_prediction, early_season_end)
  log_fh = open(output_path + '/' + log_file, 'w+')
  cyp_config.printConfig(log_fh)

"""### Load and Preprocess Data

"""

if ((test_env == 'notebook') and
    (not use_saved_predictions) and
    (not use_saved_features)):

  print('#################')
  print('# Data Loading  #')
  print('#################')

  if (run_tests):
    test_loader = TestDataLoader(spark)
    test_loader.runAllTests()

  cyp_loader = CYPDataLoader(spark, cyp_config)
  data_dfs = cyp_loader.loadAllData()

  print('#######################')
  print('# Data Preprocessing  #')
  print('#######################')

  if (run_tests):
    test_preprocessor = TestDataPreprocessor(spark)
    test_preprocessor.runAllTests()

  cyp_preprocessor = CYPDataPreprocessor(spark, cyp_config)
  data_dfs = preprocessData(cyp_config, cyp_preprocessor, data_dfs)

"""### Split Data into Training and Test Sets"""

if ((test_env == 'notebook') and
    (not use_saved_predictions) and
    (not use_saved_features)):

  print('###########################')
  print('# Training and Test Split #')
  print('###########################')

  if (run_tests):
    yield_df = data_dfs['YIELD']
    test_custom = TestCustomTrainTestSplit(yield_df)
    test_custom.runAllTests()

  prep_train_test_dfs, test_years = splitDataIntoTrainingTestSets(cyp_config, data_dfs, log_fh)

"""### Summarize Data"""

if ((test_env == 'notebook') and
    (not use_saved_predictions) and
    (not use_saved_features)):

  print('#################')
  print('# Data Summary  #')
  print('#################')

  if (run_tests):
    test_summarizer = TestDataSummarizer(spark)
    test_summarizer.runAllTests()

  cyp_summarizer = CYPDataSummarizer(cyp_config)
  summary_dfs = summarizeData(cyp_config, cyp_summarizer, prep_train_test_dfs)

"""### Create Features"""

if ((test_env == 'notebook') and
    (not use_saved_predictions) and
    (not use_saved_features)):

  print('###################')
  print('# Feature Design  #')
  print('###################')

  # WOFOST, Meteo and Remote Sensing Features
  cyp_featurizer = CYPFeaturizer(cyp_config)
  pd_feature_dfs = createFeatures(cyp_config, cyp_featurizer,
                                  prep_train_test_dfs, summary_dfs, log_fh)

  # yield trend features
  if (use_yield_trend):
    yield_train_df = prep_train_test_dfs['YIELD'][0]
    yield_test_df = prep_train_test_dfs['YIELD'][1]

    if (run_tests):
      test_yield_trend = TestYieldTrendEstimator(yield_train_df)
      test_yield_trend.runAllTests()

    cyp_trend_est = CYPYieldTrendEstimator(cyp_config)
    pd_yield_train_ft, pd_yield_test_ft = createYieldTrendFeatures(cyp_config, cyp_trend_est,
                                                                   yield_train_df, yield_test_df,
                                                                   test_years)
    pd_feature_dfs['YIELD_TREND'] = [pd_yield_train_ft, pd_yield_test_ft]

"""### Combine Features and Labels"""

if ((test_env == 'notebook') and
    (not use_saved_predictions) and
    (not use_saved_features)):

  join_cols = ['IDREGION', 'FYEAR']
  pd_train_df, pd_test_df = combineFeaturesLabels(cyp_config, sqlContext,
                                                  prep_train_test_dfs, pd_feature_dfs,
                                                  join_cols, log_fh)

"""### Apply Machine Learning using scikit learn

"""

if ((test_env == 'notebook') and
    (not use_saved_predictions)):

  if (use_saved_features):
    pd_train_df, pd_test_df = loadSavedFeaturesLabels(cyp_config, spark)

  print('\n###################################')
  print('# Machine Learning using sklearn  #')
  print('###################################')

  pd_ml_predictions = getMachineLearningPredictions(cyp_config, pd_train_df, pd_test_df, log_fh)
  save_predictions = cyp_config.savePredictions()
  if (save_predictions):
    saveMLPredictions(cyp_config, sqlContext, pd_ml_predictions)

"""### Compare Predictions with JRC Predictions"""

if (test_env == 'notebook'):
  if (use_saved_predictions):
    pd_ml_predictions = loadSavedPredictions(cyp_config, spark)

  compareWithMCYFS = cyp_config.compareWithMCYFS()
  if (compareWithMCYFS):
    comparePredictionsWithMCYFS(sqlContext, cyp_config, pd_ml_predictions, log_fh)

  log_fh.close()

"""### Python Script Main

To be used in environment supporting command-line arguments
"""

#%%writefile main.py
import sys
import argparse

def main():
  print('##################')
  print('# Configuration  #')
  print('##################')

  parser = argparse.ArgumentParser(prog='mlbaseline.py')

  # Some command-line argument names are slightly different
  # from configuration option names for brevity.
  args_dict = {
      '--crop' : { 'type' : str,
                   'default' : 'potatoes',
                   'help' : 'crop name (default: potatoes)',
                 },
      '--crosses-calendar-year' : { 'type' : str,
                                    'default' : 'N',
                                    'choices' : ['Y', 'N'],
                                    'help' : 'crop growing season crosses calendar year boundary (default: N)',
                                  },
      '--country' : { 'type' : str,
                      'default' : 'NL',
                      'choices' : ['NL', 'DE', 'FR'],
                      'help' : 'country code (default: NL)',
                    },
      '--nuts-level' : { 'type' : str,
                         'default' : 'NUTS2',
                         'choices' : ['NUTS2', 'NUTS3'],
                         'help' : 'country code (default: NL)',
                       },
      '--data-path' : { 'type' : str,
                        'default' : '.',
                        'help' : 'path to data files (default: .)',
                       },
      '--output-path' : { 'type' : str,
                          'default' : '.',
                          'help' : 'path to output files (default: .)',
                        },
      '--yield-trend' : { 'type' : str,
                          'default' : 'N',
                          'choices' : ['Y', 'N'],
                          'help' : 'estimate and use yield trend (default: N)',
                        },
      '--optimal-trend-window' : { 'type' : str,
                                   'default' : 'N',
                                   'choices' : ['Y', 'N'],
                                   'help' : 'find optimal trend window for each year (default: N)',
                                 },
      '--predict-residuals' : { 'type' : str,
                                'default' : 'N',
                                'choices' : ['Y', 'N'],
                                'help' : 'predict yield residuals instead of full yield (default: N)',
                              },
      '--early-season' : { 'type' : str,
                           'default' : 'N',
                           'choices' : ['Y', 'N'],
                           'help' : 'early season prediction (default: N)',
                         },
      '--early-season-end' : { 'type' : int,
                               'default' : 15,
                               'help' : 'early season end dekad (default: 15)',
                             },
      '--centroids' : { 'type' : str,
                        'default' : 'N',
                        'choices' : ['Y', 'N'],
                        'help' : 'use centroid coordinates and distance to coast (default: N)',
                      },
      '--remote-sensing' : { 'type' : str,
                             'default' : 'Y',
                             'choices' : ['Y', 'N'],
                             'help' : 'use remote sensing data (default: Y)',
                           },
      '--save-features' : { 'type' : str,
                            'default' : 'N',
                            'choices' : ['Y', 'N'],
                            'help' : 'save features to a CSV file (default: N)',
                          },
      '--use-saved-features' : { 'type' : str,
                                 'default' : 'N',
                                 'choices' : ['Y', 'N'],
                                 'help' : 'use features from a CSV file (default: N). Set ',
                               },
      '--save-predictions' : { 'type' : str,
                               'default' : 'Y',
                               'choices' : ['Y', 'N'],
                               'help' : 'save predictions to a CSV file (default: Y)',
                             },
      '--use-saved-predictions' : { 'type' : str,
                                    'default' : 'N',
                                    'choices' : ['Y', 'N'],
                                    'help' : 'use predictions from a CSV file (default: N)',
                                  },
      '--compare-with-mcyfs' : { 'type' : str,
                                 'default' : 'N',
                                 'choices' : ['Y', 'N'],
                                 'help' : 'compare predictions with MCYFS (default: N)',
                               },
      '--debug-level' : { 'type' : int,
                          'default' : 0,
                          'choices' : range(4),
                          'help' : 'amount of debug information to print (default: 0)',
                        },
  }

  for arg in args_dict:
    arg_config = args_dict[arg]
    # add cases if other argument settings are used
    if ('choices' in arg_config):
      parser.add_argument(arg, type=arg_config['type'], default=arg_config['default'],
                          choices=arg_config['choices'], help=arg_config['help'])
    else:
      parser.add_argument(arg, type=arg_config['type'], default=arg_config['default'],
                          help=arg_config['help'])

  if (run_tests):
    test_util = TestUtil(spark)
    test_util.runAllTests()

  args = parser.parse_args()
  cyp_config = CYPConfiguration()

  # must be in sync with args_dict used to parse args
  config_update = {
      'crop_name' : args.crop,
      'season_crosses_calendar_year' : args.crosses_calendar_year,
      'country_code' : args.country,
      'nuts_level' : args.nuts_level,
      'data_path' : args.data_path,
      'output_path' : args.output_path,
      'use_yield_trend' : args.yield_trend,
      'find_optimal_trend_window' : args.optimal_trend_window,
      'predict_yield_residuals' : args.predict_residuals,
      'use_centroids' : args.centroids,
      'use_remote_sensing' : args.remote_sensing,
      'early_season_prediction' : args.early_season,
      'early_season_end_dekad' : args.early_season_end,
      'save_features' : args.save_features,
      'use_saved_features' : args.use_saved_features,
      'save_predictions' : args.save_predictions,
      'use_saved_predictions' : args.use_saved_predictions,
      'compare_with_mcyfs' : args.compare_with_mcyfs,
      'debug_level' : args.debug_level,
  }

  cyp_config.updateConfiguration(config_update)
  crop = cyp_config.getCropName()
  country = cyp_config.getCountryCode()
  nuts_level = cyp_config.getNUTSLevel()
  debug_level = cyp_config.getDebugLevel()
  use_saved_predictions = cyp_config.useSavedPredictions()
  use_saved_features = cyp_config.useSavedFeatures()
  use_yield_trend = cyp_config.useYieldTrend()
  early_season_prediction = cyp_config.earlySeasonPrediction()
  early_season_end = cyp_config.getEarlySeasonEndDekad()

  output_path = cyp_config.getOutputPath()
  log_file = getLogFilename(crop, country, use_yield_trend,
                            early_season_prediction, early_season_end)
  log_fh = open(output_path + '/' + log_file, 'w+')
  cyp_config.printConfig(log_fh)

  if (not use_saved_predictions):
    if (not use_saved_features):
      print('#################')
      print('# Data Loading  #')
      print('#################')

      if (run_tests):
        test_loader = TestDataLoader(spark)
        test_loader.runAllTests()

      cyp_loader = CYPDataLoader(spark, cyp_config)
      data_dfs = cyp_loader.loadAllData()

      print('#######################')
      print('# Data Preprocessing  #')
      print('#######################')

      if (run_tests):
        test_preprocessor = TestDataPreprocessor(spark)
        test_preprocessor.runAllTests()

      cyp_preprocessor = CYPDataPreprocessor(spark, cyp_config)
      data_dfs = preprocessData(cyp_config, cyp_preprocessor, data_dfs)

      print('###########################')
      print('# Training and Test Split #')
      print('###########################')

      if (run_tests):
        yield_df = data_dfs['YIELD']
        test_custom = TestCustomTrainTestSplit(yield_df)
        test_custom.runAllTests()

      prep_train_test_dfs, test_years = splitDataIntoTrainingTestSets(cyp_config, data_dfs, log_fh)

      print('#################')
      print('# Data Summary  #')
      print('#################')

      if (run_tests):
        test_summarizer = TestDataSummarizer(spark)
        test_summarizer.runAllTests()

      cyp_summarizer = CYPDataSummarizer(cyp_config)
      summary_dfs = summarizeData(cyp_config, cyp_summarizer, prep_train_test_dfs)

      print('###################')
      print('# Feature Design  #')
      print('###################')

      # WOFOST, Meteo and Remote Sensing Features
      cyp_featurizer = CYPFeaturizer(cyp_config)
      pd_feature_dfs = createFeatures(cyp_config, cyp_featurizer,
                                      prep_train_test_dfs, summary_dfs, log_fh)

      # yield trend features
      if (use_yield_trend):
        yield_train_df = prep_train_test_dfs['YIELD'][0]
        yield_test_df = prep_train_test_dfs['YIELD'][1]

        if (run_tests):
          test_yield_trend = TestYieldTrendEstimator(yield_train_df)
          test_yield_trend.runAllTests()

        cyp_trend_est = CYPYieldTrendEstimator(cyp_config)
        pd_yield_train_ft, pd_yield_test_ft = createYieldTrendFeatures(cyp_config, cyp_trend_est,
                                                                       yield_train_df, yield_test_df,
                                                                       test_years)
        pd_feature_dfs['YIELD_TREND'] = [pd_yield_train_ft, pd_yield_test_ft]

      # combine features
      join_cols = ['IDREGION', 'FYEAR']
      pd_train_df, pd_test_df = combineFeaturesLabels(cyp_config, sqlContext,
                                                      prep_train_test_dfs, pd_feature_dfs,
                                                      join_cols, log_fh)

    # use saved features
    else:
      pd_train_df, pd_test_df = loadSavedFeaturesLabels(cyp_config, spark)

    print('###################################')
    print('# Machine Learning using sklearn  #')
    print('###################################')

    pd_ml_predictions = getMachineLearningPredictions(cyp_config, pd_train_df, pd_test_df, log_fh)
    save_predictions = cyp_config.savePredictions()
    if (save_predictions):
      saveMLPredictions(cyp_config, sqlContext, pd_ml_predictions)

  # use saved predictions
  else:
    pd_ml_predictions = loadSavedPredictions(cyp_config, spark)

  # compare with MCYFS
  compareWithMCYFS = cyp_config.compareWithMCYFS()
  if (compareWithMCYFS):
    comparePredictionsWithMCYFS(sqlContext, cyp_config, pd_ml_predictions, log_fh)

  log_fh.close()

if __name__ == '__main__':
    main()