import math
import torch
import torch.nn.functional as F

from fairseq import metrics, utils
from fairseq.criterions import register_criterion
from fairseq.data.data_utils import post_process, collate_tokens
from fairseq.data.data_utils import lengths_to_padding_mask
from fairseq.logging.meters import safe_round

from .label_smoothed_cross_entropy import LabelSmoothedCrossEntropyCriterion
import torch.nn as nn

@register_criterion("label_smoothed_cross_entropy_with_w2v_ctc_shrink_joint_AT")
class LabelSmoothedCrossEntropywithW2vCtcShrinkJointAT(
    LabelSmoothedCrossEntropyCriterion
):
    def __init__(self, task, sentence_avg, label_smoothing, post_process="letter", ctc_weight=0.0, contrastive_alpha=0.0,
                 contrastive_beta=0.0, contrastive_temperature=1.0, decrease_step=5000, get_similarity=False, is_shrink="",
                 train_st_without_ctc=False, use_token_contrastive=False, use_two_contrastive=False, add_proj_norm=False,
                 use_double_ctc=False, use_ctc_cluster=False, word_align=False, ban_cl_step=-1):
        super().__init__(task, sentence_avg, label_smoothing)
        self.blank_idx = task.target_dictionary.index(task.blank_symbol) if hasattr(task, 'blank_symbol') else 0
        self.pad_idx = task.target_dictionary.pad()
        self.eos_idx = task.target_dictionary.eos()
        self.report_accuracy = True

        #assert 0 <= ctc_weight <= 1
        self.ctc_weight = ctc_weight
        if self.ctc_weight >= 0:
            assert getattr(task, "src_dict", None) is not None, "CTC need a source dictionary."
            self.zero_infinity = True
            self.post_process = post_process
        self.contrastive_alpha = contrastive_alpha
        self.contrastive_beta = contrastive_beta
        self.contrastive_temperature = contrastive_temperature
        self.decrease_step = decrease_step
        self.get_similarity = get_similarity
        self.similarity_function = nn.CosineSimilarity(dim=-1)
        self.post_process = post_process
        self.is_shrink = is_shrink
        self.train_st_without_ctc = train_st_without_ctc
        self.use_token_contrastive = use_token_contrastive
        self.use_two_contrastive = use_two_contrastive
        self.add_proj_norm = add_proj_norm
        self.use_double_ctc = use_double_ctc
        self.use_ctc_cluster = use_ctc_cluster
        self.word_align = word_align
        self.ban_cl_step = ban_cl_step
        self.at_training = task.at_training
        self.at_level = task.at_level
        self.at_scale = task.at_scale
        self.at_nopad = task.at_nopad
        self.at_nomute = task.at_nomute
        self.at_win = {'size':20, 'stride':10}     # works when "at_level" is "window"                      # //need_check//
        self.embedding_l2norm = task.embedding_l2norm
        self.merge_mt_st = task.merge_mt_st
        assert (self.at_level == "sentence"
             or self.at_level == "token"
             or self.at_level == "window") , "Args: --at-level is wrong, expect \"sentence\" or \"token\" or \"window\"."      # //need_check//

    @staticmethod
    def add_args(parser):
        """Add criterion-specific arguments to the parser."""
        LabelSmoothedCrossEntropyCriterion.add_args(parser)
        parser.add_argument(
            "--zero-infinity",
            default=True,
            type=bool,
            help="zero inf loss when source length <= target length",
        )
        parser.add_argument(
            "--ctc-weight",
            default=0.0,
            type=float,
            metavar="D",
            help="weight of CTC loss",
        )
        parser.add_argument(
            "--contrastive-alpha",
            type=float,
            default=0.0,
            help="The contrastive loss weight",
        )
        parser.add_argument(
            "--contrastive-beta",
            type=float,
            default=0.0,
            help="The kd contrastive loss weight",
        )
        parser.add_argument(
            "--decrease-step",
            type=int,
            default=5000,
            help="The number of step to descend beta weight",
        )
        parser.add_argument(
            "--contrastive-temperature",
            type=float,
            default=1.0,
            help="The temperature to adjust contrastive loss",
        )
        parser.add_argument(
            "--get-similarity",
            type=bool,
            default=False,
            help="To get the similarity between acoustic and textual encoder",
        )
        parser.add_argument(
            "--is-shrink",
            type=str,
            default="",
            help="To remove the  wav2vec blank output",
        )
        parser.add_argument(
            "--post-process",
            default="letter",
            type=str,
            help="how to post process predictions into words. can be letter, "
            "wordpiece, BPE symbols, etc. "
            "See fairseq.data.data_utils.post_process() for full list of options",
        )
        parser.add_argument(
            "--train-st-without-ctc",
            default=False,
            type=bool,
            help="train st task without ctc loss but with constrastive loss",
        )
        parser.add_argument(
            "--use-token-contrastive",
            default=False,
            type=bool,
            help="use token level to contrastive learning",
        )
        parser.add_argument(
            "--use-two-contrastive",
            default=False,
            type=bool,
            help="use both token and sentence level to contrastive learning",
        )
        parser.add_argument(
            "--use-double-ctc",
            default=False,
            type=bool,
            help="use both pos and word level ctc",
        )
        parser.add_argument(
            "--add-proj-norm",
            default=False,
            type=bool,
            help="use proj norm to avoid over-fitting",
        )
        parser.add_argument(
            "--use-ctc-cluster",
            default=False,
            type=bool,
            help="use ctc cluster",
        )
        parser.add_argument(
            "--word-align",
            default=False,
            type=bool,
            help="use constrastive loss with word level",
        )
        parser.add_argument(
            "--ban-cl-step",
            default=-1,
            type=int,
            help="ban cl after a certain step",
        )

    def forward(self, model, sample, per_task, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        if per_task == "st":
            src_tokens, src_lengths, prev_output_tokens = sample["net_input"].values()
            #encoder_out = model.acoustic_encoder(src_tokens=src_tokens, src_lengths=src_lengths)
            #padding_mask = lengths_to_padding_mask(src_lengths)
            #print(src_tokens.shape,src_lengths.shape,padding_mask.shape)
            encoder_out = model.acoustic_encoder(src_tokens, src_lengths)
            net_output = model.decoder(
                prev_output_tokens=prev_output_tokens, encoder_out=encoder_out
            )
            # task_prob = model.task_net(encoder_out["encoder_out"][0])

            loss, nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
            sample_size = (
                sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
            )
            logging_output = {
                "trans_loss": utils.item(loss.data) if reduce else loss.data,
                "nll_loss": utils.item(nll_loss.data) if reduce else nll_loss.data,
                "ntokens": sample["ntokens"],
                "nsentences": sample["target"].size(0),
                "sample_size": sample_size,
            }

            if self.report_accuracy:
                n_correct, total = self.compute_accuracy(model, net_output, sample)
                logging_output["n_correct"] = utils.item(n_correct.data)
                logging_output["total"] = utils.item(total.data)

            if model.training:
                if self.ctc_weight > 0:
                    #ctc_loss = self.compute_ctc_loss(model, sample, encoder_out)
                    #logging_output["ctc_loss"] = utils.item(ctc_loss.data)
                    w2v_ctc_loss, tmp_logging_output = self.compute_ctc_loss(model, sample, encoder_out, True, self.train_st_without_ctc)
                    logging_output.update(tmp_logging_output)
                    loss = (1 - self.ctc_weight) * loss + self.ctc_weight * w2v_ctc_loss
            else:
                w2v_ctc_loss, tmp_logging_output = self.compute_ctc_loss(model, sample, encoder_out, True, self.train_st_without_ctc)
                loss = (1 - self.ctc_weight) * loss + self.ctc_weight * w2v_ctc_loss
                logging_output.update(tmp_logging_output)

            # logging_output["loss"] = utils.item(loss.data) if reduce else loss.data

            # ************/ Add L2 norm /************ #                                                         //need_check//
            if self.embedding_l2norm:
                l2_loss = model.get_acoustic_embedding_L2_norm()
                logging_output["l2_loss"] = utils.item(l2_loss.data) if reduce else l2_loss.data
                loss += l2_loss
                # print(logging_output["loss"], logging_output["l2_loss"], loss)
                # assert False

            logging_output["loss"] = utils.item(loss.data) if reduce else loss.data

            if self.merge_mt_st:
                # do mt task
                _, _, prev_output_tokens = sample["net_input"].values()
                src_tokens, src_lengths = sample["transcript"]['tokens'], sample["transcript"]['lengths']
                encoder_out = model.textual_encoder(src_tokens, src_lengths)
                net_output = model.decoder(
                    prev_output_tokens=prev_output_tokens, encoder_out=encoder_out
                )
                mt_loss, mt_nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
                mt_sample_size = (
                    sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
                )

                logging_output["trans_loss"] += (utils.item(mt_loss.data) if reduce else mt_loss.data)
                logging_output["nll_loss"] += (utils.item(mt_nll_loss.data) if reduce else mt_nll_loss.data)
                logging_output["ntokens"] += sample["ntokens"]
                logging_output["nsentences"] += sample["target"].size(0)
                logging_output["sample_size"] += mt_sample_size

                loss += mt_loss
                logging_output["loss"] = (utils.item(loss.data) if reduce else loss.data)

        # ************/ Add AT per-task /************ #
        elif per_task == 'at':                  # //need_check//

            assert self.at_training, 'BUG HERE: should not do AT task!'

            st_prob, st_id = None, None
            mt_prob, mt_id = None, None
            # st
            if sample['st'] is not None:
                st_src_tokens, st_src_lengths = sample['st']["net_input"]['src_tokens'], sample['st']["net_input"]['src_lengths']
                st_encoder_out = model.acoustic_encoder(st_src_tokens, st_src_lengths)

                # required data
                st_encoder_output = st_encoder_out["encoder_out"][0]                    # [len, bs, dim]
                st_encoder_padding_mask = st_encoder_out["encoder_padding_mask"][0] if self.at_nomute else torch.zeros(st_encoder_out["encoder_padding_mask"][0].shape).bool().cuda()     # [bs, len]
                if self.at_level == "sentence":
                    st_prob = model.task_net(st_encoder_output, st_encoder_padding_mask if self.at_nomute else None)
                    st_id = torch.zeros(st_encoder_padding_mask.shape[0], 1)
                elif self.at_level == "token":
                    encoder_output = st_encoder_output.transpose(0,1).reshape(-1, st_encoder_output.shape[-1])   # [bs * len, dim=512]
                    encoder_padding_mask = st_encoder_padding_mask.reshape(-1)  # [bs * len]
                    encoder_output = encoder_output[~encoder_padding_mask]      # 去除mask为True
                    st_prob = model.task_net(encoder_output.unsqueeze(0))           # [1, len*bs, dim=512] --> [len*bs, 1]
                    st_id = torch.zeros(st_prob.shape[0], 1)                        # [len*bs, 1]
                elif self.at_level == "window":
                    # 计算窗口个数
                    encoder_output = st_encoder_output.transpose(0,1)   # [bs, len, dim]
                    encoder_output[st_encoder_padding_mask.unsqueeze(-1).expand(encoder_output.shape)] = 0.0    # 将padding位置设为0，避免影响mean运算
                    _shape = encoder_output.shape # [bs, len, dim]
                    # print(" > 1. ", _shape, st_encoder_output.dtype)
                    length = _shape[1]
                    num_win = max(0, (length - self.at_win['size'])) // self.at_win['stride'] + 1 + (0 if (max(0, (length - self.at_win['size'])) % self.at_win['stride']) == 0 else 1)
                    # print(" > 2. ", length, self.at_win['size'], self.at_win['stride'], num_win)
                    task_input = torch.zeros(num_win, _shape[0], _shape[-1], dtype=torch.float16).cuda()  # [num_win, bs, dim]  保存结果
                    # print(" > 3. ", task_input.shape, task_input.dtype)
                    for i in range(num_win):
                        s = i * self.at_win['stride']
                        e = min(s + self.at_win['size'], length)    # 不超过length
                        calc_output = encoder_output[:,s:e,:]   # [bs, size, dim]
                        calc_mask = st_encoder_padding_mask[:, s:e] # [bs, size]
                        calc_len = (~calc_mask).long().sum(dim=1).unsqueeze(-1) # [bs, 1]
                        calc_len[calc_len==0.0] = 1 # 特殊0值处理
                        task_input[i,:,:] = (calc_output.sum(dim=1) / calc_len)     # [bs, dim=512]
                    # remove padding vectors ::Not implemented ::Not needed for now

                    # print(" > 4. ", task_input.shape)
                    task_input = task_input.transpose(0, 1).reshape(-1, task_input.shape[-1]) # [bs*num_win, dim=512]
                    st_prob = model.task_net(task_input.unsqueeze(0))           # [1, num_unzero, dim=512] --> [num_unzero, 1]
                    st_id = torch.zeros(st_prob.shape[0], 1)                    # [num_unzero, 1]
                else:
                    raise NotImplementedError

            # mt
            if sample['mt'] is not None:
                mt_src_tokens, mt_src_lengths = sample['mt']["net_input"]['src_tokens'], sample['mt']["net_input"]['src_lengths']
                mt_encoder_out = model.textual_encoder(mt_src_tokens, mt_src_lengths)

                # required data
                mt_encoder_output = mt_encoder_out["encoder_out"][0]                    # [len, bs, dim]
                mt_encoder_padding_mask = lengths_to_padding_mask(mt_src_lengths) if self.at_nopad else torch.zeros(mt_encoder_out["encoder_padding_mask"][0].shape).bool().cuda()  # [bs, len]

                if self.at_level == "sentence":
                    mt_prob = model.task_net(mt_encoder_output, mt_encoder_padding_mask if self.at_nopad else None)
                    mt_id = torch.ones(mt_encoder_padding_mask.shape[0], 1)
                elif self.at_level == "token" or self.at_level == 'window':
                    encoder_output = mt_encoder_output.transpose(0,1).reshape(-1, mt_encoder_output.shape[-1])   # [bs * len, dim=512]
                    encoder_padding_mask = mt_encoder_padding_mask.reshape(-1)  # [bs * len]
                    encoder_output = encoder_output[~encoder_padding_mask]      # 去除mask为True
                    mt_prob = model.task_net(encoder_output.unsqueeze(0))           # [1, len*bs, dim=512] --> [len*bs, 1]
                    mt_id = torch.ones(mt_prob.shape[0], 1)                         # [len*bs, 1]
                else:
                    raise NotImplementedError

            # merge st and mt
            task_prob = torch.cat((() if st_prob is None else (st_prob,)) + (() if mt_prob is None else (mt_prob,)), dim=0)
            task_id = torch.cat((() if st_id is None else (st_id,)) + (() if mt_id is None else (mt_id,)), dim=0)
            # print(" > 7. ", task_prob.shape, task_id.shape)

            # cuda
            if torch.cuda.is_available():
                task_id = task_id.cuda()

            # calculate loss
            loss, sample_size = None, None

            if self.at_level == "sentence":
                weight = torch.cat(
                        (() if sample['st'] is None else ((~st_encoder_padding_mask).long().sum(dim=1),))
                        + (() if sample['mt'] is None else ((~mt_encoder_padding_mask).long().sum(dim=1),)),
                        dim=0
                ).unsqueeze(1)                              # [bz] + [bz]  src token length
                """
                weight = torch.cat((() if sample['st'] is None else (sample['st']["target_lengths"],))
                                    + (() if sample['mt'] is None else (sample['mt']["target_lengths"],))
                                    , dim=0).unsqueeze(1)                          # target length"""
                # bce loss
                bce_loss = F.binary_cross_entropy_with_logits(
                        task_prob.float(),
                        task_id,
                        weight=weight,
                        reduction = "mean"
                )
                sentence_sample_size = task_prob.shape[0]   # [bz + bz]
                # merge result
                loss = bce_loss if loss is None else (loss + bce_loss)
                sample_size = sentence_sample_size if sample_size is None else (sample_size + sentence_sample_size)
            elif self.at_level == "token" or self.at_level == 'window':
                mean_set = ((0 if sample['st'] is None else st_encoder_padding_mask.shape[0])
                              + (0 if sample['mt'] is None else mt_encoder_padding_mask.shape[0]))    # bz + bz  =>  num of sentence
                weight = torch.cat(
                    (() if sample['st'] is None else (torch.ones(st_prob.shape[0], dtype=torch.float16) * (self.at_win['stride'] if self.at_level == "window" else 1),))
                    + (() if sample['mt'] is None else (torch.ones(mt_prob.shape[0], dtype=torch.float16),))
                ).unsqueeze(1) / mean_set  # [l, 1]

                # bce loss
                bce_loss = F.binary_cross_entropy_with_logits(
                        task_prob.float(),
                        task_id,
                        weight=weight.cuda(),        # bug
                        reduction = "sum"
                )
                token_sample_size = task_prob.shape[0]
                # merge result
                loss = bce_loss if loss is None else (loss + bce_loss)
                sample_size = token_sample_size if sample_size is None else (sample_size + token_sample_size)
            else:
                raise NotImplementedError

            loss *= self.at_scale       # set scale for at loss

            # print(st_encoder_output.shape, mt_encoder_output.shape, st_prob.shape, mt_prob.shape, loss, sample_size)
            # assert False

            logging_output = {
                "task_loss": utils.item(loss.data) if reduce else loss.data,
                "sample_size": sample_size,
            }

        elif per_task == "mt":

            assert not self.merge_mt_st, 'BUG HERE: should not do MT task!'

            src_tokens, src_lengths, prev_output_tokens = sample["net_input"].values()
            #with open("mt_token","a")as fo:
            #    l = [str(i) for i in src_tokens.tolist()]
            #    fo.write(" ".join(l)+"\n")
            #    #_, embed_out = model.textual_encoder.forward_embedding(src_tokens)
            #    embed_out = model.textual_encoder.embed_tokens(src_tokens)
            #    l = [str(i) for i in embed_out.mean(-1).tolist()]
            #    fo.write(" ".join(l)+"\n")
            #encoder_out = model.acoustic_encoder(src_tokens=src_tokens, src_lengths=src_lengths)
            #padding_mask = lengths_to_padding_mask(src_lengths)
            #print(src_tokens.shape,src_lengths.shape,padding_mask.shape)
            #max_l = int(src_tokens.shape[1] * 0.1)
            #src_lengths += max_l
            #src_lengths += 1
            #min_l = src_lengths.min()
            #batch = src_tokens.shape[0]
            #index = (torch.rand(1).cuda() * min_l).int()
            #insert = torch.zeros([batch,1],dtype=torch.int).cuda()
            #p1,p2 = src_tokens.split([index, src_tokens.shape[1] - index],dim=-1)
            #src_tokens = torch.cat([p1, insert, p2],dim=-1)
            #for it in batch:
            #
            #    src_tokens[it] =
            #print(self.task.target_dictionary.string(src_tokens))
            #src_tokens = torch.where((src_tokens == 4) | (src_tokens == 5), src_tokens, 0)
            #print(src_tokens)
            encoder_out = model.textual_encoder(src_tokens, src_lengths)
            """task_prob = model.task_net(encoder_out["encoder_out"][0])
            encoder_padding_mask =  encoder_out["encoder_padding_mask"][0]
            task_id = torch.ones(encoder_padding_mask.shape[0],1)
            if torch.cuda.is_available():
                task_id = task_id.cuda()
            #print(task_id.shape,encoder_padding_mask.shape)
            task_loss = F.binary_cross_entropy_with_logits(
                    task_prob.float(),
                    task_id,
                    #weight=(~acoustic_encoder_out["padding_mask"][0]).float().transpose(0,1),
                    reduction = "sum"
            )"""
            #print(task_prob.shape)
            #with open("mt_top","a")as fo:
            #    l = [str(i) for i in src_tokens.tolist()]
            #    fo.write(" ".join(l)+"\n")
            #    #_, embed_out = model.textual_encoder.forward_embedding(src_tokens)
            #    embed_out = encoder_out["encoder_out"][0]
            #    l = [str(i) for i in embed_out.mean(-1).tolist()]
            #    fo.write(" ".join(l)+"\n")
            net_output = model.decoder(
                prev_output_tokens=prev_output_tokens, encoder_out=encoder_out
            )

            loss, nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
            sample_size = (
                sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
            )
            logging_output = {
                "trans_loss": utils.item(loss.data) if reduce else loss.data,
                "nll_loss": utils.item(nll_loss.data) if reduce else nll_loss.data,
                "ntokens": sample["ntokens"],
                "nsentences": sample["target"].size(0),
                "sample_size": sample_size,
            }
            # logging_output["task_loss"] = utils.item(task_loss.data) if reduce else task_loss.data
            logging_output["loss"] = utils.item(loss.data) if reduce else loss.data

        elif per_task == "asr":
            src_tokens, src_lengths, prev_output_tokens = sample["net_input"].values()
            encoder_out = model.acoustic_encoder(src_tokens, src_lengths)
            w2v_ctc_loss, tmp_logging_output = self.compute_ctc_loss(model, sample, encoder_out,True)
            loss = w2v_ctc_loss
            sample_size = (
                sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
            )
            logging_output = {
                "ntokens": sample["ntokens"],
                "nsentences": sample["target"].size(0),
                "sample_size": sample_size,
            }
            #logging_output["w2v_ctc_loss"] = utils.item(w2v_ctc_loss.data) if reduce else w2v_ctc_loss.data
            logging_output.update(tmp_logging_output)
            logging_output["loss"] = utils.item(loss.data) if reduce else loss.data

        else:
            raise NotImplementedError

        return loss, sample_size, logging_output

    def compute_ctc_loss(self, model, sample, acoustic_encoder_out, wav_ctc=False, only_contrastive=False):
        transcript = sample["transcript"]

        non_padding_mask = ~acoustic_encoder_out["padding_mask"][0]
        input_lengths = non_padding_mask.long().sum(-1)

        pad_mask = (transcript["tokens"] != self.pad_idx) & (
                transcript["tokens"] != self.eos_idx
        )
        lprobs = model.get_acoustic_normalized_probs( acoustic_encoder_out, log_probs=True, ctc_contrastive=True).contiguous()  # (T, B, C) from the encoder
        if self.use_ctc_cluster:
            cluster_lprobs = model.get_acoustic_normalized_probs( acoustic_encoder_out, log_probs=True, cluster=True).contiguous()  # (T, B, C) from the encoder
            cluster_targets_flat = transcript["cluster_tokens"].masked_select(pad_mask)
        targets_flat = transcript["tokens"].masked_select(pad_mask)
        transcript_lengths = pad_mask.sum(-1)

        #print(targets_flat.shape,transcript["tokens"].masked_select(pad_mask).shape)
        #lp=lprobs.transpose(0, 1)
        #toks = lp.argmax(dim=-1)#.unique_consecutive()
        #with open("ctc_decode","a") as f:
        #    f.write("H- "+self.task.target_dictionary.string(toks.cpu())+"\n")
        #    f.write("R- "+self.task.target_dictionary.string(transcript["tokens"].cpu())+"\n")
        #    f.write("T- "+self.task.target_dictionary.string(sample["target"].cpu())+"\n")
        #print(toks)
        #mask_=toks>4

        if only_contrastive:
            loss = 0.0
        else:
            with torch.backends.cudnn.flags(enabled=False):
                loss = F.ctc_loss(
                    lprobs,
                    targets_flat,
                    input_lengths,
                    transcript_lengths,
                    blank=self.blank_idx,
                    reduction="sum",
                    zero_infinity=self.zero_infinity,
                )
            if self.use_ctc_cluster:
                with torch.backends.cudnn.flags(enabled=False):
                    cluster_loss = F.ctc_loss(
                        cluster_lprobs,
                        cluster_targets_flat,
                        input_lengths,
                        transcript_lengths,
                        blank=self.blank_idx,
                        reduction="sum",
                        zero_infinity=self.zero_infinity,
                    )
                loss += cluster_loss
                #pred_units = self.task.target_dictionary.string(targets_flat.cpu())
                #print(pred_units)

        if self.contrastive_alpha > 0 and ( model.acoustic_encoder.num_updates < self.ban_cl_step or self.ban_cl_step == -1):

            if self.decrease_step != 0:
                contrastive_beta = max(0, self.contrastive_beta - int(model.acoustic_encoder.num_updates / self.decrease_step) * 0.1)
            else:
                contrastive_beta = self.contrastive_beta

            if contrastive_beta > 1e-4:
                reverse_sample = self.swap_sample(sample)
                if self.use_token_contrastive:
                    with torch.no_grad():
                        #embed_out = model.textual_encoder.embed_tokens(reverse_sample["net_input"]["src_tokens"])
                        _, embed_out = model.textual_encoder.forward_embedding(reverse_sample["net_input"]["src_tokens"])
                    tokens_padding_mask = reverse_sample["net_input"]["src_tokens"].eq(self.padding_idx)
                    #tokens = reverse_sample["net_input"]["src_tokens"]
                    #ctc_tokens = lprobs.transpose(0, 1).argmax(dim=-1)#.unique_consecutive()
                    #kd_sample = self.gen_kd_sample(sample, lprobs)
                    #print(tokens.shape, kd_sample["net_input"]["src_tokens"].shape)
                    contrastive_loss, similarity = self.get_token_contrastive_loss(
                        acoustic_encoder_out["wav2vec_feature"][0].transpose(0, 1),
                        acoustic_encoder_out["encoder_padding_mask"][0],
                        embed_out,
                        tokens_padding_mask
                    )
                    if self.use_two_contrastive:
                        with torch.no_grad():
                            mt_out = model.textual_encoder(reverse_sample["net_input"]["src_tokens"], reverse_sample["net_input"]["src_lengths"])
                        sentence_contrastive_loss, similarity = self.get_contrastive_loss(
                            acoustic_encoder_out,
                            mt_out,
                        )
                        contrastive_loss = sentence_contrastive_loss + contrastive_loss

                else:
                    with torch.no_grad():
                        mt_out = model.textual_encoder(reverse_sample["net_input"]["src_tokens"], reverse_sample["net_input"]["src_lengths"])
                    contrastive_loss, similarity = self.get_contrastive_loss(
                        acoustic_encoder_out,
                        mt_out,
                    )
            else:
                contrastive_loss = 0.0
                similarity = 0.0
            if contrastive_beta < 1:
                #print(sample["net_input"]["src_tokens"].shape)
                kd_sample = self.gen_kd_sample(sample, lprobs, input_lengths)
                #print(kd_sample["net_input"]["src_tokens"].shape)
                if self.use_token_contrastive:
                    with torch.no_grad():
                        _, embed_out = model.textual_encoder.forward_embedding(kd_sample["net_input"]["src_tokens"])
                    #tokens=kd_sample["net_input"]["src_tokens"]
                    tokens_padding_mask = kd_sample["net_input"]["src_tokens"].eq(self.padding_idx)
                    contrastive_kd_loss, similarity = self.get_token_contrastive_loss(
                        acoustic_encoder_out["wav2vec_feature"][0].transpose(0, 1),
                        acoustic_encoder_out["encoder_padding_mask"][0],
                        embed_out,
                        tokens_padding_mask)
                    if self.use_two_contrastive:
                        with torch.no_grad():
                            kd_mt_out = model.textual_encoder(kd_sample["net_input"]["src_tokens"], kd_sample["net_input"]["src_lengths"])
                        sentence_contrastive_kd_loss, kd_similarity = self.get_contrastive_loss(
                            acoustic_encoder_out,
                            kd_mt_out,
                            self.word_align
                        )
                        contrastive_kd_loss = sentence_contrastive_kd_loss + contrastive_kd_loss
                else:
                    with torch.no_grad():
                        kd_mt_out = model.textual_encoder(kd_sample["net_input"]["src_tokens"], kd_sample["net_input"]["src_lengths"])
                    contrastive_kd_loss, kd_similarity = self.get_contrastive_loss(
                        acoustic_encoder_out,
                        kd_mt_out,
                        self.word_align
                    )
            else:
                contrastive_kd_loss = 0.0
                kd_similarity = 0.0

            nsentences = sample["transcript"]["tokens"].size(0)
            ntokens = sample["transcript"]["ntokens"]

            if self.is_shrink == "" or self.word_align:
                all_loss = loss + (contrastive_beta * contrastive_loss + (1 - contrastive_beta) * contrastive_kd_loss) * self.contrastive_alpha / ntokens
            else:
                all_loss = loss + (contrastive_beta * contrastive_loss + (1 - contrastive_beta) * contrastive_kd_loss) * self.contrastive_alpha * ntokens / nsentences
        else:
            all_loss=loss

        if not only_contrastive:
            if wav_ctc:
                logging_output = {
                    "w2v_ctc_loss": utils.item(loss.data),  # * sample['ntokens'],
                }
            else:
                logging_output = {
                    "ctc_loss": utils.item(loss.data),  # * sample['ntokens'],
                }
        else:
            if wav_ctc:
                logging_output = {
                    "w2v_ctc_loss": utils.item(0.0),  # * sample['ntokens'],
                }
            else:
                logging_output = {
                    "ctc_loss": utils.item(0.0),  # * sample['ntokens'],
                }

        if self.contrastive_alpha > 0 and ( model.acoustic_encoder.num_updates < self.ban_cl_step or self.ban_cl_step == -1):
            logging_output["contrastive_loss"] = (contrastive_beta * contrastive_loss + (1 - contrastive_beta) * contrastive_kd_loss).data
        else:
            logging_output["contrastive_loss"] = 0

        if not model.training:
            import editdistance

            with torch.no_grad():
                lprobs_t = lprobs.transpose(0, 1).float().contiguous().cpu()

                c_err = 0
                c_len = 0
                w_errs = 0
                w_len = 0
                wv_errs = 0
                for lp, t, ct, inp_l in zip(
                    lprobs_t,
                    transcript["tokens"],
                    transcript["cluster_tokens"]
                    if self.use_ctc_cluster
                    else transcript["tokens"],
                    input_lengths,
                ):
                    lp = lp[:inp_l].unsqueeze(0)

                    decoded = None
                    p = (t != self.task.target_dictionary.pad()) & (
                        t != self.task.target_dictionary.eos()
                    )
                    #if self.use_ctc_cluster:
                    #    targ = ct[p]
                    #else:
                    targ = t[p]
                    targ_units = self.task.target_dictionary.string(targ)
                    targ_words = post_process(targ_units, self.post_process).split()
                    targ_units_arr = targ.tolist()

                    toks = lp.argmax(dim=-1).unique_consecutive()
                    pred_units_arr = toks[toks != self.blank_idx].tolist()

                    c_err += editdistance.eval(pred_units_arr, targ_units_arr)
                    c_len += len(targ_units_arr)

                    pred_units = self.task.target_dictionary.string(pred_units_arr)
                    pred_words_raw = post_process(pred_units, self.post_process).split()

                    #print("ref",targ_units,)
                    #print("hyp",pred_units,"|||",lp.shape[1],"\n")
                    #tmp_toks = lp.argmax(dim=-1)
                    #tmp_words = wlp.argmax(dim=-1)
                    #print(len(targ_words))
                    #print(tmp_toks.tolist())
                    #print(tmp_words.tolist())
                    #f_toks=tmp_toks[tmp_toks != self.blank_idx].tolist()
                    #print(targ.tolist(), tmp_toks.tolist(), len(targ.tolist()), len(f_toks))

                    if decoded is not None and "words" in decoded:
                        pred_words = decoded["words"]
                        w_errs += editdistance.eval(pred_words, targ_words)
                        wv_errs += editdistance.eval(pred_words_raw, targ_words)
                    else:
                        dist = editdistance.eval(pred_words_raw, targ_words)
                        w_errs += dist
                        wv_errs += dist

                    w_len += len(targ_words)

                logging_output["wv_errors"] = wv_errs
                logging_output["w_errors"] = w_errs
                logging_output["w_total"] = w_len
                logging_output["c_errors"] = c_err
                logging_output["c_total"] = c_len
        return all_loss, logging_output
        #else:
        #    return all_loss

    def swap_sample(self, sample):
        transcript = sample["transcript"]
        transcript_tokens = transcript["tokens"]
        transcript_lengths = transcript["lengths"] - 1
        ntokens= transcript_lengths.sum()
        target_data=sample["net_input"]["src_tokens"].contiguous()
        sample_id=sample["id"]
        #prev_output_tokens = sample["net_input"]["prev_output_tokens"]
        #src_tokens = torch.cat((prev_output_tokens[:, :1], sample["net_input"]['src_tokens']), dim=-1)
        return {
            "net_input": {
                "src_tokens": transcript_tokens.contiguous(),
                "src_lengths": transcript_lengths.contiguous(),
                "ntokens": ntokens
            },
            "target": target_data,
            "id": sample_id,
        }

    def gen_kd_sample(self, sample, lprobs, encoder_length):
        assert self.blank_idx==0, "the blank idx should be 0!"
        pred = lprobs.transpose(0,1).argmax(dim=-1)
        #print(pred.cpu().tolist())
        labels = pred.chunk(lprobs.shape[1])
        target_tokens=[]
        lengths=[]
        ntokens=0
        if self.is_shrink != "":
            for label, length in zip(labels, encoder_length):
                if "uniq" in self.is_shrink:
                    label = label[:,:length].unique_consecutive()
                if "blank" in self.is_shrink:
                    label = label[label!=0]
                lengths.append(label.shape[0])
                ntokens += label.shape[0]
                target_tokens.append(label)
            target_lengths = torch.Tensor(lengths)
            if torch.cuda.is_available():
                target_lengths = target_lengths.cuda()
            target = collate_tokens(target_tokens, self.task.target_dictionary.pad())
        else:
            ntokens = pred.numel()
            target_lengths = torch.Tensor(pred.shape[1]).repeat(pred.shape[0])
            if torch.cuda.is_available():
                target_lengths = target_lengths.cuda()
            target = pred
        target_data=sample["net_input"]["src_tokens"].contiguous()
        sample_id=sample["id"]
        #prev_output_tokens = sample["net_input"]["prev_output_tokens"]
        #src_tokens = torch.cat((prev_output_tokens[:, :1], sample["net_input"]['src_tokens']), dim=-1)
        return {
            "net_input": {
                "src_tokens": target.contiguous(),
                "src_lengths": target_lengths.contiguous(),
                "ntokens": ntokens
            },
            "target": target_data,
            "id": sample_id,
        }

    def get_token_contrastive_loss(self, tokens_1, mask_1, tokens_2, mask_2):
        def _avg_pooling(tokens, padding_mask):
            mask=(~padding_mask).int()
            sentence_rep = (tokens * mask.unsqueeze(-1)).sum(dim=1) / mask.float().sum(dim=1).unsqueeze(-1)  # [batch, hidden_size]
            return sentence_rep
        anchor_feature = _avg_pooling(tokens_1, mask_1)
        contrast_feature = _avg_pooling(tokens_2, mask_2)
        batch_size = anchor_feature.shape[0]
        feature_dim = anchor_feature.shape[1]
        #print(anchor_feature.shape)
        #print(contrast_feature.shape)
        anchor_dot_contrast = self.similarity_function(anchor_feature.expand((batch_size, batch_size, feature_dim)),
                                                      torch.transpose(contrast_feature.expand((batch_size, batch_size, feature_dim)), 0, 1))

        loss = -nn.LogSoftmax(0)(torch.div(anchor_dot_contrast, self.contrastive_temperature)).diag().sum()
        return loss, 0

    def get_contrastive_loss(self, encoder_out1, encoder_out2, word_align=False):
        def _sentence_embedding(encoder_out, padding_mask):
            mask=(~padding_mask).int()
            encoder_output = encoder_out.transpose(0, 1)

            #if "src_tokens" in sample["net_input"]:
            #    src_tokens = sample["net_input"]["src_tokens"]
            #    mask = (src_tokens != self.padding_idx)
            encoder_embedding = (encoder_output * mask.unsqueeze(-1)).sum(dim=1) / mask.float().sum(dim=1).unsqueeze(-1)  # [batch, hidden_size]
            return encoder_embedding

        if self.is_shrink != "" and not word_align:
            encoder_embedding1 = _sentence_embedding(encoder_out1["encoder_out"][0], encoder_out1["encoder_padding_mask"][0])  # [batch, hidden_size]
            encoder_embedding2 = _sentence_embedding(encoder_out2["encoder_out"][0], encoder_out2["encoder_padding_mask"][0])  # [batch, hidden_size]
            batch_size = encoder_embedding2.shape[0]
            feature_dim = encoder_embedding2.shape[1]
            anchor_feature = encoder_embedding1
            contrast_feature = encoder_embedding2
            if self.get_similarity:
                similarity = self.similarity_function(encoder_out1["wav2vec_out"].mean(1),encoder_embedding2).mean(-1)
                #print(encoder_out1["wav2vec_out"].mean(1).shape)
            #else:
            #    similarity = self.similarity_function(encoder_embedding1,encoder_embedding2).mean(-1)
            anchor_dot_contrast = self.similarity_function(anchor_feature.expand((batch_size, batch_size, feature_dim)),
                                                      torch.transpose(contrast_feature.expand((batch_size, batch_size, feature_dim)), 0, 1))

            loss = -nn.LogSoftmax(0)(torch.div(anchor_dot_contrast, self.contrastive_temperature)).diag().sum()
        else:
            encoder_embedding1 = encoder_out1["encoder_out"][0]
            encoder_embedding2 = encoder_out2["encoder_out"][0]
            padding_mask = encoder_out1["encoder_padding_mask"][0]
            assert encoder_embedding1.shape == encoder_embedding2.shape
            mask=(~padding_mask).int().unsqueeze(-1).transpose(0, 1)
            encoder_embedding1 = encoder_embedding1 * mask
            encoder_embedding2 = encoder_embedding2 * mask
            batch_size = encoder_embedding2.shape[1]
            length = encoder_embedding2.shape[0]
            feature_dim = encoder_embedding2.shape[2]
            anchor_dot_contrast = self.similarity_function(encoder_embedding1.expand((length, length, batch_size, feature_dim)).transpose(0,2),
                                                           encoder_embedding2.expand((length, length, batch_size, feature_dim)).transpose(0,2))
            loss = -nn.LogSoftmax(1)(torch.div(anchor_dot_contrast, self.contrastive_temperature)).diagonal().sum()

        return loss, 0

    @staticmethod
    def reduce_metrics(logging_outputs) -> None:
        """Aggregate logging outputs from data parallel training."""
        loss_sum = utils.item(sum(log.get("loss", 0) for log in logging_outputs))
        trans_loss_sum = utils.item(
            sum(log.get("trans_loss", 0) for log in logging_outputs)
        )
        nll_loss_sum = utils.item(
            sum(log.get("nll_loss", 0) for log in logging_outputs)
        )
        w2v_ctc_loss_sum = utils.item(
            sum(log.get("w2v_ctc_loss", 0) for log in logging_outputs)
        )
        ntokens = utils.item(sum(log.get("ntokens", 0) for log in logging_outputs))
        sample_size = utils.item(
            sum(log.get("sample_size", 0) for log in logging_outputs)
        )
        metrics.log_scalar(
            "loss", loss_sum / sample_size / math.log(2), sample_size, round=3
        )
        metrics.log_scalar(
            "trans_loss", trans_loss_sum / ntokens / math.log(2), ntokens, round=3
        )
        metrics.log_scalar(
            "nll_loss", nll_loss_sum / ntokens / math.log(2), ntokens, round=3
        )
        metrics.log_scalar(
            "w2v_ctc_loss",
            w2v_ctc_loss_sum / sample_size / math.log(2),
            sample_size,
            round=3,
        )
        metrics.log_derived(
            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
        )

        nsentences = utils.item(
            sum(log.get("nsentences", 0) for log in logging_outputs)
        )
        task_loss = utils.item(
            sum(log.get("task_loss", 0) for log in logging_outputs)
        )
        metrics.log_scalar(
            "task_loss",
            task_loss / nsentences / math.log(2),
            nsentences,
            round=3,
        )
        contrastive_loss = utils.item(
            sum(log.get("contrastive_loss", 0) for log in logging_outputs)
        )
        metrics.log_scalar(
            "contrastive_loss",
            contrastive_loss / nsentences / math.log(2),
            nsentences,
            round=3,
        )

        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
        if total > 0:
            metrics.log_scalar("total", total)
            n_correct = utils.item(
                sum(log.get("n_correct", 0) for log in logging_outputs)
            )
            metrics.log_scalar("n_correct", n_correct)
            metrics.log_derived(
                "accuracy",
                lambda meters: round(
                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
                )
                if meters["total"].sum > 0
                else float("nan"),
            )

        c_errors = sum(log.get("c_errors", 0) for log in logging_outputs)
        metrics.log_scalar("_c_errors", c_errors)
        c_total = sum(log.get("c_total", 0) for log in logging_outputs)
        metrics.log_scalar("_c_total", c_total)
        w_errors = sum(log.get("w_errors", 0) for log in logging_outputs)
        metrics.log_scalar("_w_errors", w_errors)
        wv_errors = sum(log.get("wv_errors", 0) for log in logging_outputs)
        metrics.log_scalar("_wv_errors", wv_errors)
        w_total = sum(log.get("w_total", 0) for log in logging_outputs)
        metrics.log_scalar("_w_total", w_total)

        if c_total > 0:
            metrics.log_derived(
                "uer",
                lambda meters: safe_round(
                    meters["_c_errors"].sum * 100.0 / meters["_c_total"].sum, 3
                )
                if meters["_c_total"].sum > 0
                else float("nan"),
            )
        if w_total > 0:
            metrics.log_derived(
                "wer",
                lambda meters: safe_round(
                    meters["_w_errors"].sum * 100.0 / meters["_w_total"].sum, 3
                )
                if meters["_w_total"].sum > 0
                else float("nan"),
            )
            metrics.log_derived(
                "raw_wer",
                lambda meters: safe_round(
                    meters["_wv_errors"].sum * 100.0 / meters["_w_total"].sum, 3
                )
                if meters["_w_total"].sum > 0
                else float("nan"),
            )

    @staticmethod
    def logging_outputs_can_be_summed() -> bool:
        """
        Whether the logging outputs returned by `forward` can be summed
        across workers prior to calling `reduce_metrics`. Setting this
        to True will improves distributed training speed.
        """
        return True


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

'''
Script to search data for relevant tweets:
(1) build lists of keywords/hastags for each topic
(2) search for tweets:
	tweets that have at least two keywords for a single topic are relevant to that topic
	tweets that have a hashtag for a topic are relevent to that topic
(3) print relevant lists to output file

input: directory of tweet files to search
output: relevant lists for each topic

'''

import sys
import json
import os

def search_tweets(input_dir, keywords, hashtags, output_dir):

	with open(output_dir + "masks.json", 'w') as masks_out, open(output_dir + "social distancing.json", 'w') as sd_out, open(output_dir + "quarantine.json", 'w') as quaran_out:

		for in_file in os.listdir(input_dir):
			if not in_file.startswith('.'):
				print('searching: ' + in_file)
				with open(input_dir + in_file, 'r') as f:
					for line in f:
						line = line.strip()
						if line:
							tweet = json.loads(line)
							mask_words = set(tweet['tokens']).intersection(set(keywords['masks']))
							sd_words = set(tweet['tokens']).intersection(set(keywords['social distancing']))
							quarantine_words = set(tweet['tokens']).intersection(set(keywords['quarantine']))

							mask_hashtags = set(tweet['tokens']).intersection(set(hashtags['masks']))
							sd_hashtags = set(tweet['tokens']).intersection(set(hashtags['social distancing']))
							q_hashtags = set(tweet['tokens']).intersection(set(hashtags['quarantine']))

							if mask_hashtags or len(mask_words) >= 2:
								json.dump(tweet, masks_out)
								masks_out.write('\n')
							if sd_hashtags or len(sd_words) >= 2:
								json.dump(tweet, sd_out)
								sd_out.write('\n')
							if q_hashtags or quarantine_words:
								json.dump(tweet, quaran_out)
								quaran_out.write('\n')

				print('finished: ' + in_file)
				print()

if __name__ == "__main__":

	input_dir = sys.argv[1]
	output_dir = sys.argv[2]

	if not os.path.exists(output_dir):
		os.mkdir(output_dir)

	masks = ['wear', 'wear', 'wearing', 'mask', 'masks', 'cover', 'face', 'facial', 'N95', 'respirator']
	masks_tags = ['mask', 'wearamask', 'facemask']
	soc_dist = ['distancing', 'social', 'distance']
	sd_tags = ['socialdistancing']
	quarantine = ['quarantine', 'quarantined', 'quarantining', 'self-isolate', 'self-isolating', 'isolating', 'isolate']#, 'work', 'home', 'stay']
	q_tags = ['quarantine', 'workfromhome', 'stayhomesavelives', 'selfisolation', 'stayathome', 'lockdown']

	keywords = {'masks' : masks, 'social distancing' : soc_dist, 'quarantine' : quarantine}
	hashtags = {'masks' : masks_tags, 'social distancing' : sd_tags, 'quarantine' : q_tags}

	search_tweets(input_dir, keywords, hashtags, output_dir)


# 内容度量值：衡量图片之间的内容差异程度
#
# 风格度量值：衡量图片之间的风格差异程度

# 建立神经网络模型，对内容图片中的内容和风格图片的风格进行提取
#
# 以内容图片为基准将其输入建立的模型中，并不断调整内容度量和风格度量值
#
# 让他们趋近于最小！！！
import matplotlib_inline
import torch
import torchvision
from torchvision import transforms, datasets, models
from PIL import Image
import matplotlib.pyplot as plt
from torch.autograd import Variable
import copy

matplotlib_inline

#图片变换
transform = transforms.Compose([transforms.Resize([224,224]),
                                        transforms.ToTensor()])

#数据加载
def loading(path = None):
    img = Image.open(path)
    img = transform(img)
    img = img.unsqueeze(0)
    return img

#将数据转移到cuda
content_img = loading("3333.jpg")
content_img = Variable(content_img).cuda()
style_img = loading("2222.jpg")
style_img = Variable(style_img).cuda()

#内容图片损失
class Content_loss(torch.nn.Module):
    def __init__(self, weight, target):
        #weight是权重参数，用来控制内容和风格对最后合成图像的影响程度
        #target是通过卷积获取到的输入图像中的内容
        super(Content_loss).__init__()
        self.weight = weight
        self.target = target.detach()*weight
        #用来对提取到的内容进行锁定，不需要进行梯度下降？
        self.loss_fn = torch.nn.MSELoss()

    def forward(self, input):
        #用来计算输入图像和内容图像之间的损失值
        self.loss = self.loss_fn(input*self.weight, self.target)
        return input

    def backward(self):
        #根据计算得到的损失值进行反向传播，并返回损失值。
        self.loss.backward(retain_graph = True)
        return self.loss

# 格拉姆矩阵
# CNN提取的风格是由数字组成的
# 数字大小代表了图片中的风格的突出程度
# Gram矩阵实矩阵的内积运算
# 在运算过后输入到该矩阵的特征图的大的数字会变得更大（放大效果）
# 放大的风格再参与损失计算，对合成的图片产生更大的影响

class Gram_matrix(torch.nn.Module):
    # 实例参与风格损失的计算。
    def forward(self, input):
        a, b, c, d = input.size()
        feature = input.view(a * b, c * d)
        gram = torch.mm(feature, feature.t())
        return gram.div(a * b * c * d)

class Style_loss(torch.nn.Module):
    def __init__(self, weight, target):
        super(Style_loss, self).__init__()
        self.weight = weight
        self.target = target.detach() * weight
        self.loss_fn = torch.nn.MSELoss()

        self.gram = Gram_matrix()

    def forward(self, input):
        self.Gram = self.gram(input.clone())
        self.Gram.mul_(self.weight)

        self.loss = self.loss_fn(self.Gram, self.target)
        return input

    def backward(self):
        self.loss.backward(retain_graph=True)
        return self.loss

Use_gpu = torch.cuda.is_available()
cnn = models.vgg16(pretrained = True).features

if Use_gpu:
    cnn = cnn.cuda()

model = copy.deepcopy(cnn)

# 指定在第三层卷积层提取内容， 在一、二、三、四层卷积层提取风格
# why ？
content_layer = ["Conv3"]
style_layer = ["Conv_1", "Conv_2", "Conv_3", "Conv_4"]

content_losses = []
style_losses = []

content_weight = 10000000
style_weight = 10

new_model = torch.nn.Sequential()
model = copy.deepcopy(cnn)
gram = Gram_matrix()

if Use_gpu:
    print("Using GPU for training!!!!! ")
    print("显卡数量："+torch.cuda.device_count())
    print("显卡型号："+torch.cuda.get_device_name(0))
    new_model = new_model.cuda()
    gram = gram.cuda()

index = 1

for layer in list(model)[:8]:
    # 指明我们仅仅用到迁移模型特征提取部分卷积层前八层，因为内容提取和风格提取在前八层已经完成了

    if isinstance(layer, torch.nn.Conv2d):
        # 然后建立一个空的模型
        # 用torch.nn.Module 和 add_module 方法向空的模型中加入指定的层次模块

        name = "Conv_"+str(index)
        new_model.add_module(name,layer)

        if name in content_layer:
            target = new_model(content_img).clone()
            content_loss = Content_loss(content_weight, target)
            new_model.add_module("content_loss_"+str(index), content_loss)
            content_losses.append(content_loss)

        if name in style_layer:
            target = new_model(style_img).clone()
            target = gram(target)
            style_loss = Style_loss(style_weight, target)
            new_model.add_module("style_loss_" + str(index), style_loss)
            style_losses.append(style_loss)

    if isinstance(layer, torch.nn.ReLU):
        name = "Relu_"+str(index)
        new_model.add_module(name,layer)
        index = index+1

    if isinstance(layer, torch.nn.MaxPool2d):
        name = "MaxPool_"+str(index)
        new_model.add_module(name,layer)

# 优化参数部分代码
# 优化函数使用的是LBFGS
# 因为这个模型需要优化的损失值有多个并且规模比较大

input_img = content_img.clone()
parameter = torch.nn.Parameter(input_img.data)
optimizer = torch.optim.LBFGS([parameter])

epoch_n = 500
epoch = [0]

while epoch[0] <= epoch_n:
    def closure():
        optimizer.zero_grad()
        style_score = 0
        content_score = 0
        parameter.data.clamp_(0,1)
        new_model(parameter)
        for sl in style_losses:
            # style_score += sl.backward()
            style_score += sl.loss

        for cl in content_losses:
            # content_score += cl.backward()
            content_score += cl.loss

        style_score *= style_weight
        content_score *= content_weight
        loss = style_score + content_score
        loss.backward()

        epoch[0] += 1
        if epoch[0] % 50 == 0:
            print('Epoch:{} Style_loss : {:4f} Content_loss : {:4f}'.format(epoch[0], style_score, content_score))
        return style_score + content_score
    optimizer.step(closure)

#图片的输出
output = parameter.data
unloader = transforms.ToPILImage()  # reconvert into PIL image

plt.ion()
plt.figure()
def imshow(tensor, title=None):
    image = tensor.clone().cpu()  # we clone the tensor to not do changes on it
    image = image.view(3, 224, 224)  # remove the fake batch dimension
    image = unloader(image)
    plt.imshow(image)
    if title is not None:
        plt.title(title)
    plt.pause(0.001) # pause a bit so that plots are updated
imshow(output, title='Output Image')
plt.ioff()
plt.show()


from PyQt6.QtWidgets import (
    QComboBox,
    QPushButton,
    QGridLayout,
    QVBoxLayout,
    QApplication,
    QWidget,
    QMainWindow,
)
from PyQt6.QtCore import Qt
from k_means import *
from linear import *
from pca import *
from clear_layout import *
from start_spark import *
from calendar_filter import *
import sys
import os

class mainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Weather app")
        self.resize(1400, 700)

        self.layout = QGridLayout()
        self.layout.setAlignment(Qt.AlignmentFlag.AlignRight)

        self.layout.setColumnStretch(0, 3)
        self.layout.setColumnStretch(1, 1)
        self.layout.setRowStretch(0, 5)
        self.layout.setRowStretch(1, 4)

        self.graph_win = QGridLayout()
        self.layout.addLayout(self.graph_win, 0, 0)

        self.options_widget = QWidget()
        self.options_layout = QVBoxLayout()
        self.options_widget.setLayout(self.options_layout)
        self.options_widget.setMaximumWidth(350)  # TODO make this a var

        self.year_win = self.select_year()
        self.options_layout.addWidget(self.year_win)

        self.filter_win = QVBoxLayout()
        self.options_layout.addLayout(self.filter_win)

        self.layout.addWidget(self.options_widget, 0, 1)

        self.var_win = QHBoxLayout()
        self.layout.addLayout(self.var_win, 1, 0)

        self.ml_win = QGridLayout()
        self.layout.addLayout(self.ml_win, 1, 1)

        # Sets all of the buttons for the ML algorithms
        linear_button = QPushButton("Linear Regression")
        linear_button.clicked.connect(self.linear_win)

        kmeans_button = QPushButton("K-Means")
        kmeans_button.clicked.connect(self.kmeans_window)

        gaussian_button = QPushButton("Gaussian Mixture")
        gaussian_button.clicked.connect(self.gaussian_window)

        pca_button = QPushButton("Principal Component Analysis")
        pca_button.clicked.connect(self.pca_window)

        self.graph_active = None
        self.filter_active = None
        self.active_fil_layout = None
        self.var_active = None
        self.active_var_layout = None
        self.cal_win_active = None

        # Adds the buttons to the ML window
        self.ml_win.addWidget(linear_button, 0, 0)
        self.ml_win.addWidget(kmeans_button, 0, 1)
        self.ml_win.addWidget(gaussian_button, 1, 0)
        self.ml_win.addWidget(pca_button, 1, 1)

        container = QWidget()
        container.setLayout(self.layout)
        self.setCentralWidget(container)

        self.select_year()

        self.csv = "2006.csv"

    # A widget that selects what year the user wants to see.
    def select_year(self):
        years = [
            "2006",
            "2007",
            "2008",
            "2009",
            "2010",
            "2011",
            "2012",
            "2013",
            "2014",
            "2015",
            "2016",
            "2017",
            "2018",
            "2019",
            "2020",
            "2021",
            "2022",
        ]

        file_path = "Data/processed"
        self.dir_list = os.listdir(file_path)

        self.year_widget = QWidget()
        self.year_layout = QVBoxLayout()
        self.year_widget.setLayout(self.year_layout)

        self.year_label = QLabel("Select Year:")
        self.year_layout.addWidget(self.year_label)

        self.year_select = QComboBox()
        self.year_select.addItems(self.dir_list)
        # self.year_select.activated.connect(self.set_csv)
        self.year_select.currentTextChanged.connect(self.set_csv)
        self.year_layout.addWidget(self.year_select)
        self.filter_active = True

        self.min_year_label = QLabel("Starting Year:")
        self.year_layout.addWidget(self.min_year_label)
        self.min_year_select = QComboBox()
        self.min_year_select.addItems(years)
        self.min_year_select.activated.connect(self.set_min_year)
        self.year_layout.addWidget(self.min_year_select)

        self.max_year_label = QLabel("Ending Year:")
        self.year_layout.addWidget(self.max_year_label)
        self.max_year_select = QComboBox()
        self.max_year_select.addItems(years)
        self.max_year_select.activated.connect(self.set_max_year)
        self.year_layout.addWidget(self.max_year_select)

        self.cal_win = None
        self.calendar_button = QPushButton("Open Calendar")
        self.calendar_button.clicked.connect(self.open_calendar)
        self.year_layout.addWidget(self.calendar_button)

        self.year_widget.setMaximumHeight(200)
        self.year_layout.setAlignment(Qt.AlignmentFlag.AlignTop)

        return self.year_widget

    def open_calendar(self):
        if self.cal_win == None:
            self.cal_win = CalendarFilter()
            self.cal_win.show()
            self.cal_win_active = True
            self.cal_win.done.clicked.connect(self.set_csv)
        else:
            self.cal_win.close()
            self.cal_win_active = False
            self.cal_win = None

    # When a starting year is selected, the variable min_year is updated
    def set_min_year(self):
        self.min_year = self.min_year_select.currentText()

    # When an ending year is selected, the variable max_year is updated
    def set_max_year(self):
        self.max_year = self.max_year_select.currentText()

    # When an option is selected by the user in the combobox
    # it is then passed to self.csv that will be used again in
    # the graph window class.
    def set_csv(self, csv):
        if self.cal_win != None:
            self.csv = self.cal_win.start_date()
            print(self.csv)
            self.cal_win.close()
            self.cal_win == None
        if self.cal_win_active != True:
            self.csv = csv
            print(self.csv)

    # Has the user select the year and calls linear to
    # perform linear regression on daily precipitation with the chosen year
    def linear_win(self):
        if self.graph_active == True:
            clear_graph_win(self.graph_win)
        if self.filter_active == True:
            clear_fil_win(self.filter_win, self.active_fil_layout)
        if self.var_active == True:
            clear_var_win(self.var_win, self.active_var_layout)
        self.linear_enter = linear_enter()
        self.linear_cancel = linear_cancel()
        self.filter_win.addWidget(self.linear_enter)
        self.filter_win.addWidget(self.linear_cancel)

        self.linear_enter.clicked.connect(self.linear_run)
        self.linear_cancel.clicked.connect(self.canceled)

    # Takes the year and runs spark and displays the linear regression graph.
    def linear_run(self):
        if self.csv != None:
            print(self.csv)
            sdf = setup(self.csv)
            self.linear_graph = linear_reg(sdf)
            self.graph_win.addWidget(self.linear_graph, 0, 0)
            self.graph_active = True

    #  Calls k_means to display the k-means graph.
    def kmeans_window(self):
        print("Var result = ", self.var_active)
        if self.graph_active == True:
            clear_graph_win(self.graph_win)
        if self.filter_active == True:
            clear_fil_win(self.filter_win, self.active_fil_layout)
        if self.var_active == True:
            print("Should clear")
            clear_var_win(self.var_win, self.active_var_layout)

        self.kmeans_options = QVBoxLayout()
        k_means_label = QLabel("K-Means Options")
        k_means_label.setAlignment(Qt.AlignmentFlag.AlignHCenter)
        k_means_label.setFont(QFont("Times", 30))
        self.kmeans_options.addWidget(k_means_label)

        self.column_label = QLabel("Variables:")
        k_col = k_columns()
        self.column1 = None
        self.column2 = "None"
        self.column3 = "None"
        self.kmeans_var1 = QComboBox()
        self.kmeans_var2 = QComboBox()
        self.kmeans_var3 = QComboBox()
        self.kmeans_var1.addItems(k_col)
        self.kmeans_var2.addItems(k_col)
        self.kmeans_var3.addItems(k_col)
        self.kmeans_options.addWidget(self.column_label)
        self.kmeans_var1.activated.connect(self.set_kColumn1)
        self.kmeans_var2.activated.connect(self.set_kColumn2)
        self.kmeans_var3.activated.connect(self.set_kColumn3)
        self.kmeans_options.addWidget(self.kmeans_var1)
        self.kmeans_options.addWidget(self.kmeans_var2)
        self.kmeans_options.addWidget(self.kmeans_var3)

        k_num = QLabel("Number of Clusters: ")
        self.k = None
        self.k_textbox = QLineEdit()
        self.k_textbox.setValidator(QIntValidator())
        self.kmeans_options.addWidget(k_num)
        self.kmeans_options.addWidget(self.k_textbox)

        self.enter = QPushButton("Run K-Means", self)
        self.kmeans_options.addWidget(self.enter)
        self.enter.clicked.connect(self.k_check)

        self.filter_win.addLayout(self.kmeans_options)
        self.active_fil_layout = self.kmeans_options
        self.filter_active = True

    # Set the number of clusters for K-Means
    def set_clusters(self):
        self.k = int(self.k_textbox.text())

    # Sets the first column to perform K-Means on
    def set_kColumn1(self):
        self.column1 = self.kmeans_var1.currentText()

    # Sets the second column to perform K-Means on
    def set_kColumn2(self):
        self.column2 = self.kmeans_var2.currentText()
        print(self.column2)

    # Sets the third column to perform K-Means on
    def set_kColumn3(self):
        self.column3 = self.kmeans_var3.currentText()

    # Does a check to make sure the number of clusters have been selected
    # also is where the kmeans graph and silhouette graph is added to the
    # graph window.
    def k_check(self):
        print("Var result = ", self.var_active)
        if self.graph_active == True:
            clear_graph_win(self.graph_win)
        if self.var_active == True:
            print("Should clear")
            clear_var_win(self.var_win, self.active_var_layout)
        self.set_clusters()
        if self.k != None and self.csv != None:
            sdf = setup(self.csv)
            self.kgraph, self.sil_graph, self.legend = kmeans(
                sdf, self.k, self.column1, self.column2, self.column3
            )
            self.var_win.addLayout(self.legend)
            self.graph_win.addWidget(self.kgraph, 0, 0)
            self.graph_win.addWidget(self.sil_graph, 0, 1)
            self.active_var_layout = self.legend
            self.graph_active = True
            self.var_active = True

    def gaussian_window(self):
        pass

    # Calls pca and creates a visualization.
    def pca_window(self):
        if self.graph_active == True:
            clear_graph_win(self.graph_win)
        if self.filter_active == True:
            clear_fil_win(self.filter_win, self.active_fil_layout)
        if self.var_active == True:
            clear_var_win(self.var_win, self.active_var_layout)
        # Create the options box
        self.pca_options = QVBoxLayout()
        # Create the label for the number of components
        num_comp_label = QLabel("Number of Components:")
        self.pca_options.addWidget(num_comp_label)
        # Create the textbox for the number of components, and make it only take ints
        self.pca_textbox = QLineEdit()
        self.pca_textbox.setValidator(QIntValidator())
        self.pca_options.addWidget(self.pca_textbox)
        # Create the button to run PCA
        run_PCA = QPushButton("Run PCA", self)
        run_PCA.clicked.connect(self.run_pca)
        self.pca_options.addWidget(run_PCA)
        # Add the PCA options to the filter window
        self.filter_win.addLayout(self.pca_options)
        # Doesn't run PCA until runPCA button pressed
        self.filter_active = True
        self.active_fil_layout = self.pca_options

    # When the runPCA button is pressed, creates the visualization
    def run_pca(self):
        # The number of modified values to output after PCA runs
        NUM_OUTPUT_VALUES = 5  # temporary until PCA visualization gets built
        # Create the window for PCA
        self.pca_win = QVBoxLayout()
        # Set the number of components from the textbox
        self.num_comp = int(self.pca_textbox.text())

        if self.csv != None:
            sdf = setup(self.csv)

        # If the number of components
        if self.num_comp != None and self.num_comp > 0:
            # Run PCA, get the model and data out
            model, data = pca(sdf, self.num_comp, NUM_OUTPUT_VALUES)
            # Add label for text output
            pca_label = QLabel("Principal Component Analysis")
            self.pca_win.addWidget(pca_label)
            # Print out the explained variances
            variances = QLabel("Explained Variances:" + str(model.explainedVariance))
            self.pca_win.addWidget(variances)
            # Print out label for datapoints
            first_str = "First " + str(NUM_OUTPUT_VALUES) + " Data Points:"
            first_values_label = QLabel(first_str)
            self.pca_win.addWidget(first_values_label)
            # Loop through output data and print
            for out in data:
                values = QLabel(str(out.output))
                self.pca_win.addWidget(values)

            # Add to the graph window, default to bottom left
            self.graph_win.addLayout(self.pca_win, 1, 0)
            self.graph_active = True

    # If a cancel button is present in the filter window,
    # it should clear the window.
    def canceled(self):
        if self.graph_active == True:
            clear_graph_win(self.graph_win)
        if self.filter_active == True:
            clear_fil_win(self.filter_win, self.active_fil_layout)

app = QApplication(sys.argv)
window = mainWindow()
window.show()
sys.exit(app.exec())


from django.contrib.auth import get_user_model
from django.core.cache import cache
from django.test import TestCase, Client

from http import HTTPStatus

from ..models import Post, Group

User = get_user_model()

class StaticURLTests(TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.user = User.objects.create_user(username='test-user')
        cls.author = User.objects.create_user(username='author')
        cls.group = Group.objects.create(
            title='Тестовая группа',
            slug='test-slug',
            description='Тестовое описание',
        )
        cls.post = Post.objects.create(
            author=cls.author,
            text='Тестовый пост',
        )

    def setUp(self):
        cache.clear()
        self.guest_client = Client()
        self.authorized_client = Client()
        self.authorized_client.force_login(self.user)
        self.author_client = Client()
        self.author_client.force_login(self.author)

    def test_home_url_exists_at_desired_location(self):
        """Страница / доступна любому пользователю."""
        response = self.guest_client.get('/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_group_list_url_exists_at_desired_location(self):
        """Страница /group/test-slug/ доступна любому пользователю."""
        response = self.guest_client.get('/group/test-slug/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_profile_url_exists_at_desired_location(self):
        """Страница /profile/test-user/ доступна любому пользователю."""
        response = self.guest_client.get('/profile/test-user/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_post_detail_url_exists_at_desired_location(self):
        """Страница /posts/1/ доступна любому пользователю."""
        response = self.guest_client.get('/posts/1/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_post_create_url_exists_at_desired_location(self):
        """Страница /create/ доступна авторизованному пользователю."""
        response = self.authorized_client.get('/create/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_post_edit_url_exists_at_desired_location(self):
        """Страница /posts/1/edit/ доступна автору."""
        response = self.author_client.get('/posts/1/edit/')
        self.assertEqual(response.status_code, HTTPStatus.OK)

    def test_unexisting_url_exists_at_desired_location(self):
        """Страница /unexisting_page/ не доступна."""
        response = self.guest_client.get('/unexisting_page/')
        self.assertEqual(response.status_code, HTTPStatus.NOT_FOUND)

    def test_urls_uses_correct_template(self):
        """URL-адрес использует соответствующий шаблон."""
        templates_url_names = {
            '/': 'posts/index.html',
            '/group/test-slug/': 'posts/group_list.html',
            '/profile/test-user/': 'posts/profile.html',
            f'/posts/{self.post.id}/': 'posts/post_detail.html',
            '/create/': 'posts/create_post.html',
            f'/posts/{self.post.id}/edit/': 'posts/create_post.html',
        }
        for address, template in templates_url_names.items():
            with self.subTest(address=address):
                response = self.author_client.get(address)
                self.assertTemplateUsed(response, template)


import cv2
import mediapipe as mp
import time

import math

import tkinter as tk

cap = cv2.VideoCapture(0)
# Get the width of the video
videoWidth = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
videoHeight = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

mpHands = mp.solutions.hands
hands = mpHands.Hands()
mpDraw = mp.solutions.drawing_utils

# Create a tkinter window
root = tk.Tk()
root.title("Pixel by Pixel Line")

# Create a canvas widget
canvas = tk.Canvas(root, width=videoWidth, height=videoHeight)
canvas.pack()
cursor = canvas.create_oval(10, 10, 15, 15, fill="blue")

pTime = 0
cTime = 0

# Thumb and index start coordinates
xt, yt, xi, yi = 0, 0, 0, 0
indexUpdated = False
thumbUpdated = False

while True:
    success, img = cap.read()
    img = cv2.flip(img, 1)
    imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    results = hands.process(imgRGB)

    if results.multi_hand_landmarks:
        for handLms in results.multi_hand_landmarks:

            indexUpdated = False
            thumbUpdated = False

            for idx, lm in enumerate(handLms.landmark):
                # print(idx, lm)
                h, w, c = img.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                # Draw a circle for one of the points

                # 4 - thumb: xt, yt
                # 8 - index finger: xi, yi
                if idx == 4:
                    xt, yt = cx, cy
                    cv2.circle(img, (xt, yt), 15, (255, 0, 255), cv2.FILLED)
                    thumbUpdated = True
                elif idx == 8:
                    xi, yi = cx, cy
                    cv2.circle(img, (xi, yi), 15, (255, 255, 255), cv2.FILLED)
                    indexUpdated = True

                    canvas.coords(cursor, xi, yi, xi+5, yi+5)
                    root.update()
                else:
                    continue

                if indexUpdated and thumbUpdated:

                    distance = math.sqrt((xt - xi) ** 2 + (yt - yi) ** 2)

                    if distance < 30:
                        # Tkinter
                        canvas.create_rectangle(xi, yi, xi + 4, yi + 4, fill="black")
                        # root.update()
                        # print(distance)

            mpDraw.draw_landmarks(img, handLms)  # Dots
            # mpDraw.draw_landmarks(img, handLms, mpHands.HAND_CONNECTIONS)  # Connections

    cTime = time.time()
    fps = 1/(cTime - pTime)
    pTime = cTime

    cv2.putText(img, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3, (255, 0, 255), 3)

    cv2.imshow("Image", img)
    cv2.waitKey(1)


import csv
import requests
from bs4 import BeautifulSoup

url = 'https://barefootbuttons.com/product-category/version-1/'
response = requests.get(url)

html = response.text
soup = BeautifulSoup(html, 'html.parser')

products = soup.find_all('div', 'product-small box')

for product in products:
    link = product.find('a')['href']
    title = product.find('p', {'class': 'name product-title woocommerce-loop-product__title'}).text.strip()
    price = float(product.find('span', {'class': 'woocommerce-Price-amount amount'}).text.strip().replace('$', ''))
    with open('woo_products.csv', 'a', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow([title, link, price])


import torch
import numpy as np
import torch.nn as nn
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from scipy.special import kl_div
import os
from util import mix_dataset

np.random.seed(35813)
torch.manual_seed(35813)

### Direct to working environment
WORK_PATH = "/home/jia/Desktop/MSc_Project/imagetemplate/FedLIT"
MODEL_OUTPUT_PATH = "/home/jia/Desktop/MSc_Project/imagetemplate/FedLIT/output"

# def plot_latent_distribution(local_ground_truth, mixed_ground_truth, predicted, save_path):

#     # Flatten the datasets to 1D arrays
#     flattened_dataset1 = np.array(local_ground_truth).reshape(-1,1)
#     flattened_dataset2 = np.array(mixed_ground_truth).reshape(-1,1)
#     flattened_dataset3 = np.array([img.flatten() for img in predicted]).reshape(-1,1)
#     # Calculate the KDE for each dataset
#     kde_dataset1 = gaussian_kde(flattened_dataset1)
#     kde_dataset2 = gaussian_kde(flattened_dataset2)
#     kde_dataset3 = gaussian_kde(flattened_dataset3)

#     # Define the x-axis range for the plots
#     x_range = np.linspace(min(flattened_dataset1.min(), flattened_dataset2.min(), flattened_dataset3.min()),
#                           max(flattened_dataset1.max(), flattened_dataset2.max(), flattened_dataset3.max()), 50)

#     # Plot the KDE curves as line graphs
#     plt.plot(x_range, kde_dataset1(x_range), label='local test set')
#     plt.plot(x_range, kde_dataset2(x_range), label='mixed test set')
#     plt.plot(x_range, kde_dataset3(x_range), label='augmented')

#     # plt.hist(flattened_dataset1, bins=50, alpha=0.7, label='ground_truth')
#     # plt.hist(flattened_dataset2, bins=50, alpha=0.7, label='predicted')

#     plt.xlabel('Value')
#     plt.ylabel('Density')
#     plt.title('Distributions of Latent Embeddings')
#     plt.legend()
#     plt.savefig(save_path)
#     plt.show()

# def linearize_image_data(dataset):
#     new_dataset = []
#     for image, label in dataset:
#         linear_image = image.flatten()
#         new_dataset.append(linear_image)
#     return new_dataset

def plot_latent_distribution(local_ground_truth, mixed_ground_truth, predicted, save_path):

    # Flatten the datasets to 1D arrays
    flattened_dataset1 = np.concatenate(local_ground_truth).flatten()
    flattened_dataset2 = np.concatenate(mixed_ground_truth).flatten()
    flattened_dataset3 = np.concatenate(predicted).flatten()

    # Calculate the KDE for each dataset
    kde_dataset1 = gaussian_kde(flattened_dataset1)
    kde_dataset2 = gaussian_kde(flattened_dataset2)
    kde_dataset3 = gaussian_kde(flattened_dataset3)

    # Define the x-axis range for the plots
    x_range = np.linspace(min(flattened_dataset1.min(), flattened_dataset2.min(), flattened_dataset3.min()),
                          max(flattened_dataset1.max(), flattened_dataset2.max(), flattened_dataset3.max()), 50)

    # Plot the KDE curves as line graphs
    plt.plot(x_range, kde_dataset1(x_range), label='local test set')
    plt.plot(x_range, kde_dataset2(x_range), label='mixed test set')
    plt.plot(x_range, kde_dataset3(x_range), label='augmented')

    # plt.hist(flattened_dataset1, bins=50, alpha=0.7, label='ground_truth')
    # plt.hist(flattened_dataset2, bins=50, alpha=0.7, label='predicted')

    plt.xlabel('Value')
    plt.ylabel('Density')
    plt.title('Distributions of Latent Embeddings')
    plt.legend()
    plt.savefig(save_path)
    plt.show()

def mix_dataset(datset_list):
    mixed_dataset = []

    for datset in datset_list:
        mixed_dataset.extend(datset)

    np.random.shuffle(mixed_dataset)

    return mixed_dataset

# if __name__ == '__main__':

#     for number in range(5):
#         output_path = MODEL_WEIGHT_BACKUP_PATH + "/" + 'number' + str(number)
#         data_path = WORK_PATH + "/" + 'classifier_data/'
#         gt_imgs1 = np.load(data_path + 'client1/mixed_test_data.npy', allow_pickle=True)
#         gt_imgs1 = linearize_image_data(gt_imgs1)
#         gt_imgs2 = np.load(data_path + 'client2/mixed_test_data.npy', allow_pickle=True)
#         gt_imgs2 = linearize_image_data(gt_imgs2)
#         gt_imgs3 = np.load(data_path + 'client3/mixed_test_data.npy', allow_pickle=True)
#         gt_imgs3 = linearize_image_data(gt_imgs3)
#         gt_imgs4 = np.load(data_path + 'client4/mixed_test_data.npy', allow_pickle=True)
#         gt_imgs4 = linearize_image_data(gt_imgs4)
#         gt_imgs5 = np.load(data_path + 'client5/mixed_test_data.npy', allow_pickle=True)
#         gt_imgs5 = linearize_image_data(gt_imgs5)

#         mixed_gt_imgs = mix_dataset([gt_imgs1, gt_imgs2, gt_imgs3, gt_imgs4, gt_imgs5])

#         for type in ['ablated_nofed', 'ablated_fed']:
#             save_path = output_path + "/" + type
#             dist_path = save_path + "/" + "dist_imgs" + "/"
#             if not os.path.exists(dist_path):
#                 os.makedirs(dist_path)

#             aug_img_path = save_path + '/' + 'aug_imgs/'
#             aug_imgs1 = np.load(aug_img_path + 'client1_aug_imgs.npy')
#             aug_imgs2 = np.load(aug_img_path + 'client2_aug_imgs.npy')
#             aug_imgs3 = np.load(aug_img_path + 'client3_aug_imgs.npy')
#             aug_imgs4 = np.load(aug_img_path + 'client4_aug_imgs.npy')
#             aug_imgs5 = np.load(aug_img_path + 'client5_aug_imgs.npy')

#             plot_latent_distribution(gt_imgs1, mixed_gt_imgs, aug_imgs1, dist_path+'client1_latent_dist')
#             plot_latent_distribution(gt_imgs2, mixed_gt_imgs, aug_imgs2, dist_path+'client2_latent_dist')
#             plot_latent_distribution(gt_imgs3, mixed_gt_imgs, aug_imgs3, dist_path+'client3_latent_dist')
#             plot_latent_distribution(gt_imgs4, mixed_gt_imgs, aug_imgs4, dist_path+'client4_latent_dist')
#             plot_latent_distribution(gt_imgs5, mixed_gt_imgs, aug_imgs5, dist_path+'client5_latent_dist')

if __name__ == '__main__':

    for number in range(5):
        work_path = MODEL_OUTPUT_PATH + "/" + 'number' + str(number)

        gt_emb_path = MODEL_OUTPUT_PATH + f'/number{number}/ablated_nofed/gt_embs/'
        gt_embs1 = np.load(gt_emb_path + 'client1_test_embs.npy')
        gt_embs2 = np.load(gt_emb_path + 'client2_test_embs.npy')
        gt_embs3 = np.load(gt_emb_path + 'client3_test_embs.npy')
        gt_embs4 = np.load(gt_emb_path + 'client4_test_embs.npy')
        gt_embs5 = np.load(gt_emb_path + 'client5_test_embs.npy')

        mixed_gt_embs = mix_dataset([gt_embs1, gt_embs2, gt_embs3, gt_embs4, gt_embs5])

        for type in ['ablated_nofed', 'ablated_fed']:
            save_path = work_path + "/" + type
            dist_path = save_path + "/" + "dist_embs" + "/"
            if not os.path.exists(dist_path):
                os.makedirs(dist_path)

            aug_emb_path = save_path + '/' + 'aug_embs/'
            aug_embs1 = np.load(aug_emb_path + 'client1_aug_embs.npy')
            aug_embs2 = np.load(aug_emb_path + 'client2_aug_embs.npy')
            aug_embs3 = np.load(aug_emb_path + 'client3_aug_embs.npy')
            aug_embs4 = np.load(aug_emb_path + 'client4_aug_embs.npy')
            aug_embs5 = np.load(aug_emb_path + 'client5_aug_embs.npy')

            plot_latent_distribution(gt_embs1, mixed_gt_embs, aug_embs1, dist_path+'client1_latent_dist')
            plot_latent_distribution(gt_embs2, mixed_gt_embs, aug_embs2, dist_path+'client2_latent_dist')
            plot_latent_distribution(gt_embs3, mixed_gt_embs, aug_embs3, dist_path+'client3_latent_dist')
            plot_latent_distribution(gt_embs4, mixed_gt_embs, aug_embs4, dist_path+'client4_latent_dist')
            plot_latent_distribution(gt_embs5, mixed_gt_embs, aug_embs5, dist_path+'client5_latent_dist')


import torch
import torch.nn as nn

###############################################################################
# Unconditional layers

class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        super().__init__()
        self.conv_layer = torch.nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                 padding=padding, dilation=dilation, groups=groups,
                 bias=bias, padding_mode=padding_mode))

    def forward(self, input):
        out = self.conv_layer(input)
        return out

class SpectralLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.linear_layer = torch.nn.utils.spectral_norm(nn.Linear(in_features, out_features, bias=bias))

    def forward(self, input):
        out = self.linear_layer(input)
        return out

class LeakyReLU2(nn.Module):
    def __init__(self, inplace=False):
        super().__init__()
        self.lrelu = nn.LeakyReLU(0.2, inplace=inplace)

    def forward(self, input):
        out = self.lrelu(input)
        return out

conv_layers = {
    'conv': nn.Conv2d,
    'spectral_conv': SpectralConv2d
}

linear_layers = {
    'linear': nn.Linear,
    'spectral_linear': SpectralLinear
}

activation_layers = {
    'relu': nn.ReLU,
    'lrelu': LeakyReLU2
}

norm_layers = {
    'BN': nn.BatchNorm2d,
    'identity': nn.Identity
}

###############################################################################
# Layers with conditional support

class CategoricalConditionalBatchNorm(torch.nn.Module):
    # as in the chainer SN-GAN implementation, we keep per-cat weight and bias
    def __init__(self, num_features, num_cats, eps=2e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super().__init__()
        self.num_features = num_features
        self.num_cats = num_cats
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        if self.affine:
            self.weight = torch.nn.Parameter(torch.Tensor(num_cats, num_features))
            self.bias = torch.nn.Parameter(torch.Tensor(num_cats, num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        if self.track_running_stats:
            self.register_buffer('running_mean', torch.zeros(num_features))
            self.register_buffer('running_var', torch.ones(num_features))
            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        else:
            self.register_parameter('running_mean', None)
            self.register_parameter('running_var', None)
            self.register_parameter('num_batches_tracked', None)
        self.reset_parameters()

    def reset_running_stats(self):
        if self.track_running_stats:
            self.running_mean.zero_()
            self.running_var.fill_(1)
            self.num_batches_tracked.zero_()

    def reset_parameters(self):
        self.reset_running_stats()
        if self.affine:
            self.weight.data.fill_(1.0)
            self.bias.data.zero_()

    def forward(self, inp):
        input = inp[0]
        cats = inp[1]
        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            self.num_batches_tracked += 1
            if self.momentum is None:  # use cumulative moving average
                exponential_average_factor = 1.0 / self.num_batches_tracked.item()
            else:  # use exponential moving average
                exponential_average_factor = self.momentum

        out = torch.nn.functional.batch_norm(
            input, self.running_mean, self.running_var, None, None,
            self.training or not self.track_running_stats,
            exponential_average_factor, self.eps)
        if self.affine:
            shape = [input.size(0), self.num_features] + (input.dim() - 2) * [1]
            weight = self.weight.index_select(0, cats).view(shape)
            bias = self.bias.index_select(0, cats).view(shape)
            out = out * weight + bias
        return (out, cats)

    def extra_repr(self):
        return '{num_features}, num_cats={num_cats}, eps={eps}, momentum={momentum}, affine={affine}, ' \
               'track_running_stats={track_running_stats}'.format(**self.__dict__)

class CondConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        super().__init__()
        self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                 padding=padding, dilation=dilation, groups=groups,
                 bias=bias, padding_mode=padding_mode)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.conv_layer(img)
        return (out, label)

class CondConvTranspose2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, output_padding=0, groups=1, bias=True,
                 dilation=1, padding_mode='zeros'):
        super().__init__()
        self.conv_transpose_layer = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride,
                 padding=padding, output_padding=output_padding, groups=groups, bias=bias,
                 dilation=dilation, padding_mode=padding_mode)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.conv_transpose_layer(img)
        return (out, label)

class UnconditionalBatchNorm2d(nn.Module):
    def __init__(self, num_features, num_classes=10, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super().__init__()
        self.batch_norm = nn.BatchNorm2d(num_features, eps=eps, momentum=momentum, affine=affine,
                 track_running_stats=track_running_stats)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.batch_norm(img)
        return (out, label)

class CondReLU(nn.Module):
    def __init__(self, inplace=False):
        super().__init__()
        self.relu = nn.ReLU(inplace=inplace)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.relu(img)
        return (out, label)

class CondLeakyReLU(nn.Module):
    def __init__(self, inplace=False):
        super().__init__()
        self.lrelu = nn.LeakyReLU(0.2, inplace=inplace)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.lrelu(img)
        return (out, label)

class CondTanh(nn.Module):
    def __init__(self):
        super().__init__()
        self.tanh = nn.Tanh()

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.tanh(img)
        return (out, label)

class CondUpsample(nn.Module):
    def __init__(self, scale_factor=1):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=scale_factor)

    def forward(self, input):
        img = input[0]
        label = input[1]
        out = self.upsample(img)
        return (out, label)

class CondSpectralLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.linear_layer = torch.nn.utils.spectral_norm(nn.Linear(in_features, out_features, bias=bias))

    def forward(self, input):
        img, label = input
        out = self.linear_layer(img)
        return out, label

class CondLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.linear_layer = nn.Linear(in_features, out_features, bias=bias)

    def forward(self, input):
        img, label = input
        out = self.linear_layer(img)
        return out, label

cond_norm_layers = {
    'BN': UnconditionalBatchNorm2d,
    'conditionalBN': CategoricalConditionalBatchNorm
}

cond_conv_layers = {
    'conv': CondConv2d,
    'convT': CondConvTranspose2d,
    'spectral_conv': None,
    'spectral_convT': None
}

cond_activation_layers = {
    'relu': CondReLU,
    'lrelu': CondLeakyReLU,
    'tanh': CondTanh
}

cond_linear_layers = {
    'linear': CondLinear,
    'spectral_linear': CondSpectralLinear
}


import logging
import threading
import time

from rest_framework.authentication import SessionAuthentication, BasicAuthentication
from rest_framework.decorators import (
    api_view,
    authentication_classes,
    permission_classes,
)
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response

from tradesense.repository.crypto_repository import CryptoRepository
from tradesense.services.analytics_service import AnalyticsService
from tradesense.services.socket_service import SocketService
from tradesense.services.tradesense_service import TradeSenseService
from tradesense.utils import get_message

# Create your views here.
@api_view(["GET"])
@authentication_classes([SessionAuthentication, BasicAuthentication])
@permission_classes([IsAuthenticated])
def start_trading(request):
    crypto_id = int(request.GET["id"])
    if crypto_id is None:
        logging.error("crypto id null")
        raise ValueError("Please pass crypto's id to trade in")
    conn = SocketService.get_connection()

    # Start a separate thread to handle threshold setting
    threshold_thread = threading.Thread(target=handle_threshold, args=(conn, crypto_id))
    threshold_thread.start()

    while True:
        tradesense_service = TradeSenseService(crypto_id)
        arbitrage_dto = tradesense_service.start_trading()
        msg = get_message(arbitrage_dto)
        logging.info(msg)
        conn.send(str(msg).encode())  # send data to the client
        time.sleep(60)

def handle_threshold(conn, crypto_id):
    while True:
        data = conn.recv(1024).decode()
        if data:
            # set_threshold:400
            try:
                threshold = int(data.split(":")[1])
                CryptoRepository.update_crypto_threshold(crypto_id, threshold)
                conn.send("Threshold successfully updated!".encode())
            except Exception as ex:
                logging.error(str(ex))
                conn.send(str(ex).encode())

@api_view(["GET"])
@authentication_classes([SessionAuthentication, BasicAuthentication])
@permission_classes([IsAuthenticated])
def analytics_view(request):
    analytics_dto = AnalyticsService().get_trade_analytics(request.GET)
    return Response(analytics_dto.to_json())


import os
import glob
import logging
import argparse
import tqdm
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

from data import fullsize_sequence, DOWNGRADES
from model import load_model
from util import reset_session

DATASETS = ['Set5_LR_bicubic', 'Set14_LR_bicubic', 'DIV2K_valid_LR_bicubic', 'BSDS100_LR_bicubic', 'Urban100_LR_bicubic']

logger = logging.getLogger(__name__)

def model_paths(input_dir):
	if os.path.splitext(input_dir)[1] == '.h5':
		return [input_dir]
	path_pattern = os.path.join(input_dir, '**', '*.h5')
	paths = glob.glob(path_pattern, recursive=True)
	paths.sort()
	return paths

def predict_model(model_path, image_path, scale, outdir):
	print('Load model %s', model_path)
	model = load_model(model_path, compile=False)

	for dataset_name in DATASETS:
		print(dataset_name)
		target_outdir = os.path.join(outdir, dataset_name, 'X%d'%(scale))
		target_image_path = os.path.join(image_path, dataset_name, 'X%d'%(scale))
		if not os.path.exists(target_outdir):
			os.makedirs(target_outdir)

		image_filenames = sorted(glob.glob(os.path.join(target_image_path, '*.png')))
		print(os.path.join(target_image_path, '*.png'))
		if len(image_filenames) == 0:
			logger.warning('No image files. Stop prediction..')
			exit()

		logger.info('Start prediction with %s', model_path)
		for i, f in tqdm.tqdm(enumerate(image_filenames)):
			filename = os.path.split(f)[1]
			im = Image.open(f).convert('RGB')
			im = np.array(im)
			im = np.reshape(im, (1,) + im.shape)
			output = model.predict(im, batch_size=1)
			output = np.squeeze(output)
			im_out = output
			if output.shape[-1] > 3:
				im_out = output[:,:,:3]
				logvar = output[:,:,3:]
				logvar_filename = os.path.join(target_outdir, os.path.splitext(filename)[0])
				np.save(logvar_filename, logvar)
				plt.imsave(logvar_filename+'_logvar.png', logvar.mean(axis=-1), cmap='viridis')

			im_out = Image.fromarray(np.uint8(im_out))
			im_out.save(os.path.join(target_outdir, filename))

def main(args):
	"""
	Evaluate all models in a user-defined directory against the DIV2K validation set.

	The results are written to a user-defined JSON file. All models in the input
	directory must have been trained for the same downgrade operator (bicubic or
	unknown) and the same scale (2, 3 or 4).
	"""

	mps = model_paths(args.indir)

	if mps:
		reset_session(args.gpu_memory_fraction)
		predict_model(mps[-1], args.dataset, args.scale, args.outdir)

	else:
		logger.warning('No models found in %s', args.indir)

def parser():
	parser = argparse.ArgumentParser(description='Evaluation against DIV2K validation set')

	parser.add_argument('-d', '--dataset', type=str, default='/home/esoc/datasets/SuperResolution/',
						help='path to DIV2K dataset with images stored as numpy arrays')
	parser.add_argument('-i', '--indir', type=str,
						help='path to models directory')
	parser.add_argument('-o', '--outdir', type=str, default='./image_predictions',
						help='output JSON file')
	parser.add_argument('-s', '--scale', type=int, default=2, choices=[2, 3, 4],
						help='super-resolution scale')
	parser.add_argument('--downgrade', type=str, default='bicubic', choices=DOWNGRADES,
						help='downgrade operation')
	parser.add_argument('--gpu-memory-fraction', type=float, default=0.8,
						help='fraction of GPU memory to allocate')

	return parser

if __name__ == '__main__':
	logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)

	main(parser().parse_args())