#!/usr/bin/env python3

import pickle
import argparse
from load import Segment, CID
from grades import QALog
from collections import Counter
from random import shuffle
from functools import reduce
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator

parser = argparse.ArgumentParser(description='Ptakopět log processing.')
parser.add_argument('blog3', help='Path to a blog3 file')
args = parser.parse_args()

OFFSET = 0.07
NLOCATORS = 4
LEGEND = False
XAXIS_DESC = True
ENGINE_NAME = 'Czech 3'
ENGINE_CODE = 'css'

with open(args.blog3, 'rb') as f:
    data = pickle.load(f)

data = [
    x for x in data if
    (len(x.grade_f) != 0) and
    (x.score is not None) and
    not x.invalid and
    (x.cid.engine == ENGINE_CODE) and
    len([x for line in x.data if line[0] == 'TRANSLATE1']) != 0
]

c_req = {}
c_tim = {}
c_len = {}
q_req = {}
q_tim = {}
q_len = {}

for segment in data:
    avgGrade = np.round(np.average([float(x.overall)
                                    for x in segment.grade_v if x.overall is not None]))

    texts = [s[2] for s in segment.data if s[0] == 'TRANSLATE1']
    texts_len = [len(s[2]) for s in segment.data if s[0] == 'TRANSLATE1']

    if len(texts) != 0:
        c_req.setdefault(segment.score, []).append(len(texts))
        q_req.setdefault(avgGrade, []).append(len(texts))

        time = float(segment.data[-1][1])-float(segment.data[0][1])
        c_len.setdefault(segment.score, []).append(np.average(texts_len))
        q_len.setdefault(avgGrade, []).append(np.average(texts_len))

        if time <= 6*60*1000:
            c_tim.setdefault(segment.score, []).append(time/1000)
            q_tim.setdefault(avgGrade, []).append(time/1000)


for segment in data:
    avgGrade = np.round(np.average([float(x.overall)
                                    for x in segment.grade_f if x.overall is not None]))

    texts = [s[3] for s in segment.data if s[0] == 'CONFIRM_OK']
    texts_len = [len(s[3]) for s in segment.data if s[0] == 'CONFIRM_OK']
    requests_count = len([s[2] for s in segment.data if s[0] == 'TRANSLATE1'])

    if len(texts) != 0:
        c_req.setdefault(segment.score, []).append(requests_count)
        q_req.setdefault(avgGrade, []).append(requests_count)

        time = float(segment.data[-1][1])-float(segment.data[0][1])
        c_len.setdefault(segment.score, []).append(np.average(texts_len))
        q_len.setdefault(avgGrade, []).append(np.average(texts_len))

        if time <= 6*60*1000:
            c_tim.setdefault(segment.score, []).append(time/1000)
            q_tim.setdefault(avgGrade, []).append(time/1000)


SCORES = [x for x in range(1, 5+1)]
c_req = {k: np.average(v) for (k, v) in c_req.items()}
c_tim = {k: np.average(v) for (k, v) in c_tim.items()}
c_len = {k: np.average(v) for (k, v) in c_len.items()}
q_req = {k: np.average(v) for (k, v) in q_req.items()}
q_tim = {k: np.average(v) for (k, v) in q_tim.items()}
q_len = {k: np.average(v) for (k, v) in q_len.items()}

# display plot
plt.rcParams.update({'font.size': 11})

fig = plt.figure(figsize=(5, 4.5 if LEGEND else 3))

# display confidence
ax1 = fig.add_subplot()
color = 'darkgreen'
ax1.set_ylabel('Source text length (char)', color=color)
ax1.tick_params(axis='y', labelcolor=color)
sc1 = ax1.plot(
    [str(x) for x in SCORES],
    [c_len[k] for k in SCORES],
    color=color, label='Source text length',
    markersize=10, marker='o',
    alpha=0.6
)
ax1.yaxis.labelpad = -2
ax1.yaxis.set_major_locator(MaxNLocator(NLOCATORS, min_n_ticks=NLOCATORS))
ax1.set_ylim(45, 140)

ax2 = ax1.twinx()
color = 'darkred'
ax2.set_ylabel('Requests', color=color)
ax2.tick_params(axis='y', labelcolor=color)
sc2 = ax2.plot(
    [str(x) for x in SCORES],
    [c_req[k] for k in SCORES],
    color=color, label='Requests',
    markersize=10, marker='v',
    alpha=0.6
)
ax2.yaxis.set_major_locator(MaxNLocator(NLOCATORS, min_n_ticks=NLOCATORS))
ax2.yaxis.labelpad = -1
ax2.set_ylim(1, 8)

ax3 = ax1.twinx()
color = 'darkblue'
ax3.set_ylabel('Time (s)', color=color)
ax3.tick_params(axis='y', labelcolor=color)
sc3 = ax3.plot(
    [str(x) for x in SCORES],
    [c_tim[k] for k in SCORES],
    color=color, label='Time',
    markersize=10, marker='s',
    alpha=0.6
)
ax3.spines['right'].set_position(('outward', 30))
ax3.yaxis.labelpad = -5
ax3.yaxis.set_major_locator(MaxNLocator(NLOCATORS, min_n_ticks=NLOCATORS))
ax3.set_ylim(33, 100)


# display qualities
ax4 = ax1.twinx()
color = 'lightgreen'
ax4.set_ylabel('Source text length (char)', color=color)
ax4.tick_params(axis='y', labelcolor=color)
sc4 = ax4.plot(
    [str(x) for x in SCORES],
    [q_len[k] for k in SCORES],
    linestyle='--',
    color=color, label='Source text length',
    markersize=10, marker='o',
    alpha=1
)
ax4.yaxis.set_visible(False)
ax4.set_ylim(45, 150)

ax5 = ax1.twinx()
color = 'lightcoral'
ax5.set_ylabel('Requests', color=color)
ax5.tick_params(axis='y', labelcolor=color)
sc5 = ax5.plot(
    [str(x) for x in SCORES],
    [q_req[k] for k in SCORES],
    linestyle='--',
    color=color, label='Requests',
    markersize=10, marker='v',
    alpha=1
)
ax5.yaxis.set_visible(False)
ax5.set_ylim(1.0, 8)

ax6 = ax1.twinx()
color = 'lightblue'
ax6.set_ylabel('Time (s)', color=color)
ax6.tick_params(axis='y', labelcolor=color)
sc6 = ax6.plot(
    [str(x) for x in SCORES],
    [q_tim[k] for k in SCORES],
    linestyle='--',
    color=color, label='Time',
    markersize=10, marker='s',
    alpha=1
)
ax6.yaxis.set_visible(False)
ax6.set_ylim(33, 100)


if XAXIS_DESC:
    ax1.set_xlabel('Self-reported user confidence, Overall translation quality')
if LEGEND:
    lns = sc1+sc2+sc3+sc4+sc5+sc6
    lbs = [
        'Text length, conf.', 'Requests, conf.', 'Time, conf.',
        'Text length, qual.', 'Requests, qual.', 'Time, qual.'
    ]
    plt.legend(
        handles=lns, labels=lbs,
        loc='upper center',
        fancybox=True,
        bbox_to_anchor=(0.5, 1.3),
        ncol=2
    )
fig.tight_layout(rect=[-0.03, -0.05, 1.03, 1])
# fig.text(0.01, 0.43, ENGINE_NAME, rotation=90)
plt.title(ENGINE_NAME)

plt.show()
# print(Counter(requests))
# print(Counter(linears))
# print(Counter(increasings))
