import torch
import torch.nn as nn
import losses
import math
import argparse
import numpy as np
import os
from tqdm import tqdm
import torch.nn.functional as F
from utils import init_weights
import torchvision
from utils import *
import torch.utils.model_zoo as model_zoo
from torch import nn
class ArcMarginProduct(nn.Module):
    r"""Implement of large margin arc distance: :
        Args:
            in_features: size of each input sample
            out_features: size of each output sample
            s: norm of input feature
            m: margin
            cos(theta + m)
        """
    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
        super(ArcMarginProduct, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        self.easy_margin = easy_margin
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m

    def forward(self, input, label):
        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
        sine = torch.sqrt((1.0 - torch.pow(cosine, 2)).clamp(0, 1))
        phi = cosine * self.cos_m - sine * self.sin_m
        if self.easy_margin:
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
        # --------------------------- convert label to one-hot ---------------------------
        # one_hot = torch.zeros(cosine.size(), requires_grad=True, device='cuda')
        one_hot = torch.zeros(cosine.size(), device='cuda')
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)  # you can use torch.where if your torch.__version__ is 0.4
        output *= self.s
        # print(output)

        return output

class DSAN(nn.Module):
    def __init__(self, num_classes=31, bottle_neck=True,max_iter=1000):
        super(DSAN, self).__init__()
        self.feature_layers = ResBase50()
        self.lmmd_loss = losses.ori_LMMD_loss(class_num=num_classes)
        self.bottle_neck = bottle_neck
        if bottle_neck:
            bottleneck_dim = 256
            self.bottle = nn.Linear(2048, 256)
            self.cls_fc = nn.Linear(256, num_classes)
        else:
            bottleneck_dim = 2048
            self.cls_fc = nn.Linear(2048, num_classes)
        self.num_class = num_classes
        self.max_iter = max_iter
        self.iter_num = 0
        self.ad_net = AdversarialNetwork(bottleneck_dim*num_classes, 1024, max_iter=self.max_iter).cuda()
    # def init_cate_dict(self,source_loader,target_loader,num_classes,cate_len):
    #     with torch.no_grad():
    #         source_cate_dict = {idx:[] for idx in range(num_classes)}
    #         iter_data = iter(source_loader)
    #         dict_counter = 0
    #         dict_index_set = set()
    #         while True:
    #             data_source, label_source = next(iter_data)
    #             data_source,label_source = data_source.to(self.device),label_source.to(self.device)
    #             source,source_stages = self.base_network(data_source)
    #             if self.use_bottleneck:
    #                 source = self.bottleneck_layer(source)
    #             source_clf = self.classifier_layer(source)
    #             #B,D
    #             source = source[source_clf.argmax(-1)==label_source]
    #             label_source = label_source[source_clf.argmax(-1)==label_source]
    #             for source_f,class_idx in zip(source[:],label_source.int().tolist()):
    #                 if len(source_cate_dict[class_idx])<cate_len:
    #                     source_cate_dict[class_idx].append(F.normalize(source_f.detach(),dim=-1))
    #                 if len(source_cate_dict[class_idx])>=cate_len and class_idx not in dict_index_set:
    #                     dict_index_set.add(class_idx)
    #                     dict_counter += 1
    #             if dict_counter >= num_classes:
    #                 break
    #         source_cate_dict = torch.stack([torch.stack(v) for k,v in source_cate_dict.items()])

    #         target_cate_dict = {idx:[] for idx in range(num_classes)}
    #         iter_data = iter(target_loader)
    #         dict_counter = 0
    #         dict_index_set = set()
    #         thres = 0.9
    #         while True:
    #             data_target, _ = next(iter_data)
    #             data_target = data_target.to(self.device)
    #             target,target_stages = self.base_network(data_target)
    #             if self.use_bottleneck:
    #                 target = self.bottleneck_layer(target)
    #             target_clf = self.classifier_layer(target)
    #             #B,D
    #             label_target = target_clf.argmax(-1)
    #             target = target[label_target>thres]
    #             label_target = label_target[label_target>thres]
    #             for target_f,class_idx in zip(target[:],label_target.int().tolist()):
    #                 if len(target_cate_dict[class_idx])<cate_len:
    #                     target_cate_dict[class_idx].append(F.normalize(target_f.detach(),dim=-1))
    #                 if len(target_cate_dict[class_idx])>=cate_len and class_idx not in dict_index_set:
    #                     dict_index_set.add(class_idx)
    #                     dict_counter += 1
    #             if dict_counter >= num_classes:
    #                 break
    #             thres -= 0.01
    #         target_cate_dict = torch.stack([torch.stack(v) for k,v in target_cate_dict.items()])
    #         cate_dict = torch.concat([source_cate_dict,target_cate_dict],dim=1)
    #         idx = torch.randperm(cate_dict.shape[1])
    #         cate_dict = cate_dict[:,idx]
    #         return cate_dict.to(self.device)

    # def update_cate_dict(self,source_clf,source_label,source_feature,target_clf,target_feature):
    #     with torch.no_grad():
    #         new_cate_dict = torch.zeros_like(self.cate_dict)
    #         source_feature = source_feature[source_clf.argmax(-1)==source_label]
    #         source_label = source_label[source_clf.argmax(-1)==source_label]
    #         class_idx_set = set()
    #         for source_f,class_idx in zip(source_feature[:],source_label.int().tolist()):
    #             if class_idx not in class_idx_set:
    #                 class_idx_set.add(class_idx)
    #                 new_cate_dict[:,1:] = self.cate_dict[:,:-1].detach()
    #                 new_cate_dict[:,0] = F.normalize(source_f.detach(),dim=-1)
    #         target_label = target_clf[target_clf.argmax(-1)>0.8].argmax(-1)
    #         target_feature = target_feature[target_clf.argmax(-1)>0.8]
    #         for target_f,class_idx in zip(target_feature[:],target_label.int().tolist()):
    #             if class_idx not in class_idx_set:
    #                 class_idx_set.add(class_idx)
    #                 new_cate_dict[:,1:] = self.cate_dict[:,:-1].detach()
    #                 new_cate_dict[:,0] = F.normalize(target_f.detach(),dim=-1)

    #         return new_cate_dict.to(self.device)

    # def get_logits_thres(self,max_thres=0.8,min_thres=0.2):
    #     p = self.iter_num / self.max_iter
    #     return max_thres - (max_thres - min_thres) * p

    # def step(self):
    #     self.iter_num+=1

    def forward(self, source, target, s_label):
        source = self.feature_layers(source)
        if self.bottle_neck:
            source = self.bottle(source)
        s_pred = self.cls_fc(source)

        target = self.feature_layers(target)
        if self.bottle_neck:
            target = self.bottle(target)
        t_pred = self.cls_fc(target)
        t_logits = torch.nn.functional.softmax(t_pred,dim=1)

        # entropy = None
        # features = torch.cat((source, target), dim=0)
        # outputs = torch.cat((s_pred,t_pred), dim=0)
        # softmax_out = nn.Softmax(dim=1)(outputs)
        # eff = calc_coeff(self.iter_num, max_iter=self.max_iter)
        # cdan_loss = losses.CDAN([features, softmax_out], self.ad_net, entropy, eff)

        # logits_thres = self.get_logits_thres()
        # t_logits_max = t_logits.max(-1)
        # t_logits_filter_bool = t_logits_max.values>logits_thres
        # t_logits_filter_index = torch.arange(0,t_logits.size()[0])[t_logits_filter_bool]
        # t_logits = t_logits[t_logits_filter_index]
        # target = target[t_logits_filter_index]
        loss_lmmd = self.lmmd_loss.get_loss(source, target, s_label,t_logits)
        # ad_net = utils.AdversarialNetwork(args.bottleneck_dim*class_num, 1024, max_iter=args.max_iter).cuda()
        # entropy = loss.Entropy(softmax_out)
        # self.step()
        return s_pred, loss_lmmd

    def predict(self, x):
        x = self.feature_layers(x)
        if self.bottle_neck:
            x = self.bottle(x)
        return self.cls_fc(x)

class VGG16Base(nn.Module):
  def __init__(self):
    super(VGG16Base, self).__init__()
    model_vgg = torchvision.models.vgg16(pretrained=True)
    self.features = model_vgg.features
    self.classifier = nn.Sequential()
    for i in range(6):
        self.classifier.add_module("classifier"+str(i), model_vgg.classifier[i])
    self.feature_layers = nn.Sequential(self.features, self.classifier)
    self.in_features = 4096

  def forward(self, x):
    x = self.features(x)
    x = x.view(x.size(0), -1)
    x = self.classifier(x)
    return x

class ResBase18(nn.Module):
    def __init__(self):
        super(ResBase18, self).__init__()
        model_resnet = torchvision.models.resnet18(pretrained=True)
        self.conv1 = model_resnet.conv1
        self.bn1 = model_resnet.bn1
        self.relu = model_resnet.relu
        self.maxpool = model_resnet.maxpool
        self.layer1 = model_resnet.layer1
        self.layer2 = model_resnet.layer2
        self.layer3 = model_resnet.layer3
        self.layer4 = model_resnet.layer4
        self.avgpool = model_resnet.avgpool
        self.in_features = model_resnet.fc.in_features

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        return x

class ResBase34(nn.Module):
    def __init__(self):
        super(ResBase34, self).__init__()
        model_resnet = torchvision.models.resnet34(pretrained=True)
        self.conv1 = model_resnet.conv1
        self.bn1 = model_resnet.bn1
        self.relu = model_resnet.relu
        self.maxpool = model_resnet.maxpool
        self.layer1 = model_resnet.layer1
        self.layer2 = model_resnet.layer2
        self.layer3 = model_resnet.layer3
        self.layer4 = model_resnet.layer4
        self.avgpool = model_resnet.avgpool
        self.in_features = model_resnet.fc.in_features

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        return x

