import os
import shutil
import pickle
import boto3
from Training.App_Logging.logger import appLogger

class fileOperations:
	"""
	This class is used to save the model after training
	and load the saved model for prediction.
	"""
	def __init__(self, log_file):
		self.log_file = log_file
		self.logger = appLogger()
		self.model_dir = 'Models/'

	def saveModel(self, model, filename):
		"""
		This method saves the model to the directory
		:param model: model object
		:param filename: name of file to save
		:return:
		"""
		self.logger.log(self.log_file, 'Entered the saveModel method of fileOperations class ...')
		try:
			path = os.path.join(self.model_dir, filename)
			if os.path.isdir(path):
				shutil.rmtree(path)   # remove existing directory
			os.makedirs(path)

			with open(path+'/'+filename+'.sav', 'wb') as f:
				pickle.dump(model, f)
			self.logger.log(self.log_file, 'Model file saved successfully: '+filename)
			return "Success"

		except Exception as e:
			self.logger.log(self.log_file, 'Error saving model: ' + str(e))
			raise e

	def loadModel(self, model_name):
		"""

		:param model_name:
		:return:
		"""
		self.logger.log(self.log_file, 'Entered the loadModel method of fileOperations class ...')
		try:
			with open(self.model_dir + model_name + '/' + model_name + '.sav', 'rb') as f:
				self.logger.log(self.log_file, 'Model file ' + model_name + ' loaded successfully !')
				return pickle.load(f)

		except Exception as e:
			self.logger.log(self.log_file, 'Error loading model: ' + str(e))
			raise e

	def find_correct_model_file(self, cluster):
		"""
		This method finds the correct model name given cluster number
		:param cluster:
		:return: model name
		"""
		self.logger.log(self.log_file, 'Entered the find_correct_model_file method of fileOperations class ...')
		try:
			for file in os.listdir(self.model_dir):
				try:
					if file.index(str(cluster)) != -1:
						model_name = file
						break
				except:
					continue
			return model_name

		except Exception as e:
			self.logger.log(self.log_file, 'Error finding correct model: ' + str(e))
			raise e

	def delete_existing_model_dir(self):
		"""
		This method deletes existing directory containing all models
		:return:
		"""
		self.logger.log(self.log_file, 'Entered the delete_existing_model_dir method of fileOperations class ...')
		try:
			if os.path.isdir('Models/'):
				shutil.rmtree('Models/')
				self.logger.log(self.log_file, 'Deleted existing model directory successfully !')

		except Exception as e:
			self.logger.log(self.log_file, 'Error deleting existing model directory: ' + str(e))
			raise e

	def upload_prediction_results(self, bucketName, filepath, name):
		"""
		This method downloads all prediction data files and schema file from the S3 bucket's
		prediction directory.
		:param bucketName: name of the S3 bucket
		:param filename: name of the file to upload
		:return:
		"""
		self.logger.log(self.log_file, "Entered the upload_prediction_results method of fileOperations class ...")
		try:
			s3 = boto3.resource(service_name='s3')
			s3.Bucket(bucketName).upload_file(Filename=filepath, Key=name)
			self.logger.log(self.log_file, 'Prediction results uploaded successfully to S3 bucket !')

		except Exception as e:
			self.logger.log(self.log_file, 'Error uploading prediction results to S3: ' + str(e))
			raise e