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

import re

from dna2vec.multi_k_model import MultiKModel

import random

from gensim.models import Word2Vec

from pathlib import Path

import os

from tqdm import tqdm

from  torch import tensor

import gzip

from mimetypes import guess_type

from functools import partial

# fasta extensions bansed on https://en.wikipedia.org/wiki/FASTA_format

gen_seq_extensions = ['.fasta', '.fna', '.ffn', '.faa', '.frn','.fa',".gz"]

def ifnone(a: Any, b: Any) -> Any:

    "`a` if `a` is not None, otherwise `b`."

    return b if a is None else a

def gen_seq_reader(fn: PathOrStr):

    "Read the sequences in `fn`."

    ext = str(fn).split(".")[-1]

    _open = partial(gzip.open, mode='rt') if ext.lower() == "gz" else open

    with _open(fn) as f:

        return SeqIO.to_dict(SeqIO.parse(f, 'fasta'))

def seq_record(fn: PathOrStr, record_id:str):

    content = gen_seq_reader(fn)

    for record in content:

        if content[record].id == record_id:

            return content[record].seq

    return None

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

## Processors

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

class GSFileProcessor(PreProcessor):

    """`PreProcessor` Opens the fasta file listed in item,

    reads fasta and returns sequences with IDs provided by the item.

    """

    def __init__(self, ds: ItemList = None, filters=None):

        self.ds,self.filters = ds, filters

    def process_one(self, item) -> Seq:

        return seq_record(item["file"], item["id"])

    def process(self, ds: Collection) -> Collection[Seq]:

        df = pd.DataFrame(data=list(ds.items), columns=['file', 'description', "id", "name"])

        multi_fastas = df.groupby("file").agg({"id": list})

        res = []

        for row in tqdm(multi_fastas.index.values):

            content = gen_seq_reader(str(row))

            for record in content:

                if content[record].id in multi_fastas.loc[row, 'id']:

                    res.append(content[record].seq)

        ds.items = apply_filters(res,self.filters)

        ds.state = "sequence"

class GSTokenizer():

    def __init__(self, ngram=8, skip=0, n_cpus=1):

        self.ngram, self.skip,self.n_cpus = ngram, skip,n_cpus

    def tokenizer(self, t):

        if self.ngram == 1:

            toks = list(t)

            if self.skip > 0:

                toks = toks[::2] if self.skip == 1 else toks[::self.skip]

        else:

            toks = [t[i:i + self.ngram] for i in range(0, len(t), self.ngram + self.skip) if i+self.ngram < len(t)]

        return toks

    def _process_all_1(self, texts:Collection[str]) -> List[List[str]]:

        "Process a list of `texts` in one process."

        return [self.tokenizer(str(t)) for t in texts]

    def process_all(self, texts:Collection[str]) -> List[List[str]]:

        "Process a list of `texts`."

        if self.n_cpus <= 1: return self._process_all_1(texts)

        with ProcessPoolExecutor(self.n_cpus) as e:

            res = sum(e.map(self._process_all_1,

                             partition_by_cores(texts, self.n_cpus)), [])

        return res

class GSTokenizeProcessor(PreProcessor):

    "`PreProcessor` that tokenizes the texts in `ds`."

    def __init__(self, ds: ItemList = None, tokenizer: Tokenizer = None, ngram:int=8, skip:int=0, chunksize: int = 10000,

                 mark_fields: bool = False):

        self.tokenizer, self.chunksize, self.mark_fields = ifnone(tokenizer, GSTokenizer(ngram=ngram, skip=skip)), chunksize, mark_fields

    def process_one(self, sequence):

        return self.tokenizer.tokenizer(str(sequence))

    def process(self, ds):

        tokens = []

        # if len(ds.items) < self.chunksize: ds.items = self.tokenizer._process_all_1(ds.items); return

        chunks = len(ds.items) // self.chunksize + 1

        for i in tqdm(range(chunks)):

            advance = min((len(ds.items) - i * self.chunksize), self.chunksize)

            tokens += self.tokenizer.process_all(ds.items[i:i + advance])

        ds.items = tokens

        ds.state = "tokens"

