#!/usr/bin/env python
"""
plots the metric values at the start and end of each slice registration

Usage:
before_and_after_registration values_dir
"""

from sys import argv
from os import listdir
from metric_values import MetricValues

from numpy import *
import matplotlib.pyplot as plt

metric_values_dir = argv[1]
basenames = listdir(metric_values_dir)

# extract metric values
vals = MetricValues(metric_values_dir)

fig = plt.figure()
ax = fig.gca()
x = range(vals.number_of_slices())
ax.plot(x, vals.initial_values(), label="Initial values")
ax.plot(x, vals.final_values(), label="Final values")
plt.grid(axis="y")
ax.legend()
plt.show()
Rat28_root = registration_root + "results/Rat28/"
thesis_root = "/Users/Matt/Documents/DPhil/written_work/thesis/"

import sys
from os import listdir
from numpy import *
import matplotlib.pyplot as plt

# add registration graphing directory to python path
sys.path.insert(0, registration_root + "graphing")

from metric_values import MetricValues

# metric values dirs
metric_value_dir = registration_root + "results/Rat28/%s/MetricValues/"
rigid_values      = MetricValues(metric_value_dir % "rigid" + "CenteredRigid2DTransform" )
similarity_values = MetricValues(metric_value_dir % "size" + "CenteredSimilarity2DTransform")
affine_values     = MetricValues(metric_value_dir % "affine" + "CenteredAffineTransform")

# check that they all have the same number of slices
assert(rigid_values.number_of_slices() == similarity_values.number_of_slices() == affine_values.number_of_slices())
x = range(rigid_values.number_of_slices())

# plot figure
fig = plt.figure(frameon=False)
ax = fig.add_axes([0.05,0.05,0.9,0.9])
ax.plot(x, rigid_values.initial_values(), label="Initial values")
ax.plot(x, rigid_values.final_values(), label="Final rigid values")
ax.plot(x, similarity_values.final_values(), label="Final similarity values")
ax.plot(x, affine_values.final_values(), 'y', label="Final affine values")
# plot all slices
else:
    def plot_3d_lines(values, labels):
        fig = plt.figure(frameon=False)
        ax = fig.add_axes([0.0,0.0,1.0,1.0], projection='3d')
        for i, slice in enumerate(values):
            x = range(len(slice))
            y = [i] * len(slice)
            ax.plot(x, y, slice, label=labels[i])
        ax.legend()
        plt.xlabel('Iteration', fontsize='xx-large')
        plt.ylabel('Slice Number', fontsize='xx-large')
        ax.set_zlabel("Normalised Correlation", fontsize='xx-large')
        plt.show()
    
    basenames = listdir(metric_values_dir)
    labels = [basename + " (" + str(i+1) + ")" for i, basename in enumerate(basenames)]
    
    metric_values = MetricValues(metric_values_dir)
    
    # plot slices in batches
    if len(argv) == 4 and argv[2] == '--batch-size':
        batch_size = int(argv[3])
        for i in range(0,len(metric_values.values()), batch_size):
            plot_3d_lines(metric_values.values()[i:i+batch_size], labels[i:i+batch_size])
    
    # plot the whole set of slices together    
    else:
        plot_3d_lines(metric_values.values(), labels)
        plot_3d_lines(metric_values.delta_values(), labels)    
Usage:
compare_final_metric_values dir1 [dir2 [dir3...]]
"""

from sys import argv
from metric_values import MetricValues

from numpy import *
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.gca()

for i, directory in enumerate(argv[1:]):
    vals = MetricValues(directory)
    x = range(vals.number_of_slices())
    ax.plot(x, vals.final_values(), label=str(i+1) )

plt.grid(axis="y")
ax.legend()
plt.show()

# plot the differences if 2 directories are given
if len(argv) == 3:
    fig = plt.figure()
    ax = fig.gca()
    vals1 = MetricValues(argv[1]).final_values()
    vals2 = MetricValues(argv[2]).final_values()
    x = range(len(vals1))