#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# Pavel Korshunov <pavel.korshunov@idiap.ch>
# Mon 12 Oct 14:43:22 CEST 2015

"""
  AVSpoof database implementation of bob.db.verification.utils.Database interface.
  It is essentially a wrapper around a Database implemented in query.py, referred to as AVSpoofDatabase
"""

from . import __doc__ as long_description
from .query import Database as AVSpoofDatabase
import bob.db.verification.utils


class File(bob.db.verification.utils.File):
    def __init__(self, f):
        """
        Initializes this File object with our own File equivalent
        """
        bob.db.verification.utils.File.__init__(self, client_id=f.client_id, path=f.path, file_id=f.id)

        self.__f = f

    def load(self, directory=None, extension='.hdf5'):
        return self.__f.load(directory=directory, extension=extension)

    def get_client_id(self):
        return self.__f.client_id


class Database(bob.db.verification.utils.Database):
    """
    Implements verification API for querying AVspoof database.
    """
    __doc__ = long_description

    def __init__(self, original_directory=None, original_extension=None):
        # call base class constructors to open a session to the database
        bob.db.verification.utils.Database.__init__(self, original_directory=original_directory,
                                                    original_extension=original_extension)

        self.__db = AVSpoofDatabase()

    def groups(self, protocol=None):
        """Returns the names of all registered groups"""

        return self.__db.groups()

    def genders(self):
        """Returns the list of genders"""

        return self.__db.genders()

    def protocols(self):
        """Returns all registered protocols"""
        return self.__db.protocols()

    def protocol_names(self):
        """Returns all registered protocol names"""

        return self.__db.protocol_names()

    def has_protocol(self, name):
        """Tells if a certain protocol is available"""

        return self.__db.has_protocol(name)

    def protocol(self, name):
        """Returns the protocol object in the database given a certain name. Raises
        an error if that does not exist."""

        return self.__db.protocol(name)

    def models(self, groups=None, protocol=None, gender=None):
        """Returns a set of models for the specific query by the user.

        Keyword Parameters:

        protocol
            Protocol is ignored in this context, since its choice has no influence on models.

        groups
            The groups to which the subjects attached to the models belong ('dev', 'eval', 'world')

        gender
            The gender to consider ('male', 'female')

        Returns: A list containing the ids of all models belonging to the given group.
        """
        return self.__db.clients(groups=self._convert_group_names(groups), gender=gender)

    def model_ids(self, groups=None, protocol=None, gender=None):
        """Returns a set of models ids for the specific query by the user.

        Keyword Parameters:

        protocol
            Protocol is ignored in this context, since its choice has no influence on models.

        groups
            The groups to which the subjects attached to the models belong ('dev', 'eval', 'world')

        gender
            The gender to consider ('male', 'female')

        Returns: A list containing the ids of all models belonging to the given group.
        """
        return [client.id for client in self.__db.clients(groups=self._convert_group_names(groups), gender=gender)]

    def clients(self, protocol=None, groups=None, gender=None):
        """Returns a list of Clients for the specific query by the user.
           If no parameters are specified - return all clients.

        Keyword Parameters:

        protocol
            Protocol is ignored in this context, since its choice has no influence on clients.

        groups
            The groups to which the subjects attached to the models belong ('dev', 'eval', 'world')

        gender
            The gender to consider ('male', 'female')

        Returns: A list containing the ids of all models belonging to the given group.
        """

        return self.__db.clients(groups=self._convert_group_names(groups), gender=gender)

    def has_client_id(self, id):
        """Returns True if we have a client with a certain integer identifier"""

        return self.__db.has_client_id(id)

    def client(self, id):
        """Returns the Client object in the database given a certain id. Raises
        an error if that does not exist."""

        return self.__db.client(id)

    def get_client_id_from_model_id(self, model_id, **kwargs):
        """Returns the client_id attached to the given model_id

        Keyword Parameters:

        model_id
            The model_id to consider

        Returns: The client_id attached to the given model_id
        """
        return model_id

    def annotations(self, file):
        """Currently, there is no support for annotations in AVspoof database
        """
        return None

    def _convert_group_names(self, groups):
        """ convert verification groups to what is used in __db
        """
        if groups is None:
            return None
        group_alias = {'world': 'train', 'dev': 'devel', 'eval': 'test'}
        matched_groups = []
        if isinstance(groups, (tuple, list)):
            for group in groups:
                matched_groups += [group_alias[group]]
        else:
            matched_groups = [group_alias[groups]]
        return matched_groups

    def objects(self, protocol=None, purposes=None, model_ids=None,
                groups=None, gender=None, device=None):
        """Returns a set of Files for the specific query by the user.

        Keyword Parameters:

        protocol
            Only two protocols are implemented: 'licit' and 'spoof'. Any other protocol's name will be
            assumed to be and converted to 'licit'.

        purposes
            The purposes can be either 'enroll', 'probe', or their tuple.
            If 'None' is given (this is the default), it is
            considered the same as a tuple with both possible values.

        model_ids
            Only retrieves the files for the provided list of model ids (claimed
            client id).  If 'None' is given (this is the default), no filter over
            the model_ids is performed.

        groups
            One of the groups ('dev', 'eval', 'world') or a tuple with several of them.
            If 'None' is given (this is the default), it is considered the same as a
            tuple with all possible values.

        gender
            The gender to consider ('male', 'female')

        device
            The device to consider ('laptop', 'mobile'). In case of 'licit' protocol, device is recording device,
            and for 'attack' protocol - device is the device of attack.

        Returns: A set of Files with the specified properties.
        """

        matched_groups = self._convert_group_names(groups)

        # by default we assume licit protocol, which we actually express via purposes parameters
        # this is the difference in terminology. In DB, protocol is grandtest (all data), while for
        # verification experiments, there are two types licit and spoof - this is expressed through purposes variable
        if not (protocol == 'licit' or protocol == 'spoof'):
            protocol = 'licit'

        correct_purposes = purposes
        # licit protocol is for real access data only
        if protocol == 'licit':
            # by default we assume all real data
            if purposes is None:
                correct_purposes = ('enroll', 'probe')

        # spoof protocol uses real data for enrollment and spoofed data for probe
        # so, probe set is the same as attack set
        if protocol == 'spoof':
            # by default we return all data (enroll:realdata + probe:attackdata)
            if purposes is None:
                correct_purposes = ('enroll', 'attack')
            # otherwise replace 'probe' with 'attack'
            elif isinstance(purposes, (tuple, list)):
                correct_purposes = []
                for purpose in purposes:
                    if purpose == 'probe':
                        correct_purposes += ['train']
                    else:
                        correct_purposes += [purpose]
            elif purposes == 'probe':
                correct_purposes = ('attack',)

        if 'attack' in correct_purposes:
            objects = self.__db.objects(groups=matched_groups, cls=correct_purposes, clients=model_ids, gender=gender,
                                     attackdevices=device)

        else:
            objects = self.__db.objects(groups=matched_groups, cls=correct_purposes, clients=model_ids, gender=gender,
                                     devices=device)
        # make sure to return verification.utils representation of a file, not the database one
        return [File(f) for f in objects]