#!/usr/bin/env python3  # ~~~ VERIFIED ~~~ #

# region         === importations ===         NOIGER #
import math
import random

from typing import Union
from typing import Tuple
from typing import Optional
from typing import List
from dataclasses import dataclass

import pandas as pd

import torch
from torch import nn
from torch import Tensor
from torch import FloatTensor
from torch import LongTensor
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.nn.modules.loss import L1Loss
from torch.nn import CrossEntropyLoss
from torch.nn import MSELoss
from torch.nn import CTCLoss

from transformers import BartModel
from transformers import BartForConditionalGeneration
from transformers import BartConfig
from transformers import get_linear_schedule_with_warmup
from transformers.modeling_outputs import BaseModelOutput
from transformers.modeling_outputs import Seq2SeqModelOutput
from transformers.modeling_outputs import Seq2SeqLMOutput
from transformers.modeling_outputs import CausalLMOutput
from transformers.generation_utils import GenerationMixin
from transformers import logging
from transformers.models.bart.modeling_bart import BartEncoder
from transformers.models.bart.modeling_bart import BartDecoder
from transformers.models.bart.modeling_bart import shift_tokens_right
from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding, BartEncoderLayer
from transformers.models.bart.modeling_bart import _expand_mask

import pytorch_lightning as pl

try:
    from .torch_cif import cif_function
    from .utils import mask_generator
    from .datasets import DataSetCollectorGeneral
except:
    from torch_cif import cif_function
    from utils import mask_generator
    from datasets import DataSetCollectorGeneral
# endregion      === importations ===      NOIGERDNE #

# region       === aug_dataclasses ===        NOIGER #
@dataclass
class AugBaseModelOutput(BaseModelOutput):
    last_hidden_state: torch.FloatTensor = None
    out_attention_mask: Optional[torch.FloatTensor] = None
    length_loss: Optional[Tuple[torch.FloatTensor]] = None
    pred_word_lengths: Optional[torch.LongTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None

@dataclass
class AugSeq2SeqModelOutput(Seq2SeqModelOutput):
    last_hidden_state: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_last_hidden_out_attention_mask: Optional[torch.FloatTensor] = None
    encoder_length_loss: Optional[Tuple[torch.FloatTensor]] = None
    encoder_pred_word_lengths: Optional[torch.LongTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None

@dataclass
class AugSeq2SeqLMOutput(Seq2SeqLMOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_last_hidden_out_attention_mask: Optional[torch.FloatTensor] = None
    encoder_length_loss: Optional[Tuple[torch.FloatTensor]] = None
    encoder_pred_word_lengths: Optional[torch.LongTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    masked_lm_loss: Optional[torch.FloatTensor] = None
    real_length_loss: Optional[torch.FloatTensor] = None
# endregion    === aug_dataclasses ===     NOIGERDNE #

logger = logging.get_logger("transformers.models.bart.modeling_bart")

class SentBartEncoder(BartEncoder):
    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        embed_dim = config.d_model
        self.collapse_n = getattr(config, "collapse_n", 0)
        self.config.collapse_n = self.collapse_n
        if self.collapse_n == -1:
            self.skip_cif = True
        else:
            self.skip_cif = False
        self.use_self = getattr(config, "use_self", False)
        self.config.use_self = self.use_self

        self.post_initialization(embed_dim, word_extractor=cif_function)

    def post_initialization(self, embed_dim, word_extractor=cif_function):
        # == encoder == #
        self.alpha_predictor = nn.Linear(embed_dim, 1)  # TODO: check!
        self.word_extractor = word_extractor
        self.length_predictor = nn.Linear(embed_dim, 1)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
      ) -> Union[Tuple, AugBaseModelOutput]:
        encoder_outputs = super().forward(
            input_ids,
            attention_mask,
            head_mask,
            inputs_embeds,
            output_attentions,
            output_hidden_states,
            return_dict,
        )
        encoder_outputs = BaseModelOutput(
            last_hidden_state=encoder_outputs[0],
            hidden_states=encoder_outputs[1]
                if len(encoder_outputs) > 1 else None,
            attentions=encoder_outputs[2]
                if len(encoder_outputs) > 2 else None,
        )

        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        if self.skip_cif:
            hidden_states = encoder_outputs.last_hidden_state
            out_attention_mask = attention_mask
            length_loss = (
                torch.zeros(len(hidden_states)).float(),
                torch.zeros(len(hidden_states)).long(),)
            pred_word_lengths = None
        else:
            if self.collapse_n == 1:    # 1 for all
                word_length_tensor = torch.ones_like(word_length_tensor)
            elif self.collapse_n == 2:  # all for all
                word_length_tensor = (attention_mask.sum(-1).long()
                                      .to(word_length_tensor.device))
                                        # else: N for all
            (hidden_states, out_attention_mask, length_loss,
                          # out_attention_mask := (
                          #     encoder_output_attention_mask)
             pred_word_lengths,
             _, _) = self.sent_retriever(
                encoder_outputs.last_hidden_state,
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                padding_mask=1 - attention_mask if attention_mask is not None else None)
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        if not return_dict:
            return tuple(v for v in [
                hidden_states,
                out_attention_mask,
                encoder_outputs.hidden_states,
                length_loss,
                pred_word_lengths,
                encoder_outputs.attentions,
            # ] if v is not None)
            ] if v is not None)
        return AugBaseModelOutput(
            last_hidden_state=hidden_states,
            out_attention_mask=out_attention_mask,
            length_loss=length_loss,
            pred_word_lengths=pred_word_lengths,
            hidden_states=encoder_outputs.hidden_states,
            # # attentions=encoder_outputs.attentions,
            # attentions=encoder_outputs.attentions,
        )

    def sent_retriever(self,
        encoder__last_hidden_state,
        word_length_tensor=None,
        alpha_values=None,
        padding_mask=None,
        return_all=False,
        return_original=False,
      ):
        if alpha_values is None:
            alpha_values = self.alpha_predictor(encoder__last_hidden_state)
            alpha_values = alpha_values.squeeze(-1).sigmoid()  # B x S

        if word_length_tensor is None or self.use_self:
            # print("No given! self predict")
            word_length_tensor = alpha_values.sum(-1).long()
        else:
            # print("Wordlen given")
            # predicted_word_length_tensor = alpha_values.sum(-1).long()
            pass

        encoder__word_representations_CIF = (
            self.word_extractor(
                encoder__last_hidden_state,
                alpha=alpha_values,
                padding_mask=padding_mask,
                target_lengths=word_length_tensor,
            )
        )
        # TODO: Keep all CIF
        [encoder_word_representation] = encoder__word_representations_CIF['cif_out']
        [pred_word_lengths] = encoder__word_representations_CIF['alpha_sum']
        encoder_word_representation = encoder_word_representation.contiguous()
        # pred_word_lengths = pred_word_lengths.contiguous()
        # length_loss = 0.
        length_loss = ((pred_word_lengths, word_length_tensor,)
            if word_length_tensor is not None else
            (pred_word_lengths, pred_word_lengths,))
            # aliased as `encoder_word_representation`
            # FIXME: distributed problem!
            # TODO: add other CIF ouptuts!
        # length_loss = length_loss.contiguous()

        encoder_output_attention_mask = (
            # mask_generator(word_length_tensor)
            mask_generator(word_length_tensor)
            if word_length_tensor is not None else
            mask_generator(pred_word_lengths))
            # TODO: check length prediction!

        return (
            encoder_word_representation,
            encoder_output_attention_mask,
            length_loss,
            pred_word_lengths,
            encoder__word_representations_CIF if return_all else None,
            # encoder__last_hidden_state if return_original else None,
            encoder__last_hidden_state if return_original else None,
        )

class SentBart(BartModel):
    def __init__(self, config: BartConfig):
        super(BartModel, self).__init__(config)

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
        self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)

        self.encoder = SentBartEncoder(config, self.shared)
        self.decoder = BartDecoder(config, self.shared)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, AugSeq2SeqModelOutput]:

        # different to other models, Bart automatically creates decoder_input_ids from
        # input_ids if no decoder_input_ids are provided
        if decoder_input_ids is None and decoder_inputs_embeds is None:
            if input_ids is None:
                raise ValueError(
                    "If no `decoder_input_ids` or `decoder_inputs_embeds` are "
                    "passed, `input_ids` cannot be `None`. Please pass either "
                    "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."
                )

            decoder_input_ids = shift_tokens_right(
                input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
            )

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
                head_mask=head_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        # If the user passed a tuple for encoder_outputs, we wrap it in an AugBaseModelOutput when return_dict=True
        elif return_dict and not isinstance(encoder_outputs, AugBaseModelOutput):
            encoder_outputs = AugBaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                out_attention_mask=encoder_outputs[1],
                length_loss=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
                pred_word_lengths=encoder_outputs[3] if len(encoder_outputs) > 3 else None,
                hidden_states=encoder_outputs[4] if len(encoder_outputs) > 4 else None,
                # attentions=encoder_outputs[5] if len(encoder_outputs) > 5 else None,
                    attentions=encoder_outputs[5] if len(encoder_outputs) > 5 else None,
            )
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            encoder_attention_mask=encoder_outputs[1],
            # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return AugSeq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_last_hidden_out_attention_mask=encoder_outputs.out_attention_mask,
            encoder_length_loss=encoder_outputs.length_loss,
            encoder_pred_word_lengths=encoder_outputs.pred_word_lengths,
            encoder_hidden_states=encoder_outputs.hidden_states,
            # # encoder_attentions=encoder_outputs.attentions,
            # encoder_attentions=encoder_outputs.attentions,
        )

    # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    def fix_encoder_(self, to_fix=True):
        self.encoder.requires_grad_(not to_fix)
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

