import csv
from itertools import groupby, ifilter
import glob
import matplotlib.pyplot as plt
from muscles import MTest, Tracker
import numpy as np

def readcsv(fname, types = None):
    f = open(fname, 'rb')
    r = csv.DictReader(f)
    out = []
    for row in r:
        if types is not None:
            for k in types.keys():
                row[k] = types[k](row[k])
        out.append(row)
    f.close()
    return out

class DictKey:
    def __init__(self, key):
        self.key = key

    def __call__(self, q):
        return q[self.key]

def avg(iterable, keyf=None):
    sum = 0.0
    cnt = 0
    for i in iterable:
        if keyf is not None:
            i = keyf(i)
        sum += float(i)
        cnt += 1
    return sum/cnt

def plot(fname, lst, xtitle, ytitle, format = 'o-', xt = None):
    plt.figure()
    if not isinstance(lst, dict):
        lst = {'': lst}
    for name in sorted(lst.keys()):
        x, y = zip(*lst[name])
        if isinstance(x[0], str):
            xt = x
            x = range(0, len(x))
        plt.margins(0.1)
        plt.plot(x, y, format, label=name)
    if len(lst) > 1:
        plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
        rightmargin = 0.75 if max([len(i) for i in lst.keys()]) > 4 else 0.8
        plt.subplots_adjust(right=rightmargin)
    if xt is not None:
        rotation = None
        if len(xt[0]) > 4:
            rotation = 45
            plt.subplots_adjust(bottom=0.3)
        plt.xticks([float(i) for i in x], xt, rotation=rotation, ha='right')
    plt.xlabel(xtitle)
    plt.ylabel(ytitle)
    plt.grid()
    plt.savefig('images/' + fname)

def muscname(name):
    return name.replace('.Origin', '').replace('Muscle.', '')

stats = readcsv('stats.csv', types = {'window size': float, 'scale': float, 'time per frame': float})
ws = DictKey('window size')
err = DictKey('error')
scale = DictKey('scale')
test = DictKey('test')
t = DictKey('time per frame')

ws_scale_error = {}
for sc in set([scale(i) for i in stats]):
    ws_scale_error['scale=%.2f'%sc] = [(i, avg(j, err)) for i, j in groupby(sorted(ifilter(lambda x: scale(x) == sc, stats), key=ws), ws)]
plot('ws_scale_error.png', ws_scale_error, 'Window size', 'Error')

ws_scale_time = {}
for sc in set([scale(i) for i in stats]):
    ws_scale_time['scale=%.2f'%sc] = [(i, avg(j, t)*1000) for i, j in groupby(sorted(ifilter(lambda x: scale(x) == sc, stats), key=ws), ws)]
plot('ws_scale_time.png', ws_scale_time, 'Window size', 'Time (msec)')

ws_error = [(i, avg(j, err)) for i, j in groupby(sorted(stats, key=ws), ws)]
plot('ws_error.png', ws_error, 'Window size', 'Error')

scale_error = [(i, avg(j, err)) for i, j in groupby(sorted(stats, key=scale), scale)]
plot('scale_error.png', scale_error, 'Scale', 'Error')

ws_time = [(i, avg(j, t)*1000) for i, j in groupby(sorted(stats, key=ws), ws)]
plot('ws_time.png', ws_time, 'Window size', 'Time (msec)')

scale_time = [(i, avg(j, t)*1000) for i, j in groupby(sorted(stats, key=scale), scale)]
plot('scale_time.png', scale_time, 'Scale', 'Time (msec)')

# per test error

best_scale = 0.5
best_ws = 20
best_case_data = ifilter(lambda x: x['scale'] == best_scale and x['window size'] == best_ws, stats)
setup = lambda x: x['test'][2:]
test_error = [(i, avg(j, err)) for i,j in groupby(sorted(best_case_data, key=setup), key=setup)]
plot('test_error.png', test_error, 'Setup', 'Error', 'o')

# per muscle error

lst = glob.glob('tests/*_' + str(best_scale) + '_' + str(best_ws) + '_error.csv')

all_data = []
each_data = {}
for fname in lst:
    data = readcsv(fname)
    this_data = []
    for i in data:
        if i['jaw'] is None:
            break
        this_data.append(i)
    # wow
    tname = fname.replace('tests\\', '').replace('_' + str(best_scale) + '_' + str(best_ws) + '_error.csv', '')
    each_data[tname] = this_data
    all_data += this_data

each_error = {'avg': []}
for i in sorted(all_data[0].keys()):
    if i == 'frame':
        continue
    name = muscname(i)
    for j in each_data.keys():
        if j not in each_error.keys():
            each_error[j] = []
        each_error[j].append((name, avg(each_data[j], DictKey(i))))
    each_error['avg'].append((name, avg(all_data, DictKey(i))))
plot('muscle_error.png', each_error, 'Muscle', 'Error')

# max muscle error
test = MTest('m1s1')
P = test.tracker.find_P(test.tracker.prev_points)
maxerr = []
for mid in range(0, test.tracker.prev_x.shape[0]):
    x = np.asmatrix(np.zeros(test.tracker.prev_x.shape)).T
    x[0, mid] = 1.0
    maxcur = 0.0
    for i in range(0, len(test.tracker.prev_points)):
        if i == test.tracker.jaw or i in test.tracker.static_points:
            continue
        d = np.matrix([[0.0], [0.0]])
        p0 = Tracker.project(np.zeros(test.tracker.prev_x.shape), test.tracker.points[i], P, test.tracker.coeffs[i])
        maxcur = max(maxcur, Tracker.error(x, p0, test.tracker.points[i], d, P, test.tracker.coeffs[i]))
    maxerr.append((muscname(test.tracker.muscles[mid].name), float(maxcur)))
plot('max_muscle_error.png', maxerr, 'Muscle', 'Max displacement')

# movement degrade
data = readcsv('tests/m1s4_0.5_30_error.csv')
frame_err = [(int(i['frame']), avg([i[j] for j in i.keys() if j != 'frame'])) for i in data if i['jaw'] is not None]
plot('movement_error.png', frame_err, 'Frame', 'Error')

# mesh
mesh1 = readcsv('stats_mesh_m1.csv')
data = {}
rename = {'spartak': 'm1', 'new': 'm2', 'random':'random'}
for i, j in groupby(sorted(mesh1, key=DictKey('mesh')), DictKey('mesh')):
    data[rename[i]] = [(k['test'], k['error']) for k in j]
plot('mesh_error_m1.png', data, 'Test', 'Error')

mesh1 = readcsv('stats_mesh_m2.csv')
data = {}
for i, j in groupby(sorted(mesh1, key=DictKey('mesh')), DictKey('mesh')):
    data[rename[i]] = [(k['test'], k['error']) for k in j]
plot('mesh_error_m2.png', data, 'Test', 'Error')