from django.contrib import admin
from .models import User, Client, Contrat, Event
from django.contrib.auth.admin import UserAdmin
from django.contrib.admin import ModelAdmin

class UserAdminConfig(UserAdmin):
    ordering = ('-start_date',)
    list_display = ('email', 'username', 'first_name', 'groups',
                    'is_active', 'is_staff')
    list_filter = ('groups', 'is_active')
    search_fields = ('email', 'username', 'groups')
    fieldsets = [
        ('User Information',
         {
             'fields': [('email', 'username'), ('first_name', 'last_name')], }, ),
        ('Permissions',
         {
             'fields': ['is_staff', 'is_active', 'groups'], }, ),
    ]
    add_fieldsets = (
        ('User Information',
         {'classes': ('wide',),
          'fields': (
              'email', 'username', 'first_name', 'password1', 'password2'), },),
        ('Groups',
         {'fields': (
             'groups', 'is_active', 'is_staff',), },),
    )
    filter_horizontal = ()
    radio_fields = {'groups': admin.VERTICAL}

    def save_model(self, request, obj, form, change):
        """
            Check if each competition are already passed or no
            & add keyword 'valid' set on True or False

            :param: competitions: list of competitions
            :return: list of competitions
        """
        if change:
            saler = User.objects.filter(email=obj.email)
            old_team = saler[0].groups
            if old_team != obj.groups:
                clients = Client.objects.filter(sales_contact__in=saler)
                for client in clients:
                    client.sales_contact = None
                    client.save()
        super().save_model(request, obj, form, change)

class ClientAdminConfig(ModelAdmin):
    ordering = ('-date_created',)
    list_display = ('company_name', 'email', 'sales_contact', 'client_status', 'date_created', 'date_updated')
    list_filter = ('sales_contact', 'client_status')
    search_fields = ('company_name', 'email', 'sales_contact')

    fieldsets = [
        ('Client Information',
         {
             'fields': ['company_name', ('first_name', 'last_name')], },),
        ('Contact',
         {
             'fields': ['email', 'phone', 'sales_contact', 'client_status'], },),
    ]
    radio_fields = {'client_status': admin.VERTICAL}

    @admin.action(description="client status = Existant")
    def client_status_ex(self, request, query):
        query.update(client_status="EX")

    @admin.action(description="client status = Potentiel")
    def client_status_po(self, request, query):
        query.update(client_status="PO")

    actions = [client_status_ex, client_status_po]

    def formfield_for_foreignkey(self, db_field, request, **kwargs):
        if db_field.name == 'sales_contact':
            kwargs['queryset'] = User.objects.filter(groups__name__contains='sales')
        return super().formfield_for_foreignkey(db_field, request, **kwargs)

class EventAdminConfig(ModelAdmin):
    ordering = ('-date_created',)
    list_display = ('event_name', 'attendees', 'event_date', 'support_contact', 'contrat', 'event_status')
    list_display_links = ('event_name', 'support_contact', 'contrat')
    list_filter = ('event_date', 'support_contact', 'event_status')
    search_fields = ('event_name',)

    fieldsets = [
        ('Event Information',
         {
             'fields': ['event_name', 'attendees', 'event_date'], },),
        ('Links',
         {
             'fields': ['contrat', 'support_contact'], },),
        ('More',
         {
             'fields': ['event_status', 'notes'], },),
    ]
    radio_fields = {'event_status': admin.VERTICAL}

    @admin.action(description="event status = Organisation")
    def event_status_or(self, request, query):
        query.update(event_status="OR")

    @admin.action(description="event status = Prêt")
    def event_status_pr(self, request, query):
        query.update(event_status="PR")

    @admin.action(description="event status = Terminé")
    def event_status_te(self, request, query):
        query.update(event_status="TE")

    actions = [event_status_te, event_status_or, event_status_pr]

    def formfield_for_foreignkey(self, db_field, request, **kwargs):
        if db_field.name == 'support_contact':
            kwargs['queryset'] = User.objects.filter(groups__name__contains='support')
        return super().formfield_for_foreignkey(db_field, request, **kwargs)

class ContratAdminConfig(ModelAdmin):
    ordering = ('-date_created',)
    list_display = ('internal_contrat_number', 'client', 'amount', 'payment_due', 'status')
    list_filter = ('client', 'status')
    search_fields = ('client',)

    fieldsets = [
        ('Contrat Information',
         {
             'fields': ['internal_contrat_number', 'amount', 'payment_due', 'status'], },),
        ('Links',
         {
             'fields': ['client'], },),
    ]

    def formfield_for_foreignkey(self, db_field, request, **kwargs):
        if db_field.name == 'client':
            kwargs['queryset'] = Client.objects.filter(client_status='EX')
        return super().formfield_for_foreignkey(db_field, request, **kwargs)

admin.site.register(User, UserAdminConfig)
admin.site.register(Client, ClientAdminConfig)
admin.site.register(Contrat, ContratAdminConfig)
admin.site.register(Event, EventAdminConfig)


