import csv

from django.shortcuts import render
from django.http import HttpResponse


from .models import EvalSession, EvalAnswer, EvalConfig, QA

def eval_download_results(request, session_id):
    session = EvalSession.objects.get(id=session_id)
    eval_answers = EvalAnswer.objects.filter(eval_session=session).order_by('question__id')
    response = HttpResponse(content_type='text/csv')

    writer = csv.writer(response)
    writer.writerow(['Question ID', 'Question', 'Correct Answer', 'Correct Answer Index', 'Instruction', 'Assistant answer', 'Is Correct'])

    for eval_answer in eval_answers:
        writer.writerow([
            eval_answer.question.id, 
            eval_answer.question.question, 
            eval_answer.question.correct_answer, 
            eval_answer.question.correct_answer_idx, 
            eval_answer.instruction, 
            eval_answer.assistant_answer, 
            eval_answer.is_correct])

    response['Content-Disposition'] = f'attachment; filename="{session.name}.csv"'
    return response

def config_download_results(request, config_id):
    response = HttpResponse(content_type='text/csv')
    config = EvalConfig.objects.get(id=config_id)
    sessions = EvalSession.objects.filter(config=config).filter(is_active=True)
    
    extra_info_fields = []
    questions = dict()
    for qa in QA.objects.filter(dataset=config.dataset):
        questions[qa.id] = qa
        if not extra_info_fields and qa.extra_info:
            extra_info_fields = qa.extra_info.keys()
    

    models_data = {}
    for session in sessions:
        eval_answers = EvalAnswer.objects.filter(eval_session=session).order_by('question__id')
        data = {}
        for eval_answer in eval_answers:
            data[eval_answer.question.id] = 1 if eval_answer.is_correct else 0
        models_data[f"{session.llm_model.name} [{session.name}]"] = data

    writer = csv.writer(response)
    writer.writerow(['Question ID'] + ['Question Text'] +
                    [x for x in extra_info_fields] +
                    [model_name for model_name in models_data.keys()])
    
    for question_id in models_data[list(models_data.keys())[0]].keys():
            if questions[question_id].xid:
                qid = questions[question_id].xid
            else:
                qid = question_id
            row = [qid] + [questions[question_id].question] +\
            [questions[question_id].extra_info.get(field, '') for field in extra_info_fields] +\
            [data.get(question_id, '') for data in models_data.values()]
            writer.writerow(row)
    row = ['Accuracy'] + [''] + ['' for _ in extra_info_fields] + [f"{sum(data.values()) / len(data) * 100:.2f}%" for data in models_data.values()]
    writer.writerow(row)
    response['Content-Disposition'] = f'attachment; filename="{config.name}.csv"'
    return response