class ResBase50(nn.Module):
    def __init__(self):
        super(ResBase50, self).__init__()
        model_resnet50 = torchvision.models.resnet50(pretrained=True)
        self.conv1 = model_resnet50.conv1
        self.bn1 = model_resnet50.bn1
        self.relu = model_resnet50.relu
        self.maxpool = model_resnet50.maxpool
        self.layer1 = model_resnet50.layer1
        self.layer2 = model_resnet50.layer2
        self.layer3 = model_resnet50.layer3
        self.layer4 = model_resnet50.layer4
        self.avgpool = model_resnet50.avgpool
        self.in_features = model_resnet50.fc.in_features

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        return x

class ResBase101(nn.Module):
    def __init__(self):
        super(ResBase101, self).__init__()
        model_resnet101 = torchvision.models.resnet101(pretrained=True)
        self.conv1 = model_resnet101.conv1
        self.bn1 = model_resnet101.bn1
        self.relu = model_resnet101.relu
        self.maxpool = model_resnet101.maxpool
        self.layer1 = model_resnet101.layer1
        self.layer2 = model_resnet101.layer2
        self.layer3 = model_resnet101.layer3
        self.layer4 = model_resnet101.layer4
        self.avgpool = model_resnet101.avgpool
        self.in_features = model_resnet101.fc.in_features

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        return x

class ResClassifier(nn.Module):
    def __init__(self, class_num, feature_dim,bottleneck=True,bottleneck_dim=256):
        super(ResClassifier, self).__init__()
        self.bottleneck = bottleneck
        if bottleneck:
            self.bottleneck = nn.Linear(feature_dim, bottleneck_dim)
            self.fc = nn.Linear(bottleneck_dim, class_num)
        else:
            bottleneck_dim = feature_dim
            self.fc = nn.Linear(bottleneck_dim, class_num)
        # self.bottleneck.apply(init_weights)
        # self.fc.apply(init_weights)

    def forward(self, x):
        if self.bottleneck:
            x = self.bottleneck(x)
        y = self.fc(x)
        return x,y

class AdversarialNetwork(nn.Module):
  def __init__(self, in_feature, hidden_size, max_iter=10000):
    super(AdversarialNetwork, self).__init__()
    self.ad_layer1 = nn.Linear(in_feature, hidden_size)
    self.ad_layer2 = nn.Linear(hidden_size, hidden_size)
    self.ad_layer3 = nn.Linear(hidden_size, 1)
    self.relu1 = nn.ReLU()
    self.relu2 = nn.ReLU()
    self.dropout1 = nn.Dropout(0.5)
    self.dropout2 = nn.Dropout(0.5)
    self.sigmoid = nn.Sigmoid()
    self.apply(init_weights)
    self.iter_num = 0
    self.alpha = 10
    self.low = 0.0
    self.high = 1.0
    self.max_iter = max_iter

  def forward(self, x):
    if self.training:
        self.iter_num += 1
    coeff = calc_coeff(self.iter_num, self.high, self.low, self.alpha, self.max_iter)
    x = x * 1.0
    x.register_hook(grl_hook(coeff))
    x = self.ad_layer1(x)
    x = self.relu1(x)
    y = self.ad_layer3(x)
    y = self.sigmoid(y)
    return y

  def output_num(self):
    return 1
  def get_parameters(self):
    return [{"params":self.parameters(), "lr_mult":10, 'decay_mult':2}]

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.baselayer = [self.conv1, self.bn1, self.layer1, self.layer2, self.layer3, self.layer4]
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)

        return x

def resnet50(pretrained=False, **kwargs):
    """Constructs a ResNet-50 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
    return model

def resnet101(pretrained=False,**kwargs):
    """Constructs a ResNet-50 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
    return model


from fastapi                 import FastAPI, WebSocket
from fastapi.staticfiles     import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from app.routers.user        import router as user_router
from app.routers.category    import router as category_router
from app.routers.blog        import router as blog_router
from app.routers.content     import router as content_router
from app.routers.image       import router as image_router
from app.home                import router as home_router
from app.config              import *
import os

os.makedirs("static", exist_ok=True)

app = FastAPI(docs_url=DOCS_URL, redoc_url=None)
app.add_middleware(
    middleware_class=CORSMiddleware,
    allow_origins=ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)
app.mount(path="/images",    app=StaticFiles(directory="static"),    name="static")
app.mount(path="/templates", app=StaticFiles(directory="templates"), name="templates")

websocket_clients = set()
messages = []

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    websocket_clients.add(websocket)

    try:
        while True:
            text = await websocket.receive_text()

            messages.append(text)

            for client in websocket_clients:
                # for message in messages:
                #     print(message)
                await client.send_json({"message": messages})
    except Exception as e:
        print(e)
    finally:
        websocket_clients.remove(websocket)

app.include_router(home_router,     prefix="",                 tags=["Home"])
app.include_router(user_router,     prefix="/api/v1/user",     tags=["Users"])
app.include_router(category_router, prefix="/api/v1/category", tags=["Category"])
app.include_router(blog_router,     prefix="/api/v1/blog",     tags=["Blog"])
app.include_router(content_router,  prefix="/api/v1/content",  tags=["Content"])
app.include_router(image_router,    prefix="/api/v1/upload",   tags=["Image"])


# -*- coding: utf-8 -*-
#
# Author: Taylor Smith <taylor.smith@alkaline-ml.com>

import numpy as np
import pandas as pd

from .base import BaseFeatureSelector
from ..utils.validation import check_dataframe, validate_multiple_cols
from ..utils.metaestimators import timed_instance_method

__all__ = [
    'FeatureFilter',
    'MultiCorrFilter',
    'NearZeroVarianceFilter',
    'SparseFeatureFilter'
]

# TODO: add functionality for a dummy column indicating present value?
class SparseFeatureFilter(BaseFeatureSelector):
    """Drop overly sparse features.

    Retains features that are less sparse (NaN) than the provided
    threshold. Useful in situations where matrices are too sparse to
    impute reliably.

    Parameters
    ----------
    cols : array-like, shape=(n_features,), optional (default=None)
        The names of the columns on which to apply the transformation.
        If no column names are provided, the transformer will be ``fit``
        on the entire frame. Note that the transformation will also only
        apply to the specified columns, and any other non-specified
        columns will still be present after transformation.

    threshold : float, optional (default=0.5)
        The threshold of sparsity above which features will be
        deemed "too sparse" and will be dropped.

    as_df : bool, optional (default=True)
        Whether to return a Pandas ``DataFrame`` in the ``transform``
        method. If False, will return a Numpy ``ndarray`` instead.
        Since most skutil transformers depend on explicitly-named
        ``DataFrame`` features, the ``as_df`` parameter is True by default.

    Examples
    --------
    >>> import numpy as np
    >>> import pandas as pd
    >>>
    >>> nan = np.nan
    >>> X = np.array([
    ...     [1.0, 2.0, nan],
    ...     [2.0, 3.0, nan],
    ...     [3.0, nan, 1.0],
    ...     [4.0, 5.0, nan]
    ... ])
    >>>
    >>> X = pd.DataFrame.from_records(data=X, columns=['a','b','c'])
    >>> dropper = SparseFeatureFilter(threshold=0.5)
    >>> X_transform = dropper.fit_transform(X)
    >>> assert X_transform.shape[1] == 2 # drop out last column

    Attributes
    ----------
    sparsity_ : array-like, shape=(n_features,)
        The array of sparsity values

    drop_ : array-like, shape=(n_features,)
        Assigned after calling ``fit``. These are the features that
        are designated as "bad" and will be dropped in the ``transform``
        method.

    Notes
    -----
    Sometimes the presence of a value in an overly sparse column can be highly
    informative. If you're using the sparse filter, consider creating a new
    (dummy) feature indicating whether there was a value present.
    """
    def __init__(self, cols=None, threshold=0.5, as_df=True):

        super(SparseFeatureFilter, self).__init__(
            cols=cols, as_df=as_df)

        self.threshold = threshold

    @timed_instance_method(attribute_name="fit_time_")
    def fit(self, X, y=None):
        """Fit the transformer.

        Parameters
        ----------
        X : pd.DataFrame, shape=(n_samples, n_features)
            The Pandas frame to fit. The frame will only
            be fit on the prescribed ``cols`` (see ``__init__``) or
            all of them if ``cols`` is None. Furthermore, ``X`` will
            not be altered in the process of the fit.

        y : array-like or None, shape=(n_samples,), optional (default=None)
            Pass-through for ``sklearn.pipeline.Pipeline``. Even
            if explicitly set, will not change behavior of ``fit``.
        """
        X, cols = check_dataframe(X, cols=self.cols)
        thresh = self.threshold

        # validate the threshold
        if not (isinstance(thresh, float) and (0.0 <= thresh < 1.0)):
            raise ValueError('thresh must be a float between '
                             '0 (inclusive) and 1. Got %s' % str(thresh))

        # assess sparsity
        subset = X[cols]
        self.sparsity_ = subset.apply(
            lambda x: x.isnull().sum() / x.shape[0]).values  # type: np.ndarray

        mask = self.sparsity_ > thresh  # numpy boolean array
        self.drop_ = subset.columns[mask].tolist()
        return self

class FeatureFilter(BaseFeatureSelector):
    """A simple feature-dropping transformer class.

    A very simple class to be used at the beginning or any stage of a
    Pipeline that will drop the given features from the remainder of the pipe.
    This is useful if a transformer or encoder creates variables that you're
    disinterested in and would like to exclude from your modeling process.

    Parameters
    ----------
    cols : array-like, shape=(n_features,), optional (default=None)
        The features to drop. Note that ``FeatureFilter`` behaves slightly
        differently from all other ``BaseFeatureSelector`` classes in the sense
        that it will drop all of the features prescribed in this parameter.
        However, if ``cols`` is None, it will not drop any (which is counter
        to other classes, which will operate on all columns in the absence
        of an explicit ``cols`` parameter).

    as_df : bool, optional (default=True)
        Whether to return a Pandas ``DataFrame`` in the ``transform``
        method. If False, will return a Numpy ``ndarray`` instead.
        Since most skoot transformers depend on explicitly-named
        ``DataFrame`` features, the ``as_df`` parameter is True by default.

    Examples
    --------
    An example using the FeatureFilter:

    >>> import numpy as np
    >>> import pandas as pd
    >>>
    >>> X = pd.DataFrame.from_records(data=np.random.rand(3,3),
    ...                               columns=['a','b','c'])
    >>> dropper = FeatureFilter(cols=['a','b'])
    >>> X_transform = dropper.fit_transform(X)
    >>> assert X_transform.shape[1] == 1 # drop out first two columns

    Attributes
    ----------
    drop_ : array-like, shape=(n_features,)
        Assigned after calling ``fit``. These are the features that
        are designated as "bad" and will be dropped in the ``transform``
        method.
    """
    def __init__(self, cols=None, as_df=True):
        # just a pass-through for super constructor
        super(FeatureFilter, self).__init__(
            cols=cols, as_df=as_df)

    def fit(self, X, y=None):
        # check on state of X and cols
        X, cols = check_dataframe(X, cols=self.cols)

        # if the provided self.cols was None, we drop nothing. otherwise
        # we drop the specified columns
        self.drop_ = [] if not self.cols else cols
        return self