"""This function constructs a dictionary with abinitio_energies.

Input:
    fname (str) path/filename of ase.db file
    selection (list) ase.db selection
"""
import warnings
import numpy as np
from tqdm import tqdm
from ase.data import chemical_symbols
from ase.symbols import string2symbols, symbols2numbers
# get_distances requires ASE 3.16 or above.
from ase.geometry import get_layers, get_distances
from catlearn.api.ase_atoms_api import images_connectivity
from catlearn.featurize.periodic_table_data import get_radius

ads_syms = ['H', 'C', 'O', 'N', 'S', 'F', 'Cl', 'P', 'Na', 'K']

def catalysis_hub_to_info(images):
    raise NotImplementedError("Coming soon.")

def autogen_info(images):
    """Return a list of atoms objects with atomic group information
    attached to atoms.subsets.
    This information is  needed by some functions in the
    AdsorbateFingerprintGenerator.

    Parameters
    ----------
    images : list
        list of atoms objects representing adsorbates on slabs.
        No further information is required in atoms.subsets.
    """
    images = images_connectivity(images)
    for i, atoms in tqdm(enumerate(images)):
        # Pre-allocate subsets dictionary.
        if not hasattr(atoms, 'subsets'):
            atoms.subsets = {}

        # Identify adsorbate atoms.
        if 'ads_atoms' not in atoms.subsets:
            ads_atoms = detect_adsorbate(atoms)
            atoms.subsets['ads_atoms'] = ads_atoms
            if len(ads_atoms) == 0:
                continue

        # Identify slab atoms.
        if 'slab_atoms' not in atoms.subsets:
            atoms.subsets['slab_atoms'] = slab_index(atoms)
        if not (len(atoms.subsets['slab_atoms']) +
                len(atoms.subsets['ads_atoms']) ==
                len(atoms)):
            raise AssertionError("all atoms must belong to slab or adsorbate.")

        # Identify atoms at node distance 0 and 1 from the surface-ads bond.
        if ('chemisorbed_atoms' not in atoms.subsets or
            'site_atoms' not in atoms.subsets or
                'ligand_atoms' not in atoms.subsets):
            chemi, site, ligand = info2primary_index(atoms)
            atoms.subsets['chemisorbed_atoms'] = chemi
            atoms.subsets['site_atoms'] = site
            atoms.subsets['ligand_atoms'] = ligand

        # Identify surface layers.
        if ('bulk_atoms' not in atoms.subsets or
            'termination_atoms' not in atoms.subsets or
           'subsurf_atoms' not in atoms.subsets):
            bulk, term, subsurf = detect_termination(atoms)
            atoms.subsets['bulk_atoms'] = bulk
            atoms.subsets['termination_atoms'] = term
            atoms.subsets['subsurf_atoms'] = subsurf
    return images

def termination_info(images):
    """Return a list of atoms objects with attached information about
    the slab termination, the slab second outermost layer and the bulk slab
    compositions.

    Parameters
    ----------
    images : list
        list of atoms objects representing adsorbates on slabs.
        The atoms objects must have the following keys in atoms.subsets:

            - 'ads_atoms' : list
                indices of atoms belonging to the adsorbate
            - 'slab_atoms' : list
                indices of atoms belonging to the slab
    """
    for atoms in tqdm(images):
        bulk, subsurf, term = detect_termination(atoms)
        atoms.subsets['bulk_atoms'] = bulk
        atoms.subsets['termination_atoms'] = term
        atoms.subsets['subsurf_atoms'] = subsurf
    return images

def detect_adsorbate(atoms):
    """Return a list of indices of atoms belonging to an adsorbate.

    Parameters
    ----------
    atoms : object
        An ase atoms object.
    """
    if 'ads_atoms' in atoms.subsets:
        return atoms.subsets['ads_atoms']
    if 'slab_atoms' in atoms.subsets:
        return ads_index(atoms)
    if (len(np.unique(atoms.get_tags())) >= 4 and
        min(atoms.get_tags()) <= -1 and
       max(atoms.get_tags()) > 2):
        return tags2ads_index(atoms)
    # Use species chemical formula if known.
    try:
        species = atoms.info['key_value_pairs']['species']
    except KeyError:
        warnings.warn("'species' key or .info['key_value_pairs'] missing.")
        return sym2ads_index(atoms, ads_syms=ads_syms)
    if species == '':
        return []
    elif '-' in species or '+' in species or ',' in species:
        raise NotImplementedError('Co-adsorption.')
    try:
        return formula2ads_index(atoms, species)
    except AssertionError:
        try:
            return connectivity2ads_index(atoms, species)
        except AssertionError:
            return last2ads_index(atoms, species)

