import numpy as np   # Библиотека для матриц, векторов и линала
import pandas as pd  # Библиотека для табличек
import time          # Библиотека для времени

from selenium import webdriver 

browser = webdriver.Firefox()
time.sleep(0.3)
browser.implicitly_wait(0.3) 

from bs4 import BeautifulSoup
from lxml import etree
from tqdm.notebook import tqdm

def get_dom(page_link):
    browser.get(page_link)
    html = browser.page_source
    soup = BeautifulSoup(html, 'html.parser') 
    return etree.HTML(str(soup)) 

def get_listpage_links(page_no):
    # page_link = f'https://www.kinopoisk.ru/lists/movies/year--2010-2019/?page={page_no}'
    page_link = f'https://www.kinopoisk.ru/lists/movies/year--2021/?page={page_no}'
    dom = get_dom(page_link) 
    return dom.xpath("body//div[@data-tid='8a6cbb06']/a[@href]/@href")

def get_moviepage_info(movie_link):
    page_link = 'https://www.kinopoisk.ru' + movie_link
    dom = get_dom(page_link)

    elem = dom.xpath("body//div/span//span[@data-tid='939058a8']")
    rating = elem[0].text if elem else ''
    elem = dom.xpath("body//div//h1[@data-tid='f22e0093']/span")
    name = elem[0].text if elem else ''

    features = {}
    elem = dom.xpath("body//div/div[@data-test-id='encyclopedic-table' and @data-tid='bd126b5e']")[0]
    for child in elem.getchildren():
        #print(etree.tostring(child))
        feature = child.xpath('div[1]')[0].text
        ahrefs = child.xpath('div[position()>1]//a[text()] | div[position()>1]//div[text() and not(*)] | div[position()>1]//span[text()]')
        values = [ahr.text for ahr in ahrefs]
        features[feature] = values

    elem = dom.xpath("body//div/p[text() and not(*) and @data-tid='bfd38da2']")
    short_descr = elem[0].text if elem else ''

    elem = dom.xpath("body//div/p[text() and not(*) and @data-tid='bbb11238']")
    descr = ' '.join([x.text for x in elem])

    return (name, rating, short_descr, descr, features)

_df = pd.DataFrame(columns=['id', 'type', 'name', 'rating', 'short_descr', 'descr', 'features'])
for page_number in tqdm(range(1, 912), desc='List pages'):
    try:
        links = get_listpage_links(page_number)
        _df = pd.DataFrame(columns=['id', 'type', 'name', 'rating', 'short_descr', 'descr', 'features'])
        for movie_link in links:
          movie_id = movie_link.split('/')[1:3]
          name, rating, short_descr, descr, features = get_moviepage_info(movie_link)
          data_row = {'id':movie_id[1], 'type':movie_id[0], 'name':name, 'rating':rating, 'short_descr':short_descr, 'descr':descr, 'features': features}
          _df = pd.concat([_df, pd.DataFrame([data_row])], ignore_index=True)
        with open('kinopoisk_2010-2019.csv', 'a') as f:
            _df.to_csv(f, mode='a', header=f.tell()==0, index=False)
    except Exception as err:
        print(f"Unexpected {err=}, {type(err)=}")

from sklearn.preprocessing import MultiLabelBinarizer
mlb = MultiLabelBinarizer()
mlb_result = mlb.fit_transform([str(data.loc[i,'genre_multi']).split(',') for i in range(len(data))])
data = pd.concat([data, pd.DataFrame(mlb_result, columns = list(mlb.classes_))], axis=1)

target_strings = mlb.classes_

from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit
from sklearn.utils import indexable, _safe_indexing
from sklearn.utils.validation import _num_samples
from sklearn.model_selection._split import _validate_shuffle_split
from itertools import chain

