import unittest
from pygor3 import imgt
from pygor3.utils import *
from pygor3 import IgorRefGenome
from pygor3 import IgorModel

from pygor3 import get_IgorRefGenome_VDJ_from_IMGT
from pandas._testing import assert_frame_equal

import os
import tempfile
import time
import numpy as np
import pandas as pd

str_mock_VDJ_fln_genomicVs = \
""">TRBV1*01
GATACTGGAATTACCCAGACACCAAAATACCTGGTCACAGCAATGGGGAGTAAAAGGACA
ATGAAACGTGAGCATCTGGGACATGATTCTATGTATTGGTACAGACAGAAAGCTAAGAAA
TCCCTGGAGTTCATGTTTTACTACAACTGTAAGGAATTCATTGAAAACAAGACTGTGCCA
AATCACTTCACACCTGAATGCCCTGACAGCTCTCGCTTATACCTTCATGTGGTCGCACTG
CAGCAAGAAGACTCAGCTGCGTATCTCTGCACCAGCAGCCAAGA
>TRBV2*01
GAACCTGAAGTCACCCAGACTCCCAGCCATCAGGTCACACAGATGGGACAGGAAGTGATC
TTGCGCTGTGTCCCCATCTCTAATCACTTATACTTCTATTGGTACAGACAAATCTTGGGG
CAGAAAGTCGAGTTTCTGGTTTCCTTTTATAATAATGAAATCTCAGAGAAGTCTGAAATA
TTCGATGATCAATTCTCAGTTGAAAGGCCTGATGGATCAAATTTCACTCTGAAGATCCGG
TCCACAAAGCTGGAGGACTCAGCCATGTACTTCTGTGCCAGCAGTGAAGC
>TRBV2*02
GAACCTGAAGTCACCCAGACTCCCAGCCATCAGGTCACACAGATGGGACAGGAAGTGATC
TTGCACTGTGTCCCCATCTCTAATCACTTATACTTCTATTGGTACAGACAAATCTTGGGG
CAGAAAGTCGAGTTTCTGGTTTCCTTTTATAATAATGAAATCTCAGAGAAGTCTGAAATA
TTCGATGATCAATTCTCAGTTGAAAGGCCTGATGGATCAAATTTCACTCTGAAGATCCGG
TCCACAAAGCTGGAGGACTCAGCCATGTACTTCTGTGCCAGCAGT
>TRBV2*03
GAACCTGAAGTCACCCAGACTCCCAGCCATCAGGTCACACAGATGGGACAGGAAGTGATC
TTGCGCTGTGTCCCCATCTCTAATCACTTATACTTCTATTGGTACAGACAAATCTTGGGG
CAGAAAGTCGAGTTTCTGGTTTCCTTTTATAATAATGAAATCTCAGAGAAGTCTGAAATA
TTCGATGATCAATTCTCAGTTGAGAGGCCTGATGGATCAAATTTCACTCTGAAGATCCGG
TCCACAAAGCTGGAGGACTCAGCCATGTACTTCTGTGCCAGCAGTGAA
>TRBV3-1*01
GACACAGCTGTTTCCCAGACTCCAAAATACCTGGTCACACAGATGGGAAACGACAAGTCC
ATTAAATGTGAACAAAATCTGGGCCATGATACTATGTATTGGTATAAACAGGACTCTAAG
AAATTTCTGAAGATAATGTTTAGCTACAATAATAAGGAGCTCATTATAAATGAAACAGTT
CCAAATCGCTTCTCACCTAAATCTCCAGACAAAGCTCACTTAAATCTTCACATCAATTCC
CTGGAGCTTGGTGACTCTGCTGTGTATTTCTGTGCCAGCAGCCAAGA
>TRBV3-1*02
GACACAGCTGTTTCCCAGACTCCAAAATACCTGGTCACACAGATGGGAAACGACAAGTCC
ATTAAATGTGAACAAAATCTGGGCCATGATACTATGTATTGGTATAAACAGGACTCTAAG
AAATTTCTGAAGATAATGTTTAGCTACAATAACAAGGAGATCATTATAAATGAAACAGTT
CCAAATCGATTCTCACCTAAATCTCCAGACAAAGCTAAATTAAATCTTCACATCAATTCC
CTGGAGCTTGGTGACTCTGCTGTGTATTTCTGTGCCAGC
>TRBV3-2*01
GACACAGCCGTTTCCCAGACTCCAAAATACCTGGTCACACAGATGGGAAAAAAGGAGTCT
CTTAAATGAGAACAAAATCTGGGCCATAATGCTATGTATTGGTATAAACAGGACTCTAAG
AAATTTCTGAAGACAATGTTTATCTACAGTAACAAGGAGCCAATTTTAAATGAAACAGTT
CCAAATCGCTTCTCACCTGACTCTCCAGACAAAGCTCATTTAAATCTTCACATCAATTCC
CTGGAGCTTGGTGACTCTGCTGTGTATTTCTGTGCCAGCAGCCAAGA
>TRBV3-2*02
GACACAGCCGTTTCCCAGACTCCAAAATACCTGGTCACACAGATGGGAAAAAAGGAGTCT
CTTAAATGAGAACAAAATCTGGGCCATAATGCTATGTATTGGTATAAACAGGACTCTAAG
AAATTTCTGAAGACAATGTTTATCTACAGTAACAAGGAGCCAATTTTAAATGAAACAGTT
CCAAATCGCTTCTCACCTGACTCTCCAGACAAAGTTCATTTAAATCTTCACATCAATTCC
CTGGAGCTTGGTGACTCTGCTGTGTATTTCTGTGCCAGCAGCCAAGA
>TRBV3-2*03
GACACAGCCGTTTCCCAGACTCCAAAATACCTGGTCACACAGACGGGAAAAAAGGAGTCT
CTTAAATGAGAACAAAATCTGGGCCATAATGCTATGTATTGGTATAAACAGGACTCTAAG
AAATTTCTGAAGACAATGTTTATCTACAGTAACAAGGAGCCAATTTTAAATGAAACAGTT
CCAAATCGCTTCTCACCTGACTCTCCAGACAAAGTTCATTTAAATCTTCACATCAATTCC
CTGGAGCTTGGTGACTCTGCTGTGTATTTCTGTGCCAGCAGCCAAG
>TRBV4-1*02
CACCTGGTCATGGGAATGACAAATAAGAAGTCTTTGAAATGTGAACAACATATGGGGCAC
AGGGCAATGTATTGGTACAAGCAGAAAGCTAAGAAGCCACCGGAGCTCATGTTTGTCTAC
AGCTATGAGAAACTCTCTATAAATGAAAGTGTGCCAAGTCGCTTCTCACCTGAATGCCCC
AACAGCTCTCTCTTAAACCTTCACCTACACGCCCTGCAGCCAGAAGACTCAGCCCTGTAT
CTCTGCGCCAGCAGCCAAG
>TRBV4-2*01
GAAACGGGAGTTACGCAGACACCAAGACACCTGGTCATGGGAATGACAAATAAGAAGTCT
TTGAAATGTGAACAACATCTGGGGCATAACGCTATGTATTGGTACAAGCAAAGTGCTAAG
AAGCCACTGGAGCTCATGTTTGTCTACAACTTTAAAGAACAGACTGAAAACAACAGTGTG
CCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATTCCTTCACCTACACACC
CTGCAGCCAGAAGACTCGGCCCTGTATCTCTGTGCCAGCAGCCAAGA
>TRBV4-2*02
GAAACGGGAGTTACGCAGACACCAAGACACCTGGTCATGGGAATGACAAATAAGAAGTCT
TTGAAATGTGAACAACATCTGGGGCATAACGCTATGTATTGGTACAAGCAAAGTGCTAAG
AAGCCACTGGAGCTCATGTTTGTCTACAACTTTAAAGAACAGACTGAAAACAACAGTGTG
CCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATGCCTTCACCTACACACC
CTGCAGCCAGAAGACTCGGCCCTGTATCTCTGTGCCAGCACC
>TRBV4-3*01
GAAACGGGAGTTACGCAGACACCAAGACACCTGGTCATGGGAATGACAAATAAGAAGTCT
TTGAAATGTGAACAACATCTGGGTCATAACGCTATGTATTGGTACAAGCAAAGTGCTAAG
AAGCCACTGGAGCTCATGTTTGTCTACAGTCTTGAAGAACGGGTTGAAAACAACAGTGTG
CCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATTCCTTCACCTACACACC
CTGCAGCCAGAAGACTCGGCCCTGTATCTCTGCGCCAGCAGCCAAGA
>TRBV4-3*02
GAAACGGGAGTTACGCAGACACCAAGACACCTGGTCATGGGAATGACAAATAAGAAGTCT
TTGAAATGTGAACAACATCTGGGTCATAACGCTATGTATTGGTACAAGCAAAGTGCTAAG
AAGCCACTGGAGCTCATGTTTGTCTACAGTCTTGAAGAACGGGTTGAAAACAACAGTGTG
CCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATCCCTTCACCTACACACC
CTGCAGCCAGAAGACTCGGCCCTGTATCTCTGCGCCAGCAGC
>TRBV4-3*03
GAAACGGGAGTTACGCAGACACCAAGACACCTGGTCATGGGAATGACAAATAAGAAGTCT
TTGAAATGTGAACAACATCTGGGTCATAACGCTATGTATTGGTACAAGCAAAGTGCTAAG
AAGCCACTGGAGCTCATGTTTGTCTACAGTCTTGAAGAACGTGTTGAAAACAACAGTGTG
CCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATTCCTTCACCTACACACC
CTGCAGCCAGAAGACTCGGCCCTGTATCTCTGCGCCAGCAGC
>TRBV4-3*04
AAGAAGTCTTTGAAATGTGAACAACATCTGGGGCATAACGCTATGTATTGGTACAAGCAA
AGTGCTAAGAAGCCACTGGAGCTCATGTTTGTCTACAGTCTTGAAGAACGGGTTGAAAAC
AACAGTGTGCCAAGTCGCTTCTCACCTGAATGCCCCAACAGCTCTCACTTATTCCTTCAC
CTACACACCCTGCAGCCAGAAGACTCGGCCCTGTATCTCTGCGCCAGCAGC
>TRBV5-1*01
AAGGCTGGAGTCACTCAAACTCCAAGATATCTGATCAAAACGAGAGGACAGCAAGTGACA
CTGAGCTGCTCCCCTATCTCTGGGCATAGGAGTGTATCCTGGTACCAACAGACCCCAGGA
CAGGGCCTTCAGTTCCTCTTTGAATACTTCAGTGAGACACAGAGAAACAAAGGAAACTTC
CCTGGTCGATTCTCAGGGCGCCAGTTCTCTAACTCTCGCTCTGAGATGAATGTGAGCACC
TTGGAGCTGGGGGACTCGGCCCTTTATCTTTGCGCCAGCAGCTTGG
>TRBV5-1*02
AGGGCTGGGGTCACTCAAACTCCAAGACATCTGATCAAAACGAGAGGACAGCAAGTGACA
CTGGGCTGCTCCCCTATCTCTGGGCATAGGAGTGTATCCTGGTACCAACAGACCCTAGGA
CAGGGCCTTCAGTTCCTCTTTGAATACTTCAGTGAGACACAGAGAAACAAAGGAAACTTC
CTTGGTCGATTCTCAGGGCGCCAGTTCTCTAACTCTCGCTCTGAGATGAATGTGAGCACC
TTGGAGCTGGGGGACTCGGCCCTTTATCTTTGCGCCAGC
>TRBV5-2*01
GAGGCTGGAATCACCCAAGCTCCAAGACACCTGATCAAAACAAGAGACCAGCAAGTGACA
CTGAGATGCTCCCCTGCCTCTGGGCATAACTGTGTGTCCTGGTACCTACGAACTCCAAGT
CAGCCCCTCTAGTTATTGTTACAATATTGTAATAGGTTACAAAGAGCAAAAGGAAACTTG
CCTAATTGATTCTCAGCTCACCACGTCCATAACTATTACTGAGTCAAACACGGAGCTAGG
GGACTCAGCCCTGTATCTCTGTGCCAGCAACTTGATG
>TRBV5-3*01
GAGGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTCCTATCTCTGGGCACAGCAGTGTGTCCTGGTACCAACAGGCCCCGGGT
CAGGGGCCCCAGTTTATCTTTGAATATGCTAATGAGTTAAGGAGATCAGAAGGAAACTTC
CCTAATCGATTCTCAGGGCGCCAGTTCCATGACTGTTGCTCTGAGATGAATGTGAGTGCC
TTGGAGCTGGGGGACTCGGCCCTGTATCTCTGTGCCAGAAGCTTGG
>TRBV5-3*02
GAGGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTCCTATCTCTGGGCACAGCAGTGTGTCCTGGTACCAACAGGCCCCGGGT
CAGGGGCCCCAGTTTATCTTTGAATATGCTAATGAGTTAAGGAGATCAGAAGGAAACTTC
CCTAATCGATTCTCAGGGCGCCAGTTCCATGACTATTGCTCTGAGATGAATGTGAGTGCC
TTGGAGCTGGGGGACTCGGCCCTGTATCTCTGTGCCAGAAGCTTGG
>TRBV5-4*01
GAGACTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTTCTCAGTCTGGGCACAACACTGTGTCCTGGTACCAACAGGCCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATAGGGAGGAAGAGAATGGCAGAGGAAACTTC
CCTCCTAGATTCTCAGGTCTCCAGTTCCCTAATTATAGCTCTGAGCTGAATGTGAACGCC
TTGGAGCTGGACGACTCGGCCCTGTATCTCTGTGCCAGCAGCTTGG
>TRBV5-4*02
GAGACTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTTCTCAGTCTGGGCACAACACTGTGTCCTGGTACCAACAGGCCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATAGGGAGGAAGAGAATGGCAGAGGAAACTTC
CCTCCTAGATTCTCAGGTCTCCAGTTCCCTAATTATAACTCTGAGCTGAATGTGAACGCC
TTGGAGCTGGACGACTCGGCCCTGTATCTCTGTGCCAGCAGC
>TRBV5-4*03
CAGCAAGTGACACTGAGATGCTCTTCTCAGTCTGGGCACAACACTGTGTCCTGGTACCAA
CAGGCCCTGGGTCAGGGGCCCCAGTTTATCTTTCAGTATTATAGGGAGGAAGAGAATGGC
AGAGGAAACTTCCCTCCTAGATTCTCAGGTCTCCAGTTCCCTAATTATAGCTCTGAGCTG
AATGTGAACGCCTTGGAGCTGGACGACTCGGCCCTGTATCTCTGTGCCAGCAGC
>TRBV5-4*04
ACTGTGTCCTGGTACCAACAGGCCCTGGGTCAGGGGCCCCAGTTTATCTTTCAGTATTAT
AGGGAGGAAGAGAATGGCAGAGGAAACTCCCCTCCTAGATTCTCAGGTCTCCAGTTCCCT
AATTATAGCTCTGAGCTGAATGTGAACGCCTTGGAGCTGGACGACTCGGCCCTGTATCTC
TGTGCCAGCAGC
>TRBV5-5*01
GACGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTCCTATCTCTGGGCACAAGAGTGTGTCCTGGTACCAACAGGTCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATGAGAAAGAAGAGAGAGGAAGAGGAAACTTC
CCTGATCGATTCTCAGCTCGCCAGTTCCCTAACTATAGCTCTGAGCTGAATGTGAACGCC
TTGTTGCTGGGGGACTCGGCCCTGTATCTCTGTGCCAGCAGCTTGG
>TRBV5-5*02
GACGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCACGTGACT
CTGAGATGCTCTCCTATCTCTGGGCACAAGAGTGTGTCCTGGTACCAACAGGTCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATGAGAAAGAAGAGAGAGGAAGAGGAAACTTC
CCTGATCGATTCTCAGCTCGCCAGTTCCCTAACTATAGCTCTGAGCTGAATGTGAACGCC
TTGTTGCTGGGGGACTCGGCCCTGTATCTCTGTGCCAGCAGC
>TRBV5-5*03
GACGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTCCTATCTCTGAGCACAAGAGTGTGTCCTGGTACCAACAGGTCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATGAGAAAGAAGAGAGAGGAAGAGGAAACTTC
CCTGATCGATTCTCAGCTCGCCAGTTCCCTAACTATAGCTCTGAGCTGAATGTGAACGCC
TTGTTGCTGGGGGACTCGGCCCTGTATCTCTGTGCCAGCAGC
>TRBV5-6*01
GACGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGTGACT
CTGAGATGCTCTCCTAAGTCTGGGCATGACACTGTGTCCTGGTACCAACAGGCCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATGAGGAGGAAGAGAGACAGAGAGGCAACTTC
CCTGATCGATTCTCAGGTCACCAGTTCCCTAACTATAGCTCTGAGCTGAATGTGAACGCC
TTGTTGCTGGGGGACTCGGCCCTCTATCTCTGTGCCAGCAGCTTGG
>TRBV5-7*01
GACGCTGGAGTCACCCAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCACGTGACT
CTGAGATGCTCTCCTATCTCTGGGCACACCAGTGTGTCCTCGTACCAACAGGCCCTGGGT
CAGGGGCCCCAGTTTATCTTTCAGTATTATGAGAAAGAAGAGAGAGGAAGAGGAAACTTC
CCTGATCAATTCTCAGGTCACCAGTTCCCTAACTATAGCTCTGAGCTGAATGTGAACGCC
TTGTTGCTAGGGGACTCGGCCCTCTATCTCTGTGCCAGCAGCTTGG
>TRBV5-8*01
GAGGCTGGAGTCACACAAAGTCCCACACACCTGATCAAAACGAGAGGACAGCAAGCGACT
CTGAGATGCTCTCCTATCTCTGGGCACACCAGTGTGTACTGGTACCAACAGGCCCTGGGT
CTGGGCCTCCAGTTCCTCCTTTGGTATGACGAGGGTGAAGAGAGAAACAGAGGAAACTTC
CCTCCTAGATTTTCAGGTCGCCAGTTCCCTAATTATAGCTCTGAGCTGAATGTGAACGCC
TTGGAGCTGGAGGACTCGGCCCTGTATCTCTGTGCCAGCAGCTTGG
>TRBV5-8*02
AGGACAGCAAGCGACTCTGAGATGCTCTCCTATCTCTGGGCACACCAGTGTGTACTGGTA
CCAACAGGCCCTGGGTCTGGGCCTCCAGCTCCTCCTTTGGTATGACGAGGGTGAAGAGAG
AAACAGAGGAAACTTCCCTCCTAGATTTTCAGGTCGCCAGTTCCCTAATTATAGCTCTGA
GCTGAATGTGAACGCCTTGGAGCTGGAGGACTCGGCCCTGTATCTCTGTGCCAGCAGC
"""
str_mock_VDJ_fln_genomicDs = \
""">TRBD1*01
GGGACAGGGGGC
>TRBD2*01
GGGACTAGCGGGGGGG
>TRBD2*02
GGGACTAGCGGGAGGG
"""
str_mock_VDJ_fln_genomicJs = \
""">TRBJ1-1*01
TGAACACTGAAGCTTTCTTTGGACAAGGCACCAGACTCACAGTTGTAG
>TRBJ1-2*01
CTAACTATGGCTACACCTTCGGTTCGGGGACCAGGTTAACCGTTGTAG
>TRBJ1-3*01
CTCTGGAAACACCATATATTTTGGAGAGGGAAGTTGGCTCACTGTTGTAG
>TRBJ1-4*01
CAACTAATGAAAAACTGTTTTTTGGCAGTGGAACCCAGCTCTCTGTCTTGG
>TRBJ1-5*01
TAGCAATCAGCCCCAGCATTTTGGTGATGGGACTCGACTCTCCATCCTAG
>TRBJ1-6*01
CTCCTATAATTCACCCCTCCACTTTGGGAATGGGACCAGGCTCACTGTGACAG
>TRBJ1-6*02
CTCCTATAATTCACCCCTCCACTTTGGGAACGGGACCAGGCTCACTGTGACAG
>TRBJ2-1*01
CTCCTACAATGAGCAGTTCTTCGGGCCAGGGACACGGCTCACCGTGCTAG
>TRBJ2-2*01
CGAACACCGGGGAGCTGTTTTTTGGAGAAGGCTCTAGGCTGACCGTACTGG
>TRBJ2-2P*01
CTGAGAGGCGCTGCTGGGCGTCTGGGCGGAGGACTCCTGGTTCTGG
>TRBJ2-3*01
AGCACAGATACGCAGTATTTTGGCCCAGGCACCCGGCTGACAGTGCTCG
>TRBJ2-4*01
AGCCAAAAACATTCAGTACTTCGGCGCCGGGACCCGGCTCTCAGTGCTGG
>TRBJ2-5*01
ACCAAGAGACCCAGTACTTCGGGCCAGGCACGCGGCTCCTGGTGCTCG
>TRBJ2-6*01
CTCTGGGGCCAACGTCCTGACTTTCGGGGCCGGCAGCAGGCTGACCGTGCTGG
>TRBJ2-7*01
CTCCTACGAGCAGTACTTCGGGCCGGGCACCAGGCTCACGGTCACAG
>TRBJ2-7*02
CTCCTACGAGCAGTACGTCGGGCCGGGCACCAGGCTCACGGTCACAG
"""
str_mock_VDJ_fln_V_gene_CDR3_anchors = \
"""gene;anchor_index;gfunction
TRBV1*01;267;P
TRBV2*01;273;F
TRBV2*02;273;(F)
TRBV2*03;273;(F)
TRBV3-1*01;270;F
TRBV3-1*02;270;(F)
TRBV3-2*01;270;P
TRBV3-2*02;270;P
TRBV3-2*03;270;(P)
TRBV4-1*01;270;F
TRBV4-1*02;243;(F)
TRBV4-2*01;270;F
TRBV4-2*02;270;(F)
TRBV4-3*01;270;F
TRBV4-3*02;270;(F)
TRBV4-3*03;270;(F)
TRBV4-3*04;219;(F)
TRBV5-1*01;270;F
TRBV5-1*02;270;(F)
TRBV5-2*01;259;P
TRBV5-3*01;270;ORF
TRBV5-3*02;270;ORF
TRBV5-4*01;270;F
TRBV5-4*02;270;(F)
TRBV5-4*03;222;(F)
TRBV5-4*04;180;(F)
TRBV5-5*01;270;F
TRBV5-5*02;270;(F)
TRBV5-5*03;270;(F)
TRBV5-6*01;270;F
TRBV5-7*01;270;ORF
TRBV5-8*01;270;F
TRBV5-8*02;226;(F)
"""
str_mock_VDJ_fln_J_gene_CDR3_anchors = \
"""gene;anchor_index;function
TRBJ1-1*01;17;F
TRBJ1-2*01;17;F
TRBJ1-3*01;19;F
TRBJ1-4*01;20;F
TRBJ1-5*01;19;F
TRBJ1-6*01;22;F
TRBJ1-6*02;22;F
TRBJ2-1*01;19;F
TRBJ2-2*01;20;F
TRBJ2-3*01;18;F
TRBJ2-4*01;19;F
TRBJ2-5*01;17;F
TRBJ2-6*01;22;F
TRBJ2-7*01;16;F
"""

