# Copyright (c) 2021 PPViT Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
CycleMLP in Paddle
A Paddle Implementation of CycleMLP as described in:
"CycleMLP: A MLP-like Architecture for Dense Prediction"
    - Paper Link: https://arxiv.org/abs/2107.10224
"""

import os
import math
import paddle
import paddle.nn as nn
from paddle import Tensor
from paddle.vision.ops import deform_conv2d
import paddle.nn.functional as F
from droppath import DropPath

zeros_ = nn.initializer.Constant(value=0.0)
ones_ = nn.initializer.Constant(value=1.0)
trunc_normal_ = nn.initializer.TruncatedNormal(std=0.02)
kaiming_uniform_ = nn.initializer.KaimingUniform()

class Identity(nn.Layer):
    """Identity layer
    This is does nothing but passing the input as output
    """
    def __init__(self):
        super().__init__()

    def forward(self, inputs):
        return inputs

class Mlp(nn.Layer):
    """ MLP module
    Impl using nn.Linear and activation is GELU, dropout is applied.
    Ops: fc -> act -> dropout -> fc -> dropout

    Attributes:
        fc1: nn.Linear
        fc2: nn.Linear
        act: GELU
        dropout: dropout after fc
    """
    def __init__(self,
                 in_features,
                 hidden_features=None,
                 out_features=None,
                 act_layer=nn.GELU,
                 drop=0.0):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class CycleFC(nn.Layer):

    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 kernel_size,  # re-defined kernel_size, represent the spatial area of staircase FC
                 stride: int = 1,
                 padding: int = 0,
                 dilation: int = 1,
                 groups: int = 1,
                 bias: bool = True):
        super(CycleFC, self).__init__()

        if in_channels % groups != 0:
            raise ValueError("in_channels must be divisible by groups")
        if out_channels % groups != 0:
            raise ValueError("out_channels must be divisible by groups")
        if stride != 1:
            raise ValueError("stride must be 1")
        if padding != 0:
            raise ValueError("padding must be 0")

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = (stride, stride)
        self.padding = (padding, padding)
        self.dilation = (dilation, dilation)
        self.groups = groups

        self.weight = self.create_parameter(
            shape=[out_channels, in_channels // groups, 1, 1],
            default_initializer=kaiming_uniform_,
        )  # kernel size == 1

        if bias:
            bound = 1 / math.sqrt(self.weight.shape[1])
            self.bias = self.create_parameter(
                shape=[out_channels],
                default_initializer=nn.initializer.Uniform(-bound, bound),
            )
        else:
            self.bias = None
        self.register_buffer("offset", self.gen_offset())

    def gen_offset(self):
        """
        offset (Tensor[batch_size, 2 * offset_groups * kernel_height * kernel_width,
            out_height, out_width]): offsets to be applied for each position in the
            convolution kernel.
        """
        offset = paddle.empty([1, self.in_channels * 2, 1, 1])
        start_idx = (self.kernel_size[0] * self.kernel_size[1]) // 2
        assert self.kernel_size[0] == 1 or self.kernel_size[1] == 1, self.kernel_size
        for i in range(self.in_channels):
            if self.kernel_size[0] == 1:
                offset[0, 2 * i + 0, 0, 0] = 0
                offset[0, 2 * i + 1, 0, 0] = (i + start_idx) % self.kernel_size[1] - (
                    self.kernel_size[1] // 2
                )
            else:
                offset[0, 2 * i + 0, 0, 0] = (i + start_idx) % self.kernel_size[0] - (
                    self.kernel_size[0] // 2
                )
                offset[0, 2 * i + 1, 0, 0] = 0
        return offset

    def forward(self, inputs: Tensor) -> Tensor:
        """
        Args:
            input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor
        """
        B, C, H, W = inputs.shape
        deformable_groups = self.offset.shape[1] // (
            2 * self.weight.shape[2] * self.weight.shape[3])

        return deform_conv2d(inputs,
                             self.offset.expand([B, -1, H, W]),
                             self.weight,
                             self.bias,
                             stride=self.stride,
                             padding=self.padding,
                             dilation=self.dilation,
                             deformable_groups=deformable_groups)

    def extra_repr(self) -> str:
        s = self.__class__.__name__ + "("
        s += "{in_channels}"
        s += ", {out_channels}"
        s += ", kernel_size={kernel_size}"
        s += ", stride={stride}"
        s += ", padding={padding}" if self.padding != (0, 0) else ""
        s += ", dilation={dilation}" if self.dilation != (1, 1) else ""
        s += ", groups={groups}" if self.groups != 1 else ""
        s += ", bias=False" if self.bias is None else ""
        s += ")"
        return s.format(**self.__dict__)

class CycleMLP(nn.Layer):
    def __init__(self,
                 dim,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.0,
                 proj_drop=0.0):
        super().__init__()
        self.mlp_c = nn.Linear(dim, dim, bias_attr=qkv_bias)

        self.sfc_h = CycleFC(dim, dim, (1, 3), 1, 0)
        self.sfc_w = CycleFC(dim, dim, (3, 1), 1, 0)

        self.reweight = Mlp(dim, dim // 4, dim * 3)

        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, H, W, C = x.shape
        h = self.sfc_h(x.transpose([0, 3, 1, 2])).transpose([0, 2, 3, 1])
        w = self.sfc_w(x.transpose([0, 3, 1, 2])).transpose([0, 2, 3, 1])
        c = self.mlp_c(x)

        a = (h + w + c).transpose([0, 3, 1, 2]).flatten(2).mean(2)
        a = F.softmax(self.reweight(a).reshape((B, C, 3)).transpose([2, 0, 1]), axis=0)
        a = a.unsqueeze(2)
        a = a.unsqueeze(2)

        x = h * a[0] + w * a[1] + c * a[2]

        x = self.proj(x)
        x = self.proj_drop(x)

        return x

class CycleBlock(nn.Layer):
    def __init__(self,
                 dim,
                 mlp_ratio=4.0,
                 qkv_bias=False,
                 qk_scale=None,
                 drop=0.0,
                 attn_drop=0.0,
                 drop_path=0.0,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm,
                 skip_lam=1.0,
                 mlp_fn=CycleMLP):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = mlp_fn(dim, qkv_bias=qkv_bias, qk_scale=None, attn_drop=attn_drop)

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()

        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer)
        self.skip_lam = skip_lam

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x))) / self.skip_lam
        x = x + self.drop_path(self.mlp(self.norm2(x))) / self.skip_lam
        return x

class PatchEmbedOverlapping(nn.Layer):
    """2D Image to Patch Embedding with overlapping"""

    def __init__(self,
                 patch_size=16,
                 stride=16,
                 padding=0,
                 in_chans=3,
                 embed_dim=768,
                 norm_layer=None,
                 groups=1):
        super().__init__()
        patch_size = (patch_size, patch_size)
        stride = (stride, stride)
        padding = (padding, padding)
        self.patch_size = patch_size
        # remove image_size in model init to support dynamic image size

        self.proj = nn.Conv2D(in_chans,
                              embed_dim,
                              kernel_size=patch_size,
                              stride=stride,
                              padding=padding,
                              groups=groups)
        self.norm = norm_layer(embed_dim) if norm_layer else Identity()

    def forward(self, x):
        x = self.proj(x)
        return x

class Downsample(nn.Layer):
    """Downsample transition stage"""
    def __init__(self, in_embed_dim, out_embed_dim, patch_size):
        super().__init__()
        assert patch_size == 2, patch_size
        self.proj = nn.Conv2D(in_embed_dim,
                              out_embed_dim,
                              kernel_size=(3, 3),
                              stride=(2, 2),
                              padding=1)

    def forward(self, x):
        x = x.transpose([0, 3, 1, 2])
        x = self.proj(x)  # B, C, H, W
        x = x.transpose([0, 2, 3, 1])
        return x

def basic_blocks(dim,
                 index,
                 layers,
                 mlp_ratio=3.0,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.0,
                 drop_path_rate=0.0,
                 skip_lam=1.0,
                 mlp_fn=CycleMLP,
                 **kwargs):
    blocks = []

    for block_idx in range(layers[index]):
        block_dpr = (
            drop_path_rate * (block_idx + sum(layers[:index])) / (sum(layers) - 1)
        )
        blocks.append(
            CycleBlock(
                dim,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                attn_drop=attn_drop,
                drop_path=block_dpr,
                skip_lam=skip_lam,
                mlp_fn=mlp_fn,
            )
        )
    blocks = nn.Sequential(*blocks)

    return blocks

class CycleNet(nn.Layer):
    """CycleMLP Network"""
    def __init__(self,
                 layers,
                 img_size=224,
                 patch_size=4,
                 in_chans=3,
                 num_classes=1000,
                 embed_dims=None,
                 transitions=None,
                 segment_dim=None,
                 mlp_ratios=None,
                 skip_lam=1.0,
                 qkv_bias=False,
                 qk_scale=None,
                 drop_rate=0.0,
                 attn_drop_rate=0.0,
                 drop_path_rate=0.0,
                 norm_layer=nn.LayerNorm,
                 mlp_fn=CycleMLP,
                 fork_feat=False):
        super().__init__()
        if not fork_feat:
            self.num_classes = num_classes
        self.fork_feat = fork_feat
        self.patch_embed = PatchEmbedOverlapping(patch_size=7,
                                                 stride=4,
                                                 padding=2,
                                                 in_chans=3,
                                                 embed_dim=embed_dims[0])
        network = []
        for i in range(len(layers)):
            stage = basic_blocks(dim=embed_dims[i],
                                 index=i,
                                 layers=layers,
                                 mlp_ratio=mlp_ratios[i],
                                 qkv_bias=qkv_bias,
                                 qk_scale=qk_scale,
                                 attn_drop=attn_drop_rate,
                                 drop_path_rate=drop_path_rate,
                                 norm_layer=norm_layer,
                                 skip_lam=skip_lam,
                                 mlp_fn=mlp_fn)
            network.append(stage)
            if i >= len(layers) - 1:
                break
            if transitions[i] or embed_dims[i] != embed_dims[i + 1]:
                patch_size = 2 if transitions[i] else 1
                network.append(Downsample(embed_dims[i], embed_dims[i + 1], patch_size))

        self.network = nn.LayerList(network)

        if self.fork_feat:
            # add a norm layer for each output
            self.out_indices = [0, 2, 4, 6]
            for i_emb, i_layer in enumerate(self.out_indices):
                if i_emb == 0 and os.environ.get("FORK_LAST3", None):
                    layer = Identity()
                else:
                    layer = norm_layer(embed_dims[i_emb])
                layer_name = f"norm{i_layer}"
                self.add_layer(layer_name, layer)
        else:
            # Classifier head
            self.norm = norm_layer(embed_dims[-1])
            self.head = nn.Linear(embed_dims[-1], num_classes) if num_classes > 0 else Identity()
        self.apply(self.cls_init_weights)

    def cls_init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                zeros_(m.bias)
        elif isinstance(m, nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)
        elif isinstance(m, CycleFC):
            trunc_normal_(m.weight)
            zeros_(m.bias)

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=""):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else Identity()

    def forward_embeddings(self, x):
        x = self.patch_embed(x)
        # B,C,H,W-> B,H,W,C
        x = x.transpose([0, 2, 3, 1])
        return x

    def forward_tokens(self, x):
        outs = []
        for idx, block in enumerate(self.network):
            x = block(x)
            if self.fork_feat and idx in self.out_indices:
                norm_layer = getattr(self, f"norm{idx}")
                x_out = norm_layer(x)
                outs.append(x_out.transpose([0, 3, 1, 2]))
        if self.fork_feat:
            return outs

        B, H, W, C = x.shape
        x = x.reshape([B, -1, C])
        return x

    def forward(self, x):
        x = self.forward_embeddings(x)
        # B, H, W, C -> B, N, C
        x = self.forward_tokens(x)
        if self.fork_feat:
            return x

        x = self.norm(x)
        cls_out = self.head(x.mean(1))
        return cls_out

def build_cyclemlp(config):
    '''build cyclemlp model'''
    model = CycleNet(num_classes=config.MODEL.NUM_CLASSES,
                     layers=config.MODEL.LAYERS,
                     embed_dims=config.MODEL.EMBED_DIMS,
                     patch_size=config.MODEL.PATCH_SIZE,
                     transitions=config.MODEL.TRANSITIONS,
                     mlp_ratios=config.MODEL.MLP_RATIOS,
                     mlp_fn=CycleMLP)
    return model


import glob
import random
from datetime import datetime

import numpy as np
from scipy.ndimage.filters import gaussian_filter
from scipy.ndimage.interpolation import map_coordinates

from matplotlib import pyplot as plt

import skimage.transform

# CONFIG FILE
import config_vacbag_tversky as config
import cv2
import loss as loss
import paths as paths
import tensorflow as tf
import unet as unet

physical_devices = tf.config.experimental.list_physical_devices("GPU")
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
config = tf.config.experimental.set_memory_growth(physical_devices[0], True)

# The meaning of life
random.seed(42)
np.random.seed(42)

class make_gen(tf.keras.utils.Sequence):
    def __init__(
        self,
        input_paths,
        label_paths,
        batch_size,
        training_mean,
        training_std,
        shuffle_on_end=True,
        augment=True,
    ):
        self.input_paths = input_paths
        self.label_paths = label_paths
        self.batch_size = batch_size
        self.training_mean = training_mean
        self.training_std = training_std
        self.shuffle_on_end = shuffle_on_end
        self.augment = augment

    def __len__(self):
        # number of batches per epoch
        return int(np.ceil(len(self.input_paths) / float(self.batch_size)))

    def on_epoch_end(self):
        """Updates indexes after each epoch
        """
        if self.shuffle_on_end == True:
            self.inputs, self.truths = self.suffle_together(self.inputs, self.truths)

    def shuffle_together(self, inputs, truths):
        shuffle_together = list(zip(inputs, truths))
        random.shuffle(shuffle_together)
        inputs, truths = zip(*shuffle_together)
        return inputs, truths

    def gaussian_noise(self, img, mean=0, sigma=0.003):
        img = img.copy()
        noise = np.random.normal(mean, sigma, img.shape)
        mask_overflow_upper = img + noise >= 1.0
        mask_overflow_lower = img + noise < 0
        noise[mask_overflow_upper] = 1.0
        noise[mask_overflow_lower] = 0
        img = img + noise
        return img

    def random_crop_resize(self, img, label, crop_size=500):
        size_img = img.shape
        size_label = label.shape
        crop_size = random.randint(crop_size, img.shape[0] - 1)
        crop_size = (crop_size, crop_size)

        # "Crop size should be less than image size"
        assert crop_size[0] <= img.shape[0] and crop_size[1] <= img.shape[1]

        w, h = img.shape[:2]
        x, y = np.random.randint(h - crop_size[0]), np.random.randint(w - crop_size[1])

        img = img[y : y + crop_size[0], x : x + crop_size[1], :]
        img = skimage.transform.resize(img, size_img)

        label = label[y : y + crop_size[0], x : x + crop_size[1], :]
        label = skimage.transform.resize(label, size_label)
        return img, label

    def affine_transform(self, image, label, alpha_affine=0.5, random_state=None):

        if random_state is None:
            random_state = np.random.RandomState(None)

        shape = image.shape
        shape_size = shape[:2]
        center_square = np.float32(shape_size) // 2
        square_size = min(shape_size) // 3
        pts1 = np.float32(
            [
                center_square + square_size,
                [center_square[0] + square_size, center_square[1] - square_size],
                center_square - square_size,
            ]
        )
        pts2 = pts1 + random_state.uniform(
            -alpha_affine, alpha_affine, size=pts1.shape
        ).astype(np.float32)
        M = cv2.getAffineTransform(pts1, pts2)

        image = cv2.warpAffine(
            image, M, shape_size[::-1], borderMode=cv2.BORDER_REFLECT_101
        )
        image = image[..., np.newaxis]
        label = cv2.warpAffine(
            label, M, shape_size[::-1], borderMode=cv2.BORDER_REFLECT_101
        )
        return image, label

    def elastic_transform(self, image, label, alpha, sigma, random_state=None):

        if random_state is None:
            random_state = np.random.RandomState(None)

        shape = image.shape
        shape_label = label.shape

        dx = (
            gaussian_filter(
                (random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0
            )
            * alpha
        )
        dy = (
            gaussian_filter(
                (random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0
            )
            * alpha
        )
        dz = np.zeros_like(dx)

        # image
        x, y, z = np.meshgrid(
            np.arange(shape[0]), np.arange(shape[1]), np.arange(shape[2])
        )
        indices = (
            np.reshape(y + dy, (-1, 1)),
            np.reshape(x + dx, (-1, 1)),
            np.reshape(z, (-1, 1)),
        )
        image = map_coordinates(image, indices, order=1, mode="reflect").reshape(shape)

        # label
        x, y, z = np.meshgrid(
            np.arange(shape_label[0]),
            np.arange(shape_label[1]),
            np.arange(shape_label[2]),
        )
        indices = (
            np.reshape(y + dy, (-1, 1)),
            np.reshape(x + dx, (-1, 1)),
            np.reshape(z, (-1, 1)),
        )
        label = map_coordinates(label, indices, order=1, mode="reflect").reshape(
            shape_label
        )

        return image, label

    def data_augment(self, img, mask, chance=0.5):
        # flip l/r
        if random.uniform(0, 1) < 0.5:
            img = cv2.flip(img, 1)
            mask = cv2.flip(mask, 1)
            if len(img.shape) == 2:
                img = img[..., np.newaxis]
            if len(mask.shape) == 2:
                mask = mask[..., np.newaxis]

        # random crop and resize
        if random.uniform(0, 1) < chance:
            img, mask = self.random_crop_resize(img, mask)
            if len(img.shape) == 2:
                img = img[..., np.newaxis]
            if len(mask.shape) == 2:
                label = label[..., np.newaxis]

        # random affine transformation
        if random.uniform(0, 1) < chance:
            img, mask = self.affine_transform(img, mask, alpha_affine=20)
            if len(img.shape) == 2:
                img = img[..., np.newaxis]
            if len(mask.shape) == 2:
                mask = mask[..., np.newaxis]

        if random.uniform(0, 1) < chance:
            args = random.choice(((1201, 10), (1501, 12), (991, 8)))
            img, mask = self.elastic_transform(img, mask, *args)

        # random Gaussian noise
        if random.uniform(0, 1) < chance:
            sigma = random.choice(np.arange(0.1, 0.3, 0.02))
            img = self.gaussian_noise(img, mean=0, sigma=sigma)

        return img, mask

    def normalise(self, x, mean, std):
        return (x - mean) / std

    def read_array_list(self, arr_path_list):
        return np.array([np.load(arr_path) for arr_path in arr_path_list])

    def __getitem__(self, batch_index):
        # Get batch at index: batch_index

        # Handle case when not enough inputs for a full batch
        if (batch_index + 1) * self.batch_size > len(self.input_paths):
            batch_size = len(self.input_paths) - batch_index * self.batch_size
        else:
            batch_size = self.batch_size

        # Extract input paths for one batch at index: batch_index
        batch_input_paths = self.input_paths[
            batch_index * batch_size : (batch_index + 1) * batch_size
        ]
        batch_label_paths = self.label_paths[
            batch_index * batch_size : (batch_index + 1) * batch_size
        ]

        batch_imgs = []
        batch_masks = []

        for x, y in zip(batch_input_paths, batch_label_paths):
            x = np.load(x)
            y = np.load(y)
            x = self.normalise(x, self.training_mean, self.training_std)

            if self.augment is True:
                x, y = self.data_augment(x, y, chance=0.33)

            batch_imgs.append(x)
            batch_masks.append(y)

        return (
            np.array(batch_imgs, dtype=np.float32),
            np.array(batch_masks, dtype=np.float32),
        )

def get_data_statistics(train_inputs, valid_inputs):
    data = train_inputs + valid_inputs
    arr = np.array([np.load(x) for x in data])
    mean = np.mean(arr)
    std = np.std(arr)
    del arr
    return mean, std

patient_paths = paths.get_patient_paths(config.DATA_PATH)
patient_paths.sort()

img_paths = [glob.glob(path + "/img/*") for path in patient_paths]
mask_paths = [glob.glob(path + "/mask/*") for path in patient_paths]

valid = int(len(img_paths) * 0.15 // 1)
test = int(len(img_paths) * 0.1 // 1)
train = int(len(img_paths) - valid - test)

train_inputs = paths.flatten_list(img_paths[0:train])
train_truths = paths.flatten_list(mask_paths[0:train])

train_inputs.sort()
train_truths.sort()

valid_inputs = paths.flatten_list(img_paths[train : train + valid])
valid_truths = paths.flatten_list(mask_paths[train : train + valid])

valid_inputs.sort()
valid_truths.sort()

test_inputs = paths.flatten_list(img_paths[train + valid :])
test_truths = paths.flatten_list(mask_paths[train + valid :])

test_inputs.sort()
test_truths.sort()

data_mean, data_std = get_data_statistics(train_inputs, valid_inputs)

train_gen = make_gen(
    train_inputs,
    train_truths,
    config.BATCH_SIZE,
    data_mean,
    data_std,
    shuffle_on_end=True,
    augment=True,
)

valid_gen = make_gen(
    valid_inputs,
    valid_truths,
    config.BATCH_SIZE,
    data_mean,
    data_std,
    shuffle_on_end=False,
    augment=False,
)

# test_gen = make_gen(test_inputs,
#                      test_truths,
#                      config.BATCH_SIZE,
#                      data_mean,
#                      data_std,
#                      shuffle_on_end=False,
#                      augment=False)

model = unet.model(output_channels=config.OUTPUT_CHANNELS)

model.compile(config.OPTIMIZER, config.LOSS, config.METRICS)

if config.INITIAL_WEIGHTS is not None:
    model.load_weights(config.INITIAL_WEIGHTS)

checkpoint_name = config.MODEL_SAVE + "_epoch_{epoch:02d}" + ".hdf5"

early_stopping = tf.keras.callbacks.EarlyStopping(
    patience=config.STOP_PATIENCE, verbose=1, restore_best_weights=True
)
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
    factor=config.LR_SCALE, patience=config.LR_PATIENCE, verbose=1
)
model_checkpoint = tf.keras.callbacks.ModelCheckpoint(
    checkpoint_name, save_weights_only=True, verbose=1
)
csv_logger = tf.keras.callbacks.CSVLogger(
    config.MODEL_SAVE + "_csv_log", separator=",", append=False
)
tensor_board = tf.keras.callbacks.TensorBoard(
    log_dir=config.MODEL_SAVE + "_log_dir",
    histogram_freq=5,
    write_graph=True,
    write_images=True,
    embeddings_freq=5,
    update_freq="epoch",
)

callbacks = [early_stopping, model_checkpoint, reduce_lr, csv_logger, tensor_board]

steps_per_epoch = train_gen.__len__()
valid_steps = valid_gen.__len__()

train_history = model.fit(
    train_gen,
    epochs=config.EPOCHS,
    steps_per_epoch=steps_per_epoch,
    validation_steps=valid_steps,
    validation_data=valid_gen,
    callbacks=callbacks,
    verbose=1,
)


#!/usr/bin/env python3
'''
Causeway Server - key/value storage server geared toward small files with ECSDA signature auth

Usage:
    python3 server.py
'''
import os, json, random, time, string
from settings import DATABASE, PRICE, DATA_DIR, SERVER_PORT

from flask import Flask
from flask import request
from flask import abort, url_for
from flask.ext.sqlalchemy import SQLAlchemy

from two1.lib.wallet import Wallet
from two1.lib.bitserv.flask import Payment

from models import *

app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + DATABASE
db = SQLAlchemy(app)
wallet = Wallet()
payment = Payment(app, wallet)

# start time
start_time = time.time()
stored = 0

@app.route('/')
@app.route('/help')
def home():
    '''Return service, pricing and endpoint information'''
    home_obj = [{"name": "causeway/1",       # service 'causeway', version '1'
                 "pricing-type": "per-mb",   # pricing is listed per 1000000 bytes
                 "pricing" : [{"rpc": "buy",
                               "per-req": 0,
                               "per-unit": PRICE,
                               "description": "1 MB hosting, 50 MB bandwidth, 1 year expiration"
                              },
                              {"rpc": "get",
                               "per-req": 0,
                               "per-mb": 0
                              },
                              {"rpc": "put",
                               "per-req": 0,
                               "per-mb": 0
                              },

                              # default
                              {"rpc": True,        # True indicates default
                               "per-req": 0,
                               "per-mb": 0
                              }],
                  "description": "This Causeway server provides microhosting services. Download the "\
                  "client and server at https://github.com/jgarzik/playground21/archive/master.zip"
                }
               ]

    body = json.dumps(home_obj, indent=2)

    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
                       )

@app.route('/status')
def status():
    '''Return general info about server instance. '''
    uptime = str(int(time.time() - start_time))
    st = os.statvfs(DATA_DIR)
    free = st.f_bavail * st.f_frsize
    body = json.dumps({'uptime': uptime,
                       'stored': str(stored),
                       'free': str(free),
                       'price': str(PRICE)
                      }, indent=2
                     )
    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
           )

@app.route('/price')
def price():
    '''Return price for 1MB storage with bundled 50MB transfer.'''
    body = json.dumps({'price': PRICE})
    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
           )

@app.route('/buy')
@payment.required(PRICE)
def buy_hosting():
    '''Registers one hosting bucket to account on paid request.'''
    # extract account address from client request
    owner = request.args.get('address')
    contact = request.args.get('contact')

    # check if user exists
    o = db.session.query(Owner).get(owner)
    if o is None:
        # create them
        o = Owner(owner)
        db.session.add(o)
        db.session.commit()

    # owner should now exist,  create sale record for address
    s = Sale(owner, contact, 1, 30, PRICE)
    db.session.add(s)
    db.session.commit()

    body = json.dumps({'result': 'success',
                       'buckets': s.get_buckets()}, indent=2)
    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
           )

@app.route('/put', methods=['POST'])
def put():
    '''Store a key-value pair.'''
    # get size of file sent
    # Validate JSON body w/ API params
    try:
        body = request.data.decode('utf-8')
        in_obj = json.loads(body)
    except:
        return ("JSON Decode failed", 400, {'Content-Type':'text/plain'})

    k = in_obj['key']
    v = in_obj['value']
    o = in_obj['address']
    n = in_obj['nonce']
    s = in_obj['signature']

    # check signature
    owner = Owner.query.filter_by(address=o).first()
    if owner.nonce not in n or wallet.verify_bitcoin_message(k + v + o + n, s, o):
        body = json.dumps({'error': 'Incorrect signature.'})
        code = 401
    else:
        size = len(k) + len(v)

        # check if owner has enough free storage
        # get free space from each of owner's buckets
        result = db.engine.execute('select * from sale where julianday("now") - \
                    julianday(sale.created) < sale.term order by sale.created desc')
        # choose newest bucket that has enough space
        sale_id = None
        for row in result:
            if (row[7] + size) < (1024 * 1024):
                sale_id = row[0]

        if sale_id is None:     # we couldn't find enough free space
            body = json.dumps({'error': 'Insufficient storage space.'})
            code = 403
        else:
            # check if key already exists and is owned by the same owner
            kv = db.session.query(Kv).filter_by(key=k).filter_by(owner=o).first()

            if kv is None:
                kv = Kv(k, v, o, sale_id)
                db.session.add(kv)
                db.session.commit()
            else:
                kv.value = v
                db.session.commit()

            s = db.session.query(Sale).get(sale_id)
            s.bytes_used = s.bytes_used + size
            db.session.commit()
            body = json.dumps({'result': 'success'})
            code = 201

    return (body, code, {'Content-length': len(body),
                        'Content-type': 'application/json',
                        }
           )

@app.route('/delete', methods=['POST'])
def delete():
    '''Delete a key-value pair.'''
    # Validate JSON body w/ API params
    try:
        body = request.data.decode('utf-8')
        in_obj = json.loads(body)
    except:
        return ("JSON Decode failed", 400, {'Content-Type':'text/plain'})

    k = in_obj['key']
    o = in_obj['address']
    n = in_obj['nonce']
    s = in_obj['signature']

    # check signature
    owner = Owner.query.filter_by(address=o).first()
    if owner.nonce not in n or wallet.verify_bitcoin_message(k + o + n, s, o):
        body = json.dumps({'error': 'Incorrect signature.'})
        code = 401
    else:
        # check if key already exists and is owned by the same owner
        kv = db.session.query(Kv).filter_by(key=k).filter_by(owner=o).first()
        if kv is None:
            body = json.dumps({'error': 'Key not found or not owned by caller.'})
            code = 404
        else:
            # free up storage quota and remove kv
            size = len(kv.value)
            sale_id = kv.sale
            s = db.session.query(Sale).get(sale_id)
            s.bytes_used = s.bytes_used - size
            db.session.delete(kv)
            db.session.commit()
            body = json.dumps({'result': 'success'})
            code = 200

    return (body, code, {'Content-length': len(body),
                         'Content-type': 'application/json',
                        }
           )

@app.route('/get')
def get():
    '''Get a key-value pair.'''

    key = request.args.get('key')

    kv = Kv.query.filter_by(key=key).first()

    if kv is None:
        body = json.dumps({'error': 'Key not found.'})
        code = 404
    else:
        body = json.dumps({'key': key, 'value': kv.value})
        code = 200

    # calculate size and check against quota on kv's sale record
    return (body, code, {'Content-length': len(body),
                        'Content-type': 'application/json',
                        }
           )

@app.route('/nonce')
def nonce():
    '''Return 32-byte nonce for generating non-reusable signatures..'''
    # check if user exists
    o = db.session.query(Owner).get(request.args.get('address'))
    if o is None:
        return abort(500)

    # if nonce is set for user return it, else make a new one
    if o.nonce and len(o.nonce) == 32:
        body = json.dumps({'nonce': o.nonce})
    # if not, create one and store it
    else:
        print("storing")
        n = ''.join(random.SystemRandom().choice(string.hexdigits) for _ in range(32))
        o.nonce = n.lower()
        db.session.commit()
        body = json.dumps({'nonce': o.nonce})

    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
           )

@app.route('/address')
def get_deposit_address():
    '''Return new or unused deposit address for on-chain funding.'''
    # check if user exists
    o = db.session.query(Owner).get(request.args.get('address'))
    if o is None:
        return abort(500)

    address = request.args.get('address')
    message = request.args.get('contact') + "," + address
    signature = request.args.get('signature')

    print(len(signature))
    if len(signature) == 88 and wallet.verify_bitcoin_message(message, signature, address):
        body = json.dumps({'address': 'hereyago'})
    else:
        body = json.dumps({'error': 'Invalid signature'})

    return (body, 200, {'Content-length': len(body),
                        'Content-type': 'application/json',
                       }
           )

def has_no_empty_params(rule):
    '''Testing rules to identify routes.'''
    defaults = rule.defaults if rule.defaults is not None else ()
    arguments = rule.arguments if rule.arguments is not None else ()
    return len(defaults) >= len(arguments)

@app.route('/info')
def info():
    '''Returns list of defined routes.'''
    links = []
    for rule in app.url_map.iter_rules():
        # Filter out rules we can't navigate to in a browser
        # and rules that require parameters
        if "GET" in rule.methods and has_no_empty_params(rule):
            url = url_for(rule.endpoint, **(rule.defaults or {}))
            links.append(url)

    return json.dumps(links, indent=2)

if __name__ == '__main__':
    app.debug = True
    app.run(host='0.0.0.0', port=SERVER_PORT)


from torch.utils.data import Dataset, DataLoader
from pycocotools.coco import COCO

import numpy as np
import cv2
import os

class CustomDataset(Dataset):
    """COCO format"""
    def __init__(self, data_dir, ann_file, categories, mode = 'train', transform = None):
        super().__init__()
        self.data_dir = data_dir
        self.categories = categories
        self.mode = mode
        self.transform = transform

        ann_path = os.path.join(data_dir, ann_file)
        self.coco = COCO(ann_path)

    def __getitem__(self, index: int):
        # dataset이 index되어 list처럼 동작
        image_id = self.coco.getImgIds(imgIds=index)
        image_infos = self.coco.loadImgs(image_id)[0]

        # cv2 를 활용하여 image 불러오기
        images = cv2.imread(os.path.join(self.data_dir, image_infos['file_name']))
        images = cv2.cvtColor(images, cv2.COLOR_BGR2RGB).astype(np.float32)
        images /= 255.0

        if (self.mode in ('train', 'val')):
            ann_ids = self.coco.getAnnIds(imgIds=image_infos['id'])
            anns = self.coco.loadAnns(ann_ids)

            # Load the categories in a variable
            cat_ids = self.coco.getCatIds()
            cats = self.coco.loadCats(cat_ids)

            # masks : size가 (height x width)인 2D
            # 각각의 pixel 값에는 "category id" 할당
            # Background = 0
            masks = np.zeros((image_infos["height"], image_infos["width"]))
            # General trash = 1, ... , Cigarette = 10
            anns = sorted(anns, key=lambda idx : idx['area'], reverse=True)
            for i in range(len(anns)):
                className = self.get_classname(anns[i]['category_id'], cats)
                pixel_value = self.categories.index(className)
                masks[self.coco.annToMask(anns[i]) == 1] = pixel_value
            masks = masks.astype(np.int8)

            # transform -> albumentations 라이브러리 활용
            if self.transform is not None:
                transformed = self.transform(image=images, mask=masks)
                images = transformed["image"]
                masks = transformed["mask"]
            return images, masks, image_infos

        if self.mode == 'test':
            # transform -> albumentations 라이브러리 활용
            if self.transform is not None:
                transformed = self.transform(image=images)
                images = transformed["image"]
            return images, image_infos

    def __len__(self) -> int:
        # 전체 dataset의 size를 return
        return len(self.coco.getImgIds())

    def get_classname(self, classID, cats):
        for i in range(len(cats)):
            if cats[i]['id']==classID:
                return cats[i]['name']
        return "None"

if __name__ == "__main__":
    import json
    import pandas as pd

    # collate_fn needs for batch
    def collate_fn(batch):
        return tuple(zip(*batch))

    cfg = json.load(open("../cfg.json", "r"))
    train = CustomDataset(data_dir=cfg["data_dir"], ann_file=cfg["ann_file"]["train"], categories = cfg["categories"], mode='train', transform=None)
    train_loader = DataLoader(dataset=train,
                                        batch_size=8,
                                        shuffle=True,
                                        num_workers=4,
                                        collate_fn=collate_fn,
                                        drop_last=True)
    val = CustomDataset(data_dir=cfg["data_dir"], ann_file=cfg["ann_file"]["val"], categories = cfg["categories"], mode='val', transform=None)
    test = CustomDataset(data_dir=cfg["data_dir"], ann_file=cfg["ann_file"]["test"], categories = cfg["categories"], mode='test', transform=None)

    batch = next(iter(train_loader))
    for i, (image, mask, info) in enumerate(zip(*batch)):
        image*=255

        image = image.astype(np.uint8)

        mask = np.expand_dims(mask, axis=2)
        mask = mask.astype(np.uint8)

        color_map = np.array([
            [0, 0, 0],
            [255, 0, 0],
            [0, 255, 0],
            [0, 0, 255],
            [255, 255, 0],
            [0, 255, 255],
            [255, 0, 255],
            [192, 128, 64],
            [192, 192, 128],
            [64, 64, 128],
            [128, 0, 192],
        ])

        ## segment의 bit 이미지 제작
        bit_mask = mask.copy()
        bit_mask[bit_mask>0] = 255
        bit_mask = cv2.cvtColor(bit_mask, cv2.COLOR_GRAY2RGB)

        ## 마스크, 세그먼트 원본, 배경 원본 작성
        mask = np.array(list(map(lambda x: color_map[x], mask)), dtype=np.uint8).squeeze()
        segment = cv2.bitwise_and(image, bit_mask)
        bg = cv2.subtract(image, bit_mask)

        masked_segment = cv2.addWeighted(segment, 0.5, mask, 0.5, 0)
        viz = cv2.bitwise_or(masked_segment, bg)

        cv2.imwrite(f"src{i}.jpg", viz)


from common.ui_utils import ui_utils
from hawkular.hawkular_api import hawkular_api
from common.db import db
from common.navigate import navigate

class deployments():
    web_session = None
    ui_utils = None
    hawkular_api = None
    db = None

    def __init__(self, web_session):
        self.web_session = web_session
        self.ui_utils = ui_utils(self.web_session)
        self.hawkular_api = hawkular_api(self.web_session)
        self.db = db(self.web_session)

    def validate_deployment_details(self):
        navigate(self.web_session).get("{}/middleware_deployment/show_list".format(self.web_session.MIQ_URL))
        self.ui_utils.waitForTextOnPage("Deployment Name", 15)
        deployments_ui = self.ui_utils.get_list_table()
        assert deployments_ui, "No UI Deployments found."

        deployments_hawk = self.hawkular_api.get_hawkular_deployments()
        assert deployments_hawk, "No Hawkular Deployments found."

        deployments_db = self.db.get_deployments()
        assert deployments_db, "No DB Deployments found."

        # Test only a random sample of the entire Deployments list
        for dep in self.ui_utils.get_random_list(deployments_ui, 3):
            deployment_name = dep.get('Deployment Name')
            self.web_session.logger.info("Validate Deployment {}.".format(deployment_name))

            navigate(self.web_session).get("{}/middleware_deployment/show_list".format(self.web_session.MIQ_URL))
            assert self.ui_utils.waitForTextOnPage("Deployment Name", 15)

            self.ui_utils.click_on_row_containing_text(deployment_name)
            self.ui_utils.waitForTextOnPage("Nativeid", 15)

            dep_details_ui = self.ui_utils.get_generic_table_as_dict()
            assert dep_details_ui, "UI Deployment {} not found.".format(deployment_name)
            self.web_session.logger.info("dep_details_ui: {}".format(dep_details_ui))

            dep_details_hawk = self.ui_utils.find_row_in_list(deployments_hawk, 'Name', deployment_name)
            assert dep_details_hawk, "Hawkular Deployment {} not found.".format(deployment_name)
            self.web_session.logger.info("dep_details_hawk: {}".format(dep_details_hawk))

            dep_details_db = self.ui_utils.find_row_in_list(deployments_db, 'name', deployment_name)
            assert dep_details_db, "DB Deployment {} not found.".format(deployment_name)
            self.web_session.logger.info("dep_details_hawk: {}".format(dep_details_db))

            assert dep_details_ui.get('Middleware Provider') == self.web_session.HAWKULAR_PROVIDER_NAME

            assert deployment_name in dep_details_hawk.get('Name') and deployment_name == dep_details_db.get('name')

            ui_nativeid = dep_details_ui.get('Nativeid')
            assert ui_nativeid == dep_details_hawk.get('Nativeid') and ui_nativeid == dep_details_db.get('nativeid')

        return True

    def validate_deployments_list(self):
        navigate(self.web_session).get("{}/middleware_deployment/show_list".format(self.web_session.MIQ_URL))
        self.ui_utils.waitForTextOnPage("Deployment Name", 15)
        deployments_ui = self.ui_utils.get_list_table()
        assert deployments_ui, "No UI Deployments found."
        deployments_hawk = self.hawkular_api.get_hawkular_deployments()
        assert deployments_hawk, "No Hawkular Deployments found."

        for dep_ui in deployments_ui:
            deployment_name = dep_ui.get('Deployment Name')
            dep_hawk = self.ui_utils.find_row_in_list(deployments_hawk, 'Name', deployment_name)

            assert dep_hawk, "Deployment Name {} not found".format(deployment_name)
            assert (deployment_name in dep_hawk.get("Name")), \
                "Deployment Name mismatch ui:{}, hawk:{}".format(deployment_name, dep_hawk.get("Name"))
            self.web_session.logger.info(
                "UI Deployment name is: {}, and Hawkular deployment is: {} ".format(deployment_name,
                                                                                dep_hawk.get("Name")))

        return True


from selenium.webdriver.common.by import By

class SignInPageLocator(object):
    email = (By.XPATH, '//*[@id="email"]')
    password = (By.XPATH, '//*[@id="pass"]')
    signInBtn = (By.XPATH, '//button[@data-testid="royal_login_button"]')

class common_locators(object):
    def pagination(page_number):
        page_locator = (By.XPATH, f"//a[contains(text(),'{page_number}')]")
        return page_locator

class Daraz(object):
    advertisement = (By.XPATH, '//a[@class="modal-close is-attached is-outer d"]')
    SearchField = (By.XPATH, "//input[@id='q']")
    SearchButton = (By.XPATH, "//button[@class='search-box__button--1oH7']")
    single_product_div = (By.XPATH, '//div[contains(@class,"grid")]')
    product_name_h1 = (By.XPATH, ".//div[contains(@class,'title')]//a")
    product_price_span = (By.XPATH, ".//div[contains(@class,'price')]//span")
    product_image_url = (By.XPATH, ".//div[contains(@class,'mainPic')]//img")

class GadgetAndGear(object):
    advertisement = (By.XPATH, '//a[@class="modal-close is-attached is-outer d"]')
    SearchField = (By.XPATH, "(//input[@id='searchItem'])[1]")
    SearchButton = (By.XPATH, "(//button[@type='submit'])[1]")
    single_product_div = (By.XPATH, '//div[@class="row"]//ul//li')
    product_name_h1 = (By.XPATH, ".//p[@class='product-name d-block mb-0']")
    product_price_span = (By.XPATH, ".//p[contains(@class,'product-price')]")
    product_image_url = (By.XPATH, ".//div[contains(@class,'product-image')]//img")

    @staticmethod
    def pagination(page_number):
        page_locator = (By.XPATH, f"//a[contains(text(),'{page_number}')]")
        return page_locator


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re

from tensorflow.python.client import session
from tensorflow.python.estimator import estimator as estimator_lib
from tensorflow.python.estimator import export as export_lib
from tensorflow.python.estimator import model_fn as model_fn_lib
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
from tensorflow.python.framework import tensor_util
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import models
from tensorflow.python.keras import optimizers
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import metrics as metrics_module
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.training import checkpoint_management
from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.training import optimizer as tf_optimizer_module
from tensorflow.python.training import saver as saver_lib
from tensorflow.python.training import training_util

_DEFAULT_SERVING_KEY = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY

def _cast_tensor_to_floatx(x):
  """Cast tensor to keras's floatx dtype if it is not already the same dtype."""
  if x.dtype == K.floatx():
    return x
  else:
    return math_ops.cast(x, K.floatx())