def detect_termination(atoms):
    """ Returns three lists, the first containing indices of bulk atoms and
    the second containing indices of atoms in the second outermost layer, and
    the last denotes atoms in the outermost layer or termination or the slab.

    Parameters
    ----------
    atoms : object.
        The atoms object must have the following keys in atoms.subsets:

        'slab_atoms' : list
            indices of atoms belonging to the slab
    """
    if (len(np.unique(atoms.get_tags())) >= 4 and
        min(atoms.get_tags()) <= 0 and
       max(atoms.get_tags()) > 2):
        bulk, term, subsurf = tags_termination(atoms)
    elif len(atoms.constraints) == 1:
        bulk, term, subsurf = constraints_termination(atoms)
    elif 'key_value_pairs' not in atoms.info:
        bulk, term, subsurf = connectivity_termination(atoms)
    elif 'layers' in atoms.info['key_value_pairs']:
        bulk, term, subsurf = layers_termination(atoms)
    else:
        bulk, term, subsurf = connectivity_termination(atoms)
    try:
        sx, sy = atoms.info['key_value_pairs']['supercell'].split('x')
        sx = int(''.join(i for i in sx if i.isdigit()))
        sy = int(''.join(i for i in sy if i.isdigit()))
        if int(sx) * int(sy) != len(term):
            msg = str(len(term)) + ' termination atoms identified.' + \
                ' ' + atoms.info['key_value_pairs']['facet'] + \
                ' ' + atoms.info['key_value_pairs']['supercell'] + \
                '. id=' + str(atoms.info['id'])
            warnings.warn(msg)
    except KeyError:
        pass
    if 'site_atoms' in atoms.subsets:
        for a_s in atoms.subsets['site_atoms']:
            if a_s not in term:
                msg = 'site atom not in term.'
                try:
                    msg += ' ' + str(atoms.info['key_value_pairs']['phase'])
                except KeyError:
                    pass
                try:
                    msg += ' ' + str(atoms.info['key_value_pairs']['facet'])
                    msg += ' ' + str(
                            atoms.info['key_value_pairs']['supercell'])
                except KeyError:
                    pass
                if 'id' in atoms.info:
                    msg += ' ' + str(atoms.info['id'])
                warnings.warn(msg)
                break
    return bulk, term, subsurf

def check_reconstructions(image_pairs):
    """Return a list of database ids, for adsorbate/slab structures, which
    has a reconstructed slab with respect to the reference slab.

    Parameters
    ----------
    image_pairs : list
        List of tuples containing pairs of ASE atoms objects.
        The first element in each tuple must represent an adsorbate*slab
        structure and the second element must represent a slab.
    """
    reconstructed = []
    for i, row in enumerate(image_pairs):
        identical = compare_slab_connectivity(row[0], row[1])
        if not identical:
            reconstructed.append(i)
    return reconstructed

def compare_slab_connectivity(atoms, reference_atoms):
    """Return a boolean for whether an adsorbate has caused a slab to
    reconstruct and change it's connectivity.

    Parameters
    ----------
    atoms : object
        ASE atoms object with connectivity and 'slab_atoms' subsets attached.
        This represents an adsorbate*slab structure.
    reference_atoms : object
        ASE atoms object with connectivity and 'slab_atoms' subsets attached.
        This represents a slab structure.

    Returns
    ----------
    identical : boolean
        Are the connectivities within the slabs identical or not.
    """
    slab_atoms = atoms.subsets['slab_atoms']
    slab_cm = atoms.connectivity[:, slab_atoms][slab_atoms, :]
    reference_cm = reference_atoms.connectivity
    return np.allclose(slab_cm, reference_cm)

def slab_index(atoms):
    """ Returns a list of indices of atoms belonging to the slab.
    These are defined as atoms that are not belonging to the adsorbate.

    Parameters
    ----------
    atoms : ase atoms object
        The atoms object must have the key 'ads_atoms' in atoms.subsets:

            - 'ads_atoms' : list
                indices of atoms belonging to the adsorbate
    """
    chemi = [a.index for a in atoms if a.index not in
             atoms.subsets['ads_atoms']]
    return chemi

def ads_index(atoms):
    """ Returns a list of indices of atoms belonging to the adsorbate.
    These are defined as atoms that are not belonging to the slab.

    Parameters
    ----------
    atoms : ase atoms object
        The atoms object must have the key 'ads_atoms' in atoms.subsets:

            - 'slab_atoms' : list
                indices of atoms belonging to the adsorbate
    """
    ads = [a.index for a in atoms if a.index not in
           atoms.subsets['slab_atoms']]
    return ads

def sym2ads_index(atoms, ads_syms):
    """Return the indexes of atoms from the global list of adsorbate symbols.

    Parameters
    ----------
    atoms : object
        An ase atoms object.
    """
    ads_atoms = [a.index for a in atoms if a.symbol in ads_syms]

    return ads_atoms

