#!/bin/env python

#######################################################################
# Copyright (C) 2019 Julian Dosch
#
# This file is part of FAS.
#
#  FAS is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  FAS is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with FAS.  If not, see <http://www.gnu.org/licenses/>.
#
#######################################################################


from greedyFAS.mainFAS.fasWeighting import w_weight_const_rescale


def sf_calc_score(path, protein, weights, search_features, query_features, seed_proteome, clan_dict, query_clans,
                  option):
    """Used to be main Scoring function, starts the functions for the scores and calculates the complete FAS score,
    used durin query linearization when priority mode is enabled (linearized query versus full seed)

    :param path: Path to score
    :param protein: protein id
    :param weights: feature weights
    :param query_features: all features to be linearized in the query protein
    :param search_features: all features to be linearized in the seed protein
    :param seed_proteome: dictionary that contains the feature architecture of all seed proteins
    :param clan_dict: dictionary that maps features to clans
    :param query_clans: (Pfam)-clans of the current query protein
    :param option: dictionary that contains the main option variables of FAS
    :return: score_ms, score_ps, score_cs, final_score
    """
    adjusted_weights = {}
    tmp_weight = {}
    if option["weight_const"]:
        adjusted_weights = w_weight_const_rescale(path, weights, search_features, False, option)
        tmp_weight = {}
        for i in adjusted_weights:
            tmp_weight[i] = weights[i]
            weights[i] = adjusted_weights[i]
    score_cs = round(sf_cs_score(path, clan_dict, query_clans, query_features), 4)
    tmp = sf_ms_score(path, protein, seed_proteome, query_features, option)
    score_ms = round(tmp[0], 4)
    score_ps = round(
        sf_ps_score(path, tmp[2], protein, query_features, seed_proteome, option), 4)
    final_score = ((score_ms * option["score_weights"][0]) + (score_cs * option["score_weights"][1])
                   + (score_ps * option["score_weights"][2]))
    if option["weight_const"]:
        for i in adjusted_weights:
            weights[i] = tmp_weight[i]
    return score_ms, score_ps, score_cs, final_score


def sf_entire_calc_score(path, query_path, weights, search_features, a_s_f, query_features, a_q_f, clan_dict, option):
    """Main Scoring function, starts the functions for the scores and calculates the complete FAS score

    :param path: Path to score
    :param query_path: Path to score (query)
    :param weights: feature weights
    :param query_features: all features to be linearized in the query protein
    :param search_features: all features to be linearized in the seed protein
    :param a_s_f: additional seed features [non linearized]
    :param a_q_f: additional query features [non linearized]
    :param clan_dict: dictionary that maps features to clans
    :param option: dictionary that contains the main option variables of FAS
    :return: score_ms, score_ps, score_cs, final_score, path_weight, common_feature, score_ls
    """
    adjusted_weights = {}
    tmp_weight = {}
    if option["weight_const"]:
        adjusted_weights = w_weight_const_rescale(path, weights, search_features, False, option)
        tmp_weight = {}
        for i in adjusted_weights:
            tmp_weight[i] = weights[i]
            weights[i] = adjusted_weights[i]
    score_cs = round(sf_entire_cs_score(path, query_path, query_features, clan_dict, search_features), 4)
    tmp = sf_entire_ms_score(path, query_path, search_features, a_s_f, query_features, a_q_f, weights, option)
    score_ms = round(tmp[0], 4)
    path_weight = tmp[3]
    common_feature = tmp[4]
    ps_tmp = sf_entire_ps_score(path, tmp[2], query_path, search_features, a_s_f, query_features, a_q_f, weights,
                                option)
    score_ps = round(ps_tmp[0], 4)
    score_ls = round(ps_tmp[1], 4)
    final_score = ((score_ms * option["score_weights"][0]) + (score_cs * option["score_weights"][1])
                   + (score_ps * option["score_weights"][2]))
    if option["weight_const"]:
        for i in adjusted_weights:
            weights[i] = tmp_weight[i]
    if len(path) == 0 and len(query_path) == 0 and option['empty_as_1']:
        return 1.0, 1.0, 1.0, 1.0, path_weight, common_feature, 1.0
    else:
        return score_ms, score_ps, score_cs, float(final_score), path_weight, common_feature, score_ls


def sf_cs_score(path, clan_dict, query_clans, features):
    """Calculates clan score

    :param path: Path to score
    :param clan_dict: dictionary that maps features to clans
    :param query_clans: Clans in the query
    :param features: features (seed or query) [dict]
    :return: score (PS)
    """

    counter = 0.0
    path_clans = {}
    score = 0.0

    # counting clans in path
    # path_clans: contains counts for clans in path
    for i in path:
        feature = features[i]

        if feature[0] in clan_dict:
            if clan_dict[feature[0]] in path_clans:
                path_clans[clan_dict[feature[0]]] += 1
            else:
                path_clans[clan_dict[feature[0]]] = 1
    for clan in path_clans:
        counter += 1.0
        if clan in query_clans:
            score += float(path_clans[clan] * query_clans[clan])\
                     / float(max(path_clans[clan], query_clans[clan]) * max(path_clans[clan], query_clans[clan]))
    if counter == 0:
        score = 0.0
    else:
        score = score / counter
    return float(score)