def _convert_tensor(x):
  """Create or cast tensor if needed."""
  if not tensor_util.is_tensor(x):
    # x is a numpy array
    x = sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(x)
  if check_ops.is_numeric_tensor(x):
    # is_numeric_tensor returns False if provided with a numpy array
    x = _cast_tensor_to_floatx(x)
  return x

def _any_weight_initialized(keras_model):
  """Check if any weights has been initialized in the Keras model.

  Args:
    keras_model: An instance of compiled keras model.

  Returns:
    boolean, True if at least one weight has been initialized, else False.
    Currently keras initialize all weights at get_session().
  """
  if keras_model is None:
    return False
  for layer in keras_model.layers:
    for weight in layer.weights:
      if hasattr(weight, '_keras_initialized'):
        return True
  return False

def _convert_estimator_io_to_keras(keras_model, features, labels):
  """Converts estimator features and labels to keras input and target tensors.

  Args:
    keras_model: a compiled `tf.keras.Model` instance, used to determine the
      order of the returned lists.
    features: Dict of tensors or `None`.
    labels: Dict of tensors, a single tensor, or `None`.

  Returns:
    Tuple of (
      list of input tensors or `None`,
      list of target tensors or `None`)
    The order of tensors is determined by the order set in the keras model.
  """

  def _to_ordered_tensor_list(obj, key_order, obj_name, order_name):
    """Convert obj to an ordered list of tensors.

    Args:
      obj: List, dict, or single tensor. May be `None`.
      key_order: List of strings with the order to return (used if obj is a
        dict).
      obj_name: String name of object (e.g. "features" or "labels")
      order_name: String name of the key order (e.g. "inputs" or "outputs")

    Returns:
      List of tensors, or `None`

    Raises:
      KeyError: If obj has invalid keys.
    """
    if obj is None:
      return None
    elif isinstance(obj, (list, tuple)):
      return [_convert_tensor(x) for x in obj]
    elif isinstance(obj, dict):
      # Ensure that the obj keys and keys in key_order are exactly the same.
      different_keys = set(obj.keys()) ^ set(key_order)

      if different_keys:
        raise KeyError(
            'The dictionary passed into {obj_name} does not have the expected '
            '{order_name} keys defined in the keras model.'
            '\n\tExpected keys: {order_keys}'
            '\n\t{obj_name} keys: {obj_keys}'
            '\n\tDifference: {different_keys}'.format(
                order_name=order_name, order_keys=set(key_order),
                obj_name=obj_name, obj_keys=set(obj.keys()),
                different_keys=different_keys))

      return [_convert_tensor(obj[key]) for key in key_order]
    else:  # Assume obj is a tensor.
      return [_convert_tensor(obj)]

  input_names = None
  output_names = None
  if isinstance(features, dict):
    input_names = (
        keras_model.input_names if keras_model._is_graph_network else
        ['input_%d' % i for i in range(1, len(features) + 1)])
  if isinstance(labels, dict):
    output_names = (
        keras_model.output_names if keras_model._is_graph_network else
        ['output_%d' % i for i in range(1, len(labels) + 1)])

  input_tensors = _to_ordered_tensor_list(
      features, input_names, 'features', 'inputs')
  target_tensors = _to_ordered_tensor_list(
      labels, output_names, 'labels', 'outputs')

  return input_tensors, target_tensors

