"""Data about clusters generated by the agreement module"""

ARGPARSE_PARAMS = ("-D", "--clusterdata", "Data about clusters generated by "
                   "the agreement module.")
MODULE_NAME = "clusterdata"

import os
import numpy as np
import pandas as pd

NUM_ITEMS = 10


def num_hits(ranking, hits):
    cnt = 0
    for item_id in ranking:
        if item_id in hits:
            cnt += 1
    return cnt


def gen_frame(dsr, agr_res):
    user_ratings = dsr.raw.groupby("user_id")
    user_cnt = user_ratings.count()['rating']
    user_mean = user_ratings.mean()['rating']

    tuples = []

    alg_names = dsr.algs

    for p, algs in agr_res.items():
        part = dsr.parts[p]
        hits_lists = part.hits.lists
        gpra_lists = part.gpra.lists
        for alg, (m, user_ids, centroids, labels) in algs.items():
            alg_lists = part.algs[alg].lists
            for u_i, (label, user_id) in enumerate(zip(labels, user_ids)):
                t = [
                    p,
                    alg,
                    label,
                    np.count_nonzero(m[u_i, :NUM_ITEMS]),
                    user_mean[user_id],
                    user_cnt[user_id]
                ]
                if user_id in dsr.parts[p].hits.lists:
                    alg_hits = [num_hits(part.algs[alg_name].lists[user_id][:NUM_ITEMS],
                                         hits_lists[user_id])
                                for alg_name in alg_names]
                    gpra_hits = [num_hits(gpra_lists[user_id][:NUM_ITEMS],
                                          hits_lists[user_id])]
                else:
                    alg_hits = [0 for alg in alg_names]
                    gpra_hits = [0]
                tuples.append(tuple(t + gpra_hits + alg_hits))

    return pd.DataFrame(tuples,
                        columns=(['fold', 'alg', 'label', 'hits',
                                  'mean_rating', 'rating_cnt', 'gpra'] +
                                 alg_names))


def label_group_by_num_users(frame):
    sorted_labels = frame.label.value_counts().keys()
    label_to_pos = {label: pos for pos, label in enumerate(sorted_labels)}
    frame.loc[:, 'label'] = frame.label.apply(lambda label:
                                              label_to_pos[label])
    return frame


def generate(dsr, results, conf):
    agr_res = results['agreement']

    frame = gen_frame(dsr, agr_res)
    grouped = frame.groupby(['fold', 'alg'])
    relabeled = grouped.apply(label_group_by_num_users)

    fc = relabeled.groupby(['fold', 'label'])

    return fc.mean(), fc.std(), fc.median()


def plot(res, dsr, output_dir, conf, ext='pdf'):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    out_file = os.path.join(output_dir, "out.txt")
    f = open(out_file, "w")

    iters = zip(*[x.iterrows() for x in res])

    meas_names = res[0].columns
    meas_fmt = "{0} {3} ({1} std, {2} mdn)\n"

    for (idx, mean), (_, std), (_, median) in iters:
        f.write("{0}\n".format(idx))
        for meas, m_mean, m_std, m_median in zip(meas_names, mean,
                                                 std, median):
            f.write(meas_fmt.format(m_mean, m_std, m_median, meas))
        f.write("\n")

    f.close()