class SentBartForConditionalGeneration(BartForConditionalGeneration):
    def __init__(self, config: BartConfig):
        super(BartForConditionalGeneration, self).__init__(config)
        self.model = SentBart(config)
        self.register_buffer(
            "final_logits_bias",
            torch.zeros((1, self.model.shared.num_embeddings)))
        self.lm_head = nn.Linear(
            config.d_model,
            self.model.shared.num_embeddings,
            bias=False)

        self.weight_len = getattr(config, "weight_len", None)
        self.config.weight_len = self.weight_len
        self.minimize_len = getattr(config, "minimize_len", False)
        self.config.minimize_len = self.minimize_len

        # Initialize weights and apply final processing
        self.post_init()

    # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    def fix_encoder_(self, to_fix=True):
        self.model.fix_encoder_(to_fix)
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, AugSeq2SeqLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            if use_cache:
                logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
            use_cache = False
            if decoder_input_ids is None and decoder_inputs_embeds is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        # print(self.model.encoder.layers[5].fc1.weight[0][:6] * 100)
        # print(self.model.encoder.alpha_predictor.weight[0][:6] * 100)
        # print(self.model.decoder.layers[0].fc1.weight[0][:6] * 100)

        if self.minimize_len:
            word_length_tensor = None

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            word_length_tensor=word_length_tensor,
            alpha_values=alpha_values,
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        lm_logits = self.lm_head(outputs[0]) + self.final_logits_bias

        masked_lm_loss = None
        if labels is not None:
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            loss_fct = CrossEntropyLoss()
            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
            # loss_fct = nn.CTCLoss(1)
            # masked_lm_loss = loss_fct(
            #     lm_logits, # $$$
            #     labels,
            #     attention_mask.sum(-1),
            #     (labels != 1).sum(-1))
            # breakpoint()
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        real_length_loss = None
        if word_length_tensor is not None and self.weight_len is not None:
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            loss_fct = nn.L1Loss()
            # real_length_loss = torch.linalg.vector_norm(outputs.encoder_length_loss.float(), ord=1)
            real_length_loss = loss_fct(*(outputs.encoder_length_loss))
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        # elif self.minimize_len and word_length_tensor is None:
        elif self.minimize_len:
            loss_fct = nn.L1Loss()
            real_length_loss = loss_fct(
                outputs.encoder_length_loss[0],
                torch.zeros_like(outputs.encoder_length_loss[0]),
            )

        loss = ((masked_lm_loss if masked_lm_loss is not None else 0.)
              + ((self.weight_len if self.weight_len is not None else 0.
                 ) * (real_length_loss if real_length_loss is not None else 0.))
        ) if ((masked_lm_loss is not None) or (real_length_loss is not None)) else None

        if not return_dict:
            output = (lm_logits,) + outputs[1:] + (masked_lm_loss, real_length_loss,)
            return (
                ((loss,) + output)
                if loss is not None else
                output)

        return AugSeq2SeqLMOutput(
            loss=loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_last_hidden_out_attention_mask=outputs.encoder_last_hidden_out_attention_mask,
            encoder_length_loss=outputs.encoder_length_loss,
            encoder_pred_word_lengths=outputs.encoder_pred_word_lengths,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
            masked_lm_loss=masked_lm_loss,
            real_length_loss=real_length_loss,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        word_length_tensor=None,
        alpha_values=None,
        past=None,
        attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        **kwargs
    ):
        output_dict = super().prepare_inputs_for_generation(
            decoder_input_ids,
            past,
            attention_mask,
            head_mask,
            decoder_head_mask,
            cross_attn_head_mask,
            use_cache,
            encoder_outputs,
            **kwargs
        )

        return {
            "word_length_tensor": word_length_tensor,
            "alpha_values": alpha_values,
            **output_dict,
        }

    @staticmethod
    def _expand_inputs_for_generation(
        input_ids: torch.LongTensor,
        expand_size: int = 1,
        is_encoder_decoder: bool = False,
        attention_mask=None,
        encoder_outputs=None,
        **model_kwargs,
    ):
        expanded_return_idx = (
            torch.arange(input_ids.shape[0])
                 .view(-1, 1)
                 .repeat(1, expand_size)
                 .view(-1)
                 .to(input_ids.device)
        )
        input_ids, model_kwargs = (
            GenerationMixin._expand_inputs_for_generation(
                input_ids,
                expand_size,
                is_encoder_decoder,
                attention_mask,
                encoder_outputs,
                **model_kwargs,
        ))
        if "word_length_tensor" in model_kwargs and model_kwargs.get("word_length_tensor") is not None:
            word_length_tensor = model_kwargs["word_length_tensor"]
            model_kwargs["word_length_tensor"] = word_length_tensor.index_select(0, expanded_return_idx)
        elif "word_length_tensor" in model_kwargs and model_kwargs.get("word_length_tensor") is None:
            model_kwargs.pop("word_length_tensor")
        # TODO: maybe check model arch?
        if "alpha_values" in model_kwargs and model_kwargs.get("alpha_values") is not None:
            alpha_values = model_kwargs["alpha_values"]
            model_kwargs["alpha_values"] = alpha_values.index_select(0, expanded_return_idx)
        elif "alpha_values" in model_kwargs and model_kwargs.get("alpha_values") is None:
            model_kwargs.pop("alpha_values")
        # TODO: maybe check model arch?

        if "encoder_outputs" in model_kwargs:
            if "out_attention_mask" in encoder_outputs:
                out_attention_mask = model_kwargs["encoder_outputs"]["out_attention_mask"]
                model_kwargs["encoder_outputs"]["out_attention_mask"] = out_attention_mask.index_select(0, expanded_return_idx)

        return input_ids, model_kwargs
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

class PreSentBartEncoder(BartEncoder):
    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        embed_dim = config.d_model
        self.collapse_n = getattr(config, "collapse_n", 0)
        self.config.collapse_n = self.collapse_n
        if self.collapse_n == -1:
            self.skip_cif = True
        else:
            self.skip_cif = False
        self.use_self = getattr(config, "use_self", False)
        self.config.use_self = self.use_self

        self.post_initialization(embed_dim, word_extractor=cif_function)

    def post_initialization(self, embed_dim, word_extractor=cif_function):
        # == encoder == #
        self.alpha_predictor = nn.Linear(embed_dim, 1)  # TODO: check!
        self.word_extractor = word_extractor
        self.length_predictor = nn.Linear(embed_dim, 1)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
      ) -> Union[Tuple, AugBaseModelOutput]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # retrieve input_ids and inputs_embeds
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale

        embed_pos = self.embed_positions(input_shape)

        hidden_states = inputs_embeds + embed_pos
        hidden_states = self.layernorm_embedding(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        # @@@@@@@@@@@@@@@@@@
        # # expand attention_mask
        # if attention_mask is not None:
        #     # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        #     attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
        # @@@@@@@@@@@@@@@@@@

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        # check if head_mask has a correct number of layers specified if desired
        if head_mask is not None:
            if head_mask.size()[0] != (len(self.layers)):
                raise ValueError(
                    f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
                )

        # ############################ #
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        if self.skip_cif:
            hidden_states = hidden_states
            out_attention_mask = attention_mask
            length_loss = (
                torch.zeros(len(hidden_states)).float(),
                torch.zeros(len(hidden_states)).long(),)
            pred_word_lengths = None
        else:
            if self.collapse_n == 1:    # 1 for all
                word_length_tensor = torch.ones_like(word_length_tensor)
            elif self.collapse_n == 2:  # all for all
                word_length_tensor = (attention_mask.sum(-1).long()
                                      .to(word_length_tensor.device))
                                        # else: N for all
            (hidden_states, out_attention_mask, length_loss,
                          # out_attention_mask := (
                          #     encoder_output_attention_mask)
             pred_word_lengths,
             _, __) = self.sent_retriever(
                hidden_states,
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                padding_mask=1 - attention_mask if attention_mask is not None else None)
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        # expand attention_mask
        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            attention_mask = _expand_mask(out_attention_mask, inputs_embeds.dtype)
        # ############################ #

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if self.training and (dropout_probability < self.layerdrop):  # skip the layer
                layer_outputs = (None, None)
            else:
                if self.gradient_checkpointing and self.training:

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs, output_attentions)

                        return custom_forward

                    layer_outputs = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(encoder_layer),
                        hidden_states,
                        attention_mask,
                        (head_mask[idx] if head_mask is not None else None),
                    )
                else:
                    layer_outputs = encoder_layer(
                        hidden_states,
                        attention_mask,
                        layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                        output_attentions=output_attentions,
                    )

                hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            encoder_outputs = tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        else:
            encoder_outputs = BaseModelOutput(
                last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
            )

        if not return_dict:
            return tuple(v for v in [
                hidden_states,
                out_attention_mask,
                encoder_outputs.hidden_states,
                length_loss,
                pred_word_lengths,
                encoder_outputs.attentions,
            # ] if v is not None)
            ] if v is not None)
        return AugBaseModelOutput(
            last_hidden_state=hidden_states,
            out_attention_mask=out_attention_mask,
            length_loss=length_loss,
            pred_word_lengths=pred_word_lengths,
            hidden_states=encoder_outputs.hidden_states,
            # # attentions=encoder_outputs.attentions,
            # attentions=encoder_outputs.attentions,
        )

    def sent_retriever(self,
        encoder__last_hidden_state,
        word_length_tensor=None,
        alpha_values=None,
        padding_mask=None,
        return_all=False,
        return_original=False,
      ):
        if alpha_values is None:
            alpha_values = self.alpha_predictor(encoder__last_hidden_state)
            alpha_values = alpha_values.squeeze(-1).sigmoid()  # B x S

        if word_length_tensor is None or self.use_self:
            # print("No given! self predict")
            word_length_tensor = alpha_values.sum(-1).long()
        else:
            # print("Wordlen given")
            # predicted_word_length_tensor = alpha_values.sum(-1).long()
            pass

        encoder__word_representations_CIF = (
            self.word_extractor(
                encoder__last_hidden_state,
                alpha=alpha_values,
                padding_mask=padding_mask,
                target_lengths=word_length_tensor,
            )
        )
        # TODO: Keep all CIF
        [encoder_word_representation] = encoder__word_representations_CIF['cif_out']
        [pred_word_lengths] = encoder__word_representations_CIF['alpha_sum']
        encoder_word_representation = encoder_word_representation.contiguous()
        # pred_word_lengths = pred_word_lengths.contiguous()
        # length_loss = 0.
        length_loss = ((pred_word_lengths, word_length_tensor,)
            if word_length_tensor is not None else
            (pred_word_lengths, pred_word_lengths,))
            # aliased as `encoder_word_representation`
            # FIXME: distributed problem!
            # TODO: add other CIF ouptuts!
        # length_loss = length_loss.contiguous()

        encoder_output_attention_mask = (
            # mask_generator(word_length_tensor)
            mask_generator(word_length_tensor)
            if word_length_tensor is not None else
            mask_generator(pred_word_lengths))
            # TODO: check length prediction!

        return (
            encoder_word_representation,
            encoder_output_attention_mask,
            length_loss,
            pred_word_lengths,
            encoder__word_representations_CIF if return_all else None,
            # encoder__last_hidden_state if return_original else None,
            encoder__last_hidden_state if return_original else None,
        )

