data_generator

b_get_machine_operation_any_parameters(OperationName, Ids) :-
    b_get_machine_operation(OperationName, _, _, Body),
    b_get_machine_operation_any_parameters_aux([], Body, Ids),
    !.

b_get_machine_operation_any_parameters(_, []).

b_get_machine_operation_any_parameters_aux(Acc, b(Node,_,_), Ids) :-
    Node =.. [Functor,AnyIds|Args],
    ( Functor = any
    ; Functor = let
    ),
    !,
    maplist(b_get_machine_operation_any_parameters_aux([]), Args, TArgIds),
    flatten(TArgIds, ArgIds),
    append([Acc,ArgIds,AnyIds], Ids).

b_get_machine_operation_any_parameters_aux(Acc, b(Node,_,_), Ids) :-
    Node =.. [_|Args],
    maplist(b_get_machine_operation_any_parameters_aux([]), Args, TArgIds),
    flatten(TArgIds, ArgIds),
    append(Acc, ArgIds, Ids).

b_get_machine_operation_any_parameters_aux(Acc, _, Acc).

contains_symbolic_value(LoE) :-
    % TO DO: extend for lambda functions?
    member(E, LoE),
    contains_symbolic_value_aux(E).

contains_symbolic_value_aux(b(value(closure(_,_,b(member(_,_),_,_))),_,_)).

contains_symbolic_value_aux(b(Node,_,_)) :-
    Node =.. [_|Args],
    member(Arg, Args),
    contains_symbolic_value_aux(Arg).

create_example_triple(Env, Ast, (Value,NType,NVarName)) :-
    % (
    Env = [_,NormalizedSets,_,_],
    b_machine_identifier_normalization:normalize_ids_in_b_ast(Env, Ast, NAst, _), % TODO update env?
    % ; trace , normalize_ids_in_b_ast(NormalizedSets,NormalizedIds,Ast,NAst)) ,
    NAst = b(_,Type,NInfo),
    normalize_id_type(NormalizedSets, Type, NType),
    member(synthesis(machinevar,NVarName), NInfo),
    translate_bexpression(NAst, Value).

generate_data_from_operation_substitutions(_, _, _, [], _, Acc, Acc) :-
    !.

generate_data_from_operation_substitutions(Env, AR, AbsoluteFilePath-MachineName, [(OperationName,OpReturnVars,OperationBody)|T], NewEnv, Acc, OperationData) :-
    format('.', []),
    b_get_machine_operation_typed_parameters(OperationName, Parameters),
    b_get_machine_operation_any_parameters(OperationName, AnyParameters),
    append(Parameters, AnyParameters, TNParameters),
    findall(NP, ( member(P, TNParameters),
    add_texpr_infos(P, [parameter], NP)
    ), NParameters),
    get_ground_truth_from_operation_body(OpReturnVars, OperationBody, NParameters, GroundTruth),
    !,
    get_number_random_list(AR, [], RandomA),
    length(NParameters, ParameterAmount),
    select((global_ground_truth_params,TGlobalGt), GroundTruth, RestGt),
    extend_ground_truth(parameter, ParameterAmount, TGlobalGt, GlobalGt),
    enumerate_ids([], 'P', NParameters, EParameters),
    Env = [NormalizedStrings,NormalizedSets,NormalizedIds,NOperationNames],
    append(EParameters, NormalizedIds, NNormalizedIds),
    Env1 = [NormalizedStrings,NormalizedSets,NNormalizedIds,NOperationNames],
    normalize_ids_in_ground_truth(NormalizedIds, [(global_ground_truth_params,GlobalGt)|RestGt], NGroundTruth),
    generate_data_from_operation_substitutions_aux(Env1, 0, RandomA, AbsoluteFilePath-MachineName, OperationName, Env2, [], Data),
    generate_data_from_operation_substitutions(Env2, AR, AbsoluteFilePath-MachineName, T, NewEnv, [(Data,OperationName,NGroundTruth)|Acc], OperationData),
    !.

generate_data_from_operation_substitutions(Env, AR, AbsoluteFilePath-MachineName, [(_,_,_)|T], NewEnv, Acc, OperationData) :-
    generate_data_from_operation_substitutions(Env, AR, AbsoluteFilePath-MachineName, T, NewEnv, Acc, OperationData),
    !.