def multilabel_train_test_split(*arrays,
                                test_size=None,
                                train_size=None,
                                random_state=None,
                                shuffle=True,
                                stratify=None):
    """
    Train test split for multilabel classification. Uses the algorithm from: 
    'Sechidis K., Tsoumakas G., Vlahavas I. (2011) On the Stratification of Multi-Label Data'.
    """
    if stratify is None:
        return train_test_split(*arrays, test_size=test_size,train_size=train_size,
                                random_state=random_state, stratify=None, shuffle=shuffle)
    
    assert shuffle, "Stratified train/test split is not implemented for shuffle=False"
    
    n_arrays = len(arrays)
    arrays = indexable(*arrays)
    n_samples = _num_samples(arrays[0])
    n_train, n_test = _validate_shuffle_split(
        n_samples, test_size, train_size, default_test_size=0.25
    )
    cv = MultilabelStratifiedShuffleSplit(test_size=n_test, train_size=n_train, random_state=random_state)
    train, test = next(cv.split(X=arrays[0], y=stratify))

    return list(
        chain.from_iterable(
            (_safe_indexing(a, train), _safe_indexing(a, test)) for a in arrays
        )
    )

import nltk
nltk.download('stopwords')
stop_words = nltk.corpus.stopwords.words('russian')
stop_words.append('фильм')
#word_tokenizer = nltk.WordPunctTokenizer()

import re
regex = re.compile(r'[А-Яа-яA-zёЁ-]+')

def words_only(text, regex=regex):
    try:
        return " ".join(regex.findall(text)).lower()
    except:
        return ""

from pymystem3 import Mystem
from string import punctuation

mystem = Mystem() 

#Preprocess function
def preprocess_text(text):
    text = words_only(text)
    tokens = mystem.lemmatize(text.lower())
    tokens = [token for token in tokens if token not in stop_words\
              and token != " " \
              and token.strip() not in punctuation]
    
    text = " ".join(tokens)
    
    return text

pipe = Pipeline([
    ('tfidf', TfidfVectorizer(max_features=1700, min_df=0.0011, max_df=0.35, norm='l2')),
    ('logregr', MultiOutputClassifier(estimator= LogisticRegression(max_iter=10000, class_weight='balanced', multi_class='multinomial', C=0.009, penalty='l2'))),
])

pipe.fit(train_texts, train_y)

pred_y = pipe.predict(test_texts)
print(classification_report(y_true=test_y, y_pred=pred_y, target_names=target_strings))

from catboost import CatBoostClassifier

pipe2 = Pipeline([
    ('tfidf', TfidfVectorizer(max_features=1700, min_df=0.0031, max_df=0.4, norm='l2')),
    ('gboost', MultiOutputClassifier(estimator= CatBoostClassifier(task_type='CPU', verbose=False))),
])

pipe2.fit(train_texts, train_y)

pred_y = pipe2.predict(test_texts)
print(classification_report(y_true=test_y, y_pred=pred_y, target_names=target_strings))

pipe3 = Pipeline([
    ('tfidf', TfidfVectorizer(max_features=1700, min_df=0.0031, max_df=0.4, norm='l2')),
    ('gboost', CatBoostClassifier(task_type='CPU', loss_function='MultiLogloss',  class_names=target_strings, verbose=False)),
])

pipe3.fit(train_texts, train_y)

pred_y = pipe3.predict(test_texts)
print(classification_report(y_true=test_y, y_pred=pred_y, target_names=target_strings))

import matplotlib.pyplot as plt
from wordcloud import WordCloud

def gen_wordcloud(words, importances):
    d = {}
    for i, word in enumerate(words):
        d[word] = abs(importances[i])

    wordcloud = WordCloud()
    wordcloud.generate_from_frequencies(frequencies=d)

    return wordcloud

