import json
import math

import numpy as np
import pandas as pd
from loguru import logger
from scipy.optimize import differential_evolution

from scripts.constants.app_configuration import PARAMETER_INFO, PARAMETER_INPUT, POSTGRES_URI, KAIROS_DB_HOST, \
    data_conf, MlflowMetaData, MODEL
from scripts.core.data.data_import import DataPuller
from scripts.core.model_factory.mlflow_util import ModelReTrainer
from scripts.core.model_factory.model_loader import ModelLoader
from scripts.core.recommender.reset_recommendation_range import check_recommendations

# from tabulate import tabulate


class ParameterOptimization(object):
    def __init__(self, parameter_info=PARAMETER_INFO):
        self.original_bound = None
        self.parameter_info = parameter_info
        # self.FI_5263A = "7302011030 Recovery Column-A.FI-5263A"
        # self.FIC_5102 = "Feed Drum.FIC-5102"
        # self.FIC_5111 = "Feed Drum.FIC-5111"
        # self.FIC_5161 = "Feed Drum.FIC-5161"
        try:
            self.mlflow_obj = ModelReTrainer(MlflowMetaData.EXPERIMENT_NAME, MlflowMetaData.RUN_NAME,
                                             MlflowMetaData.MODEL_NAME)
            self.model, self.df_profile = self.mlflow_obj.get_latest_model()
        except Exception as e:
            logger.warning(f"Failed to Retrieve model from MLFLOW because of error: {e}")
            logger.warning("Proceeding with locally saved model")
            self.model = ModelLoader(model_info=MODEL).load_model()
            
            self.data = "conf/profile.csv"
            logger.warning(f"Loading Profile from {self.data}")
            self.df_profile = pd.read_csv(self.data)

    def calculate_current_bounds(self):
        """
        This function loads profile data from a CSV file and extracts the bounds for each column in the dataset.

        Returns:
        df_profile_temp (pandas.DataFrame): DataFrame containing the profile data.
        bounds_data (dict): Dictionary where keys represent column names and values represent the bounds for each column.
                            Bounds are stored as dictionaries with "min" and "max" keys.
        """
        # df_profile_temp = self.df_profile
        data = "conf/profile.csv"
        df_profile_temp = pd.read_csv(data)
        
        data_designated_profile = df_profile_temp.to_dict(orient="records")
        bounds_data = dict()
        for each_record in data_designated_profile:
            bounds_data[each_record["column"]] = {
                "min": each_record["golden_batch_minimum"],
                "max": each_record["golden_batch_maximum"]
            }
        return df_profile_temp, bounds_data

    def calculate_bounds(self, live_data, feed_input_columns):
        """
        Calculate bounds for the data based on the provided live_data and feed_input_columns.

        Parameters:
        live_data (dict): A dictionary representing a single row of data from the dataframe.
                          It should be in the format orient="records".
        feed_input_columns (list): A list of columns of the type "Feed Drum Inputs".

        Returns:
        list: A list containing tuples of (min, max) bounds for each parameter.
              The order of tuples corresponds to the order of parameters in PARAMETER_INPUT.
        """
        logger.info("Calculating Bounds for the data")
        try:
            temporary_bound_data = self.original_bound.copy()
            if self.parameter_info["non_controllable_params"] is not None:
                for each_parameter in self.parameter_info["non_controllable_params"]:
                    temporary_bound_data[each_parameter] = {"min": int(live_data[each_parameter]),
                                                            "max": int(live_data[each_parameter])}
            if feed_input_columns:
                for each_parameter in feed_input_columns:
                    temporary_bound_data[each_parameter] = {"min": int(live_data[each_parameter]) - 1,
                                                            "max": int(live_data[each_parameter]) + 1}

            formatted_bound_data = dict()
            for each_key in PARAMETER_INPUT:
                formatted_bound_data[each_key] = (
                    int(temporary_bound_data[each_key]["min"]), int(temporary_bound_data[each_key]["max"]))
            logger.info("Formatted Bound Data", formatted_bound_data)
            return list(formatted_bound_data.values())
        except Exception as e:
            logger.error("Unable to calculate bounds for the data : {}".format(str(e)))

    def verify_operational_bounds(self, data_dict):
        lower_key = "lower"
        upper_key = "upper"
        original_key = "original"
        temporary_bound_data = self.original_bound.copy()
        temporary_recommendation_data = data_dict.copy()
        operational_bound_flag = True
        for each_key in PARAMETER_INPUT:
            parameter = each_key
            parameter_lower = f"{parameter}_{lower_key}"
            parameter_upper = f"{parameter}_{upper_key}"
            parameter_original = f"{parameter}_{original_key}"
            if (data_dict[parameter_original] < temporary_bound_data[parameter]["min"]) or (
                    data_dict[parameter_original] > temporary_bound_data[parameter]["max"]):
                operational_bound_flag = False
            if data_dict[parameter_lower] < temporary_bound_data[parameter]["min"]:
                temporary_recommendation_data[parameter_lower] = temporary_bound_data[parameter]["min"]
            if data_dict[parameter_upper] > temporary_bound_data[parameter]["max"]:
                temporary_recommendation_data[parameter_upper] = temporary_bound_data[parameter]["max"]
        print("operational_bound_flag --> ", operational_bound_flag)
        print("temporary_recommendation_data --> ", temporary_recommendation_data)
        return operational_bound_flag

    def _objective_fun(self, x):
        temp_df = pd.DataFrame(columns=PARAMETER_INPUT, data=[x])
        return 0 - self.model.predict(temp_df)[0]

    def get_recommendation(self, live_processed_data_dict, bounds):
        logger.info(f"live processed data dict is {live_processed_data_dict}")
        rs = differential_evolution(
            self._objective_fun,
            bounds=bounds,
            seed=42,
            workers=1)
        return rs.x

    def get_original_recommendation_values(self, data_dict, add_predict):
        lower_key = "lower"
        upper_key = "upper"
        original_key = "original"
        logger.info("Recommending original trend")
        
        df_validate = pd.DataFrame(data_dict)
        original_data = df_validate.to_dict(orient="records")[0]
        df_validate = df_validate[PARAMETER_INPUT]

        live_data = df_validate.to_dict(orient="records")[0]
        live_processed_data = df_validate.iloc[0, :].to_numpy().reshape(1, -1)
        
        #Predict on live data
        live_data_prediction = self.model.predict(live_processed_data)[0]
        
        #Keep live data as recommended data
        logger.debug("Started Calculating the recommendation for the data!")
        df_recommendation = df_validate
        logger.debug("Recommendation Calculated and received!")
        
        recommendation = df_validate.iloc[0, :].to_numpy().reshape(1, -1)
        rec_array = np.array(recommendation).reshape(1, -1)
        
        #Predict on recommended data
        prediction = self.model.predict(rec_array)[0]

        logger.info("Recommendations generated for last received live data.")
        recommended_data = df_recommendation.to_dict(orient="records")[0]
        
        updated_recommendation = dict()
        for each_column in recommended_data:
            if each_column in self.original_bound:
                column_delta = self.original_bound[each_column]["max"] - self.original_bound[each_column]["min"]
                if self.original_bound[each_column]["min"] == 0:
                    percentage_delta_lower = 0
                else:
                    percentage_delta_lower = int(100 * (column_delta / self.original_bound[each_column]["min"]))
                if self.original_bound[each_column]["max"] == 0:
                    percentage_delta_upper = 0
                else:
                    percentage_delta_upper = int(100 * (column_delta / self.original_bound[each_column]["max"]))
                percentage_delta = percentage_delta_lower - percentage_delta_upper
                calculated_delta = (0.15 * percentage_delta) / 100

                if each_column in self.original_bound:
                    lower_recommendation = recommended_data[each_column] * (1 - calculated_delta)
                    upper_recommendation = recommended_data[each_column] * (1 + calculated_delta)
                    original_recommendation = recommended_data[each_column]
                else:
                    lower_recommendation = recommended_data[each_column] * 0.95
                    upper_recommendation = recommended_data[each_column] * 1.05
                    original_recommendation = recommended_data[each_column]
                updated_recommendation["{}_{}".format(each_column, lower_key)] = math.floor(lower_recommendation)
                updated_recommendation["{}_{}".format(each_column, upper_key)] = math.ceil(upper_recommendation)
                updated_recommendation["{}_{}".format(each_column, original_key)] = round(original_recommendation, 2)

        predicted_result = dict()
        totalizer_max_value = 60000000000000
        predicted_totalizer_value = original_data[add_predict] + prediction
        if predicted_totalizer_value > totalizer_max_value:
            predicted_totalizer_value = predicted_totalizer_value - totalizer_max_value
        predicted_result[add_predict] = predicted_totalizer_value
        
        predicted_result['ai_yield_ach6_predicted'] = prediction
        predicted_result['ai_yield_ach6_live'] = live_data_prediction
        predicted_result['ai_yield_ach6_delta'] = prediction - live_data_prediction

        if predicted_result['ai_yield_ach6_delta'] >= 0:
            recommendation_flag = True
        else:
            recommendation_flag = False
        logger.info("Recalculated recommendation flag is {}".format(str(recommendation_flag)))
        return recommendation_flag, updated_recommendation, predicted_result

    def update_recommendation_lookup(self, model_recommendation, engine, profile_data, recommendation_flag,recommendation_text):
        logger.info(recommendation_flag)
        df_recommendation = profile_data
        recommendation_data = df_recommendation.to_dict(orient="records")

        updated_recommendation_data = list()
        for each_data in recommendation_data:
            temp_recommendation_data = each_data.copy()
            parameter = each_data["column"]
            parameter_lower = "{}_lower".format(parameter)
            parameter_upper = "{}_upper".format(parameter)
            parameter_original = "{}_original".format(parameter)
            if parameter_lower in model_recommendation and parameter_upper in model_recommendation:
                temp_recommendation_data["rec_lower"] = model_recommendation[parameter_lower]
                temp_recommendation_data["rec_upper"] = model_recommendation[parameter_upper]
                temp_recommendation_data["rec_original"] = model_recommendation[parameter_original]
                updated_recommendation_data.append(temp_recommendation_data)

        col = ["rec_lower", "rec_upper", "rec_original"]
        df_updated_recommendation = pd.DataFrame(updated_recommendation_data)
        df_updated_recommendation[col] = df_updated_recommendation[col].round(2)
        df_updated_recommendation['rec_status'] = recommendation_text
        df_updated_recommendation.to_csv("ach6_rec_lookup.csv", index=False)
        # df_updated_recommendation.to_sql('gbr_lookup_r5', engine, if_exists="replace")
        # self.df_profile.to_sql('gbr_full_r5', engine, if_exists="replace")

    def find_optimum(self, data_dict, add_predict, controllable=False):
        logger.info(f"Controllable is {controllable}")
        logger.info("Finding optimum data for the last received live data.")
        # print("Live Data -->", data_dict)
        df_validate = pd.DataFrame(data_dict)
        print("Current timestamp is: ",df_validate['timestamp'])
        cols_temp = ['timestamp',
                    'Alcohol Feed delta',
                    'Column C2 Bottom Temperature',
                    'Carburetor C4 temperature PV',
                    'Reactor 6A Vapour temperature',
                    'Reactor 6A Inlet temperature',
                    'Reactor 6A Outlet temperature',
                    'Reactor A Ethanol + Air flow',
                    'Reactor 6B Vapour temperature',
                    'Reactor 6B Inlet temperature',
                    'Reactor 6B Outlet temperature',
                    'Reactor B Ethanol + Air flow',
                    'Reactor 6C Vapour temperature',
                    'Reactor 6C Inlet temperature',
                    'Reactor 6C Outlet temperature',
                    'Reactor C Ethanol + Air flow',
                    'Reactor 6D Vapour temperature',
                    'Reactor 6D Inlet temperature',
                    'Reactor 6D Outlet temperature',
                    'Reactor D Ethanol + Air flow',
                    'Column_C9_Top temperature',
                    'Column_C9A_Top temperature',
                    'Product Output flow to Tank Totalizer',
                    'Totalizer Alcohol',
                    'Product Output flow to Tank',
                    'ACH6 delta']
        df_validate[cols_temp].to_csv("df_validate.csv")
        # print('df_validate',df_validate.head(3))
        
        original_data = df_validate.to_dict(orient="records")[0]
        # print("original_data",original_data)
        
        df_validate = df_validate[PARAMETER_INPUT]

        live_data = df_validate.to_dict(orient="records")[0]
        # print("Live Data -->", live_data)
        
        live_processed_data = df_validate.iloc[0, :].to_numpy().reshape(1, -1)
        
        # print("Live Processed Data -->", live_processed_data)
        live_data_prediction = self.model.predict(df_validate)[0]
        # print("PREDICTION ON LIVE DATA", live_data_prediction)
        # print('live_data_prediction---', live_data_prediction,)
        
        df_profile, live_bounds = self.calculate_current_bounds()
        # print("live_bounds", live_bounds)
        self.original_bound = live_bounds
        
        feed_input_cols = self.parameter_info["feed_input"]
        bounds = self.calculate_bounds(live_data=live_data, feed_input_columns=feed_input_cols)
        # print("BOUNDS BASED ON PROFILE GBA_MIN_MAX", bounds)
        
        logger.debug("Started Calculating the recommendation for the data!")
        recommendation = (list(map(float, self.get_recommendation(
            live_processed_data_dict=live_processed_data, bounds=bounds))))
        logger.debug("Recommendation Calculated and received!")
        
        df_recommendation = pd.DataFrame(columns=list(df_validate.columns), data=[recommendation])
        # print("RECOMMENDATION BASED ON DIFFERENTIAL EVOLUTION", df_recommendation)

        prediction = self.model.predict(df_recommendation)[0]
        # print("PREDICTION ON RECOMMENDATION DATA", prediction)
        # print("prediction", prediction)
        logger.info("Recommendations generated for last received live data.")
        
        recommended_data = df_recommendation.to_dict(orient="records")[0]
        # print(f"Recommended Data --> {recommended_data}")
        df_recommendation.to_csv('recommended_data.csv', index=False)
        

        updated_recommendation = dict()
        updated_recommendation = get_recommendation_dict(recommended_data=recommended_data,
                                                        original_bound=self.original_bound,
                                                        updated_recommendation=updated_recommendation)
        # print("UPDATED REOMMENDATION BASED ON CALCULATIONS->",updated_recommendation)
        
        # df_updated_recommendation = pd.DataFrame(updated_recommendation.items(), columns=['Key', 'Value'])
        # df_updated_recommendation.to_csv('updated_recommendation.csv', index=False)
        
        # updated_recommendation = check_recommendations(recommendations_dict=updated_recommendation,
        #                                                      df_profile=df_profile)
        # print("Updated Recommendation:-", updated_recommendation)
        
        predicted_result = dict()
        # predicted_result, prediction, live_data_prediction, ai_yield_ach6_predicted,
        #                     ai_yield_ach6_live, ai_yield_ach6_delta, reactor, updated_recommendation
        totalizer_max_value = 60000000000000 #max value of totalizer
        predicted_totalizer_value = original_data[add_predict] + prediction
        if predicted_totalizer_value > totalizer_max_value:
            predicted_totalizer_value = predicted_totalizer_value - totalizer_max_value
        predicted_result[add_predict] = predicted_totalizer_value
        
        predicted_result['ai_yield_ach6_predicted'] = prediction
        predicted_result['ai_yield_ach6_live'] = live_data_prediction
        predicted_result['ai_yield_ach6_delta'] = prediction - live_data_prediction
        
        if predicted_result['ai_yield_ach6_delta'] >= 0:
            recommendation_flag = True
            final_recommendation_flag = True
            recommendation_text = "Operator to make changes based on recommendations!"
        else:
            recommendation_flag = False
            final_recommendation_flag = False
            recommendation_text = "Operator to make changes based on recommendations!"
        
        logger.info("Recommendation flag is {}".format(str(recommendation_flag)))
        
        if not recommendation_flag:
            logger.info("Recommending current trend as recommendation is false")
            final_recommendation_flag, corrected_updated_recommendation, predicted_result = \
                self.get_original_recommendation_values(data_dict=data_dict, add_predict=add_predict)
            logger.debug("Verifying if the live data was within the operational limits")
            operational_bound_flag = self.verify_operational_bounds(data_dict=corrected_updated_recommendation)
            if operational_bound_flag:
                logger.debug("Plant is operating better in true bound. Hence, current data is the advised data!")
                updated_recommendation = corrected_updated_recommendation #??
                recommendation_text = "Operator to not make any changes!"
            else:
                logger.debug("Recommending model's data as plant data is out of bound!")
                recommendation_text = "Operating outside operational range!"
        logger.info("Updating lookup with the latest recommendation")
        
        logger.debug("Attempting to write to the Postgres DB")
        self.update_recommendation_lookup(model_recommendation=updated_recommendation,
                                        engine=POSTGRES_URI,
                                        recommendation_flag=recommendation_flag,
                                        recommendation_text=recommendation_text,
                                        profile_data=df_profile)
        
        return  final_recommendation_flag, updated_recommendation, predicted_result
    
