import numpy as np
import PIL.Image
import scipy
import torch
import torch.nn.functional as F

from tada.augment_pipes.edm import (matrix, rotate2d_inv, rotate3d, scale2d,
                                    scale2d_inv, scale3d, translate2d,
                                    translate2d_inv, translate3d)
from tada.augment_pipes.utils import (AVAILABLE_AUGMENTATIONS, GEOMETRIC, COLOR,
                                      EDM_AUGMENTATIONS, TADA_AUGMENTATIONS)

_constant_cache = dict()


def constant(value, shape=None, dtype=None, device=None, memory_format=None):
    value = np.asarray(value)
    if shape is not None:
        shape = tuple(shape)
    if dtype is None:
        dtype = torch.get_default_dtype()
    if device is None:
        device = torch.device('cpu')
    if memory_format is None:
        memory_format = torch.contiguous_format

    key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)
    tensor = _constant_cache.get(key, None)
    if tensor is None:
        tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)
        if shape is not None:
            tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))
        tensor = tensor.contiguous(memory_format=memory_format)
        _constant_cache[key] = tensor
    return tensor


#----------------------------------------------------------------------------
# Coefficients of various wavelet decomposition low-pass filters.
wavelets = {
    'haar': [0.7071067811865476, 0.7071067811865476],
    'db1':  [0.7071067811865476, 0.7071067811865476],
    'db2':  [-0.12940952255092145, 0.22414386804185735, 0.836516303737469, 0.48296291314469025],
    'db3':  [0.035226291882100656, -0.08544127388224149, -0.13501102001039084, 0.4598775021193313, 0.8068915093133388, 0.3326705529509569],
    'db4':  [-0.010597401784997278, 0.032883011666982945, 0.030841381835986965, -0.18703481171888114, -0.02798376941698385, 0.6308807679295904, 0.7148465705525415, 0.23037781330885523],
    'db5':  [0.003335725285001549, -0.012580751999015526, -0.006241490213011705, 0.07757149384006515, -0.03224486958502952, -0.24229488706619015, 0.13842814590110342, 0.7243085284385744, 0.6038292697974729, 0.160102397974125],
    'db6':  [-0.00107730108499558, 0.004777257511010651, 0.0005538422009938016, -0.031582039318031156, 0.02752286553001629, 0.09750160558707936, -0.12976686756709563, -0.22626469396516913, 0.3152503517092432, 0.7511339080215775, 0.4946238903983854, 0.11154074335008017],
    'db7':  [0.0003537138000010399, -0.0018016407039998328, 0.00042957797300470274, 0.012550998556013784, -0.01657454163101562, -0.03802993693503463, 0.0806126091510659, 0.07130921926705004, -0.22403618499416572, -0.14390600392910627, 0.4697822874053586, 0.7291320908465551, 0.39653931948230575, 0.07785205408506236],
    'db8':  [-0.00011747678400228192, 0.0006754494059985568, -0.0003917403729959771, -0.00487035299301066, 0.008746094047015655, 0.013981027917015516, -0.04408825393106472, -0.01736930100202211, 0.128747426620186, 0.00047248457399797254, -0.2840155429624281, -0.015829105256023893, 0.5853546836548691, 0.6756307362980128, 0.3128715909144659, 0.05441584224308161],
    'sym2': [-0.12940952255092145, 0.22414386804185735, 0.836516303737469, 0.48296291314469025],
    'sym3': [0.035226291882100656, -0.08544127388224149, -0.13501102001039084, 0.4598775021193313, 0.8068915093133388, 0.3326705529509569],
    'sym4': [-0.07576571478927333, -0.02963552764599851, 0.49761866763201545, 0.8037387518059161, 0.29785779560527736, -0.09921954357684722, -0.012603967262037833, 0.0322231006040427],
    'sym5': [0.027333068345077982, 0.029519490925774643, -0.039134249302383094, 0.1993975339773936, 0.7234076904024206, 0.6339789634582119, 0.01660210576452232, -0.17532808990845047, -0.021101834024758855, 0.019538882735286728],
    'sym6': [0.015404109327027373, 0.0034907120842174702, -0.11799011114819057, -0.048311742585633, 0.4910559419267466, 0.787641141030194, 0.3379294217276218, -0.07263752278646252, -0.021060292512300564, 0.04472490177066578, 0.0017677118642428036, -0.007800708325034148],
    'sym7': [0.002681814568257878, -0.0010473848886829163, -0.01263630340325193, 0.03051551316596357, 0.0678926935013727, -0.049552834937127255, 0.017441255086855827, 0.5361019170917628, 0.767764317003164, 0.2886296317515146, -0.14004724044296152, -0.10780823770381774, 0.004010244871533663, 0.010268176708511255],
    'sym8': [-0.0033824159510061256, -0.0005421323317911481, 0.03169508781149298, 0.007607487324917605, -0.1432942383508097, -0.061273359067658524, 0.4813596512583722, 0.7771857517005235, 0.3644418948353314, -0.05194583810770904, -0.027219029917056003, 0.049137179673607506, 0.003808752013890615, -0.01495225833704823, -0.0003029205147213668, 0.0018899503327594609],
}