def connectivity2ads_index(atoms, species):
    """Return the indexes of atoms from the global list of adsorbate symbols.

    Parameters
    ----------
    atoms : object
        ASE atoms object with connectivity attached.
        This represents an adsorbate*slab structure.
    species : str
        chemical formula of the adsorbate.
    """
    composition = string2symbols(species)
    ads_atoms = []
    for symbol in composition:
        if (composition.count(symbol) ==
           atoms.get_chemical_symbols().count(symbol)):
            ads_atoms += [atom.index for atom in atoms if
                          atom.symbol == symbol]

    if len(ads_atoms) == len(composition):
        return ads_atoms
    elif len(ads_atoms) == 0:
        raise AssertionError("Formula adsorbate identification failed.")

    # If an atom in species also occurs in the slab, infer by connectivity.
    connected_atoms = []
    for atom in ads_atoms:
        edges = atoms.connectivity[atom, :]
        connected_atoms += [i for i, bonds in enumerate(edges) if
                            bonds > 0 and
                            atoms[i].symbol in composition and
                            atoms[i].symbol != atoms[atom].symbol]
    ads_atoms += connected_atoms

    # Final check.
    ua_ads, uc_ads = np.unique(
        np.array(atoms.get_chemical_symbols())[ads_atoms].sort(),
        return_counts=True)
    ua_comp, uc_comp = np.unique(composition.sort(), return_counts=True)
    if ua_ads != ua_comp:
        msg = str(ua_ads) + " != " + str(ua_comp)
        raise AssertionError(msg)
    elif uc_ads != uc_comp:
        msg = str(uc_ads) + " != " + str(uc_comp)
        raise AssertionError(msg)

    return list(np.unique(ads_atoms))

def slab_positions2ads_index(atoms, slab, species):
    """Return the indexes of adsorbate atoms identified by comparing positions
    to a reference slab structure.

    Parameters
    ----------
    atoms : object
    """
    composition = string2symbols(species)
    ads_atoms = []
    for symbol in composition:
        if (composition.count(symbol) ==
           atoms.get_chemical_symbols().count(symbol)):
            ads_atoms += [atom.index for atom in atoms if
                          atom.symbol == symbol]

    ua_ads, uc_ads = np.unique(ads_atoms, return_counts=True)
    ua_comp, uc_comp = np.unique(composition, return_counts=True)
    if ua_ads == ua_comp and uc_ads == uc_comp:
        return ads_atoms

    p_a = atoms.get_positions()
    p_r = slab.get_positions()

    for s in composition:
        if s in np.array(atoms.get_chemical_symbols())[ads_atoms]:
            continue
        symbol_count = composition.count(s)
        index_a = np.where(np.array(atoms.get_chemical_symbols()) == s)[0]
        index_r = np.where(np.array(slab.get_chemical_symbols()) == s)[0]
        _, dist = get_distances(p_a[index_a, :], p2=p_r[index_r, :],
                                cell=atoms.cell, pbc=True)
        # Assume all slab atoms are closest to their reference counterpart.
        deviations = np.min(dist, axis=1)
        # Sort deviations.
        ascending = np.argsort(deviations)
        # The highest deviations are assumed to be new atoms.
        ads_atoms += list(index_a[ascending[-symbol_count:]])

    # Final check.
    ua_ads, uc_ads = np.unique(
        np.array(atoms.get_chemical_symbols())[ads_atoms].sort(),
        return_counts=True)
    ua_comp, uc_comp = np.unique(composition.sort(), return_counts=True)
    if ua_ads != ua_comp:
        msg = str(ua_ads) + " != " + str(ua_comp)
        raise AssertionError(msg)
    elif uc_ads != uc_comp:
        msg = str(uc_ads) + " != " + str(uc_comp)
        raise AssertionError(msg)

    return ads_atoms

def tags2ads_index(atoms):
    """Return the indexes of atoms from the global list of adsorbate symbols.

    Parameters
    ----------
    atoms : object
        An ase atoms object. `atoms.tags` must label adsorbate atoms with 0 or
        negative numbers.
    """
    ads_atoms = [a.index for a in atoms if a.tag < 0]

    if 'key_value_pairs' in atoms.info:
        if 'species' in atoms.info['key_value_pairs']:
            species = atoms.info['key_value_pairs']['species']
            numbers = sorted(symbols2numbers(species))
            if numbers != sorted(atoms[ads_atoms].numbers):
                msg = 'Incorrect tags!'
                raise AssertionError(msg)
    return ads_atoms

def last2ads_index(atoms, species):
    """Return the indexes of the last n atoms in the atoms object, where n is
    the length of the composition of the adsorbate species. This function will
    work on atoms objects, where the slab was set up first,
    and the adsorbate was added after.

    Parameters
    ----------
    atoms : ase atoms object.
        atoms.info must be a dictionary containing the key 'key_value_pairs',
        which is expected to contain CATMAP standard adsorbate structure
        key value pairs. See the ase db to catmap module in catmap.
        the key value pair 'species' must be the
        chemical formula of the adsorbate.
    species : str
        chemical formula of the adsorbate.
    """
    n_ads = len(string2symbols(species))
    natoms = len(atoms)
    ads_atoms = list(range(natoms - n_ads, natoms))
    if sorted(symbols2numbers(species)) != sorted(atoms[ads_atoms].numbers):
        raise AssertionError("last index adsorbate identification failed.")
    warnings.warn("Adsorbate identified by last index.")
    return ads_atoms

