import math
import warnings
import traceback
from loguru import logger
from datetime import datetime

import numpy as np
import pandas as pd

from scripts.constants.constants import PickupConstants

warnings.filterwarnings("ignore")


def get_mixer_batch_date(raw_df, index_number):
    try:
        logger.info('Getting mixer batch date dictionary')
        mixer_df = raw_df[PickupConstants.pick_imp_mixer_cols]
        mixer_df = mixer_df.sort_values(by='Time Stamp')
        numeric_cols = mixer_df.select_dtypes(include=['int', 'float']).columns

        # Convert numeric columns to float
        mixer_df[numeric_cols] = mixer_df[numeric_cols].astype(float)
        mixer_df['day'] = mixer_df['Time Stamp'].dt.date
        mixer_df = mixer_df[mixer_df["Size No (INDEX No).3"] == index_number]
        mixer_df = mixer_df[mixer_df["Mixing batch number"] != 0]
        mixer_df['time_min'] = mixer_df['Time Stamp']
        mixer_df['time_max'] = mixer_df['Time Stamp']
        aggregation_dict = {
            'time_min': 'min',
            'time_max': 'max',
        }
        group_by = ['day', 'Mixing batch number']
        df_mixer_grouped = mixer_df.groupby(group_by).agg(aggregation_dict).reset_index()
        df_mixer_grouped['mixer_section_time_diff_second'] = df_mixer_grouped['time_max'] - df_mixer_grouped['time_min']
        df_mixer_grouped['mixer_section_time_diff_second'] = df_mixer_grouped[
            'mixer_section_time_diff_second'].dt.total_seconds()
        df_mixer_grouped['batch-date'] = 'Batch_' + df_mixer_grouped['Mixing batch number'].astype(str) + '_' + \
                                         df_mixer_grouped['day'].astype(str)
        date_dict = {}
        batch_lis = list(df_mixer_grouped['batch-date'].unique())
        for each_bt in batch_lis:
            df_nw = df_mixer_grouped[df_mixer_grouped['batch-date'] == each_bt]
            date_dict[each_bt] = {"start_time": str(list(df_nw['time_min'])[0]),
                                  'end_time': str(list(df_nw['time_max'])[0])}
        return date_dict
    except Exception as err:
        logger.error(f'Error while getting mixer time and forming date dict: {str(err)}')
        logger.error(traceback.format_exc())
        raise Exception(str(err))