def _clone_and_build_model(mode,
                           keras_model,
                           custom_objects,
                           features=None,
                           labels=None):
  """Clone and build the given keras_model.

  Args:
    mode: training mode.
    keras_model: an instance of compiled keras model.
    custom_objects: Dictionary for custom objects.
    features: Dict of tensors.
    labels: Dict of tensors, or single tensor instance.

  Returns:
    The newly built model.
  """
  # Set to True during training, False for inference or testing.
  K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
  input_tensors, target_tensors = _convert_estimator_io_to_keras(
      keras_model, features, labels)

  compile_clone = (mode != model_fn_lib.ModeKeys.PREDICT)

  global_step = None
  if compile_clone:
    # Set iterations to the global step created by tf.train.create_global_step()
    # which is automatically run in the estimator framework.
    global_step = training_util.get_or_create_global_step()
    K.track_variable(global_step)

  clone = models.clone_and_build_model(
      keras_model, input_tensors, target_tensors, custom_objects,
      compile_clone=compile_clone,
      in_place_reset=(not keras_model._is_graph_network),
      optimizer_iterations=global_step)

  return clone

def _convert_keras_metrics_to_estimator(model):
  """Convert metrics from a Keras model to ops used by the Estimator framework.

  Args:
    model: A `tf.keras.Model` object.

  Returns:
    Dictionary mapping metric names to tuples of (value, update) ops. May return
    `None` if the model does not contain any metrics.
  """
  if not getattr(model, 'metrics', None):
    return None

  # TODO(psv/fchollet): support stateful metrics
  eval_metric_ops = {}
  # When each metric maps to an output
  if isinstance(model.metrics, dict):
    for i, output_name in enumerate(model.metrics.keys()):
      metric_name = model.metrics[output_name]
      if callable(metric_name):
        metric_name = metric_name.__name__
      # When some outputs use the same metric
      if list(model.metrics.values()).count(metric_name) > 1:
        metric_name += '_' + output_name
      eval_metric_ops[metric_name] = metrics_module.mean(
          model.metrics_tensors[i - len(model.metrics)])
  else:
    for i, metric_name in enumerate(model.metrics):
      if callable(metric_name):
        metric_name = metric_name.__name__
      eval_metric_ops[metric_name] = metrics_module.mean(
          model.metrics_tensors[i])
  return eval_metric_ops

