"""
@package projects.common.clusters.crystals
cluster generators for some common bulk crystals

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-19 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

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

import math
import numpy as np
import os.path
import periodictable as pt
import logging

import pmsco.cluster as cluster
import pmsco.dispatch as dispatch
import pmsco.project as project
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)


class ZincblendeCluster(cluster.ClusterGenerator):
    def __init__(self, proj):
        super(ZincblendeCluster, self).__init__(proj)
        self.atomtype1 = 30
        self.atomtype2 = 16
        self.bulk_lattice = 1.0
        self.surface = (1, 1, 1)

    @classmethod
    def check(cls, outfilename=None, model_dict=None, domain_dict=None):
        """
        function to test and debug the cluster generator.

        to use this function, you don't need to import or initialize anything but the class.
        though the project class is used internally, the result does not depend on any project settings.

        @param outfilename: name of output file for the cluster (XYZ format).
            the file is written to the same directory where this module is located.
            if empty or None, no file is written.

        @param model_dict: dictionary of model parameters to override the default values.

        @param domain_dict: dictionary of domain parameters to override the default values.

        @return: @ref pmsco.cluster.Cluster object
        """
        proj = project.Project()
        dom = project.ModelSpace()
        dom.add_param('dlat', 10.)
        dom.add_param('rmax', 5.0)
        if model_dict:
            dom.start.update(model_dict)

        try:
            proj.domains[0].update({'zrot': 0.})
        except IndexError:
            proj.add_domain({'zrot': 0.})
        if domain_dict:
            proj.domains[0].update(domain_dict)
        proj.add_scan("", 'C', '1s')

        clu_gen = cls(proj)
        index = dispatch.CalcID(0, 0, 0, -1, -1)
        clu = clu_gen.create_cluster(dom.start, index)

        if outfilename:
            project_dir = os.path.dirname(os.path.abspath(__file__))
            outfilepath = os.path.join(project_dir, outfilename)
            clu.save_to_file(outfilepath, fmt=cluster.FMT_XYZ, comment="{0} {1} {2}".format(cls, index, str(dom.start)))

        return clu

    def count_emitters(self, model, index):
        return 1

    def create_cluster(self, model, index):
        """
        calculate a specific set of atom positions given the optimizable parameters.

        @param model  (dict)          optimizable parameters
            @arg    model['dlat']     bulk lattice constant in Angstrom
            @arg    model['rmax']     cluster radius
            @arg    model['phi']      azimuthal rotation angle in degrees

        @param dom (dict)             domain
            @arg    dom['term']       surface termination
        """
        clu = cluster.Cluster()
        clu.comment = "{0} {1}".format(self.__class__, index)
        clu.set_rmax(model['rmax'])
        a_lat = model['dlat']
        dom = self.project.domains[index]
        try:
            term = int(dom['term'])
        except ValueError:
            term = pt.elements.symbol(dom['term'].strip().number)

        if self.surface == (0, 0, 1):
            # identity matrix
            m = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
        elif self.surface == (1, 1, 1):
            # this will map the [111] direction onto the z-axis
            m1 = np.array([1, -1, 0]) * math.sqrt(1/2)
            m2 = np.array([0.5, 0.5, -1]) * math.sqrt(2/3)
            m3 = np.array([1, 1, 1]) * math.sqrt(1/3)
            m = np.array([m1, m2, m3])
        else:
            raise ValueError("unsupported surface specification")

        # lattice vectors
        a1 = np.matmul(m, np.array((1.0, 0.0, 0.0)) * a_lat)
        a2 = np.matmul(m, np.array((0.0, 1.0, 0.0)) * a_lat)
        a3 = np.matmul(m, np.array((0.0, 0.0, 1.0)) * a_lat)

        # basis
        b1 = [np.array((0.0, 0.0, 0.0)), (a2 + a3) / 2, (a3 + a1) / 2, (a1 + a2) / 2]
        if term == self.atomtype1:
            d1 = np.array((0, 0, 0))
            d2 = (a1 + a2 + a3) / 4
        else:
            d1 = -(a1 + a2 + a3) / 4
            d2 = np.array((0, 0, 0))
        for b in b1:
            clu.add_bulk(self.atomtype1, b + d1, a1, a2, a3)
            clu.add_bulk(self.atomtype2, b + d2, a1, a2, a3)

        return clu