generate_data_from_operation_substitutions_aux(Env, _, [], _, _, Env, Acc, Acc) :-
    !.

generate_data_from_operation_substitutions_aux(Env, _, [AmountOfExamples|AT], AbsoluteFilePath-MachineName, OperationName, NewEnv, Acc, Data) :-
    get_valid_inputs_for_operation(OperationName, AmountOfExamples, InputExamples),
    remove_dups(InputExamples, InputExamplesSet),
    InputExamplesSet \== [],
    InputExamplesSet \== [[]],
    get_outputs_for_input_examples(Env, InputExamples, OperationName, InputOutputTuples, Env1),
    !,
    length(InputOutputTuples, Size),
    min_amount_of_examples(MinExamples),
    ( ( InputOutputTuples \== [],
    Size >= MinExamples
    )
    ->
    AtLeastOneSolution = 1,
    NewAcc = [InputOutputTuples|Acc]
    ; % skip this operation if no output examples could be generated or no ground truth can be inferred
    AtLeastOneSolution = 0,
    NewAcc = Acc
    ),
    generate_data_from_operation_substitutions_aux(Env1, AtLeastOneSolution, AT, AbsoluteFilePath-MachineName, OperationName, NewEnv, NewAcc, Data).

generate_data_from_operation_substitutions_aux(Env, AtLeastOneSolution, _, AbsoluteFilePath-_, OperationName, Env, Acc, Data) :-
    AtLeastOneSolution == 0,
    !,
    format('~n~nWarning: No valid inputs could be generated for operation ~w in machine ~w.~n~n', [OperationName,AbsoluteFilePath]),
    Data = Acc.

generate_data_from_operation_substitutions_aux(Env, _, _, _, _, Env, Acc, Acc).

generate_operation_data_from_machine_path(AbsoluteFilePath, Augmentations, SolverTimeoutMs, OperationData) :-
    retractall(tools:id_counter(_)),
    load_machine_for_data_generation(AbsoluteFilePath),
    get_normalized_id_name_mapping(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames),
    get_all_operation_bodies_from_current_machine(OperationBodies),
    b_machine_name(MachineName),
    Env = [[],NormalizedSets,NormalizedIds,NOperationNames,NRecordFieldNames],
    set_preference(time_out, SolverTimeoutMs),
    generate_data_from_operation_substitutions(Env, Augmentations, AbsoluteFilePath-MachineName, OperationBodies, _NewEnv, [], OperationData).

get_all_operation_bodies_from_current_machine(NOperationBodies) :-
    findall((OperationName,OpReturnVars,OperationBody), b_get_machine_operation(OperationName, OpReturnVars, _, OperationBody), OperationBodies),
    skip_operations_using_records(OperationBodies, [], NOperationBodies).

get_number_random_list(0, Acc, Acc) :-
    !.

get_number_random_list(C, Acc, L) :-
    C1 is C-1,
    get_random_amount_of_examples(R),
    \+ member(R, Acc),
    !,
    get_number_random_list(C1, [R|Acc], L).

get_number_random_list(C, Acc, L) :-
    get_number_random_list(C, Acc, L).

get_outputs_for_input_examples(Env, InputExamples, OperationName, InputOutputTuples, NewEnv) :-
    b_get_machine_operation(OperationName, _, _, Substitution),
    get_precondition_from_operation_body(Substitution, Precondition),
    b_get_typed_invariant_from_machine(Invariant),
    conjunct_predicates([Invariant,Precondition], InvariantAndPrecondition),
    get_outputs_for_input_examples(Env, InvariantAndPrecondition, InputExamples, OperationName, [], InputOutputTuples, NewEnv).

get_outputs_for_input_examples(Env, _, [], _, Acc, Acc, Env).

