from base64 import b64encode
from io import BytesIO

import numpy as np
from flask_login import current_user
from matplotlib.figure import Figure
from matplotlib.ticker import MaxNLocator
from sqlalchemy import select

from .models import MAX_MARK, MIN_MARK, Assistant, Semester, User, db


def html_fig(fig):
    buf = BytesIO()
    fig.savefig(buf, format="png")

    return b64encode(buf.getbuffer()).decode("ascii")


def mark_hist(data, title):
    fig = Figure()
    ax = fig.subplots()
    ax.set_xlim(MIN_MARK - 0.5, MAX_MARK + 0.5)
    ax.set_xticks(np.arange(MAX_MARK + 1))
    # Only integer ticks
    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    ax.set_xlabel("Mark")

    N = data.size
    title += f"\nN = {N}"

    if N > 0:
        ax.hist(
            data,
            bins=np.arange(MAX_MARK + 2) - 0.5,
        )
        title += f" | mean = {round(np.mean(data), 1)}"

    ax.set_title(title)

    return html_fig(fig)


def get_experiment_marks(assistant, attr):
    data = []

    for experiment_mark in assistant.experiment_marks:
        mark = getattr(experiment_mark, attr)
        if mark is not None:
            data.append(mark)

    return np.array(data)


def mark_hists(markType, active_assistants):
    attr = markType.lower() + "_mark"
    mark_type_title_addition = f" | {markType} marks"
    marks = [get_experiment_marks(assistant, attr) for assistant in active_assistants]

    hists = [
        mark_hist(
            data=marks[i],
            title=str(active_assistants[i]) + mark_type_title_addition,
        )
        for i in range(len(marks))
    ]

    hists.append(
        mark_hist(
            data=np.hstack(marks),
            title="All" + mark_type_title_addition,
        )
    )

    return hists


def assistant_marks_analysis(cls):
    active_assistants = db.session.scalars(select(Assistant).join(User).where(User.active is True)).all()

    oral_mark_hists = mark_hists("Oral", active_assistants)
    protocol_mark_hists = mark_hists("Protocol", active_assistants)

    return cls.render(
        "analysis/assistant_marks.jinja.html",
        hist_indices=range(len(oral_mark_hists)),
        oral_mark_hists=oral_mark_hists,
        protocol_mark_hists=protocol_mark_hists,
    )


def get_final_part_marks(part):
    data = []

    for part_student in part.part_students:
        mark = part_student.final_part_mark
        if mark is not None:
            data.append(mark)

    return np.array(data)


def final_part_marks_analysis(cls):
    parts = current_user.active_semester.parts

    active_semester_final_part_marks_hists = [
        mark_hist(
            data=get_final_part_marks(part),
            title=part.str(),
        )
        for part in parts
    ]

    semesters = db.session.scalars(select(Semester)).all()
    mean_final_part_marks = np.array(
        [np.mean(np.hstack([get_final_part_marks(part) for part in semester.parts])) for semester in semesters]
    )

    fig = Figure()
    len_mean_final_part_marks = mean_final_part_marks.size

    ax = fig.subplots()
    x = range(len_mean_final_part_marks)
    ax.plot(
        x,
        mean_final_part_marks,
        marker="d",
    )
    ax.set_xticks(x, [semester.str() for semester in semesters])

    ax.set_xlabel("Semester")
    ax.set_ylabel("Mean final experiment mark")

    ax.set_title("Mean final experiment mark over all semesters")

    mean_final_part_mark_plot = html_fig(fig)

    return cls.render(
        "analysis/final_part_marks.jinja.html",
        active_semester_final_part_marks_hists=active_semester_final_part_marks_hists,
        mean_final_part_mark_plot=mean_final_part_mark_plot,
    )