from scout.parse.ensembl import (parse_ensembl_exons)

from scout.constants import CHROMOSOMES

def test_exons_file_37(exons_handle):
    """This test exists to make sure that the demo exons file is on correct format"""
    # GIVEN the path to an ensembl formated exons file
    
    # WHEN parsing the lines
    header = next(exons_handle).rstrip().split('\t')
    
    ## THEN assert that the information is on correct format
    for exon_line in exons_handle:
        exon_line = exon_line.rstrip().split('\t')
        exon_info = dict(zip(header,exon_line))
        assert exon_info['Chromosome/scaffold name'] in CHROMOSOMES
        assert exon_info['Exon stable ID'].startswith('ENSE')
        assert exon_info['Gene stable ID'].startswith('ENSG')
        assert exon_info['Transcript stable ID'].startswith('ENST')

        assert int(exon_info['Exon region start (bp)'])
        assert int(exon_info['Exon region end (bp)'])

def test_exons_file_38(exons_38_handle):
    """This test exists to make sure that the demo exons file is on correct format"""
    # GIVEN the path to an ensembl formated exons file
    
    # WHEN parsing the lines
    header = next(exons_38_handle).rstrip().split('\t')
    
    ## THEN assert that the information is on correct format
    for exon_line in exons_38_handle:
        exon_line = exon_line.rstrip().split('\t')
        exon_info = dict(zip(header,exon_line))
        assert exon_info['Chromosome/scaffold name'] in CHROMOSOMES
        assert exon_info['Exon stable ID'].startswith('ENSE')
        assert exon_info['Gene stable ID'].startswith('ENSG')
        assert exon_info['Transcript stable ID'].startswith('ENST')

        assert int(exon_info['Exon region start (bp)'])
        assert int(exon_info['Exon region end (bp)'])
    

def test_parse_ensembl_exons(exons_handle):
    """Test to parse a small dataframe line of ensembl transcript"""
    ## GIVEN a iterable with exon information from ensembl

    ## WHEN parsing the exons
    parsed_exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(parsed_exons)

    ## THEN assert the parsed transcript is as expected
    assert parsed_exon['chrom']
    assert parsed_exon['gene']
    assert parsed_exon['transcript']

def test_parse_ensembl_exons_iterable():
    """Test to parse all ensembl exons"""
    ## GIVEN an iterable with ensembl exon data
    exons_handle = [
        "Chromosome/scaffold name\tGene stable ID\tTranscript stable ID\tExon stable ID\tExon"\
        " region start (bp)\tExon region end (bp)\t5' UTR start\t5' UTR end\t3' UTR start\t3'"\
        " UTR end\tStrand\tExon rank in transcript",
        "1\tENSG00000176022\tENST00000379198\tENSE00001439793\t1167629\t1170421\t1167629\t1"\
        "167658\t1168649\t1170421\t1\t1"
    ]
    ## WHEN parsing the exons in that file
    exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(exons)

    ## THEN assert that the exon is correctly parsed
    assert parsed_exon['chrom'] == '1'
    assert parsed_exon['ens_exon_id'] == 'ENSE00001439793'
    assert parsed_exon['transcript'] == 'ENST00000379198'
    assert parsed_exon['gene'] == 'ENSG00000176022'
    assert parsed_exon['exon_chrom_start'] == 1167629
    assert parsed_exon['exon_chrom_end'] == 1170421
    assert parsed_exon['5_utr_start'] == 1167629
    assert parsed_exon['5_utr_end'] == 1167658
    assert parsed_exon['3_utr_start'] == 1168649
    assert parsed_exon['3_utr_end'] == 1170421
    assert parsed_exon['rank'] == 1
    assert parsed_exon['strand'] == 1
    ## THEN assert start is max(5_utr_end, exon_chrom_start) since strand is 1
    assert parsed_exon['start'] == 1167658
    ## THEN assert end is min(3_utr_start, exon_chrom_end) since strand is 1
    assert parsed_exon['end'] == 1168649

