from collections import Counter
import itertools


class FMTL_iterator():
    """
    Simple indexed iterator on FMTL
    """

    def __init__(self, fmtl, idxs):
        self.fmtl = fmtl
        self.idxs = idxs

    def __getitem__(self, i):
        return self.fmtl[self.idxs[i]]

    def __len__(self):
        return len(self.idxs)

    def __iter__(self):
        self.iter_idx = self.idxs.__iter__()
        return self

    def __next__(self):
        idx = self.iter_idx.__next__()
        return self.fmtl[idx]


class FMTL():
    """
    A field-mappable tuple list
    Each field's atomic unit can be mapped to a value.
    """

    def __init__(self, tuplelist, rows=None):
        if type(rows) is tuple or list:
            rows = {x: i for i, x in enumerate(rows)}

        self.tuplelist = tuplelist
        self.mappings = {}
        self.unknown = {}
        self.rows = rows
        self.special_transform = lambda x: x

    def __len__(self):
        return len(self.tuplelist)

    def __getitem__(self, index):
        """
        Maps field[x]:
            -> if field is a tuple/list maps each element inside, keeping structure
            -> else directly maps
            -> if mapping function return error, tries to map with 'unk' value.
            (see self._rec_apply)
        """
        t = self.special_transform(list(self.tuplelist[index]))

        if len(self.mappings) == 0:
            return tuple(t)
        else:

            for i, m in self.mappings.items():
                if type(m) is dict:
                    t[i] = self._rec_apply(m.__getitem__, t[i], self.unknown.get(i))
                else:
                    try:
                        t[i] = m(t[i])
                    except:
                        if self.unknown.get(i) is not None:
                            return self.unknown.get(i)
                        else:
                            raise KeyError("No mapping or placeholder for value: {}".format(i))

            return tuple(t)

    def __iter__(self):
        self.iter_idxs = range(len(self.tuplelist)).__iter__()
        return self

    def __next__(self):
        return self[self.iter_idxs.__next__()]

    def _f2i(self, field):
        if type(field) == int:
            return field

        if type(field) == str:
            return self.rows[field]

        if type(field) == str and self.rows == None:
            raise IndexError("field {} index is unknown, no rows attribute provided".format(field))

        raise IndexError("field {} doesn't exist".format(field))

    def _rec_apply(self, f, item, unk=None):
        if isinstance(item, list) or isinstance(item, tuple):
            return type(item)(map(lambda x: self._rec_apply(f, x, unk), item))
        else:
            try:
                return f(item)
            except:
                if unk is not None:
                    return unk
                else:
                    raise KeyError("No mapping or placeholder for value: {}".format(item))

    def set_mapping(self, field, mapping, unk=None):
        """
        Sets a mapping for a tuple field. Mappings are functions of a field value or dict
        """
        field = self._f2i(field)
        # print(mapping)
        # print(field)
        self.mappings[field] = mapping

        if unk is not None:
            self.unknown[field] = unk

        return mapping

    def field_gen(self, field, key_iter=None):
        field = self._f2i(field)

        if key_iter is None:
            key_iter = range(len(self))

        for idx in key_iter:
            yield self[idx][field]

    def indexed_iter(self, idxs):
        """
        returns a FMTL_iterator object which is the same FMTL instance which only iterates on the idxs slice.
        """
        return FMTL_iterator(self, idxs)

    def get_stats(self, field, key_iter=None, verbose=False):
        """
        helper for stats on a field. Field values should be hashable.
        """
        field = self._f2i(field)
        d = dict(Counter(self.field_gen(field, key_iter=key_iter)))
        sumv = sum([v for k, v in d.items()])
        class_per = {k: (v / sumv) for k, v in d.items()}

        if verbose:
            print(d)
            print(class_per)

        return d, class_per

    def get_field_dict(self, field, offset=0, max_count=-1, key_iter=None, iter_func=None):
        """
        Helper to create a dict from a field.
        """
        field_gen = self.field_gen(self._f2i(field), key_iter=key_iter)

        if iter_func is not None:
            field_gen = itertools.chain.from_iterable((x for x in iter_func(field_gen)))

        d = Counter(field_gen)

        if max_count > -1:
            d_keys = (k for k, v in d.most_common(max_count))
        else:
            d_keys = d.keys()

        d2k = {c: i for i, c in enumerate(d_keys, offset)}
        return d2k



