import json
import pandas as pd


file_names = [#'./llm_for_code/datasets/codecontests/CN_py800_catmix_5rand_4fold_full_oripairs_matched.json'
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_oripairs_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_oripairs_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m2_oripairs_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m3_oripairs_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m4_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_0_2_oripairs.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_rem_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_rem2_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_0_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_101_2_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_101_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_110_oripairs.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_2_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_3_oripairs_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_100_4_oripairs_temp.json',
#"./dataset_lfs/allpairs_data_val_900_unperturbed.json",
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_2_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_3_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_4_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_5_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_1_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_101_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_2_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_3_oripairs.json',
#'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_3_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_500_1_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_500_2_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0s_conc_1_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_1_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_0s_500_0_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_pos_2_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_pos_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_5_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_1_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_10s_500_0_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_10s_500_2_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_3_oripairs.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_oripairs_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_5_oripairs_temp.json',



]

file_names_per = [#'./dataset_lfs/CC_900_py_catmix_5rand_4fold_pairs_0_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_pairs_0_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m2_pairs_0_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m3_pairs_0_temp.json',
# './dataset_lfs/CN_py800_catmix_5rand_4fold_200_7m4_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_0_2_pairs_0.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_rem_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_101_3_rem2_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_0_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_101_2_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_101_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_200_110_pairs_0.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_2_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_1s_1k_3_pairs_0_temp.json',
# './dataset_lfs/CC_900_py_catmix_5rand_4fold_100_4_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_2_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_3_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_4_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_5_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_1_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_101_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_2_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_3_pairs_0.json',
#'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_3_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_500_1_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_500_2_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0s_conc_1_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_1_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_0_conc_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_0s_500_0_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_pos_2_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_pos_pairs_0_temp.json',

'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_5_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_1_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_10s_500_0_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_10s_500_2_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_2_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_3_pairs_0.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_all_conc_pairs_0_temp.json',
'./dataset_lfs/CC_900_py_catmix_5rand_4fold_train_conc_5_pairs_0_temp.json',
]
file1 = './llm_for_code/datasets/codecontests/CC_900_py_catmix_5rand_4fold_oripairs_temp_fulltry.json'
file2 = './llm_for_code/datasets/codecontests/CC_900_py_catmix_5rand_4fold_pairs_0_temp_fulltry.json'

filt1 = []
with open(file1, 'r') as file:
    for line in file:
        filt1.append(json.loads(line))

filt2 = []
with open(file2, 'r') as file:
    for line in file:
        filt2.append(json.loads(line))

filtdf = pd.DataFrame(filt1)
filtdf_per = pd.DataFrame(filt2)
original_columns_df2 = filtdf_per.columns.tolist()
filtdf_per.columns = [f"{col}_df2" for col in filtdf_per.columns]
filtdf_combined = pd.concat([filtdf, filtdf_per], axis=1)        
data = []
data_per = [[], [], [], []]
df_per_list = []
df_per = None
for fn in file_names:
    with open(fn, 'r') as file:
        for line in file:
            data.append(json.loads(line))

print(len(data), len(data_per))            
df = pd.DataFrame(data)

df_combined = df.copy(deep=True)
for per_id in range(4):            
    for file in file_names_per:
        fn = file[:file.rfind('0')] + str(per_id) + file[file.rfind('0')+1:]
        #print(fn)
        with open(fn, 'r') as file:
            for line in file:
                data_per[per_id].append(json.loads(line))
    df_per = pd.DataFrame(data_per[per_id])           
    df_per_list.append(df_per)
    original_columns_df2 = df_per.columns.tolist()

    df_per.columns = [f"{col}_df{per_id+2}" for col in df_per.columns]
    df_combined = pd.concat([df_combined, df_per], axis=1)
print(df_combined.head(), len(df_combined))
df_combined = df_combined[~df_combined.apply(tuple, axis=1).isin(filtdf_combined.apply(tuple, axis=1))]
prohib_probid= [
    x.get('probid') for x in filtdf['prog1'] 
    if isinstance(x, dict) and 'probid' in x
] 
prohib_probid += [
    x.get('probid') for x in filtdf['prog2'] 
    if isinstance(x, dict) and 'probid' in x
]