def test_parse_ensembl_exons_missing_five_utr_start():
    """Test to parse all ensembl exons"""
    ## GIVEN an iterable with ensembl exon data
    exons_handle = [
        "Chromosome/scaffold name\tGene stable ID\tTranscript stable ID\tExon stable ID\tExon"\
        " region start (bp)\tExon region end (bp)\t5' UTR start\t5' UTR end\t3' UTR start\t3'"\
        " UTR end\tStrand\tExon rank in transcript",
        "1\tENSG00000176022\tENST00000379198\tENSE00001439793\t1167629\t1170421\t\t1"\
        "167658\t1168649\t1170421\t1\t1"
    ]
    ## WHEN parsing the exons in that file
    exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(exons)

    ## THEN assert that the exon is correctly parsed
    assert parsed_exon['chrom'] == '1'
    assert parsed_exon['ens_exon_id'] == 'ENSE00001439793'
    assert parsed_exon['transcript'] == 'ENST00000379198'
    assert parsed_exon['gene'] == 'ENSG00000176022'
    assert parsed_exon['5_utr_start'] == None

def test_parse_ensembl_exons_missing_five_utr_end():
    """Test to parse all ensembl exons"""
    ## GIVEN an iterable with ensembl exon data
    exons_handle = [
        "Chromosome/scaffold name\tGene stable ID\tTranscript stable ID\tExon stable ID\tExon"\
        " region start (bp)\tExon region end (bp)\t5' UTR start\t5' UTR end\t3' UTR start\t3'"\
        " UTR end\tStrand\tExon rank in transcript",
        "1\tENSG00000176022\tENST00000379198\tENSE00001439793\t1167629\t1170421\t\t"\
        "\t1168649\t1170421\t1\t1"
    ]
    ## WHEN parsing the exons in that file
    exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(exons)

    ## THEN assert that the exon is correctly parsed
    assert parsed_exon['chrom'] == '1'
    assert parsed_exon['ens_exon_id'] == 'ENSE00001439793'
    assert parsed_exon['transcript'] == 'ENST00000379198'
    assert parsed_exon['gene'] == 'ENSG00000176022'
    assert parsed_exon['5_utr_start'] == None
    assert parsed_exon['5_utr_end'] == None

def test_parse_ensembl_exons_missing_three_utr_start():
    """Test to parse all ensembl exons"""
    ## GIVEN an iterable with ensembl exon data
    exons_handle = [
        "Chromosome/scaffold name\tGene stable ID\tTranscript stable ID\tExon stable ID\tExon"\
        " region start (bp)\tExon region end (bp)\t5' UTR start\t5' UTR end\t3' UTR start\t3'"\
        " UTR end\tStrand\tExon rank in transcript",
        "1\tENSG00000176022\tENST00000379198\tENSE00001439793\t1167629\t1170421\t\t"\
        "\t\t1170421\t1\t1"
    ]
    ## WHEN parsing the exons in that file
    exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(exons)

    ## THEN assert that the exon is correctly parsed
    assert parsed_exon['chrom'] == '1'
    assert parsed_exon['ens_exon_id'] == 'ENSE00001439793'
    assert parsed_exon['transcript'] == 'ENST00000379198'
    assert parsed_exon['gene'] == 'ENSG00000176022'
    assert parsed_exon['5_utr_start'] == None
    assert parsed_exon['5_utr_end'] == None
    assert parsed_exon['3_utr_start'] == None

def test_parse_ensembl_exons_missing_three_utr_end():
    """Test to parse all ensembl exons"""
    ## GIVEN an iterable with ensembl exon data
    exons_handle = [
        "Chromosome/scaffold name\tGene stable ID\tTranscript stable ID\tExon stable ID\tExon"\
        " region start (bp)\tExon region end (bp)\t5' UTR start\t5' UTR end\t3' UTR start\t3'"\
        " UTR end\tStrand\tExon rank in transcript",
        "1\tENSG00000176022\tENST00000379198\tENSE00001439793\t1167629\t1170421\t\t"\
        "\t\t\t1\t1"
    ]
    ## WHEN parsing the exons in that file
    exons = parse_ensembl_exons(exons_handle)
    parsed_exon = next(exons)

    ## THEN assert that the exon is correctly parsed
    assert parsed_exon['chrom'] == '1'
    assert parsed_exon['ens_exon_id'] == 'ENSE00001439793'
    assert parsed_exon['transcript'] == 'ENST00000379198'
    assert parsed_exon['gene'] == 'ENSG00000176022'
    assert parsed_exon['5_utr_start'] == None
    assert parsed_exon['5_utr_end'] == None
    assert parsed_exon['3_utr_start'] == None
    assert parsed_exon['3_utr_end'] == None
