/*************** Functions used in ProteinGTRFit.bf ******************/
/** set model definitions **/
protein_gtr.convergence_options = {{"LogL", "Assess REV fit convergence by comparing log likelihood scores"}, {"RMSE", "[Recommended] Assess REV fit convergence by comparing RMSE between fitted matrices."}};
protein_gtr.rate_variation_options = {{"Yes", "Use a four-category discrete gamma distribution when optimizing branch lengths."}, {"No", "Do not consider rate variation when optimizing branch lengths."}};

function protein_gtr.load_cached_options() {
    protein_gtr.convergence_type = (protein_gtr.analysis_results[utility.getGlobalValue("terms.json.options")])[utility.getGlobalValue("protein_gtr.options.convergence_type")];
    protein_gtr.tolerance = (protein_gtr.analysis_results[utility.getGlobalValue("terms.json.options")])[utility.getGlobalValue("protein_gtr.options.tolerance")];
    protein_gtr.baseline_model = (protein_gtr.analysis_results[utility.getGlobalValue("terms.json.options")])[utility.getGlobalValue("protein_gtr.options.baseline_model")];
    protein_gtr.use_rate_variation = (protein_gtr.analysis_results[utility.getGlobalValue("terms.json.options")])[utility.getGlobalValue("protein_gtr.options.rate_variation")];




    /*** Cache options can in theory be saved as numbers if options piped in, although not typically. Remap to strings. ****/
    if (Type(protein_gtr.convergence_type) == "Number"){
        protein_gtr.convergence_type = utility.SwapKeysAndValues(utility.MatrixToDict(protein_gtr.convergence_options))[protein_gtr.convergence_type*2];
    }
    if (Type(protein_gtr.baseline_model) == "Number"){
        protein_gtr.baseline_model = utility.SwapKeysAndValues(utility.MatrixToDict(models.protein.empirical_models))[protein_gtr.baseline_model*2];
   }
    if (Type(protein_gtr.protein_gtr.use_rate_variation) == "Number"){
        protein_gtr.use_rate_variation = utility.SwapKeysAndValues(utility.MatrixToDict(protein_gtr.rate_variation_options))[protein_gtr.use_rate_variation*2];
    }
    return 0;
}

function protein_gtr.save_options() {
    protein_gtr.analysis_results[utility.getGlobalValue("terms.json.options")] = {utility.getGlobalValue("protein_gtr.options.convergence_type"): protein_gtr.convergence_type,
                                                                                  utility.getGlobalValue("protein_gtr.options.tolerance"): protein_gtr.tolerance,
                                                                                  utility.getGlobalValue("protein_gtr.options.baseline_model"): protein_gtr.baseline_model,
                                                                                  utility.getGlobalValue("protein_gtr.options.rate_variation") : protein_gtr.use_rate_variation};

    protein_gtr.analysis_results[utility.getGlobalValue("terms.json.input")] = {utility.getGlobalValue("terms.json.file"): protein_gtr.listfile,
                                                                                "number of datasets": protein_gtr.file_list_count};

    io.WriteCacheToFile (^"protein_gtr.cache_file", ^"protein_gtr.analysis_results");
    return 0;
}


/* Model definitions, in particular for models with rate variation */
//------------------------------------------------------------------------------------------------------------------------
/**
 * @name models.protein.REV.ModelDescription.withGamma
 * @description Define REV model with four-category gamma rate variation
 */
lfunction models.protein.REV.ModelDescription.withGamma (options) {
	def = models.protein.REV.ModelDescription (options);
	def [utility.getGlobalValue("terms.model.rate_variation")] = rate_variation.types.Gamma.factory ({utility.getGlobalValue("terms.rate_variation.bins") : 4});
	return def;
};
/**
 * @name protein_gtr.REV.ModelDescription.withGamma
 * @description Define a REV model with Gamma rate variation
 */