class GSVocab(Vocab):

    def __init__(self, itos):

        self.itos = itos

        self.stoi = collections.defaultdict(int, {v: k for k, v in enumerate(self.itos)})

    @classmethod

    def create(cls, tokens, max_vocab, min_freq):

        freq = Counter(p for o in tokens for p in o)

        itos = [o for o, c in freq.most_common(max_vocab) if c >= min_freq]

        itos.insert(0, 'pad')

        return cls(itos)

class GSNumericalizeProcessor(PreProcessor):

    "`PreProcessor` that numericalizes the tokens in `ds`."

    def __init__(self, ds: ItemList = None, vocab: Vocab = None, max_vocab: int = 80000, min_freq: int = 3):

        vocab = ifnone(vocab, ds.vocab if ds is not None else None)

        self.vocab, self.max_vocab, self.min_freq = vocab, max_vocab, min_freq

    def process_one(self, item): return np.array(self.vocab.numericalize(item), dtype=np.int64)

    def process(self, ds):

        if self.vocab is None: self.vocab = GSVocab.create(ds.items, self.max_vocab, self.min_freq)

        ds.vocab = self.vocab

        super().process(ds)

        ds.state="numericalized"

class Dna2VecProcessor(PreProcessor):

    "`PreProcessor` that tokenizes the texts in `ds`."

    def __init__(self, ds: ItemList = None, agg:Callable=partial(np.mean, axis=0), emb=None, n_cpu=7):

        self.agg, self.n_cpu = agg, n_cpu

        self.emb = None if emb is None else emb if isinstance(emb, Word2Vec) else Word2Vec.load_word2vec_format(emb)

    def process_one(self, tokens):

        if self.emb is None: raise ValueError("Provide path to embedding or Word2Vec object using  ```emb``` instance variable ")

        tokens= list(filter(lambda x: set(x) == set('ATGC'), tokens))

        vectors = np.asarray([[0.] * 100, [0.] * 100])

        while len(tokens) > 0:

            try:

                vectors = np.asarray(self.emb[tokens])

                break

            except KeyError as e:

                tokens.remove(e.args[0])  # remove k-mer absent in the embedding

        return vectors if self.agg is None else self.agg(vectors)

    def _process_all_1(self, tokens:Collection[str]) -> List[List[str]]:

        return [self.process_one(t) for t in tokens]

    def process(self, ds):

        self.emb = ds.emb if (hasattr(ds, "emb") and ds.emb is not None) else self.emb

        res =[]

        with ProcessPoolExecutor(self.n_cpu) as e:

            res = sum(e.map(self._process_all_1,

                             partition_by_cores(ds.items, self.n_cpu)), [])

        ds.items = res

        ds.state = "vector"

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

## DataBunch

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

class GSUDataBunch(DataBunch):

    "DataBunch suitable for unsupervised learning from fasta data"

    @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 = 70000, min_freq: int = 2, mark_fields: bool = False, include_bos: bool = True,

                    filters:Collection[Callable] = None,

                    include_eos: bool = False, n_cpus: int = None, ngram: int = 8, skip: int = 0, **kwargs):

        "Create a unsupervised learning data bunch from fasta  files in folders."

        path = Path(path).absolute()

        tok = Tokenizer(tok_func=partial(GSTokenizer, ngram=ngram, skip=skip), n_cpus=n_cpus)

        processor = [GSFileProcessor(),

                      GSTokenizeProcessor(tokenizer=tok, chunksize=chunksize, mark_fields=mark_fields),

                      GSNumericalizeProcessor(vocab=vocab, max_vocab=max_vocab, min_freq=min_freq)]

        src = ItemLists(path, NumericalizedGSList.from_folder(path=path, filters = filters, processor=processor),

                        ItemList(items=[],ignore_empty=True))

        src=src.label_empty()

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

        return src.databunch(**kwargs)

