# %%
import torch
from transformers import (
    BertTokenizer,
    BertForMaskedLM,
    AutoModelForMaskedLM,
    AutoTokenizer,
    BertModel,
)
import numpy as np
import random
from itertools import islice
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW, SGD
from tqdm import tqdm
import os


def index_to_onehot(l, length):
    # l=[1, 5], len=6 -> [0,1,0,0,0,1]
    return [1 if i in l else 0 for i in range(length)]


def get_punctuation_position(tokenized_text, tokenizer):
    # adjust comma_pos and period_pos
    count = 0
    comma_pos = []
    period_pos = []
    punctuation_removed_text = []
    comma_id = tokenizer.convert_tokens_to_ids("、")
    period_id = tokenizer.convert_tokens_to_ids("。")

    for i, c in enumerate(tokenized_text):
        if c == comma_id:
            comma_pos.append(i - count - 1)
            count += 1
        elif c == period_id:
            period_pos.append(i - count - 1)
            count += 1
        else:
            punctuation_removed_text.append(c)

    if len(punctuation_removed_text) < 512:
        punctuation_removed_text += [tokenizer.pad_token_id] * (
            512 - len(punctuation_removed_text)
        )

    return (
        torch.tensor(punctuation_removed_text),
        [
            index_to_onehot(comma_pos, 512),
            index_to_onehot(period_pos, 512),
        ],
    )


# %%
# get_punctuation_position("今日は、いい天気です。")
# # %%
# index_to_onehot([1, 2, 3, 4, 5], 7)
# tokenizer = BertTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-char")
# tokenized_text = tokenizer(
#     "今 日 は 、 い い 天 気 で す 。",
#     max_length=512,
#     padding="max_length",
#     truncation=True,
#     return_tensors="pt",
# )
# inputs, label = get_punctuation_position(tokenized_text["input_ids"][0], tokenizer)
# print(inputs)  # ->tensor([  2, 732,  48,  12,  19,  19, 411, 343,  17,  46,   3,   0,   0,   0, ...])
# print(label)  # -> [[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, ...], # 点の位置(最初に[SOS]が入るため、1つずれる)
#  -> [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, ...]] # 丸の位置


# %%
class PunctuationPositionDataset(torch.utils.data.Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer

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

    def __getitem__(self, idx):
        text = self.data[idx]
        text = " ".join(list(text))
        inputs = self.tokenizer(
            text,
            max_length=512,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        # if idx % 100 == 0:
        #     print(masked_text, label)
        input_ids, label = get_punctuation_position(
            inputs["input_ids"][0], self.tokenizer
        )

        label = torch.tensor(label, dtype=torch.float32).transpose(0, 1)

        return (input_ids, inputs.attention_mask.squeeze(), label.squeeze(), text)


# %%
model_name = "tohoku-nlp/bert-base-japanese-char-v3"
tokenizer = BertTokenizer.from_pretrained(model_name)
base_model = BertModel.from_pretrained(model_name)


# %%
class punctuation_predictor(torch.nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model
        self.dropout = torch.nn.Dropout(0.2)
        self.linear = torch.nn.Linear(768, 2)

    def forward(self, input_ids, attention_mask):
        last_hidden_state = self.base_model(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state
        # get last hidden state token by token and apply linear layer
        return self.linear(self.dropout(last_hidden_state))


model = punctuation_predictor(base_model)
# %%
# a = tokenizer("今 日 は い い 天 気 で す 。",max_length=512,
#             padding="max_length",
#             truncation=True,
#             return_tensors="pt",)
# %%
with open("data/train.txt", "r") as f:
    texts = f.readlines()

dataset = PunctuationPositionDataset(texts, tokenizer)
# %%
data_loader = DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,
    num_workers=8,
)
# %%
# set lr to 5e-5 to base model

optimizer = AdamW(
    [
        {"params": model.base_model.parameters(), "lr": 5e-5},
        {"params": model.linear.parameters(), "lr": 1e-3},
    ],
)

criteria = torch.nn.BCEWithLogitsLoss()
# %%
model.train()
model.to("cuda")
for epoch in range(10):
    epoch_loss = 0.0
    progress_bar = tqdm(data_loader, desc=f"Epoch {epoch+1}")
    for batch in progress_bar:
        input_ids, attention_masks, labels, text = batch
        input_ids = input_ids.to("cuda")
        attention_masks = attention_masks.to("cuda")
        labels = labels.to("cuda")

        outputs = model(input_ids=input_ids, attention_mask=attention_masks)
        loss = criteria(outputs, labels)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        epoch_loss += loss.item()
        progress_bar.set_postfix({"loss": epoch_loss / len(data_loader)})
# %%
torch.save(model.state_dict(), "weight/punctuation_position_model.pth")