from pucktrick.utils import *

def duplicate(train_df, strategy=None, original_df=None ):

    affected_features = strategy.get("affected_features")
    selection_criteria = strategy.get("selection_criteria")
    percentage = strategy.get("percentage")
    mode = strategy.get("mode")
    function=strategy.get('function')
    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
    isAffectedSelected=0
    if affected_features and not isinstance(affected_features, list):
        affected_features = [affected_features]
    
    if selection_criteria=='all':
            if original_df is None:
                original_df = train_df.copy()
                original1_df = train_df.copy()

            else:
                original1_df=original_df.copy()
    else:
        if mode=="extended":
            selection_criteria,isAffectedSelected=affectedSelected(affected_features,selection_criteria)
        else:
           isAffectedSelected=0 
        if isAffectedSelected==1:
            if original_df is not None:
                df_pref=train_df[affected_features].add_prefix("puck_")
                df_pref2=original_df[affected_features].add_prefix("puck_")
                original_df=pd.concat([original_df,df_pref2],axis=1)
                train_df=pd.concat([train_df,df_pref],axis=1)
                
        if original_df is None:
                original_df = train_df.copy()
                original1_df = train_df.copy()
        else:
            original1_df=train_df.copy()
            original_df=original_df.query(selection_criteria).copy() 
        train_df=train_df.query(selection_criteria).copy()
    if function in [
        "class",
        "upper_lower",
        "replace_punctuation",
        "remove_replace",
        "abbreviate_text",
        "shuffle_words",
    ]:
        # Costruzione dei filtri
        filters = []

        if value is None or not any(value):
            filters = [
                (original_df[col].notnull() if mode == "extended" else train_df[col].notnull())
                for col in affected_features
            ]
        else:
            for col, val in zip(affected_features, value):
                if val is not None:
                    filters.append(
                        (original_df[col] == val) if mode == "extended" else (train_df[col] == val)
                    )
                else:
                    filters.append(
                        (original_df[col].notnull()) if mode == "extended" else (train_df[col].notnull())
                    )

        # Combina i filtri con logica "or" o "and"
        if condition_logic == "and":
            combined_filter = pd.Series(True, index=train_df.index if mode != "extended" else original_df.index)
            for f in filters:
                combined_filter &= f
        elif condition_logic == "or":
            combined_filter = pd.Series(False, index=train_df.index if mode != "extended" else original_df.index)
            for f in filters:
                combined_filter |= f
        elif condition_logic is None:
            # Default: "and" logic
            combined_filter = pd.Series(True, index=train_df.index if mode != "extended" else original_df.index)
            for f in filters:
                combined_filter &= f
        else:
            raise ValueError(f"Unsupported condition_logic '{condition_logic}'")

        # Filtrare le righe
        target_df = original_df if mode == "extended" else train_df
        filtered_df = target_df[combined_filter].reset_index(drop=True)
        rowsToChange = int(len(filtered_df) * percentage)
        if mode == "extended":
            old_duplicated = len(train_df) - len(original_df)
            new_rowsToChange = rowsToChange - old_duplicated
      
        if new_rowsToChange <= 0:
                return 1,train_df
        percentage = new_rowsToChange / len(filtered_df)
        sampled_indices = sample_indices(distr1, filtered_df, percentage, params,number=None)
        sampled_rows = filtered_df.iloc[sampled_indices].copy()
        for col in affected_features:
            if selection_criteria == "shuffle_words":
                sampled_rows[col] = sampled_rows[col].astype(str).apply(shuffle_words)
            elif selection_criteria == "abbreviate_text":
                sampled_rows[col] = sampled_rows[col].astype(str).apply(abbreviate_text)
            elif selection_criteria == "replace_punctuation":
                sampled_rows[col] = sampled_rows[col].astype(str).apply(replace_punctuation)
            elif selection_criteria == "remove_replace":
                sampled_rows[col] = sampled_rows[col].astype(str).apply(remove_or_replace)
            elif selection_criteria == "upper_lower":
                sampled_rows[col] = random_upper_lower(sampled_rows[col])

        noise_df = pd.concat([train_df, sampled_rows], ignore_index=True)

    if mode == "extended":
            rowsToChange = len(original_df) * percentage
            
            if selection_criteria=='all':
                 old_duplicated = original_df.duplicated().sum()
            else:
                 old_duplicated = len(original1_df.query(selection_criteria))-len(original_df)
            #new_duplicated = train_df.duplicated().sum()
            #new_duplicated =len(train_df)
            #diff = new_duplicated - old_duplicated
            #new_rowsToChange = rowsToChange - diff
            new_rowsToChange = rowsToChange - old_duplicated
           
            if new_rowsToChange <= 0:
                return 1,train_df
            new_percentage = new_rowsToChange / len(original_df)
            noise_df = train_df.copy()
            sampled_indices = sample_indices(distr1, original_df, new_percentage, number, params)
            sampled_rows = original_df.iloc[sampled_indices]
            #noise_df = pd.concat([noise_df, sampled_rows], ignore_index=True)
            noise_df=pd.concat([original1_df,sampled_rows],ignore_index=True)
            
    else:
            noise_df = original1_df.copy()
            sampled_indices = sample_indices(distr1, train_df, percentage, number, params)
            sampled_rows = train_df.iloc[sampled_indices]
            noise_df = pd.concat([noise_df, sampled_rows], ignore_index=True)
    if "date" in noise_df.columns:
        noise_df["date"] = noise_df["date"].astype(str).str.replace(r"\s00:00:00$", "", regex=True)
    
    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)
   
    return 0,noise_df

