from pucktrick.utils import *
def labels(train_df, strategy, original_df=None):

    target = strategy.get("affected_features")
    selection_criteria = strategy.get("selection_criteria")
    percentage = strategy.get("percentage")
    mode = strategy.get("mode")
    perturbate_config = strategy.get("perturbate_data", {})
    distr1 = perturbate_config.get("distribution")
    value = perturbate_config.get("value", [])
    condition_logic = perturbate_config.get("condition_logic", None)
    params= perturbate_config.get("param", {})
    number=None
    target_col = target[0] if isinstance(target, list) else target
    unique_values = train_df[target_col].dropna().unique()
    isAffectedSelected=0
    if selection_criteria=="all":
         
        if original_df is None:
            original_df = train_df.copy()
            original1_df=original_df
        else:
            original1_df=original_df.copy()
        
    else:
        if mode=="extended":
            selection_criteria,isAffectedSelected=affectedSelected(target,selection_criteria)
        else:
           isAffectedSelected=0 
        if isAffectedSelected==1:
            if original_df is not None:
                df_pref2=original_df[target].add_prefix("puck_")
                original_df=pd.concat([original_df,df_pref2],axis=1)
                train_df=pd.concat([train_df,df_pref2],axis=1)
         
        if original_df is None:
            original1_df=  train_df.copy()
            original_df = train_df.query(selection_criteria).copy()         
        else:
            original1_df=original_df.copy()
            original_df=original_df.query(selection_criteria).copy()
        train_df=train_df.query(selection_criteria).copy()
        

    #unique_values = train_df[target_col].dropna().unique()  
    if len(unique_values) == 2:
        var_type = "binary"
    elif len(unique_values) > 2:
        var_type = "categorical"
    else:
        raise ValueError(f"Unable to determine the type of target column '{target}'.")
    
   
    if var_type == "binary":
        #print("Target is binary")
        if mode == 'extended':
            noise_df= noiseBinaryExt(original_df, train_df, target_col,distr1, percentage,number,params)
        else:
            noise_df = noiseBinaryNew(train_df, target_col, distr1, percentage,params,number) 

    elif var_type == "categorical":
        if mode == 'extended':
            noise_df = noiseCategoricalIntExtendedExistingValues(original_df, train_df, target_col,distr1, percentage,number,params)

        else:
            noise_df = noiseCategoricalIntNewExistingValues(train_df, target_col, distr1, percentage,number,params) 
    for col in noise_df.select_dtypes(include=['int64']).columns:
        noise_df[col] = noise_df[col].astype('Int64')  
    if isAffectedSelected==1:
        puck_cols=[c for c in noise_df.columns if c.startswith("puck_")]
        noise_df=noise_df.drop(columns=puck_cols)
        original1_df=original1_df.drop(columns=puck_cols)
    
    original1_df.loc[noise_df.index, noise_df.columns]=noise_df
    return 0,original1_df
