if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv(dotenv_path='config.env')
import warnings
from loguru import logger
import tracemalloc
import gc
import numpy as np
from scripts.core.engine.mppt_data import GetData
from scripts.core.engine.tags_data import get_tags_data
from scripts.utils.start_end_date import KairosStartEndDate
from scripts.core.engine.final_tags import GetFinalDf
from scripts.core.engine.data_training_and_inference import Inference
from scripts.core.engine.final_predictions import ai_modelling
from scripts.core.engine.raw_predicted_tags import get_raw_predicted_tags
from scripts.core.data_puller_push.kafka_push import CalculatedDataPush
warnings.filterwarnings("ignore")


start_date, end_date, start_timestamp, end_timestamp = KairosStartEndDate().start_end_date()


def orchestrator():
    try:
        tracemalloc.start()
        print(f'memory allocation - {tracemalloc.get_traced_memory()}')
        gc.collect()
        print(f'memory allocation - {tracemalloc.get_traced_memory()}')
        tracemalloc.clear_traces()
        print(f'memory allocation - {tracemalloc.get_traced_memory()}')
        df_raw_tags, df_predicted_tags, df_coefficients = get_raw_predicted_tags()
        logger.info(f'raw tags dataframe shape - {df_raw_tags.shape}')
        logger.info(f'predicted tags dataframe shape - {df_predicted_tags.shape}')

        for inv_id in list(df_raw_tags['inv_id'].unique()):
            df = df_raw_tags[df_raw_tags['inv_id'] == inv_id]
            for mppt_id in list(df_raw_tags['mppt_id'].unique()):
                print(f'1st memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                print(f'2nd memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'3rd memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

                df_mppt_level = df[df['mppt_id'] == mppt_id]
                df_kairos_data = get_tags_data(df_input_tags=df_mppt_level, start_timestamp=start_timestamp,
                                               end_timestamp=end_timestamp, inv_id=inv_id, mppt_id=mppt_id)

                logger.info(f'Shape of final df - {df_kairos_data.shape}')

                print(f'4th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                print(f'5th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'6th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

                mppt_data = GetData()
                df_mppt = mppt_data.associate_inv_mppt_id(df=df_kairos_data)

                print(f'7th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                print(f'8th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'9th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

                df_coefficient_multiply = mppt_data.multiply_mppt_coefficients(df_mppt=df_mppt,
                                                                               df_coefficients=df_coefficients)

                print(f'10th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                print(f'11th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'12th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

                x_test, y_test, predictions = Inference(df=df_coefficient_multiply).data_inference(inv_id=inv_id, mppt_id=mppt_id)
                get_final_df = GetFinalDf()
                df_result = get_final_df.get_final_data(x_test=x_test,
                                                        y_test=y_test,
                                                        predictions=predictions)
                final_dict = get_final_df.get_final_predicted_tags(df_predicted_current_tags=df_predicted_tags,
                                                                   inv_id=inv_id, mppt_id=mppt_id)

                df_result["timestamp"] = df_result["datetime"].values.astype(np.int64) / 10 ** 9
                df_result["timestamp"] = df_result["timestamp"].astype('int')

                CalculatedDataPush(df_result=df_result, final_tags_dict=final_dict).kafka_data_push()
                logger.info(f'{final_dict}')

                del df_kairos_data
                del df_mppt
                del df_coefficient_multiply

                print(f'16th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                print(f'17th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'18th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

            print(f'final allocation for {inv_id} - {tracemalloc.get_traced_memory()}')
        tracemalloc.stop()
        logger.info(f'data inference has been completed !')
    except Exception as e:
        logger.exception(f'Exception - {e}')


if __name__ == '__main__':
    orchestrator()

