import os
import pandas as pd
import random
import datetime
import time


def neg_under_sample(data_dir, sample_data_dir, dt):
    # 欠采样，对负样本按照30%的概率进行舍弃，pos:neg ≈ 1:10
    data_df = pd.read_csv(f"{data_dir}/{dt}.csv")
    pos_df = data_df[data_df['share_status'] == 1].copy()
    neg_df = data_df[data_df['share_status'] == 0].copy()
    neg_df['rand'] = [random.uniform(0, 1) for i in range(neg_df.shape[0])]
    rate = 1 - pos_df.shape[0] / neg_df.shape[0]
    sample_neg_df = neg_df[neg_df['rand'] > rate]
    sample_neg_df = sample_neg_df.drop(columns=['rand'])
    sample_df = pd.concat([pos_df, sample_neg_df])
    sample_df = sample_df.sample(frac=1.0).reset_index(drop=True)
    # 写入csv
    if not os.path.exists(sample_data_dir):
        os.makedirs(sample_data_dir)
    sample_df.to_csv(f"{sample_data_dir}/{dt}.csv", index=False)
    print(f"pos num: {pos_df.shape[0]}")
    print(f"neg num: {neg_df.shape[0]}")
    print(f"neg_sample num: {sample_neg_df.shape[0]}")
    print(f"initial data num: {data_df.shape[0]}")
    print(f"sample data num: {sample_df.shape[0]}")


if __name__ == '__main__':
    st_time = time.time()
    data_dir = './data/train_data'
    sample_data_dir = './data/sample_train_data'
    now_date = datetime.datetime.today()
    print(f"now_date: {datetime.datetime.strftime(now_date, '%Y%m%d')}")
    dt = datetime.datetime.strftime(now_date - datetime.timedelta(days=15), '%Y%m%d')
    print(f"update data dt: {dt}")
    neg_under_sample(data_dir=data_dir, sample_data_dir=sample_data_dir, dt=dt)
    print(f"{time.time() - st_time}s")


    # for days in range(3, 19):
    #     cur_dt = datetime.datetime.strftime(now_date - datetime.timedelta(days=days), '%Y%m%d')
    #     print(f"cur_dt = {cur_dt}")
    #     neg_under_sample(data_dir=data_dir, sample_data_dir=sample_data_dir, dt=cur_dt)