def return_batch_no_bof_df(raw_df, viscosity_df, date_dict, index_number):
    try:
        logger.info('Getting bof batch number')
        raw_df['day'] = raw_df['Time Stamp'].dt.date
        raw_df['day'] = raw_df['day'].astype('str')

        raw_df['Mixing batch number'] = raw_df['Mixing batch number'].astype('float')
        raw_df['batch-date'] = 'Batch_' + raw_df['Mixing batch number'].astype(
            'str') + '_' + raw_df['day'].astype('str')

        bof_add_cols = ['Size No (INDEX No).5', 'length passed through', 'Time Stamp', 'day', 'lower door open']
        bof_df = raw_df[bof_add_cols]

        sorted_bof_df = bof_df.sort_values(by="Time Stamp", ascending=True)
        sorted_bof_df = sorted_bof_df[sorted_bof_df['Size No (INDEX No).5'] == index_number]
        dt_list = list(sorted_bof_df['day'].unique())

        day_length_dic = {}
        for each_day in dt_list:
            day_df = sorted_bof_df[sorted_bof_df['day'] == each_day]
            if (day_df['length passed through'].max() - day_df['length passed through'].min()) <= 0:
                value = 0
            else:
                value = day_df['length passed through'].max() - day_df['length passed through'].min()
            day_length_dic[each_day] = value

        sorted_viscosity_df = viscosity_df.sort_values(by="Mixing date", ascending=True)
        sorted_viscosity_df['day'] = sorted_viscosity_df['Mixing date'].dt.date
        sorted_viscosity_df['day'] = sorted_viscosity_df['day'].astype('str')

        extrud_visc_df = sorted_viscosity_df[['Batch No.', 'Input rubber weight(0.1kg)', 'day', 'Mixing date']]
        extrud_visc_df['length_from_extruder'] = extrud_visc_df['day'].map(day_length_dic)
        extrud_visc_df['length_from_extruder'] = extrud_visc_df['length_from_extruder'].fillna(0)
        daily_sum_weight = extrud_visc_df.groupby('day')['Input rubber weight(0.1kg)'].sum() / 10
        # Add a new column 'm/kg' by dividing 'length_from_extruder' by the sum for each day
        extrud_visc_df['m/kg'] = extrud_visc_df.apply(
            lambda row: row['length_from_extruder'] / daily_sum_weight[row['day']], axis=1)
        extrud_visc_df['batch_length'] = extrud_visc_df.apply(
            lambda row: row['m/kg'] * row['Input rubber weight(0.1kg)'] / 10, axis=1).astype('float64')

        extrud_visc_df['batch_length'] = extrud_visc_df['batch_length'].apply(math.ceil)
        extrud_visc_df['cumulative_length'] = extrud_visc_df.groupby('day')['batch_length'].cumsum()

        discharge_dict = extrud_visc_df.groupby('day').apply(
            lambda group: group.set_index('Batch No.').to_dict()['cumulative_length']).to_dict()

        test_sorted_extr_df = sorted_bof_df
        test_df = test_sorted_extr_df

        # Initialize an empty list to store batch numbers
        batch_numbers = []

        # Iterate through each row in the DataFrame
        for index, row in test_df.iterrows():
            day = row['day']
            discharge_length = row['length passed through']
            if discharge_length == 0:
                batch_numbers.append(0)
            else:
                # Check if the day is in the dictionary
                if day in discharge_dict:
                    # Check if discharge length is less than or equal to the corresponding batch length
                    batch_length_dict = discharge_dict[day]
                    for batch_no, batch_length in batch_length_dict.items():
                        if discharge_length <= batch_length:
                            batch_numbers.append(batch_no)
                            break
                    else:
                        # If no match is found in the dictionary, assign NaN to batch number
                        batch_numbers.append(batch_numbers[-1])
                else:
                    # If day is not in the dictionary, assign NaN to batch number
                    batch_numbers.append(np.nan)

        # Add the 'batch_no' column to the DataFrame
        test_df['batch_no'] = batch_numbers

        batch_number = 0
        batch_list = []

        started_with_one = False
        current_day = None

        for value, day in zip(list(test_df['lower door open']), list(test_df['day'])):

            if current_day != day:
                current_day = day
                batch_number = 0

            if value == 1:
                if not started_with_one:
                    batch_number += 1
                    started_with_one = True
                batch_list.append(batch_number)
            else:
                batch_list.append(batch_number)
                started_with_one = False

        batch_number = 0
        batch_list = []

        started_with_one = False

        for value in test_df['lower door open']:
            if value == 1:
                if not started_with_one:
                    batch_number += 1
                    started_with_one = True
                batch_list.append(batch_number)
            else:
                batch_list.append(batch_number)
                started_with_one = False

        test_df['batch_no'] = test_df['batch_no'].astype('float')
        test_df['bof_batch_date'] = 'Batch_' + test_df['batch_no'].astype('str') + '_' + test_df['day'].astype('str')
        extruder_flag_list = []
        extrud_flg_vms = []
        for i, value in test_df.iterrows():
            if value['batch_no'] == 0.0:
                extruder_flag_list.append('false')
                extrud_flg_vms.append(0)
            else:
                # start_time = np.datetime64(date_dict.get(value['bof_batch_date']).get('start_time'))
                # end_time = np.datetime64(date_dict.get(value['bof_batch_date']).get('end_time'))

                start_time = date_dict.get(value["bof_batch_date"]).get("start_time")
                end_time = date_dict.get(value["bof_batch_date"]).get("end_time")

                if (datetime.strptime(str(value["Time Stamp"]).split('+')[0], '%Y-%m-%d %H:%M:%S') > datetime.strptime(
                        start_time.split('+')[0], '%Y-%m-%d %H:%M:%S')) & \
                    (datetime.strptime(str(value["Time Stamp"]).split('+')[0], '%Y-%m-%d %H:%M:%S') < datetime.strptime(
                        end_time.split('+')[0], '%Y-%m-%d %H:%M:%S')):

                    extruder_flag_list.append('true')
                    extrud_flg_vms.append(1)
                else:
                    extruder_flag_list.append('false')
                    extrud_flg_vms.append(0)

        test_df['bof_flag'] = extruder_flag_list
        test_df['bof_batch_diff'] = extrud_flg_vms
        test_df['updated_bt_list'] = batch_list

        test_df['bof_batch_number'] = test_df['batch_no'] - test_df['bof_batch_diff'].astype('float')
        test_df = test_df.rename(columns={'bof_batch_date': 'batch-date'})
        return test_df
    except Exception as err:
        logger.error(f'Error while adding batch to bof section: {str(err)}')
        logger.error(traceback.format_exc())
        raise Exception(str(err))