class MultiCorrFilter(BaseFeatureSelector):
    """Remove highly correlated features.

    Multi-collinear data (features which are not independent from one another)
    can pose problems in coefficient stability for parametric models, or
    feature importance scores for non-parametric models.

    This class filters out features with a correlation greater than the
    provided threshold. When a pair of correlated features is identified, the
    mean absolute correlation (MAC) of each feature is considered, and the
    feature with the highest MAC is discarded.

    Parameters
    ----------
    cols : array-like, shape=(n_features,), optional (default=None)
        The names of the columns on which to apply the transformation.
        If no column names are provided, the transformer will be ``fit``
        on the entire frame. Note that the transformation will also only
        apply to the specified columns, and any other non-specified
        columns will still be present after transformation.

    threshold : float, optional (default=0.85)
        The threshold above which to filter correlated features

    method : str, optional (default='pearson')
        The method used to compute the correlation,
        one of ('pearson', 'kendall', 'spearman').

    as_df : bool, optional (default=True)
        Whether to return a Pandas ``DataFrame`` in the ``transform``
        method. If False, will return a Numpy ``ndarray`` instead.
        Since most skoot transformers depend on explicitly-named
        ``DataFrame`` features, the ``as_df`` parameter is True by default.

    Examples
    --------
    The following demonstrates a simple multi-correlation filter
    applied to the iris dataset.

    >>> from skoot.datasets import load_iris_df
    >>>
    >>> X = load_iris_df(include_tgt=False)
    >>> mcf = MultiCorrFilter(threshold=0.85)
    >>> mcf.fit_transform(X).head()
       sepal length (cm)  sepal width (cm)  petal width (cm)
    0                5.1               3.5               0.2
    1                4.9               3.0               0.2
    2                4.7               3.2               0.2
    3                4.6               3.1               0.2
    4                5.0               3.6               0.2

    Attributes
    ----------
    drop_ : array-like, shape=(n_features,)
        Assigned after calling ``fit``. These are the features that
        are designated as "bad" and will be dropped in the ``transform``
        method.

    mean_abs_correlations_ : list, float
        The corresponding mean absolute correlations of each ``drop_`` name
    """

    def __init__(self, cols=None, threshold=0.85,
                 method='pearson', as_df=True):

        super(MultiCorrFilter, self).__init__(
            cols=cols, as_df=as_df)

        self.threshold = threshold
        self.method = method

    def fit(self, X, y=None):
        """Fit the multi-collinearity filter.

        Parameters
        ----------
        X : pd.DataFrame, shape=(n_samples, n_features)
            The Pandas frame to fit. The frame will only
            be fit on the prescribed ``cols`` (see ``__init__``) or
            all of them if ``cols`` is None. Furthermore, ``X`` will
            not be altered in the process of the fit.

        y : array-like or None, shape=(n_samples,), optional (default=None)
            Pass-through for ``sklearn.pipeline.Pipeline``. Even
            if explicitly set, will not change behavior of ``fit``.
        """
        # check on state of X and cols. Also need all columns to be finite!
        X, cols = check_dataframe(X, cols=self.cols, assert_all_finite=True)

        # we need to make sure there's more than 1 column!
        validate_multiple_cols(self.__class__.__name__, cols)

        # Generate correlation matrix
        c = X[cols].corr(method=self.method)

        # get drops list
        # TODO: write a _find_correlations_exact for smaller matrices
        self.drop_, self.mean_abs_correlations_ = \
            self._find_correlations_fast(c, self.threshold)

        return self

    @staticmethod
    def _find_correlations_fast(c, threshold):
        """Filter highly correlated features.

        This function identifies the correlations between features that
        are greater than the provided threshold, and identifies which
        to drop on the basis of mean absolute correlation. This function is
        based on Caret's ``findCorrelation_fast`` method [1], which is very
        efficient for large matrices.

        Parameters
        ----------
        c : pd.DataFrame
            The pre-computed correlation matrix.

        threshold : float
            The threshold above which to filter features which
            are multi-collinear in nature.

        Returns
        -------
        drop_names : list, shape=(n_features,)
            The feature names that should be dropped

        average_corr : array-like, shape=(n_features,)
            The mean absolute correlations between
            the features.

        References
        ----------
        .. [1] Caret findCorrelations.R (findCorrelation_fast)
               https://bit.ly/2E1AMcJ
        """
        # get the average absolute column correlations
        c_abs = c.abs()  # type: pd.DataFrame
        average_corr = c_abs.mean().values  # type: np.ndarray

        # get the sort order
        average_corr_order = np.argsort(average_corr)

        # set the lower tri to NaN so we don't consider anymore
        # first, get the lower tri indices and then zip them
        lower_tri = np.tril_indices(n=c.shape[0], k=0)
        c.values[lower_tri] = np.nan

        # get those above cutoff (re-compute abs on amended array)
        c_abs = c.abs()  # type: pd.DataFrame
        combs_above_thresh = np.where(c_abs > threshold)
        rows_to_check, cols_to_check = combs_above_thresh

        cols_to_discard = (average_corr_order[cols_to_check] >
                           average_corr_order[rows_to_check])
        rows_to_discard = ~cols_to_discard

        # append each set of discard rows/cols, get the distinct
        drop_cols = np.unique(
            np.concatenate([cols_to_check[cols_to_discard],
                            rows_to_check[rows_to_discard]]))

        # the names to drop
        drop_names = c.columns[drop_cols].tolist()
        return drop_names, average_corr

class NearZeroVarianceFilter(BaseFeatureSelector):
    r"""Identify near zero variance predictors.

    Diagnose and remove any features that have one unique value
    (i.e., are zero variance predictors) or that are have both of the
    following characteristics: they have very few unique values relative
    to the number of samples and the ratio of the frequency of the most
    common value to the frequency of the second most common value is large.

    A note of caution: if you attempt to run this over large, continuous data,
    it might take a long time. Since for each column in ``cols`` it will
    compute ``value_counts``, applying to continuous data could be a bad idea.

    Parameters
    ----------
    cols : array-like, shape=(n_features,), optional (default=None)
        The names of the columns on which to apply the transformation.
        If no column names are provided, the transformer will be ``fit``
        on the entire frame. Note that the transformation will also only
        apply to the specified columns, and any other non-specified
        columns will still be present after transformation.

    freq_cut : float, optional (default=95/5)
        The cutoff for the ratio of the most common value to the second most
        common value. That is, if the frequency of the most common value is
        >= ``freq_cut`` times the frequency of the second most, the feature
        will be dropped.

    as_df : bool, optional (default=True)
        Whether to return a Pandas ``DataFrame`` in the ``transform``
        method. If False, will return a Numpy ``ndarray`` instead.
        Since most skutil transformers depend on explicitly-named
        ``DataFrame`` features, the ``as_df`` parameter is True by default.

    Examples
    --------
    An example of the near zero variance filter on a completely
    constant column:

    >>> import pandas as pd
    >>> import numpy as np
    >>>
    >>> X = pd.DataFrame.from_records(
    ...     data=np.array([[1,2,3], [4,5,3], [6,7,3], [8,9,3]]),
    ...     columns=['a','b','c'])
    >>> flt = NearZeroVarianceFilter(freq_cut=25)
    >>> flt.fit_transform(X)
       a  b
    0  1  2
    1  4  5
    2  6  7
    3  8  9

    An example on a column with two unique values represented at 2:1. Also
    shows how we can extract the fitted ratios and drop names:

    >>> X = pd.DataFrame.from_records(
    ...     data=np.array([[1,2,3], [4,5,3], [6,7,5]]),
    ...     columns=['a','b','c'])
    >>> nzv = NearZeroVarianceFilter(freq_cut=2.)
    >>> nzv.fit_transform(X)
       a  b
    0  1  2
    1  4  5
    2  6  7
    >>> nzv.ratios_
    array([ 1.,  1.,  2.])
    >>> nzv.drop_
    ['c']

    Attributes
    ----------
    drop_ : array-like, shape=(n_features,)
        Assigned after calling ``fit``. These are the features that
        are designated as "bad" and will be dropped in the ``transform``
        method.

    ratios_ : array-like, shape=(n_features,)
        The ratios of the counts of the most populous classes to the second
        most populated classes for each column in ``cols``.

    References
    ----------
    .. [1] Kuhn, M. & Johnson, K. "Applied Predictive
           Modeling" (2013). New York, NY: Springer.

    .. [2] Caret (R package) nearZeroVariance R code
           https://bit.ly/2J0ozbM
    """
    def __init__(self, cols=None, freq_cut=95. / 5., as_df=True):

        super(NearZeroVarianceFilter, self).__init__(
            cols=cols, as_df=as_df)

        self.freq_cut = freq_cut

    def fit(self, X, y=None):
        """Fit the near-zero variance filter.

        Parameters
        ----------
        X : pd.DataFrame, shape=(n_samples, n_features)
            The Pandas frame to fit. The frame will only
            be fit on the prescribed ``cols`` (see ``__init__``) or
            all of them if ``cols`` is None. Furthermore, ``X`` will
            not be altered in the process of the fit.

        y : array-like or None, shape=(n_samples,), optional (default=None)
            Pass-through for ``sklearn.pipeline.Pipeline``. Even
            if explicitly set, will not change behavior of ``fit``.
        """
        # check on state of X and cols
        X, cols = check_dataframe(X, self.cols)

        # get the freq cut and validate it is an appropriate value...
        freq_cut = self.freq_cut
        if not (isinstance(freq_cut, (int, float)) and 1. < freq_cut):
            raise ValueError("freq_cut must be a float > 1.0")

        # make sure it's cast to a float if not already
        freq_cut = float(freq_cut)

        # get a mask of which should be dropped
        subset = X[cols]
        ratios = subset.apply(self._filter_freq_cut).values
        self.drop_ = subset.columns[ratios >= freq_cut].tolist()
        self.ratios_ = ratios

        return self

    @staticmethod
    def _filter_freq_cut(series):
        """Filter above a frequency cut.

        Parameters
        ----------
        series : pd.Series
            One series in the specified column list.

        Returns
        -------
        ratio : float
            The ratio of the count of the most populated class to the second
            most populated class. If there is only one class, will return
            infinity.
        """
        vc = series.value_counts()
        n_levels = vc.shape[0]

        # base case 1: vc len is 1 (single value, no variance at all)
        if n_levels == 1:
            return np.inf

        # get the first two levels and counts
        first_two = vc.values[:2].astype(float)
        return first_two[0] / first_two[1]