def sf_entire_cs_score(path, query_path, query_features, clan_dict, search_features):
    """calculates clan score

    :param query_path:
    :param path: Path to score
    :param query_features: all features to be linearized in the query protein
    :param search_features: all features to be linearized in the seed protein
    :param clan_dict: dictionary that maps features to clans
    :return: score (PS)
    """

    counter = 0.0
    s_clans = {}  # clans from path
    q_clans = {}  # clans from query_path

    score = 0.0
    for i in path:
        if i in search_features:
            feature = search_features[i]
            if feature[0] in clan_dict:
                if clan_dict[feature[0]] in s_clans:
                    s_clans[clan_dict[feature[0]]] += 1
                else:
                    s_clans[clan_dict[feature[0]]] = 1
    for i in query_path:
        if i in query_features:
            feature = query_features[i]
            if feature[0] in clan_dict:
                if clan_dict[feature[0]] in q_clans:
                    q_clans[clan_dict[feature[0]]] += 1
                else:
                    q_clans[clan_dict[feature[0]]] = 1

    for clan in s_clans:
        counter += 1.0
        if clan in q_clans:
            score += float(s_clans[clan] * q_clans[clan]) / float(
                max(s_clans[clan], q_clans[clan]) * max(s_clans[clan], q_clans[clan]))
    if counter == 0:
        score = 0.0
    else:
        score = score / counter
    return float(score)


def sf_ms_score(path, protein, proteome, features, option):
    """Calculates multiplicity score, only used for priority mode now

    :param path : Path to score
    :param protein : protein id
    :param proteome : seed_proteome is a dictionary that contains the feature architecture of all seed proteins
    :param features : features (seed or query) [dict]
    :param option : specifies the behavior of the MS calculation
    :return: final_score (MS), search_domains, scale
    """
    domains = {}
    scale = 0
    scores = []
    final_score = 0.0
    tools = option["input_linearized"] + option["input_normal"]
    for i in path:
        feature = features[i]
        if feature[0] in domains:
            domains[feature[0]] += 1
        else:
            domains[feature[0]] = 1
    for feature in domains:
        scale += 1
        p_score = 0.0
        for tool in tools:
            if feature in proteome[protein][tool]:
                e_feature = False
                try:
                    if proteome[protein][tool][feature]["evalue"] <= option["eFeature"]:
                        e_feature = True
                except TypeError:
                    e_feature = True
                if e_feature:
                    s_length = 0
                    for instance in proteome[protein][tool][feature]["instance"]:
                        e_instance = False
                        try:
                            if instance[2] <= option["eInstance"]:
                                s_length += 1
                        except TypeError:
                            s_length += 1
                    p_score = min(float(domains[feature] * s_length) / float(domains[feature] * domains[feature]), 1.0)
        scores.append((feature, p_score))
    if scale > 0:
        scale = 1.0 / float(scale)
    for score in scores:
        final_score += score[1] * scale
    return float(final_score), domains, scale


def sf_entire_ms_score(path, query_path, search_features, a_s_f, query_features, a_q_f, weights, option):
    """Calculates multiplicity score

    :param path: Path to score
    :param query_path: Path to score (query)
    :param weights: feature weights
    :param query_features: all features to be linearized in the query protein
    :param search_features: all features to be linearized in the seed protein
    :param a_s_f: additional seed features [non linearized]
    :param a_q_f: additional query features [non linearized]
    :param option: dictionary that contains the main option variables of FAS
    :return: final_score (MS), search_domains, scale, final_weight, common_feature
    """
    search_domains = {}
    query_domains = {}
    scale = 0
    scores = []
    final_score = 0.0
    final_weight = 0.0
    main_features_s = {}
    main_features_q = []
    common_feature = False

    for i in path:
        if i in search_features:
            feature = search_features[i]
            main_features_s[feature[0]] = True
        else:
            feature = a_s_f[i]

        if feature[0] in search_domains:
            search_domains[feature[0]] += 1
        else:
            search_domains[feature[0]] = 1

    for i in query_path:
        if i in query_features:
            feature = query_features[i]
            main_features_q.append(feature[0])
        else:
            feature = a_q_f[i]
        if feature[0] in query_domains:
            query_domains[feature[0]] += 1
        else:
            query_domains[feature[0]] = 1

    # check for common features
    for y in main_features_q:
        try:
            if main_features_s[y]:
                common_feature = True
        except KeyError:
            pass
    for feature in search_domains:
        if option["MS_uni"] == 0:
            scale += weights[feature]
        else:
            scale += 1
        if feature in query_domains:
            s_length = query_domains[feature]
            p_score = min(float(search_domains[feature] * s_length) /
                          float(search_domains[feature] * search_domains[feature]), 1.0)
            scores.append((feature, p_score))
        else:
            scores.append((feature, 0.0))
    if scale > 0:
        scale = 1.0 / float(scale)
    for score in scores:
        if option["MS_uni"] == 0:
            final_score += score[1] * scale * weights[score[0]]
            final_weight += weights[score[0]]
        else:
            final_score += score[1] * scale
    return float(final_score), search_domains, scale, final_weight, common_feature