get_outputs_for_input_examples(Env, InvariantAndPrecondition, [InputExampleAst|T], OperationName, Acc, InputOutputTuples, NewEnv) :-
    create_equality_nodes_from_example(InputExampleAst, EqualityNodes),
    conjunct_predicates(EqualityNodes, EqualityPredicate),
    before_after_predicate_for_operation(OperationName, BeforeAfterPredicate),
    conjunct_predicates([InvariantAndPrecondition,EqualityPredicate,BeforeAfterPredicate], BeforeAfterPredicateWithEqualities),
    solve_predicate(BeforeAfterPredicateWithEqualities, _, 1, [force_evaluation], Solution),
    Solution = solution(SolutionBindings),
    get_output_nodes_from_bindings(SolutionBindings, OutputExampleAst),
    \+ contains_symbolic_value(OutputExampleAst),
    maplist(create_example_triple(Env), InputExampleAst, InputExample),
    findall(Name, member((_,_,Name), InputExample), InputVarNames),
    b_machine_identifier_normalization:map_normalize_ids_in_b_ast(Env, OutputExampleAst, OutputExampleAst2, Env1),
    reduce_input_to_violating_vars(InputVarNames, OutputExampleAst2, OutputExampleAst3),
    findall(N, ( member(b(_,_,I), OutputExampleAst3),
    member(synthesis(machinevar,N), I)
    ), OutputVarNames),
    subtract(InputVarNames, OutputVarNames, Sub),
    ( Sub == []
    ->
    maplist(create_example_triple(Env1), OutputExampleAst3, OutputExample),
    get_outputs_for_input_examples(Env1, InvariantAndPrecondition, T, OperationName, [(InputExample,OutputExample)|Acc], InputOutputTuples, NewEnv)
    ; format("Error: Missing Output value for at least one variable. Probably a bug in normalize_ids_in_b_ast/4.~n", []),
    InputOutputTuples = Acc,
    NewEnv = Env1
).

get_outputs_for_input_examples(Env, _, _, _, Acc, Acc, Env).

get_random_amount_of_examples(AmountOfExamples) :-
    min_amount_of_examples(MinAmountOfExamples),
    max_amount_of_examples(TempMaxAmountOfExamples),
    MaxAmountOfExamples is TempMaxAmountOfExamples+1,
    random(MinAmountOfExamples, MaxAmountOfExamples, AmountOfExamples),
    !.

is_classical_b_machine_path(Path) :-
    atom(Path),
    atom_concat(_, Ext, Path),
    Ext == '.mch'.

load_b_or_eventb_machine(AbsoluteFilePath) :-
    is_classical_b_machine_path(AbsoluteFilePath),
    !,
    b_load_machine_from_file(AbsoluteFilePath),
    b_machine_precompile.

load_b_or_eventb_machine(AbsoluteFilePath) :-
    b_load_eventb_project(AbsoluteFilePath),
    b_machine_precompile.

load_machine_for_data_generation(AbsoluteFilePath) :-
    b_machine_reset,
    set_preference(try_kodkod_on_load, false),
    load_b_or_eventb_machine(AbsoluteFilePath),
    \+ current_machine_uses_records,
    b_get_machine_variables(MachineVars),
    MachineVars \== [].

max_amount_of_examples(16).

min_amount_of_examples(3).

normalize_ids_in_ground_truth(_, [], []).

normalize_ids_in_ground_truth(NormalizedIds, [(VarName,GroundTruth)|T], [(NVarName,GroundTruth)|NT]) :-
    member((b(identifier(VarName),_,_),b(identifier(NVarName),_,_)), NormalizedIds),
    !,
    normalize_ids_in_ground_truth(NormalizedIds, T, NT).

normalize_ids_in_ground_truth(NormalizedIds, [(VarName,GroundTruth)|T], [(VarName,GroundTruth)|NT]) :-
    !,
    normalize_ids_in_ground_truth(NormalizedIds, T, NT).

normalize_ids_in_ground_truth(_, [H|_], _) :-
    format("~nError: Ground truth accumulator ~w is not properly wrapped. Abort.~n~n", [H]),
    !,
    fail.

skip_operations_using_records([], Acc, Acc).

skip_operations_using_records([OperationBody|T], Acc, NOperationBodies) :-
    \+ contains_record_type(OperationBody),
    !,
    skip_operations_using_records(T, [OperationBody|Acc], NOperationBodies).

skip_operations_using_records([_|T], Acc, NOperationBodies) :-
    skip_operations_using_records(T, Acc, NOperationBodies).