# Extract 'A1' from df and filter
df_combined = df_combined[ ~df['prog1'].apply(lambda x: x.get('probid') if isinstance(x, dict) else None).isin(prohib_probid)]
df_combined = df_combined[ ~df['prog2'].apply(lambda x: x.get('probid') if isinstance(x, dict) else None).isin(prohib_probid)]
df_combined = df_combined[ ~df['prog1'].apply(lambda x: x.get('probid') if isinstance(x, dict) else None).str.startswith('p0').fillna(False)]

print("filtered, ", len(df_combined))
# Step 2: Update only the 'cat' column
#df_combined['category'] = df_combined['category'].str[:-1] + df_combined['category_df2'].str[-1]
# Define a custom function to compare dictionaries in prog1 and prog2
def compare_dicts(row):
    return (str(row['prog1']), str(row['prog2']))  # Convert dictionaries to tuple of strings for comparison

# Apply the custom function and mark duplicates

df_combined['unique_check'] = df_combined.apply(compare_dicts, axis=1)

# Drop duplicates based on the unique_check column
df_unique = df_combined.drop_duplicates(subset=['unique_check'])

# Drop the helper column used for comparison
df_unique = df_unique.drop(columns=['unique_check'])
print(len(data),len(df_unique) ," are unique")     
#df = pd.DataFrame(data)
print(df['category'].value_counts())

def get_sample_size(category):
    #catlen = {'2_0_0_1': 0, '2_0_0_0': 73,'1_1_0_1': 151 , '1_1_0_0': 0, '0_1_0_0': 0, '0_1_0_1': 2, '1_0_0_0' : 112, '1_0_0_1':0 }
    #return catlen[category]
    # For example: sample 10% of the group size, with a max of 50
    # if category == '0_1_0_1':
        # return 3600
    # if category == '1_1_0_1':
        # return 3600
    # # if category[0] == '1':
        # # return 1000
    # return 1200
    return 400
# Apply the custom sampling
sampled_df = (
    df_unique.groupby('category', group_keys=False)
    .apply(lambda x: x.sample(n=min(get_sample_size(x['category'].iloc[0]),len(x)), replace=False))
)
#sampled_df1  = df[~(df['category'].str[0] == '1')]
#sampled_df = sampled_df[sampled_df['category'].str.startswith('2')]


# # Step 1: Identify and separate unique and duplicate rows
# df_unique = df_combined.drop_duplicates(subset=['unique_check'])




# df_duplicates = df_combined[~df_combined.index.isin(df_unique.index)]
# df_duplicates = df_duplicates.drop(columns=['unique_check'])
# # Step 2: Sample from df_unique first

# # Step 3: Identify categories that still need more samples
# remaining_counts = (
    # sampled_df['category'].value_counts()
    # .reindex(df_combined['category'].unique(), fill_value=0)
# )

# print(remaining_counts)
# # Step 4: Fill missing samples from df_duplicates
# additional_samples = []
# for category, count in remaining_counts.items():
    # required = get_sample_size(category) - count
    # if required > 0:
        # extra_samples = df_duplicates[df_duplicates['category'] == category].sample(
            # n=min(required, len(df_duplicates[df_duplicates['category'] == category])),
            # replace=False
        # )
        # additional_samples.append(extra_samples)

# # Concatenate the extra samples with the originally sampled DataFrame
# if additional_samples:
    # additional_samples_df = pd.concat(additional_samples)
    # sampled_df = pd.concat([sampled_df, additional_samples_df])

# Final sampled dataframe
sampled_df = sampled_df.reset_index(drop=True)





print(sampled_df['category'].value_counts(), len(sampled_df), sampled_df.columns)    

df1_final = sampled_df.iloc[:, :df.shape[1]]
df2_final_list = []
for per_id in range(4):
    df2_final = sampled_df.iloc[:, df.shape[1]+df_per.shape[1]*(per_id):df.shape[1]+df_per.shape[1]*(per_id+1)]
    print(df1_final.head(), df2_final.head())
    df2_final.columns = original_columns_df2
    df2_final_list.append(df2_final)
df2_final = pd.concat(df2_final_list, ignore_index=True)
print(len(df2_final))
with open(file_names[0][:-5]+"_fulltrydup_train_2nem.json" , 'w+') as outfile:  
    for _, sample in df1_final.iterrows():
        outfile.write(f"{json.dumps(sample.to_dict())}\n")
       
with open(file_names_per[0][:-5]+"_fulltrydup_train_2nem.json" , 'w+') as outfile:  
    for _, sample in df2_final.iterrows():
        outfile.write(f"{json.dumps(sample.to_dict())}\n")