def get_bof_batch_date(bof_batch_df, index_number):
    try:
        logger.info('Getting bof batch date dictionary')
        bof_cols = PickupConstants.pick_imp_bof_cols
        bof_df = bof_batch_df[bof_cols]
        bof_df = bof_df.sort_values(by='Time Stamp')
        numeric_cols = bof_df.select_dtypes(include=['int', 'float']).columns

        # Convert numeric columns to float
        bof_df[numeric_cols] = bof_df[numeric_cols].astype(float)
        bof_df['day'] = bof_df['Time Stamp'].dt.date
        bof_df = bof_df[bof_df["Size No (INDEX No).5"] == index_number]
        bof_df = bof_df[bof_df["bof_batch_number"] != 0]
        bof_df['time_min'] = bof_df['Time Stamp']
        bof_df['time_max'] = bof_df['Time Stamp']
        aggregation_dict = {
            'time_min': 'min',
            'time_max': 'max',
        }
        group_by = ['day', 'bof_batch_number']
        df_bof_grouped = bof_df.groupby(group_by).agg(aggregation_dict).reset_index()
        df_bof_grouped['mixer_section_time_diff_second'] = df_bof_grouped['time_max'] - df_bof_grouped['time_min']
        df_bof_grouped['mixer_section_time_diff_second'] = df_bof_grouped[
            'mixer_section_time_diff_second'].dt.total_seconds()
        df_bof_grouped['batch-date'] = 'Batch_' + df_bof_grouped['bof_batch_number'].astype(str) + '_' + df_bof_grouped[
            'day'].astype(str)
        bof_date_dict = {}
        batch_lis = list(df_bof_grouped['batch-date'].unique())
        for each_bt in batch_lis:
            df_nw = df_bof_grouped[df_bof_grouped['batch-date'] == each_bt]
            bof_date_dict[each_bt] = {"start_time": str(list(df_nw['time_min'])[0]),
                                      'end_time': str(list(df_nw['time_max'])[0])}
        return bof_date_dict
    except Exception as err:
        logger.error(f'Error while getting bof time and forming bof date dict: {str(err)}')
        logger.error(traceback.format_exc())
        raise Exception(str(err))