function protein_gtr.REV.ModelDescription.withGamma (type) {
    result = models.protein.REV.ModelDescription.withGamma(type);
    if (Type (protein_gtr.shared_EFV) == "Matrix") {
        result [terms.model.frequency_estimator] = "protein_gtr.REV.ModelDescription.freqs";
    }
    return result;
}
/**
 * @name protein_gtr.REV.ModelDescription
 * @description Define a REV model with constant site rates
 */
function protein_gtr.REV.ModelDescription (type) {
    result = models.protein.REV.ModelDescription(type);
    if (Type (protein_gtr.shared_EFV) == "Matrix") {
        result [terms.model.frequency_estimator] = "protein_gtr.REV.ModelDescription.freqs";
    }
    return result;
}
/**
 * @name protein_gtr.REV.ModelDescription.freqs
 * @description Define REV model frequencies as empirical
 */
function protein_gtr.REV.ModelDescription.freqs (model, namespace, datafilter) {
    model[terms.efv_estimate] = protein_gtr.shared_EFV;
    model[terms.model.efv_estimate_name] = terms.frequencies.predefined;
    (model[terms.parameters])[terms.model.empirical] = 0;
    return model;
}


lfunction protein_gtr.plusF.ModelDescription(type) {
    models.protein.plusF.ModelDescription.model_definition = models.protein.empirical.ModelDescription(type);
    models.protein.plusF.ModelDescription.model_definition [utility.getGlobalValue("terms.model.empirical_rates")] = utility.getGlobalValue("protein_gtr.normalized_qij");
    models.protein.plusF.ModelDescription.model_definition [utility.getGlobalValue("terms.model.frequency_estimator")] = "protein_gtr.plusF.frequencies";
    models.protein.plusF.ModelDescription.model_definition [utility.getGlobalValue("terms.parameters")] = {utility.getGlobalValue("terms.global"): {}, utility.getGlobalValue("terms.local"): {}, utility.getGlobalValue("terms.model.empirical"): 19};
    models.protein.plusF.ModelDescription.model_definition [utility.getGlobalValue("terms.model.q_ij")] = "protein_gtr.plusF._GenerateRate";
    return models.protein.plusF.ModelDescription.model_definition;
}

lfunction protein_gtr.plusF.ModelDescription.withGamma (options) {
	def = protein_gtr.plusF.ModelDescription (options);
	def [utility.getGlobalValue("terms.model.rate_variation")] = rate_variation.types.Gamma.factory ({utility.getGlobalValue("terms.rate_variation.bins") : 4});
	return def;
};


lfunction protein_gtr.plusF._GenerateRate (from,to,namespace,modelType) {
    return models.protein.empirical._GenerateRate (utility.getGlobalValue("protein_gtr.normalized_qij") , from,to,namespace,modelType);
}

lfunction protein_gtr.plusF.frequencies (model, namespace, datafilter) {
    model[utility.getGlobalValue("terms.efv_estimate")] = utility.getGlobalValue("protein_gtr.empirical_frequencies");
    model[utility.getGlobalValue("terms.model.efv_estimate_name")] = utility.getGlobalValue("terms.frequencies.predefined");
    (model[utility.getGlobalValue("terms.parameters")])[utility.getGlobalValue("terms.model.empirical")] = 19;
}
//------------------------------------------------------------------------------------------------------------------------





/**
 * @name protein_gtr.fitBaselineToFile
 * @description Fits an empirical amino acid model to dataset for branch length optimization
 * @param {String} filename - The name of the file containing the dataset to which the amino acid model will be fitted
 * @return the fitted MLE
 */
