import gc

import pandas as pd


target_col = 'HasDetections'


gbdt_csvs = [
    './data/output/submit_lgb_gbdt_seed1.csv.gz',
    './data/output/submit_lgb_gbdt_seed2.csv.gz',
    './data/output/submit_lgb_gbdt_seed3.csv.gz',
    './data/output/submit_lgb_gbdt_seed4.csv.gz',
    './data/output/submit_lgb_gbdt_seed5.csv.gz',
    './data/output/submit_lgb_gbdt_seed6.csv.gz',
    './data/output/submit_lgb_gbdt_seed7.csv.gz',
    './data/output/submit_lgb_gbdt_seed8.csv.gz',
]


dart_csvs = [
    './data/output/submit_lgb_dart_seed2020.csv.gz',
    './data/output/submit_lgb_dart_seed2021.csv.gz',
    './data/output/submit_lgb_dart_seed2022.csv.gz',
    './data/output/submit_lgb_dart_seed2023.csv.gz',
    './data/output/submit_lgb_dart_seed2024.csv.gz',
]


stage1_csvs = {
    'gbdt': './data/output/submit_lgb_gbdt_seed_avg.csv.gz',
    'dart': './data/output/submit_lgb_dart_seed_avg.csv.gz',
    'nn': './data/output/submit_nn_seed12345.csv.gz',
}


def seed_avg(files, output_filename):
    sub = pd.read_csv(files[0])
    for csv_file in files[1:]:
        data = pd.read_csv(csv_file)
        sub[target_col] += data[target_col]
    sub[target_col] /= len(files)
    print(sub.head())
    sub.to_csv(output_filename, index=False)


def rank_avg(files, weights, output_filename):
    assert len(files) == len(weights)

    data = {}
    for i, csv_file in enumerate(files):
        data[str(i)] = pd.read_csv(csv_file)

    ranks = pd.DataFrame(columns=data.keys())
    for k in data.keys():
        ranks[k] = data[k][target_col].rank(method='min')

    ranks['score'] = ranks[list(data.keys())].mul(weights).sum(1) / ranks.shape[0]

    del data
    gc.collect()

    sub = pd.read_csv(files[0])
    sub[target_col] = ranks['score']
    print(sub.head())
    sub.to_csv(output_filename, index=False)


def print_rank_corr():
    data = {}
    for k, csv_file in stage1_csvs.items():
        data[k] = pd.read_csv(csv_file)[target_col].rank()

    base_preds = pd.DataFrame({k: v for k, v in data.items()})
    print(base_preds.corr(method='spearman'))


if __name__ == '__main__':
    seed_avg(gbdt_csvs, stage1_csvs['gbdt'])
    seed_avg(dart_csvs, stage1_csvs['dart'])

    #print_rank_corr()

    weights = [0.4, 0.35, 0.25]
    rank_avg(list(stage1_csvs.values()), weights, './data/output/final_submission.csv.gz')

