from pyensembl import genome_for_reference_name

from tinytimer import benchmark

ensembl = genome_for_reference_name("GRCh38")
contigs = [str(i + 1) for i in range(22)] + ["X", "Y"]


def make_repeat_lookup_fn(lookup_fn, n_positions):
    """
    Make a thunk which calls the lookup_fn at a number of loci
    for each human chromosome (excluding MT).
    """

    def repeat_lookup_fn():
        for contig in contigs:
            for position in [10**6 + i * 10**6 for i in range(n_positions)]:
                lookup_fn(contig, position)

    return repeat_lookup_fn


def run_benchmark(lookup_fn, n_positions_per_contig=20, time_limit=60.0):
    """
    Take a lookup functions (such as EnsemblRelease.genes_at_locus) and
    time how long it takes across multiple loci.
    """
    repeat_lookup_fn = make_repeat_lookup_fn(lookup_fn, n_positions_per_contig)
    n_loci = n_positions_per_contig * len(contigs)
    name = lookup_fn.__name__
    average_time = benchmark(repeat_lookup_fn, name="%s for %d loci" % (name, n_loci))
    print("-- %s : %0.4fs" % (name, average_time))
    assert average_time < time_limit, "%s took too long for %s loci: %0.4fs" % (
        name,
        n_loci,
        average_time,
    )
    return average_time


def test_timing_genes_at_locus():
    run_benchmark(ensembl.genes_at_locus)


def test_timing_transcripts_at_locus():
    run_benchmark(ensembl.transcripts_at_locus)


def test_timing_exons_at_locus():
    run_benchmark(ensembl.exons_at_locus)


def test_timing_transcript_sequences_at_locus():
    def transcript_sequences_at_locus(contig, position):
        sequences = []
        for transcript in ensembl.transcripts_at_locus(contig, position):
            sequences.append(transcript.sequence)
        return sequences

    run_benchmark(transcript_sequences_at_locus)


def test_timing_transcript_coding_sequences_at_locus():
    def transcript_coding_sequences_at_locus(contig, position):
        sequences = []
        for transcript in ensembl.transcripts_at_locus(contig, position):
            if transcript.sequence and transcript.complete:
                sequences.append(transcript.coding_sequence)
        return sequences

    run_benchmark(transcript_coding_sequences_at_locus)


def run_all_benchmarks():
    import types

    # run all local test functions to see their timings printed
    global_variables = globals()
    for variable_name in global_variables:
        if "test_" in variable_name:
            f = global_variables[variable_name]
            if isinstance(f, types.FunctionType):
                f()


if __name__ == "__main__":
    run_all_benchmarks()