def formula2ads_index(atoms, species):
    """Return the indexes of atoms, which have symbols matching the chemical
    formula of the adsorbate. This function will not work for adsorbates
    containing the same elements as the slab.

    Parameters
    ----------
    atoms : ase atoms object.
        atoms.info must be a dictionary containing the key 'key_value_pairs',
        which is expected to contain CatMAP standard adsorbate structure
        key value pairs. See the ase db to catmap module in catmap.
        the key value pair 'species' must be the
        chemical formula of the adsorbate.
    species : str
        chemical formula of the adsorbate.
    """
    try:
        composition = string2symbols(species)
    except ValueError:
        print(species)
        raise
    ads_atoms = [a.index for a in atoms if a.symbol in composition]
    if sorted(symbols2numbers(species)) != sorted(atoms[ads_atoms].numbers):
        raise AssertionError("ads atoms identification by formula failed.")
    return ads_atoms

def layers2ads_index(atoms, species):
    """ Returns the indexes of atoms in layers exceeding the number of layers
    stored in the key value pair 'layers'.

    Parameters
    ----------
    atoms : ase atoms object.
        atoms.info must be a dictionary containing the key 'key_value_pairs',
        which is expected to contain CatMAP standard adsorbate structure
        key value pairs. See the ase db to catmap module in catmap.
        the key value pair 'species' must be the
        chemical formula of the adsorbate and 'layers' must be an integer.
    species : str
        chemical formula of the adsorbate.
    """
    li, lz = auto_layers(atoms)
    layers = int(atoms.info['key_value_pairs']['layers'])
    ads_atoms = [a.index for a in atoms if li[a.index] > layers - 1]
    if sorted(symbols2numbers(species)) != sorted(atoms[ads_atoms].numbers):
        raise AssertionError("ads atoms identification by layers failed.")
    return ads_atoms

def z2ads_index(atoms, species):
    """ Returns the indexes of the n atoms with the highest position
    in the z direction,
    where n is the number of atoms in the chemical formula from the species
    key value pair.

    Parameters
    ----------
    atoms : ase atoms object.
        atoms.info must be a dictionary containing the key 'key_value_pairs',
        which is expected to contain CatMAP standard adsorbate structure
        key value pairs. See the ase db to catmap module in catmap.
        the key value pair 'species'.
    species : str
        chemical formula of the adsorbate.
    """
    composition = string2symbols(species)
    n_ads = len(composition)
    z = atoms.get_positions()[:, 2]
    ads_atoms = np.argsort(z)[-n_ads:]
    if sorted(symbols2numbers(species)) != sorted(atoms[ads_atoms].numbers):
        raise AssertionError("adsorbate identification by z coord failed. " +
                             str(species))
    return ads_atoms

def info2primary_index(atoms):
    """ Returns lists identifying the nearest neighbors of the adsorbate atoms.

    Parameters
    ----------
    atoms : ase atoms object.
        The atoms object must have the following keys in atoms.subsets:

        'ads_atoms' : list
            indices of atoms belonging to the adsorbate
        'slab_atoms' : list
            indices of atoms belonging to the slab
    """
    slab_atoms = atoms.subsets['slab_atoms']
    ads_atoms = atoms.subsets['ads_atoms']

    cm = atoms.connectivity
    chemi = []
    site = []
    ligand = []
    for a_a in ads_atoms:
        for a_s in slab_atoms:
            if cm[a_a, a_s] > 0:
                site.append(a_s)
                chemi.append(a_a)
    chemi = list(np.unique(chemi))
    site = list(np.unique(site))
    for j in site:
        for a_s in slab_atoms:
            if cm[a_s, j] > 0:
                ligand.append(a_s)
    ligand = [lig for lig in list(np.unique(ligand)) if lig not in site]
    if len(chemi) is 0 or len(site) is 0:
        print(chemi, site, ligand)
        msg = 'No adsorption site detected.'
        if 'id' in atoms.info:
            msg += ' id=' + str(atoms.info['id'])
        warnings.warn(msg)
    elif len(ligand) < 2:
        msg = 'Site has less than 2 nearest neighbors.'
        if 'id' in atoms.info:
            msg += ' id=' + str(atoms.info['id'])
        warnings.warn(msg)
    return chemi, site, ligand

def tags_termination(atoms):
    """Return lists bulk, term and subsurf containing atom indices belonging to
    those subsets of a surface atoms object. CatKit and ase.build contain
    functions that by default store this information in tags.

    Parameters
    ----------
    atoms : object
        the termination atoms should have tag=1 and subsequent layers should be
        tagged in increasing order.
    """
    term = [a.index for a in atoms if a.tag == 1]
    subsurf = [a.index for a in atoms if a.tag == 2]
    bulk = [a.index for a in atoms if a.tag > 1]
    return bulk, term, subsurf