def _create_keras_model_fn(keras_model, custom_objects=None):
  """Creates model_fn for keras Estimator.

  Args:
    keras_model: an instance of compiled keras model.
    custom_objects: Dictionary for custom objects.

  Returns:
    The model_fn for a keras Estimator.
  """

  def model_fn(features, labels, mode):
    """model_fn for keras Estimator."""
    # Raise an error when users use DistributionStrategy with native Keras
    # optimizers. Currently we only support native TensorFlow optimizers.
    if distribution_strategy_context.has_distribution_strategy() and \
        not isinstance(keras_model.optimizer,
                       (tf_optimizer_module.Optimizer, optimizers.TFOptimizer)):
      raise ValueError('Only TensorFlow native optimizers are supported with '
                       'DistributionStrategy.')

    model = _clone_and_build_model(mode, keras_model, custom_objects, features,
                                   labels)
    model_output_names = []
    # We need to make sure that the output names of the last layer in the model
    # is the same for each of the cloned models. This is required for mirrored
    # strategy when we call regroup.
    if distribution_strategy_context.has_distribution_strategy():
      for name in model.output_names:
        name = re.compile(r'_\d$').sub('', name)
        model_output_names.append(name)
    else:
      model_output_names = model.output_names

    # Get inputs to EstimatorSpec
    predictions = dict(zip(model_output_names, model.outputs))

    loss = None
    train_op = None
    eval_metric_ops = None

    # Set loss and metric only during train and evaluate.
    if mode is not model_fn_lib.ModeKeys.PREDICT:
      if mode is model_fn_lib.ModeKeys.TRAIN:
        model._make_train_function()  # pylint: disable=protected-access
      else:
        model._make_test_function()  # pylint: disable=protected-access
      loss = model.total_loss

      eval_metric_ops = _convert_keras_metrics_to_estimator(model)

    # Set train_op only during train.
    if mode is model_fn_lib.ModeKeys.TRAIN:
      train_op = model.train_function.updates_op

    if not model._is_graph_network:
      # Reset model state to original state,
      # to avoid `model_fn` being destructive for the initial model argument.
      models.in_place_subclassed_model_state_restoration(keras_model)
    return model_fn_lib.EstimatorSpec(
        mode=mode,
        predictions=predictions,
        loss=loss,
        train_op=train_op,
        eval_metric_ops=eval_metric_ops,
        export_outputs={
            _DEFAULT_SERVING_KEY:
            export_lib.export_output.PredictOutput(predictions)
        })

  return model_fn