for idx, x in enumerate(target_strings):
    c1 = pipe['logregr'].estimators_[idx].coef_[0]
    words1 = pipe['tfidf'].get_feature_names_out()
    wc1 = gen_wordcloud(words1, c1)

    c2 = pipe2['gboost'].estimators_[idx].feature_importances_
    words2 = pipe2['tfidf'].get_feature_names_out()
    wc2 = gen_wordcloud(words2, c2)

    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 10))

    ax1.imshow(wc1, interpolation="bilinear")
    ax1.set_title(f'Log regr - {x}')
    ax1.axis('off')
    ax2.imshow(wc2, interpolation="bilinear")
    ax2.set_title(f'Catboost - {x}')
    ax2.axis('off')

    plt.tight_layout()

from transformers import BertTokenizer, AutoTokenizer


selected_model = 'ai-forever/ruBert-base'
# Load the tokenizer.
tokenizer = AutoTokenizer.from_pretrained(selected_model)

from torch.utils.data import TensorDataset


def make_dataset(texts, labels):
    # Tokenize all of the sentences and map the tokens to thier word IDs.
    input_ids = []
    attention_masks = []
    token_type_ids = []

    # For every sentence...
    for sent in texts:
        # `encode_plus` will:
        #   (1) Tokenize the sentence.
        #   (2) Prepend the `[CLS]` token to the start.
        #   (3) Append the `[SEP]` token to the end.
        #   (4) Map tokens to their IDs.
        #   (5) Pad or truncate the sentence to `max_length`
        #   (6) Create attention masks for [PAD] tokens.
        encoded_dict = tokenizer.encode_plus(
                            sent,                      # Sentence to encode.
                            add_special_tokens = True, # Add '[CLS]' and '[SEP]'
                            max_length = 500,          # Pad & truncate all sentences.
                            padding='max_length',
                            return_attention_mask = True,   # Construct attn. masks.
                            return_tensors = 'pt',     # Return pytorch tensors.
                            truncation=True,
                            return_token_type_ids=True
                    )

        # Add the encoded sentence to the list.
        input_ids.append(encoded_dict['input_ids'])

        token_type_ids.append(encoded_dict['token_type_ids'])

        # And its attention mask (simply differentiates padding from non-padding).
        attention_masks.append(encoded_dict['attention_mask'])

    # Convert the lists into tensors.
    input_ids = torch.cat(input_ids, dim=0)
    token_type_ids = torch.cat(token_type_ids, dim=0)
    attention_masks = torch.cat(attention_masks, dim=0)
    labels = torch.tensor(labels.values)
    dataset = TensorDataset(input_ids, token_type_ids, attention_masks, labels)

    return dataset

import transformers

model = transformers.AutoModelForSequenceClassification.from_pretrained(
    selected_model, 
    problem_type="multi_label_classification",
    num_labels = 26, # The number of output labels--2 for binary classification.
                    # You can increase this for multi-class tasks.
    output_attentions = False, # Whether the model returns attentions weights.
    output_hidden_states = False, # Whether the model returns all hidden-states.
)

from torch.utils.tensorboard import SummaryWriter
from sklearn import metrics

def log_metrics(writer, loss, outputs, targets, postfix):
    print(outputs)
    outputs = np.array(outputs)
    predictions = np.zeros(outputs.shape)
    predictions[np.where(outputs >= 0.5)] = 1
    outputs = predictions
    accuracy = metrics.accuracy_score(targets, outputs)
    f1_score_micro = metrics.f1_score(targets, outputs, average='micro')
    f1_score_macro = metrics.f1_score(targets, outputs, average='macro')
    recall_score_micro = metrics.recall_score(targets, outputs, average='micro')
    recall_score_macro = metrics.recall_score(targets, outputs, average='macro')
    precision_score_micro = metrics.precision_score(targets, outputs, average='micro', zero_division=0.0)
    precision_score_macro = metrics.precision_score(targets, outputs, average='macro', zero_division=0.0)

    writer.add_scalar(f'Loss/{postfix}', loss, epoch)
    writer.add_scalar(f'Accuracy/{postfix}', accuracy, epoch)
    writer.add_scalar(f'F1 (Micro)/{postfix}', f1_score_micro, epoch)
    writer.add_scalar(f'F1 (Macro)/{postfix}', f1_score_macro, epoch)
    writer.add_scalar(f'Recall (Micro)/{postfix}', recall_score_micro, epoch)
    writer.add_scalar(f'Recall (Macro)/{postfix}', recall_score_macro, epoch)
    writer.add_scalar(f'Precision (Micro)/{postfix}', precision_score_micro, epoch)
    writer.add_scalar(f'Precision (Macro)/{postfix}', precision_score_macro, epoch)