def layers_termination(atoms, miller=(0, 0, 1)):
    """Return lists bulk, term and subsurf containing atom indices belonging to
    those subsets of a surface atoms object.
    This function relies on ase.atoms.get_layers, default atomic radii, and a
    slab oriented in the xy plane, where the termination in the z+ direction
    is the surface.

    Parameters
    ----------
    atoms : object
        The atoms object must have the following keys in atoms.subsets:

        'slab_atoms' : list
            indices of atoms belonging to the slab.
    """
    il, zl = auto_layers(atoms, miller=miller)
    il_slab = list(il)
    for index in sorted(atoms.subsets['ads_atoms'], reverse=True):
        del il_slab[index]
    nlayers = len(np.unique(il_slab))
    try:
        db_layers = atoms.info['key_value_pairs']['layers']
        if db_layers != nlayers:
            msg = str(db_layers) + ' != ' + str(nlayers)
            if 'id' in atoms.info:
                msg += ' id=' + str(atoms.info['id'])
            warnings.warn(msg)
    except(KeyError):
        if nlayers < 3:
            msg = 'nlayers = ' + str(nlayers)
            if 'id' in atoms.info:
                msg += ' id=' + str(atoms.info['id'])
            warnings.warn(msg)
    # bulk atoms can include subsurface atoms.
    term = [a.index for a in atoms if il[a.index] >= nlayers - 1 and
            a.index not in atoms.subsets['ads_atoms']]
    bulk = [a.index for a in atoms if il[a.index] <= nlayers - 2]
    subsurf = [a.index for a in atoms if il[a.index] == nlayers - 2]
    return bulk, term, subsurf

def constraints_termination(atoms):
    """Return lists bulk, term and subsurf containing atom indices belonging to
    those subsets of a surface atoms object.
    This function relies on the connectivity of the atoms and it assumes that
    bulk atoms are those that have are constrained in the first constraint.

    Parameters
    ----------
    atoms : object
        atoms.connectivity should be a connectivity matrix.
        The atoms object must have the following keys in atoms.subsets:

        'slab_atoms' : list
            indices of atoms belonging to the slab.
    """
    cm = atoms.connectivity.copy()
    np.fill_diagonal(cm, 0)
    # List coordination numbers of slab atoms.
    slab_atoms = atoms.subsets['slab_atoms']
    coord = np.count_nonzero(cm[slab_atoms, :][:, slab_atoms], axis=1)
    bulk = []
    term = []
    max_coord = np.max(coord)
    for j, c in enumerate(coord):
        a_s = atoms.subsets['slab_atoms'][j]
        if a_s in atoms.constraints[0].get_indices():
            bulk.append(a_s)
        elif c >= max_coord - 2:
            bulk.append(a_s)
        else:
            term.append(a_s)
    subsurf = []
    for a_b in bulk:
        for a_t in term:
            if cm[a_t, a_b] > 0:
                subsurf.append(a_b)
    subsurf = list(np.unique(subsurf))
    return bulk, term, subsurf

def connectivity_termination(atoms):
    """Return lists bulk, term and subsurf containing atom indices belonging to
    those subsets of a surface atoms object.
    This function relies on the connectivity of the atoms.

    Parameters
    ----------
    atoms : object
        atoms.connectivity should be a connectivity matrix.
        The atoms object must have the following keys in atoms.subsets:

        'slab_atoms' : list
            indices of atoms belonging to the slab
    """
    raise NotImplementedError("Todo: identify and remove backside termination")
    cm = atoms.connectivity.copy()
    np.fill_diagonal(cm, 0)
    # List coordination numbers of slab atoms.
    slab_atoms = atoms.subsets['slab_atoms']
    coord = np.count_nonzero(cm[slab_atoms, :][:, slab_atoms], axis=1)
    bulk = []
    term = []
    max_coord = np.max(coord)
    bulk = atoms.subsets['slab_atoms']
    for j, c in enumerate(coord):
        a_s = atoms.subsets['slab_atoms'][j]
        if c >= max_coord - 2:
            bulk.append(a_s)
        else:
            term.append(a_s)
    subsurf = []
    for a_b in bulk:
        for a_t in term:
            if cm[a_t, a_b] > 0:
                subsurf.append(a_b)
    subsurf = list(np.unique(subsurf))
    return bulk, term, subsurf

def auto_layers(atoms, miller=(0, 0, 1)):
    """Returns two arrays describing which layer each atom belongs
    to and the distance between the layers and origo.
    Assumes the tolerance corresponds to the average atomic radii of the slab.

    Parameters
    ----------
    atoms : object
        The atoms object must have the following keys in atoms.subsets:

        'slab_atoms' : list
            indices of atoms belonging to the slab
    """
    radii = [get_radius(z) for z in atoms.numbers[atoms.subsets['slab_atoms']]]
    radius = np.average(radii) / 2.
    li, lz = get_layers(atoms, miller=miller, tolerance=radius)
    return li, lz