def _save_first_checkpoint(keras_model, custom_objects, config):
  """Save first checkpoint for the keras Estimator.

  Args:
    keras_model: an instance of compiled keras model.
    custom_objects: Dictionary for custom objects.
    config: Estimator config.

  Returns:
    The path where keras model checkpoint is saved.
  """
  # save checkpoint into subdirectory to allow warm start
  keras_model_dir = os.path.join(config.model_dir, 'keras')
  # Load weights and save to checkpoint if there is no checkpoint
  latest_path = checkpoint_management.latest_checkpoint(keras_model_dir)
  if not latest_path:
    keras_weights = None
    if _any_weight_initialized(keras_model):
      keras_weights = keras_model.get_weights()
    if not gfile.IsDirectory(keras_model_dir):
      gfile.MakeDirs(keras_model_dir)
    with ops.Graph().as_default():
      random_seed.set_random_seed(config.tf_random_seed)
      training_util.create_global_step()
      model = _clone_and_build_model(model_fn_lib.ModeKeys.TRAIN, keras_model,
                                     custom_objects)
      # save to checkpoint
      with session.Session(config=config.session_config) as sess:
        if keras_weights:
          model.set_weights(keras_weights)
        # Make update ops and initialize all variables.
        if not model.train_function:
          # pylint: disable=protected-access
          model._make_train_function()
          K._initialize_variables(sess)
          # pylint: enable=protected-access
        saver = saver_lib.Saver()
        latest_path = os.path.join(keras_model_dir, 'keras_model.ckpt')
        saver.save(sess, latest_path)
  return latest_path

def model_to_estimator(keras_model=None,
                       keras_model_path=None,
                       custom_objects=None,
                       model_dir=None,
                       config=None):
  """Constructs an `Estimator` instance from given keras model.

  For usage example, please see:
  [Creating estimators from Keras
  Models](https://tensorflow.org/guide/estimators#model_to_estimator).

  Args:
    keras_model: A compiled Keras model object. This argument is mutually
      exclusive with `keras_model_path`.
    keras_model_path: Path to a compiled Keras model saved on disk, in HDF5
      format, which can be generated with the `save()` method of a Keras model.
      This argument is mutually exclusive with `keras_model`.
    custom_objects: Dictionary for custom objects.
    model_dir: Directory to save `Estimator` model parameters, graph, summary
      files for TensorBoard, etc.
    config: `RunConfig` to config `Estimator`.

  Returns:
    An Estimator from given keras model.

  Raises:
    ValueError: if neither keras_model nor keras_model_path was given.
    ValueError: if both keras_model and keras_model_path was given.
    ValueError: if the keras_model_path is a GCS URI.
    ValueError: if keras_model has not been compiled.
  """
  if not (keras_model or keras_model_path):
    raise ValueError(
        'Either `keras_model` or `keras_model_path` needs to be provided.')
  if keras_model and keras_model_path:
    raise ValueError(
        'Please specity either `keras_model` or `keras_model_path`, '
        'but not both.')

  if not keras_model:
    if keras_model_path.startswith(
        'gs://') or 'storage.googleapis.com' in keras_model_path:
      raise ValueError(
          '%s is not a local path. Please copy the model locally first.' %
          keras_model_path)
    logging.info('Loading models from %s', keras_model_path)
    keras_model = models.load_model(keras_model_path)
  else:
    logging.info('Using the Keras model provided.')
    keras_model = keras_model

  if not hasattr(keras_model, 'optimizer') or not keras_model.optimizer:
    raise ValueError(
        'The given keras model has not been compiled yet. '
        'Please compile the model with `model.compile()` '
        'before calling `model_to_estimator()`.')

  config = estimator_lib.maybe_overwrite_model_dir_and_session_config(config,
                                                                      model_dir)

  keras_model_fn = _create_keras_model_fn(keras_model, custom_objects)
  if _any_weight_initialized(keras_model):
    # Warn if config passed to estimator tries to update GPUOptions. If a
    # session has already been created, the GPUOptions passed to the first
    # session sticks.
    if config.session_config.HasField('gpu_options'):
      logging.warning(
          'The Keras backend session has already been set. '
          'The _session_config passed to model_to_estimator will not be used.')
  else:
    # Pass the config into keras backend's default session.
    sess = session.Session(config=config.session_config)
    K.set_session(sess)

  warm_start_path = None
  if keras_model._is_graph_network:
    warm_start_path = _save_first_checkpoint(keras_model, custom_objects,
                                             config)
  elif keras_model.built:
    logging.warning('You are creating an Estimator from a Keras model manually '
                    'subclassed from `Model`, that was already called on some '
                    'inputs (and thus already had weights). We are currently '
                    'unable to preserve the model\'s state (its weights) as '
                    'part of the estimator in this case. Be warned that the '
                    'estimator has been created using a freshly initialized '
                    'version of your model.\n'
                    'Note that this doesn\'t affect the state of the model '
                    'instance you passed as `keras_model` argument.')

  estimator = estimator_lib.Estimator(keras_model_fn,
                                      config=config,
                                      warm_start_from=warm_start_path)

  return estimator


from transformers.models.clip.modeling_clip import (
    CLIPEncoderLayer,
    CLIPTextTransformer,
    CLIPVisionTransformer,
)

from parallelformers.policies.base import Layer, Policy
from parallelformers.transformers.modeling_clip import CLIPAttention_
from parallelformers.utils import AllReduceLinear