class PreSentBart(SentBart):
    def __init__(self, config: BartConfig):
        super(SentBart, self).__init__(config)

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
        self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)

        self.encoder = PreSentBartEncoder(config, self.shared)
        self.decoder = BartDecoder(config, self.shared)

        # Initialize weights and apply final processing
        self.post_init()

class PreSentBartForConditionalGeneration(SentBartForConditionalGeneration):
    def __init__(self, config: BartConfig):
        super(SentBartForConditionalGeneration, self).__init__(config)
        self.model = PreSentBart(config)
        self.register_buffer(
            "final_logits_bias",
            torch.zeros((1, self.model.shared.num_embeddings)))
        self.lm_head = nn.Linear(
            config.d_model,
            self.model.shared.num_embeddings,
            bias=False)

        self.weight_len = getattr(config, "weight_len", None)
        self.config.weight_len = self.weight_len
        self.minimize_len = getattr(config, "minimize_len", False)
        self.config.minimize_len = self.minimize_len

        # Initialize weights and apply final processing
        self.post_init()


from jedis import *
import random

def main():

    jedi_council = [Jedi(input("name of the jedi: ")),Jedi(input("name of the jedi: ")),Jedi(input("name of the jedi: "))]
    sith_lords = [Sith(input("name of the sith: ")),Sith(input("name of the sith: ")),Sith(input("name of the sith: "))]

    number_of_fights = 0
    while len(sith_lords) >0:

        jedi_fighter = random.randint(1,3)-1
        sith_fighter = random.randint(1,len(sith_lords))-1

        fight = jedi_council[jedi_fighter].fight_method(sith_lords[sith_fighter])
        if fight == 0:
            pass
        elif fight == 1: #jedi wins
            jedi_council[jedi_fighter].train()
            del sith_lords[sith_fighter]
        else:
            jedi_council[jedi_fighter].train()
            sith_lords[sith_fighter].train()

        number_of_fights+=1

        if number_of_fights == 100:
            print("the sith has one, everyone came to the dark side for cookies")

    print("the light has peresevered over darkness")

if __name__ == "__main__":
    main()


from functions import Range, range_small, range_normal, range_simulation

if __name__ == '__main__':
    """
    Enunciado:
    Catorce ciudades, aproximadamente del mismo tamaño, se eligen para un
    estudio de seguridad vial. Siete de ellas se eligen al azar y durante un
    mes aparecen en los periódicos locales artículos relativos a la seguridad
    vial. Los números de accidentes de tránsito del mes posterior a la campaña
    son los siguientes:

    Grupo de tratamiento: 19 31 39 45 47 66 75
    Grupo de control: 28 36 44 49 52 72 72

    a) Calcular el p−valor exacto de este conjunto de datos, correspondiente a
       la hipótesis de que en ambos grupos se tienen resultados idénticos
       (es decir, los artículos no tuvieron ningún efecto).
    b) Calcular el p−valor aproximado en base a una aproximación normal,
    c) Calcular el p−valor aproximado en base a una simulación.
    """

    # Sample1 es mas chica
    sample_1 = [19, 31, 39, 45, 47, 66, 75]
    sample_2 = [28, 36, 44, 49, 52, 72, 72]

    r = Range(sample_1, sample_2)
    print("Rango: {}".format(r))

    n = len(sample_2)
    m = len(sample_1)

    p_value_1 = range_small(n, m, r)
    print("a) p-valor exacto")
    print("P_Value: {}\n".format(p_value_1))
    Iter = 10000

    p_value_2 = range_normal(n, m, r)
    print("b) p-valor normal")
    print("P_Value: {}\n".format(p_value_2))

    N = 10000
    p_value_3 = range_simulation(sample_1, sample_2, r, N)
    print("b) p-valor simulación")
    print("P_Value: {}\n".format(p_value_3))


import discord

import feedparser
import json
import time

import modules.convert_time as ct
from modules.log import log
import modules.misc as misc

def scan_feed(url):
    """Simply scans a feed and returns it. Does NOT check whether

    Args:
        url (string): the url to scan
    Returns:
        The feedparser object incl. all feed items.
    """
    return feedparser.parse(url)

def get_new_items_from_feed(*, feed=False, url=False, last_checked):
    """Scans a feed for new items and returns a list of them.

    Args:
        feed: an already scanned feed
        url (string): url to scan in case a feed isn't given directly
        last_checked (time.struct_time object): last time it got checked for new items
    Returns:
        list of all new items
    """
    # if
    if feed == False:
        feed = scan_feed(url)

    # check if the feed exists - otherwise return False to indicate that the page likely doesn't exist (anymore)
    try:
        feed.entries[0]
    except (NameError, AttributeError, IndexError):
        return False

    final = []
    # dont stop scanning after the items are older since some new items might've been lower in the list (like for Manga, some special in-between chapters might've been released later on and as such are lower in the list where they technically belong despite being new)
    for i in feed.entries:
        if i.published_parsed > ct.ms_to_struct(last_checked):
            final.append(i)

    # TODO: add the last 10 titles or something to the database as a double check

    return final

def get_embed_from_item(*, item, db_item, feed, channel_item, client):
    """Creates an embed from an rss feed.

    Args:
        item: the rss feed item that should be used
        db_item: the corresponding feed in the database (used in case name & url can't be parsed)
        feed: the full rss feed, not just the item (used for thumbnail)
        channel_item: the 'channels' item from the database, used to get 'guildID'
    Returns:
        discord.Embed object
    """
    embed = {}
    date_format = json.load(open('settings/config.json', 'r'))['rss']['dateFormat']

    if hasattr(item, 'title'):
        embed['title'] = item.title
    else:
        embed['title'] = db_item['name']

    if hasattr(item, 'summary'):
        # cut if summary is above 350 characters
        embed['description'] = item.summary[:350]
    else:
        embed['description'] = '[no description found]'

    if hasattr(item, 'link'):
        embed['url'] = item.link
    else:
        embed['url'] = db_item['url']

    embed['color'] = misc.get_embed_color(default_color=client.get_guild(channel_item['guildID']).me.color)

    embed = discord.Embed(**embed)

    if hasattr(item, 'published_parsed'):
        embed.set_footer(text=time.strftime(date_format, item.published_parsed))
    else:
        embed.set_footer(text='[could not parse time published]')

    image_found = False

    try:
        for i in item.media_content:
            if i['medium'] == 'image':
                embed.set_image(url = i['url'])
                image_found = True
    except (NameError, AttributeError, IndexError):
        pass

    # if no image was found, set the thumbnail as image instead
    try:
        if image_found == True:
            embed.set_thumbnail(url = feed.feed.image.href)
        else:
            embed.set_image(url = feed.feed.image.href)
    except (NameError, AttributeError, IndexError):
        pass

    return embed

def already_checked(last_checked, interval, current_time=ct.get_current_time()):
    """Checks if feed has already scanned in the last interval.

    Args:
        last_checked (number): the time it has last been checked in ms since epoch
        interval (number): the interval in seconds
        current_time (time.time_struct object): overwrite the current time
    """
    # multiply interval by 1000 to convert from seconds to ms
    # subtract 5s from interval to make sure that a scheduled scan doesn't hit it
    last_checked = last_checked + (interval * 1000 - 5000)
    last_checked = ct.ms_to_struct(last_checked)

    return last_checked > current_time

