import csv
from collections import defaultdict
import numpy as np
from statistics import mean

#def avg_scores(scores, num_queries_per_video, num_videos):
#    avg_scores = []
#    sum_scores = [0.0] * int(num_videos)
#    for q_idx in range(len(scores)):
#        if q_idx % int(num_queries_per_video) == 0 and q_idx != 0:
#            avg_scores.append([score / num_queries_per_video for score in sum_scores])
#            sum_scores = [0.0] * int(num_videos)
#        for v_idx in range(int(num_videos)):
#            sum_scores[v_idx] += scores[q_idx][v_idx]

#    avg_scores.append([score / num_queries_per_video for score in sum_scores])
#    print("Average scores by aggregating queries in each video:")
#    print(avg_scores)


def avg_scores(scores, num_queries_per_video, num_videos):
    output_csv = "sims.csv"
    video_list = "data/MSRVTT/high-quality/structured-symlinks/test_list_full.txt"

    with open(video_list, "r") as f_in:
        lines = f_in.readlines()

    videos = [line.strip() for line in lines]

    avg_scores = []
    sum_scores = [0.0] * int(num_videos)
    for q_idx in range(len(scores)):
        if q_idx % int(num_queries_per_video) == 0 and q_idx != 0:
            avg_scores.append(['%.2f' % (score / num_queries_per_video) for score in sum_scores])
            sum_scores = [0.0] * int(num_videos)
        for v_idx in range(int(num_videos)):
            sum_scores[v_idx] += scores[q_idx][v_idx]

    avg_scores.append(['%.2f' % (score / num_queries_per_video) for score in sum_scores])
    print("Average scores by aggregating queries in each video:")
    print(avg_scores)

    with open(output_csv, 'w', newline='') as csvfile:
        csv_writer = csv.writer(csvfile, dialect='excel')

        csv_writer.writerow(['Videos'] + videos)
        for i, avg_score in enumerate(avg_scores):
            csv_writer.writerow([videos[i]] + avg_score)


def text_rank(csv_file, output_file):
    with open(csv_file, 'r', newline='') as csvfile:
        csv_reader = csv.reader(csvfile, dialect='excel')
        headers = next(csv_reader)
        video_list = headers[1:]
        video2text = defaultdict(list)
        for row in csv_reader:
            for h, v in zip(video_list, row[1:]):
                video2text[h].append(float(v))

    ranked_text = {}
    for video in video_list:
        text_scores = video2text[video]
        ranks = np.argsort(text_scores)[::-1]
        ranked_text[video] = [video_list[i] for i in ranks]

    #print(ranked_text)
    with open(output_file, 'w', newline='') as csvfile:
        csv_writer = csv.writer(csvfile, dialect='excel')
        for video, ranks in ranked_text.items():
            csv_writer.writerow([video+": "] + ranks[:50])


def overall_sim(csv_file):
    sim_scores = []
    with open(csv_file, 'r', newline='') as csvfile:
        csv_reader = csv.reader(csvfile, dialect='excel')
        headers = next(csv_reader)
        for i, row in enumerate(csv_reader):
            scores = row[1:]
            sim_scores.append(float(scores[i]))

    print("Overall similarity scores for videos: ", '%.3f' % (mean(sim_scores)))



if __name__ == "__main__":
    #text_rank('sims.csv', 'text_rank.csv')
    overall_sim('/vulcanscratch/lzhao/repos/collaborative-experts/sims_person_pos.csv')