class CLIPTextPolicy(Policy):
    @staticmethod
    def replace_arguments(config, world_size):
        arguments = {
            f"encoder.layers[{i}].self_attn.embed_dim": config.text_config.hidden_size
            // world_size
            for i in range(config.text_config.num_hidden_layers)
        }

        arguments.update(
            {
                f"encoder.layers[{i}].self_attn.num_heads": config.text_config.num_attention_heads
                // world_size
                for i in range(config.text_config.num_hidden_layers)
            }
        )

        return arguments

    @staticmethod
    def original_layer_class():
        return CLIPTextTransformer

class CLIPVisionPolicy(Policy):
    @staticmethod
    def replace_arguments(config, world_size):
        arguments = {
            f"encoder.layers[{i}].self_attn.embed_dim": config.vision_config.hidden_size
            // world_size
            for i in range(config.vision_config.num_hidden_layers)
        }

        arguments.update(
            {
                f"encoder.layers[{i}].self_attn.num_heads": config.vision_config.num_attention_heads
                // world_size
                for i in range(config.vision_config.num_hidden_layers)
            }
        )

        return arguments

    @staticmethod
    def original_layer_class():
        return CLIPVisionTransformer

class CLIPLayerPolicy(Policy):
    @staticmethod
    def replace_modules():
        return {
            "CLIPAttention": CLIPAttention_,
        }

    @staticmethod
    def attn_qkv():
        return [
            Layer(
                weight="self_attn.q_proj.weight",
                bias="self_attn.q_proj.bias",
            ),
            Layer(
                weight="self_attn.k_proj.weight",
                bias="self_attn.k_proj.bias",
            ),
            Layer(
                weight="self_attn.v_proj.weight",
                bias="self_attn.v_proj.bias",
            ),
        ]

    @staticmethod
    def attn_out():
        return [
            Layer(
                weight="self_attn.out_proj.weight",
                bias="self_attn.out_proj.bias",
                replace=AllReduceLinear,
            ),
        ]

    @staticmethod
    def mlp_in():
        return [
            Layer(
                weight="mlp.fc1.weight",
                bias="mlp.fc1.bias",
            ),
        ]

    @staticmethod
    def mlp_out():
        return [
            Layer(
                weight="mlp.fc2.weight",
                bias="mlp.fc2.bias",
                replace=AllReduceLinear,
            ),
        ]

    @staticmethod
    def original_layer_class():
        return CLIPEncoderLayer


import logging
import unittest
import numpy as np
from os.path import join, dirname
from chmpy.crystal import Crystal, UnitCell
from tempfile import TemporaryDirectory

LOG = logging.getLogger(__name__)
_ICE_II = join(dirname(__file__), "iceII.cif")

class UnitCellTestCase(unittest.TestCase):
    def test_unit_cell_lattice(self):
        c = UnitCell.cubic(2.0)
        np.testing.assert_allclose(c.lattice, 2.0 * np.eye(3), atol=1e-8)
        np.testing.assert_allclose(c.reciprocal_lattice, 0.5 * np.eye(3), atol=1e-8)

    def test_coordinate_transforms(self):
        c = UnitCell.cubic(2.0)
        np.testing.assert_allclose(
            c.to_fractional(np.eye(3)), 0.5 * np.eye(3), atol=1e-8
        )
        np.testing.assert_allclose(c.to_cartesian(np.eye(3)), 2 * np.eye(3), atol=1e-8)

    def test_handle_bad_angles(self):
        # should warn
        c = UnitCell.from_lengths_and_angles([2.0] * 3, [90] * 3)

    def test_repr(self):
        c = UnitCell.cubic(2.0)
        self.assertTrue(str(c) == "<UnitCell: cubic (2.000)>")

    def test_cell_types(self):
        cubic = UnitCell.cubic(2.0)
        orthorhombic = UnitCell.orthorhombic(3.0, 4.0, 5.0)
        tetragonal = UnitCell.tetragonal(3.0, 4.0, unit="degrees")
        tetragonal = UnitCell.from_unique_parameters((3.0, 4.0), cell_type="tetragonal")
        monoclinic = UnitCell.monoclinic(3.0, 4.0, 5.0, 75, unit="degrees")
        monoclinic = UnitCell.monoclinic(3.0, 4.0, 5.0, 1.5)
        triclinic = UnitCell.triclinic(3.0, 4.0, 5.0, 45, 75, 90, unit="degrees")
        rhombohedral = UnitCell.rhombohedral(3.0, 97, unit="degrees")
        rhombohedral = UnitCell.rhombohedral(3.0, 1.35)
        hexagonal = UnitCell.hexagonal(3.0, 4.0)
        hexagonal = UnitCell.from_lengths_and_angles(
            [3.0, 3.0, 5.0], [90, 90, 120], unit="degrees"
        )
        self.assertTrue(cubic.cell_type == "cubic")
        self.assertTrue(orthorhombic.cell_type == "orthorhombic")
        self.assertTrue(tetragonal.cell_type == "tetragonal")
        self.assertTrue(monoclinic.cell_type == "monoclinic")
        self.assertTrue(triclinic.cell_type == "triclinic")
        self.assertTrue(rhombohedral.cell_type == "rhombohedral")
        self.assertTrue(hexagonal.cell_type == "hexagonal")

    def test_angles(self):
        c = UnitCell.cubic(2.0)
        self.assertFalse(c.angles_different)
        self.assertTrue(c.is_cubic)
        self.assertFalse(c.is_triclinic)
        self.assertFalse(c.is_monoclinic)
        self.assertFalse(c.is_tetragonal)
        self.assertFalse(c.is_rhombohedral)
        self.assertFalse(c.is_hexagonal)
        self.assertFalse(c.is_orthorhombic)

        np.testing.assert_allclose([c.alpha_deg, c.beta_deg, c.gamma_deg], 90.0)
        np.testing.assert_allclose(c.parameters, [2, 2, 2, 90, 90, 90])


import numpy as np
execfile('/mnt/data2/rumbaugh/EVLA/11A-138/scripts/utils_11A-138.py')

set_obs_arr()

FILE = open('/mnt/data2/rumbaugh/EVLA/11A-138/scripts/obsreflist.txt','w')

for EorL in ['Early','Late']:
    for groupnum in np.sort(EVLA_obs_dict[EorL].keys()):
        for SBnum in np.sort(EVLA_obs_dict[EorL][groupnum].keys()):
            FILE.write('%5s %i %2s\n'%(EorL,groupnum,str(SBnum)))
FILE.close()


from dfa import DFA

class NFA:
    def __init__(self, states, input_symbols, transitions, initial_state, final_states):
        self.states = states
        self.input_symbols = input_symbols
        self.transitions = transitions
        self.initial_state = initial_state
        self.final_states = final_states

        # nfa = NFA(
        #     states={'q0', 'q1', 'q2'},
        #     input_symbols={'0', '1'},
        #     transitions={
        #         'q0': {'0': [q0], '': [q1]},
        #         'q1': {'': [q0,q2], '1': [q2]},
        #         'q2': {'0': [q2], '1': [q1]}
        #     },
        #     initial_state='q0',
        #     final_states={'q1'}
        # )

    def eliminate_nondeterminism(self):

        # all symbols must be in each state transitions dictionary
        for key in self.transitions.keys():
            for symbol in [""] + self.input_symbols:
                if symbol not in self.transitions[key].keys():
                    self.transitions[key][symbol] = []

        # eliminate lambda-transitions
        for state in self.transitions.keys():
            for symbol in self.input_symbols:
                states_set = set()
                for s in NFA.lambda_closure(set([state]), self.transitions):
                    states_set = states_set.union(set(self.transitions[s][symbol]))
                states_set = NFA.lambda_closure(states_set, self.transitions)
                self.transitions[state][symbol] = list(states_set)

        # remove lambda transitions from transitions
        for state in self.transitions.keys():
            self.transitions[state].pop("")

        # Using the Subset Construction to eliminate Nondeterminism
        is_finished = False
        while not is_finished:
            is_finished = True
            to_be_added = {}
            for key in self.transitions.keys():
                for symbol in self.input_symbols:
                    if "".join(self.transitions[key][symbol]) not in self.states:
                        is_finished = False
                        # if the destination set is empty create a "empty set" state
                        if "".join(str(self.transitions[key][symbol])) == "":
                            new_state = ""
                            self.states.append(new_state)
                            to_be_added[new_state] = {}
                            for sym in self.input_symbols:
                                to_be_added[new_state][sym] = ""
                        # if the destination set is not empty create a new state corresponding to that set
                        else:
                            new_state = "".join(self.transitions[key][symbol])
                            self.states.append(new_state)
                            to_be_added[new_state] = {}
                            for sym in self.input_symbols:
                                destinations = []
                                for start in self.transitions[key][symbol]:
                                    destinations = destinations + list(
                                        self.transitions[start][sym]
                                    )
                                to_be_added[new_state][sym] = list(set(destinations))

            # Merge initial and new states
            self.transitions = {**self.transitions, **to_be_added}

        # change destinations type from set to string
        for key in self.transitions.keys():
            for symbol in self.input_symbols:
                if len(self.transitions[key][symbol]) == 0:
                    self.transitions[key][symbol] = "phi"
                else:
                    self.transitions[key][symbol] = "".join(
                        self.transitions[key][symbol]
                    )

        # replace state '' with state 'phi'
        self.transitions["phi"] = self.transitions.pop("")
        self.states.remove("")
        self.states.append("phi")

        # find final states
        for final_state in self.final_states:
            for state in self.states:
                if final_state in state and state not in self.final_states:
                    self.final_states.append(state)

        # return minimized DFA
        return DFA(
            self.states,
            self.input_symbols,
            self.transitions,
            self.initial_state,
            self.final_states,
        ).minimize()

    def lambda_closure(states_set, transitions):
        # returns the lambda-closure set of a set of states
        visited = set()
        while len(states_set) > 0:
            state = states_set.pop()
            visited.add(state)
            for s in transitions[state][""]:
                if s not in visited:
                    states_set.add(s)
        return visited

    def __str__(self):
        print("nfa: ")
        for i in self.transitions:
            print("\t" + i + ": ", self.transitions[i])
        print("initial states:", self.initial_state)
        print("final states", self.final_states)
        return ""


import scipy.special as ssp
import numpy as np
import scipy.stats as ss
import random as rnd

Zt = lambda p, zt, n: np.array(
    [0 if i > 2 * p else (zt if i < p else -zt) for i in (rnd.random() for i in range(n))])  # szum dyskretny

autocov = lambda X, lag: [
    1 / len(X) * (np.sum((np.array(X[:len(X) - i]) - np.mean(X)) * (np.array(X[i:]) - np.mean(X)))) for i in
    range(lag[0], lag[1] + 1)]  # klasyczna autokowariancja

autocorr = lambda X, lag: [i / autocov(X, [0, 1])[0] for i in autocov(X, lag)]  # klasyczna autokorelacja

def Q(X, k):
    values = [abs(X[j] - X[i]) for i in range(len(X)) for j in range(i)]
    values = sorted(values)
    return k * values[int(np.floor((ssp.binom(len(X), 2) + 2) / 4))]

def autocov1(X, lag, k):
    acvf = []
    for h in range(lag[0], lag[1] + 1):
        if h == 0:
            u = np.array(X)
        else:
            u = np.array(X[:-h])
        v = np.array(X[h:])
        acvf.append(1 / 4 * ((Q(u + v, k)) ** 2 - (Q(u - v, k)) ** 2))
    return acvf