async def scan_all_feeds(*, client):
    """Scans all RSS feeds

    Args:
        client: the discord.py client object required for logging and posting new items
    """

    with open('settings/database.json', 'r+') as f:
        database = json.load(f)
        feeds = database['feeds']
        config = json.load(open('settings/config.json', 'r'))

        combine_posts = config['rss']['combinePosts']
        oldest_posts_first = config['rss']['oldestFirst']
        interval = config['rss']['interval'] * 60

        if interval < 600:
            interval = 600

        # convert both to seconds
        scan_delay = config['rss']['scanDelay'] / 1000
        post_delay = config['rss']['postDelay'] / 1000

        if scan_delay < 0:
            scan_delay = 0

        if post_delay < 0:
            post_delay = 0

        current_time = ct.get_current_time()

        last_checked = database['general']['lastChecked']

        # check if already scanned during the last interval
        if already_checked(last_checked, interval, current_time):
            await log(f" Skip scanning all feeds ({int(interval / 60)}min interval)..", 'info', client=client)
            return

        await log('Start scanning all feeds.', 'info', client=client)

        # go through all feeds that have to be scanned
        # use range() because the original feeds variable has to be edited
        for i in range(len(feeds)):
            time.sleep(scan_delay)

            last_checked = feeds[i]['lastChecked']

            # check if already scanned during the last interval
            if already_checked(last_checked, interval, current_time):
                await log(f" Skip scanning {feeds[i]['name']} ({int(interval / 60)}min interval).", 'spamInfo', client=client)
                continue

            await log(f"Start scanning {feeds[i]['name']}.", 'spamInfo', client=client)

            try:
                feed = scan_feed(feeds[i]['url'])
                results = get_new_items_from_feed(feed=feed, url=feeds[i]['url'], last_checked=feeds[i]['lastChecked'])
            except:
                await log(f"Failed to scan {feeds[i]['url']}.", 'warn', client=client)
                continue;

            # if feed is unavailable, mark it as such with the current time
            if results == False:
                await log(f"{feeds[i]['name']} is not available (url: '{feeds[i]['url']}'.", 'spamInfo', client=client)
                # only mark it with the current time if it wasn't already marked as unavailable before
                if feeds[i]['unavailable'] == False:
                    feeds[i]['unavailable'] = ct.struct_to_ms(current_time)

                    for c in feeds[i]['channels']:
                        time.sleep(post_delay)

                        channel = client.get_channel(c['channelID'])
                        try:
                            await channel.send(f"**{c['feedName']}** is not available (this message won't be posted again until it's available again). Feed URL: {feeds[i]['url']}")
                        except:
                            await log(f"Could not send {feeds[i]['name']} as the channel is likely deleted.", 'warn', client=client)

                continue

            # mark feeds as available
            feeds[i]['unavailable'] = False

            # check if posts should be combined
            if len(results) > combine_posts and combine_posts != 0:
                if oldest_posts_first == True:
                    first_item = results[len(results) - 1]
                else:
                    first_item = results[0]

                for c in feeds[i]['channels']:
                    time.sleep(post_delay)

                    embed = get_embed_from_item(item=first_item, db_item=feeds[i], feed=feed, channel_item=c, client=client)

                    channel = client.get_channel(c['channelID'])
                    try:
                        await channel.send(embed=embed)
                        await channel.send(f"...and {len(results) - combine_posts} more from {c['feedName']}.")
                    except:
                        await log(f"Could not send {feeds[i]['name']} as the channel is likely deleted.", 'warn', client=client)

            # post all items
            else:

                # reverse the results list so the oldest items come first
                if oldest_posts_first == True:
                    results = results[::-1]

                for c in feeds[i]['channels']:
                    time.sleep(post_delay)

                    channel = client.get_channel(c['channelID'])
                    for r in results:
                        embed = get_embed_from_item(item=r, db_item=feeds[i], feed=feed, channel_item=c, client=client)
                        try:
                            await channel.send(embed=embed)
                        except:
                            await log(f"Could not send {feeds[i]['name']} as the channel is likely deleted.", 'warn', client=client)

            feeds[i]['lastChecked'] = ct.struct_to_ms(ct.get_current_time())
            await log(f" Done scanning {feeds[i]['name']}.", 'spamInfo', client=client)

        database['feeds'] = feeds
        database['general']['lastChecked'] = ct.struct_to_ms(ct.get_current_time())
        # reset file position to the beginning - stackoverflow copy, dont ask
        f.seek(0)
        json.dump(database, f, indent=4)
        f.truncate()
        await log('All feeds scanned.', 'info', client=client)
        return


from typing import Tuple

import torch
import torch.jit
import torch.nn.functional as F

@torch.jit.script
def scale_mat_2d(scale: torch.Tensor) -> torch.Tensor:
    assert scale.dim() == 2 and scale.shape[1] == 2
    mat = torch.eye(3, device=scale.device, dtype=scale.dtype).unsqueeze(0).repeat(scale.shape[0], 1, 1)
    mat[:, 0, 0] = scale[:, 0]
    mat[:, 1, 1] = scale[:, 1]
    return mat

@torch.jit.script
def rotate_mat_2d(angle: torch.Tensor) -> torch.Tensor:
    assert angle.dim() == 1
    mat = torch.eye(3, device=angle.device, dtype=angle.dtype).unsqueeze(0).repeat(angle.shape[0], 1, 1)
    mat[:, 0, 0] = angle.cos()
    mat[:, 0, 1] = -angle.sin()
    mat[:, 1, 0] = angle.sin()
    mat[:, 1, 1] = angle.cos()
    return mat

@torch.jit.script
def translate_mat_2d(trans: torch.Tensor) -> torch.Tensor:
    assert trans.dim() == 2 and trans.shape[1] == 2
    mat = torch.eye(3, device=trans.device, dtype=trans.dtype).unsqueeze(0).repeat(trans.shape[0], 1, 1)
    mat[:, 0, 2] = trans[:, 0]
    mat[:, 1, 2] = trans[:, 1]
    return mat

@torch.jit.script
def trs_matrix_2d(trans: torch.Tensor, angle: torch.Tensor, scale: torch.Tensor, input_size: Tuple[int, int]) -> torch.Tensor:
    """
    :param trans: [b, (x, y)]
    :param angle: [b]
    :param scale: [b, (x, y)]
    :param input_size: tuple(w, h)
    :return: [b, 3, 3]
    """
    B = trans.shape[0]
    assert trans.shape == (B, 2)
    assert angle.shape == (B,)
    assert scale.shape == (B, 2)
    assert len(input_size) == 2
    input_size = torch.tensor([[input_size[1] / input_size[0], 1.0]], dtype=scale.dtype, device=scale.device)
    mat = translate_mat_2d(trans) @ scale_mat_2d(input_size) @ rotate_mat_2d(angle) @ scale_mat_2d(scale / input_size)
    return mat

@torch.jit.script
def crop_images(images, affine, output_size: Tuple[int, int]):
    grid = F.affine_grid(affine[:, :2], (images.shape[0], images.shape[1], output_size[1], output_size[0]), align_corners=False)
    crop = F.grid_sample(images, grid, align_corners=False)
    return crop

@torch.jit.script
def random_crop(images, crop_fraction: float = 0.9):
    im_size = images.shape[-2], images.shape[-1]
    tr = (torch.rand((images.shape[0], 2), device=images.device) * 2 - 1) * (1 - crop_fraction)
    angle = torch.zeros(images.shape[0], device=images.device)
    scale = torch.full((images.shape[0], 2), crop_fraction, device=images.device)
    trs = trs_matrix_2d(tr, angle, scale, im_size)
    return crop_images(images, trs, im_size)


import argparse
import os
import sys

import numpy as np
import torch
import yaml

from runners.real_data_runner import train, semisupervised, transfer, compute_representations, plot_transfer, \
    plot_representation, compute_mcc

def parse():
    parser = argparse.ArgumentParser(description='')

    parser.add_argument('--config', type=str, default='mnist.yaml', help='Path to the config file')
    parser.add_argument('--run', type=str, default='run', help='Path for saving running related data.')
    parser.add_argument('--n-sims', type=int, default=0, help='Number of simulations to run')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--baseline', action='store_true', help='Run the script for the baseline')
    parser.add_argument('--transfer', action='store_true',
                        help='Run the transfer learning experiments after pretraining')
    parser.add_argument('--semisupervised', action='store_true', help='Run semi-supervised experiments')
    parser.add_argument('--representation', action='store_true',
                        help='Run CCA representation validation across multiple seeds')
    parser.add_argument('--mcc', action='store_true', help='compute MCCs -- '
                                                           'only relevant for representation experiments')
    parser.add_argument('--second-seed', type=int, default=0, help='Second random seed for computing MCC -- '
                                                                   'only relevant for representation experiments')
    parser.add_argument('--subset-size', type=int, default=6000,
                        help='Number of data points per class to consider -- '
                             'only relevant for transfer learning if not run with --all flag')
    parser.add_argument('--all', action='store_true',
                        help='Run transfer learning experiment for many seeds and subset sizes -- '
                             'only relevant for transfer and representation experiments')
    parser.add_argument('--plot', action='store_true',
                        help='Plot selected experiment for the selected dataset')
    # THIS OPTIONS ARE FOR DEBUG ONLY --- WILL BE REMOVED
    parser.add_argument('-a', action='store_true')
    parser.add_argument('-p', action='store_true')
    parser.add_argument('-z', type=int, default=0)
    args = parser.parse_args()
    return args

def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace

def make_and_set_dirs(args, config):
    """call after setting args.doc to set and create necessary folders"""
    args.dataset = config.data.dataset.split('_')[0]  # take into account baseline datasets e.g.: mnist_transferBaseline
    if 'doc' in vars(args).keys():
        if config.model.positive:
            args.doc += 'p'
        if config.model.augment:
            args.doc += 'a'
        if config.model.final_layer:
            args.doc += str(config.model.feature_size)
        args.doc += config.model.architecture.lower()
    else:
        # args has no attribute doc
        args.doc = config.model.architecture.lower()
    args.doc = os.path.join(args.dataset, args.doc)  # group experiments by dataset
    if 'doc2' in vars(args).keys():
        # add second level doc folder
        args.doc2 = os.path.join(args.doc, args.doc2)
    else:
        # if not defined, set to level 1 doc
        args.doc2 = args.doc
    os.makedirs(args.run, exist_ok=True)
    args.log = os.path.join(args.run, 'logs', args.doc2)
    os.makedirs(args.log, exist_ok=True)
    args.checkpoints = os.path.join(args.run, 'checkpoints', args.doc2)
    os.makedirs(args.checkpoints, exist_ok=True)
    args.output = os.path.join(args.run, 'output', args.doc)
    os.makedirs(args.output, exist_ok=True)

    if 'doc_baseline' in vars(args).keys():
        if config.model.positive:
            args.doc_baseline += 'p'
        if config.model.augment:
            args.doc_baseline += 'a'
        if config.model.final_layer:
            args.doc_baseline += str(config.model.feature_size)
        args.doc_baseline += config.model.architecture.lower()
        args.doc_baseline = os.path.join(args.dataset, args.doc_baseline)
        args.checkpoints_baseline = os.path.join(args.run, 'checkpoints', args.doc_baseline)
        os.makedirs(args.checkpoints_baseline, exist_ok=True)
        args.output_baseline = os.path.join(args.run, 'output', args.doc_baseline)
        os.makedirs(args.output_baseline, exist_ok=True)

