import torch
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    HfArgumentParser
)

from train_classifier import ClassifierModelArguments
from shared import CATEGORIES, DatasetArguments
from tqdm import tqdm

from preprocess import get_words, clean_text
from segment import extract_segment
import os
import json
import numpy as np


def softmax(_outputs):
    maxes = np.max(_outputs, axis=-1, keepdims=True)
    shifted_exp = np.exp(_outputs - maxes)
    return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True)


def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ClassifierModelArguments, DatasetArguments))
    model_args, dataset_args = parser.parse_args_into_dataclasses()

    model = AutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path)
    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)

    processed_db_path = os.path.join(
        dataset_args.data_dir, dataset_args.processed_database)
    with open(processed_db_path) as fp:
        data = json.load(fp)

    mapped_categories = {
        str(v).lower(): k for k, v in enumerate(CATEGORIES)
    }

    for video_id, segments in tqdm(data.items()):

        words = get_words(video_id)

        if not words:
            continue  # No/empty transcript for video_id

        valid_segments = []
        texts = []
        for segment in segments:
            segment_words = extract_segment(
                words, segment['start'], segment['end'])
            text = clean_text(' '.join(x['text'] for x in segment_words))

            duration = segment['end'] - segment['start']
            wps = len(segment_words)/duration if duration > 0 else 0
            if wps < 1.5:
                continue

            # Do not worry about those that are locked or have enough votes
            if segment['locked']:  # or segment['votes'] > 5:
                continue

            texts.append(text)
            valid_segments.append(segment)

        if not texts:
            continue  # No valid segments

        model_inputs = tokenizer(
            texts, return_tensors='pt', padding=True, truncation=True)

        with torch.no_grad():
            model_outputs = model(**model_inputs)
            outputs = list(map(lambda x: x.numpy(), model_outputs['logits']))

        scores = softmax(outputs)

        for segment, text, score in zip(valid_segments, texts, scores):
            predicted_index = score.argmax().item()

            if predicted_index == mapped_categories[segment['category']]:
                continue  # Ignore correct segments

            a = {k: round(float(score[i]), 3)
                 for i, k in enumerate(CATEGORIES)}

            del segment['submission_time']
            segment.update({
                'predicted': str(CATEGORIES[predicted_index]).lower(),
                'text': text,
                'scores': a
            })

            print(json.dumps(segment))


if __name__ == "__main__":
    main()