from math import *

import numpy as np


# Grid generator for direct and reciprocal space
class UnitCell:
    def __init__(self, N, basis):
        self.N = N
        self.d = basis / N # displacement basis of grid points


    # Generate NxNxN grid to represent functions inside the cell
    def grid_points(self):
        gx = np.zeros((self.N, self.N, self.N))
        gy = np.zeros_like(gx)
        gz = np.zeros_like(gx)

        if self.N % 2 == 0:
            ns = range(-self.N // 2, self.N // 2)
        else:
            ns = range(-(self.N - 1) // 2, (self.N - 1) // 2 + 1)

        for i in range(self.N):
            for j in range(self.N):
                for k in range(self.N):
                    v = ns[i] * self.d[0] + ns[j] * self.d[1] + ns[k] * self.d[2]
                    gx[i, j, k] = v[0]
                    gy[i, j, k] = v[1]
                    gz[i, j, k] = v[2]

        return gx, gy, gz


    # Volume of smallest parallelepiped representable on this grid
    def differential_volume(self):
        dV = np.dot(self.d[0], np.cross(self.d[1], self.d[2]))
        return abs(dV)



# Base class for periodic structures, don't use it directly
class Crystal:
    def __init__(self, N, basis):
        self.a = basis
        self.b = self.get_reciprocal_basis()
        self.cell_direct  = UnitCell(N, self.a)
        self.cell_fourier = UnitCell(N, self.b * N)


    def get_reciprocal_basis(self):
        V = np.dot(self.a[0], np.cross(self.a[1], self.a[2]))
        b1 = (2 * pi / V) * np.cross(self.a[1], self.a[2])
        b2 = (2 * pi / V) * np.cross(self.a[2], self.a[0])
        b3 = (2 * pi / V) * np.cross(self.a[0], self.a[1])
        return np.array([b1, b2, b3])


    def get_grid_direct(self):
        rx, ry, rz = self.cell_direct.grid_points()
        dr = self.cell_direct.differential_volume()
        return rx, ry, rz, dr


    def get_grid_fourier(self):
        qx, qy, qz = self.cell_fourier.grid_points()
        dq = self.cell_fourier.differential_volume()
        return qx, qy, qz, dq


    # Convert sequence of letters to coordinates for band structure plots.
    # The letters' meanings are stored in a dictionary self.bz_points.
    def get_brillouin_zone_path(self, chars):
        path = []
        for c in chars:
            k = self.bz_points[c] # should probably catch KeyError here
            p = k[0] * self.b[0] + k[1] * self.b[1] + k[2] * self.b[2]
            path.append(p)

        return path



# Derived class for simple cubic (sc) crystals
class CubicSimple(Crystal):
    def __init__(self, N, a):
        a1 = np.array((a, 0, 0))
        a2 = np.array((0, a, 0))
        a3 = np.array((0, 0, a))
        super().__init__(N, np.array([a1, a2, a3]))

    # Taken from http://lampx.tugraz.at/~hadley/ss1/bzones/sc.php
    bz_points = {
        "G" : (0,   0,   0),
        "R" : (0.5, 0.5, 0.5),
        "X" : (0,   0.5, 0),
        "M" : (0.5, 0.5, 0)
    }



# Derived class for body-centered cubic (bcc) crystals
class CubicBodyCentered(Crystal):
    def __init__(self, N, a):
        a1 = np.array((-a/2,  a/2,  a/2))
        a2 = np.array(( a/2, -a/2,  a/2))
        a3 = np.array(( a/2,  a/2, -a/2))
        super().__init__(N, np.array([a1, a2, a3]))

    # Taken from http://lampx.tugraz.at/~hadley/ss1/bzones/bcc.php
    bz_points = {
        "G" : ( 0,    0,    0),
        "H" : (-0.5,  0.5,  0.5),
        "P" : ( 0.25, 0.25, 0.25),
        "N" : ( 0,    0.5,  0)
    }



# Derived class for face-centered cubic (fcc) crystals
class CubicFaceCentered(Crystal):
    def __init__(self, N, a):
        a1 = np.array((0,   a/2, a/2))
        a2 = np.array((a/2,   0, a/2))
        a3 = np.array((a/2, a/2,   0))
        super().__init__(N, np.array([a1, a2, a3]))

    # Taken from http://lampx.tugraz.at/~hadley/ss1/bzones/fcc.php
    bz_points = {
        "G" : (0,     0,     0),
        "X" : (0,     0.5,   0.5),
        "L" : (0.5,   0.5,   0.5),
        "W" : (0.25,  0.75,  0.5),
        "U" : (0.25,  0.625, 0.625),
        "K" : (0.375, 0.75,  0.375)
    }