def return_pick_batch_no_df(raw_df, viscosity_df, bof_date_dict, bof_batch_num_df, index_number):
    try:
        logger.info('Getting pickup batch date dataframe')
        raw_df['day'] = raw_df['Time Stamp'].dt.date
        raw_df['day'] = raw_df['day'].astype('str')

        raw_df['Mixing batch number'] = raw_df['Mixing batch number'].astype('float')
        raw_df['batch-date'] = 'Batch_' + raw_df['Mixing batch number'].astype('str') + '_' + raw_df['day'].astype(
            'str')
        pick_add_cols = PickupConstants.pick_cols + PickupConstants.pick_additional_cols
        pick_df = raw_df[pick_add_cols]

        sorted_pick_df = pick_df.sort_values(by="Time Stamp", ascending=True)
        sorted_pick_df = sorted_pick_df[sorted_pick_df['Size No (INDEX No).6'] == index_number]
        dt_list = list(sorted_pick_df['day'].unique())

        day_length_dic = {}
        for each_day in dt_list:
            day_df = sorted_pick_df[sorted_pick_df['day'] == each_day]
            if day_df['length passed through.1'].max() - day_df['length passed through.1'].min() <= 0:
                value = 0
            else:
                value = day_df['length passed through.1'].max() - day_df['length passed through.1'].min()
            day_length_dic[each_day] = value

        sorted_viscosity_df = viscosity_df.sort_values(by="Mixing date", ascending=True)
        sorted_viscosity_df['day'] = sorted_viscosity_df['Mixing date'].dt.date
        sorted_viscosity_df['day'] = sorted_viscosity_df['day'].astype('str')

        extrud_visc_df = sorted_viscosity_df[['Batch No.', 'Input rubber weight(0.1kg)', 'day', 'Mixing date']]
        extrud_visc_df['length_from_pickup'] = extrud_visc_df['day'].map(day_length_dic)
        extrud_visc_df['length_from_pickup'] = extrud_visc_df['length_from_pickup'].fillna(0)
        daily_sum_weight = extrud_visc_df.groupby('day')['Input rubber weight(0.1kg)'].sum() / 10
        # Add a new column 'm/kg' by dividing 'length_from_pickup' by the sum for each day
        extrud_visc_df['m/kg'] = extrud_visc_df.apply(
            lambda row: row['length_from_pickup'] / daily_sum_weight[row['day']], axis=1)
        extrud_visc_df['batch_length'] = extrud_visc_df.apply(
            lambda row: row['m/kg'] * row['Input rubber weight(0.1kg)'] / 10, axis=1).astype('float64')

        extrud_visc_df['batch_length'] = extrud_visc_df['batch_length'].apply(math.ceil)
        extrud_visc_df['cumulative_length'] = extrud_visc_df.groupby('day')['batch_length'].cumsum()

        discharge_dict = extrud_visc_df.groupby('day').apply(
            lambda group: group.set_index('Batch No.').to_dict()['cumulative_length']).to_dict()

        test_sorted_extr_df = sorted_pick_df
        test_pick_df = test_sorted_extr_df

        # Initialize an empty list to store batch numbers
        batch_numbers = []

        # Iterate through each row in the DataFrame
        for index, row in test_pick_df.iterrows():
            day = row['day']
            discharge_length = row['length passed through.1']
            if discharge_length == 0:
                batch_numbers.append(0)
            else:
                # Check if the day is in the dictionary
                if day in discharge_dict:
                    # Check if discharge length is less than or equal to the corresponding batch length
                    batch_length_dict = discharge_dict[day]
                    for batch_no, batch_length in batch_length_dict.items():
                        if discharge_length <= batch_length:
                            batch_numbers.append(batch_no)
                            break
                    else:
                        # If no match is found in the dictionary, assign NaN to batch number
                        batch_numbers.append(batch_numbers[-1])
                else:
                    # If day is not in the dictionary, assign NaN to batch number
                    batch_numbers.append(np.nan)

        # Add the 'batch_no' column to the DataFrame
        test_pick_df['batch_no'] = batch_numbers
        test_pick_df['batch_no'] = test_pick_df['batch_no'].astype('float')
        test_pick_df['pickup_batch_date'] = 'Batch_' + test_pick_df['batch_no'].astype('str') + '_' + test_pick_df[
            'day'].astype('str')
        test_pick_df['bof_batch_number'] = bof_batch_num_df['bof_batch_number']
        extruder_flag_list = []
        extrud_flg_vms = []
        for i, value in test_pick_df.iterrows():
            if value['batch_no'] == 0.0:
                extruder_flag_list.append('false')
                extrud_flg_vms.append(0)
            else:
                # start_time = np.datetime64(bof_date_dict.get(value['pickup_batch_date']).get('start_time'))
                # end_time = np.datetime64(bof_date_dict.get(value['pickup_batch_date']).get('end_time'))

                start_time = bof_date_dict.get(value["pickup_batch_date"]).get("start_time")
                end_time = bof_date_dict.get(value["pickup_batch_date"]).get("end_time")

                if (datetime.strptime(str(value["Time Stamp"]).split('+')[0], '%Y-%m-%d %H:%M:%S') > datetime.strptime(start_time.split('+')[0], '%Y-%m-%d %H:%M:%S')) & \
                    (datetime.strptime(str(value["Time Stamp"]).split('+')[0], '%Y-%m-%d %H:%M:%S') < datetime.strptime(end_time.split('+')[0], '%Y-%m-%d %H:%M:%S')):
                    extruder_flag_list.append('false')
                    extrud_flg_vms.append(0)
                else:
                    extruder_flag_list.append('false')
                    extrud_flg_vms.append(0)

        test_pick_df['pickup_flag'] = extruder_flag_list
        test_pick_df['pickup_batch_diff'] = extrud_flg_vms

        test_pick_df['pickup_batch_verify_number'] = test_pick_df['batch_no'] - test_pick_df[
            'pickup_batch_diff'].astype('float')

        actual_pickup_bt_num = []
        for i, value in test_pick_df.iterrows():
            pickup_batch_number = value['pickup_batch_verify_number']
            bof_batch_num = value['bof_batch_number']
            if pickup_batch_number <= bof_batch_num:
                actual_pickup_bt_num.append(pickup_batch_number)
            else:
                actual_pickup_bt_num.append(bof_batch_num)
        test_pick_df['pickup_batch_number'] = actual_pickup_bt_num
        test_pick_df['batch-date'] = 'Batch_' + test_pick_df['pickup_batch_number'].astype('str') + '_' + \
                                     test_pick_df['day'].astype('str')

        return test_pick_df
    except Exception as err:
        logger.error(f'Error in adding batch to pick section: {str(err)}')
        logger.error(traceback.format_exc())


