"""
File contains functions used to train Linear Support Vector Machine.

Author: Dominik Macko
"""

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

import numpy as np
import pandas as pd
from sklearn.svm import LinearSVC
import joblib

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

def linear_svm_best_params_surroundings(best_params: Dict[str, Any]) -> Dict[str, Any]:
    """Get best parameters surroundings for random forest."""
    
    C = best_params["C"]
    return {
        "C": [1/2 * C, C, 3/2 * C]
    }

def train_linear_svm(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[LinearSVC, pd.DataFrame]:
    """Trains linear support vector machine by searching for optimal l1 ratio.
    
    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 = {
        "penalty": ["l1", "l2"],
        "loss": ["hinge", "squared_hinge"],
        "C": [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]
    }
    return train_model(
        LinearSVC(
            dual=False,
            class_weight="balanced",
            random_state=seed,
            max_iter=1000
        ),
        train_gridsearchcv_model,
        train_gridsearchcv_model,
        grid,
        linear_svm_best_params_surroundings,
        train_X,
        train_y,
        scoring=scoring,
        n_jobs=n_jobs,
        verbose=verbose,
        cv_splits=cv_splits
    )

def train_svm_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[LinearSVC, Optional[pd.DataFrame]]:
    
    print("Training: SVM: Starting training.")
    if is_multiclass(train_y):
        metric = "f1_macro"
    else:
        metric = "f1"
    model, _ = train_linear_svm(train_X.values, train_y.values, scoring=metric, cv_splits=cv_splits, n_jobs=threads)
    if task_config.output_model_path:
        print("Training: SVM: Saving model.")
        joblib.dump(model, task_config.output_model_path)

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