from sklearn.datasets import make_regression
from mllib.supervised.LinearRegression.regressions import *
from mllib.utilities.losses import *
from mllib.utilities.lr_schedulers import *
np.random.seed(1)
plt.rcParams['axes.grid'] = True

if __name__=='__main__':
    N_FEATURES = 1
    ANIMATE = True
    LEARNING_RATE = 1e-1
    assert (not ANIMATE and N_FEATURES>=1) or (ANIMATE and N_FEATURES==1), 'Cannot animate more than 1 dimension'

    X, y = make_regression(n_features=N_FEATURES, noise=15)
    X = np.array([[0.05, 0.1,0.5,0.51,0.51,0.52,0.8, 0.9, 1.2, 1.3,1.55,1.75,1.76,2.1,2.25,2.4,2.5,2.8,2.9, 2.9]]).reshape(-1,1)
    y = np.array([0.25,1.1,1.05,2.25,0.8,1.1,2.5,0.6,1.4,1.2,2.2,1.05,0.95,1.5,2,1.5,1.85,1.4,2.5,2.6])

    # X = 2*np.array(range(0, 11)) + np.random.normal()
    # X = X.reshape(-1,1)
    # y = 2*X**2 - 2*X
    # y = y.flatten()
    # sc = StandardScaler()
    # X = sc.fit_transform(X)
    scheduler = configure_power_scheduler(lr0=LEARNING_RATE, every=30)
    # scheduler = configure_exponential_decay(lr0=LEARNING_RATE, every=50)

    # lr = LinearRegression(loss=MSE, learning_rate=LEARNING_RATE, scheduler_fn=None)
    # lr.fit(X, y, epochs=30,animate=ANIMATE)

    lr = LassoRegression(loss=MSE, alpha=0.9, learning_rate=LEARNING_RATE, scheduler_fn=None)
    lr.fit(X, y, epochs=30,animate=ANIMATE)

    # lasso = Lasso(alpha=0.9)
    # lr = PolynomialRegression(loss=MSE,  learning_rate=LEARNING_RATE, scheduler_fn=scheduler, degree=6, penalty=lasso)
    # lr.fit(X, y, epochs=200,animate=ANIMATE)

    # nr = NormalEquation(alpha=0.9)
    # nr.fit(X, y, animate=ANIMATE)
    # ridge_reg = ElasticNetRegression(l1_ratio=0.9, alpha=0.9, learning_rate=1e-1, loss=MSE)
    # ridge_reg.fit(X,y,epochs=2000, animate=ANIMATE)

    # lr = ElasticNetRegression(loss=MSE, alpha=0.9, learning_rate=LEARNING_RATE, scheduler_fn=scheduler, l1_ratio=0.9)
    # lr.fit(X, y, epochs=50, animate=True)


# Given two strings s and t of lengths m and n respectively, return the minimum window
# substring of s such that every character in t (including duplicates) is included in
# the window. If there is no such substring, return the empty string "".

# The testcases will be generated such that the answer is unique.

# Constraints:
#   m == s.length
#   n == t.length
#   1 <= m, n <= 105
#   s and t consist of uppercase and lowercase English letters.

# Follow up: Could you find an algorithm that runs in O(m + n) time?

from collections import Counter

TEST_CASES = [
  ["ADOBECODEBANC", "ABC", "BANC"],
  ["a", "a", "a"],
  ["a", "aa", ""],
]

def all_letters_used(substr, target):
    # could be improved with uniqueness
    for char in target:
        if substr.count(char) < target.count(char):
            return False

    return True

def minimum_window_substring(s, t):
    ans = ""

    if len(s) < len(t):
        return ans

    left = 0
    right = len(t)

    # build first window
    substr = s[left:right]

    # check if all letters present
    if all_letters_used(substr, t):
        # case where t is the leftmost substr of s
        return substr

    right += 1
    while right <= len(s):
        substr = s[left:right]
        while all_letters_used(substr, t):
            if len(substr) <= len(ans) or ans == "":
                ans = substr

            left += 1
            substr = s[left:right]

        right += 1

    # return answer
    return ans

def minimum_window_substring_optimized(s, t):
    # print("Start:", s, t)
    ans = ""

    if len(s) < len(t):
        return ans

    # m
    characters_needed = Counter(t)
    # print('A:', characters_needed)

    left = 0
    right = 0
    while right < len(s):
        # expand window right until all characters are accounted for
        while not all(count <= 0 for count in characters_needed.values()):
            char = s[right]
            if char in characters_needed:
                characters_needed[char] -= 1

            right += 1

        # shrink window from the left until we need more characters
        while all(count <= 0 for count in characters_needed.values()):
            substr = s[left:right]

            # check if shortest valid substr so far
            if len(substr) < len(ans) or ans == "":
                # if so, replace ans
                ans = substr

            char = s[left]
            if char in characters_needed:
                characters_needed[char] += 1
            left += 1

    return ans

# def minimum_window_substring_best(s, t):
#     tDict = defaultdict(int)
#     for c in t:
#         tDict[c]+=1
#     print(tDict)
#     m,n=len(s),len(t)
#     need=len(tDict)
#     have=0
#     l,r=0,0
#     mS=""
#     ma=float('inf')
#     sDict=defaultdict(int)
#     for r in range(m):
#         # print(sDict)
#         sDict[s[r]]+=1
#         if tDict.get(s[r],-1) == sDict[s[r]]:
#             have+=1
#             # print(have)
#         while have==need and l<=r:
#             if r-l+1<ma:
#                 ma=r-l+1
#                 mS=s[l:r+1]
#             if sDict[s[l]]==tDict.get(s[l],-1):
#                 have-=1
#             sDict[s[l]]-=1
#             l+=1
#     return mS

def minimum_window_substring_best(s, t):
    # count letters in t; this is O(n)
    t_counts = Counter(t)

    n = len(s)
    # how many character counts have to match, or, how many unique characters in t
    need = len(t_counts)

    # how many character counts we currently have that do match
    have = 0

    # initialize left and right pointers
    left, right = 0,0

    ans = ""
    len_ans = float('inf')
    s_counts = Counter()

    # for character in s; this is O(m)
    for right in range(n):
        # expand the window to the right
        right_char = s[right]
        s_counts[right_char] += 1

        # if adding that character makes our count for that character match the target
        if t_counts.get(right_char, -1) == s_counts[right_char]:
            # increment the number of matching character counts
            have += 1

        # whenever we have as many matching character counts as we need,
        # we have a valid substring solution
        while have == need and left <= right:
            # check if current substr len is better than ans len
            substring_len = right - left + 1
            if substring_len < len_ans:
                # if so, update the answer
                len_ans = substring_len
                ans = s[left:right+1]

            left_char = s[left]

            # if the character we're removing is one in the target,
            if s_counts[left_char] == t_counts.get(left_char,-1):
                # decrement the number of matching character counts
                have -= 1

            # shorten the window from the left
            s_counts[left_char] -= 1
            left += 1

    return ans

if __name__ == "__main__":
    for case in TEST_CASES:
        print(minimum_window_substring_best(case[0], case[1]) == case[2])


import numpy as np
from keras.models import Sequential, Model, model_from_json
from keras.layers import Input
from keras.layers.core import Dense, Activation, Dropout, Merge, RepeatVector
from keras.layers.recurrent import LSTM
from keras.layers.embeddings import Embedding
from keras.preprocessing import sequence
from keras.utils.visualize_util import plot
import keras.backend as K
import json
import random
#from IPython.display import SVG
#from keras.utils.visualize_util import plot
from pymongo import MongoClient
from math import ceil

#fix random number for repeatability
#np.random.seed(17)

input_dim = 3 #input dimensions of fused sensor data
nb_timesteps = 2000 #maximum amount of samples/timesteps per shot
nb_output_class = 3     #slap, snap, wrist, none
nb_batch_size = 36 #number of samples per batch
train_test_ratio = 0.6
nb_input_multi = 6
model_filename = 'LSTM_silkymitties.json'
weights_filename = 'LSTM_silkymitties_weights.h5'