def attach_cations(atoms, anion_number=8):
    """Attaches list of cation and anion atomic indices.

    Parameters
    ----------
    atoms : object
        ase.Atoms object.
    anion_number : int
        Atomic number of the anion of this chalcogenide.
    """
    if anion_number not in atoms.numbers:
        raise ValueError('Anion ' + chemical_symbols[anion_number] +
                         'not in atoms')
    atoms.subsets['cation_atoms'] = [a.index for a in atoms if
                                     a.number != anion_number]
    atoms.subsets['anion_atoms'] = [a.index for a in atoms if
                                    a.number == anion_number]


#!/usr/bin/env python
# coding: utf-8

# # YOLO v3 Object Detection
#
# Let's see how to use the state of the art in object detection! Please make sure to watch the video, there is no code along here, since we can't reasonably train the YOLOv3 network ourself, instead we will use a pre-established version.
#
# CODE SOURCE: https://github.com/xiaochus/YOLOv3
#
# REFERENCE (for original YOLOv3):
#
#         @article{YOLOv3,
#               title={YOLOv3: An Incremental Improvement},
#               author={J Redmon, A Farhadi },
#               year={2018}
# --------
# ----------
# ## YOU MUST WATCH THE VIDEO LECTURE TO PROPERLY SET UP THE MODEL AND WEIGHTS. THIS NOTEBOOK WON'T WORK UNLESS YOU FOLLOW THE EXACT SET UP SHOWN IN THE VIDEO LECTURE.
# -------
# -------

# In[3]:

import os
import time
import cv2
import numpy as np
from model.yolo_model import YOLO

# In[4]:

def process_image(img):
    """Resize, reduce and expand image.

    # Argument:
        img: original image.

    # Returns
        image: ndarray(64, 64, 3), processed image.
    """
    image = cv2.resize(img, (416, 416),
                       interpolation=cv2.INTER_CUBIC)
    image = np.array(image, dtype='float32')
    image /= 255.
    image = np.expand_dims(image, axis=0)

    return image

# In[5]:

def get_classes(file):
    """Get classes name.

    # Argument:
        file: classes name for database.

    # Returns
        class_names: List, classes name.

    """
    with open(file) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]

    return class_names

# In[6]:

