import numpy as np


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


class SCC_Lattice(Lattice):
    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


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_b_r = 4.554
    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)

        print("A_r: ", offset_a_r, "C_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

    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