# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

from collections import defaultdict
import json
import codecs
import pymongo
import bson
import hashlib
import operator

from fastcache import clru_cache



class FileDataset(object):

    def __init__(self, file_path=None, file_path_in=None, file_path_out=None):

        self.file_path_in = file_path_in or file_path
        self.file_path_out = file_path_out or file_path
        if self.file_path_in is None:
            raise ValueError("Must provide 'file_path' or 'file_path_in'")
        if self.file_path_out is None:
            raise ValueError("Must provide 'file_path' or f'ile_path_out'")

        self._attributes = {}
        self._hits = defaultdict(set)
        self._classifications = {}
        self.classificator = None

        if self.file_path_in:
            self.load(self.file_path_in)


    def load(self, file_path):
        try:
            with codecs.open(file_path, "rb", encoding="utf-8") as fd:
                raw_data = json.load(fd)
                self._attributes = raw_data["attributes"]
                self._hits = defaultdict(set, {k: set(v) for k, v in raw_data["hits"].iteritems()})
                self._classifications = raw_data["classifications"]
            return True
        except IOError:
            return False

    def register_attribute(self, attribute):
        if attribute in self._attributes:
            return False
        count = len(self._attributes)
        self._attributes[attribute] = count
        return True

    def hit(self, instance, attribute):
        position = self._attributes[attribute]
        self._hits[instance].add(position)
        if self.classificator and instance not in self._classifications:
            self.classify(instance, self.classificator.classify(instance))

    def classify(self, instance, value):
        self._classifications[instance] = value

    def save(self, file_path=None):
        file_path = file_path or self.file_path_out
        with codecs.open(file_path, "wb", encoding="utf-8") as fd:
            json.dump(self._to_raw_data(), fd)

    def _to_raw_data(self):
        return {
            "attributes": self._attributes,
            "hits": {k: [i for i in sorted(v)] for k, v in self._hits.iteritems()},
            "classifications": self._classifications,
        }

    @property
    def attributes(self):
        sorted_attributes = sorted(self._attributes.items(), key=operator.itemgetter(1))
        # <DEBUG>
        for i, (a, p) in enumerate(sorted_attributes):
            assert p == i
        # </DEBUG>
        return (attr for attr, pos in sorted_attributes)

    def attributes_len(self):
        return len(self._attributes)

    @property
    def individuals(self):
        # return ((k, v, self._classifications[k]) for k, v in self._hits.iteritems())
        return ((k, v, self._classifications[k]) for k, v in sorted(self._hits.iteritems(), key=operator.itemgetter(0)))


class DbDataset(object):

    ATTRIBUTES_COLLECTION = "attributes"
    HITS_COLLECTION = "hits"
    INSTANCES_COLLECTION = "instances"
    COUNTERS_COLLECTION = "counters"
    COUNTER = "last_attribute"

    def __init__(self, host, port=27017, database=None, username=None, password=None, auth_source=None):
        self._client = pymongo.MongoClient(host, port)
        self._db = self._client[database] if database else self._client.get_default_database()
        if username and password and auth_source:
            self._db.authenticate(username, password, auth_source)

        self.classificator = None
        self._hash_cache = {}
        
    def create_indexes(self):
        self._db[self.ATTRIBUTES_COLLECTION].create_index("name", unique=True)
        self._db[self.INSTANCES_COLLECTION].create_index("name", unique=True)

        # Add the field "position" in each attribute:
        i = 0
        for attr in self._db[self.ATTRIBUTES_COLLECTION].find().sort("name", pymongo.ASCENDING):
            attr["position"] = i
            i += 1
            self._db[self.ATTRIBUTES_COLLECTION].replace_one({"_id": attr["_id"]}, attr, upsert=False)

        self._db[self.ATTRIBUTES_COLLECTION].create_index("position", unique=True)
    
    def register_attribute(self, attribute):
        attribute_id = self._hash(attribute)
        try:
            self._db[self.ATTRIBUTES_COLLECTION].insert_one({
                "_id": attribute_id,
                "name": attribute,
            })
        except pymongo.errors.DuplicateKeyError:
            return False
        else:
            return True


    def hit(self, instance, attribute):
        hit_id = self._hash("{}@{}".format(attribute, instance))
        instance_id = self._hash(instance)
        instance_ref = bson.dbref.DBRef(self.ATTRIBUTES_COLLECTION, instance_id)
        attribute_id = self._hash(attribute)
        attribute_ref = bson.dbref.DBRef(self.ATTRIBUTES_COLLECTION, attribute_id)
        try:
            self._db[self.HITS_COLLECTION].insert_one({
                "_id": hit_id,
                "instance": instance_ref,
                "attribute": attribute_ref,
            })
        except pymongo.errors.DuplicateKeyError:
            return False

        if self.classificator and not self._classified(instance_id):
            self.classify(instance, self.classificator.classify(instance))
        return True

    def classify(self, instance, value):
        instance_id = self._hash(instance)
        self._db[self.INSTANCES_COLLECTION].insert_one({
            "_id": instance_id,
            "name": instance,
            "classification": value
        })

    def _classified(self, instance_id):
        return self._db[self.INSTANCES_COLLECTION].find_one({"_id": instance_id}) is not None

    def save(self, file_path=None):
        pass

    def _hash(self, obj):
        if str(obj) in self._hash_cache:
            return self._hash_cache[str(obj)]
        h = hashlib.sha1(str(obj)).hexdigest()
        self._hash_cache[str(obj)] = h
        return h

    @property
    def attributes(self):
        for i, attr in enumerate(self._db[self.ATTRIBUTES_COLLECTION].find().sort("position", pymongo.ASCENDING)):
            assert attr["position"] == i, "{} != {}".format(attr["position"], i)
            yield attr["name"]

    # @property
    # def attributes(self):
    #     for attr in self._db[self.ATTRIBUTES_COLLECTION].find().sort("name", pymongo.ASCENDING):
    #         yield attr["name"]

    def attributes_len(self):
        return self._db[self.ATTRIBUTES_COLLECTION].count()

    @property
    def individuals(self):
        for instance in self._db[self.INSTANCES_COLLECTION].find().sort("name", pymongo.ASCENDING):
            _id = instance["_id"]

            # hits = [self._db[self.ATTRIBUTES_COLLECTION].find_one({"_id": hit["attribute"].id})["position"]
            #  for hit in self._db[self.HITS_COLLECTION].find({"instance.$id": _id})]

            hits = []
            for hit in self._db[self.HITS_COLLECTION].find({"instance.$id": _id}):
                attr_id = hit["attribute"].id
                try:
                    position = self._get_attribute_position(attr_id)
                    hits.append(position)
                except AssertionError:
                    raise AssertionError(
                        "Can't find attr with id = '{}'. Related with hit with _id = '{}'. Related with instance with id = '{}'".format(
                            attr_id, hit["_id"], _id))
            yield instance["name"], hits, instance["classification"]

    # @clru_cache(maxsize=1024*10, typed=False)
    # def _get_attribute_position(self, attr_id):
    #     i = 0
    #     for attr in self._db[self.ATTRIBUTES_COLLECTION].find().sort("name", pymongo.ASCENDING):
    #         if attr["_id"] == attr_id:
    #             return i
    #         i += 1
    #     raise AssertionError("Unable to find attribute '{}'".format(attr_id))

    @clru_cache(maxsize=1024 * 10, typed=False)
    def _get_attribute_position(self, attr_id):
        return self._db[self.ATTRIBUTES_COLLECTION].find_one({"_id":attr_id})["position"]