def draw(image, boxes, scores, classes, all_classes):
    """Draw the boxes on the image.

    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    for box, score, cl in zip(boxes, scores, classes):
        x, y, w, h = box

        top = max(0, np.floor(x + 0.5).astype(int))
        left = max(0, np.floor(y + 0.5).astype(int))
        right = min(image.shape[1], np.floor(x + w + 0.5).astype(int))
        bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int))

        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(all_classes[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    2 ,(0, 0, 255), 3,
                    cv2.LINE_AA)

       # print('class: {0}, score: {1:.2f}'.format(all_classes[cl], score))
        #print('box coordinate x,y,w,h: {0}'.format(box))

    print()

# In[7]:

# function to find if
# given point lies inside
# a given rectangle or not.
def FindPoint2(left, bottom, right, top, x, y,w,h) :
    px1, py1, px2, py2 = left, top, right, bottom
    hx1, hy1, hx2, hy2 = x,y,x+w,y+h
    if (hx1 >= px1 and hy1 >= py1) or (hx2 <= px2 and hy2 <=py2):
        print("True!")
        return True
    else:
        return False

def FindPoint(left, bottom, right, top, x, y,w,h) :
    try:
        cat_x=int((2*x+w)/2)
        cat_y=int((2*y+h)/2)
        if cat_x>left and cat_y<top and cat_x<right and cat_y>bottom:
            print("True!")
            return True
        else:
            return False
    except:
        print("something wrong")
        return False

def only_if_max(curr_val, new_val):

    if curr_val==None or new_val > curr_val :

        return new_val
    return curr_val

def only_if_min(curr_val, new_val):
    if curr_val==None or new_val < curr_val:
        return new_val
    return curr_val

# In[8]:

# In[22]:

def draw2(image, boxes, scores, classes, all_classes, bad_guy, forbbiden_place):
    """Draw the boxes on the image.

    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    cat_exsists=0
    toilet_exists=0
    toilet_left=None
    toilet_right=None
    toilet_bottom=None
    toilet_top=None
    tags=[]
    for box, score, cl in zip(boxes, scores, classes):
        x, y, w, h = box

        top = max(0, np.floor(x + 0.5).astype(int))
        left = max(0, np.floor(y + 0.5).astype(int))
        right = min(image.shape[1], np.floor(x + w + 0.5).astype(int))
        bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int))
        tags.append(all_classes[cl])
        if all_classes[cl]==bad_guy:
            cv2.rectangle(image, (top, left), (right, bottom), (0, 0, 255), 2)
            cv2.putText(image, '{0} {1:.2f}'.format(all_classes[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1 ,(0, 0, 255), 1,
                    cv2.LINE_AA)

        if all_classes[cl]==forbbiden_place:
            toilet_exists=1
            toilet_exists=0
            toilet_left=only_if_min(toilet_left, left)
            toilet_right=only_if_max(toilet_right,right)
            toilet_bottom=only_if_max(toilet_bottom, bottom)
            toilet_top=only_if_min(toilet_top, top)
            text_hieght=  int((toilet_bottom+toilet_top)/2)
            cat_x=int((2*x+w)/2)
            cat_y=int((2*y+h)/2)
            #print("toilet")

        if toilet_left: #prevent nonetype
            cv2.rectangle(image, (toilet_top, toilet_left), (toilet_right, toilet_bottom), (200, 0, 0), 2)
        if forbidden_place and bad_guy in tags:
            if bad_guy in tags:
                if FindPoint(toilet_left,toilet_top,toilet_right,toilet_bottom, x,y,w,h):
                    #image = cv2.circle(image, (cat_y,cat_x) , 2, (0, 0, 100), 2 )
                    text = str(str(bad_guy).upper()+" IS ON THE "+str(forbbiden_place).upper())
                    print(text)
                    #print("CAT IN THE TOILET")
                    image = cv2.circle(image, (cat_x,cat_y) , 4, (0, 0, 255), 4 )
                    text_y=text_hieght
                    text_x=int((cat_x/2))
                    # font
                    font = cv2.FONT_HERSHEY_SIMPLEX

                    # org
                    org = (text_x-100, text_y )

                    # fontScale
                    fontScale = 1

                    # Blue color in BGR
                    color = (50, 0, 255)

                    # Line thickness of 2 px
                    thickness = 2

                    # Using cv2.putText() method

                    y0, dy = text_y, 150
                    #for i, line in enumerate(text.split('\n')):
                    #        y = y0 + i*dy
                    #        cv2.putText(image, line, (text_x, y ), cv2.FONT_HERSHEY_SIMPLEX, fontScale, color, thickness, cv2.LINE_AA)
                    image = cv2.putText(image,  text, org, font, fontScale, color, thickness, cv2.LINE_AA)

                    #cv2.putText(image,'{0} {1:.2f}'.format( "CAT IN THE TOILET"),(bottom, left +6), cv2.FONT_HERSHEY_SIMPLEX, 10 ,(0, 0, 255), 8, cv2.LINE_AA)

        #print('class: {0}, score: {1:.2f}'.format(all_classes[cl], score))
        #print('box coordinate x,y,w,h: {0}'.format(box))

# In[17]:

def detect_image(image, yolo, all_classes, bad_guy, forbbiden_place):
    """Use yolo v3 to detect images.

    # Argument:
        image: original image.
        yolo: YOLO, yolo model.
        all_classes: all classes name.

    # Returns:
        image: processed image.
    """
    pimage = process_image(image)

    start = time.time()
    boxes, classes, scores = yolo.predict(pimage, image.shape)
    end = time.time()

    print('frame process time: {0:.2f}s'.format(end - start))

    if boxes is not None:
        draw2(image, boxes, scores, classes, all_classes, bad_guy, forbbiden_place)

    return image

# In[18]:

def detect_video(video, yolo, all_classes, bad_guy, forbbiden_place):
    """Use yolo v3 to detect video.

    # Argument:
        video: video file.
        yolo: YOLO, yolo model.
        all_classes: all classes name.
    """
    print(video)
    # FILE SOURCE
    #video_path = os.path.join("videos", "test", video)
    #camera = cv2.VideoCapture(video_path)
    #cv2.namedWindow("detection", cv2.WINDOW_AUTOSIZE)

    # STREAM source
    camera = cv2.VideoCapture(0)

    # Prepare for saving the detected video
    sz = (int(camera.get(cv2.CAP_PROP_FRAME_WIDTH)), int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fourcc = cv2.VideoWriter_fourcc(*'mpeg')
    fps=20

    vout = cv2.VideoWriter()
    vout.open(os.path.join("videos", "res", video), fourcc, fps, sz, True)

    while True:
        res, frame = camera.read()

        if not res:
            break

        image = detect_image(frame, yolo, all_classes, bad_guy, forbbiden_place)
        #cv2.imshow("detection", image)

        # Save the video frame by frame
        vout.write(image)

        if cv2.waitKey(1) & 0xff == 27:
                break

    vout.release()
    camera.release()
    print("done")

# In[19]:

yolo = YOLO(0.25, 0.5)
file = 'data/coco_classes.txt'
all_classes = get_classes(file)

# ### Detecting Images

# In[35]:

# DEFINE MISSION:
# choose objects from classes list: data/coco_classes.txt
bad_guy="cat"
forbidden_place="toilet"

with open("data/coco_classes.txt") as f:
    if bad_guy and forbidden_place in f.read():
        print("Objects are valid")
    else:
        print("CHANGE OBJECT NAMES")

# In[20]:

f = 'oscar_toilet.jpg'
path = 'images/'+f
image = cv2.imread(path)
image = detect_image(image, yolo, all_classes,bad_guy, forbidden_place)
cv2.imwrite('images/res/' + f, image)

# # Detecting on Video

# In[36]:

# # detect videos one at a time in videos/test folder

video = 'cat_flushing2.mp4'
detect_video(video, yolo, all_classes, bad_guy, forbidden_place)

# In[ ]:

# In[ ]: