import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import datasets
import nltk  # Here to have a nice missing dependency error message early on
import numpy as np
from datasets import load_dataset, load_metric

import transformers
from filelock import FileLock
from transformers import (
    AutoConfig,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    MBart50Tokenizer,
    MBart50TokenizerFast,
    MBartTokenizer,
    MBartTokenizerFast,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    set_seed,
)
from transformers.file_utils import is_offline_mode
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version
from transformers.utils.versions import require_version

from model import XMCLIPBARTForConditionalGeneration
from data import Data
from trainer import XMTrainer, SelfCriticalTrainer

from scipy.special import softmax

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
# check_min_version("4.17.0.dev0")

# require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt")

logger = logging.getLogger(__name__)

try:
    nltk.data.find("tokenizers/punkt")
except (LookupError, OSError):
    if is_offline_mode():
        raise LookupError(
            "Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files"
        )
    with FileLock(".lock") as lock:
        nltk.download("punkt", quiet=True)

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={
            "help": "Path to pretrained model or model identifier from huggingface.co/models"
        }
    )
    config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained config name or path if not the same as model_name"
        },
    )
    tokenizer_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained tokenizer name or path if not the same as model_name"
        },
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={
            "help": "Where to store the pretrained models downloaded from huggingface.co"
        },
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={
            "help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."
        },
    )
    model_revision: str = field(
        default="main",
        metadata={
            "help": "The specific model version to use (can be a branch name, tag name or commit id)."
        },
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
            "with private models)."
        },
    )
    resize_position_embeddings: Optional[bool] = field(
        default=None,
        metadata={
            "help": "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
            "the model's position embeddings."
        },
    )

    # self critical
    train_self_critical: bool = field(default=False)
    reward_model: Optional[str] = field(default=None)
    rl_weight: Optional[float] = field(default=0.998)
    rouge_key: Optional[str] = field(default=None)

    cbs_weight: Optional[float] = field(default=1.0)
    rouge_weight: Optional[float] = field(default=1.0)