function protein_gtr.fitBaselineToFile (filename) {
    protein_gtr.file_info = alignments.ReadNucleotideDataSet ("protein_gtr.msa",
                                                              filename);
    protein_gtr.name_mapping = protein_gtr.file_info[utility.getGlobalValue("terms.data.name_mapping")];
    if (None == protein_gtr.name_mapping) { /** create a 1-1 mapping if nothing was done */
        protein_gtr.name_mapping = {};
        utility.ForEach (alignments.GetSequenceNames ("protein_gtr.msa"), "_value_", "`&protein_gtr.name_mapping`[_value_] = _value_");
    }
    utility.ToggleEnvVariable ("GLOBAL_FPRINTF_REDIRECT", "/dev/null");
    ExecuteCommands ('protein_gtr.partitions_and_trees = trees.LoadAnnotatedTreeTopology.match_partitions (protein_gtr.file_info[terms.data.partitions], protein_gtr.name_mapping)',
                     {"0" : "Y"});
    utility.ToggleEnvVariable ("GLOBAL_FPRINTF_REDIRECT", None);

    protein_gtr.partition_count      = Abs (protein_gtr.partitions_and_trees);
    io.CheckAssertion ("protein_gtr.partition_count==1", "This analysis can only handle a single partition");


    /*********** Store dataset information *************/
    protein_gtr.output_data_info = { utility.getGlobalValue("terms.original_name"): {},
                                     utility.getGlobalValue("terms.json.sequences"): protein_gtr.file_info[utility.getGlobalValue("terms.data.sequences")],
                                     utility.getGlobalValue("terms.json.sites"): protein_gtr.file_info[utility.getGlobalValue("terms.data.sites")]
                                   };



    utility.ForEach (utility.Keys (protein_gtr.name_mapping), "branch_name",
                             "utility.EnsureKey (protein_gtr.output_data_info[terms.original_name], branch_name)");

    utility.ForEach (utility.Keys (protein_gtr.name_mapping), "branch_name",
                             "(protein_gtr.output_data_info[terms.original_name])[branch_name] = protein_gtr.name_mapping[branch_name]");


    /*****************************************************/

    protein_gtr.filter_specification = alignments.DefineFiltersForPartitions (protein_gtr.partitions_and_trees,
                                                                            "protein_gtr.msa" ,
                                                                            "protein_gtr.filter.",
                                                                            protein_gtr.file_info);


    protein_gtr.full_trees = utility.Map (protein_gtr.partitions_and_trees, "_value_", '_value_[terms.data.tree]');
    protein_gtr.full_data_filter = utility.Map (protein_gtr.filter_specification, "_value_", "_value_[terms.data.name]");



    /******************** NORMALIZE MATRIX WITH DATA FREQUENCIES HERE *************************/
    utility.ToggleEnvVariable("COUNT_GAPS_IN_FREQUENCIES", 0);
    protein_gtr.empirical_frequencies = frequencies._aux.empirical.collect_data(protein_gtr.full_data_filter, 1, 1, 1);
    utility.ToggleEnvVariable("COUNT_GAPS_IN_FREQUENCIES", None);

    protein_gtr.normalized_qij = plusF_helper.BuildCustomNormalizedQ(protein_gtr.empirical_frequencies, protein_gtr.baseline_Rij, models.protein.alphabet);

    protein_gtr.baseline_mle = estimators.FitSingleModel_Ext(protein_gtr.full_data_filter,
                                                        protein_gtr.full_trees,
                                                        protein_gtr.final_baseline_model,
                                                        None,
                                                        None);

    protein_gtr.baseline_mle - terms.global; // delete empty key
    return protein_gtr.baseline_mle;
}



/**
 * @name protein_gtr.fitGTRtoFileList
 * @description Fits an average GTR amino acid model to a set of datasets
 * @param {Dict} current_results - the current analysis results, used for accurate file indexing
 * @param {Dict} previous_values - the fitted REV model from previous fit iteration, to use as initial values for this fit if they exist. If this is the first fit, then this is simply 0.
 * @param {String} phase - the key for this fit iteration
 * @param {Bool} final - True if this is the final tuning after convergence achieved, False if another fit iteration.
 * @return the fitted MLE
 */