def sf_ps_score(path, scale, protein, features, seed_proteome, option):
    """Calculates positional score

    :param path: Path to score
    :param scale: contains scaling factor for each weight or number of feature types for uniform weighting
    :param protein: protein id
    :param features: features (seed or query) [dict]
    :param seed_proteome: dictionary that contains the feature architecture of all seed proteins
    :return: final_score (PS)
    """
    count = {}
    final_score = 0.0
    scores = {}
    tools = option["input_linearized"] + option["input_normal"]
    best_match = 0.0
    for i in path:
        feature = features[i]
        for tool in tools:
            if feature[0] in seed_proteome[protein][tool]:
                e_feature = False
                try:
                    if seed_proteome[protein][tool][feature[0]]["evalue"] <= option["eFeature"]:
                        e_feature = True
                except TypeError:
                    e_feature = True
                if e_feature:
                    best_match = 0.0
                    if not feature[0] in scores:
                        scores[feature[0]] = 0.0
                        count[feature[0]] = 0
                    for instance in seed_proteome[protein][tool][feature[0]]["instance"]:
                        e_instance = False
                        try:
                            if instance[2] <= option["eInstance"]:
                                e_instance = True
                        except TypeError:
                            e_instance = True
                        if e_instance:
                            pos = (float(instance[0]) + float(instance[1])) / 2.0 / float(
                                seed_proteome[protein]["length"])
                            match = 1.0 - float(abs(feature[1]) - pos)
                            if best_match < match:
                                best_match = match
                scores[feature[0]] += best_match
                count[feature[0]] += 1

    for f_score in scores:
        final_score += scores[f_score] / count[f_score] * scale
    return float(final_score)


def sf_entire_ps_score(path, scale, query_path, search_features, a_s_f, query_features, a_q_f, weights, option):
    """Calculates positional score

    :param path: Path to score
    :param scale: contains scaling factor for each weight or number of feature types for uniform weighting
    :param query_path: Path to score (query)
    :param weights: feature weights
    :param query_features: all features to be linearized in the query protein
    :param search_features: all features to be linearized in the seed protein
    :param a_s_f: additional seed features [non linearized]
    :param a_q_f: additional query features [non linearized]
    :param option: dictionary that contains the main option variables of FAS
    :return: final_score (PS), final_ls_score
    """
    count = {}
    final_score = 0.0
    final_ls_score = 0.0
    scores = {}
    ls_scores = {}
    local_query_protein = {}
    ls = 0.0
    # get current features from query path
    for i in query_path:
        if i in query_features:
            feature = query_features[i]
        else:
            feature = a_q_f[i]
        if feature[0] in local_query_protein:
            length = feature[3] - feature[2] + 1.0
            local_query_protein[feature[0]].append((feature[1], length))
        else:
            local_query_protein[feature[0]] = [(feature[1], feature[3] - feature[2] + 1.0)]

    # compare features in path with features form query path
    for i in path:
        if i in search_features:
            feature = search_features[i]
        else:
            feature = a_s_f[i]
        if feature[0] in local_query_protein:
            best_match = 0.0
            if not feature[0] in scores:
                count[feature[0]] = 0
                scores[feature[0]] = 0.0
                ls_scores[feature[0]] = 0.0
            for instance in local_query_protein[feature[0]]:
                match = 1.0 - float(abs(feature[1] - instance[0]))
                if best_match < match:
                    best_match = match
                    if instance[1] <= (feature[3] - feature[2]):
                        ls = float(instance[1]) / float(feature[3] - feature[2] + 1.0)
                    else:
                        ls = float(feature[3] - feature[2] + 1.0) / float(instance[1])

            scores[feature[0]] += best_match
            ls_scores[feature[0]] += ls
            count[feature[0]] += 1

    for f_score in scores:
        if option["MS_uni"] == 0:
            final_score += scores[f_score] / count[f_score] * scale * weights[f_score]
            final_ls_score += ls_scores[f_score] / count[f_score] * scale * weights[f_score]
        else:
            final_score += scores[f_score] / count[f_score] * scale
            final_ls_score += ls_scores[f_score] / count[f_score] * scale

    return float(final_score), float(final_ls_score)