def get_recommendation_dict(recommended_data, original_bound, updated_recommendation):
    try:
        for each_column in recommended_data:
            if each_column in original_bound:
                column_delta = original_bound[each_column]["max"] - original_bound[each_column]["min"]
                try:
                    percentage_delta_lower = float(100 * (column_delta / original_bound[each_column]["min"]))
                    percentage_delta_upper = float(100 * (column_delta / original_bound[each_column]["max"]))
                except:
                    percentage_delta_lower = column_delta * 2
                    percentage_delta_upper = column_delta
                percentage_delta = percentage_delta_lower - percentage_delta_upper
                calculated_delta = (0.15 * percentage_delta) / 100

                if each_column in original_bound:
                    lower_recommendation = max(original_bound[each_column]["min"],
                                            recommended_data[each_column] * (1 - calculated_delta))
                    upper_recommendation = min(original_bound[each_column]["max"],
                                            recommended_data[each_column] * (1 + calculated_delta))
                    original_recommendation = recommended_data[each_column]
                else:
                    lower_recommendation = recommended_data[each_column] * 0.95
                    upper_recommendation = recommended_data[each_column] * 1.05
                    original_recommendation = recommended_data[each_column]
                updated_recommendation["{}_lower".format(each_column)] = math.floor(lower_recommendation)
                updated_recommendation["{}_upper".format(each_column)] = math.ceil(upper_recommendation)
                updated_recommendation["{}_original".format(each_column)] = round(original_recommendation, 2)

        return updated_recommendation
    except Exception as e:
        logger.exception(f'Exception - {e}')
        