@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    # lang: str = field(default=None, metadata={"help": "Language id for summarization."})
    data_dir: Optional[str] = field(default=None)
    image_dir: Optional[str] = field(default=None)
    visual_feat_dim: Optional[int] = field(default=512)

    # dataset_name: Optional[str] = field(
    #     default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    # )
    # dataset_config_name: Optional[str] = field(
    #     default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    # )
    text_column: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the column in the datasets containing the full texts (for summarization)."
        },
    )
    summary_column: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the column in the datasets containing the summaries (for summarization)."
        },
    )
    overwrite_cache: bool = field(
        default=False,
        metadata={"help": "Overwrite the cached training and evaluation sets"},
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_source_length: Optional[int] = field(
        default=1024,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    max_target_length: Optional[int] = field(
        default=128,
        metadata={
            "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    val_max_target_length: Optional[int] = field(
        default=None,
        metadata={
            "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
            "during ``evaluate`` and ``predict``."
        },
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": "Whether to pad all samples to model maximum sentence length. "
            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
            "efficient on GPU but very bad for TPU."
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
            "value if set."
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
            "value if set."
        },
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
            "which is used during ``evaluate`` and ``predict``."
        },
    )
    ignore_pad_token_for_loss: bool = field(
        default=True,
        metadata={
            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
        },
    )
    source_prefix: Optional[str] = field(
        default="",
        metadata={
            "help": "A prefix to add before every source text (useful for T5 models)."
        },
    )

    forced_bos_token: Optional[str] = field(
        default=None,
        metadata={
            "help": "The token to force as the first generated token after the decoder_start_token_id."
            "Useful for multilingual models like mBART where the first generated token"
            "needs to be the target language token (Usually it is the target language token)"
        },
    )

    max_images: Optional[int] = field(default=10)

    def __post_init__(self):
        if self.val_max_target_length is None:
            self.val_max_target_length = self.max_target_length

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser(
        (ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)
    )
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(
            json_file=os.path.abspath(sys.argv[1])
        )
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    if data_args.source_prefix is None and model_args.model_name_or_path in [
        "t5-small",
        "t5-base",
        "t5-large",
        "t5-3b",
        "t5-11b",
    ]:
        logger.warning(
            "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with "
            "`--source_prefix 'summarize: ' `"
        )

    # Detecting last checkpoint.
    last_checkpoint = None
    if (
        os.path.isdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif (
            last_checkpoint is not None and training_args.resume_from_checkpoint is None
        ):
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name
        if model_args.config_name
        else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    config.visual_feat_dim = data_args.visual_feat_dim
    config.max_images = data_args.max_images

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name
        if model_args.tokenizer_name
        else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = XMCLIPBARTForConditionalGeneration.from_pretrained(
    # model = AutoModelForSeq2SeqLM.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
        ignore_mismatched_sizes=True,
    )

    model.resize_token_embeddings(len(tokenizer))

    print(sum(p.numel() for p in model.parameters()))

    if model.config.decoder_start_token_id is None:
        raise ValueError(
            "Make sure that `config.decoder_start_token_id` is correctly defined"
        )

    if (
        hasattr(model.config, "max_position_embeddings")
        and model.config.max_position_embeddings < data_args.max_source_length
    ):
        if model_args.resize_position_embeddings is None:
            logger.warning(
                f"Increasing the model's number of position embedding vectors from {model.config.max_position_embeddings} "
                f"to {data_args.max_source_length}."
            )
            model.resize_position_embeddings(data_args.max_source_length)
        elif model_args.resize_position_embeddings:
            model.resize_position_embeddings(data_args.max_source_length)
        else:
            raise ValueError(
                f"`--max_source_length` is set to {data_args.max_source_length}, but the model only has {model.config.max_position_embeddings}"
                f" position encodings. Consider either reducing `--max_source_length` to {model.config.max_position_embeddings} or to automatically "
                "resize the model's position encodings by passing `--resize_position_embeddings`."
            )

    prefix = data_args.source_prefix if data_args.source_prefix is not None else ""

    if training_args.label_smoothing_factor > 0 and not hasattr(
        model, "prepare_decoder_input_ids_from_labels"
    ):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    data = Data(data_args, tokenizer)
    data.load_dataset(
        training_args.do_train, training_args.do_eval, training_args.do_predict
    )

    train_dataset = data.train_dataset
    eval_dataset = data.eval_dataset
    predict_dataset = data.predict_dataset
    data_collator = data.get_data_collator(model, training_args.fp16)

    # Metric
    metric = load_metric("rouge")

    def postprocess_text(preds, labels):
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]

        # rougeLSum expects newline after each sentence
        preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
        labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]

        return preds, labels

    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        print(decoded_preds[:5])
        if data_args.ignore_pad_token_for_loss:
            # Replace -100 in the labels as we can't decode them.
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        result = metric.compute(
            predictions=decoded_preds, references=decoded_labels, use_stemmer=True
        )
        # Extract a few results from ROUGE
        result = {key: value.mid.fmeasure * 100 for key, value in result.items()}

        prediction_lens = [
            np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds
        ]
        result["gen_len"] = np.mean(prediction_lens)
        result = {k: round(v, 4) for k, v in result.items()}
        return result

    # Initialize our Trainer
    if model_args.train_self_critical:
        trainer = SelfCriticalTrainer(
        reward_model=model_args.reward_model,
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics
        if training_args.predict_with_generate
        else None,
        rl_weight=model_args.rl_weight,
        rouge_key=model_args.rouge_key,
        cbs_weight=model_args.cbs_weight,
        rouge_weight=model_args.rouge_weight,
    )
    else:
        trainer = XMTrainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset if training_args.do_train else None,
            eval_dataset=eval_dataset if training_args.do_eval else None,
            tokenizer=tokenizer,
            data_collator=data_collator,
            compute_metrics=compute_metrics
            if training_args.predict_with_generate
            else None,
        )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples
            if data_args.max_train_samples is not None
            else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_target_length
    )
    num_beams = (
        data_args.num_beams
        if data_args.num_beams is not None
        else training_args.generation_num_beams
    )
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(
            max_length=max_length, num_beams=num_beams, metric_key_prefix="eval"
        )
        max_eval_samples = (
            data_args.max_eval_samples
            if data_args.max_eval_samples is not None
            else len(eval_dataset)
        )
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")

        predict_results = trainer.predict(
            predict_dataset,
            metric_key_prefix="predict",
            max_length=max_length,
            num_beams=num_beams,
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples
            if data_args.max_predict_samples is not None
            else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        # parse predictions
        seq_pred = predict_results.predictions
        img_pred = None
        if isinstance(predict_results.predictions, tuple):
            seq_pred = predict_results.predictions[0]
            img_pred = predict_results.predictions[1]

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(
                    seq_pred,
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=True,
                )
                predictions = [pred.strip() for pred in predictions]
                output_prediction_file = os.path.join(
                    training_args.output_dir, "generated_predictions.txt"
                )
                with open(output_prediction_file, "w") as writer:
                    writer.write("\n".join(predictions))

            if img_pred is not None:
                # predictions = np.argmax(img_pred, axis=-1)
                img_pred = softmax(img_pred, axis=-1)
                output_predict_file = os.path.join(
                    training_args.output_dir, "image_predictions.txt"
                )
                predict_dataset.reset_format()  # so that we do not load image embeds
                with open(output_predict_file, "w") as writer:
                    for pred, imgs in zip(img_pred, predict_dataset["images"]):
                        pred = [p[1] for p in pred[: len(imgs)]]
                        pred = np.argsort(pred)[::-1][:3]
                        rec_images = [imgs[p] for p in pred]
                        writer.write(" ".join(rec_images) + "\n")

    return results

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

if __name__ == "__main__":
    main()


import json

def search_staff(firstname, lastname, middlename):
	file = open('data/staff.json', 'r')
	staffs = json.load(file)['data']

	results = []
	for staff in staffs:
		if firstname  != None and firstname  != staff['firstname'].lower():  continue
		if lastname   != None and lastname   != staff['lastname'].lower():   continue
		if middlename != None and middlename != staff['middlename'].lower(): continue
		results.append(staff)
	return results


#oq e o Yfinance
import yfinance as yf
import pandas as pd
from datetime import datetime

import streamlit as st
#como usar?
st.markdown('# Analisando empresas')

st.text_input('Preencha o Ticker Code: ',key ='tickercode',value = 'GOOG')
st.markdown(f'## Ultimas noticias da {st.session_state.tickercode} :')

ticker = st.session_state.tickercode

data = yf.Ticker(ticker)

data_news=pd.DataFrame(data.news)

data_news2 = data_news[['title','publisher','link','relatedTickers']]
st.dataframe(data_news2)

end_date = datetime.now().strftime('%Y-%m-%d')
data_hist = data.history(period='max',start='2019-03-16',end =end_date,interval='5d' )
data_hist = data_hist.reset_index()

st.markdown('# Contrua seu grafico: ')
ey= st.selectbox('Eixo y:',data_hist.columns)
ex = st.selectbox('Eixo x:',data_hist.columns)

st.markdown(f'## Grafico {ey} x {ex}:')
st.line_chart(data_hist,x = ex,y=ey)

print(data_hist)


# License AGPL-3.0 or later (https://www.gnu.org/licenses/agpl).

from odoo import api, fields, models
from datetime import datetime

class AccountPaymentOrder(models.Model):
    _inherit = 'account.payment.order'

    amount_untaxed_invoice = fields.Float(
        string='Amount untaxed invoice'
    )
    invoice_id = fields.Many2one(
        comodel_name='account.invoice',
        string='Invoice'
    )

    @api.multi
    def generated2uploaded(self):
        if self.payment_type == 'inbound' \
                and self.amount_untaxed_invoice > 0 \
                and not self.invoice_id.id:
            product_id = int(
                self.env['ir.config_parameter'].sudo().get_param(
                    'account_payment_order_invoice_product_id'
                )
            )
            if product_id > 0:
                product = self.env['product.product'].browse(product_id)
                vals = {
                    'partner_id': self.company_partner_bank_id.bank_id.partner_id.id,
                    'date': datetime.today(),
                    'date_invoice': datetime.today(),
                    'state': 'draft',
                    'type': 'in_invoice',
                    'comment': ' '
                }
                invoice_obj = self.env['account.invoice'].sudo().create(vals)
                self.invoice_id = invoice_obj.id
                vals = {
                    'invoice_id': self.invoice_id.id,
                    'product_id': product.id,
                    'name': product.name,
                    'quantity': 1,
                    'price_unit': self.amount_untaxed_invoice,
                    'account_id': product.property_account_expense_id.id,
                }
                invoice_line_obj = self.env['account.invoice.line'].sudo().create(vals)
                invoice_line_obj._onchange_product_id()
                invoice_line_obj.compute_taxes()
        # super
        return super(AccountPaymentOrder, self).generated2uploaded()


import numpy as np
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt

theta1 = pd.read_csv('theta1.csv')
theta2 = pd.read_csv('theta2.csv')
theta3 = pd.read_csv('theta3.csv')

theta1 = theta1.values[:, 1:]
theta2 = theta2.values[:, 1:]
theta3 = theta3.values[:, 1:]

#show picture
path = "dataset2/5/5_62057.jpg"
img = Image.open(path)
data = img.getdata()
data = np.matrix(data)
q = data.reshape(45,45)
print(q)
reim  = Image.fromarray(q)
plt.imshow(reim)
plt.show()

#recognize
def sigmoid(z):
    return 1 / (1 + np.exp(-z))
q = q.reshape(1, 2025)
q = np.insert(q, 0, values=np.ones(q.shape[0]), axis=1)
z2 = np.dot(q, theta1.T)
a2 = np.insert(sigmoid(z2), 0, values=np.ones(q.shape[0]), axis=1)
z3 = np.dot(a2, theta2.T)
a3 = np.insert(sigmoid(z3), 0, values=np.ones(q.shape[0]), axis=1)
z4 = np.dot(a3, theta3.T)
a4 = sigmoid(z4)
print(a4)
words = ["+", "-", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "=", "div", "x"]
result1  = np.argmax(a4)
print("数字为："+words[result1])


import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import sys
from einops import rearrange, repeat

from warpformer.Layers import EncoderLayer
from warpformer.Modules import Attention
from warpformer.WarpingLayer import Almtx

PAD = 0

def get_len_pad_mask(seq):
    """ Get the non-padding positions. """
    assert seq.dim() == 2
    seq = seq.ne(PAD)
    seq[:,0] = 1
    return seq.type(torch.float)

def get_attn_key_pad_mask_K(seq_k, seq_q, transpose=False):
    """ For masking out the padding part of key sequence. """
    # [B,L_q,K]
    if transpose:
        seq_q = rearrange(seq_q, 'b l k -> b k l 1')
        seq_k = rearrange(seq_k, 'b l k -> b k 1 l')
    else:
        seq_q = rearrange(seq_q, 'b k l -> b k l 1')
        seq_k = rearrange(seq_k, 'b k l -> b k 1 l')

    return torch.matmul(seq_q, seq_k).eq(PAD)

def get_subsequent_mask(seq):
    """ For masking out the subsequent info, i.e., masked self-attention. """
    sz_b, len_s, type_num = seq.size()
    subsequent_mask = torch.triu(
        torch.ones((len_s, len_s), device=seq.device, dtype=torch.uint8), diagonal=1)

    subsequent_mask = rearrange(subsequent_mask, 'l l -> b k l l', b=sz_b, k=type_num)
    return subsequent_mask

class FFNN(nn.Module):
    def __init__(self, input_dim, hid_units, output_dim):
        self.hid_units = hid_units
        self.output_dim = output_dim
        super(FFNN, self).__init__()

        self.linear = nn.Linear(input_dim, hid_units)
        self.W = nn.Linear(hid_units, output_dim, bias=False)

    def forward(self, x):
        x = self.linear(x)
        x = self.W(torch.tanh(x))
        return x

class Value_Encoder(nn.Module):
    def __init__(self, hid_units, output_dim, num_type):
        self.hid_units = hid_units
        self.output_dim = output_dim
        self.num_type = num_type
        super(Value_Encoder, self).__init__()

        self.encoder = nn.Linear(1, output_dim)

    def forward(self, x, non_pad_mask):
        non_pad_mask = rearrange(non_pad_mask, 'b l k -> b l k 1')
        x = rearrange(x, 'b l k -> b l k 1')
        x = self.encoder(x)
        return x * non_pad_mask

class Event_Encoder(nn.Module):
    def __init__(self, d_model, num_types):
        super(Event_Encoder, self).__init__()
        self.event_emb = nn.Embedding(num_types+1, d_model, padding_idx=PAD)

    def forward(self, event):
        # event = event * self.type_matrix
        event_emb = self.event_emb(event.long())
        return event_emb

class Time_Encoder(nn.Module):
    def __init__(self, embed_time, num_types):
        super(Time_Encoder, self).__init__()
        self.periodic = nn.Linear(1, embed_time - 1)
        self.linear = nn.Linear(1, 1)
        self.k_map = nn.Parameter(torch.ones(1,1,num_types,embed_time))

    def forward(self, tt, non_pad_mask):
        non_pad_mask = rearrange(non_pad_mask, 'b l k -> b l k 1')
        if tt.dim() == 3:  # [B,L,K]
            tt = rearrange(tt, 'b l k -> b l k 1')
        else: # [B,L]
            tt = rearrange(tt, 'b l -> b l 1 1')

        out2 = torch.sin(self.periodic(tt))
        out1 = self.linear(tt)
        out = torch.cat([out1, out2], -1) # [B,L,1,D]
        out = torch.mul(out, self.k_map)
        # return out * non_pad_mask # [B,L,K,D]
        return out

class MLP_Tau_Encoder(nn.Module):
    def __init__(self, embed_time, num_types, hid_dim=16):
        super(MLP_Tau_Encoder, self).__init__()
        self.encoder = FFNN(1, hid_dim, embed_time)
        self.k_map = nn.Parameter(torch.ones(1,1,num_types,embed_time))

    def forward(self, tt, non_pad_mask):
        non_pad_mask = rearrange(non_pad_mask, 'b l k -> b l k 1')
        if tt.dim() == 3:  # [B,L,K]
            tt = rearrange(tt, 'b l k -> b l k 1')
        else: # [B,L]
            tt = rearrange(tt, 'b l -> b l 1 1')

        # out1 = F.gelu(self.linear1(tt))
        tt = self.encoder(tt)
        tt = torch.mul(tt, self.k_map)
        return tt * non_pad_mask # [B,L,K,D]

class Encoder(nn.Module):
    """ A encoder model with self attention mechanism. """

    def __init__(
            self, opt,
            num_types, d_model, d_inner,
            n_layers, n_head, d_k, d_v, dropout):

        super().__init__()
        self.opt = opt
        self.d_model = d_model
        self.embed_time = d_model

        # event type embedding
        self.event_enc = Event_Encoder(d_model, num_types)
        self.type_matrix = torch.tensor([int(i) for i in range(1,num_types+1)]).to(opt.device)
        self.type_matrix = rearrange(self.type_matrix, 'k -> 1 1 k')

        self.num_types = num_types

        # self.enc_tau = opt.enc_tau
        self.a = nn.Parameter(torch.ones(1,num_types,1,1))
        self.b = nn.Parameter(torch.ones(1,num_types,1,1))
        self.sigma = nn.Parameter(torch.ones(1,num_types,1,1))
        self.pi = torch.tensor(math.pi)

        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout, opt=opt)
            for _ in range(n_layers)])

        self.value_enc = Value_Encoder(hid_units=d_inner, output_dim=d_model, num_type=num_types)
        self.learn_time_embedding = Time_Encoder(self.embed_time, num_types)

        self.w_t = nn.Linear(1, num_types, bias=False)

        self.tau_encoder = MLP_Tau_Encoder(self.embed_time, num_types)

        self.agg_attention = Attention_Aggregator(d_model, task=opt.task)

    def forward(self, event_time, event_value, non_pad_mask, tau=None, diffseg=None):
        """ Encode event sequences via masked self-attention. """
        '''
        non_pad_mask: [B,L,K]
        slf_attn_mask: [B,K,LQ,LK], the values to be masked are set to True
        len_pad_mask: [B,L], pick the longest length and mask the remains
        '''

        tem_enc_k = self.learn_time_embedding(event_time, non_pad_mask)  # [B,L,1,D], [B,L,K,D]
        tem_enc_k = rearrange(tem_enc_k, 'b l k d -> b k l d') # [B,K,L,D]

        value_emb = self.value_enc(event_value, non_pad_mask)
        value_emb = rearrange(value_emb, 'b l k d -> b k l d') # [B,K,L,D]

        self.type_matrix = self.type_matrix.to(non_pad_mask.device)
        event_emb = self.type_matrix * non_pad_mask
        event_emb = self.event_enc(event_emb)
        event_emb = rearrange(event_emb, 'b l k d -> b k l d') # [B,K,L,D]

        tau_emb = self.tau_encoder(tau, non_pad_mask)
        tau_emb = rearrange(tau_emb, 'b l k d -> b k l d')
        k_output = value_emb + tau_emb + event_emb + tem_enc_k

        non_pad_mask = rearrange(non_pad_mask, 'b l k -> b k l') # [B,K,L]

        for enc_layer in self.layer_stack:
            k_output, _, _ = enc_layer(
                k_output,
                non_pad_mask=non_pad_mask)

        non_pad_mask = rearrange(non_pad_mask, 'b k l -> b k l 1') # [B,K,L]
        output = self.agg_attention(k_output, non_pad_mask) # [B,D]
        # k_output = rearrange(k_output, 'b k l d -> b l k d')
        return output

class Hie_Encoder(nn.Module):
    """ A encoder model with self attention mechanism. """

    def __init__(
            self, opt,
            num_types, d_model, d_inner,
            n_layers, n_head, d_k, d_v, dropout):

        super().__init__()
        self.opt = opt
        self.d_model = d_model
        self.embed_time = d_model

        # event type embedding
        self.event_enc = Event_Encoder(d_model, num_types)
        self.type_matrix = torch.tensor([int(i) for i in range(1,num_types+1)]).to(opt.device)
        self.type_matrix = rearrange(self.type_matrix, 'k -> 1 1 k')
        self.num_types = num_types
        self.task = opt.task
        if len(opt.warp_num) > 0:
            warp_num = opt.warp_num
        elif opt.task == "mor" or opt.task == "wbm":
            warp_num = [0,12]
        else:
            warp_num = [0,6]

        if sum(opt.warp_num) == 0:
            self.no_warping = True
        else:
            self.no_warping = False
        self.full_attn = opt.full_attn

        if not opt.hourly and len(opt.warp_num) == 0:
            warp_layer_num = 1
        elif opt.hourly:
            warp_layer_num = 2
        else:
            warp_layer_num = len(warp_num)

        print("warp_num: ", str(warp_num), "\t warp_layer_num:", str(warp_layer_num))

        self.warpformer_layer_stack = nn.ModuleList([
            Warpformer(int(warp_num[i]), n_head, d_k, d_v, dropout, opt)
            for i in range(warp_layer_num)])

        self.value_enc = Value_Encoder(hid_units=d_inner, output_dim=d_model, num_type=num_types)
        self.learn_time_embedding = Time_Encoder(self.embed_time, num_types)

        self.w_t = nn.Linear(1, num_types, bias=False)

        self.tau_encoder = MLP_Tau_Encoder(self.embed_time, num_types)

        self.agg_attention = Attention_Aggregator(d_model, task=opt.task)
        self.linear = nn.Linear(d_model*warp_layer_num, d_model)

    def forward(self, event_time, event_value, non_pad_mask, tau=None, return_almat=False):
        """ Encode event sequences via masked self-attention. """
        '''
        non_pad_mask: [B,L,K]
        slf_attn_mask: [B,K,LQ,LK], the values to be masked are set to True
        len_pad_mask: [B,L], pick the longest length and mask the remains
        '''
        # embedding
        tem_enc_k = self.learn_time_embedding(event_time, non_pad_mask)  # [B,L,1,D], [B,L,K,D]
        tem_enc_k = rearrange(tem_enc_k, 'b l k d -> b k l d') # [B,K,L,D]

        value_emb = self.value_enc(event_value, non_pad_mask)
        value_emb = rearrange(value_emb, 'b l k d -> b k l d') # [B,K,L,D]

        self.type_matrix = self.type_matrix.to(non_pad_mask.device)
        # event_emb = self.type_matrix * non_pad_mask
        event_emb = self.type_matrix
        event_emb = self.event_enc(event_emb)
        event_emb = rearrange(event_emb, 'b l k d -> b k l d') # [B,K,L,D]

        tau_emb = self.tau_encoder(tau, non_pad_mask)
        tau_emb = rearrange(tau_emb, 'b l k d -> b k l d')

        if self.opt.remove_rep == 'abs':
            h0 = value_emb + tau_emb + event_emb
        elif self.opt.remove_rep == 'type':
            h0 = value_emb + tau_emb + tem_enc_k
        elif self.opt.remove_rep == 'rel':
            h0 = value_emb + event_emb + tem_enc_k
        elif self.opt.remove_rep == 'tem':
            h0 = value_emb + event_emb
        else:
            h0 = value_emb + tau_emb + event_emb + tem_enc_k

        amlt_list = []

        if self.opt.input_only:
            # b, k, l, d = h0.shape
            # z0 = output = self.agg_attention(h0, rearrange(non_pad_mask, 'b l k -> b k l 1'))
            # z0 = rearrange(h0, 'b k l d -> (b l) d k')
            # z0 = F.max_pool1d(z0, k).squeeze()
            # z0 = rearrange(z0, '(b l) d -> b d l', b=b, l=l)
            # z0 = F.max_pool1d(z0, l).squeeze(-1)
            z0 = torch.mean(h0, dim=1)
            if self.task != 'active':
                z0 = torch.mean(z0, dim=1)
        elif self.opt.dec_only:
            z0 = self.agg_attention(h0, rearrange(non_pad_mask, 'b l k -> b k l 1'))
        else:
            non_pad_mask = rearrange(non_pad_mask, 'b l k -> b k l') # [B,K,L]

            # if self.task != "active":
            #     z0 = self.agg_attention(h0, rearrange(non_pad_mask, 'b k l -> b k l 1')) # [B D]
            # else:
            #     z0 = h0 # [B K L D]
            z0 = None
            # warpformer, the first layer is identical warping
            idwarp=True
            for i, enc_layer in enumerate(self.warpformer_layer_stack):
                if enc_layer.new_l == 0:
                    idwarp=True
                else:
                    idwarp=False
                if i > 0 and self.no_warping and self.full_attn:
                    non_pad_mask = torch.ones_like(non_pad_mask).to(non_pad_mask.device)

                h0, non_pad_mask, almat = enc_layer(h0, non_pad_mask, event_time, id_warp=idwarp)

                output = self.agg_attention(h0, rearrange(non_pad_mask, 'b k l -> b k l 1'))

                if z0 is not None and z0.shape == output.shape:
                    z0 = z0 + output
                else:
                    z0 = output

                if almat is not None:
                    amlt_list.append(almat.detach().cpu())

        if return_almat:
            return z0, amlt_list
        else:
            return z0

