from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
from sklearn.model_selection import StratifiedKFold
import numpy as np
import torch
from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset
from transformers import AdamW
from torch.utils.data import DataLoader
import os
from tqdm import tqdm
from fast_fine_tuna.models import MiniModel
from torch import nn

class FastFineTuna:

    def __init__(self, model_name, tokenizer_name):
        self.model_name = model_name
        self.tokenizer_name = tokenizer_name
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    def cross_validate_fit(self, texts, labels, splits=5, epochs=5, batch_size=16, learning_rate=5e-5):

        config = AutoConfig.from_pretrained(self.model_name, num_labels=len(set(labels)),
                                            finetuning_task="custom")
        tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
        texts = np.array(texts)
        labels = np.array(labels)

        skf = StratifiedKFold(n_splits=splits)

        original = []
        predicted = []

        for train_index, test_index in skf.split(texts, labels):
            model = AutoModelForSequenceClassification.from_pretrained(self.model_name, config=config)

            X_train, X_test = texts[train_index].tolist(), texts[test_index].tolist()
            y_train, y_test = labels[train_index].tolist(), labels[test_index].tolist()

            # not the smartest way to do this, but faster to code up
            tokenized_train = tokenizer(X_train, truncation=True, padding=True)
            tokenized_test = tokenizer(X_test, truncation=True, padding=True)

            train_dataset = MainDataset(tokenized_train, y_train)
            test_dataset = MainDataset(tokenized_test, y_test)

            model.to(self.device)
            model.train()

            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

            optim = AdamW(model.parameters(), lr=learning_rate)

            pbar = tqdm(total=epochs, position=0, leave=True)
            for epoch in range(epochs):
                pbar.update(1)
                for batch in train_loader:
                    optim.zero_grad()
                    input_ids = batch['input_ids'].to(self.device)
                    attention_mask = batch['attention_mask'].to(self.device)
                    lab = batch['labels'].to(self.device)
                    outputs = model(input_ids, attention_mask=attention_mask, labels=lab)

                    loss = outputs[0]
                    loss.backward()
                    optim.step()
            pbar.close()
            model.eval()

            loader = DataLoader(test_dataset, batch_size=batch_size)
            original.extend(y_test)
            with torch.no_grad():
                for batch in loader:
                    input_ids = batch['input_ids'].to(self.device)
                    attention_mask = batch['attention_mask'].to(self.device)

                    outputs = model(input_ids, attention_mask=attention_mask)
                    predicted.extend(torch.argmax(outputs["logits"], axis=1).cpu().numpy().tolist())
            del model
        return original, predicted

    def train_and_save(self, texts, labels, path, epochs=5, batch_size=16, learning_rate=5e-5):

        config = AutoConfig.from_pretrained(self.model_name, num_labels=len(set(labels)),
                                            finetuning_task="custom")

        model = AutoModelForSequenceClassification.from_pretrained(self.model_name, config=config)
        tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)

        tokenized_train = tokenizer(texts, truncation=True, padding=True)

        train_dataset = MainDataset(tokenized_train, labels)

        model.to(self.device)
        model.train()

        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

        optim = AdamW(model.parameters(), lr=learning_rate)

        pbar = tqdm(total=epochs, position=0, leave=True)
        for epoch in range(epochs):
            pbar.update(1)
            for batch in train_loader:
                optim.zero_grad()
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                lab = batch['labels'].to(self.device)
                outputs = model(input_ids, attention_mask=attention_mask, labels=lab)
                loss = outputs[0]
                loss.backward()
                optim.step()
        pbar.close()

        os.makedirs(path)
        model.save_pretrained(path)
        tokenizer.save_pretrained(path)


