import warnings
import pandas as pd

from scripts.constants.constants import ViscosityConstants

warnings.filterwarnings("ignore")


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)
    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
