if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv(dotenv_path='config.env')
import warnings
from loguru import logger
import tracemalloc
import gc
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.utils.preprocessing import DataPreprocessing
from scripts.core.engine.data_training_and_inference import Training
from scripts.core.engine.raw_predicted_tags import get_raw_predicted_tags
warnings.filterwarnings("ignore")


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


def orchestrator():
    try:
        tracemalloc.start()
        print(f'initial memory usage - {tracemalloc.get_traced_memory()}')
        gc.collect()
        tracemalloc.reset_peak()
        tracemalloc.clear_traces()
        print(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()
                tracemalloc.reset_peak()
                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)

                print(f'4th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                tracemalloc.reset_peak()
                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()}')

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

                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()
                tracemalloc.reset_peak()
                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()
                tracemalloc.reset_peak()
                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()}')

                data_preprocessing = DataPreprocessing()
                df_clean = data_preprocessing.remove_outliers(df=df_coefficient_multiply,
                                                              param_list=['tilt_irradiance', 'voltage_mppt',
                                                                          'current_mppt'])
                print(f'13th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                gc.collect()
                tracemalloc.reset_peak()
                print(f'14th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                tracemalloc.clear_traces()
                print(f'15th memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')

                Training(df=df_clean).data_training(inv_id=inv_id, mppt_id=mppt_id)

                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()
                tracemalloc.reset_peak()
                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} & {mppt_id} - {tracemalloc.get_traced_memory()}')
        logger.info(f'data training for {inv_id}, {mppt_id} has been completed !')
        tracemalloc.stop()
    except Exception as e:
        logger.exception(f'Exception - {e}')


if __name__ == '__main__':
    orchestrator()