class Dna2VecDataBunch(DataBunch):

    "DataBunch of tokenized genomic sequences for use with dna2vec embedding"

    @classmethod

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

                    classes: Collection[Any] = None, tokenizer: Tokenizer = None,

                    chunksize: int = 1000, mark_fields: bool = False,

                    filters:Collection[Callable] = None, labeler:Callable=None, n_cpus: int = 1,

                    ngram: int = 8, skip: int = 0, agg:Callable=None, emb = None, **kwargs):

        path = Path(path).absolute()

        tok = ifnone(tokenizer, GSTokenizer(ngram=ngram, skip=skip, n_cpus=n_cpus))

        processors = [GSFileProcessor(),

                     GSTokenizeProcessor(tokenizer=tok, chunksize=chunksize, mark_fields=mark_fields),

                     Dna2VecProcessor(emb=emb, agg=agg)]

        train_items = Dna2VecList.from_folder(path=path / train, filters=filters, processor=processors)

        valid_items = Dna2VecList.from_folder(path=path / valid, filters=filters, processor=processors)

        src = ItemLists(path, train_items, valid_items)

        tl,cl = train_items.label_from_description(labeler)

        vl,_ = valid_items.label_from_description(labeler, labels=cl)

        src=src.label_from_lists(train_labels=tl, valid_labels=vl,label_cls=CategoryList, classes = cl)

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

        return src.databunch(**kwargs)

def regex_filter(items:Collection, rx:str= "", target:str= "description", search=True, keep:bool=True) -> Collection:

    if rx== "": return items

    rx = re.compile(rx)

    if search: return list(filter(lambda x: rx.search(x[target]) if keep else  not rx.search(x[target]), items))

    return list(filter(lambda x: rx.match(x[target]) if keep else not rx.match(x[target]), items))

def id_filter(items:Collection, ids:Collection)->Collection:

    return [i for i in list(items) if i['id'] in ids]

def name_filter(items:Collection, names:Collection)->Collection:

    return [i for i in list(items) if i['name'] in names]

def count_filter(items:Collection, num_fastas:tuple=(1, None), keep:int=None, sample:str= "first") -> Collection:

    df = pd.DataFrame(data=list(items), columns=['file', 'description', "id", "name"])

    df_agg = df.groupby("file").agg({"id": list})

    selected_ids = []

    for file in iter(df_agg.index):

        ids = df_agg.loc[file,"id"]

        if len(ids) < num_fastas[0]: continue

        if num_fastas[1] is not None and len(ids) > num_fastas[1]: continue

        if keep is None:

            selected_ids += ids

        else:

            selected_ids += ids[:min([keep, len(ids)])] if sample == "first" else ids[-min([keep, len(ids)]):]

    res= id_filter(items=items, ids=selected_ids)

    return res

def seq_len_filter(items:Collection, len:tuple=(1, None), keep:bool=True) -> Collection:

    """filters sequence by length. ```len``` tuple is (min,max) values for filtering, ```keep``` controls """

    selected_ids=[i["id"] for i in items if (i["len"] >= len[0] and (len[1] is None or i["len"] < len[1])) == keep]

    res= id_filter(items=items, ids=selected_ids)

    return res

def total_count_filter(items:Collection, parser:Callable,num_fastas:tuple=(1, None), balance:bool=True) -> Collection:

    """Counts items for category extracted by parser.

    Subsamples overrepresented categories to match max amount of samples in the least represented category  """

    pass

def describe(items:Collection) -> dict:

    """compute statistics for items in the list"""

    pass

def apply_filters(dicts:Collection, filters:Union[Callable, Collection[Callable]]=None) -> Collection:

    if filters is None: return dicts

    if callable(filters): return filters(items=dicts)

    for f in filters: dicts = f(items=dicts)

    return dicts

class NumericalizedGSList(ItemList):

    "`ItemList`of numericalised genomic sequences."

    _bunch, _processor = GSUDataBunch, [GSFileProcessor, GSTokenizeProcessor, GSNumericalizeProcessor]

    def __init__(self, items:Iterator, vocab:Vocab=None, pad_idx:int=1, **kwargs):

        super().__init__(items, **kwargs)

        self.vocab,self.pad_idx = vocab,pad_idx

        self.copy_new += ['vocab', 'pad_idx']

    @classmethod

    def from_folder(cls, path: PathOrStr = '.', extensions: Collection[str] = None,

                    filters:Union[Callable, Collection[Callable]]=None, vocab:GSVocab=None, **kwargs) -> ItemList:

        "Get the list of files in `path` that have an image suffix. `recurse` determines if we search subfolders."

        extensions = ifnone(extensions, gen_seq_extensions)

        this = super().from_folder(path=path, extensions=extensions, **kwargs)

        return cls(items=fasta_content(this,filters), path=path, vocab=vocab, **kwargs)