#---------------------------------------------------------
def _map_augmentations(augmentations, p):
    augmentations = augmentations.split(",")

    aug_dict = {aug: 0.0 for aug in AVAILABLE_AUGMENTATIONS}
    if augmentations[0] == "none":
        return aug_dict
    elif augmentations[0] == "all":
        for aug in AVAILABLE_AUGMENTATIONS:
            aug_dict[aug] = p
        augmentations.pop(0)
    elif augmentations[0] == "geo":
        for aug in GEOMETRIC:
            aug_dict[aug] = p
        augmentations.pop(0)
    elif augmentations[0] == "color":
        for aug in COLOR:
            aug_dict[aug] = p
        augmentations.pop(0)
    elif augmentations[0] == "edm":
        aug_dict["hflip"] = 1.0  # NOTE: always apply hflip
        for aug in EDM_AUGMENTATIONS:
            aug_dict[aug] = p
        augmentations.pop(0)
    elif augmentations[0] == "tada":
        for aug in TADA_AUGMENTATIONS:
            aug_dict[aug] = p
        augmentations.pop(0)
    elif augmentations[0].split("=")[0] not in AVAILABLE_AUGMENTATIONS:
        raise ValueError(f"Unknown augmentation: {augmentations[0]}")

    for aug in augmentations:
        if "=" in aug:
            aug, aug_p = aug.split("=")
        else:
            aug_p = p
        if aug not in AVAILABLE_AUGMENTATIONS:
            raise ValueError(f"Unknown augmentation: {aug}")
        else:
            aug_dict[aug] = float(aug_p)
    return aug_dict


