import pickle
import random
from datetime import datetime, timedelta

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

import util


def sentences_published_on_date(articles, date, tolerance_days=2, num_first_sentences=2):

    date_to_pub = pickle.load(open('data/ghalandari_date_to_pub.pkl', 'rb'))
    return date_to_pub[date]

    result = []
    for article in articles:
        pub_date = datetime.strptime(article['pub_date'], '%Y-%m-%d')
        start_date = datetime.strptime(date, '%Y-%m-%d')
        end_date = start_date + timedelta(days=tolerance_days)
        if start_date <= pub_date <= end_date:
            result += [sentence['text'] for sentence in article['sentences'][:num_first_sentences] if len(sentence['mentioned_dates']) == 0]
    return result


def sentences_mentioning_date(articles, date):
    date_to_ment = pickle.load(open('data/ghalandari_date_to_ment.pkl', 'rb'))
    return date_to_ment[date]
    result = []
    for article in articles:
        result += [sentence['text'] for sentence in article['sentences'] if date in sentence['mentioned_dates']]
    return result


def candidate_sentences(articles, date, keywords):
    pub_sentences = sentences_published_on_date(articles, date)
    ment_sentences = sentences_mentioning_date(articles, date)

    # handle the cases where there are no sentences to process
    if len(pub_sentences) == 0 and len(ment_sentences) == 0:
        return []
    if len(pub_sentences) == 0:
        pub_sentences = ment_sentences
    if len(ment_sentences) == 0:
        ment_sentences = pub_sentences

    # train TFIDF vectorizer on all sentences (not just the ones for this date)
    all_sentences = [sentence['text'] for article in articles for sentence in article['sentences']]
    vectorizer = TfidfVectorizer(stop_words='english', lowercase=True)
    vectorizer.fit(all_sentences)

    # get TFIDF vectors for sentences published on date and sentences mentioning date
    pub_vectors = vectorizer.transform(pub_sentences)
    ment_vectors = vectorizer.transform(ment_sentences)

    # average them to a single vector each
    pub_mean_vector = pub_vectors.mean(0).tolist()[0]
    ment_mean_vector = ment_vectors.mean(0).tolist()[0]

    # get date vector (weighted average of publication vector and mention vector)
    pub_weight = 1 / len(pub_sentences)
    ment_weight = 1 / len(ment_sentences)
    date_vector = [(pub_weight * pub_mean_vector[i] + ment_weight * ment_mean_vector[i]
                    if pub_mean_vector[i] > 0 and ment_mean_vector[i] > 0
                    else 0)
                   for i in range(len(pub_mean_vector))]

    # build candidate pool from only those sentences that contain a keyword
    candidate_sentence_pool = list(set(pub_sentences + ment_sentences))
    candidate_sentence_pool = list(filter(lambda s: util.contains_any(s, keywords), candidate_sentence_pool))

    # if the keyword filtering leaves no candidate sentences, abort
    if len(candidate_sentence_pool) == 0:
        return [], date_vector

    # select candidate sentences as those most similar to the date vector
    sent_vectors = vectorizer.transform(candidate_sentence_pool).toarray().tolist()
    similarities = cosine_similarity([date_vector], sent_vectors)[0]

    sorted_indices = sorted(list(range(len(candidate_sentence_pool))), key=lambda i: similarities[i], reverse=True)
    sorted_sentences = [{'text': candidate_sentence_pool[i], 'vector': sent_vectors[i], 'sim': similarities[i]} for i in sorted_indices]
    sorted_scores = [similarities[i] for i in sorted_indices]

    # only consider sentences above the "knee point"
    cutoff_index = knee_point(sorted_scores)

    return sorted_sentences[:cutoff_index + 1], date_vector


def knee_point(values):

    if len(values) <= 1:
        return 0

    # get coordinates of all the points
    n_points = len(values)
    all_coords = np.vstack((range(n_points), values)).T

    # get the first point
    first_point = all_coords[0]

    # get vector between first and last point - this is the line
    line_vec = all_coords[-1] - all_coords[0]
    line_vec_norm = line_vec / np.sqrt(np.sum(line_vec ** 2))
    vec_from_first = all_coords - first_point
    scalar_prod = np.sum(vec_from_first * np.tile(line_vec_norm, (n_points, 1)), axis=1)
    vec_from_first_parallel = np.outer(scalar_prod, line_vec_norm)
    vec_to_line = vec_from_first - vec_from_first_parallel

    # distance to line is the norm of vec_to_line
    dist_to_line = np.sqrt(np.sum(vec_to_line ** 2, axis=1))

    # knee/elbow is the point with max distance value
    best_idx = np.argmax(dist_to_line)
    return best_idx
