# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

from __future__ import print_function

import numpy as np

from ase.lattice.cubic import Diamond, FaceCenteredCubic, SimpleCubic

###

def set_groups(a, n, skin_x, skin_y, central_x=-1./2, central_y=-1./2,
               invert_central=False):
    nx, ny, nz = n
    sx, sy, sz = a.cell.diagonal()
    print('skin_x = {0}*a0, skin_y = {1}*a0'.format(skin_x, skin_y))
    skin_x = skin_x*sx/nx
    skin_y = skin_y*sy/ny
    print('skin_x = {0}, skin_y = {1}'.format(skin_x, skin_y))
    r = a.positions

    g = np.ones(len(a), dtype=int)
    mask = np.logical_or(
               np.logical_or(
                   np.logical_or(
                       r[:, 0]/sx < (1.-central_x)/2,
                       r[:, 0]/sx > (1.+central_x)/2),
                   r[:, 1]/sy < (1.-central_y)/2),
               r[:, 1]/sy > (1.+central_y)/2)
    if invert_central:
        mask = np.logical_not(mask)
    g = np.where(mask, g, 2*np.ones_like(g))

    mask = np.logical_or(
               np.logical_or(
                   np.logical_or(
                       r[:, 0] < skin_x, r[:, 0] > sx-skin_x),
                   r[:, 1] < skin_y),
               r[:, 1] > sy-skin_y)
    g = np.where(mask, np.zeros_like(g), g)

    a.set_array('groups', g)


def cluster(el, a0, n, crack_surface=[1,1,0], crack_front=[0,0,1],
            lattice=None, shift=None):
    nx, ny, nz = n
    third_dir = np.cross(crack_surface, crack_front)
    directions = [ third_dir, crack_surface, crack_front ]
    if np.linalg.det(directions) < 0:
        third_dir = -third_dir
    directions = [ third_dir, crack_surface, crack_front ]
    unitcell = lattice(el, latticeconstant=a0, size=[1, 1, 1], 
                       directions=directions  )
    if shift is not None:
        unitcell.translate(np.dot(shift, unitcell.cell))

    # Center cluster in unit cell
    x, y, z = (unitcell.get_scaled_positions()%1.0).T
    x += (1.0-x.max()+x.min())/2 - x.min()
    y += (1.0-y.max()+y.min())/2 - y.min()
    z += (1.0-z.max()+z.min())/2 - z.min()
    unitcell.set_scaled_positions(np.transpose([x, y, z]))

    a = unitcell.copy()
    a *= (nx, ny, nz)
    #a.info['unitcell'] = unitcell

    a.set_pbc([False, False, True])

    return a


def diamond(*args, **kwargs):
    kwargs['lattice'] = Diamond
    return cluster(*args, **kwargs)


def fcc(*args, **kwargs):
    kwargs['lattice'] = FaceCenteredCubic
    return cluster(*args, **kwargs)


def sc(*args, **kwargs):
    kwargs['lattice'] = SimpleCubic
    return cluster(*args, **kwargs)
