import logging
import os
import random
import sys

import numpy as np
import pandas as pd
import torch
from sklearn import model_selection


class AverageMeter:
    """
    Computes and stores the average and current value
    """
    def __init__(self):
        self.val = None
        self.avg = None
        self.sum = None
        self.count = None
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


class EarlyStopping:
    def __init__(self, patience=7, mode="max", delta=0.001):
        self.patience = patience
        self.counter = 0
        self.mode = mode
        self.best_score = None
        self.early_stop = False
        self.delta = delta
        if self.mode == "min":
            self.val_score = np.Inf
        else:
            self.val_score = -np.Inf

    def __call__(self, epoch_score, model, model_path):

        if self.mode == "min":
            score = -1.0 * epoch_score
        else:
            score = np.copy(epoch_score)

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(epoch_score, model, model_path)
        elif score < self.best_score + self.delta:
            self.counter += 1
            print('EarlyStopping counter: {} out of {}'.format(self.counter, self.patience))
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(epoch_score, model, model_path)
            self.counter = 0
        return self.best_score

    def save_checkpoint(self, epoch_score, model, model_path):
        if epoch_score not in [-np.inf, np.inf, -np.nan, np.nan]:
            print('Validation score improved ({} --> {}). Saving model!'.format(self.val_score, epoch_score))
            torch.save(model.state_dict(), model_path)
        self.val_score = epoch_score


def jaccard(str1, str2):
    a = set(str1.lower().split())
    b = set(str2.lower().split())
    c = a.intersection(b)
    return float(len(c)) / (len(a) + len(b) - len(c))


def bert_output_to_string(original_tweet, sentiment_val, idx_start, idx_end, offsets):
    if idx_end < idx_start:
        idx_end = idx_start

    filtered_output = ""
    for ix in range(idx_start, idx_end + 1):
        filtered_output += original_tweet[offsets[ix][0]: offsets[ix][1]]
        if (ix + 1) < len(offsets) and offsets[ix][1] < offsets[ix + 1][0]:
            filtered_output += " "

    if sentiment_val == "neutral" or len(original_tweet.split()) < 2:
        filtered_output = original_tweet

    return filtered_output


def create_folds(config):
    df = pd.read_csv(config.TRAIN_FILE)
    df = df.dropna().reset_index(drop=True)
    df["kfold"] = -1

    df = df.sample(frac=1).reset_index(drop=True)

    kf = model_selection.StratifiedKFold(n_splits=5)

    for fold, (trn_, val_) in enumerate(kf.split(X=df, y=df.sentiment.values)):
        print(len(trn_), len(val_))
        df.loc[val_, 'kfold'] = fold

    df.to_csv(config.TRAINING_FILE, index=False)


def device():
    device_num = cuda_num()
    return torch.device(f"cuda:{device_num}" if torch.cuda.is_available() else "cpu")


def cuda_num():
    return os.environ.get('CUDA_NUM', 0)


def get_logger():
    root = logging.getLogger()
    root.setLevel(logging.DEBUG)

    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    root.addHandler(handler)
    return root


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def batch_jaccard(start_positions, end_positions, batch):
    orig_tweet = batch['orig_tweet']
    orig_selected = batch['orig_selected']
    sentiment = batch['sentiment']
    offsets = batch['offsets']

    jaccard_scores = []
    filtered_outputs_ = []

    wrong_order_count = 0

    for px, tweet in enumerate(orig_tweet):
        selected_tweet = orig_selected[px]
        tweet_sentiment = sentiment[px]

        if end_positions[px] < start_positions[px]:
            wrong_order_count += 1

        f_output = bert_output_to_string(tweet, tweet_sentiment, start_positions[px], end_positions[px], offsets[px])

        jaccard_score = jaccard(selected_tweet.strip(), f_output.strip())
        jaccard_scores.append(jaccard_score)
        filtered_outputs_.append(f_output)

    return jaccard_scores, filtered_outputs_, wrong_order_count
