#!/usr/bin/env python3
# Author: Armit
# Create Time: 2022/11/22 

from argparse import ArgumentParser

from sklearnex import patch_sklearn ; patch_sklearn()
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

from data import get_data, FEATURE_NUM
from utils import show_clf_metrics


def lr(args):
  X, y = get_data(args.limit, FEATURE_NUM)
  X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.7, random_state=42)
  print(f'dataset: {len(X_train)} for train. {len(X_test)} samples for test')

  model = LogisticRegression(solver=args.solver, max_iter=233,
                             tol=1e-3, C=args.C, penalty='l2',
                             verbose=1, n_jobs=4, random_state=42)

  model.fit(X_train, y_train)
  y_pred = model.predict(X_test)

  show_clf_metrics(y_test, y_pred)


if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('-M', '--solver', default='lbfgs', choices=['newton-cg', 'lbfgs', 'liblinear', 'sag', 'saga'])
  parser.add_argument('-C', default=0.8, type=float, help='smaller values specify stronger regularization')
  parser.add_argument('-N', '--limit', default=20000, type=int, help='limit dataset size')
  args = parser.parse_args()

  lr(args)