function protein_gtr.fitGTRtoFileList (current_results, previous_values, phase, final) {

    io.ReportProgressMessageMD ("Protein GTR Fitter", phase,
                                "Fitting the REV model using constrained branch lengths proportions (" + phase + ")");

    file_list = utility.Keys (current_results);
    file_count = utility.Array1D (file_list);

    partition_info = {};
    filter_info    = {};
    trees = {};
    initial_values = {terms.global : {}, terms.branch_length : {}};
    proportional_scalers = {};
    index_to_file_name   = {};



    for (file_index = 0; file_index < file_count; file_index += 1) {
        file_path = file_list [file_index];
        index_to_file_name [file_path] = file_index;
        dataset_name = "protein_gtr.msa." + file_index;
        partition_info [file_path] = alignments.ReadNucleotideDataSet (dataset_name, file_path);
        partition_specification = { "0" : {terms.data.name : "all", terms.data.filter_string : "", term.data.tree : ((current_results[file_path])[terms.fit.trees])[0]}};

        filter_info [file_index] = (alignments.DefineFiltersForPartitions (partition_specification,
                                                                            dataset_name ,
                                                                            dataset_name,
                                                                            partition_info [file_path]))[0];
        trees [file_index] = {terms.trees.newick :  ((current_results[file_path])[terms.fit.trees])[0]};
        (initial_values[terms.branch_length])[file_index] = ((current_results[file_path])[terms.branch_length])[0];
        if (!final) {
            scaler = "protein_gtr.gtr_scaler_" + file_index;
            parameters.DeclareGlobalWithRanges (scaler, 1, 0, 1000);
            proportional_scalers[file_index] = scaler;
        }
    }


    utility.ToggleEnvVariable ("VERBOSITY_LEVEL", 1);
    utility.ToggleEnvVariable ("AUTO_PARALLELIZE_OPTIMIZE", 1);
    utility.ToggleEnvVariable ("OPTIMIZATION_METHOD", 0);

    if (! final) {
        // Set initial values to the previous fit
        if (utility.Has (previous_values, terms.global, "AssociativeList")) {
            initial_values[terms.global] = previous_values[terms.global];
        }
        // Set initial values to the specified baseline model ****r_ij****
        else {
            for (l1 = 0; l1 < 20; l1 += 1) {
                for (l2 = l1 + 1; l2 < 20; l2 += 1) {
                    //(initial_values[terms.global]) [terms.aminoacidRate (models.protein.alphabet[l1],models.protein.alphabet[l2])] = {terms.fit.MLE : 1}; // set all to 1
                    //(initial_values[terms.global]) [terms.aminoacidRate (models.protein.alphabet[l1],models.protein.alphabet[l2])] = {terms.fit.MLE : Random (0.5,2)}; // set all to random number
                    (initial_values[terms.global]) [terms.aminoacidRate (models.protein.alphabet[l1],models.protein.alphabet[l2])] = {terms.fit.MLE : protein_gtr.baseline_Rij[l1][l2]};
                }
            }
        }
        // fit the model
        protein_gtr.rev.mle = estimators.FitSingleModel_Ext (
                                            utility.Map (filter_info, "_value_", "_value_[terms.data.name]"),
                                            trees,
                                            protein_gtr.rev_model_branch_lengths,
                                            initial_values,
                                            {terms.run_options.proportional_branch_length_scaler : proportional_scalers}
                                       );


        // save the model fit to be associated with its filename index
        protein_gtr.rev.mle [protein_gtr.filename_to_index] = index_to_file_name;
    }
    // FINAL TUNING
    else {
        protein_gtr.rev.mle = estimators.FitSingleModel_Ext (
                                        utility.Map (filter_info, "_value_", "_value_[terms.data.name]"),
                                        trees,
                                        protein_gtr.rev_model_branch_lengths,
                                        previous_values,
                                        None
                                   );

    }



    // Save the rev.mle into the analysis_results, and cache it.
    (^"protein_gtr.analysis_results")[phase] = protein_gtr.rev.mle;
    io.WriteCacheToFile (^"protein_gtr.cache_file", ^"protein_gtr.analysis_results");

    console.log (""); // clear past the optimization progress line
    utility.SetEnvVariable ("VERBOSITY_LEVEL", 0);
    utility.ToggleEnvVariable ("AUTO_PARALLELIZE_OPTIMIZE", None);
    utility.ToggleEnvVariable ("OPTIMIZATION_METHOD", None);


    return protein_gtr.rev.mle;

}