optimizer = torch.optim.AdamW(model.parameters(),
                  lr = 2e-5, # args.learning_rate - default is 5e-5, our notebook had 2e-5
                  eps = 1e-8 # args.adam_epsilon  - default is 1e-8.
                )

from transformers import get_linear_schedule_with_warmup

# Number of training epochs. The BERT authors recommend between 2 and 4.
# We chose to run for 4, but we'll see later that this may be over-fitting the
# training data.
epochs = model_setup[selected_model]['epochs']

# Total number of training steps is [number of batches] x [number of epochs].
# (Note that this is not the same as the number of training samples).
total_steps = len(train_dataloader) * epochs

# Create the learning rate scheduler.
scheduler = get_linear_schedule_with_warmup(optimizer,
                                            num_warmup_steps = 0, # Default value in run_glue.py
                                            num_training_steps = total_steps)

from tqdm import tqdm

def train(epoch):
    # print(f'Epoch {epoch+1} training started.')
    total_train_loss = 0
    model.train()
    fin_targets=[]
    fin_outputs=[]
    with tqdm(train_dataloader, unit="batch") as tepoch:
        for data in tepoch:
            tepoch.set_description(f"Epoch {epoch+1}")
            ids = data[0].to(device, dtype = torch.long)
            mask = data[2].to(device, dtype = torch.long)
            token_type_ids = data[1].to(device, dtype = torch.long)
            targets = data[3].to(device, dtype = torch.float)

            res = model(ids,
                             token_type_ids=None,
                             attention_mask=mask,
                             labels=targets)
            loss = res['loss']
            logits = res['logits']

            optimizer.zero_grad()
            total_train_loss += loss.item()
            fin_targets.extend(targets.cpu().detach().numpy().tolist())
            fin_outputs.extend(torch.sigmoid(logits).cpu().detach().numpy().tolist())

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            tepoch.set_postfix(loss=loss.item())
    
    return total_train_loss / len(train_dataloader), fin_outputs, fin_targets


def validate(epoch):
    model.eval()
    fin_targets=[]
    fin_outputs=[]
    total_loss = 0.0
    with torch.no_grad():
        for step, data in enumerate(test_dataloader, 0):
            ids = data[0].to(device, dtype = torch.long)
            mask = data[2].to(device, dtype = torch.long)
            token_type_ids = data[1].to(device, dtype = torch.long)
            targets = data[3].to(device, dtype = torch.float)
            res = model(ids,
                             token_type_ids=None,
                             attention_mask=mask,
                             labels=targets)
            loss = res['loss']
            logits = res['logits']
            total_loss += loss.item()
            fin_targets.extend(targets.cpu().detach().numpy().tolist())
            fin_outputs.extend(torch.sigmoid(logits).cpu().detach().numpy().tolist())
    return total_loss/len(test_dataloader), fin_outputs, fin_targets

writer = SummaryWriter(comment= '-' + selected_model.replace('/', '-'))

for epoch in range(epochs):
    avg_train_loss, outputs, targets = train(epoch)

    log_metrics(writer, avg_train_loss, outputs, targets, 'train')
    
    avg_val_loss, outputs, targets = validate(epoch)
    log_metrics(writer, avg_val_loss, outputs, targets, 'val')