class Warpformer(nn.Module):
    def __init__(self, new_l, n_head, d_k, d_v, dropout, opt):
        super().__init__()
        self.new_l = new_l
        self.num_types = opt.num_types
        self.opt = opt
        d_inner = opt.d_inner_hid
        d_model = opt.d_model

        if not opt.hourly:
            self.get_almtx = Almtx(opt, self.new_l)

        else:
            self.time_split = [i for i in range(self.new_l+1)] # for hour aggregation

        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout, opt=opt)
            for _ in range(opt.n_layers)])

    def hour_aggregate(self, event_time, h0, non_pad_mask):
        # time_split: [new_l]
        # event_time: [B,L]
        # h0: [B,K,L,D]
        # non_pad_mask: [B,K,L]
        new_l = len(self.time_split)-1
        b, k, l, dim = h0.shape

        event_time_k = repeat(event_time, 'b l -> b k l', k=self.num_types)
        new_event_time = torch.zeros((b,self.num_types,new_l)).to(h0.device)
        new_h0 = torch.zeros((b,self.num_types,new_l, dim)).to(h0.device)
        new_pad_mask = torch.zeros((b,self.num_types,new_l)).to(h0.device)
        almat = torch.zeros((b,l,new_l)).to(h0.device) # [B,L,S]

        # for each time slot
        for i in range(len(self.time_split)-1):
            idx = (event_time_k.ge(self.time_split[i]) & event_time_k.lt(self.time_split[i+1])) # [B,K,L]
            total = torch.sum(idx, dim=-1) # [B,K]
            total[total==0] = 1

            tmp_h0 = h0 * idx.unsqueeze(-1) # [B,K,L,D]
            tmp_h0 = rearrange(tmp_h0, 'b k l d -> (b k) d l')
            tmp_h0 = F.max_pool1d(tmp_h0, tmp_h0.size(-1)).squeeze() # [BK,D,1]
            new_h0[:,:,i,:] = rearrange(tmp_h0, '(b k) d -> b k d', b=b)
            almat[:,:,i] = (event_time.ge(self.time_split[i]) & event_time.lt(self.time_split[i+1])) # [B,L]

            new_event_time[:,:,i] = torch.sum(event_time_k * idx, dim=-1) / total
            new_pad_mask[:,:,i] = torch.sum(non_pad_mask * idx, dim=-1) / total

        almat = repeat(almat, 'b l s -> b k l s', k=k)
        return new_h0, new_event_time, new_pad_mask, almat

    def almat_aggregate(self, event_time, h0, non_pad_mask):
        # K: the number of clusters
        # event_time: [B,L]
        # h0: [B,K,L,D]
        # non_pad_mask: [B,K,L]

        b, k, l, dim = h0.shape
        new_event_time = None

        bound_mask, almat = self.get_almtx(h0, mask=non_pad_mask) # [bk, s, l], [bk, s, l]

        # almat = almat.to(h0.device)

        almat = rearrange(almat, '(b k) s l -> b k s l', k=k)
        bound_mask = rearrange(bound_mask, '(b k) s l -> b k s l', k=k)
        new_h0 = torch.matmul(almat, h0) # [b k s d]

        # # calculate new_event_time
        # event_time = repeat(event_time, 'b l -> b k l', k=(k*self.new_l))
        # event_time = repeat(event_time, 'b (k s) l -> b k l s', k=k, s=self.new_l)
        # new_event_time = event_time * hard_amlt
        # new_event_time = torch.sum(new_event_time, axis=-2)
        # new_event_time = new_event_time / torch.sum(hard_amlt, axis=-2)
        # new_event_time = torch.nan_to_num(new_event_time)

        # non_pad_mask = repeat(non_pad_mask, 'b k l -> b k l s', s=self.new_l)
        new_pad_mask = torch.sum(bound_mask, dim=-1)
        new_pad_mask[new_pad_mask > 0] = 1
        new_pad_mask = torch.nan_to_num(new_pad_mask)

        return new_h0, new_event_time, new_pad_mask, almat

    def forward(self, h0, non_pad_mask, event_time, id_warp=False):
        if id_warp:
            z0 = h0
            new_pad_mask = non_pad_mask
            almat = None
        else:
            # get alignment matrix
            if self.opt.hourly:
                z0, _, new_pad_mask, almat = self.hour_aggregate(event_time, h0, non_pad_mask)
            else:
                z0, _, new_pad_mask, almat = self.almat_aggregate(event_time, h0, non_pad_mask)

        for enc_layer in self.layer_stack:
            z0, _, _ = enc_layer(z0, non_pad_mask=new_pad_mask)

        return z0, new_pad_mask, almat