def main():
    if torch.cuda.is_available():
        dev = torch.cuda.device_count() - 1
        print("Running on gpu:{}".format(dev))
        # torch.cuda.set_device(dev)
    else:
        print("Running on cpu")
    args = parse()
    # load config
    with open(os.path.join('configs', args.config), 'r') as f:
        print('loading config file: {}'.format(os.path.join('configs', args.config)))
        config_raw = yaml.load(f, Loader=yaml.FullLoader)
    config = dict2namespace(config_raw)
    config.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    # print(config)
    # set random seeds
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    # FOR DEBUG ONLY
    if args.a:
        config.model.augment = True
    if args.p:
        config.model.positive = True
    if args.z > 0:
        config.model.final_layer = True
        config.model.feature_size = args.z

    # TRANSFER LEARNING EXPERIMENTS
    # 1- no special flag: pretrain icebeem on 0-7 // --doc should be different between datasets
    # 2- --transfer: train only g on 8-9 // --doc should be the same as in step 1
    # 3- no special flags but with baseline config: train icebeem on 8-9
    # 4- --transfer --baseline: train icebeem on 8-9 (same as above)
    # steps 2, 3 and 4 are for many seeds and many subset sizes: the user can do them manually, or add the flag --all
    # and the script will perform the loop
    # step 3 is only for debug and shouldn't be used in practice

    if not args.transfer and not args.semisupervised and not args.baseline and not args.plot and not args.representation:
        print('Training an ICE-BeeM on {}'.format(config.data.dataset))
        args.doc = 'transfer'
        make_and_set_dirs(args, config)
        train(args, config)

    if args.transfer and not args.baseline and not args.plot:
        if not args.all:
            print(
                'Transfer for {} - subset size: {} - seed: {}'.format(config.data.dataset, args.subset_size, args.seed))
            args.doc = 'transfer'
            make_and_set_dirs(args, config)
            transfer(args, config)
        else:
            new_args = argparse.Namespace(**vars(args))
            for n in [0, 500, 1000, 2000, 3000, 4000, 5000, 6000]:
                for seed in range(args.seed, args.n_sims + args.seed):
                    print('Transfer for {} - subset size: {} - seed: {}'.format(config.data.dataset, n, seed))
                    new_args.subset_size = n
                    new_args.seed = seed
                    # change random seed
                    np.random.seed(seed)
                    torch.manual_seed(seed)
                    new_args.doc = 'transfer'
                    make_and_set_dirs(new_args, config)
                    transfer(new_args, config)

    if config.data.dataset.lower() in ['mnist_transferbaseline', 'cifar10_transferbaseline',
                                       'fashionmnist_transferbaseline', 'cifar100_transferbaseline']:
        # this is just here for debug, shouldn't be run, use --baseline --transfer instead
        if not args.all:
            print('Transfer baseline for {} - subset size: {} - seed: {}'.format(config.data.dataset.split('_')[0],
                                                                                args.subset_size, args.seed))
            args.doc = 'transferBaseline'
            args.doc2 = 'size{}_seed{}'.format(args.subset_size, args.seed)
            make_and_set_dirs(args, config)
            train(args, config)
        else:
            new_args = argparse.Namespace(**vars(args))
            for n in [0, 500, 1000, 2000, 3000, 4000, 5000, 6000]:
                for seed in range(args.seed, args.n_sims + args.seed):
                    print('Transfer baseline for {} - subset size: {} - seed: {}'.format(
                        config.data.dataset.split('_')[0], n, seed))
                    new_args.subset_size = n
                    new_args.seed = seed
                    # change random seed
                    np.random.seed(seed)
                    torch.manual_seed(seed)
                    new_args.doc = 'transferBaseline'
                    new_args.doc2 = 'size{}_seed{}'.format(n, seed)
                    make_and_set_dirs(new_args, config)
                    train(new_args, config)

    if args.transfer and args.baseline and not args.plot:
        # update args and config
        new_args = argparse.Namespace(**vars(args))
        new_args.config = os.path.splitext(args.config)[0] + '_baseline' + os.path.splitext(args.config)[1]
        with open(os.path.join('configs', new_args.config), 'r') as f:
            print('loading baseline config file: {}'.format(os.path.join('configs', new_args.config)))
            config_raw = yaml.load(f, Loader=yaml.FullLoader)
        config = dict2namespace(config_raw)
        config.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        if not args.all:
            print(
                'Transfer baseline for {} - subset size: {} - seed: {}'.format(config.data.dataset.split('_')[0],
                                                                               new_args.subset_size, new_args.seed))
            new_args.doc = 'transferBaseline'
            new_args.doc2 = 'size{}_seed{}'.format(args.subset_size, args.seed)
            make_and_set_dirs(new_args, config)
            train(new_args, config)
        else:
            for n in [0, 500, 1000, 2000, 3000, 4000, 5000, 6000]:
                for seed in range(args.seed, args.n_sims + args.seed):
                    print('Transfer baseline for {} - subset size: {} - seed: {}'.format(
                        config.data.dataset.split('_')[0], n, seed))
                    new_args.subset_size = n
                    new_args.seed = seed
                    # change random seed
                    np.random.seed(seed)
                    torch.manual_seed(seed)
                    new_args.doc = 'transferBaseline'
                    new_args.doc2 = 'size{}_seed{}'.format(n, seed)
                    make_and_set_dirs(new_args, config)
                    train(new_args, config)

    # PLOTTING TRANSFER LEARNING
    # 1- just use of the flag --plot and --transfer AND NO other flag (except --config of course)
    if args.plot and not args.baseline and not args.semisupervised and args.transfer and not args.representation:
        print('Plotting transfer experiment for {}'.format(config.data.dataset))
        args.doc = 'transfer'
        args.doc_baseline = 'transferBaseline'
        make_and_set_dirs(args, config)
        plot_transfer(args, config)

    # SEMI-SUPERVISED EXPERIMENTS
    # 1- no special flag: pretrain icebeem on 0-7 // same as 1- above
    # 2- --semisupervised: classify 8-9 using pretrained icebeem // --doc should be the same as from step 1-
    # 3- --baseline: pretrain unconditional ebm  on 0-7: IT IS VERY IMPORTANT HERE TO SPECIFY A --doc THAT IS
    # DIFFERENT FROM WHEN RUN FOR ICEBEEM
    # 4- --semisupervised --baseline: classify 8-9 using unconditional ebm // --doc should be the same as from step 3-

    if args.baseline and not args.semisupervised and not args.transfer and not args.representation and not args.plot:
        print('Training a baseline EBM on {}'.format(config.data.dataset))
        args.doc = 'semisupervisedBaseline'
        make_and_set_dirs(args, config)
        train(args, config, conditional=False)

    if args.semisupervised and not args.baseline and not args.plot:
        print('Computing semi-supervised accuracy for ICE-BeeM on {}'.format(config.data.dataset))
        args.doc = 'transfer'  # first step of semi-supervised learning is the same as first step in transfer learning
        make_and_set_dirs(args, config)
        semisupervised(args, config)

    if args.semisupervised and args.baseline and not args.plot:
        print('Computing semi-supervised accuracy for baseline EBM on {}'.format(config.data.dataset))
        args.doc = 'semisupervisedBaseline'
        make_and_set_dirs(args, config)
        semisupervised(args, config)

    # COMPARE QUALITY OF REPRESENTATIONS ON REAL DATA
    # 1- --representation: trains ICE-BeeM on train dataset, and save learnt rep of test data for
    # different seeds
    # 2- --representation --baseline: trains unconditional EBM on train dataset, and save learnt rep of test data for
    # different seeds
    # 3- --representation --mcc: compute pairwaise mccs for ICE-BeeM:
    #       --all: do it for all random seeds used in step 1-
    #       --seed X --second-seed Y: compute mcc between seeds X and Y
    # 4- --representation --mcc --baseline: same as 3- for baseline
    # 5- --representation --plot: plot boxplot of MCCs in and out of sample
    #       --n-sims: only consider n_sims seeds and not all if n_ims < n seeds used in 1-

    if args.representation:
        config.n_labels = 10 if config.data.dataset.lower().split('_')[0] != 'cifar100' else 100
        if not args.mcc and not args.baseline and not args.plot:
            for seed in range(args.seed, args.n_sims + args.seed):
                print('Learning representation for {} - seed: {}'.format(config.data.dataset, seed))
                new_args = argparse.Namespace(**vars(args))
                new_args.seed = seed
                np.random.seed(args.seed)
                torch.manual_seed(args.seed)
                new_args.doc = 'representation'
                new_args.doc2 = 'seed{}'.format(seed)
                make_and_set_dirs(new_args, config)
                compute_representations(new_args, config)

        if args.baseline and not args.mcc and not args.plot:
            for seed in range(args.seed, args.n_sims + args.seed):
                print('Learning baseline representation for {} - seed: {}'.format(config.data.dataset, seed))
                new_args = argparse.Namespace(**vars(args))
                new_args.seed = seed
                np.random.seed(args.seed)
                torch.manual_seed(args.seed)
                new_args.doc = 'representationBaseline'
                new_args.doc2 = 'seed{}'.format(seed)
                make_and_set_dirs(new_args, config)
                compute_representations(new_args, config, conditional=False)

        if args.mcc and not args.baseline and not args.plot:
            if args.all:
                for seed in range(args.seed, args.n_sims + args.seed - 1):
                    for second_seed in range(seed + 1, args.n_sims + args.seed):
                        print('Computing MCCs for {} - seeds: {} and {}'.format(config.data.dataset, seed,
                                                                                second_seed))
                        new_args = argparse.Namespace(**vars(args))
                        new_args.seed = seed
                        new_args.second_seed = second_seed
                        np.random.seed(args.seed)
                        torch.manual_seed(args.seed)
                        new_args.doc = 'representation'
                        make_and_set_dirs(new_args, config)
                        compute_mcc(new_args, config)
            else:
                assert 'second_seed' in vars(args).keys()
                print('Computing MCCs for {} - seeds: {} and {}'.format(config.data.dataset, args.seed,
                                                                        args.second_seed))
                args.doc = 'representation'
                make_and_set_dirs(args, config)
                compute_mcc(args, config)

        if args.mcc and args.baseline and not args.plot:
            if args.all:
                for seed in range(args.seed, args.n_sims + args.seed - 1):
                    for second_seed in range(seed + 1, args.n_sims + args.seed):
                        print('Computing baseline MCCs for {} - seeds: {} and {}'.format(config.data.dataset, seed,
                                                                                second_seed))
                        new_args = argparse.Namespace(**vars(args))
                        new_args.seed = seed
                        new_args.second_seed = second_seed
                        np.random.seed(args.seed)
                        torch.manual_seed(args.seed)
                        new_args.doc = 'representationBaseline'
                        make_and_set_dirs(new_args, config)
                        compute_mcc(new_args, config)
            else:
                assert 'second_seed' in vars(args).keys()
                print('Computing baseline MCCs for {} - seeds: {} and {}'.format(config.data.dataset, args.seed,
                                                                                args.second_seed))
                args.doc = 'representationBaseline'
                make_and_set_dirs(args, config)
                compute_mcc(args, config)

        if args.plot:
            print('Plotting representation experiment for {}'.format(config.data.dataset))
            args.doc = 'representation'
            args.doc_baseline = 'representationBaseline'
            make_and_set_dirs(args, config)
            plot_representation(args, config)

