# Copyright (c) 2015. Mount Sinai School of Medicine
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

import gtfparse


def _get_gtf_column(column_name, gtf_path, dictionary):
    """
    Helper function which returns a dictionary column or raises an ValueError
    abou the absence of that column in a GTF file.
    """
    values = dictionary.get(column_name)
    if not values:
        raise ValueError(
            "Missing '%s' in columns of %s, available: %s" % (
                column_name,
                gtf_path,
                list(dictionary.keys())))
    return values


def load_transcript_fpkm_dict_from_gtf(
        gtf_path,
        transcript_id_column_name="reference_id",
        fpkm_column_name="FPKM",
        feature_column_name="feature"):
    """
    Load a GTF file generated by StringTie which contains transcript-level
    quantification of abundance. Returns a dictionary mapping Ensembl
    IDs of transcripts to FPKM values.
    """
    columns = gtfparse.read_gtf_as_dict(
        gtf_path,
        column_converters={fpkm_column_name: float})
    transcript_ids = _get_gtf_column(transcript_id_column_name, gtf_path, columns)
    fpkm_values = _get_gtf_column(fpkm_column_name, gtf_path, columns)
    features = _get_gtf_column(feature_column_name, gtf_path, columns)
    logging.info("Loaded %d rows from %s" % (len(transcript_ids), gtf_path))
    logging.info("Found %s transcript entries" % sum(
        feature == "transcript" for feature in features))
    result = {
        transcript_id: float(fpkm)
        for (transcript_id, fpkm, feature)
        in zip(transcript_ids, fpkm_values, features)
        if transcript_id is not None
        and len(transcript_id) > 0
        and feature == "transcript"
    }
    logging.info("Keeping %d transcript rows with reference IDs" % (
        len(result),))
    return result