def preprocess_pickup_section(raw_df, index_number, viscosity_df):
    try:
        logger.info('Preprocessing and getting aggregated pickup dataframe')
        raw_df['Time Stamp'] = pd.to_datetime(raw_df['Time Stamp'])
        raw_df['day'] = raw_df['Time Stamp'].dt.date
        raw_df['day'] = raw_df['day'].astype('str')
        extr_cols = PickupConstants.pick_cols + PickupConstants.pick_additional_cols
        pick_df = raw_df[extr_cols]

        sorted_pick_df = pick_df.sort_values(by="Time Stamp", ascending=True)
        sorted_pick_df = sorted_pick_df[sorted_pick_df['Size No (INDEX No).6'] == index_number]

        dt_list = list(sorted_pick_df['day'].unique())
        day_length_dic = {}
        for each_day in dt_list:
            day_df = sorted_pick_df[sorted_pick_df['day'] == each_day]
            day_length_dic[each_day] = day_df['length passed through.1'].max() - day_df['length passed through.1'].min()

        ''' Reading viscosity file with skipping 2 rows '''
        viscosity_df['Mixing date'] = pd.to_datetime(viscosity_df['Mixing date'])
        sorted_viscosity_df = viscosity_df.sort_values(by="Mixing date", ascending=True)
        sorted_viscosity_df['day'] = sorted_viscosity_df['Mixing date'].dt.date
        sorted_viscosity_df['day'] = sorted_viscosity_df['day'].astype('str')
        sorted_viscosity_df['batch-date'] = 'Batch_' + viscosity_df['Batch No.'].astype(
            'float').astype(str) + '_' + sorted_viscosity_df['day'].astype(str)
        sorted_viscosity_df = sorted_viscosity_df[sorted_viscosity_df['Index No'] == index_number]

        weight_date_dict = {}
        weight_batch_dict = {}
        for each_day in dt_list:
            day_df = sorted_viscosity_df[sorted_viscosity_df['day'] == each_day]
            summed = day_df['Input rubber weight(0.1kg)'].astype('float64').sum()
            weight_date_dict[each_day] = summed
            weight_batch_dict[each_day] = summed

        date_dict = get_mixer_batch_date(raw_df, index_number)
        bof_test_df = return_batch_no_bof_df(raw_df, sorted_viscosity_df, date_dict, index_number)
        bof_date_dict = get_bof_batch_date(bof_test_df, index_number)
        pick_merged_batch_df = return_pick_batch_no_df(raw_df, sorted_viscosity_df, bof_date_dict,
                                                       bof_test_df, index_number)

        # Merging pick data with viscosity data on date-batch column
        pickup_merged_df_final = pd.merge(pick_merged_batch_df, sorted_viscosity_df[['batch-date', 'viscosity']],
                                          on='batch-date', how='left')
        # Removing batch 0
        pickup_merged_df_final = pickup_merged_df_final[pickup_merged_df_final['pickup_batch_number'] != 0]
        # Grouping with aggregated data
        df_pickup_grouped = pickup_merged_df_final.groupby(PickupConstants.pick_grouped_cols).agg(
            PickupConstants.pick_aggregate_dict).reset_index()
        col_renamer = {}
        for col, col_agg in PickupConstants.pick_aggregate_dict.items():
            if col not in ['viscosity', 'time_min', 'time_max', 'Mixing Weight (Integrated Value)_diff',
                           'max_rpm_count']:
                renamed_col = f'{col.replace("(", "").replace(")", "").replace(" ", "_")}_{col_agg}'.lower()
                col_renamer[col] = renamed_col
            else:
                col_renamer[col] = col
        df_pickup_grouped = df_pickup_grouped.rename(columns=col_renamer)
        df_pickup_grouped_visc = df_pickup_grouped.drop('viscosity', axis=1)
        return df_pickup_grouped_visc

    except Exception as err:
        logger.error(f'Error while performing main function for pickup section {str(err)}')
        logger.error(traceback.format_exc())
        raise Exception(str(err))