str_mock_VDJ_fln_dict = dict()
str_mock_VDJ_fln_dict['fln_genomicVs'] = str_mock_VDJ_fln_genomicVs
str_mock_VDJ_fln_dict['fln_genomicDs'] = str_mock_VDJ_fln_genomicDs
str_mock_VDJ_fln_dict['fln_genomicJs'] = str_mock_VDJ_fln_genomicJs
str_mock_VDJ_fln_dict['fln_V_gene_CDR3_anchors'] = str_mock_VDJ_fln_V_gene_CDR3_anchors
str_mock_VDJ_fln_dict['fln_J_gene_CDR3_anchors'] = str_mock_VDJ_fln_J_gene_CDR3_anchors


class MyTestCase(unittest.TestCase):

    def setUp(self) -> None:
        self.tmp_dir = tempfile.TemporaryDirectory(dir=".", prefix="models")
        # self.fln_dict = get_default_fln_names_for_directory(self.tmp_dir.name)
        self.fln_dict = get_default_ref_genome_fln_paths(self.tmp_dir.name + "/ref_genome")
        self.ref_genome_path_dir = self.tmp_dir.name + "/ref_genome"
        os.makedirs(self.ref_genome_path_dir)
        for fln_key in self.fln_dict.keys():
            if fln_key in str_mock_VDJ_fln_dict.keys():
                with open(self.fln_dict[fln_key], mode='w') as ofile:
                    ofile.write(str_mock_VDJ_fln_dict[fln_key])
            # self.assertTrue(os.path.isfile(self.fln_dict[fln_key]))
            # print(self.fln_dict[fln_key],  )

    def test_check_files(self):
        for fln_key in self.fln_dict.keys():
            self.assertTrue(os.path.isfile(self.fln_dict[fln_key]))

    def test_IgorRefGenome(self):
        ref_geno = IgorRefGenome()
        ref_geno.path_ref_genome = self.ref_genome_path_dir
        ref_geno.update_fln_names(path_ref_genome=ref_geno.path_ref_genome)
        ref_geno.load_dataframes_from_ref_genome_files()
        print(ref_geno.to_dict())
        # Check equality of properties
        self.assertTrue(True, True)

    def test_IgorRefGenome_from_path(self):
        ref_genome = IgorRefGenome.load_from_path(self.ref_genome_path_dir)
        ref_genome.load_J_anchors_from_file(ref_genome.fln_J_gene_CDR3_anchors)
        print(ref_genome.df_J_anchors)
        row = ref_genome.df_J_ref_genome.loc[9]
        # print(row)
        self.assertEqual(row['name'], "TRBJ2-2P*01")
        self.assertEqual(row['value'], "CTGAGAGGCGCTGCTGGGCGTCTGGGCGGAGGACTCCTGGTTCTGG")

        self.assertTrue(np.isnan(row['anchor_index']))
        self.assertTrue(np.isnan(row['function']))

        # time.sleep(5)
        # ref_genome = IgorRefGenome.load_from_dataframe_genomics_dict(ref_genome_pd_dict)

    def test_IgorRefGenome_clean_empty_anchors(self):
        ref_genome = IgorRefGenome.load_from_path(self.ref_genome_path_dir)
        row_9 = ref_genome.df_J_ref_genome.loc[9].copy()
        row_10 = ref_genome.df_J_ref_genome.loc[10].copy()

        self.assertEqual(row_9['name'], "TRBJ2-2P*01")
        self.assertEqual(row_9['value'], "CTGAGAGGCGCTGCTGGGCGTCTGGGCGGAGGACTCCTGGTTCTGG")

        self.assertTrue(np.isnan(row_9['anchor_index']))
        self.assertTrue(np.isnan(row_9['function']))

        ref_genome.clean_empty_anchors()
        row_new_9 = ref_genome.df_J_ref_genome.loc[9].copy()

        self.assertEqual(row_10['name'], row_new_9['name'])
        self.assertEqual(row_10['value'], row_new_9['value'])
        self.assertEqual(row_10['anchor_index'], row_new_9['anchor_index'])
        self.assertEqual(row_10['function'], row_new_9['function'])

    def test_IgorRefGenome_from_dataframe_genomics_dict(self):
        ref_genome_pd_dict = dict()
        ref_genome_pd_dict['V'] = get_dataframe_from_fasta_and_csv_anchors(
            self.fln_dict['fln_genomicVs'], self.fln_dict['fln_V_gene_CDR3_anchors'])

        ref_genome_pd_dict['J'] = get_dataframe_from_fasta_and_csv_anchors(
            self.fln_dict['fln_genomicJs'], self.fln_dict['fln_J_gene_CDR3_anchors'])

        ref_genome_pd_dict['D'] = get_dataframe_from_fasta_and_csv_anchors(
            self.fln_dict['fln_genomicDs'])

        ref_genome = IgorRefGenome.load_from_dataframe_genomics_dict(ref_genome_pd_dict)
        row = ref_genome.df_J_ref_genome.loc[9]
        self.assertEqual(row['name'], "TRBJ2-2P*01")
        self.assertEqual(row['value'], "CTGAGAGGCGCTGCTGGGCGTCTGGGCGGAGGACTCCTGGTTCTGG")

        self.assertTrue(np.isnan(row['anchor_index']))
        self.assertTrue(np.isnan(row['function']))

    def test_IgorRefGenome_default(self):
        species = "human"
        chain = "tcr_beta"
        ref_genome_files_dict = get_default_fln_dict_ref_genomes_species_chain(species, chain)
        genomes = IgorRefGenome(**ref_genome_files_dict)
        genomes_default = IgorRefGenome.load_default(species, chain)
        self.assertEqual(genomes.V.loc[0].values[0], genomes_default.V.loc[0].values[0])
        self.assertEqual(genomes.V.loc[0].values[1], genomes_default.V.loc[0].values[1])

        # self.assertEqual(genomes_default.df_V_ref_genome, genomes_default.V)

    def test_IgorRefGenome_write_ref_genome(self):
        # 1. Make a temporal directory
        tmp_test_dir = tempfile.TemporaryDirectory(".", prefix="testando")
        species = "human"
        chain = "tcr_beta"

        # 2. load default RefGenome
        genomes = IgorRefGenome.load_default(species, chain)
        print(genomes.df_genomicVs)
        print(genomes.df_V_anchors)
        print(genomes.df_V_ref_genome)

        # 3. write ref_genome in temporary directory.
        genomes.write_ref_genome_dir(tmp_test_dir.name)

        # 4. Assert if files in directory
        tmp_test_fln_dict = dict()
        tmp_test_fln_dict['fln_genomicVs'] = tmp_test_dir.name + "/" + "genomicVs.fasta"
        tmp_test_fln_dict['fln_genomicDs'] = tmp_test_dir.name + "/" + "genomicDs.fasta"
        tmp_test_fln_dict['fln_genomicJs'] = tmp_test_dir.name + "/" + "genomicJs.fasta"
        tmp_test_fln_dict['fln_V_gene_CDR3_anchors'] = tmp_test_dir.name + "/" + "V_gene_CDR3_anchors.csv"
        tmp_test_fln_dict['fln_J_gene_CDR3_anchors'] = tmp_test_dir.name + "/" + "J_gene_CDR3_anchors.csv"
        for fln_key in self.fln_dict.keys():
            self.assertTrue(os.path.isfile(tmp_test_fln_dict[fln_key]))

        # 5. Remove temporal directory.
        tmp_test_dir.cleanup()

    def test_property_dict_genomicXs(self):

        self.assertTrue(True, False)

    def test_IgorRefGenome_from_imgt(self):
        print(IgorRefGenome.get_imgt_list_species())
        hs_trb_imgt_ref_genome = IgorRefGenome.load_VDJ_from_IMGT_website("Homo+sapiens", "TRB")
        self.assertIsInstance(hs_trb_imgt_ref_genome, IgorRefGenome)

        hs_tra_imgt_ref_genome = IgorRefGenome.load_VJ_from_IMGT_website("Homo+sapiens", "TRA")
        self.assertIsInstance(hs_tra_imgt_ref_genome, IgorRefGenome)

    def test_something(self):
        ref_genome = get_IgorRefGenome_VDJ_from_IMGT("Homo+sapiens", "TRB")


    def tearDown(self) -> None:
        self.tmp_dir.cleanup()


if __name__ == '__main__':
    unittest.main()