if __name__ == '__main__':
    sys.exit(main())


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

# # Urban Sound Classification

# In[89]:

# In[95]:

import IPython.display as ipd
ipd.Audio('D:\\Urban Sound Classification\\train\\Train\\2022.wav')

# In[90]:

import numpy as np
import pandas as pd
import librosa
import librosa.display
import glob
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import os

# In[2]:

test = pd.read_csv("D:\\Urban Sound Classification\\test\\test.csv", index_col = 0)
train = pd.read_csv("D:\\Urban Sound Classification\\train\\train.csv", index_col = 0)
exclude = pd.read_csv("C:\\Users\\Zara\\Desktop\\exclude.csv", index_col = 0)
exclude_test = pd.read_csv("C:\\Users\\Zara\\Desktop\\exclude_test.csv", index_col = 0)

# In[3]:

#masking mismatched or missing indexes from the train and test csv files

a_index = train.index
b_index = exclude.index
mask = ~a_index.isin(b_index)
train = train.loc[mask]
train = pd.DataFrame(train)
train = train.reset_index()
train = train.reset_index(drop = True)
train = pd.DataFrame(train)
train.head(7)

# In[4]:

a_index = test.index
b_index = exclude_test.index
mask = ~a_index.isin(b_index)
test = test.loc[mask]
test = pd.DataFrame(test)
test = test.reset_index()
test = test.reset_index(drop = True)
test = pd.DataFrame(test)
test.head(7)

# In[93]:

y, sr = librosa.load('D:\\Urban Sound Classification\\train\\Train\\2022.wav')

# In[94]:

plt.figure(figsize=(12, 4))
plt.title('Waveplot: Children Playing')
librosa.display.waveplot(y, sr=sr)

# In[75]:

#counts of items in each category

train['Class'].value_counts()
plt.bar(train['Class'].value_counts().index, train['Class'].value_counts(), align='center', data=train)
plt.title('Class Distribution')
plt.xticks(rotation='vertical')
plt.xlabel('Class')
plt.ylabel('Frequency')
print('Class Distribution:\n', train['Class'].value_counts())

# In[82]:

#MFCC of children playing
mfccs = librosa.feature.mfcc(y, sr, n_mfcc=40)
plt.figure(figsize=(10,4))
librosa.display.specshow(mfccs, x_axis='time')
plt.colorbar()
plt.title('MFCC')
plt.tight_layout()

# In[81]:

#melspectrogram of children playing
melspectrogram =librosa.feature.melspectrogram(y=y, sr=sr, n_mels=40,fmax=8000)
plt.figure(figsize=(10,4))
librosa.display.specshow(librosa.power_to_db(melspectrogram,ref=np.max),y_axis='mel', fmax=8000,x_axis='time')
plt.colorbar(format='%+2.0f dB')
plt.title('Mel spectrogram')
plt.tight_layout()

#For explanation purpose.
#this function outputs the indexes of the files that cannot be read or are corrupt
#the output of the function was used to create the exclusion and exclusion_test files

def parser(row):
   # function to load files and extract features
    file_name = os.path.join(os.path.abspath(data_dir),str(row.ID)+'.wav')

   # handle exception to check if there isn't a file which is corrupted
    try:
      # here kaiser_fast is a technique used for faster extraction
        X, sample_rate = librosa.load(file_name, res_type='kaiser_fast')
      # we extract mfcc feature from data
        mfccs = np.mean(librosa.feature.mfcc(y=X, sr=sample_rate, n_mfcc=40).T,axis=0)
    except Exception as e:
        print("Error encountered while parsing file: ", file)
        return None, None

err = train.apply(parser, axis=1)
err.columns = ['feature', 'label']
# In[5]:

#extracting features from train set
#taking a mean of the mfccs

features=[]
labels=[]
for i in range(len(train)):
    data_dir = "D:\\Urban Sound Classification\\train\\Train\\"
    filename = os.path.join(os.path.abspath(data_dir)+'\\'+str(train.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    features.append(np.mean(librosa.feature.mfcc(x, sr=sample_rate, n_mfcc=40).T,axis=0))
    labels.append(train.Class[i])

# In[6]:

#extracting features from test set
features_test=[]
for i in range(len(test)):
    test_dir = "D:\\Urban Sound Classification\\test\\Test\\"
    filename=os.path.join(os.path.abspath(test_dir)+'\\'+str(test.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    features_test.append(np.mean(librosa.feature.mfcc(x, sr=sample_rate, n_mfcc=40).T,axis=0))

# In[7]:

#label encoding the categorical data (labels)
#making test and train

x=np.array(features)
x_test=np.array(features_test)

#encoding the labels
from sklearn.preprocessing import LabelEncoder
from keras.utils import to_categorical

le = LabelEncoder()
y = to_categorical(le.fit_transform(labels))

# In[8]:

print('Shape of Features(Train): ',x.shape)
print('Shape of Features(Test): ',x_test.shape)
print('Shape of Labels(Train): ',y.shape)

# # Nueral Network

# In[9]:

import tensorflow as tf

model1=tf.keras.models.Sequential([

    tf.keras.layers.Dense(256,activation='relu',input_shape=(40,)),
    #using dropout regularization to randomly drop nodes , reduce overfitting, and improve generalization
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),

    #10 nodes since we have 10 labels
    tf.keras.layers.Dense(10,activation='softmax')
])

# In[10]:

#using tensorflow callbacks to reduce the learning rate if loss does not improve

reduce =tf. keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, mode='auto')

model1.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])

history1=model1.fit(x,y, batch_size=32, epochs=100, validation_split=0.1,callbacks=[reduce])

# In[11]:

#plot of training and validation accuracy/loss per epoch

acc=history1.history['accuracy']
val_acc=history1.history['val_accuracy']
loss=history1.history['loss']
val_loss=history1.history['val_loss']

epochs=range(len(acc))

plt.plot(epochs,acc,'b',label='Training Accuracy')
plt.plot(epochs,val_acc,'r',label='Validation Accuracy')
plt.legend()
plt.figure()

plt.plot(epochs,loss,'b',label='Training Loss')
plt.plot(epochs,val_loss,'r',label='Validation Loss')
plt.legend()
plt.show()

# In[13]:

nn_pred=model1.predict(x_test)
nn_max=np.argmax(nn_pred,axis=1)
nn_pred_test=le.inverse_transform(nn_max)

# In[14]:

proba_nn = model1.predict_proba(x_test).ravel()

# In[22]:

nn_test_label = pd.DataFrame({'ID':test['ID'].values})
nn_test_label.insert(loc=0, column='Class', value=nn_pred_test)

# In[24]:

nn_test_label.head()

# # CNN

# In[25]:

max_pad_len=174 #defining a maximum length for each column

file_name = "D:\\Urban Sound Classification\\train\\Train\\8279.wav"
audio, sample_rate = librosa.load(file_name, res_type='kaiser_fast')
mfcc=librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
print('Shape of feature matrix: ',mfcc.shape)
pad_width = max_pad_len - mfcc.shape[1]
print('Pad width: ',pad_width)
mfcc = np.pad(mfcc, pad_width=((0, 0), (0, pad_width)), mode='constant')
print('Shape of feature matrix after padding: ',mfcc.shape)

# In[26]:

#CNN does not support varying column length.fix by setting max length and padding
features_padded=[]
labels=[]
pad=174 # defining max column length as 174 for variation