def _get_files(parent, p, f, extensions):

    "Return  sequence of files in a folder files including gzipped format for given extensions"

    p = Path(p)#.relative_to(parent)

    if isinstance(extensions,str): extensions = [extensions]

    low_extensions = [e.lower() for e in extensions] if extensions is not None else None

    res = [p/o for o in f if not o.startswith('.')

           and (extensions is None or f'.{o.split(".")[-1].lower()}' in low_extensions

           or (o.split(".")[-1].lower() == "gz" and  f'.{o.split(".")[-2].lower()}' in low_extensions))]

    return res

class Dna2VecList(ItemList):

    "`ItemList` of Kmer tokens vectorized by dna2vec embedding"

    _bunch, _processor = Dna2VecDataBunch, [GSFileProcessor, GSTokenizeProcessor,Dna2VecProcessor]

    def __init__(self, items:Iterator, path, ngram:int=8, n_cpus:int=7,

                 emb:Union[Path,str,Word2Vec]=None,

                 agg:Callable=partial(np.mean, axis=0), #mean values of dna2vec embedding vectors for all k-mers in genome

                **kwargs):

        super().__init__(items, path, **kwargs)

        self.ngram,self.agg,self.n_cpus = ngram,agg,n_cpus

        self.emb = emb if isinstance(emb, Word2Vec) else Word2Vec.load_word2vec_format(emb) if emb is not None else None

        self.descriptions, self.ids, self.names, self.files, self.lengths= None, None, None, None, None

        self.state = "initial"

    def get_metadata(self, filters):

        dicts = []

        for file in tqdm(self.items):

            content = gen_seq_reader(file)

            dicts += [

                {"file": str(file),

                 'description': content[r].description,

                 'id': content[r].id,

                 'name': content[r].name,

                 "len":len(content[r].seq)}

                for r in content.keys()]

        self.items = apply_filters(dicts, filters)

        self.descriptions = [item['description'] for item in list(self.items)]

        self.ids = [item['id'] for item in list(self.items)]

        self.names = [item['name'] for item in list(self.items)]

        self.files = [item['file'] for item in list(self.items)]

        self.lengths = [item["len"] for item in list(self.items)]

        return self

    @property

    def c(self):

        return len(self.items[0])

    def get(self, i) ->Any:

        return self.items[i]

    def process_one(self, i):

        return self.items[i]

    def analyze_pred(self, pred):

        _, index = ensor.max()

        return index

    @classmethod

    def get_files(cls, path: PathOrStr, extensions: Collection[str] = None, recurse: bool = False,

                  include: Optional[Collection[str]] = None) -> FilePathList:

        "Return list of files in `path` that have a suffix in `extensions`; optionally `recurse`."

        def _get_files(parent, p, f, extensions):

            "Return  sequence of files in a folder files including gzipped format for given extensions"

            p = Path(p)  # .relative_to(parent)

            if isinstance(extensions, str): extensions = [extensions]

            low_extensions = [e.lower() for e in extensions] if extensions is not None else None

            res = [p / o for o in f if not o.startswith('.')

                   and (extensions is None or f'.{o.split(".")[-1].lower()}' in low_extensions

                        or (o.split(".")[-1].lower() == "gz" and f'.{o.split(".")[-2].lower()}' in low_extensions))]

            return res

        if recurse:

            res = []

            for i, (p, d, f) in enumerate(os.walk(path)):

                # skip hidden dirs

                if include is not None and i == 0:

                    d[:] = [o for o in d if o in include]

                else:

                    d[:] = [o for o in d if not o.startswith('.')]

                res += _get_files(path, p, f, extensions)

            return res

        else:

            f = [o.name for o in os.scandir(path) if o.is_file()]

            return _get_files(path, path, f, extensions)

    def label_from_description(self, labeler:Callable=None, labels:Collection=None):

        assert labeler is not None, "must provide labeler"

        lbls=list(map(labeler, self.descriptions))

        cl = list(set(lbls)) if labels is None else labels

        def _oh(i, cat_cnt):

            res=np.zeros(cat_cnt,dtype=int); res[i] = 1

            return res

        # return [_oh(cl.index(x), len(cl)) for x in lbls], cl

        # return [cl.index(x) for x in lbls],cl

        return lbls,cl

    @classmethod

    def from_folder(cls, path: PathOrStr = '.', extensions: Collection[str] = None,

                    filters:Collection[Callable]=None, ngram:int=8, n_cpus=1, agg:Callable=None, recurse:bool = False, **kwargs) -> ItemList:

        "Get the list of files in `path` that have an sequence suffix. `recurse` determines if we search subfolders."

        extensions = ifnone(extensions, gen_seq_extensions)

        this = cls(items=cls.get_files(path, extensions, recurse),path=path, **kwargs)

        return this.get_metadata(filters)

    @classmethod

    def store_by_label_class(self,path):

        """store fasta into multi-fasta files labeled by class """

        pass

    def __repr__(self):

        return f"{self.__class__.__name__} {len(self.items)}  items; {self.ngram}-mer tokensation\n" \

            f" {self.emb}, reducer:{self.agg}" \

            f"\n Head: \n  {self.descriptions[0]}\n  {self.items[0]}" \

            f"\n Tail: \n  {self.descriptions[-1]}\n  {self.items[-1]}"