class DoubleFastFineTuna:

    def __init__(self, model_name, tokenizer_name):
        self.model_name = model_name
        self.tokenizer_name = tokenizer_name
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    def cross_validate_fit(self, texts, labels_A, labels_B, splits=5, epochs=5, batch_size=16, learning_rate=5e-5,
                           ):

        tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
        texts = np.array(texts)
        labels_A = np.array(labels_A)
        labels_B = np.array(labels_B)

        skf = StratifiedKFold(n_splits=splits)

        original_A = []
        original_B = []
        predicted_A = []
        predicted_B = []

        for train_index, test_index in skf.split(texts, labels_A, labels_B):
            model = MiniModel(self.model_name, len(set(labels_A)), len(set(labels_B)))

            X_train, X_test = texts[train_index].tolist(), texts[test_index].tolist()
            y_A_train, y_A_test = labels_A[train_index].tolist(), labels_A[test_index].tolist()
            y_B_train, y_B_test = labels_B[train_index].tolist(), labels_B[test_index].tolist()

            # not the smartest way to do this, but faster to code up
            tokenized_train = tokenizer(X_train, truncation=True, padding=True)
            tokenized_test = tokenizer(X_test, truncation=True, padding=True)

            train_dataset = MainDatasetDouble(tokenized_train, y_A_train, y_B_train)
            test_dataset = MainDatasetDouble(tokenized_test, y_A_test, y_B_test)

            model.to(self.device)
            model.train()

            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

            optim = AdamW(model.parameters(), lr=learning_rate)

            pbar = tqdm(total=epochs, position=0, leave=True)
            for epoch in range(epochs):
                pbar.update(1)
                for batch in train_loader:
                    optim.zero_grad()
                    input_ids = batch['input_ids'].to(self.device)
                    attention_mask = batch['attention_mask'].to(self.device)
                    lab_A = batch['labels_A'].to(self.device)
                    lab_B = batch['labels_B'].to(self.device)

                    outputs = model(input_ids, attention_mask=attention_mask)
                    loss = nn.CrossEntropyLoss()

                    loss_A = loss(outputs[0], lab_A)
                    loss_B = loss(outputs[1], lab_B)

                    loss = loss_A + loss_B

                    loss.backward()
                    optim.step()
            pbar.close()
            model.eval()

            loader = DataLoader(test_dataset, batch_size=batch_size)
            original_A.extend(y_A_test)
            original_B.extend(y_B_test)
            with torch.no_grad():
                for batch in loader:
                    input_ids = batch['input_ids'].to(self.device)
                    attention_mask = batch['attention_mask'].to(self.device)

                    outputs = model(input_ids, attention_mask=attention_mask)
                    predicted_A.extend(torch.argmax(outputs[0], axis=1).cpu().numpy().tolist())
                    predicted_B.extend(torch.argmax(outputs[1], axis=1).cpu().numpy().tolist())
            del model
        return original_A, original_B, predicted_A, predicted_B

    def train_and_save(self, texts, labels, path, epochs=5, batch_size=16, learning_rate=5e-5):

        config = AutoConfig.from_pretrained(self.model_name, num_labels=len(set(labels)),
                                            finetuning_task="custom")

        model = AutoModelForSequenceClassification.from_pretrained(self.model_name, config=config)
        tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)

        tokenized_train = tokenizer(texts, truncation=True, padding=True)

        train_dataset = MainDataset(tokenized_train, labels)

        model.to(self.device)
        model.train()

        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

        optim = AdamW(model.parameters(), lr=learning_rate)

        pbar = tqdm(total=epochs, position=0, leave=True)
        for epoch in range(epochs):
            pbar.update(1)
            for batch in train_loader:
                optim.zero_grad()
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                lab = batch['labels'].to(self.device)
                outputs = model(input_ids, attention_mask=attention_mask, labels=lab)
                loss = outputs[0]
                loss.backward()
                optim.step()
        pbar.close()

        os.makedirs(path)
        model.save_pretrained(path)
        tokenizer.save_pretrained(path)

