import matplotlib.pyplot as plt
import numpy as np
from cache import timeit
from abc import ABC, abstractmethod


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


class Lattice(ABC):

    @abstractmethod
    def get_phase(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        pass

    @abstractmethod
    def parse_mask(self, mask: np.ndarray) -> np.ndarray:
        pass

    @abstractmethod
    def get_phases(self) -> list[tuple[np.ndarray, np.ndarray]]:
        pass

    @abstractmethod
    def get_spots(self) -> list[tuple[np.ndarray, np.ndarray]]:
        pass


class SCC_Lattice(Lattice):

    X: np.ndarray
    Y: np.ndarray

    def __init__(self, x_len: int, y_len: int):
        x = np.arange(x_len) * 5
        y = np.arange(x_len) * 4.5
        self.X, self.Y = np.meshgrid(x, y)

    def get_phase(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        return (self.X, self.Y)

    def get_phases(self) -> list[tuple[np.ndarray, np.ndarray]]:
        return [self.get_phase(0)]

    def get_spots(self) -> list[tuple[np.ndarray, np.ndarray]]:
        x = np.arange(-3, 4) * 0.2
        y = np.arange(-3, 4) * 1./4.5
        X, Y = np.meshgrid(x, y)
        return [(X.flatten(), Y.flatten())]

    def parse_mask(self, mask: np.ndarray) -> np.ndarray:
        return mask


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 __init__(self, x_len: int, y_len: int):
        self.X, self.Y = self._generate_vec(x_len * 2, y_len * 2)

    def parse_mask(self, mask: np.ndarray) -> np.ndarray:
        maske = np.empty((mask.shape[0]*2, mask.shape[1]*2))
        maske[0::2, 0::2] = mask
        maske[1::2, 0::2] = mask
        maske[0::2, 1::2] = mask
        maske[1::2, 1::2] = mask
        return maske

    def get_phase(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        if index == 0:
            return self._get_rutile()
        else:
            return self._get_mono()

    def get_phases(self) -> list[tuple[np.ndarray, np.ndarray]]:
        return [self.get_phase(0), self.get_phase(1)]

    def get_spots(self) -> list[tuple[np.ndarray, np.ndarray]]:
        p1 = self._reci_rutile()
        p2 = self._reci_mono()
        p3 = self._reci_mono_2()
        return [p1, p2, p3]
        pass

    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 = self.X * self.base_c_r + \
            np.mod(self.Y, 4) * 0.5 * self.base_c_r
        y = self.Y * 0.5 * self.base_a_r
        return (x, y)

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

        # offset_c_m = -offset_c_m
        # offset_a_m = -offset_a_m

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

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

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

    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 _reci_rutile(self):
        len_x = np.max(self.X)
        len_y = np.max(self.Y)
        x = np.arange(-len_x, len_x+1)
        y = np.arange(-len_y, len_y+1)
        X, Y = np.meshgrid(x, y)
        x = X * 1./self.base_c_r
        y = Y * 2 * 1./self.base_a_r + np.mod(X, 2) * 1./self.base_a_r
        return x.flatten(), y.flatten()

    def _reci_mono(self):
        x, y = self._reci_rutile()
        return x + 0.5 * 1. / self.base_c_r, y + 0.5 * 1./self.base_a_r

    def _reci_mono_2(self):
        x, y = self._reci_rutile()
        return x - 0.5 * 1. / self.base_c_r, y + 0.5 * 1./self.base_a_r


class VO2_New(VO2_Lattice):
    # def parse_mask(self, mask: np.ndarray) -> np.ndarray:
    #    maske = np.empty((mask.shape[0]*2, mask.shape[1]*2))
    #    maske[0::2, 0::2] = mask
    #    maske[1::2, 0::2] = mask
    #    maske[0::2, 1::2] = mask
    #    maske[1::2, 1::2] = mask
    #    maske[0::4, :] = np.roll(maske[0::4, :], axis=1, shift=1)
    #    maske[1::4, :] = np.roll(maske[1::4, :], axis=1, shift=1)
    #    return maske
    # def parse_mask(self, mask: np.ndarray) -> np.ndarray:
    #    print(mask.shape)
    #    maske = np.empty((mask.shape[0]*2, mask.shape[1]))
    #    maske[0::2, :] = mask
    #    maske[1::2, :] = mask
    #    print(maske.shape)
    #    return maske
    def parse_mask(self, mask: np.ndarray):
        return mask