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)

[(24, 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (25, 'numpy.array', 'np.array', (['texts'], {}), True, 'import numpy as np\n'), (26, 'numpy.array', 'np.array', (['labels'], {}), True, 'import numpy as np\n'), (28, 'sklearn.model_selection.StratifiedKFold', 'StratifiedKFold', ([], {'n_splits': 'splits'}), False, 'from sklearn.model_selection import StratifiedKFold\n'), (86, 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (87, 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (91, 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'labels'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (96, 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (100, 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), False, 'from tqdm import tqdm\n'), (114, 'os.makedirs', 'os.makedirs', (['path'], {}), False, 'import os\n'), (129, 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (130, 'numpy.array', 'np.array', (['texts'], {}), True, 'import numpy as np\n'), (131, 'numpy.array', 'np.array', (['labels_A'], {}), True, 'import numpy as np\n'), (132, 'numpy.array', 'np.array', (['labels_B'], {}), True, 'import numpy as np\n'), (134, 'sklearn.model_selection.StratifiedKFold', 'StratifiedKFold', ([], {'n_splits': 'splits'}), False, 'from sklearn.model_selection import StratifiedKFold\n'), (204, 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (205, 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.tokenizer_name'], {}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (209, 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'labels'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (214, 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (218, 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), False, 'from tqdm import tqdm\n'), (232, 'os.makedirs', 'os.makedirs', (['path'], {}), False, 'import os\n'), (18, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (18, 'torch.device', 'torch.device', (['"""cuda"""'], {}), False, 'import torch\n'), (18, 'torch.device', 'torch.device', (['"""cpu"""'], {}), False, 'import torch\n'), (34, 'transformers.AutoModelForSequenceClassification.from_pretrained', 'AutoModelForSequenceClassification.from_pretrained', (['self.model_name'], {'config': 'config'}), False, 'from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig\n'), (43, 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_train', 'y_train'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (44, 'fast_fine_tuna.dataset.MainDataset', 'MainDataset', (['tokenized_test', 'y_test'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (49, 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (53, 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), False, 'from tqdm import tqdm\n'), (69, 'torch.utils.data.DataLoader', 'DataLoader', (['test_dataset'], {'batch_size': 'batch_size'}), False, 'from torch.utils.data import DataLoader\n'), (124, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (124, 'torch.device', 'torch.device', (['"""cuda"""'], {}), False, 'import torch\n'), (124, 'torch.device', 'torch.device', (['"""cpu"""'], {}), False, 'import torch\n'), (152, 'fast_fine_tuna.dataset.MainDatasetDouble', 'MainDatasetDouble', (['tokenized_train', 'y_A_train', 'y_B_train'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (153, 'fast_fine_tuna.dataset.MainDatasetDouble', 'MainDatasetDouble', (['tokenized_test', 'y_A_test', 'y_B_test'], {}), False, 'from fast_fine_tuna.dataset import MainDatasetDouble, MainDataset\n'), (158, 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (162, 'tqdm.tqdm', 'tqdm', ([], {'total': 'epochs', 'position': '(0)', 'leave': '(True)'}), False, 'from tqdm import tqdm\n'), (185, 'torch.utils.data.DataLoader', 'DataLoader', (['test_dataset'], {'batch_size': 'batch_size'}), False, 'from torch.utils.data import DataLoader\n'), (71, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (188, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (173, 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), False, 'from torch import nn\n'), (77, 'torch.argmax', 'torch.argmax', (["outputs['logits']"], {'axis': '(1)'}), False, 'import torch\n'), (194, 'torch.argmax', 'torch.argmax', (['outputs[0]'], {'axis': '(1)'}), False, 'import torch\n'), (195, 'torch.argmax', 'torch.argmax', (['outputs[1]'], {'axis': '(1)'}), False, 'import torch\n')]