/**
 * @name protein_gtr.fitGTRtoFile
 * @description Use a previously-fitted average REV amino acid model to a file, specifically for branch length optimization under this model
 * @param {String} filename - the filename of the dataset to be fitted
 * @param {Dict} rates - the rates for the GTR model used in fitting
 * @param {Dict} branch_lengths - the current branch length values for this dataset
 * @return the fitted MLE
 */
function protein_gtr.fitGTRtoFile (filename, rates, branch_lengths) {

    protein_gtr.file_info = alignments.ReadNucleotideDataSet ("protein_gtr.msa",
                                                              filename);

    utility.ToggleEnvVariable ("GLOBAL_FPRINTF_REDIRECT", "/dev/null");

    ExecuteCommands ('protein_gtr.partitions_and_trees = trees.LoadAnnotatedTreeTopology.match_partitions (protein_gtr.file_info [utility.getGlobalValue("terms.data.partitions")], protein_gtr.name_mapping)',
                     {"0" : "Y"});

    utility.ToggleEnvVariable ("GLOBAL_FPRINTF_REDIRECT", None);
    protein_gtr.filter_specification = alignments.DefineFiltersForPartitions (protein_gtr.partitions_and_trees,
                                                                            "protein_gtr.msa" ,
                                                                            "protein_gtr.filter.",
                                                                            protein_gtr.file_info);

    protein_gtr.rev_file_mle = {terms.global : {}};

    for (l1 = 0; l1 < 20; l1 += 1) {
        for (l2 = l1 + 1; l2 < 20; l2 += 1) {
            rate_term = terms.aminoacidRate (models.protein.alphabet[l1],models.protein.alphabet[l2]);
            (protein_gtr.rev_file_mle[terms.global]) [rate_term] =
                {terms.fit.MLE : (rates[rate_term])[terms.fit.MLE] , terms.fix : TRUE};
        }
    }

    protein_gtr.rev_file_mle [terms.branch_length] = { "0" : branch_lengths };

    utility.SetEnvVariable ("VERBOSITY_LEVEL", 1);
    protein_gtr.rev_file_mle = estimators.FitSingleModel_Ext (
                                        utility.Map (protein_gtr.filter_specification, "_value_", "_value_[terms.data.name]"), // value => value['name']
                                        utility.Map (protein_gtr.partitions_and_trees, "_value_", "_value_[terms.data.tree]"), // value => value['tree']
                                        protein_gtr.rev_model_branch_lengths, // pre-defined with or without gamma.
                                        protein_gtr.rev_file_mle,
                                        None
                                   );
   utility.SetEnvVariable ("VERBOSITY_LEVEL", 0);
   console.log (""); // clear past the optimization progress line

    protein_gtr.rev_file_mle - terms.global; // delete redundant keys

    return protein_gtr.rev_file_mle;

}





/**
 * @name protein_gtr.run_gtr_iteration_branch_lengths
 * @description Optimizes branch lengths for all datasets using the REV model fitted in the current iteration
 * @return Dictionary containing summed LogL values from branch length optimizations and the phase index for this iteration
 */
