"""
File contains functions used to traing Logistic Regression.

Author: Dominik Macko
"""

from typing import Dict, Any, Union, Callable, Tuple, Optional

import numpy as np
import pandas as pd
#from sklearn.linear_model import LogisticRegression
# this should be then easier to interpret - just a simple scikit-learn wrapper
from interpret.glassbox import LogisticRegression
import joblib

from androidcrypto.analysis.training.config import TrainingTaskConfig
from androidcrypto.analysis.training.utils import is_multiclass
from androidcrypto.analysis.training.generic import train_model, train_gridsearchcv_model

def logistic_regression_best_params_surroundings(best_params: Dict[str, Any]) -> Dict[str, Any]:
    """Get best parameters surroundings for random forest."""
    
    alpha = best_params["l1_ratio"]
    C = best_params["C"]
    return {
        "l1_ratio": [max(0.0, alpha - 0.05), alpha, min(alpha + 0.05, 1.0)],
        "C": [1/2 * C, C, 3/2 * C]
    }

def train_logistic_regression(train_X: np.array,
                              train_y: np.array,
                              scoring: Union[str, Callable[[Any, np.array, np.array], int]]="f1_macro",
                              n_jobs: int=8,
                              verbose: int=3,
                              seed: int=42,
                              cv_splits: int=5
                             ) -> Tuple[LogisticRegression, pd.DataFrame]:
    """Trains logistic regression by searching for optimal l1 ratio and inverse regularization strength C.
    
    train_X - training set features
    train_y - training set targets
    scoring - scikit scoring function to use
    n_jobs - threads to use
    verbose - scikit verbose level
    seed - random seed to use
    
    returns (model, history dataframe)
    """
    
    grid = {
        "l1_ratio": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0],
        "C": [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]
    }
    return train_model(
        LogisticRegression(
            penalty="elasticnet",
            solver="saga", # because we want to be faster + possibly multinomial + l1 and l2
            class_weight="balanced",
            random_state=seed,
            max_iter=1000
        ),
        train_gridsearchcv_model,
        train_gridsearchcv_model,
        grid,
        logistic_regression_best_params_surroundings,
        train_X,
        train_y,
        scoring=scoring,
        n_jobs=n_jobs,
        verbose=verbose,
        cv_splits=cv_splits
    )

def train_logistic_regression_based_on_config(train_X: pd.DataFrame, 
                                        train_y: pd.Series, 
                                        test_X: Optional[pd.DataFrame], 
                                        task_config: TrainingTaskConfig,
                                        cv_splits: int,
                                        threads: int
                                        ) -> Tuple[LogisticRegression, Optional[pd.DataFrame]]:
    
    print("Training: Logistic Regression: Starting training.")
    if is_multiclass(train_y):
        metric = "f1_macro"
    else:
        metric = "f1"
    model, _ = train_logistic_regression(train_X.values, train_y.values, scoring=metric, cv_splits=cv_splits, n_jobs=threads)
    if task_config.output_model_path:
        print("Training: Logistic Regression: Saving model.")
        joblib.dump(model, task_config.output_model_path)

    test_pred = None
    if test_X is not None:
        print("Training: Logistic Regression: Predicting test labels.")
        test_pred = model.predict(test_X)
        if task_config.output_prediction_path:
            print("Training: Logistic Regression: Saving predicted test labels.")
            pd.Series(test_pred).to_csv(task_config.output_prediction_path)
    
    return model, test_pred
    