class AugReg:
    def __init__(
        self, p=0.5, augmentations="all",
        translate_int_max=0.125,
        scale_std=0.2, rotate_frac_max=1, aniso_std=0.2, aniso_rotate_prob=0.5, translate_frac_std=0.125,
        brightness_std=0.2, contrast_std=0.5, hue_max=1, saturation_std=1,
        imgfilter_bands=[1, 1, 1, 1], imgfilter_std=1,
        cutout_size=0.5,
        *args, **kwargs,
    ):
        super().__init__()
        self.probs              = _map_augmentations(augmentations, float(p))
        self.probs["hflip"]     = 1.0  # NOTE: always apply hflip
        print(self.probs)

        # Pixel blitting.
        self.translate_int_max  = float(translate_int_max)  # Range of integer translation, relative to image dimensions.

        # Geometric transformations.
        self.scale_std          = float(scale_std)          # Log2 standard deviation of isotropic scaling.
        self.rotate_frac_max    = float(rotate_frac_max)    # Range of fractional rotation, 1 = full circle.
        self.aniso_std          = float(aniso_std)          # Log2 standard deviation of anisotropic scaling.
        self.aniso_rotate_prob  = float(aniso_rotate_prob)  # Probability of doing anisotropic scaling w.r.t. rotated coordinate frame.
        self.translate_frac_std = float(translate_frac_std) # Standard deviation of frational translation, relative to image dimensions.

        # Color transformations.
        self.brightness_std     = float(brightness_std)     # Standard deviation of brightness.
        self.contrast_std       = float(contrast_std)       # Log2 standard deviation of contrast.
        self.hue_max            = float(hue_max)            # Range of hue rotation, 1 = full circle.
        self.saturation_std     = float(saturation_std)     # Log2 standard deviation of saturation.

        # Image-space filtering.
        self.imgfilter_bands    = list(imgfilter_bands)     # Probability multipliers for individual frequency bands.
        self.imgfilter_std      = float(imgfilter_std)      # Log2 standard deviation of image-space filter amplification.

        # Image-space corruptions.
        self.cutout_size        = float(cutout_size)        # Size of the cutout rectangle, relative to image dimensions.

        # Construct filter bank for image-space filtering.
        Hz_lo = np.asarray(wavelets['sym2'])            # H(z)
        Hz_hi = Hz_lo * ((-1) ** np.arange(Hz_lo.size)) # H(-z)
        Hz_lo2 = np.convolve(Hz_lo, Hz_lo[::-1]) / 2    # H(z) * H(z^-1) / 2
        Hz_hi2 = np.convolve(Hz_hi, Hz_hi[::-1]) / 2    # H(-z) * H(-z^-1) / 2
        Hz_fbank = np.eye(4, 1)                         # Bandpass(H(z), b_i)
        for i in range(1, Hz_fbank.shape[0]):
            Hz_fbank = np.dstack([Hz_fbank, np.zeros_like(Hz_fbank)]).reshape(Hz_fbank.shape[0], -1)[:, :-1]
            Hz_fbank = scipy.signal.convolve(Hz_fbank, [Hz_lo2])
            Hz_fbank[i, (Hz_fbank.shape[1] - Hz_hi2.size) // 2 : (Hz_fbank.shape[1] + Hz_hi2.size) // 2] += Hz_hi2
        self.Hz_fbank = Hz_fbank

    def __call__(self, image, *args, **kwargs):
        assert isinstance(image, PIL.Image.Image), \
            f"Invalid image type ({type(image)}). Must be PIL.Image.Image."
        image = np.asarray(image).astype(np.float32)
        image = (image / 127.5) - 1.0  # Normalize to [-1, 1].

        H, W, C = image.shape
        labels = []

        # ---------------
        # Pixel blitting.
        # ---------------

        if np.random.rand() < self.probs["hflip"]:
            w = np.random.randint(0, 2)
            if w == 1:
                image = np.fliplr(image)
            labels.append(float(w))
        else:
            labels.append(0.0)

        if np.random.rand() < self.probs["vflip"]:
            w = np.random.randint(0, 2)
            if w == 1:
                image = np.flipud(image)
            labels.append(float(w))
        else:
            labels.append(0.0)

        # if np.random.rand() < self.probs["rotate90"]:
        #     w = np.random.randint(0, 4)
        #     if w > 0:
        #         image = np.rot90(image, k=w)
        #     labels.append(w/3.0)
        # else:
        #     labels.append(0.0)

        # if np.random.rand() < self.probs["translate_int"]:
        #     raise NotImplementedError
        # else:
        #     labels += [0.0, 0.0]

        # ------------------------------------------------
        # Select parameters for geometric transformations.
        # ------------------------------------------------

        image = torch.from_numpy(image.copy()).permute(2, 0, 1)
        image.unsqueeze_(0)
        N, C, H, W = image.shape
        device = image.device
        I_3 = torch.eye(3, device=device)
        G_inv = I_3

        if np.random.rand() < self.probs["scale"]:
            w = np.random.randn(1)
            s = np.exp2(w * self.scale_std)
            s = torch.from_numpy(s).to(torch.float32)
            G_inv = G_inv @ scale2d_inv(s, s)
            labels.append(w[0])
        else:
            labels.append(0.0)

        if np.random.rand() < self.probs["rotate_frac"]:
            w = (np.random.rand(1) * 2 - 1) * (np.pi * self.rotate_frac_max)
            w_torch = torch.from_numpy(w).to(torch.float32)
            G_inv = G_inv @ rotate2d_inv(-w_torch)
            labels.append((np.cos(w) - 1)[0])
            labels.append(np.sin(w)[0])
        else:
            labels += [0.0, 0.0]

        if np.random.rand() < self.probs["aniso"]:
            w = np.random.randn(1)
            if np.random.rand() < self.aniso_rotate_prob:
                r = (np.random.rand(1) * 2 - 1) * np.pi
            else:
                r = np.zeros(1)
            s = np.exp2(w * self.aniso_std)
            s = torch.from_numpy(s).to(torch.float32)
            r_torch = torch.from_numpy(r).to(torch.float32)
            G_inv = G_inv @ rotate2d_inv(r_torch) @ scale2d_inv(s, 1 / s) @ rotate2d_inv(-r_torch)
            labels.append((w * np.cos(r))[0])
            labels.append((w * np.sin(r))[0])
        else:
            labels += [0.0, 0.0]

        if np.random.rand() < self.probs["translate_frac"]:
            w = np.random.randn(2, 1)
            w_torch = torch.from_numpy(w).to(torch.float32)
            G_inv = G_inv @ translate2d_inv(w_torch[0].mul(W * self.translate_frac_std), w_torch[1].mul(H * self.translate_frac_std))
            labels.append(w[0,0])
            labels.append(w[1,0])
        else:
            labels += [0.0, 0.0]

        # ----------------------------------
        # Execute geometric transformations.
        # ----------------------------------

        if G_inv is not I_3:
            cx = (W - 1) / 2
            cy = (H - 1) / 2
            cp = matrix([-cx, -cy, 1], [cx, -cy, 1], [cx, cy, 1], [-cx, cy, 1], device=device) # [idx, xyz]
            cp = G_inv @ cp.t() # [batch, xyz, idx]
            Hz = np.asarray(wavelets['sym6'], dtype=np.float32)
            Hz_pad = len(Hz) // 4
            margin = cp[:, :2, :].permute(1, 0, 2).flatten(1) # [xy, batch * idx]
            margin = torch.cat([-margin, margin]).max(dim=1).values # [x0, y0, x1, y1]
            margin = margin + constant([Hz_pad * 2 - cx, Hz_pad * 2 - cy] * 2, device=device)
            margin = margin.max(constant([0, 0] * 2, device=device))
            margin = margin.min(constant([W - 1, H - 1] * 2, device=device))
            mx0, my0, mx1, my1 = margin.ceil().to(torch.int32)

            # Pad image and adjust origin.
            image = F.pad(input=image, pad=[mx0,mx1,my0,my1], mode='reflect')
            G_inv = translate2d((mx0 - mx1) / 2, (my0 - my1) / 2) @ G_inv

            # Upsample.
            conv_weight = constant(Hz[None, None, ::-1], dtype=image.dtype, device=image.device).tile([image.shape[1], 1, 1])
            conv_pad = (len(Hz) + 1) // 2
            image = torch.stack([image, torch.zeros_like(image)], dim=4).reshape(N, C, image.shape[2], -1)[:, :, :, :-1]
            image = F.conv2d(image, conv_weight.unsqueeze(2), groups=image.shape[1], padding=[0,conv_pad])
            image = torch.stack([image, torch.zeros_like(image)], dim=3).reshape(N, C, -1, image.shape[3])[:, :, :-1, :]
            image = F.conv2d(image, conv_weight.unsqueeze(3), groups=image.shape[1], padding=[conv_pad,0])
            G_inv = scale2d(2, 2, device=device) @ G_inv @ scale2d_inv(2, 2, device=device)
            G_inv = translate2d(-0.5, -0.5, device=device) @ G_inv @ translate2d_inv(-0.5, -0.5, device=device)

            # Execute transformation.
            shape = [N, C, (H + Hz_pad * 2) * 2, (W + Hz_pad * 2) * 2]
            G_inv = scale2d(2 / image.shape[3], 2 / image.shape[2], device=device) @ G_inv @ scale2d_inv(2 / shape[3], 2 / shape[2], device=device)
            grid = F.affine_grid(theta=G_inv[:,:2,:], size=shape, align_corners=False)
            image = F.grid_sample(image, grid, mode='bilinear', padding_mode='zeros', align_corners=False)

            # Downsample and crop.
            conv_weight = constant(Hz[None, None, :], dtype=image.dtype, device=image.device).tile([image.shape[1], 1, 1])
            conv_pad = (len(Hz) - 1) // 2
            image = F.conv2d(image, conv_weight.unsqueeze(2), groups=image.shape[1], stride=[1,2], padding=[0,conv_pad])[:, :, :, Hz_pad : -Hz_pad]
            image = F.conv2d(image, conv_weight.unsqueeze(3), groups=image.shape[1], stride=[2,1], padding=[conv_pad,0])[:, :, Hz_pad : -Hz_pad, :]

        # --------------------------------------------
        # Select parameters for color transformations.
        # --------------------------------------------

        I_4 = torch.eye(4, device=device)
        M = I_4
        luma_axis = constant(np.asarray([1, 1, 1, 0]) / np.sqrt(3), device=device)

        if np.random.rand() < self.probs["brightness"]:
            w = np.random.randn(1)
            b = w * self.brightness_std
            b = torch.from_numpy(b).to(torch.float32)
            M = translate3d(b, b, b) @ M
            labels.append(w[0])
        else:
            labels.append(0.0)

        if np.random.rand() < self.probs["contrast"]:
            w = np.random.randn(1)
            c = np.exp2(w * self.contrast_std)
            c = torch.from_numpy(c).to(torch.float32)
            M = scale3d(c, c, c) @ M
            labels.append(w[0])
        else:
            labels.append(0.0)

        # TODO: lumaflip range
        if np.random.rand() < self.probs["lumaflip"]:
            w = np.random.randint(2, size=[1, 1, 1])
            w_torch = torch.from_numpy(w).to(torch.float32)
            M = (I_4 - 2 * luma_axis.ger(luma_axis) * w_torch) @ M
            labels.append(w[0,0,0])
        else:
            labels.append(0.0)

        if False and np.random.rand() < self.probs["hue"]:
            w = (np.random.rand(1) * 2 - 1) * (np.pi * self.hue_max)
            w_torch = torch.from_numpy(w).to(torch.float32)
            M = rotate3d(luma_axis, w_torch) @ M
            labels.append(np.cos(w)[0] - 1)
            labels.append(np.sin(w)[0])
        else:
            labels += [0.0, 0.0]

        if np.random.rand() < self.probs["saturation"]:
            w = np.random.randn(1, 1, 1)
            s = np.exp2(w * self.saturation_std)
            s = torch.from_numpy(s).to(torch.float32)
            M = (luma_axis.ger(luma_axis) + (I_4 - luma_axis.ger(luma_axis)) * s) @ M
            labels.append(w[0,0,0])
        else:
            labels.append(0.0)

        # ------------------------------
        # Execute color transformations.
        # ------------------------------
        if M is not I_4:
            image = image.reshape([N, C, H * W])
            if C == 3:
                image = M[:, :3, :3] @ image + M[:, :3, 3:]
            elif C == 1:
                M = M[:, :3, :].mean(dim=1, keepdims=True)
                image = image * M[:, :, :3].sum(dim=2, keepdims=True) + M[:, :, 3:]
            else:
                raise ValueError('Image must be RGB (3 channels) or L (1 channel)')
            image = image.reshape([N, C, H, W])

        # ----------------------
        # Image-space filtering.
        # ----------------------
        if np.random.rand() < self.probs["imgfilter"]:
            num_bands = self.Hz_fbank.shape[0]
            assert len(self.imgfilter_bands) == num_bands
            expected_power = np.array([10, 1, 1, 1]) / 13  # Expected power spectrum (1/f).

            # Apply amplification for each band with probability (imgfilter * strength * band_strength).
            g = np.ones([1, num_bands])  # Global gain vector (identity).
            for i, band_strength in enumerate(self.imgfilter_bands):
                # t_i = torch.exp2(torch.randn([N], device=device) * imgfilter_std)
                w = np.random.randn(1)
                t_i = np.exp2(w * self.imgfilter_std)
                if np.random.rand() < self.probs["imgfilter"] * band_strength:
                    labels.append(t_i[0])
                else:
                    t_i = 1.0
                    labels.append(0.0)

                t = np.ones([1, num_bands])
                t[:, i] = t_i  # Replace i'th element.
                t = t / np.sqrt((expected_power * np.square(t)).sum(axis=-1, keepdims=True))  # Normalize power.
                g = g * t

            # Construct combined amplification filter.
            Hz_prime = g @ self.Hz_fbank  # [1, 4]@[4, tap] -> [1, tap]
            Hz_prime = np.expand_dims(Hz_prime, axis=1)  # [1, tap] -> [1, 1, tap]
            Hz_prime = np.repeat(Hz_prime, C, axis=1)  # [1, 1, tap] -> [1, channels, tap]
            Hz_prime = Hz_prime.reshape([C, 1, -1])  # [channels, 1, tap]
            Hz_prime = torch.from_numpy(Hz_prime).to(torch.float32)

            # Apply filter.
            p = self.Hz_fbank.shape[1] // 2
            image = image.reshape([1, C, H, W])
            image = F.pad(input=image, pad=[p,p,p,p], mode='reflect')
            image = F.conv2d(input=image, weight=Hz_prime.unsqueeze(2), groups=C)
            image = F.conv2d(input=image, weight=Hz_prime.unsqueeze(3), groups=C)
        else:
            labels += [0.0, 0.0, 0.0, 0.0]
        image.squeeze_(0)


        # ------------------------
        # Image-space corruptions.
        # ------------------------
        # if np.random.rand() < self.probs["cutout"]:
        #     w = np.random.rand(2)
        #     size = w * self.cutout_size
        #     size = (round(size[0]*H), round(size[1]*W))
        #     mask = torch.ones_like(image)  # [C, H, W]
        #     coord_y = np.random.randint(0, H-size[0]+1)
        #     coord_x = np.random.randint(0, W-size[1]+1)
        #     mask[:, coord_y:coord_y+size[0], coord_x:coord_x+size[1]] = 0.0
        #     image = image * mask
        #     labels.append((coord_y+1)/H)  # top
        #     labels.append((coord_x+1)/W)  # left
        #     labels.append(w[0])  # height
        #     labels.append(w[1])  # width
        # else:
        #     labels += [0.0, 0.0, 0.0, 0.0]

        image = torch.clip(image, -1.0, 1.0)
        labels = np.asarray(labels, dtype=np.float32)
        return image, labels
