from fastai import *

from fastai.text import *

from Bio import Seq

from Bio.Seq import Seq

from Bio import SeqIO

from Bio.SeqRecord import SeqRecord

from Bio.SeqFeature import FeatureLocation, CompoundLocation

""" 

Sample 

data = (AlignmentsItemList.from_folder(bam_sam_folder)

        #Where are the sequences? -> in ```bam_sam_folder``` and its subfolders

        .generate_missing_index(),

        #Alignment files need indexes. Run samtools to generate index if missing

        .take_by_name(list_of_named_alignments),

        #select metachondria, chromosome, etc.

        .toFasta(),

        #generate fasta sequence from alignment

        .do_not_label(),

        #create empty labels for unsupervised learning tasks

        .databunch(bs=16, collate_fn=bb_pad_collate))

        #Finally we convert to a DataBunch, use a batch size of 16,

        # and we use bb_pad_collate to collate the data into a mini-batch

"""

##=====================================

## ItemBase classes

##=====================================

class AlignmentIndexBase(ItemBase):

    pass

class AlignmentItemBase(ItemBase):

    """an alignment item contains an alignment record tracked in bai index file as n_ref"""

    def __init__(self,n_ref:str, seq:str, qual:str,  cigar:object, meta:dict):

        self.n_ref, self.seq, self.qual, self.cigar, self.meta = n_ref, seq, cigar, meta

    def loadAlighment(self, file):

        pass

    def loadAlignmentIndex(self, aignmentFileName):

        with open(aignmentFileName) as

        if self.generate_index

    def _getNamedPart(self, part:str):

    ```select a named parts e.g chromosome, metachondria etc.```

    def  toFasta(self, item:Collection[str], position:slice) -> str:

        pass

    def toVariants(self, item:Collection[str], position:slice) -> Collection[str]:

        pass

class AlignmentFileProcessor(PreProcessor):

    "`PreProcessor` that opens the filenames and read alignment files."

    def process_one(self,item):

        return bam_reader(item) if isinstance(item, Path) else item

class AlignmentItemList(ItemList):

    "Special `ItemList` for BAM/SAM alignment files"

    _bunch = AlignmentDataBunch

    _sequencer,_indexer, _cygar_processor =  SequenceProcessor, IndexProcessor, CygarProcessor

    _is_compressed = True

    _in_memory = False

    def from_file(self, file):

    def do_not_label(self, **kwargs):

        "A special labelling method for unsupervised learning"

        self.__class__ = UnlabeledAlignemtList

        kwargs['label_cls'] = UnlabeledAlignemtList

        return self.label_empty(**kwargs)

    def named_parts(self):

        """Retreave named parts from alignment index"""

        pass

##=====================================

## DataBunch

##=====================================

class AlignmentDataBunch(DataBunch):

    "DataBunch suitable for generic sequence processing."

    @classmethod

    def from_folder(cls, path:PathOrStr, train:str='train', valid:str='valid', test:Optional[str]=None,

                    classes:Collection[Any]=None, tokenizer:Tokenizer=None, vocab:Vocab=None, chunksize:int=10000, max_vocab:int=60000,

                    min_freq:int=2, mark_fields:bool=False, include_bos:bool=True, include_eos:bool=False, **kwargs):

        "Create a `AlignmentDataBunch` from text files in folders."

        path = Path(path).absolute()

        processor = [AlignmentFileProcessor()]  +

                    _get_processor(tokenizer=tokenizer, vocab=vocab, chunksize=chunksize, max_vocab=max_vocab,

                    min_freq=min_freq, mark_fields=mark_fields, include_bos=include_bos, include_eos=include_eos)

        src = (AlignmentItemList.from_folder(path, processor=processor)

                       .split_by_folder(train=train, valid=valid))

        src = src.label_for_clustering() if cls==UnlabeledAlignementDataBunch else src.label_from_folder(classes=classes)

        if test is not None: src.add_test_folder(path/test)

        return src.databunch(**kwargs)

##=====================================

## Unlabeled Alignment Data Bunch

##=====================================

class UnlabeledAlignemtList(AlignementItemList):

    "Special `ItemList` for a language model."

    _bunch = AlignmentDataBunch

class UnlabeledAlignementDataBunch(AlignmentDataBunch):

    "DataBunch suitable for unsupervised learning over alignment data"

    def label_for_clustering(self, **kwargs):

            "A special labelling method for unsupervised learning"

            self.__class__ = UnlabeledAlignemtList

            kwargs['label_cls'] = UnlabeledAlignemtList

            return self.label_empty(**kwargs)