function protein_gtr.run_gtr_iteration_branch_lengths () {
    protein_gtr.fit_phase += 1;

    protein_gtr.queue = mpi.CreateQueue ({  "Headers"   : utility.GetListOfLoadedModules ("libv3/") ,
                                            "Functions" :
                                            {
                                                {"protein_gtr.REV.ModelDescription",
                                                 "protein_gtr.REV.ModelDescription.withGamma",
                                                 "protein_gtr.REV.ModelDescription.freqs"
                                                }
                                            },
                                            "Variables" : {{
                                                "protein_gtr.shared_EFV",
                                                "protein_gtr.final_baseline_model",
                                                "protein_gtr.rev_model_branch_lengths",
                                                "protein_gtr.baseline_Rij",
                                                "protein_gtr.filename_to_index"
                                            }}
                                         });

    protein_gtr.phase_key = "Rev-local-Phase-" + protein_gtr.fit_phase;

    io.ReportProgressMessageMD ("Protein GTR Fitter", protein_gtr.phase_key, "Retuning branch lengths (" + protein_gtr.phase_key + ")");

    for (file_index = 0; file_index < protein_gtr.file_list_count; file_index += 1) {

       if (utility.Has (protein_gtr.analysis_results, {{ protein_gtr.file_list[file_index],  protein_gtr.phase_key}}, None)) {
            logL = ((protein_gtr.analysis_results[protein_gtr.file_list[file_index]])[protein_gtr.phase_key])[terms.fit.log_likelihood];

            io.ReportProgressMessageMD ("Protein GTR Fitter", protein_gtr.phase_key,
                                        "Loaded cached results for '" + protein_gtr.file_list[file_index] + "' " + (file_index+1) + "/" + protein_gtr.file_list_count + ". Log(L) = " + logL);

        } else {
            io.ReportProgressMessageMD ("Protein GTR Fitter", protein_gtr.phase_key,
                                        "Dispatching file '" + protein_gtr.file_list[file_index] + "' " + (file_index+1) + "/" + protein_gtr.file_list_count);

            mpi.QueueJob (protein_gtr.queue, "protein_gtr.fitGTRtoFile", {"0" : protein_gtr.file_list[file_index],
                                                                          "1" : protein_gtr.current_gtr_fit[terms.global],
                                                                          "2" : (protein_gtr.current_gtr_fit[terms.branch_length])[(protein_gtr.current_gtr_fit[protein_gtr.filename_to_index])[protein_gtr.file_list[file_index]]]},
                                                                "protein_gtr.handle_gtr_callback");
        }
    }
    mpi.QueueComplete (protein_gtr.queue);

    protein_gtr.run_gtr_iteration_branch_lengths.logL = math.Sum (utility.Map (utility.Filter (protein_gtr.analysis_results, "_value_", "_value_/protein_gtr.phase_key"), "_value_", "(_value_[protein_gtr.phase_key])[terms.fit.log_likelihood]"));

    io.ReportProgressMessageMD ("Protein GTR Fitter", protein_gtr.phase_key,
                            "Overall Log(L) = " + protein_gtr.run_gtr_iteration_branch_lengths.logL);

    return { protein_gtr.logl : protein_gtr.run_gtr_iteration_branch_lengths.logL , protein_gtr.phase : protein_gtr.phase_key};

}

/**
 * @name protein_gtr.handle_gtr_callback
 * @description Handle MPI callback after fitting a REV model and save fit information to cache file
 */
function protein_gtr.handle_gtr_callback (node, result, arguments) {

    utility.EnsureKey (^"protein_gtr.analysis_results", arguments[0]);

    ((^"protein_gtr.analysis_results")[arguments[0]])[utility.getGlobalValue("protein_gtr.phase_key")] = result;

    io.ReportProgressMessageMD ("Protein GTR Fitter", "* Rev-local-Phase-" + ^"protein_gtr.fit_phase",
                                "Received file '" + arguments[0] + "' from node " + node + ". LogL = " + result[terms.fit.log_likelihood]);

    io.WriteCacheToFile (^"protein_gtr.cache_file", ^"protein_gtr.analysis_results");

}


/**
 * @name protein_gtr.handle_baseline_callback
 * @param node - node name which processed the given data
 * @param {Dict} result - Dictionary of fitted information for given data
 * @param {Dict} arguments - Dictionary with single key:value :: 0:datafile name
 * @description Handle MPI callback after fitting a baseline amino acid model (for initial branch length optimization) and save fit information to cache file
 */
function protein_gtr.handle_baseline_callback (node, result, arguments) {

    utility.EnsureKey (^"protein_gtr.analysis_results", arguments[0]);

    ((^"protein_gtr.analysis_results")[arguments[0]])[utility.getGlobalValue("protein_gtr.phase_key")] = result;


    io.ReportProgressMessageMD ("Protein GTR Fitter", "Initial branch length fit",
                                "Received file '" + arguments[0] + "' from node " + node + ". LogL = " + result[terms.fit.log_likelihood]);



    io.WriteCacheToFile (^"protein_gtr.cache_file", ^"protein_gtr.analysis_results");
}
