import torch
import torch.nn as nn
from torch.utils.data import Dataset
from transformers import AutoModel


# =========================================================
# DATASET
# =========================================================

class ChunkedReadmissionDataset(Dataset):

    def __init__(
        self,
        texts,
        labels,
        tokenizer,
        max_length=256,
        max_chunks=10
    ):

        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.max_chunks = max_chunks

    def __len__(self):
        return len(self.texts)

    # =====================================================
    # SMART SENTENCE SPLITTING
    # =====================================================

    def _split_text(
        self,
        text,
        max_words=200
    ):

        import re

        sentences = re.split(
            r'(?<=[.!?])\s+',
            str(text)
        )

        chunks = []
        current_chunk = []
        current_len = 0

        for sentence in sentences:

            sentence = sentence.strip()

            if len(sentence) == 0:
                continue

            words = sentence.split()

            # new chunk
            if current_len + len(words) > max_words:

                if len(current_chunk) > 0:
                    chunks.append(
                        " ".join(current_chunk)
                    )

                current_chunk = [sentence]
                current_len = len(words)

            else:

                current_chunk.append(sentence)
                current_len += len(words)

        # last chunk
        if len(current_chunk) > 0:
            chunks.append(
                " ".join(current_chunk)
            )

        return chunks

    # =====================================================
    # TOKENIZE CHUNKS
    # =====================================================

    def _chunk(self, text):

        chunks = self._split_text(text)

        input_ids = []
        attention_masks = []

        # limit number of chunks
        chunks = chunks[:self.max_chunks]

        if len(chunks) == 0:
            chunks = [str(text)]

        for chunk in chunks:

            enc = self.tokenizer(
                chunk,
                max_length=self.max_length,
                truncation=True,
                padding="max_length",
                return_tensors="pt"
            )

            input_ids.append(
                enc["input_ids"].squeeze(0)
            )

            attention_masks.append(
                enc["attention_mask"].squeeze(0)
            )

        input_ids = torch.stack(input_ids)
        attention_masks = torch.stack(attention_masks)

        return input_ids, attention_masks

    # =====================================================
    # GET ITEM
    # =====================================================

    def __getitem__(self, idx):

        input_ids, attention_mask = self._chunk(
            self.texts[idx]
        )

        return {

            "input_ids": input_ids,

            "attention_mask": attention_mask,

            "label": torch.tensor(
                self.labels[idx],
                dtype=torch.float
            )
        }


# =========================================================
# COLLATE FUNCTION
# =========================================================

def chunked_collate_fn(batch, tokenizer=None):

    batch_size = len(batch)

    max_chunks = max(
        x["input_ids"].size(0)
        for x in batch
    )

    seq_len = batch[0]["input_ids"].size(1)

    pad_token_id = (
        tokenizer.pad_token_id
        if tokenizer is not None
        else 0
    )

    input_ids = torch.full(
        (batch_size, max_chunks, seq_len),
        pad_token_id,
        dtype=torch.long
    )

    attention_mask = torch.zeros(
        (batch_size, max_chunks, seq_len),
        dtype=torch.long
    )

    chunk_mask = torch.zeros(
        (batch_size, max_chunks),
        dtype=torch.float
    )

    labels = torch.zeros(
        batch_size,
        dtype=torch.float
    )

    for i, item in enumerate(batch):

        n_chunks = item["input_ids"].size(0)

        input_ids[i, :n_chunks] = item["input_ids"]

        attention_mask[i, :n_chunks] = item["attention_mask"]

        chunk_mask[i, :n_chunks] = 1.0

        labels[i] = item["label"]

    return {

        "input_ids": input_ids,

        "attention_mask": attention_mask,

        "chunk_mask": chunk_mask,

        "labels": labels
    }


# =========================================================
# MODEL
# =========================================================

class ChunkedReadmissionModel(nn.Module):

    def __init__(
        self,
        model_name,
        mlp_head_config=None,
        dropout_rate=0.3
    ):

        super().__init__()

        # =================================================
        # TRANSFORMER ENCODER
        # =================================================

        self.encoder = AutoModel.from_pretrained(
            model_name
        )

        hidden_size = self.encoder.config.hidden_size

        # =================================================
        # CHUNK ATTENTION
        # =================================================

        self.chunk_attention = nn.Sequential(

            nn.Linear(hidden_size, hidden_size),

            nn.GELU(),

            nn.Dropout(dropout_rate),

            nn.Linear(hidden_size, 1)
        )

        self.dropout = nn.Dropout(dropout_rate)

        # =================================================
        # MLP HEAD
        # =================================================

        hidden_dims = []

        # support dict
        if isinstance(mlp_head_config, dict):

            hidden_dims = mlp_head_config.get(
                "hidden_dims",
                []
            )

        # support list
        elif isinstance(mlp_head_config, list):

            for item in mlp_head_config:

                if isinstance(item, dict):

                    dims = item.get(
                        "hidden_dims",
                        []
                    )

                    hidden_dims.extend(dims)

        layers = []

        in_dim = hidden_size

        for hidden_dim in hidden_dims:

            layers.extend([

                nn.Linear(in_dim, hidden_dim),

                nn.ReLU(),

                nn.Dropout(dropout_rate)
            ])

            in_dim = hidden_dim

        # final binary output
        layers.append(
            nn.Linear(in_dim, 1)
        )

        self.classifier = nn.Sequential(
            *layers
        )

    # =====================================================
    # FORWARD
    # =====================================================

    def forward(
        self,
        input_ids,
        attention_mask,
        chunk_mask
    ):

        batch_size, num_chunks, seq_len = input_ids.shape

        # flatten chunks
        input_ids = input_ids.view(
            batch_size * num_chunks,
            seq_len
        )

        attention_mask = attention_mask.view(
            batch_size * num_chunks,
            seq_len
        )

        # =================================================
        # TRANSFORMER
        # =================================================

        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        # CLS token
        cls_embeddings = outputs.last_hidden_state[:, 0, :]

        hidden_size = cls_embeddings.size(-1)

        # reshape back
        cls_embeddings = cls_embeddings.view(
            batch_size,
            num_chunks,
            hidden_size
        )

        # IMPORTANT:
        # removed F.normalize()
        # because it hurts classification performance

        cls_embeddings = self.dropout(
            cls_embeddings
        )

        # =================================================
        # ATTENTION POOLING
        # =================================================

        attention_scores = self.chunk_attention(
            cls_embeddings
        ).squeeze(-1)

        # mask padded chunks
        attention_scores = attention_scores.masked_fill(
            chunk_mask == 0,
            -1e9
        )

        attention_weights = torch.softmax(
            attention_scores,
            dim=1
        ).unsqueeze(-1)

        # weighted document embedding
        document_embedding = torch.sum(
            cls_embeddings * attention_weights,
            dim=1
        )

        document_embedding = self.dropout(
            document_embedding
        )

        # =================================================
        # CLASSIFIER
        # =================================================

        logits = self.classifier(
            document_embedding
        )

        return logits, document_embedding

    # =====================================================
    # PREDICT PROBA
    # =====================================================

    def predict_proba(
        self,
        input_ids,
        attention_mask,
        chunk_mask
    ):

        logits, _ = self.forward(
            input_ids,
            attention_mask,
            chunk_mask
        )

        probs_pos = torch.sigmoid(
            logits
        )

        probs_neg = 1.0 - probs_pos

        probs = torch.cat(
            [probs_neg, probs_pos],
            dim=1
        )

        return probs