def load_training_data():
    client = MongoClient('localhost', 27017)
    db = client['restdb']
    handedness = []
    speed = []
    accuracy = []
    shotTypes = []
    y = []
    labelled_values = db['testlblfusedshots'].find()
    nb_input_samples = labelled_values.count()
    x_upper = np.empty(shape=[nb_input_samples,nb_timesteps,input_dim])
    x_lower = np.empty(shape=[nb_input_samples,nb_timesteps,input_dim])
    handedness = np.zeros(shape=[nb_input_samples,2], dtype='float')
    #speed = np.empty(shape=[nb_input_samples])
    #accuracy = np.empty(shape=[nb_input_samples])
    shotTypes = np.zeros(shape=[nb_input_samples,nb_output_class], dtype='float')
    #slapshotTypes = np.empty(shape=[nb_input_samples])
    #y = np.empty(shape=[nb_input_samples,nb_output_class])
    index = 0
    for target in labelled_values:
        upperTheta = np.vstack((target['upperTheta']['uthetaX'], target['upperTheta']['uthetaY'], target['upperTheta']['uthetaZ']))
        lowerTheta = np.vstack((target['lowerTheta']['lthetaX'], target['lowerTheta']['lthetaY'], target['lowerTheta']['lthetaZ']))
        normalizedUpperTheta = upperTheta / 180.0
        normalizedLowerTheta = lowerTheta / 180.0
        x_upper[index] = sequence.pad_sequences(normalizedUpperTheta, maxlen=nb_timesteps, dtype='float', padding='post', truncating='post', value=0.).T
        x_lower[index] = sequence.pad_sequences(normalizedLowerTheta, maxlen=nb_timesteps, dtype='float', padding='post', truncating='post', value=0.).T
        shotTypes[index,shotTypeToInt(target['shotType'])] = 1.0
        #slapshotTypes[index] = isSlapShot(target['shotType'])
        handedness[index,isLeftHanded(target['handedness'])] = 1.0
        #speed = nb.append(speed, target['speed']], axis=0)
        #accuracy = nb.append(accuracy, ['accuracy'], axis=0)
        index+=1

    #for size in range(20, nb_input_samples, 20):
    #    trainIndex = round(size * train_test_ratio)
    #    nb_epochs = ceil(size / nb_batch_size)
    #    trainShotTypeCompileFit(nb_epochs, handedness[0:trainIndex], handedness[trainIndex:size], \
    #            x_upper[0:trainIndex], x_lower[0:trainIndex], shotTypes[0:trainIndex,:], \
    #            x_upper[trainIndex:size], x_lower[trainIndex:size], shotTypes[trainIndex:size,:])

    #Shuffle the samples in unison to decrease data clustering
    s_handedness, s_x_upper, s_x_lower, s_shotTypes = unison_shuffle(handedness,
        x_upper, x_lower, shotTypes)

    trainIndex = round(nb_input_samples * train_test_ratio)
    nb_epochs = ceil(nb_input_samples / nb_batch_size)
    #trainSlapShotCompileFit(nb_epochs, handedness[0:trainIndex], handedness[trainIndex:], \
    #    x_upper[0:trainIndex], x_lower[0:trainIndex], slapshotTypes[0:trainIndex], \
    #    x_upper[trainIndex:], x_lower[trainIndex:], slapshotTypes[trainIndex:])
    trainShotTypeCompileFit(nb_epochs, s_handedness[0:trainIndex], s_handedness[trainIndex:], \
        s_x_upper[0:trainIndex], s_x_lower[0:trainIndex], s_shotTypes[0:trainIndex], \
        s_x_upper[trainIndex:], s_x_lower[trainIndex:], s_shotTypes[trainIndex:])

def unison_shuffle(a, b, c, d):
    p = np.random.permutation(len(a))
    return a[p], b[p], c[p], d[p]