for i in range(len(train)):
    file_name = "D:\\Urban Sound Classification\\train\\Train\\"
    filename=os.path.join(os.path.abspath(file_name)+'\\'+str(train.ID[i]) +'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    mfcc=librosa.feature.mfcc(y=x, sr=sample_rate, n_mfcc=40)
    padding = pad - mfcc.shape[1]
    mfcc = np.pad(mfcc, pad_width=((0, 0), (0, padding)), mode='constant')
    features_padded.append(mfcc)
    labels.append(train.Class[i])

# In[27]:

features_test_padded=[]

for i in range(len(test)):
    file_name = "D:\\Urban Sound Classification\\test\\Test\\"
    filename=os.path.join(os.path.abspath(file_name)+'\\'+str(test.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    mfcc=librosa.feature.mfcc(y=x, sr=sample_rate, n_mfcc=40)
    padding = pad - mfcc.shape[1]
    mfcc = np.pad(mfcc, pad_width=((0, 0), (0, padding)), mode='constant')
    features_test_padded.append(mfcc)

# In[28]:

x_new=np.array(features_padded)
x_new=x_new.reshape(x_new.shape[0], 40, 174, 1) #reshaping the array for the convolution layer input

x_new_test=np.array(features_test_padded)
x_new_test=x_new_test.reshape(x_new_test.shape[0], 40, 174, 1)

#encoding the labels
from sklearn.preprocessing import LabelEncoder
from keras.utils import to_categorical

le = LabelEncoder()
y = to_categorical(le.fit_transform(labels))

# In[29]:

print('Shape of Features(Train): ',x_new.shape)
print('Shape of Features(Test): ',x_new_test.shape)
print('Shape of Labels(Train): ',y.shape)

# In[30]:

import tensorflow as tf

model2=tf.keras.models.Sequential([

    tf.keras.layers.Conv2D(16,(3,3),activation='relu',input_shape=(40,174,1)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(32,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(64,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(128,(3,3),activation='relu',padding='same'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.4),

    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),

    tf.keras.layers.Dense(10,activation='softmax')
])

# In[31]:

reduce =tf. keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, mode='auto')

model2.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
history2=model2.fit(x_new,y, batch_size=32, epochs=100, validation_split=0.1,callbacks=[reduce])

# In[32]:

acc=history2.history['accuracy']
val_acc=history2.history['val_accuracy']
loss=history2.history['loss']
val_loss=history2.history['val_loss']

epochs=range(len(acc)) #No. of epochs

#Plot training and validation accuracy per epoch
import matplotlib.pyplot as plt
plt.plot(epochs,acc,'b',label='Training Accuracy')
plt.plot(epochs,val_acc,'r',label='Validation Accuracy')
plt.legend()
plt.figure()

#Plot training and validation loss per epoch
plt.plot(epochs,loss,'b',label='Training Loss')
plt.plot(epochs,val_loss,'r',label='Validation Loss')
plt.legend()
plt.show()

# In[44]:

cnn_pred=model2.predict(x_new_test)
cnn_max=np.argmax(cnn_pred,axis=1)
cnn_pred_test=le.inverse_transform(cnn_max)

# In[45]:

proba_cnn = model2.predict_proba(x_new_test).ravel()

# In[47]:

cnn_test_label= pd.DataFrame({'ID':test['ID'].values})
cnn_test_label.insert(loc=0, column='Class', value=cnn_pred_test)

# In[48]:

cnn_test_label.head()


from collections import defaultdict

from dungeonsheets import features, weapons
from dungeonsheets.classes.classes import CharClass, SubClass

# PHB
class BerserkerPath(SubClass):
    """For some barbarians, rage is a means to an end--that end being
    violence. The Path of the Berserker is a path of untrammeled fury, slick
    with blood. As you enter the berserker's rage, you thrill in the chaos of
    battle, heedless of your own health or well-being.

    """
    name = "Path of the Berserker"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.Frenzy]
    features_by_level[6] = [features.MindlessRage]
    features_by_level[10] = [features.IntimidatingPresence]
    features_by_level[14] = [features.Retaliation]

class TotemWarriorPath(SubClass):
    """The Path of the Totem Warrior is a spiritual journey, as the barbarian
    accepts a spirit animal as guide, protector, and inspiration. In battle,
    your totem spirit fills you with supernatural might, adding magical fuel to
    your barbarian rage.

    Most barbarian tribes consider a totem animal to be kin to a particular
    clan. In such cases, it is unusual for an individual to have more than one
    totem animal spirit, though exceptions exist.

    """
    name = "Path of the Totem Warrior"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.SpiritSeeker, features.TotemSpirit]
    features_by_level[6] = [features.BeastAspect]
    features_by_level[10] = [features.SpiritWalker]
    features_by_level[14] = [features.TotemicAttunement]

# SCAG
class BattleragerPath(SubClass):
    """Known as Kuldjargh (literally "axe idiot") in Dwarvish, battleragers are
    dwarf followers of the gods of war and take the Path of the
    Battlerager. They specialize in wearing bulky, spiked armor and throwing
    themselves into combat, striking with their body itself and giving
    themselves over to the fury of battle.

    """
    name = "Path of the Battlerager"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.BattleragerArmor]
    features_by_level[6] = [features.RecklessAbandon]
    features_by_level[10] = [features.BattleragerCharge]
    features_by_level[14] = [features.SpikedRetribution]

# XGTE
class AncestralGuardianPath(SubClass):
    """Some barbarians hail from cultures that revere their ancestors. These
    tribes teach that the warriors of the past linger in the world as mighty
    spirits, who can guide and protect the living. When a barbarian who follows
    this path rages, the barbarian contacts the spirit world and calls on these
    guardian spirits for aid.

    Barbarians who draw on their ancestral guardians can better fight to
    protect their tribes and their allies. In order to cement ties to their
    ancestral guardians, barbarians who follow this path cover themselves in
    elabo- rate tattoos that celebrate their ancestors' deeds. These tattoos
    tell sagas of victories against terrible monsters and other fearsome
    rivals.

    """
    name = "Path of the Ancestral Guardian"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.AncestralProtectors]
    features_by_level[6] = [features.SpiritShield]
    features_by_level[10] = [features.ConsultTheSpirits]
    features_by_level[14] = [features.VengefulAncestors]

class StormHeraldPath(SubClass):
    """All barbarians harbor a fury within. Their rage grants them superior
    strength, durability, and speed. Barbarians who follow the Path of the
    Storm Herald learn to transform that rage into a mantle of primal magic,
    which swirls around them. When in a fury, a barbarian ofthis path taps into
    the forces of nature to create powerful magical effects.

    Storm heralds are typically elite champions who train alongside druids,
    rangers, and others sworn to protect nature. Other storm heralds hone their
    craft in lodges in regions wracked by storms, in the frozen reaches at the
    world's end, or deep in the hottest deserts.

    """
    name = "Path of the Storm Herald"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.StormAura]
    features_by_level[6] = [features.StormSoul]
    features_by_level[10] = [features.ShieldingStorm]
    features_by_level[14] = [features.RagingStorm]

class ZealotPath(SubClass):
    """Some deities inspire their followers to pitch themselves into a ferocious
    battle fury. These barbarians are zealots-warriors who channel their rage
    into powerful disn plays of divine power.

    A variety of gods across the worlds of D&D inspire their followers to
    embrace this path. Tempus from the Forgotten Realms and Hextor and Erythnul
    of Greyhawk are all prime examples. In general, the gods who inspire
    zealots are deities of combat, destruction, and violence. Not all are evil,
    but few are good

    """
    name = "Path of the Zealot"
    features_by_level = defaultdict(list)
    features_by_level[3] = [features.DivineFury, features.WarriorOfTheGods]
    features_by_level[6] = [features.FanaticalFocus]
    features_by_level[10] = [features.ZealousPresence]
    features_by_level[14] = [features.RageBeyondDeath]

class Barbarian(CharClass):
    """
    Barbarians, different as they might be, are defined by their rage: unbridled, unquenchable, and unthinking fury.
    More than a mere emotion, their anger is the ferocity of a cornered predator, the unrelenting assault of a storm,
    the churning turmoil of the sea.

    For some, their rage springs from a communion with fierce animal spirits. Others draw from a roiling reservoir of
    anger at a world full of pain. For every barbarian, rage is a power that fuels not just a battle frenzy but also
    uncanny reflexes, resilience, and feats of strength.

    # Primal Instinct
    People of towns and cities take pride in how their civilized ways set them apart from animals, as if denying one’s
    own nature was a mark of superiority. To a barbarian, though, civilization is no virtue, but a sign of weakness.
    The strong embrace their animal nature—keen instincts, primal physicality, and ferocious rage. Barbarians are
    uncomfortable when hedged in by walls and crowds. They thrive in the wilds of their homelands: the tundra, jungle,
    or grasslands where their tribes live and hunt.

    Barbarians come alive in the chaos of combat. They can enter a berserk state where rage takes over, giving them
    superhuman strength and resilience. A barbarian can draw on this reservoir of fury only a few times without resting,
    but those few rages are usually sufficient to defeat whatever threats arise.

    # A Life of Danger
    Not every member of the tribes deemed “barbarians” by scions of civilized society has the barbarian class. A true
    barbarian among these people is as uncommon as a skilled fighter in a town, and he or she plays a similar role as a
    protector of the people and a leader in times of war. Life in the wild places of the world is fraught with peril:
    rival tribes, deadly weather, and terrifying monsters. Barbarians charge headlong into that danger so that their
    people don’t have to.

    Their courage in the face of danger makes barbarians perfectly suited for adventuring. Wandering is often a way of
    life for their native tribes, and the rootless life of the adventurer is little hardship for a barbarian. Some
    barbarians miss the close-knit family structures of the tribe, but eventually find them replaced by the bonds formed
    among the members of their adventuring parties.
    """
    name = 'Barbarian'
    hit_dice_faces = 12
    subclass_select_level = 3
    saving_throw_proficiencies = ('strength', 'constitution')
    primary_abilities = ('strength',)
    weapon_proficiencies = (weapons.SimpleWeapon, weapons.MartialWeapon)
    _proficiencies_text = ('light armor', 'medium armor', 'shields',
                           'simple weapon_list', 'martial weapon_list')
    multiclass_weapon_proficiencies = weapon_proficiencies
    _multiclass_proficiencies_text = ('shields', 'simple weapon_list',
                                      'martial weapon_list')
    class_skill_choices = ('Animal Handling', 'Athletics',
                           'Intimidation', 'Nature', 'Perception', 'Survival')
    subclasses_available = (BerserkerPath, TotemWarriorPath, BattleragerPath,
                            AncestralGuardianPath, StormHeraldPath, ZealotPath)
    features_by_level = defaultdict(list)
    features_by_level[1] = [features.BarbarianAbilityScoreImprovement, features.Rage, features.UnarmoredDefenseBarbarian]
    features_by_level[2] = [features.RecklessAttack, features.DangerSense]
    features_by_level[5] = [features.ExtraAttackBarbarian,
                            features.FastMovement]
    features_by_level[7] = [features.FeralInstinct]
    features_by_level[9] = [features.BrutalCritical]
    features_by_level[11] = [features.RelentlessRage]
    features_by_level[15] = [features.PersistentRage]
    features_by_level[18] = [features.IndomitableMight]
    features_by_level[20] = [features.PrimalChampion]