[((955, 1005), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), '(self.tokenizer_name)\n', (984, 1005), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((1022, 1037), 'numpy.array', 'np.array', (['texts'], {}), '(texts)\n', (1030, 1037), True, 'import numpy as np\n'), ((1055, 1071), 'numpy.array', 'np.array', (['labels'], {}), '(labels)\n', (1063, 1071), True, 'import numpy as np\n'), ((1087, 1119), 'sklearn.model_selection.StratifiedKFold', 'StratifiedKFold', ([], {'n_splits': 'splits'}), '(n_splits=splits)\n', (1102, 1119), False, 'from sklearn.model_selection import StratifiedKFold\n'), ((3583, 3670), 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), '(self.model_name, config=\n    config)\n', (3633, 3670), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((3686, 3736), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), '(self.tokenizer_name)\n', (3715, 3736), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((3837, 3873), 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'labels'], {}), '(tokenized_train, labels)\n', (3848, 3873), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((3951, 4013), 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), '(train_dataset, batch_size=batch_size, shuffle=True)\n', (3961, 4013), False, 'from torch.utils.data import DataLoader\n'), ((4091, 4133), 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), '(total=epochs, position=0, leave=True)\n', (4095, 4133), False, 'from tqdm import tqdm\n'), ((4671, 4688), 'os.makedirs', 'os.makedirs', (['path'], {}), '(path)\n', (4682, 4688), False, 'import os\n'), ((5192, 5242), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), '(self.tokenizer_name)\n', (5221, 5242), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((5259, 5274), 'numpy.array', 'np.array', (['texts'], {}), '(texts)\n', (5267, 5274), True, 'import numpy as np\n'), ((5294, 5312), 'numpy.array', 'np.array', (['labels_A'], {}), '(labels_A)\n', (5302, 5312), True, 'import numpy as np\n'), ((5332, 5350), 'numpy.array', 'np.array', (['labels_B'], {}), '(labels_B)\n', (5340, 5350), True, 'import numpy as np\n'), ((5366, 5398), 'sklearn.model_selection.StratifiedKFold', 'StratifiedKFold', ([], {'n_splits': 'splits'}), '(n_splits=splits)\n', (5381, 5398), False, 'from sklearn.model_selection import StratifiedKFold\n'), ((8433, 8520), 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), '(self.model_name, config=\n    config)\n', (8483, 8520), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((8536, 8586), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), '(self.tokenizer_name)\n', (8565, 8586), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((8687, 8723), 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'labels'], {}), '(tokenized_train, labels)\n', (8698, 8723), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((8801, 8863), 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), '(train_dataset, batch_size=batch_size, shuffle=True)\n', (8811, 8863), False, 'from torch.utils.data import DataLoader\n'), ((8941, 8983), 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), '(total=epochs, position=0, leave=True)\n', (8945, 8983), False, 'from tqdm import tqdm\n'), ((9521, 9538), 'os.makedirs', 'os.makedirs', (['path'], {}), '(path)\n', (9532, 9538), False, 'import os\n'), ((618, 643), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (641, 643), False, 'import torch\n'), ((594, 614), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (606, 614), False, 'import torch\n'), ((649, 668), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (661, 668), False, 'import torch\n'), ((1252, 1339), 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), '(self.model_name, config=\n    config)\n', (1302, 1339), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), ((1767, 1804), 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'y_train'], {}), '(tokenized_train, y_train)\n', (1778, 1804), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((1832, 1867), 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_test', 'y_test'], {}), '(tokenized_test, y_test)\n', (1843, 1867), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((1957, 2019), 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), '(train_dataset, batch_size=batch_size, shuffle=True)\n', (1967, 2019), False, 'from torch.utils.data import DataLoader\n'), ((2105, 2147), 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), '(total=epochs, position=0, leave=True)\n', (2109, 2147), False, 'from tqdm import tqdm\n'), ((2772, 2819), 'torch.utils.data.DataLoader', 'DataLoader', (['test_dataset'], {'batch_size': 'batch_size'}), '(test_dataset, batch_size=batch_size)\n', (2782, 2819), False, 'from torch.utils.data import DataLoader\n'), ((4974, 4999), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (4997, 4999), False, 'import torch\n'), ((4950, 4970), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (4962, 4970), False, 'import torch\n'), ((5005, 5024), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (5017, 5024), False, 'import torch\n'), ((6199, 6255), 'fast_fine_tuna.dataset.MainDatasetDouble', 'MainDatasetDouble', (['tokenized_train', 'y_A_train', 'y_B_train'], {}), '(tokenized_train, y_A_train, y_B_train)\n', (6216, 6255), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((6283, 6336), 'fast_fine_tuna.dataset.MainDatasetDouble', 'MainDatasetDouble', (['tokenized_test', 'y_A_test', 'y_B_test'], {}), '(tokenized_test, y_A_test, y_B_test)\n', (6300, 6336), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), ((6426, 6488), 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), '(train_dataset, batch_size=batch_size, shuffle=True)\n', (6436, 6488), False, 'from torch.utils.data import DataLoader\n'), ((6574, 6616), 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), '(total=epochs, position=0, leave=True)\n', (6578, 6616), False, 'from tqdm import tqdm\n'), ((7458, 7505), 'torch.utils.data.DataLoader', 'DataLoader', (['test_dataset'], {'batch_size': 'batch_size'}), '(test_dataset, batch_size=batch_size)\n', (7468, 7505), False, 'from torch.utils.data import DataLoader\n'), ((2873, 2888), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (2886, 2888), False, 'import torch\n'), ((7603, 7618), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (7616, 7618), False, 'import torch\n'), ((7143, 7164), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), '()\n', (7162, 7164), False, 'from torch import nn\n'), ((3187, 3226), 'torch.argmax', 'torch.argmax', (["outputs['logits']"], {'axis': '(1)'}), "(outputs['logits'], axis=1)\n", (3199, 3226), False, 'import torch\n'), ((7919, 7951), 'torch.argmax', 'torch.argmax', (['outputs[0]'], {'axis': '(1)'}), '(outputs[0], axis=1)\n', (7931, 7951), False, 'import torch\n'), ((8015, 8047), 'torch.argmax', 'torch.argmax', (['outputs[1]'], {'axis': '(1)'}), '(outputs[1], axis=1)\n', (8027, 8047), False, 'import torch\n')]