import pymor.basic as pmb
from pymor.models.interface import Model
from tqdm import tqdm

import os
import argparse

import model as md


def load_remote(mod: Model, instance_count: int, load_dir: str):
    '''Load solutions and returns them
        Parameters:   
            mod (pymor.Model): The model whose solution_space the solutions lie in
            instance_count (int): The number of instance that calculated the 
                solutions, e.g. the number of file the solution are in
            load_dir (str). The direcory the solutions are in. They must be named u{i}_c.h5 
    '''
    U_list = mod.solution_space.empty()
    for i in tqdm(range(instance_count)):
        if f'u{i}_c.h5' in os.listdir(load_dir):
            U_list.append(md.load_sol_list(f'{load_dir}/u{i}_c.h5', U_list.space, radica='u'))
    return U_list


def make_parser():
    parser = argparse.ArgumentParser(description='Build a reduced order model for the model' \
                                     'specified in input with previously computed solution',
                                     prefix_chars='-')
    parser.add_argument('-n', '--nb_files',
                        help='number of result files',
                        type=int,
                        default=30)
    parser.add_argument('-i', '--input_dir',
                        help='input directory for the fom solutions. the files should be dir/u0_c.h5, ..., dir/u{n}_c.h5',
                        type=str,
                        default='train')
    parser.add_argument('-o', '--output_dir',
                        help='output directory for the model. Output files are dir/rom.out and dir/basis.h5',
                        type=str,
                        default='.')
    parser.add_argument('-m', '--modes',
                        help='number of mode (vectors) to keep in the reduction basis',
                        type=int,
                        default=100)

    return parser


# Build a reduced order model for the model specified in input with previously computed solution
# options are:
#   -n  number of result files (defaults to 30). If some files are missing because instances failed, 
# the script will carry on.
#   -i  input directory for the solutions. the files should be dir/u0_c.h5, ..., dir/u{n}_c.h5 (defaults to 'train')
#   -o  output directory for the model. Output files are dir/rom.out and dir/basis.h5. Defaults to '.'
#   -m  number of mode (vectors) to keep in the reduction basis (defaults to 100)
if __name__=='__main__':
    fom = md.make_fom('mesh.xml')
    
    args = make_parser().parse_args()
    
    print('\nloading solutions')
    U_train = load_remote(fom, args.nb_files, args.input_dir)
    print(f'loaded {len(U_train)} vectors of dim {U_train.dim}')
    
    # Create the reduction basis with POD algorithm
    pod_basis, pod_singular_values = pmb.pod(U_train, product=fom.h1_0_semi_product, modes=args.modes)
    
    # Project the model
    reductor = pmb.InstationaryRBReductor(fom, RB=pod_basis, product=fom.h1_0_semi_product)
    rom = reductor.reduce()
    
    # Convert to Parametric mass model
    rom = md.InstationaryParametricMassModel(rom.T, 
                                        initial_data=rom.initial_data,
                                        operator=rom.operator,
                                        rhs=rom.rhs,
                                        mass=rom.mass,
                                        time_stepper=rom.time_stepper,
                                        num_values=rom.num_values,
                                        output_functional=rom.output_functional,
                                        products=rom.products,
                                        error_estimator=rom.error_estimator,
                                        visualizer=rom.visualizer,
                                        name=rom.name)

    # Save the model
    output_dir = args.output_dir
    pmb.dump(rom, open(f'{output_dir}/rom.out', 'wb'))
    md.dump_sol_list(f'{output_dir}/pod_basis.out', pod_basis, radica='basis')
