"""
Tests for methods which return collections of transcript IDs that aren't
converting from some type of name or ID.
"""
from __future__ import absolute_import

from pyensembl import genome_for_reference_name
from .common import eq_
from .common import run_multiple_genomes

grch38 = genome_for_reference_name("GRCh38")

# subset of transcript IDs for HLA-A
HLA_A_TRANSCRIPT_IDS = [
    "ENST00000396634",
    "ENST00000376809",
    "ENST00000376806",
    "ENST00000376802",
    "ENST00000496081",
    "ENST00000495183",
    "ENST00000461903",
    "ENST00000479320",
]


def test_transcript_ids_ensembl_grch38_hla_a():
    # chr6:29,945,884  is a position for HLA-A
    # based on:
    # http://useast.ensembl.org/Homo_sapiens/Gene/
    # Summary?db=core;g=ENSG00000206503;r=6:29941260-29945884
    transcript_ids = grch38.transcript_ids_at_locus(6, 29941260, 29945884)
    for transcript_id in HLA_A_TRANSCRIPT_IDS:
        assert transcript_id in transcript_ids, (
            "Transcript %s of HLA-A not found overlapping locus" % transcript_id
        )


KNOWN_TRANSCRIPT_IDS = HLA_A_TRANSCRIPT_IDS + [
    "ENST00000398417",  # transcript ID of SMAD4-001
    "ENST00000334701",  # transcript ID of HSP90AA1-001
    "ENST00000599837",  # transcript ID of CTAG1A-002
]


# TODO: add release 54 after transcript IDs for older GTFs are filled in
# See https://github.com/hammerlab/pyensembl/issues/20
@run_multiple_genomes(75, grch38.release)
def test_all_transcript_ids(genome):
    transcript_ids = set(genome.transcript_ids())
    for transcript_id in KNOWN_TRANSCRIPT_IDS:
        assert transcript_id in transcript_ids, "Missing transcript ID %s from %s" % (
            transcript_id,
            genome,
        )


def test_transcript_id_of_protein_id_CCR2():
    # Looked up on Oct 9 2021:
    # CCR2-203 ENST00000445132.3 maps to ENSP00000399285.2
    # Ensembl release 104, GRCh38.p13
    transcript_id = grch38.transcript_id_of_protein_id("ENSP00000399285")
    eq_("ENST00000445132", transcript_id)


def test_protein_ids_at_locus_grch38_hla_a():
    # Regression test for https://github.com/openvax/pyensembl/issues/286:
    # protein_ids_at_locus previously queried the "transcript" feature,
    # which stores an empty protein_id, so results were a list of empty
    # strings instead of real protein IDs.
    # chr6:29,942,555 falls inside the first CDS of HLA-A transcripts.
    protein_ids = set(grch38.protein_ids_at_locus(6, 29942555))
    assert protein_ids, "Expected non-empty protein IDs at HLA-A CDS locus"
    assert "" not in protein_ids, (
        "protein_ids_at_locus should not return empty strings: %s" % (protein_ids,)
    )
    assert "ENSP00000365998" in protein_ids, (
        "Expected HLA-A protein ENSP00000365998 at locus, got: %s" % (protein_ids,)
    )