def autocorr1(X, lag, k):
    acf = []
    for h in range(lag[0], lag[1] + 1):
        if h == 0:
            u = np.array(X)
        else:
            u = np.array(X[:-h])
        v = np.array(X[h:])
        acf.append(((Q(u + v, k)) ** 2 - (Q(u - v, k)) ** 2) / ((Q(u + v, k)) ** 2 + (Q(u - v, k)) ** 2))
    return acf

g = lambda alfa, n: int(np.floor(n * alfa))

def L(X, t, alfa):
    Y = sorted(X)
    n = len(X)
    _g = g(alfa, n)
    if Y[_g] < X[t] and X[t] < Y[n - _g]:
        return 1
    return 0

def mean_X(X, alfa):
    suma1 = sum([L(X, i, alfa) for i in range(len(X))])
    suma2 = sum([X[i] * L(X, i, alfa) for i in range(len(X))])
    return 1 / suma1 * suma2

def autocov2(X, lag, alfa):
    n = len(X)
    L_vector = [L(X, i, alfa) for i in range(n)]
    acvf = []
    for h in range(lag[0], lag[1] + 1):
        if h == 0:
            suma1 = sum(L_vector)
            suma2 = sum(np.array(X) ** 2 * np.array(L_vector))
        else:
            suma1 = sum(np.array(L_vector[:-h]) * np.array(L_vector[h:]))
            suma2 = sum(np.array(X[:-h] - mean_X(X, alfa)) * np.array(X[h:] - mean_X(X, alfa)) * np.array(
                L_vector[:-h]) * np.array(L_vector[h:]))
        acvf.append(suma2 / suma1)
    return acvf

def autocorr2(X, lag, alfa):
    acvf = autocov2(X, lag, alfa)[0]
    return [i / acvf for i in autocov2(X, lag, alfa)]

def autocorr4(X, lag):
    X = np.array(X) - np.mean(X)
    acf = [1.0]
    denominator = np.median([i**2 for i in X])
    for h in range(lag[0] + 1, lag[1] + 1):
        acf.append(np.median(np.array(X[:-h]) * np.array(X[h:]))/denominator)
    return acf

def J(x):
    return ss.norm.ppf(x)

def c(X):
    n = len(X)
    R = sorted(X)
    return 1 / np.sum([(J((list(X).index(R[i]) + 1) / (n + 1)))**2 for i in range(n)])

def autocorr3(X, lag):
    acf = []
    n = len(X)
    R = sorted(X)
    for h in range(lag[0], lag[1] + 1):
        suma = 0
        for i in range(n-h):
            suma += J((R.index(list(X)[i]) + 1)/(n+1)) * J((R.index(list(X)[i + h]) + 1)/ (n+1))
        acf.append(c(X) * suma)
    return acf

def autocorr5(X, lag):
    n = len(X)
    return [2/((n - h)*(n - h - 1)) * np.sum([np.sign((X[i] - X[j]) * (X[i + h] - X[j + h]))
                                              for i in range(n - h) for j in range(i)]) for h in range(lag[0], lag[1] + 1)]

def autocorr6(X, lag):
    n = len(X)
    mediana = np.median(X)
    acf = [1/(n - h) * np.sum(np.sign((np.array(X[:-h]) - mediana) *
                                      (np.array(X[h:]) - mediana))) for h in range(lag[0] + 1, lag[1] + 1)]
    return [1, *acf]

def zmatrix(X, k):
    n = len(X)
    # Zk = np.zeros((k + 1, n + k))
    # for j in range(k + 1):
    #     Zk[j, j:j + n - 1] = X
    # return Zk
    return np.matrix([[*[0 for i in range(j)], *[X[i] for i in range(n)], *[0 for i in range(k - j)]] for j in range(k+1)])

def gamma_k(X, k):
    z_matr = zmatrix(X, k)
    return (z_matr * z_matr.transpose())/len(X)

def Xi(X, k, i, j):
    gamm = gamma_k(X, k)
    return gamm[i - 1, j - 1]/np.sqrt(gamm[i - 1, i - 1] * gamm[j - 1, j - 1])

def autocorr7(X, lag):
    k = lag[1]
    return [1/(k - h + 1)*sum([Xi(X, k, i, i + h) for i in range(1, k - h + 2)]) for h in range(lag[0], lag[1] + 1)]

def autocorr8(X, lag):
    X = np.array(X)
    return [(np.var(X+X) - np.var(X - X))/(np.var(X+X) + np.var(X - X)) if h == 0 else
            (np.var(X[h:]+X[:-h]) - np.var(X[h:] - X[:-h]))/(np.var(X[h:]+X[:-h]) + np.var(X[h:] - X[:-h]))
            for h in range(lag[0], lag[1] + 1)]

def MAD2(X, const=1.4826):
    return (const * np.median([abs(i - np.median(X)) for i in X]))**2

def IQR2(X, const=0.7413):
    return (const * (np.quantile(X, 0.75) - np.quantile(X, 0.25)))**2

def autocorr9(X, lag):
    X = np.array(X)
    return [(MAD2(X+X) - MAD2(X - X))/(MAD2(X+X) + MAD2(X - X)) if h == 0 else
            (MAD2(X[h:]+X[:-h]) - MAD2(X[h:] - X[:-h]))/(MAD2(X[h:]+X[:-h]) + MAD2(X[h:] - X[:-h]))
            for h in range(lag[0], lag[1] + 1)]

def autocorr10(X, lag):
    X = np.array(X)
    return [(IQR2(X+X) - IQR2(X - X))/(IQR2(X+X) + IQR2(X - X)) if h == 0 else
            (IQR2(X[h:]+X[:-h]) - IQR2(X[h:] - X[:-h]))/(IQR2(X[h:]+X[:-h]) + IQR2(X[h:] - X[:-h]))
            for h in range(lag[0], lag[1] + 1)]


from pathlib import Path

from Bio.SVDSuperimposer import SVDSuperimposer
import numpy as np
import pandas as pd

from protein import Protein, AverageProtein
from utils import rotate_points, get_shared_indices, get_mutation_position

