from __future__ import print_function
import importlib
import random
import pandas as pd
import time
import pickle
from datetime import datetime
from Utils import makeDir,write_to_csv,copyfiles,copyFile
from Envelope import getEnvelopeResult
from ast import literal_eval
import logging
import sys
import statistics
from typing import List
import numpy as np
from joblib import Parallel, delayed
import multiprocessing
import math


#generate train and test sets
def generate_data(data_path: str, n_folds: int, make_cv_data, configurations):
    train_data, train_labels, test_data, test_labels,min_max_dict = make_cv_data(data_path, n_folds, configurations['weight_files'], configurations['fold_data_dir'])
    return train_data, train_labels, test_data, test_labels,min_max_dict


#train a neural network based on the training data
def train_NN(train_data, train_labels,fold, configurations, train):
    NN_model, MIP_model = train(train_data, train_labels,fold,configurations['weight_files'],configurations['model_parameters']['layer_size'],configurations['model_parameters']['hidden_size'],configurations['model_parameters']['epochs'],configurations['model_parameters']['batch_size'],configurations['model_parameters']['lr'])
    return NN_model, MIP_model


#given a size, randomly assing index of trainset to a batch
def make_batch(train_data, train_labels, configurations):
    #select a batch size, e.g., 64
    data_size = train_data.shape
    # batch_size = 64
    batch_size = configurations['model_parameters']['batch_size']
    batches = []
    batches_labels = []
    data_indexes = list(range(0, data_size[0]))
    # random.shuffle(data_indexes)
    start = 0
    while (start+batch_size)<data_size[0]:
        batch= data_indexes[start:start+batch_size]
        batches.append(batch)
        start += batch_size
    batch = data_indexes[start:]
    batches.append(batch)
    return batches

def evaluate_model(NN_model, test_data, test_labels, configurations, evaluate):
    test_dataframe = pd.DataFrame(test_data)
    label_dataframe = pd.DataFrame(test_labels)
    score = evaluate(NN_model, test_dataframe, label_dataframe)
    return score

#fit the model based on the new counter batch and return a new model
def update_model(counter_batch, batch_labels, NN_model, MIP_model,fold, configurations, update_batch):
    # weight_dir = configurations['weight_files']
    # NN_model, MIP_model = update(NN_model, MIP_model, counter_batch, batch_labels,fold,weight_dir)
    # return NN_model, MIP_model
    weight_dir = configurations['weight_files']
    NN_model, MIP_model = update_batch(NN_model, MIP_model, counter_batch, batch_labels,fold,weight_dir,counter_batch.shape[0])
    return NN_model, MIP_model