import pulumi
import pulumi_equinix as equinix

config = pulumi.Config()
device_id = config.require("deviceId")
subnet_cidr = config.get("subnetCidr")
if subnet_cidr is None:
    subnet_cidr = "147.229.10.152/31"
ip_attach_resource = equinix.metal.IpAttachment("ipAttach",
    device_id=device_id,
    cidr_notation=subnet_cidr)
pulumi.export("ipAttach", ip_attach_resource.id)
pulumi.export("ipNetmask", ip_attach_resource.netmask)


from inspect import getmembers, isfunction
from testconfig import config
from fixture import PesterFixture

tests_to_run = ["test_connection", "test_display_name"]
fixture = False

def setup_module():
    print("setup_module")
    browsers = ['chrome', 'chrome']
    print(config)
    if 'browser1' in config:
        browsers[0] = config['browser1']
    if 'browser2' in config:
        browsers[1] = config['browser2']

    global fixture
    fixture = PesterFixture(browsers=browsers)
    fixture.connect_all()

def teardown_module():
    fixture.close_all()

def test_generate_tests():
    import importlib

    for module_name in tests_to_run:
        module = importlib.import_module(module_name)
        if hasattr(module, 'get_tests'):
            functions = module.get_tests()
            for test in functions:
                test_generate_tests.compat_func_name = module_name + '.' +\
                                                       test.__name__;
                yield test
        else:
            functions = [o for o in getmembers(module) if
                         isfunction(o[1]) and o[0].startswith('test')]
            for test in functions:
                test_generate_tests.compat_func_name = module_name + '.' + \
                                                       test[0]
                yield test[1]


from queue import PriorityQueue
from typing import Tuple

class Documento:
    def __init__(self, nombre: str, tipo: str):
        self.nombre = nombre
        self.tipo = tipo

    def __lt__(self, other):
        return False

def imprimir_cola(cola: PriorityQueue[Tuple[int, Documento]]):
    print("\nDocumentos en la cola de impresión:")
    while not cola.empty():
        documento = cola.get()
        print(f"Documento: {documento[1].nombre} (Tipo: {documento[1].tipo})")

cola_impresion: PriorityQueue[Tuple[int, Documento]] = PriorityQueue()

cola_impresion.put((1, Documento("Doc1 (Empleado)", "Empleado")))
cola_impresion.put((1, Documento("Doc2 (Empleado)", "Empleado")))
cola_impresion.put((1, Documento("Doc3 (Empleado)", "Empleado")))

primer_documento = cola_impresion.get()
print(f"Documento impreso: {primer_documento[1].nombre}")

cola_impresion.put((2, Documento("Doc4 (TI)", "TI")))
cola_impresion.put((2, Documento("Doc5 (TI)", "TI")))

cola_impresion.put((3, Documento("Doc6 (Gerente)", "Gerente")))

segundo_documento = cola_impresion.get()
tercer_documento = cola_impresion.get()
print(f"Segundo documento impreso: {segundo_documento[1].nombre}")
print(f"Tercer documento impreso: {tercer_documento[1].nombre}")

cola_impresion.put((1, Documento("Doc7 (Empleado)", "Empleado")))
cola_impresion.put((1, Documento("Doc8 (Empleado)", "Empleado")))
cola_impresion.put((3, Documento("Doc9 (Gerente)", "Gerente")))

imprimir_cola(cola_impresion)


import logging
_logger = logging.getLogger(__name__)

import glob
import itertools

from gunpowder import ArrayKey, ArraySpec, Hdf5Source

from fuse import Z5Source

RAW_KEY                = ArrayKey('RAW')
DEFECT_MASK_KEY        = ArrayKey('DEFECT_MASK')
LABELS_KEY             = ArrayKey('LABELS')
NEURON_IDS_KEY         = ArrayKey('NEURON_IDS')
NEURON_IDS_NO_GLIA_KEY = ArrayKey('NEURON_IDS_NOGLIA')
MASK_KEY               = ArrayKey('MASK')
TRAINING_MASK_KEY      = ArrayKey('TRAINING_MASK')
LOSS_GRADIENT_KEY      = ArrayKey('LOSS_GRADIENT')
AFFINITIES_KEY         = ArrayKey('AFFINITIES')
GT_AFFINITIES_KEY      = ArrayKey('GT_AFFINITIES')
AFFINITIES_MASK_KEY    = ArrayKey('AFFINITIES_MASK')
AFFINITIES_SCALE_KEY   = ArrayKey('AFFINITIES_SCALE')
AFFINITIES_NN_KEY      = ArrayKey('AFFINITIES_NN')
GLIA_MASK_KEY          = ArrayKey('GLIA_MASK')
GLIA_KEY               = ArrayKey('GLIA')
GLIA_SCALE_KEY         = ArrayKey('GLIA_SCALE')
GT_GLIA_KEY            = ArrayKey('GT_GLIA')

_LOCALS = dict(locals())
_lower_key_identifier_to_key = {v.identifier.lower() : v for _, v in _LOCALS.items() if isinstance(v, ArrayKey)}
_logger.info('Lower key identifier to key mapping: %s', _lower_key_identifier_to_key)

DEFAULT_PATHS = {
    RAW_KEY.identifier.lower()                : 'volumes/raw',
    LABELS_KEY.identifier.lower()             : 'volumes/labels/neuron_ids-downsampled',
    NEURON_IDS_KEY.identifier.lower()         : 'volumes/labels/neuron_ids-downsampled',
    MASK_KEY.identifier.lower()               : 'volumes/labels/mask-downsampled',
    GT_GLIA_KEY.identifier.lower()            : 'volumes/labels/glia-downsampled',
    GLIA_MASK_KEY.identifier.lower()          : 'volumes/labels/mask-downsampled',
    NEURON_IDS_NO_GLIA_KEY.identifier.lower() : 'volumes/labels/neuron_ids_noglia-downsampled'
}

def make_data_providers(
        *provider_strings,
        required_paths=('raw', 'labels', 'mask'),
        default_paths=DEFAULT_PATHS,
        lower_key_identifier_to_key=_lower_key_identifier_to_key):
    return tuple(itertools.chain.from_iterable(tuple(make_data_provider(s, required_paths, default_paths, lower_key_identifier_to_key) for s in provider_strings)))

def make_data_provider(provider_string, required_paths, default_paths, lower_key_identifier_to_key):
    data_providers = []
    # data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/realigned'
    # file_pattern = '*merged*fixed-offset-fixed-mask.h5'

    pattern = provider_string.split(':')[0]
    paths   = {r: default_paths[r] for r in required_paths if r in default_paths}
    paths.update(**{entry.split('=')[0].lower() : entry.split('=')[1] for entry in provider_string.split(':')[1:]})

    for r in required_paths:
        assert r in paths, "Required path {} not provided in {}".format(r, paths)

    for k, _ in paths.items():
        assert k in lower_key_identifier_to_key, "Identifier %s not in array key map %s" % (k, lower_key_identifier_to_key)

    datasets = {lower_key_identifier_to_key[k] : v for k, v in paths.items()}

    specs = {
        MASK_KEY:      ArraySpec(interpolatable=False),
        GLIA_MASK_KEY: ArraySpec(interpolatable=False)
    }

    for data in glob.glob(pattern):
        if data.endswith('h5') or data.endswith('hdf'):
            source = Hdf5Source(data, datasets=datasets, array_specs=specs)
        else:
            source = Z5Source(data, datasets=datasets, array_specs=specs, revert=False)
        data_providers.append(source)

    return tuple(data_providers)


# -*- coding: utf-8 -*-
"""
Created on Mon Feb  7 23:29:10 2022

@author: huangguo
"""
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pickle
from SIS_sim import SIS_sim
from me_analysis import me_analysis
N = 100
k = 10
p = 0.1
#p_n = 0.2
p_ws = 0.5
dir1 = 'npz文件\\'
dir2 = 'png文件\\'
for NetworkType in ['er','ws','re']:
    for nn in range(1,11):
        for lambda1 in np.arange(0.5,3.5,0.5):
            network={'NetworkType':NetworkType,
             'N':N,
             'k':k,
             'p_ws':p_ws,
             #'p_n':p_n,
             'strategy':'random'
             #'p_ws':p_ws,
             #'p':p,
            # 'p_n':p_n
            }
            param ={'rho0':0.1,
                        't0':1000,
                        'lambda1':lambda1,
                        #'beta_1':0.2,
                        #'lambda2':lambda2,
                        #'beta_2':0.2,
                        'mu':0.3,
                        'p_n':0.5,
                        'simplex':1,
                        'gamma1':1,
                        'gamma2':1,
                        'alpha1':0
                        }
            sis = SIS_sim(network,param)
            sis.run()
            pickle.dump(sis, open(dir1+"random1_"+NetworkType+"_"+str(round(lambda1,1))+'_'+str(nn)+".pkl", 'wb'))
            print("random1_"+NetworkType+"_"+str(round(lambda1,1))+'_'+str(nn))

代码/分析_1_一阶传播_随机_随机配置.py

分析_1_一阶传播_随机_随机配置.py