RequireVersion("2.3");

LoadFunctionLibrary("libv3/UtilityFunctions.bf");
LoadFunctionLibrary("libv3/IOFunctions.bf");
LoadFunctionLibrary("libv3/convenience/regexp.bf");

LoadFunctionLibrary("libv3/all-terms.bf");

//------------------------------------------------------------------------------
lfunction selection.io.defineBranchSets(partition_info) {

    available_models = {};
    branch_set = {};

    tree_for_analysis = None;



    if (Abs (partition_info) == 1) { // allow more refined selections if there is only one tree
        tree_for_analysis = (partition_info[0])[utility.getGlobalValue("terms.data.tree")];
        utility.ForEach (tree_for_analysis[utility.getGlobalValue("terms.trees.model_map")], "_value_", "`&available_models`[_value_] += 1");
        list_models   = utility.sortStrings(utility.Keys   (available_models)); // get keys
        option_count  = Abs (available_models);

    } else {
        option_count = 0;
    }

    selectTheseForTesting = {
        option_count + 3, 2
    };

    selectTheseForTesting[0][0] = "All";
    selectTheseForTesting[0][1] = "Include all branches in the analysis";
    selectTheseForTesting[1][0] = "Internal";
    selectTheseForTesting[1][1] = "Include all internal branches in the analysis";
    selectTheseForTesting[2][0] = "Leaves";
    selectTheseForTesting[2][1] = "Include all leaf branches in the analysis";

    for (k = 0; k < option_count; k += 1) {
        if (list_models[k] != "") {
            selectTheseForTesting[3 + k][0] = list_models[k];
            selectTheseForTesting[3 + k][1] = "Set " + list_models[k] + " with " + available_models[list_models[k]] + " branches";
        } else {
            selectTheseForTesting[3 + k][0] = "Unlabeled branches";
            selectTheseForTesting[3 + k][1] = "Set of " + available_models[list_models[k]] + " unlabeled branches";
        }
    }
    
    ChoiceList(testSet, "Choose the set of branches to test for selection", 1, NO_SKIP, selectTheseForTesting);

    io.CheckAssertion ("`&testSet` >= 0", "User cancelled branch selection; analysis terminating");

    return_set = {};

    for (k = 0; k < Abs (partition_info); k+=1) {
        tree_configuration = {};
        tree_for_analysis = (partition_info[k])[utility.getGlobalValue("terms.data.tree")];
        if (testSet < 3) {
            if (testSet == 0) {
                tag = None;
            } else {
                if (testSet == 1) {
                    tag = utility.getGlobalValue("terms.tree_attributes.internal");
                } else {
                    tag = utility.getGlobalValue("terms.tree_attributes.leaf");
                }
            }
            utility.ForEachPair (tree_for_analysis[utility.getGlobalValue("terms.trees.partitioned")], "_key_", "_value_", '
            	if (`&tag` == _value_ || `&tag` == None) {
            		`&tree_configuration`[_key_] = utility.getGlobalValue("terms.tree_attributes.test");
            	}
            	else {
            		`&tree_configuration`[_key_] = utility.getGlobalValue("terms.tree_attributes.background");
            	}
            ');

        } else {
            tag = SELECTION_STRINGS;
            if (tag == "Unlabeled branches") {
                tag = "";
            }
            utility.ForEachPair (tree_for_analysis[utility.getGlobalValue("terms.trees.model_map")], "_key_", "_value_", "if (`&tag` == _value_ || `&tag` == None) {`&tree_configuration`[_key_] = utility.getGlobalValue('terms.tree_attributes.test');} else {`&tree_configuration`[_key_] = utility.getGlobalValue('terms.tree_attributes.background');}");
        }
        return_set + tree_configuration;
    }
    return return_set;

}

//------------------------------------------------------------------------------
lfunction selection.io.SelectAllBranches(partition_info) {
    return_set = {};

    for (k = 0; k < Abs (partition_info); k+=1) {
        tree_configuration = {};
        tree_for_analysis = (partition_info[k])[utility.getGlobalValue("terms.data.tree")];
        utility.ForEachPair (tree_for_analysis[utility.getGlobalValue("terms.trees.partitioned")], "_key_", "_value_",
        '
            `&tree_configuration`[_key_] = utility.getGlobalValue("terms.tree_attributes.test");
        ');
        return_set + tree_configuration;
    }
    return return_set;

}

//------------------------------------------------------------------------------

lfunction selection.io.startTimer(timers, key, sort_index) {
    timers[key] = {
        utility.getGlobalValue("terms.timers.timer"): Time(1),
        utility.getGlobalValue("terms.timers.order"): sort_index
    };

}

//------------------------------------------------------------------------------

lfunction selection.io.stopTimer(timers, key) {
    (timers[key])[utility.getGlobalValue("terms.timers.timer")] = Time(1) - (timers[key])[utility.getGlobalValue("terms.timers.timer")];
}

//------------------------------------------------------------------------------

lfunction selection.io.extract_global_MLE(fit, id) {
    return ((fit[utility.getGlobalValue("terms.global")])[id])[utility.getGlobalValue("terms.fit.MLE")];
}

//------------------------------------------------------------------------------

lfunction selection.io.report_distribution (distribution) {

    // TODO remove hardcoding
    settings = {"header" : TRUE, "column-widths": {
            "0": 35,
            "1": 15,
            "2": 35
            }};

    fprintf (stdout, "\n", io.FormatTableRow ({{"Rate", "Proportion, %", "Notes"}}, settings));
    rates = Rows (distribution);
    settings ["header"] = FALSE;
    for (i = 0; i < rates; i += 1) {
        row_matrix = {1,3};

        row_matrix [0] =  Format (distribution [i][0], 6, 3);
        row_matrix [1] =  Format (distribution [i][1] * 100, 6, 3);
        row_matrix [2] = "";

        if (distribution [i][1] < 1e-4) {
            row_matrix [2] = "Not supported by data";
        } else {
            if (i > 0) {
                if (distribution [i-1][0] == distribution [i][0] || parameters.NormalizeRatio (Abs(distribution [i-1][0] - distribution [i][0]), Max (distribution [i-1][0], distribution [i][0])) < 0.05) {
                    row_matrix [2] = "Collapsed rate class";
                }
            }
        }

        fprintf (stdout, io.FormatTableRow (row_matrix, settings));

    }
    fprintf (stdout, "\n");
    return row_matrix;
}

//------------------------------------------------------------------------------

lfunction selection.io.report_dnds (distribution) {

    // TODO remove hardcoding
    settings = {"header" : TRUE, "column-widths": {
            "0": 35,
            "1": 15,
            "2": 10,
            "3": 35
            }};

    fprintf (stdout, "\n", io.FormatTableRow ({{"Selection mode", "dN/dS", "Proportion, %", "Notes"}}, settings));
    rates = Rows (distribution);
    settings ["header"] = FALSE;
    for (i = 0; i < rates; i += 1) {
        row_matrix = {1,4};

        if (distribution [i][0] < 1) {
            row_matrix [0] = "Negative selection";
        } else {
            if (distribution [i][0] == 1) {
                row_matrix [0] = "Neutral evolution";
            } else {
               row_matrix [0]  = "Diversifying selection";
            }
        }

        row_matrix [1] = Format (distribution [i][0], 6, 3);
        row_matrix [2] =  Format (distribution [i][1] * 100, 6, 3);
        row_matrix [3] = "";

        if (distribution [i][1] < 1e-4) {
            row_matrix [3] = "Not supported by data";
        } else {
            if (i > 0) {
                if (distribution [i-1][0] == distribution [i][0] || parameters.NormalizeRatio (Abs(distribution [i-1][0] - distribution [i][0]), Max (distribution [i-1][0], distribution [i][0])) < 0.05) {
                    row_matrix [3] = "Collapsed rate class";
                }
            }
        }

        fprintf (stdout, io.FormatTableRow (row_matrix, settings));

    }
    fprintf (stdout, "\n");
    return row_matrix;
}


//------------------------------------------------------------------------------

lfunction selection.io.extract_global_MLE_re (fit, regexp) {
     matches = utility.Filter (utility.Keys(fit[utility.getGlobalValue("terms.global")]), "_tag_", "regexp.Find (_tag_, `&regexp`)");
    //TODO: a single 'MLE' in the return line cannot be replaced by terms.
     return utility.Map (matches, "_tag_", "{utility.getGlobalValue('terms.description'): _tag_, 'MLE' : ((`&fit`[utility.getGlobalValue('terms.global')])[_tag_])[utility.getGlobalValue('terms.fit.MLE')]}");
}

//------------------------------------------------------------------------------

lfunction selection.io.branch.length(branch_info) {
    return branch_info[utility.getGlobalValue("terms.fit.MLE")];
}

//------------------------------------------------------------------------------------------------------------------------

lfunction selection.io.getIC(logl, params, samples) {
    return -2 * logl + 2 * samples / (samples - params - 1) * params;
}

//------------------------------------------------------------------------------

function selection.io.report_fit (fit, xp, ss) {
    return "Log(L) = " +
            Format (fit[terms.fit.log_likelihood], 8, 2) +
            ", AIC-c = " +
            Format (selection.io.getIC(fit[terms.fit.log_likelihood], fit[terms.parameters] + xp, ss), 8, 2) +
            " (" + (xp+fit[terms.parameters]) + " estimated parameters)";
}

//------------------------------------------------------------------------------

function selection.io.report_lrt (lrt) {

    return "p-value = " +
            Format (lrt[terms.p_value], 8, 4) +
            ", LRT = " +
            Format (lrt[terms.LRT], 8, 2);
}


//------------------------------------------------------------------------------


lfunction selection.io.ensure_not_null(object) {
    if (Type (object) == "Undefined") {
        return {};
    }
    return object;
}
//------------------------------------------------------------------------------


function selection.io.json_store_lf(json, name, ll, df, ss, rate_distribution, display_order) {


    utility.EnsureKey(json, terms.json.fits);

    selection.io.json_store_key_value_pair(json, terms.json.fits, name, {
            terms.json.log_likelihood: ll,
            terms.json.parameters: df,
            terms.json.AICc: selection.io.getIC(ll, df, ss),
            terms.json.rate_distribution: selection.io.ensure_not_null(rate_distribution),
            terms.json.display_order: display_order
    });

}

//------------------------------------------------------------------------------


function selection.io.json_store_lf_withEFV (json, name, ll, df, ss, rate_distribution, frequencies, display_order) {

    utility.EnsureKey(json, terms.json.fits);

    selection.io.json_store_key_value_pair(json, terms.json.fits, name, {
            terms.json.log_likelihood: ll,
            terms.json.parameters: df,
            terms.json.AICc: selection.io.getIC(ll, df, ss),
            terms.json.frequencies: frequencies,
            terms.json.rate_distribution: selection.io.ensure_not_null(rate_distribution),
            terms.json.display_order: display_order
    });

}




//------------------------------------------------------------------------------


function selection.io.json_store_lf_spool (file, json, name, ll, df, ss, rate_distribution) {

    selection.io.json_store_lf(json, name, ll, df, ss, rate_distribution);
    io.SpoolJSON (json, file);

}

//------------------------------------------------------------------------------

lfunction selection.io.adjust_partition_string (string, sites) {
    contig = regexp.Split (string, "-");
    if (Abs (contig) == 2) {
        start = +contig[0];
        end   = +contig[1];


        if (start % 3) { // need to adjust the start position
            if (start % 3 == 2) { // move upstream
                start += 1;
            } else {
                start += (-1);
            }
        }
        if (end % 3 ! = 2) { // need to adjust the ending position
            if (end % 3 == 1) { // move up if possible
                end += 1;
            } else {
                end += (-1);
            }
            if (end >= sites) {
                end = sites$3*3-1;
            }
        }
        io.CheckAssertion ("`&start` < `&end`", "Could not adjust the partition '" + string + "' to fall on codon boundaries and remain non-empty");
        string = "" + start + "-" + end;
    }

    return string;
}

//------------------------------------------------------------------------------

function selection.io.json_store_key_value_pair(json, dict_key, key, value) {

    if (None == dict_key) {
        json[key] = value;
    } else {
        utility.EnsureKey(json, dict_key);
        (json[dict_key])[key] = value;
    }
}

//------------------------------------------------------------------------------

function selection.io.json_store_branch_attribute(json, attribute_name, attribute_type, display_order, partition, values) {
    utility.EnsureKey(json, terms.json.branch_attributes);
    utility.EnsureKey(json[terms.json.branch_attributes], partition);
    utility.EnsureKey(json[terms.json.branch_attributes], terms.json.attribute);
    utility.EnsureKey((json[terms.json.branch_attributes])[terms.json.attribute], attribute_name);
    ((json[terms.json.branch_attributes])[terms.json.attribute])[attribute_name] = {terms.json.attribute_type : attribute_type,
                                                                                            terms.json.display_order: display_order};

    utility.ForEach (utility.Keys (values), "selection.io.json_store_branch_attribute.branch_name",
                             "utility.EnsureKey ((json[terms.json.branch_attributes])[partition], selection.io.json_store_branch_attribute.branch_name)");

    utility.ForEach (utility.Keys (values), "selection.io.json_store_branch_attribute.branch_name",
                             "  if (values / selection.io.json_store_branch_attribute.branch_name) {
                                    (((json[terms.json.branch_attributes])[partition])[selection.io.json_store_branch_attribute.branch_name])[attribute_name] = values[selection.io.json_store_branch_attribute.branch_name];
                                }
                              ");
 }


//------------------------------------------------------------------------------------------------------------------------

function selection.io._aux.extract_branch_info.callback(key, value) {
    selection.io.extract_branch_info_result[key] = Call (callback, value);
}

function selection.io.extract_branch_info(branch_spec, callback) {
    selection.io.extract_branch_info_result = {};
    branch_spec["selection.io._aux.extract_branch_info.callback"][""];
    return selection.io.extract_branch_info_result;
}

