try:
    import model
except:
    import input.model as model

import pymor.basic as pmb

import os
from tqdm import tqdm
import logging
import argparse


def get_env():
    '''Return instance number id instance number count when running on Qarnot platform'''
    return int(os.environ.get('INSTANCE_ID', '0')), int(os.environ.get('INSTANCE_COUNT', '1'))


def load_training_set(options, parameter_space, instance_id, instance_count):
    '''Load parameter set and divide it among several instances 
    Arguments:
        options: parsed arguments containing input_param_file and nb_param
        parameter_space (pymor.ParameterSpace): the space the parameters are in
        instance_id (int): instance number
        instance_count (int): number of instances'''
    param_set = []
    if options.input_param_file:
        param_set = pmb.load(open(options.input_param_file, 'rb'))
    elif options.nb_param > 0:
        param_set = parameter_space.sample_randomly(options.nb_param)
    else:
        param_set = parameter_space.sample_randomly(instance_count)
        
    p_per_inst = n = len(param_set)
    if instance_count > 1:
        p_per_inst = n//(instance_count-1)

    return param_set[instance_id*p_per_inst:(instance_id+1)*p_per_inst]


def set_log_level(inst_id):
    # Only instance 0 should display logs for easier reading
    if inst_id > 0:
        logging.getLogger('FFC').setLevel(logging.WARNING)
        pmb.set_log_levels({'pymor': 'WARN'})


def make_parser():
    parser = argparse.ArgumentParser(description='It creates the model, solves it for the specified parameter and write the result to a hdf5 file.',
                                     prefix_chars='-')
    parser.add_argument('-i', '--input_param_file',
                        help='set input parameter file',
                        type=str,
                        default='')
    parser.add_argument('-n', '--nb_param',
                        help='set total number of parameter to be sampled randomly (overwritten by i)',
                        type=int,
                        default=0)
    parser.add_argument('-o', '--output_file',
                        help='set output file name. Defaults to u',
                        type=str,
                        default='u')

    return parser


# Script that will be executed on Qarnot's node. 
# It creates the model, solves it for the specified parameter and write the result to a hdf5 file.
# It is made to run on multiple Qarnot instance.
# Options are the following:
#   -i   set input parameter file 
#   -n   set total number of parameter to be sampled randomly (overwritten by i). Defaults to instance count
#   -o   set output file name. Defaults to 'u'
if __name__ == '__main__':
    iid, icount = get_env()
    set_log_level(iid)
    
    fom = model.make_fom('mesh.xml')
    param_space = model.make_param_space()
    
    args = make_parser().parse_args()
    
    training_set = load_training_set(args, param_space, iid, icount)
    
    U_train = fom.solution_space.empty()
    t_used = model.used_time()
    
    # instance number 0 displays task advancement through tqdm for better monitoring
    if iid == 0:
        for mu in tqdm(training_set):
            U_train.append(fom.solve(mu)[t_used])
    else:
        for mu in training_set:
            U_train.append(fom.solve(mu)[t_used])
    output_file = args.output_file
    model.dump_sol_list(f'{output_file}{iid}.h5', U_train, 'u')
