import pickle
import gzip
import glob
from pathlib import Path
from typing import Dict
import json
from multiprocessing import Process
import argparse
import os, sys

def get_sample(filename):
    with gzip.open(filename, 'rb') as f:
        raw_data = pickle.load(f)
    return raw_data

def getInstanceToSamples(files, id = 0):
    instance_to_sample: Dict[str, list]= {}
    count = 0
    maxCount = len(files)
    with open('logs/log_{}.txt'.format(id), 'w') as f:
        f.write('')
    for sample_name in files:
        count += 1
        if count % 250 == 0:
            with open('logs/log_{}.txt'.format(id), 'a') as f:
                f.write('Instance to samples: {} / {}\n'.format(count, maxCount))
        sample = get_sample(sample_name)
        instance_name: str = sample['instance'].name
        if instance_name not in instance_to_sample:
            instance_to_sample[instance_name] = []
        instance_to_sample[instance_name].append({
            'sample_name': sample_name,
            'action_number':  len(sample['data'][2])
        })
    return instance_to_sample

def storeDepthInfo(instance_to_sample):
    for instance_name, samples in instance_to_sample.items():
        l = list(map(lambda x:x['action_number'], samples))
        max_action_number = max(l)
        min_action_number = min(l)
        for sample in samples:
            info = {
                'depth': max_action_number - sample['action_number'],
                'max_depth': max_action_number - min_action_number
            }
            info_file_name = Path(sample['sample_name']).with_suffix('.json')
            with open(str(info_file_name), 'w') as f:
                json.dump(info, f)

def ProceedDepth(files, ind=0):
    storeDepthInfo(getInstanceToSamples(files, ind))

def ProceedDepthMulti(files, n_jobs=8):
    blockSize = len(files) // n_jobs
    processList = []
    for i in range(n_jobs):
        begin = i * blockSize
        end = min((i+1) * blockSize, len(files))
        file_block = files[begin: end]

        process = Process(target=ProceedDepth, args=(file_block,i,))
        process.start()
        processList.append(process)

    for process in processList:
        process.join()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'problem',
        help='MILP instance type to process.',
        choices=['item_placement', 'load_balancing', 'anonymous'],
    )
    args = parser.parse_args()

    Path('logs/').mkdir(exist_ok=True)
    # get sample directory
    if args.problem == 'item_placement':
        train_files = glob.glob('train_files/samples/1_item_placement/train/sample_*.pkl')
        valid_files = glob.glob('train_files/samples/1_item_placement/valid/sample_*.pkl')

    elif args.problem == 'load_balancing':
        train_files = glob.glob('train_files/samples/2_load_balancing/train/sample_*.pkl')
        valid_files = glob.glob('train_files/samples/2_load_balancing/valid/sample_*.pkl')

    elif args.problem == 'anonymous':
        train_files = glob.glob('train_files/samples/3_anonymous/train/sample_*.pkl')
        valid_files = glob.glob('train_files/samples/3_anonymous/valid/sample_*.pkl')

    ProceedDepthMulti(train_files + valid_files, n_jobs=16)