#!/usr/bin/env python
import pickle

import numpy as np

from Main.MainArguments import MainArguments
from Main.Parameter import Parameter
from Model.ModelBuilder import ModelBuilder
from Utils import FilenameGenerator

"""
-model /d/test/tmp/
-total_part 2
-ion 0
"""


def main():
    deal_with_arguments()
    model = ModelBuilder.load_from_file(FilenameGenerator.feature_data(Parameter.model),
                                        load_training_data=False)

    model_num = model.get_total_feature_num()
    batch_size = model_num // Parameter.total_part + 1

    fo = open(FilenameGenerator.merged_peak_model_data(Parameter.model), 'wb')
    overview = []
    for cur_part in range(Parameter.total_part):
        model_start_num = batch_size * cur_part

        filename_output = FilenameGenerator.peak_model_data(Parameter.model) + '/peak_model_' + str(model_start_num)
        fi_data = open(filename_output + '.data', 'rb')

        cur_file_place = fo.tell()
        fo.write(fi_data.read())
        cur_overview = pickle.load(open(filename_output + '.overview', 'rb'))
        cur_overview = np.array(cur_overview)
        cur_overview[:, 1] += cur_file_place
        overview.append(cur_overview)
        pass
    overview = np.concatenate(overview, axis=0)
    overview = overview[:, [1, 2]]
    model.add_data_overview(overview)
    del model.mol_id
    del model.spec_id
    del model.uniq_mol_id
    pickle.dump(model, open(FilenameGenerator.merged_peak_model_info(Parameter.model), 'wb'), 1)
    pass


def deal_with_arguments():
    parser = MainArguments(description='')

    parser.add_argument('-total_part', required=False, type=int, dest='total_part', default=1,
                        metavar='', help='')

    parser.parse_args()


if __name__ == '__main__':
    main()