def trainSlapShotCompileFit(epoch_count, train_handedness, test_handedness,
    x_train_upper, x_train_lower, y_train, x_test_upper, x_test_lower, y_test):
    #Upper hand LSTM input
    encoder_a = Sequential()
    encoder_a.add(LSTM(output_dim=nb_input_multi*input_dim,
        batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid'))
    #Lower hand LSTM input
    encoder_b = Sequential()
    encoder_b.add(LSTM(output_dim=nb_input_multi*input_dim,
        batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid'))
    encoded_handedness = Sequential()
    encoded_handedness.add(keras.layers.core.RepeatVector(nb_timesteps))
    #Merge both LSTM units with concatenation
    merged = Merge([encoded_handedness, encoder_a, encoder_b], mode='concat')
    decoder = Sequential()
    decoder.add(merged)
    decoder.add(Dropout(0.1))
    decoder.add(Dense(input_dim*nb_input_multi*2 + 2, activation='relu'))
    decoder.add(Dropout(0.2))
    decoder.add(Dense(input_dim*2, activation='relu'))
    decoder.add(Dropout(0.3))
    #sigmoid activation instead of softmax to avoid normalizing to 1.0
    #1 output signal for the binary classification likelihood
    decoder.add(Dense(1, activation='sigmoid'))

    decoder.compile(loss='binary_crossentropy', optimizer='rmsprop', metrics=['accuracy'])
    decoder.fit([train_handedness, x_train_upper, x_train_lower], y_train,
    batch_size=nb_batch_size, nb_epoch=epoch_count,
    validation_data=([test_handedness, x_test_upper, x_test_lower], y_test))

    printSummary(decoder,test_handedness,x_test_upper, x_test_lower, y_test)
    return decoder

def trainShotTypeCompileFit(epoch_count, train_handedness, test_handedness,
    x_train_upper, x_train_lower, y_train, x_test_upper, x_test_lower, y_test):
    #Upper hand LSTM input
    encoder_a = Sequential()
    encoder_a.add(LSTM(4*input_dim,
        batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid',
        return_sequences=True))
    encoder_a.add(Dropout(0.2))
    encoder_a.add(LSTM(4*input_dim, return_sequences=True,
        activation='sigmoid', inner_activation='hard_sigmoid'))
    encoder_a.add(Dropout(0.25))
    encoder_a.add(LSTM(8*input_dim, batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid'))

    #Lower hand LSTM input
    encoder_b = Sequential()
    encoder_b.add(LSTM(4*input_dim,
        batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid',
        return_sequences=True))
    encoder_b.add(Dropout(0.2))
    encoder_b.add(LSTM(4*input_dim, return_sequences=True,
        activation='sigmoid', inner_activation='hard_sigmoid'))
    encoder_b.add(Dropout(0.25))
    encoder_b.add(LSTM(8*input_dim, batch_input_shape=(nb_batch_size, nb_timesteps, input_dim),
        activation='sigmoid', inner_activation='hard_sigmoid'))

    encoded_handedness = Sequential()
    encoded_handedness.add(Dense(2, batch_input_shape=(nb_batch_size, 2)))

    #Merge both LSTM units with concatenation
    merged = Merge([encoded_handedness, encoder_a, encoder_b], mode='concat')
    decoder = Sequential()
    decoder.add(merged)
    #decoder.add(Dropout(0.25))
    #Use CNNs to expand then shrink to desired output signal
    decoder.add(Dropout(0.25))
    decoder.add(Dense(input_dim*8, activation='sigmoid'))
    decoder.add(Dropout(0.25))
    decoder.add(Dense(output_dim=(2*nb_output_class), activation='sigmoid'))
    decoder.add(Dropout(0.3))
    decoder.add(Dense(nb_output_class, activation='softmax'))

    decoder.compile(loss='categorical_crossentropy', optimizer='rmsprop', metrics=['accuracy'])
    decoder.fit([train_handedness, x_train_upper, x_train_lower], y_train,
    batch_size=nb_batch_size, nb_epoch=epoch_count,
    validation_data=([test_handedness, x_test_upper, x_test_lower], y_test))

    printSummary(decoder, test_handedness, x_test_upper, x_test_lower, y_test)
    saveCompiledShotTypeModel(decoder)
    return decoder

def saveCompiledShotTypeModel(decoder):
    saved_model = decoder.to_json()
    with open('LSTM_silkymitties_ShotType.json', 'w') as outfile:
        json.dump(saved_model, outfile)
    decoder.save_weights('LSTM_silkymitties_ShotType_weights.h5')

def loadCompiledShotTypeModel():
    with open('LSTM_silkymitties_ShotType.json', 'r') as infile:
        architecture = json.load(infile)
    model = model_from_json(architecture)
    model.load_weights('LSTM_silkymitties_ShotType_weights.h5')
    return model

def predictShotTypeResult(fusedShotID):
    client = MongoClient('localhost', 27017)
    db = client['restdb']
    shot = db['fusedshots'].find_one(fusedShotID)
    x_upper = np.empty(shape=[1,nb_timesteps,input_dim])
    x_lower = np.empty(shape=[1,nb_timesteps,input_dim])
    raw_x_upper = np.vstack((shot['upperTheta']['uthetaX'],
            shot['upperTheta']['uthetaY'], shot['upperTheta']['uthetaZ']))
    raw_x_lower = np.vstack((shot['lowerTheta']['lthetaX'],
            shot['lowerTheta']['lthetaY'], shot['lowerTheta']['lthetaZ']))
    normalizedUpperTheta = raw_x_upper / 180.0
    normalizedLowerTheta = raw_x_lower / 180.0
    x_upper[0] = sequence.pad_sequences(normalizedUpperTheta, maxlen=nb_timesteps,
                dtype='float', padding='post', truncating='post', value=0.).T
    x_lower[0] = sequence.pad_sequences(normalizedLowerTheta, maxlen=nb_timesteps,
                dtype='float', padding='post', truncating='post', value=0.).T
    handedness = np.zeros(shape=[1,2])
    handedness[0,isLeftHanded(shot['handedness'])] = 1.0

    print("Loading Model")
    model = loadCompiledShotTypeModel()
    print("Loaded Model Succesfully")
    result = model.predict([handedness, x_upper, x_lower], batch_size=1)
    print("Result: " + str(result))
    resultIndex = result.argmax()
    print(resultIndex)
    shotTypeResult = shotTypeToString(resultIndex)
    print(shotTypeResult)
    return shotTypeResult

def trainShotSpeedCompileFit(speed, test_handedness, x_train_upper, x_train_lower, y_train, x_test_upper, x_test_lower, y_test):
    #Upper hand LSTM input
    encoder_a = Sequential()
    encoder_a.add(LSTM(output_dim=2*input_dim,
                    input_shape=(nb_timesteps, input_dim),
                    activation='sigmoid', inner_activation='hard_sigmoid'))
    #Lower hand LSTM input
    encoder_b = Sequential()
    encoder_b.add(LSTM(output_dim=2*input_dim,
                    input_shape=(nb_timesteps, input_dim),
                    activation='sigmoid', inner_activation='hard_sigmoid'))

    #Merge both LSTM units with concatenation
    merged = Merge([speed, encoder_a, encoder_b], mode='concat')
    decoder = Sequential()
    decoder.add(merged)
    #Use CNNs to reduce to desired intermediate shape
    decoder.add(Dense(output_dim=2*input_dim, activation='relu'))
    #sigmoid activation instead of softmax to avoid normalizing to 1.0
    #1 output signal for the binary classification likelihood
    decoder.add(Dense(1, activation='sigmoid'))

    decoder.compile(loss='mae', optimizer='rmsprop', metrics=['accuracy'])
    decoder.fit([x_train_upper, x_train_lower], y_train,
                batch_size=4*input_dim, nb_epoch=nb_epoch,
                validation_data=([x_test_upper, x_test_lower], y_test))

    printSummary(decoder,x_test_upper, x_test_lower, y_test)
    return decoder

def saveCompiledShotSpeedModel(decoder):
    saved_model = decoder.to_json()
    with open('LSTM_silkymitties_ShotSpeed.json', 'w') as outfile:
        json.dump(saved_model, outfile)
    decoder.save_weights('LSTM_silkymitties_ShotSpeed_weights.h5')

def loadCompiledShotSpeedModel():
    with open('LSTM_silkymitties_ShotSpeed.json', 'r') as infile:
        architecture = json.load(infile)
    model = model_from_json(architecture)
    model.load_weights('LSTM_silkymitties_ShotSpeed_weights.h5')
    return model

def trainShotAccuracyCompileFit(accuracy, test_handedness, x_train_upper,
                x_train_lower, y_train, x_test_upper, x_test_lower, y_test):
    #Upper hand LSTM input
    encoder_a = Sequential()
    encoder_a.add(LSTM(output_dim=2*input_dim,
                input_shape=(nb_timesteps, input_dim),
                activation='sigmoid', inner_activation='hard_sigmoid'))
    #Lower hand LSTM input
    encoder_b = Sequential()
    encoder_b.add(LSTM(output_dim=2*input_dim,
                input_shape=(nb_timesteps, input_dim),
                activation='sigmoid', inner_activation='hard_sigmoid'))

    #Merge both LSTM units with concatenation
    merged = Merge([accuracy, encoder_a, encoder_b], mode='concat')
    decoder = Sequential()
    decoder.add(merged)
    #Use CNNs to reduce to desired output classes
    decoder.add(Dense(output_dim=2*input_dim, activation='relu'))
    #decoder.add(Dropout(0.5))
    decoder.add(Dense(1, activation='softmax'))

    decoder.compile(loss='mse', optimizer='rmsprop', metrics=['accuracy'])
    decoder.fit([x_train_upper, x_train_lower], y_train,
                batch_size=4*input_dim, nb_epoch=nb_epoch,
                validation_data=([x_test_upper, x_test_lower], y_test))

    printSummary(decoder,x_test, y_test)
    return decoder

def saveCompiledShotAccuracyModel(decoder):
    saved_model = decoder.to_json()
    with open('LSTM_silkymitties_ShotAccuracy.json', 'w') as outfile:
        json.dump(saved_model, outfile)
    decoder.save_weights('LSTM_silkymitties_ShotAccuracy_weights.h5')

def loadCompiledShotAccuracyModel():
    with open('LSTM_silkymitties_ShotAccuracy.json', 'r') as infile:
        architecture = json.load(infile)
    model = model_from_json(architecture)
    model.load_weights('LSTM_silkymitties_ShotAccuracy_weights.h5')
    return model

def mean_pred(y_true, y_pred):
    return K.mean(y_pred)

def isLeftHanded(handedness):
    if str(handedness) == "L":
        return 1
    else:
        return 0

def isSlapShot(shotType):
    if str(shotType) == "slap" or str(shotType) == "Slap":
        return 1
    else:
        return 0

def shotTypeToInt(shotType):
    stringShot = str(shotType)
    #if stringShot == "notashot" or stringShot == "none" or stringShot == "NoShot":
    #    return 0
    if stringShot == "slap" or stringShot == "Slap":
        return 0
    elif stringShot == "wrist" or stringShot == "Wrist":
        return 1
    elif stringShot == "snap" or stringShot == "Snap":
        return 2

def shotTypeToString(shotType):
    #if stringShot == "notashot" or stringShot == "none" or stringShot == "NoShot":
    #    return 0
    if shotType == 0:
        return "Slap"
    elif shotType == 1:
        return "Wrist"
    elif shotType == 2:
        return "Snap"

def printSummary(decoder, test_handedness, x_test_upper, x_test_lower, y_test):
    #print(decoder.summary())
    scores = decoder.evaluate([test_handedness, x_test_upper, x_test_lower],
                y_test, verbose=0, batch_size=nb_batch_size)
    print("Accuracy: %2f%%" % (scores[1]*100))
    print(scores)
    plot(decoder, to_file='model.png', show_shapes=True)
    #plot(decoder, to_file='model.png')

if __name__=='__main__':
    #loadFromFile(True)
    #fuseLabelledShots()
    load_training_data()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This module is intended as a quick and dirty drop in replacement of kombu module, covering only
the functionality we use in zmon.
"""

import re
from collections import namedtuple

Connection = namedtuple('Connection', 'hostname port virtual_host')

def parse_redis_conn(conn_str):
    '''
    Emulates kombu.connection.Connection that we were using only to parse redis connection string
    :param conn_str: example 'redis://localhost:6379/0'
    :return: namedtuple(hostname, port, virtual_host)

    >>> parse_redis_conn('localhost:6379')
    Connection(hostname='localhost', port=6379, virtual_host='0')

    >>> parse_redis_conn('localhost:6379/0')
    Connection(hostname='localhost', port=6379, virtual_host='0')
    '''

    if '://' not in conn_str:
        conn_str = 'redis://' + conn_str

    conn_regex = r'redis://([-.a-zA-Z0-9_]+):([0-9]+)(/[0-9]+)?'
    m = re.match(conn_regex, conn_str)
    if not m:
        raise Exception('unable to parse redis connection string: {}'.format(conn_str))
    return Connection(m.group(1), int(m.group(2)), m.group(3).lstrip('/') if m.group(3) else '0')


import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
import pathlib
import models.simple_model
import models.alexnet
from data_loader.loader import get_dataloaders
from utility.util import DataSetName
from utility.learning_utils import *
from byol_pytorch import BYOL
import wandb
import torchvision.models as predefined_models

def freeze_convs(soft_model, sig_model):
    freeze(soft_model)
    freeze(sig_model)
    return soft_model, sig_model

def freeze(model):
    for module in model.layer1:
        freeze_module(module)
    for module in model.layer2:
        freeze_module(module)
    for module in model.layer3:
        freeze_module(module)
    for module in model.layer4:
        freeze_module(module)

def freeze_module(module):
    for param in module.parameters():
        param.requires_grad = False

def get_summary(dist):
    summary = []
    for key, val in dist.items():
        avg = np.mean(np.array(val), axis=0)
        summary.append((key, avg))
    return summary

def get_models(dataset_name):
    model = None
    if dataset_name.name == DataSetName.CIFAR100.name:
        model = models.simple_model.ConvModel(100, in_chanel=3)
    elif dataset_name.name == DataSetName.CIFAR10.name:
        # model = models.simple_model.ConvModel(10, in_chanel=3)
        # model = models.alexnet.AlexNet()
        model = predefined_models.vgg16(pretrained=True, num_classes=10)
    elif dataset_name.name == DataSetName.MNIST.name or dataset_name.name == DataSetName.FashionMNIST.name:
        model = models.simple_model.ConvModel(10, in_chanel=1)

    model.name = 'soft_model'
    return model

def perform_cil_tasks(tasks, model, dataset_name, epochs=2, lr=0.0001, report_step=20, approach='dll'):
    pre_testset, pre_valset = None, None
    is_first_task = True
    test_stat, validation_stat, running_losses = list(), list(), list()
    classes_avg = dict()
    for task in tasks:
        print(f"{Color.RED.value}task: {task}{Color.END.value}")

        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if device == 'cuda':
            num_workers = 2
        else:
            num_workers = 0

        trainloader, validationloader, testloader, pre_testset, pre_valset = get_dataloaders(dataset_name=dataset_name,
                                                                                pre_testset=pre_testset,
                                                                                pre_valset=pre_valset,
                                                                                task_lbls=task, num_workers=num_workers)

        if type(model) == type(None):
            model = get_models(dataset_name=dataset_name)

        if approach == 'dll':
            test_stat, validation_stat, running_losses, model = train(model, trainloader, validationloader,
                                                                      testloader,task, test_stat, validation_stat,
                                                                      running_losses, epochs=epochs, lr=lr,
                                                                      report_step=report_step)

        if approach == 'dll':
            preds, y_true, (acc, f1, p) = predict(model, testloader, task)
            print(f'acc:{acc:.4f}, f1:{f1:.4f}, p:{p:.4f}')
            # plot_confusion_matrix(y_true, preds, classes=task, task=task)

            # plt.plot(range(1, len(running_losses) + 1), running_losses, alpha=.6)
            # plt.xlabel("steps")
            # plt.ylabel("loss")
            # plt.title("Running Loss")
            # plt.show()

    plt.plot(range(1, len(validation_stat) + 1), validation_stat, alpha=.6)
    plt.xlabel("steps")
    plt.ylabel("accuracy")
    plt.title("Validation Accuracy")
    plt.show()

    plt.plot(range(1, len(running_losses) + 1), running_losses, alpha=.6)
    plt.xlabel("epochs")
    plt.ylabel("accuracy")
    plt.title("Test Accuracy")
    plt.show()

    if approach == 'dll':
        running_losses = np.array(running_losses)
    validation_stat = np.array(validation_stat)
    test_stat = np.array(test_stat)

    path = os.path.join(f'../results/batch_cil/{dataset_name}')
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)
    if approach == 'dll':
        running_losses_path = os.path.join(path, 'running_losses')
        np.save(running_losses_path, running_losses)

    validation_stat_path = os.path.join(path, 'validation_stat')
    test_stat_path = os.path.join(path, 'test_stat')
    np.save(validation_stat_path, validation_stat)
    np.save(test_stat_path, test_stat)

def perform_ssl_cil_tasks(tasks, model, dataset_name, epochs=2, lr=0.0001, report_step=20, ssl_dict = {}, ratio = 1.0, batch_size = 16):
    pre_testset, pre_valset = None, None
    is_first_task = True
    test_stat, validation_stat, running_losses = list(), list(), list()
    classes_avg = dict()

    task_validation_loaders = []
    task_test_loaders = []
    tasks_classes_seen = []
    task_cum_accs = []

    for i, task in enumerate(tasks):
        print(f"{Color.RED.value}task: {task}{Color.END.value}")

        tasks_classes_seen.append(task)

        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if device == 'cuda':
            num_workers = 0
        else:
            num_workers = 0

        trainloader, validationloader, testloader, _, _ = get_dataloaders(dataset_name=dataset_name,
                                                                                pre_testset=None,
                                                                                pre_valset=None,
                                                                                task_lbls=task, num_workers=num_workers, train_batch_size=batch_size)
        task_validation_loaders.append(validationloader)
        task_test_loaders.append(validationloader)

        test_stat, validation_stat, running_losses, model = train_ssl(model, trainloader, validationloader,
                                                                      testloader,(i,task), test_stat, validation_stat,
                                                                      running_losses, epochs=epochs, lr=lr,
                                                                      report_step=report_step, ssl_dict = ssl_dict, ratio = ratio)

        accs = []
        for task_id, task_loader in enumerate(task_validation_loaders):

            preds, y_true, (acc, f1, p) = predict(model, task_loader, (task_id,tasks_classes_seen[task_id]))
            print('Task {d} val acc : {acc}'.format(d=task_id+1, acc = acc))
            accs.append(acc)
            # plot_confusion_matrix(y_true, preds, classes=task, task=task)

        # plt.plot(range(1, len(running_losses) + 1), running_losses, alpha=.6)
        # plt.xlabel("steps")
        # plt.ylabel("loss")
        # plt.title("Running Loss")
        # plt.show()

        cum_acc = np.array(accs).mean()

        #wandb logging
        wandb.log(
            {
                "Cumulative accuracy": cum_acc,
                "Task ID": task_id,
            }
                )

        task_cum_accs.append(cum_acc)

        print('Cumulative accuracy: {cum_acc}'.format(cum_acc=cum_acc))

    plt.plot(range(len(task_cum_accs)), task_cum_accs)
    plt.xlabel("steps")
    plt.ylabel("accuracy")
    plt.title("Validation Accuracy")
    plt.show()

    running_losses = np.array(running_losses)
    validation_stat = np.array(validation_stat)
    test_stat = np.array(test_stat)

    path = os.path.join(f'../results/batch_cil/{dataset_name}')
    pathlib.Path(path).mkdir(parents=True, exist_ok=True)

    running_losses_path = os.path.join(path, 'running_losses')
    np.save(running_losses_path, running_losses)

    validation_stat_path = os.path.join(path, 'validation_stat')
    test_stat_path = os.path.join(path, 'test_stat')
    np.save(validation_stat_path, validation_stat)
    np.save(test_stat_path, test_stat)


from tkinter import E
import numpy as np
import torch
import pickle
import torch.nn as nn
from model_paras_count.pre_process import *
from model_paras_count.count import *
from network.plain_network import *
from network.batchcrypt import BatchCrypt
from socket import socket, AF_INET, SOCK_STREAM
from train_params import *
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import math
import time
import logging

def server_func():
    logging.basicConfig(level=LOG_LEVEL,
                        format='%(asctime)s %(filename)s %(levelname)s %(message)s',
                        datefmt='%a %d %b %Y %H:%M:%S',
                        filename="{}/server.log".format(LOG_DIR),
                        filemode="w")
    logging.debug("Debug mode")

    """ Init public key """
    model = PlainNetwork(encry_method=BACKEND)
    encryptor = model.encryptor

    """ Init socket """

    ServerSocket = socket()
    ServerSocket.bind((SERVER_IP, SERVER_PORT))
    ServerSocket.listen(5)

    leader_conn, address = ServerSocket.accept()

    """ Send public key """
    sendData = pickle.dumps(encryptor)
    connections = {}
    for _ in range(CLIENT_NUM):
        conn, address = ServerSocket.accept()
        conn_id = pickle.loads(conn.recv(BUFF_SIZE))
        connections[conn_id] = conn
        conn.send(sendData)

    """** Local Anchors **"""

    logging.debug("init success")
    logging.info("Params: MODE:{} CLIENTS:{} K:{} SPARSE:{} DATASET:{}".\
                 format(SWITCH_MODE,CLIENT_NUM,K,SPARSE,DATASET))
    protocol_switch = '1'
    switch_back = False
    batch_num = math.ceil(TRAINDATA_SIZE // CLIENT_NUM / BATCH_SIZE) if CLIENT_NUM < 50 else math.ceil(TRAINDATA_SIZE // 25 / BATCH_SIZE)
    model_chose = np.zeros((EPOCH,batch_num,CLIENT_NUM)) # save distribution of selected clients
    for epoch in range(EPOCH):
        for b in range(batch_num):
            logging.debug("epoch: {} batch seq: {}".format(epoch,b))
            print("epoch:",epoch,"step:",b)

            '''recv switch from leader and broadcast'''
            # recv switch info from leader
            protocol_switch = pickle.loads(leader_conn.recv(BUFF_SIZE))

            #leader_conn.send(pickle.dumps(COMP_MSG))
            logging.debug("receive switch from leader")
            time.sleep(0.5)

            # broadcast switch info to clients
            for i in range(1,CLIENT_NUM):
                conn = connections[i]
                conn.send(pickle.dumps(protocol_switch))

            logging.debug("broadcast switch")
            #time.sleep(0.8)

            '''trans process'''
            if PARASCOUNT and protocol_switch == '2':
                # recieve hash of anchors
                hash_list = None
                for i in range(CLIENT_NUM):
                    conn = connections[i]
                    remain = pickle.loads(conn.recv(BUFF_SIZE))
                    conn.send(pickle.dumps(READY_MSG))
                    recvData = b""
                    while remain > 0:
                        each_data = conn.recv(BUFF_SIZE)
                        remain -= len(each_data)
                        recvData += each_data
                    hashs = pickle.loads(recvData)
                    if i == 0:
                        hash_list = hashs
                    else:
                        hash_list = np.vstack((hash_list,hashs))

                logging.debug("receive clients' hash list")

                # Model Paras Count
                hash_list = hash_list.T
                alloc,selected,sparse_index = pick_represents(hash_list,K,SPARSE,ENCRYPT_BATCHSIZE)
                alloc_num = len(alloc.keys())
                selected_index = []
                selected_info = dict()
                selected_freq = np.zeros(CLIENT_NUM).astype(int)
                begin,end = 0,0
                for c in selected:
                    begin = end
                    end = begin + len(alloc[c])
                    selected_info[c] = {'index':alloc[c],'begin':begin,'end':end}
                    selected_index += alloc[c]
                    selected_freq[c] = len(alloc[c])
                print([(key,len(selected_info[key]['index'])) for key in selected_info.keys()])
                selected_index += sparse_index
                restore_mask = np.argsort(np.array(selected_index))
                selected_vec = one_hot(selected,CLIENT_NUM)
                model_chose[epoch,b,:] = selected_freq
                # If only one client in selection, switch back to protocol 1
                if SWITCH_MODE == 'thre' and len(selected) <= 1:
                    switch_back = True
                    print("switch back to protocol 1")
                else:
                    switch_back = False

                np.save(CHOICE_PATH,model_chose)

                # send indexes to selected clients
                # send selected info to clients 2022.4.26
                for i in range(CLIENT_NUM):
                    if selected_vec[i]:
                        sendData = pickle.dumps({'switch_back':switch_back,'selected':True,'remask':restore_mask,'sparse':sparse_index,'selected_info':selected_info[i]})
                    else:
                        sendData = pickle.dumps({'switch_back':switch_back,'selected':False,'remask':restore_mask,'sparse':sparse_index})
                    remain = len(sendData)
                    conn = connections[i]
                    conn.send(pickle.dumps(remain))
                    resp = conn.recv(BUFF_SIZE)
                    if pickle.loads(resp) != READY_MSG:
                        raise ValueError("receive unrecognized message!")
                    conn.send(sendData)

                logging.debug("send alloc info to clients")

                if not switch_back:
                    # receive anchors and residues
                    grad_anchors,grad_residues = None, None
                    first_conn = min(list(alloc.keys()))
                    for i in range(CLIENT_NUM):
                        conn = connections[i]
                        if selected_vec[i]:
                            # recv anchors
                            remain = pickle.loads(conn.recv(BUFF_SIZE))
                            conn.send(pickle.dumps(READY_MSG))
                            recvData = b""
                            while remain > 0:
                                each_data = conn.recv(BUFF_SIZE)
                                remain -= len(each_data)
                                recvData += each_data

                            anch_res = pickle.loads(recvData)
                            if i == first_conn:
                                grad_anchors = np.array(anch_res['anchors'])
                                grad_residues = np.array(anch_res['residues'])
                            else:
                                grad_anchors += np.array(anch_res['anchors'])
                                grad_residues += np.array(anch_res['residues'])
                            #conn.send(pickle.dumps(COMP_MSG))

                    logging.debug("receive local anchors & residues from clients")

                    # broadcast global gradients
                    grad_residues /= alloc_num
                    global_grads = [grad_anchors,grad_residues]
                    sendData = pickle.dumps(global_grads)
                    remain = len(sendData)
                    for i in range(CLIENT_NUM):
                        conn = connections[i]
                        conn.send(pickle.dumps(remain))
                        resp = conn.recv(BUFF_SIZE)
                        if pickle.loads(resp) != READY_MSG:
                            raise ValueError("receive unrecognized message!")
                        conn.send(sendData)
                    logging.debug("send global anchors & residues to clients")
                else:
                    logging.debug("switch back to protocol 1 in epoch {} step {}".format(epoch,b))
            elif protocol_switch == '1':
                # receive local encrypted gradients
                agg_grad = None
                for i in range(CLIENT_NUM):
                    conn = connections[i]
                    remain = pickle.loads(conn.recv(BUFF_SIZE))
                    conn.send(pickle.dumps(READY_MSG))
                    recvData = b""
                    while remain > 0:
                        each_data = conn.recv(BUFF_SIZE)
                        remain -= len(each_data)
                        recvData += each_data
                    grads = pickle.loads(recvData)
                    if i == 0:
                        agg_grad = np.array(grads)
                    else:
                        agg_grad += np.array(grads)
                logging.debug("receive local gradients from clients")

                # send global encrypted gradients
                sendData = pickle.dumps(agg_grad)
                remain = len(sendData)
                for i in range(CLIENT_NUM):
                    conn = connections[i]
                    conn.send(pickle.dumps(remain))
                    resp = conn.recv(BUFF_SIZE)
                    if pickle.loads(resp) != READY_MSG:
                        raise ValueError("receive unrecognized message!")
                    conn.send(sendData)
                logging.debug("send global gradients to clients")
            else:
                print('unknown protocol!')
                assert 0


from tqdm import tqdm

from MCTS import MCTS
from BinaryTree import BinaryTree
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(15)

def run_experiment(max_iterations, dynamic_c=False):
    """
    Run a single experiment of a sequence of MCTS searches to find the optimal path.
    :param max_iterations: Number of iterations to run the MCTS.
    :param dynamic_c: Boolean indicating whether to use a dynamic c or not.
    :return: value of the optimal path found from the search
    """
    tree = BinaryTree(depth=12, b=20, tau=3)
    best_leaf = max(tree.leaves)
    mcts = MCTS(max_iterations=max_iterations, c=2)

    optimal_path = ""
    while tree.depth > 0:
        # search the best direction
        direction = mcts.search(tree, dynamic_c=dynamic_c, verbose=False)
        optimal_path += direction

        # update the tree
        tree.update_root(direction)

    # return the distance of the optimal path found from the search wrt the best leaf
    return sum(1 for a, b in zip(optimal_path, best_leaf.address) if a != b)

def main():
    # compute statistics for static c and dynamic c
    n_iterations = np.logspace(0.7, 3, num=18, base=10, dtype=int)
    values_static_c = [run_experiment(max_iterations=n, dynamic_c=False) for n in tqdm(n_iterations, desc='Execute MCTS with c=2', unit=' experiment')]
    values_dynamic_c = [run_experiment(max_iterations=n, dynamic_c=True) for n in tqdm(n_iterations, desc='Execute MCTS with dynamic c', unit=' experiment')]

    # plot the results
    plt.figure(figsize=(8, 4))
    plt.plot(n_iterations, values_dynamic_c, '-o', label="MCTS with dynamic c")
    plt.plot(n_iterations, values_static_c, '-o', label="MCTS with c=2")
    plt.xlabel("Number of iterations")
    plt.ylabel("Distance of the optimal path from the best leaf")
    plt.title("Compare the value of the optimal path found by MCTS with and without dynamic c")
    plt.grid(linestyle='--', linewidth=1)
    plt.xscale("log")
    plt.xticks(n_iterations, n_iterations)
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()


from flask import Flask
from flask_bootstrap import Bootstrap
from flask_mail import Mail
from flask_moment import Moment
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager
from flask_pagedown import PageDown
from config import config
from flask_admin import Admin, BaseView, expose
from flask_admin.contrib.fileadmin import FileAdmin
from flask_babelex import Babel
import os.path as op

bootstrap = Bootstrap()
mail = Mail()
moment = Moment()
db = SQLAlchemy()
pagedown = PageDown()
flas_admin = Admin(name='smart-iiot')
babel = Babel()

# flask-login
login_manager = LoginManager()
login_manager.login_view = 'auth.login'

# flask-admin add views
from app.admin import UserAdminView, TestAdminView, UserModelView, LoraGatewayModelView, LoraNodeModelView, NodeMqttTransFuncModelView, PowerIoModelView, RelayCurrentRs485FuncModelView, \
GrainStorehouseModelView, GrainBarnModelView, GrainTempModelView, AlarmLevelSettingModelView, AlarmStatusModelView, AlarmTypesModelView, AlarmRecordsModelView

#
# flas_admin.add_view(UserAdminView(name='UserAdmin', category='UserAdmin'))
# flas_admin.add_view(TestAdminView(name='test', endpoint='test', category='UserAdmin'))

flas_admin.add_view(GrainStorehouseModelView(db.session, name='GrainStorehouse', endpoint='grain_storehouse', category='GrainAdmin'))
flas_admin.add_view(GrainBarnModelView(db.session, name='GrainBarn', endpoint='grain_barn', category='GrainAdmin'))
flas_admin.add_view(GrainTempModelView(db.session, name='GrainTemp', endpoint='grain_temps', category='GrainAdmin'))

flas_admin.add_view(LoraGatewayModelView(db.session, name='LoraGateway', endpoint='lora_gateway', category='LoraAdmin'))
flas_admin.add_view(LoraNodeModelView(db.session, name='LoraNode', endpoint='lora_node', category='LoraAdmin'))
flas_admin.add_view(NodeMqttTransFuncModelView(db.session, name='NodeMqttTransFunc', endpoint='node_mqtt_trans_func', category='LoraAdmin'))
flas_admin.add_view(PowerIoModelView(db.session, name='PowerIo', endpoint='power_io', category='LoraAdmin'))
flas_admin.add_view(RelayCurrentRs485FuncModelView(db.session, name='RelayCurrentRs485Func', endpoint='relay_current_rs485_func', category='LoraAdmin'))

flas_admin.add_view(AlarmStatusModelView(db.session, name='AlarmStatus', endpoint='alarm_status', category='AlarmAdmin'))
flas_admin.add_view(AlarmTypesModelView(db.session, name='AlarmTypes', endpoint='alarm_types', category='AlarmAdmin'))
flas_admin.add_view(AlarmRecordsModelView(db.session, name='AlarmRecords', endpoint='alarm_records', category='AlarmAdmin'))
flas_admin.add_view(AlarmLevelSettingModelView(db.session, name='AlarmLevelSetting', endpoint='alarm_level_setting', category='AlarmAdmin'))

flas_admin.add_view(UserModelView(db.session, name='User', endpoint='user', category='UserAdmin'))

path = op.join(op.dirname(__file__), 'static')
print(path)
flas_admin.add_view(FileAdmin(path, '/static/', name='Static Files'))

def create_app(config_name):
    app = Flask(__name__)
    app.config.from_object(config[config_name])
    config[config_name].init_app(app)
    # bable config for i18n
    app.config['BABEL_DEFAULT_LOCALE'] = 'zh_CN'

    if app.config['SSL_REDIRECT']:
        from flask_sslify import SSLify
        sslify = SSLify(app)

    configure_extensions(app)

    register_blueprints(app)

    return app

def configure_extensions(app):
    """configure flask extensions
    """
    bootstrap.init_app(app)
    mail.init_app(app)
    moment.init_app(app)
    db.init_app(app)
    login_manager.init_app(app)
    pagedown.init_app(app)
    babel.init_app(app)
    flas_admin.init_app(app)

def register_blueprints(app):
    """register all blueprints for application
    """

    from .main import main as main_blueprint
    app.register_blueprint(main_blueprint)
    from .auth import auth as auth_blueprint
    app.register_blueprint(auth_blueprint, url_prefix='/auth')
    from .api import api as api_blueprint
    app.register_blueprint(api_blueprint, url_prefix='/api/v1')


from collections import deque, defaultdict
def diff(str1, str2):
    diffNum = 0
    for c1, c2 in zip(str1, str2):
        if c1 != c2:
            diffNum += 1

        if diffNum > 1:
            return False
    return True

def bfs(begin, target, words, graph):
    visited = dict.fromkeys(words, False)
    queue = deque([[begin, 0]])
    visited[begin] = True
    answer = 0

    while queue:
        before, result = queue.popleft()
        if before == target:
            answer = result
            return answer

        for word in graph[before]:
            if not visited[word]:
                queue.append([word, result + 1])
                visited[word] = True
    return answer

def solution(begin, target, words):
    words.insert(0, begin)
    graph = defaultdict(list)

    for str1 in words:
        for str2 in words:
            if str1 == str2 or str2 == begin:
                continue

            if diff(str1, str2):
                graph[str1].append(str2)
    return bfs(begin, target, words, graph)

print(solution("hit", "cog", ["hot", "dot", "dog", "lot", "log", "cog"]))
print(solution("hit", "cog", ["hot", "dot", "dog", "lot", "log"]))


import asyncio, aiohttp, collections, tqdm
from aiohttp import web
from flags2_common import main, HTTPStatus, Result, save_flag

DEFAULT_CONCUR_REQ = 5
MAX_CONCUR_REQ = 1000

class FetchError(Exception):
    def __init__(self, country_code):
        self.country_code = country_code

async def get_flag(base_url, cc):
    url = '{}/{cc}/{cc}.gif'.format(base_url, cc=cc.lower())
    resp = await aiohttp.request('GET', url)
    if resp.status == 200:
        image = await resp.read()
        return image
    elif resp.status == 404:
        raise web.HTTPNotFound()
    else:
        raise aiohttp.HttpProcessingError(code=resp.status, message=resp.reason, headers=resp.headers)

async def download_one(cc, base_url, semaphore, verbose):
    try:
        with (await semaphore):
            image = await get_flag(base_url, cc)
    except web.HTTPNotFound:
        status = HTTPStatus.not_found
        msg = 'not found'
    except Exception as exc:
        raise FetchError(cc) from exc
    else:
        save_flag(image, cc.lower() + '.gif')
        status = HTTPStatus.ok
        msg = 'ok'
    if verbose and msg:
        print(cc, msg)
    return Result(status, cc)

def download_many(cc_list):
    loop = asyncio.get_event_loop()
    to_do = [download_one(cc) for cc in sorted(cc_list)]
    wait_coro = asyncio.wait(to_do)
    res, _ = loop.run_until_complete(wait_coro)
    loop.close()
    return len(res)

if __name__ == '__main__':
    main(download_many)