from loguru import logger
import tracemalloc
import gc
from scripts.core.engine.tags_data import get_tags_data
from scripts.core.engine.mppt_data import GetData
from scripts.utils.preprocessing import DataPreprocessing
from scripts.core.engine.data_training_and_inference import Training


class AiModelling:
    def __init__(self, df_raw_tags, df_coefficients, start_timestamp, end_timestamp):
        self.df_raw_tags = df_raw_tags
        self.df_coefficients = df_coefficients
        self.start_timestamp = start_timestamp
        self.end_timestamp = end_timestamp

    def all_calculations(self):
        try:
            for inv_id in list(self.df_raw_tags['inv_id'].unique()):
                df = self.df_raw_tags[self.df_raw_tags['inv_id'] == inv_id]
                for mppt_id in list(self.df_raw_tags['mppt_id'].unique()):
                    print(f'memory allocation for {inv_id} & {mppt_id} - {tracemalloc.get_traced_memory()}')
                    data_modelling = DataModelling(start_timestamp=self.start_timestamp,
                                                   end_timestamp=self.end_timestamp,
                                                   df_coefficients=self.df_coefficients)
                    data_modelling.get_data(inv_id=inv_id, mppt_id=mppt_id, df=df)
                    del data_modelling

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

    def __del__(self):
        try:
            print('destructor called, AiModelling die!')
        except Exception as e:
            logger.exception(f'Exception - {e}')


class DataModelling:
    def __init__(self, start_timestamp, end_timestamp, df_coefficients):
        self.start_timestamp = start_timestamp
        self.end_timestamp = end_timestamp
        self.df_coefficients = df_coefficients

    def get_data(self, inv_id, mppt_id, df):
        try:
            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=self.start_timestamp,
                                           end_timestamp=self.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=self.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()}')
        except Exception as e:
            logger.exception(f'Exception - {e}')

    def __del__(self):
        try:
            print('destructor called, DataModelling die!')
        except Exception as e:
            logger.exception(f'Exception - {e}')