if __name__ == '__main__':

    #test DataBunch

    # DB = "/home/serge/database/data/genomes/ncbi-genomes-2019-04-07/Bacillus"

    # # DB="/data/genomes/GenSeq_fastas"

    # bunch = Dna2VecDataBunch.from_folder(DB,

    #                                      filters=None,  #[partial(count_filter, keep=3, sample="last")],

    #                                      emb="../pretrained/embeddings/dna2vec-20190611-1940-k8to8-100d-10c-4870Mbp-sliding-LmP.w2v",

    #                                      ngram=8, skip=0,

    #                                      labeler=lambda x: " ".join(x.split()[1:3]),

    #                                      n_cpus=7,agg=partial(np.mean, axis=0),one_hot=True)

    # print(bunch.train_ds.y)

    #test preprocessing for embedding training

    # Dna2VecList.preprocess_for_dna2vec_training(out_path="/data/genomes/dna2vec_train",

    #                                                    path="/data/genomes/GenSeq_fastas",

    #                                                    filters=[partial(regex_filter, rx="plasmid", keep=False),

    #                                                             partial(seq_len_filter, len=(100000,None))])

    #test labeling

    # data.label_from_description(lambda x: x.split()[1], from_item_lists=True)

    # print(data)

    #test get item

    # data = Dna2VecList.from_folder("/data/genomes/GenSeq_fastas/valid",filters=[partial(regex_filter, rx="plasmid")])

    # print(data.get(0))

    #test process all itmes

    # data = Dna2VecList.from_folder("/data/genomes/GenSeq_fastas", filters=[partial(regex_filter, rx="plasmid",keep=False)],

    #                                emb='/data/genomes/embeddings/dna2vec-20190614-1532-k11to11-100d-10c-4870Mbp-sliding-X6c.w2v')

    # # print(data.get(0)))

    # processors = [GSFileProcessor(),GSTokenizeProcessor(tokenizer=GSTokenizer(ngram=11, skip=0, n_cpus=7)),Dna2VecProcessor()]

    # for p in processors: p.process(data)

    # print(data)

    #test gzipped fastas

    data = Dna2VecList.from_folder("/data/genomes/gzipped", filters=[],

                                    emb='/data/genomes/embeddings/dna2vec-20190611-1940-k8to8-100d-10c-4870Mbp-sliding-LmP.w2v')

    print(data.get(0))