class Deformation:
    def __init__(self, protein_1, protein_2, **kwargs):
        self.prot1 = protein_1
        self.prot2 = protein_2
        self.proteins = [self.prot1, self.prot2]

        self.default_method = ["strain"]
        self.all_methods = ["mut_dist", "strain", "shear", "non-affine", "ldd", "lddt", "neighborhood_dist", "rmsd"]
        self.lddt_cutoffs = kwargs.get("lddt_cutoffs", [0.5, 1, 2, 4])
        self.method = kwargs.get('method', self.default_method.copy())
        self.neigh_cut = kwargs.get('neigh_cut', 13.0)

        self.force_cutoff =  kwargs.get('force_cutoff', False)
        self.force_norm =  kwargs.get('force_norm', False)
        self.force_nonorm =  kwargs.get('force_nonorm', False)
        self.force_relative =  kwargs.get('force_relative', False)
        self.force_absolute =  kwargs.get('force_absolute', False)
        self.force_l1norm =  kwargs.get('force_l1norm', False)
        self.force_l2norm =  kwargs.get('force_l2norm', False)

        self.deformation = {}
        self.mutations = []
        self.mutation_idx = []

        self.parse_input()
        self.parse_method()

    def parse_input(self):
        # Check input types
        for prot in self.proteins:
            if not isinstance(prot, (AverageProtein, Protein)):
                raise Exception(f"Input object type {type(prot)} is not supported.")

        # Check that neighbor cutoff definitions are consistent.
        # This method also loads neighborhoods if they are not loaded already.
        self._check_neighborhoods()

        # Check that proteins are the same length
        if self.prot1.seq_len != self.prot2.seq_len:
            print("WARNING! Protein sequences are not the same length: " + \
                  f"Protein A has {self.prot1.seq_len} residues, while " + \
                  f"Protein B has {self.prot2.seq_len} residues.")

    def _update_protein_neighborhood(self, prot, neigh_cut):
        # If not calculated yet... calculate
        if len(prot.neigh_idx) == 0:
            prot.neigh_cut = neigh_cut
            prot.get_local_neighborhood()

        # If neigh cut is wrong... calculate
        elif prot.neigh_cut != neigh_cut:
            print(f"Recalculating Protein with new neighbor cutoff = {neigh_cut}")
            prot.neigh_cut = neigh_cut
            prot.get_local_neighborhood()

    def _update_averageProtein_neighborhood(self, prot, neigh_cut):
        # If not calculated yet... calculate
        if len(prot.neigh_idx) == 0:
            prot.neigh_cut = neigh_cut
            prot.get_average_structure()

        # If neigh cut is wrong... calculate
        elif prot.neigh_cut != neigh_cut:
            print(f"Recalculating AverageProtein with new neighbor cutoff = {neigh_cut}")
            prot.neigh_cut = neigh_cut
            prot.recalculate_average_structure()

    # Check for consistency in the use of neighbor cutoffs
    def _check_neighborhoods(self):
        neigh_cut = [prot.neigh_cut for prot in self.proteins]

        # Check for consistency between Protein objects.
        # If they are inconsistent, either force them to use Deformation.neigh_cutoff,
        # or exit
        if neigh_cut[0] != neigh_cut[1]:
            if self.force_cutoff:
                for prot in self.proteins:
                    if isinstance(prot, Protein):
                        self._update_protein_neighborhood(prot, self.neigh_cut)
                    if isinstance(prot, AverageProtein):
                        self._update_averageProtein_neighborhood(prot, self.neigh_cut)
            else:
                raise Exception("AverageProtein / Protein objects were created with different neighbor cutoffs!" + \
                                "\n\tYou need to use the same neighbor cutoff for each structure," + \
                                "\n\tor to automatically recalculate neighborhoods using Deformation.neigh_cutoff," +\
                                " use Deformation(..., force_cutoff=True)")
            return

        # If not "force_cutoff", then ensure Deformation.neigh_cutoff equals that of the Protein objects
        if not self.force_cutoff:
            if self.neigh_cut != neigh_cut[0]:
                self.neigh_cut = neigh_cut[0]
                print(f"WARNING! Resetting neighbour cutoff to {self.neigh_cut}, since this " + \
                       "value was used for the AverageProtein structure." + \
                       "\n\tTo override this, use Deformation(..., force_cutoff=True)")

        for i, prot in enumerate(self.proteins):
            if isinstance(prot, Protein):
                self._update_protein_neighborhood(prot, self.neigh_cut)
            if isinstance(prot, AverageProtein):
                self._update_averageProtein_neighborhood(prot, self.neigh_cut)
            if self.force_cutoff:
                for i, prot in enumerate(self.proteins):
                    prot.neigh_cut = self.neigh_cut
                    prot.recalculate_average_structure()
            else:
                self.neigh_cut = neigh_cut[0]

    # Parse method, and ensure methods are acceptable
    def parse_method(self):
        if isinstance(self.method, (list, np.ndarray, set, tuple)):
            if len(self.method) == 1:
                if self.method[0] == 'all':
                    self.method = self.all_methods.copy()

            else:
                method_list = []
                for method in self.method:
                    if method in self.all_methods:
                        method_list.append(method)
                    else:
                        print(f"WARNING! {method} is not an accepted method")
                self.method = method_list
        else:
            self.method = []

        if len(self.method) == 0:
            raise Exception("No acceptable method found!")

    # Set method, and run through parse to check method validity
    def set_method(self, value):
        self.method = value
        self.parse_method()

    def save_output(self, path_out):
        # Load any deformation that was calculated
        deform = {}
        for m in self.all_methods:
            if hasattr(self, m):
                if m != 'rmsd':
                    deform[m] = getattr(self, m)
                else:
                    # Since RMSD is a scalar, we create a list to match the output format
                    deform[m] = [getattr(self, m)] * self.prot1.seq_len

        # Only save output if deformation was calculated
        if not len(deform):
            raise Exception("ERROR! Cannot save output if there is none!")

        # Add sequence indices, residue names
        else:
            type_names = ['residue_index', 'protA_resname', 'protB_resname']
            res_data = [self.prot1.idx, self.prot1.sequence, self.prot2.sequence]
            output = {k: d for k, d in zip(type_names, res_data) if not isinstance(d, type(None))}
            ########################################
            ### NEED TO ADD NUMBER OF NEIGHBORS
            ########################################

            output.update(deform)

            pd.DataFrame(output).to_csv(path_out, index=False)

    def calculate_deformation(self):
        # Calculate deformation based on the specified method
        self.deformation = {}
        for method in self.method:
            self._run_analysis(method)

    def _run_analysis(self, method):
        analyses = {'mut_dist': self.calculate_dist_from_mutation,
                    'strain': self.calculate_strain,
                    'shear': self.calculate_shear,
                    'non-affine': self.calculate_non_affine,
                    'ldd': self.calculate_ldd,
                    'lddt': self.calculate_lddt,
                    'neighborhood_dist': self.calculate_neighborhood_dist,
                    'rmsd': self.calculate_rmsd}
        analyses[method]()

    # Calculate distance from closest mutation
    def calculate_dist_from_mutation(self):
        try:
            self.sub_pos = get_mutation_position(self.prot1.sequence, self.prot2.sequence)
        except AttributeError as E:
            raise AttributeError("Sequence is not defined for Protein object")

        # If none differ, then return np.nan
        if not len(self.sub_pos):
            print("WARNING! Trying to calculate distance from mutation, while comparing identical sequences")
            return np.zeros(self.prot1.seq_len) * np.nan

        # Calculate mindist using the full array (inc. nan)
        mut_dist1 = self.prot1.dist_mat[:,self.sub_pos]
        mut_dist2 = self.prot2.dist_mat[:,self.sub_pos]

        # Average the distance across both structures,
        # and get the minimum distance per residue to a mutated position
        self.mut_dist = np.nanmin(0.5 * (mut_dist1 + mut_dist2), axis=1)

    def calculate_lddt(self):
        lddt = np.zeros(self.prot1.seq_len, float) * np.nan
        for i in range(self.prot1.seq_len):
            # If no data for residue, leave np.nan
            if (not len(self.prot1.neigh_idx[i])) | (not len(self.prot2.neigh_idx[i])):
                continue

            # Get shared indices
            i1, i2 = get_shared_indices(self.prot1.neigh_idx[i], self.prot2.neigh_idx[i])

            # If no shared indices, leave np.nan
            if not len(i1):
                continue

            # Get local distance vectors
            v1 = np.linalg.norm(self.prot1.neigh_tensor[i][i1], axis=1)
            v2 = np.linalg.norm(self.prot2.neigh_tensor[i][i2], axis=1)

            # Get local distance difference vector
            dv = v2 - v1

            lddt[i] = np.sum([np.sum(dv<=cut) for cut in self.lddt_cutoffs]) / (len(self.lddt_cutoffs) * len(dv))

        self.lddt = lddt

    def calculate_ldd(self):
        ldd = np.zeros(self.prot1.seq_len, float) * np.nan
        for i in range(self.prot1.seq_len):
            # If no data for residue, leave np.nan
            if (not len(self.prot1.neigh_idx[i])) | (not len(self.prot2.neigh_idx[i])):
                continue

            # Get shared indices
            i1, i2 = get_shared_indices(self.prot1.neigh_idx[i], self.prot2.neigh_idx[i])

            # If no shared indices, leave np.nan
            if not len(i1):
                continue

            # Get local distance vectors
            v1 = np.linalg.norm(self.prot1.neigh_tensor[i][i1], axis=1)
            v2 = np.linalg.norm(self.prot2.neigh_tensor[i][i2], axis=1)

            # Get local distance difference vector
            dv = v2 - v1

            if self.force_relative:
                dv = dv / v1

            # Calculate local distance difference
            if self.force_norm:
                ldd[i] = np.linalg.norm(dv) / len(i1)
            else:
                ldd[i] = np.linalg.norm(dv)

        self.ldd = ldd

    def calculate_neighborhood_dist(self):
        nd = np.zeros(self.prot1.seq_len, float) * np.nan
        for i in range(self.prot1.seq_len):
            # If no data for residue, leave np.nan
            if (not len(self.prot1.neigh_idx[i])) | (not len(self.prot2.neigh_idx[i])):
                continue

            # Get shared indices
            i1, i2 = get_shared_indices(self.prot1.neigh_idx[i], self.prot2.neigh_idx[i])

            # If no shared indices, leave np.nan
            if not len(i1):
                continue

            # Get neighbourhood tensors
            c1 = self.prot1.neigh_tensor[i][i1]
            c2 = self.prot2.neigh_tensor[i][i2]

            # Rotate neighbourhood tensor and calculate Euclidean distance
            if self.force_norm:
                nd[i] = np.linalg.norm(rotate_points(c2, c1) - c1) / len(i1)
            else:
                nd[i] = np.linalg.norm(rotate_points(c2, c1) - c1)

        self.neighbor_distance = nd

    def calculate_shear(self):
        shear = np.zeros(self.prot1.seq_len, float) * np.nan
        for i in range(self.prot1.seq_len):
            # If no data for residue, leave np.nan
            if (not len(self.prot1.neigh_idx[i])) | (not len(self.prot2.neigh_idx[i])):
                continue

            # Get shared indices
            i1, i2 = get_shared_indices(self.prot1.neigh_idx[i], self.prot2.neigh_idx[i])

            # If no shared indices, leave np.nan
            if not len(i1):
                continue

            # Get neighbourhood tensors
            u1 = self.prot1.neigh_tensor[i][i1]
            u2 = self.prot2.neigh_tensor[i][i2]
            try:
                duu = u2 @ u2.T - u1 @ u1.T
                uu = np.linalg.inv(u1.T @ u1)
                C = 0.5 * (uu @ u1.T @ duu @ u1 @ uu)
                shear[i] = 0.5 * np.sum(np.diag(C@C) - np.diag(C)**2)
            except Exception as e:
                continue

        self.shear = shear

    def calculate_strain(self):
        strain = np.zeros(self.prot1.seq_len, float) * np.nan
        for i in range(self.prot1.seq_len):
            # If no data for residue, leave np.nan
            if (not len(self.prot1.neigh_idx[i])) | (not len(self.prot2.neigh_idx[i])):
                continue

            # Get shared indices
            i1, i2 = get_shared_indices(self.prot1.neigh_idx[i], self.prot2.neigh_idx[i])

            # If no shared indices, leave np.nan
            if not len(i1):
                continue

            # Get neighbourhood tensors
            c1 = self.prot1.neigh_tensor[i][i1]
            c2 = self.prot2.neigh_tensor[i][i2]

            # Rotate neighbourhood tensor and calculate Euclidean distance
            c3 = rotate_points(c2, c1)
            s = 0.0
            for j in range(len(c1)):
                if self.force_absolute:
                    s += np.linalg.norm(c3[j] - c1[j])
                else:
                    s += np.linalg.norm(c3[j] - c1[j]) / np.linalg.norm(c1[j])

            if not self.force_nonorm:
                s /= len(i1)

            strain[i] = s

        self.strain = strain

    def calculate_non_affine(self):
        pass

    #######################################################
    ### Root-Mean-Square Deviation
    ### Superimpose structures and calculate RMSD
    def calculate_rmsd(self):
        # Cannot calculate RMSD with averaged neighborhoods,
        # so if an AverageProtein object is passsed, we simply
        # take the first Protein object
        coords = []
        for prot in self.proteins:
            if isinstance(prot, Protein):
                coords.append(prot.coord)
            elif isinstance(prot, AverageProtein):
                coords.append(prot.proteins[0].coord)
                print("WARNING! Trying to calculate RMSD with an AverageProtein object. " + \
                      "Since this is not possible, an embedded Protein object is used instead.")
        c1, c2 = coords

        # Remove NAN values
        idx = ~np.any(np.isnan(c1) | np.isnan(c2), axis=1)
        c1, c2 = c1[idx], c2[idx]

        sup = SVDSuperimposer()
        sup.set(c1, c2)
        sup.run()
        self.rmsd = sup.get_rms()


"""
  Code for working with GTSRB dataset.

  Note that these are RGB images; however, we convert to grayscale for compatability
  with an anticipated target platform.

  References:
     http://benchmark.ini.rub.de/?section=gtsrb&subsection=dataset
"""

import os
import numpy as np

from subimage import Subimage

def default_train_test_split(si):
    """ GTSRB has a train/test set; however, the test data comes without labels.
        This function can be used to split up the training set (e.g. into train and validation).

        si : a Subimage object
    """
    np.random.seed(1066)
    return si.train_test_split_by_group(0.2, max_per_class=None)

def parse_train(train_dir):
    """ The training data is organized as one subdirectory per class.
    """
    train_dir = os.path.expanduser(train_dir)
    all_subdirs = [f for f in os.listdir(train_dir) if os.path.isdir(os.path.join(train_dir,f))]

    # process each subdirectory
    si = Subimage(to_grayscale=True)
    for subdir in all_subdirs:
        csvfile = os.path.join(train_dir, subdir, 'GT-%s.csv' % subdir)
        _parse_annotations(csvfile, si)

    return si

def parse_test(test_csv):
    """  The test data is not organized into subirectories like the training.
         Also, the filenames are different.
    """
    si = Subimage(to_grayscale=True)
    _parse_annotations(test_csv, si)
    return si

def _parse_annotations(csvfile, si):
    """   Adds data from a directory to a subimage object
    """

    # annotations are stored in a master CSV (well, semicolon-delimited) file
    csvfile = os.path.expanduser(csvfile)
    csv = open(csvfile, 'r')
    csv.readline() # discard header
    csv = csv.readlines()

    # base path to actual filenames.
    base_path = os.path.dirname(csvfile)

    # do it
    for idx, line in enumerate(csv):
        fields = line.split(';')

        im_filename = fields[0]
        width, height = int(fields[1]), int(fields[2])
        x0, y0 = int(fields[3]), int(fields[4])
        x1, y1 = int(fields[5]), int(fields[6])
        bbox = [x0,y0,x1,y1]

        # test data has no labels
        if len(fields) > 7:
            y = int(fields[7])
        else:
            y = -1

        # Training data filenames have track and frame information.
        # Test data does not.
        if '_' in im_filename:
            # toss the .ppm and pull out track and frame id
            track_id, frame_id = im_filename[:-5].split('_')

            # The group id is a combination of the track id and the class label.
            gid = track_id + '_' + str(y)
        else:
            # There is no notion of a "track" in the test data
            gid = None

        si.add(os.path.join(base_path, im_filename), bbox, y, gid)

    return si