import numpy as np
from cache import timeit


def deg_2_rad(winkel):
    return winkel / 180.0 * np.pi


# all units in angstrom


class Lattice:
    def __init__(self, x_len, y_len):
        pass

    def get_from_mask(self, maske):
        pass

    def get_both(self):
        pass

    def reci(self):
        pass


class SCC_Lattice(Lattice):
    @timeit
    def __init__(self, x_len, y_len):
        x = np.arange(x_len) * 5
        y = np.arange(x_len) * 5
        self.X, self.Y = np.meshgrid(x, y)

    def get_from_mask(self, maske):
        return self.X, self.Y

    def get_both(self):
        return [(self.X, self.Y), (self.X, self.Y)]

    def reci(self):
        x = np.arange(-3, 4) * 0.2
        y = np.arange(-3, 4) * 0.2
        X, Y = np.meshgrid(x, y)
        return [(X, Y)]


class VO2_Lattice(Lattice):
    base_a_m = 5.75
    base_b_m = 4.5
    base_c_m = 5.38

    base_c_r = 2.856
    base_c_r = 2.8
    base_b_r = 4.554
    base_b_r = 4.5
    base_a_r = base_b_r

    alpha_m = 122.64  # degree

    def _mono_2_rutile(self, c_m, a_m):
        a_r = np.cos(deg_2_rad(self.alpha_m - 90)) * c_m * self.base_c_m
        c_r = (a_m) * self.base_a_m + \
            np.sin(deg_2_rad(self.alpha_m - 90)) * c_m * self.base_c_m
        return a_r, c_r

    def _get_rutile(self, X, Y):
        self.atom_x_rut = X * self.base_c_r + \
            np.mod(Y, 4) * 0.5 * self.base_c_r
        self.atom_y_rut = Y * 0.5 * self.base_a_r

    def _get_mono(self, X, Y):
        offset_a_m = 0.25 - 0.23947
        offset_c_m = 0.02646

        offset_a_r, offset_c_r = self._mono_2_rutile(offset_c_m, offset_a_m)

        res = 0.05
        offset_a_r = res * int(offset_a_r/res)
        offset_c_r = res * int(offset_c_r/res)

        print(offset_a_r, offset_c_r)

        self.atom_x_mono = offset_a_r + X * \
            self.base_c_r + np.mod(Y, 4) * 0.5 * self.base_c_r
        self.atom_x_mono[np.mod(X, 2) == 0] -= 2 * offset_a_r

        self.atom_y_mono = offset_c_r + 0.5 * Y * self.base_a_r
        self.atom_y_mono[np.mod(X, 2) == 0] -= 2 * offset_c_r

    def _generate_vec(self, x_len: int, y_len: int):
        x = np.arange(x_len)
        y = np.arange(y_len)
        X, Y = np.meshgrid(x, y)
        X[np.mod(Y, 4) == 3] = X[np.mod(Y, 4) == 3] - 1
        X[np.mod(Y, 4) == 2] = X[np.mod(Y, 4) == 2] - 1
        assert np.mod(x.size, 2) == 0
        assert np.mod(y.size, 2) == 0

        return X, Y

    @timeit
    def __init__(self, x_len: int, y_len: int):
        X, Y = self._generate_vec(x_len * 2, y_len * 2)
        self._get_mono(X, Y)
        self._get_rutile(X, Y)

    def get_from_mask(self, maske: np.array):
        inplace_pos_x = np.zeros_like(self.atom_x_mono)
        inplace_pos_y = np.zeros_like(self.atom_x_mono)

        mask = np.empty_like(self.atom_x_mono, dtype=bool)
        mask[0::2, 0::2] = maske
        mask[1::2, 0::2] = maske
        mask[0::2, 1::2] = maske
        mask[1::2, 1::2] = maske

        inplace_pos_x[mask] = self.atom_x_rut[mask]
        inplace_pos_y[mask] = self.atom_y_rut[mask]
        mask = np.invert(mask)
        inplace_pos_x[mask] = self.atom_x_mono[mask]
        inplace_pos_y[mask] = self.atom_y_mono[mask]
        return inplace_pos_x, inplace_pos_y

    def get_both(self):
        return [(self.atom_x_rut, self.atom_y_rut), (self.atom_x_mono, self.atom_y_mono)]

    def reci_rutile(self):
        num = 20
        x = np.arange(-num, num + 1)
        y = np.arange(-num, num + 1)
        X, Y = np.meshgrid(x, y)
        return (X * 0.22 + Y * 0.44).flatten(), (X * 0.349).flatten()

    def reci_mono(self):
        x, y = self.reci_rutile()
        return x + 0.1083, y + 0.1719

    def reci_mono_2(self):
        x, y = self.reci_rutile()
        return x - 0.1083, y + 0.1719

    def reci(self):
        cutoff = 5.
        x, y = self.reci_rutile()
        mask = np.logical_and(np.abs(x) < cutoff, np.abs(y) < cutoff)
        p1 = (x[mask], y[mask])
        x, y = self.reci_mono()
        mask = np.logical_and(np.abs(x) < cutoff, np.abs(y) < cutoff)
        p2 = (x[mask], y[mask])
        x, y = self.reci_mono_2()
        mask = np.logical_and(np.abs(x) < cutoff, np.abs(y) < cutoff)
        p3 = (x[mask], y[mask])
        return [p1, p2, p3]