import warnings
import pandas as pd

from scripts.constants.constants import ViscosityConstants

warnings.filterwarnings("ignore")


def create_status_column(df, type_col_name, columns_list):
    status_col = []

    for i, val in enumerate(df[type_col_name]):
        if val == 0:
            status_col.append(False)
        else:
            if any(df[column].iloc[i] == 0 for column in columns_list):
                status_col.append(False)
            else:
                status_col.append(True)

    return status_col


def preprocess_viscosity_section(viscosity_df, index_number):
    # adding date col to the viscosity df
    viscosity_df = viscosity_df.sort_values(by='Mixing date')
    raw_viscosity_df = viscosity_df.sort_values(by='Mixing date')
    viscosity_df['date'] = viscosity_df['Mixing date'].dt.date
    viscosity_df['batch-date'] = 'Batch_' + viscosity_df['Batch No.'].astype(str) + '_' + viscosity_df['date'].astype(
        str)
    viscosity_df = viscosity_df[viscosity_df['Index No'] == index_number]
    rubber_cols = ViscosityConstants.rubber_cols
    # Replace '-' with 0 for numerical and float columns
    viscosity_df[rubber_cols] = viscosity_df[rubber_cols].replace('-', 0)
    viscosity_df[rubber_cols] = viscosity_df[rubber_cols].apply(pd.to_numeric, errors='coerce')

    # Identify numerical and float columns
    numerical_cols = viscosity_df.columns[
        viscosity_df.dtypes.apply(lambda x: pd.api.types.is_numeric_dtype(x) or pd.api.types.is_float_dtype(x))]
    integer_cols = viscosity_df.columns[viscosity_df.dtypes == 'int64']

    # Convert integer columns to float
    viscosity_df[integer_cols] = viscosity_df[integer_cols].astype(float)
    # Calculate weights
    viscosity_df['Weight_type1'] = round(viscosity_df['Quantity using type1 bale'] / (
            viscosity_df['Quantity using type1 bale'] + viscosity_df['Quantity using type2 bale']), 2)
    viscosity_df['Weight_type2'] = round(viscosity_df['Quantity using type2 bale'] / (
            viscosity_df['Quantity using type1 bale'] + viscosity_df['Quantity using type2 bale']), 2)
    viscosity_df['Weighted_PO_type'] = (
            viscosity_df['PO_type1'] * viscosity_df['Weight_type1'] + viscosity_df[f'PO_type2'] * viscosity_df[
        'Weight_type2'])
    viscosity_df['Weighted_DIRT_type'] = (
            viscosity_df['DIRT_type1'] * viscosity_df['Weight_type1'] + viscosity_df['DIRT_type1.1'] * viscosity_df[
        'Weight_type2'])
    viscosity_df['Weighted_ASH_type'] = (
            viscosity_df['ASH_type1'] * viscosity_df['Weight_type1'] + viscosity_df['ASH_type2'] * viscosity_df[
        'Weight_type2'])
    viscosity_df['Weighted_VM_type'] = (
            viscosity_df['VM_type1'] * viscosity_df['Weight_type1'] + viscosity_df['VM_type2'] * viscosity_df[
        'Weight_type2'])

    viscosity_df['Weighted_PRI_type'] = (
            viscosity_df['PRI_type1'] * viscosity_df['Weight_type1'] + viscosity_df[f'PRI_type2'] * viscosity_df[
        'Weight_type2'])

    viscosity_df['Weighted_NITROGEN_type'] = (
            viscosity_df['NITROGEN_type1'] * viscosity_df['Weight_type1'] + viscosity_df['NITROGEN_type2'] *
            viscosity_df['Weight_type2'])

    viscosity_df['Weighted_Temperature during transportation_type[℃]'] = (
            viscosity_df['Temperature during transportation_type1[℃]'] * viscosity_df['Weight_type1'] +
            viscosity_df['Temperature during transportation_type2[℃]'] * viscosity_df['Weight_type2'])
    viscosity_df['Weighted_Humidity during transportation__type[%]'] = (
            viscosity_df['Humidity during transportation_type1[%]'] * viscosity_df['Weight_type1'] +
            viscosity_df['Humidity during transportation__type2[%]'] * viscosity_df['Weight_type2'])
    viscosity_df['Weighted Sum'] = viscosity_df['Weighted_PO_type'] + viscosity_df['Weighted_DIRT_type'] + viscosity_df[
        'Weighted_ASH_type'] + viscosity_df['Weighted_VM_type'] + viscosity_df['Weighted_PRI_type'] + viscosity_df[
                                       'Weighted_NITROGEN_type']
    column_to_keep_at_end = 'viscosity'

    # Reorder columns
    new_order = [col for col in viscosity_df.columns if col != column_to_keep_at_end] + [column_to_keep_at_end]
    viscosity_df = viscosity_df[new_order]
    viscosity_df['batch-date'] = 'Batch_' + viscosity_df['Batch No.'].astype(str) + '_' + viscosity_df['date'].astype(
        str)
    # Added Status to check rubber
    # Rubber Type 1
    rubber_1_cols = [
                     'DIRT_type1',
                     'ASH_type1',
                     'VM_type1',
                     'PRI_type1',
                     'NITROGEN_type1'
                     ]
    # Rubber Type 2
    rubber_2_cols = [
                     'PO_type2',
                     'DIRT_type1',
                     'ASH_type2',
                     'VM_type2',
                     'PRI_type2',
                     'NITROGEN_type2'
                     ]
    viscosity_df['rubber_status_1'] = create_status_column(viscosity_df, 'Weight_type1', rubber_1_cols)
    viscosity_df['rubber_status_2'] = create_status_column(viscosity_df, 'Weight_type2', rubber_2_cols)
    viscosity_df['status'] = viscosity_df['rubber_status_1'] | viscosity_df['rubber_status_2']
    req_cols = ViscosityConstants.req_cols
    final_viscosity_df = viscosity_df[req_cols]
    final_viscosity_df = round(final_viscosity_df, 6)
    return final_viscosity_df, raw_viscosity_df