class Pool_Classifier(nn.Module):

    def __init__(self, dim, cls_dim):
        super(Pool_Classifier, self).__init__()
        self.classifier = nn.Linear(dim, cls_dim)

    def forward(self, ENCoutput):
        """
        input: [B,L,K,D]
        """
        b, l, k = ENCoutput.size(0), ENCoutput.size(1), ENCoutput.size(2)
        ENCoutput = rearrange(ENCoutput, 'b l k d -> (b l) d k')
        ENCoutput = F.max_pool1d(ENCoutput, k).squeeze()
        ENCoutput = rearrange(ENCoutput, '(b l) d -> b d l', b=b, l=l)
        ENCoutput = F.max_pool1d(ENCoutput, l).squeeze(-1)
        return self.classifier(ENCoutput)

class Attention_Aggregator(nn.Module):
    def __init__(self, dim, task):
        super(Attention_Aggregator, self).__init__()
        self.task = task
        self.attention_len = Attention(dim*2, dim)
        self.attention_type = Attention(dim*2, dim)

    def forward(self, ENCoutput, mask):
        """
        input: [B,K,L,D], mask: [B,K,L]
        """
        if self.task == "active":
            mask = rearrange(mask, 'b k l 1 -> b l k 1')
            ENCoutput = rearrange(ENCoutput, 'b k l d -> b l k d')
            ENCoutput, _ = self.attention_type(ENCoutput, mask) # [B L D]
        else:
            ENCoutput, _ = self.attention_len(ENCoutput, mask) # [B,K,D]
            ENCoutput, _ = self.attention_type(ENCoutput) # [B,D]
        return ENCoutput

class Classifier(nn.Module):

    def __init__(self, dim, type_num, cls_dim, activate=None):
        super(Classifier, self).__init__()
        # self.linear1 = nn.Linear(dim, type_num)
        # # self.activate = nn.Sigmoid()
        # self.linear2 = nn.Linear(type_num, cls_dim)
        self.linear = nn.Linear(dim, cls_dim)

    def forward(self, ENCoutput):
        """
        input: [B,L,K,D], mask: [B,L,K]
        """
        # ENCoutput = self.linear1(ENCoutput)
        # # if self.activate:
        # #     ENCoutput = self.activate(ENCoutput)
        # ENCoutput = self.linear2(ENCoutput)
        ENCoutput = self.linear(ENCoutput)
        return ENCoutput