#!/usr/bin/env python

import logging

import datetime

import sys

import json

import warnings

sys.path.append('../')

warnings.filterwarnings("ignore")

import pandas as pd

from scipy import stats

from sklearn.feature_extraction.text import TfidfVectorizer

from sklearn.model_selection import RandomizedSearchCV

import lightgbm

from utils.splits import set_group_splits

from args import args

logger = logging.getLogger(__name__)

logger.setLevel(logging.INFO)

sh = logging.StreamHandler()

sh.setFormatter(logging.Formatter('%(levelname)s:%(name)s: %(message)s'))

logger.addHandler(sh)

if __name__ == '__main__':

  if len(sys.argv) != 2:

    logger.error(f"Usage: {sys.argv[0]} task_name (ia|ps)")

    sys.exit(1)

  task = sys.argv[1]

  if task not in ['ia', 'ps']:

    logger.error("Task values are either ia (imminent admission) or ps (prolonged stay)")

    sys.exit(1)

  ori_df = pd.read_csv(args.dataset_csv, usecols=args.cols, parse_dates=args.dates)

  if task == 'ia':

    logger.info(f"Running hyperparameter search for Imminent Admission Prediction task")

    task_df = ori_df.loc[(ori_df['imminent_adm_label'] != -1)][args.imminent_adm_cols].reset_index(drop=True)

    label = 'imminent_adm_label'

  if task == 'ps':

    logger.info(f"Running hyperparameter search for Prolonged Stay Prediction task ")

    task_df = ori_df.loc[(ori_df['chartinterval'] != 0)][args.prolonged_stay_cols].reset_index(drop=True)

    label = 'prolonged_stay_label'

  df = set_group_splits(task_df.copy(), group_col='hadm_id', seed=643)

  vectorizer = TfidfVectorizer(min_df=args.min_freq, analyzer=str.split, sublinear_tf=True, ngram_range=(2,2))

  x_train = vectorizer.fit_transform(df.loc[(df['split'] == 'train')]['processed_note'])

  y_train = df.loc[(df['split'] == 'train')][label].to_numpy()

  clf_params = {

      'objective': 'binary',

      'metric': 'binary_logloss',

      'is_unbalance': True,

  }

  clf = lightgbm.LGBMClassifier(**clf_params)

  param_space = {

    'num_leaves': stats.randint(30, 60),

    'bagging_fraction': stats.uniform(0.2, 0.7),

    'learning_rate': stats.uniform(0.1, 0.9),

    'min_data_in_leaf': stats.randint(2, 20),

    'max_bin': stats.randint(3, 20),

    'boosting': ['gbdt', 'dart'],

    'bagging_freq': stats.randint(3, 31),

    'max_depth': stats.randint(0, 11),

    'feature_fraction': stats.uniform(0.2, 0.7),

    'lambda_l1': stats.uniform(0, 10),

    'num_iterations': stats.randint(100, 200),

  }

  random_search = RandomizedSearchCV(clf, param_space, n_iter=200, cv=5, iid=False, verbose=1, n_jobs=32)

  logger.info("Starting random search...")

  t1 = datetime.datetime.now()

  random_search.fit(x_train, y_train)

  dt = datetime.datetime.now() - t1

  params_file = args.workdir/f'{task}_best_params.json'

  logger.info(f"Random search completed. Took {dt.days} days, {dt.seconds//3600} hours, and {(dt.seconds//60)%60} minutes. Writing best params to {params_file}")

  json.dump(random_search.best_params_, params_file.open('w'))