#!/usr/bin/env python2

"""Program to investigate the transition distribution $p(z_{n+1} \mid z_n)$
for the available map trajectories.

Usage: <program name> <trajectory> [<trajectory> [...]]"""

from collections import defaultdict
from sys import argv

from matplotlib import pyplot as plt

import numpy as np

from observation import parse_map_trajectory

if __name__ == '__main__':
    assert len(argv) > 1, "Need at least one input file"

    pairs = defaultdict(lambda: [])

    for filename in argv[1:]:
        with open(filename) as fp:
            traj = list(parse_map_trajectory(fp))
            for past, present in zip(traj, traj[1:]):
                for key in past.data.iterkeys():
                    pairs[key].append((past[key], present[key]))

    vf_pairs = np.array(pairs['vf'])
    vf_now = vf_pairs[:, 0]
    vf_diff = vf_pairs[:, 1] - vf_now
    hist, x_edges, y_edges = np.histogram2d(vf_now, vf_diff, bins=50)
    pc_x, pc_y = np.meshgrid(x_edges, y_edges)
    # Normalise columns
    # sums = np.sum(hist, axis=0)
    # sums[sums == 0] = 1
    # hist /= sums

    plt.pcolormesh(pc_x, pc_y, hist)
    plt.xlabel('$v_f$ (m/s)')
    plt.ylabel("$v_f' - v_f$ (m/s)")
    plt.xlim(x_edges[0], x_edges[-1])
    plt.ylim(y_edges[0], y_edges[-1])
    plt.title("Transition frequencies for forward velocity ($v_f$)")

    plt.show()

    # Now do a plot of ve and vn
    ve_pairs = np.array(pairs['ve'])
    ve_deltas = ve_pairs[:, 1] - ve_pairs[:, 0]

    vn_pairs = np.array(pairs['vn'])
    vn_deltas = vn_pairs[:, 1] - vn_pairs[:, 0]

    full_data = np.zeros((vn_deltas.size, 2))
    full_data[:, 0] = ve_deltas
    full_data[:, 1] = vn_deltas

    print("Covariance of data:")
    print(np.cov(full_data.T))
    print("Mean of data:")
    print(np.mean(full_data, axis=0))

    hist, x_edges, y_edges = np.histogram2d(ve_deltas, vn_deltas, bins=100)
    pc_x, pc_y = np.meshgrid(x_edges, y_edges)

    plt.pcolormesh(pc_x, pc_y, hist)
    plt.xlabel("$v_e' - v_e$ (m/s)")
    plt.ylabel("$v_n' - v_n$ (m/s)")
    plt.xlim(x_edges[0], x_edges[-1])
    plt.ylim(y_edges[0], y_edges[-1])
    plt.title("Relative frequencies for changes in velocity")

    plt.show()

@@ -190,16 +190,40 @@
 .%5D%5D%22%22%22%0A%0A
+from bz2 import BZ2File%0A
 from col
@@ -250,16 +250,16 @@
 ultdict%0A
-
 from sys
@@ -540,33 +540,159 @@
     
-with open(filename) as fp
+if filename.endswith('.bz2'):%0A            trajectory_fp = BZ2File(filename)%0A        else:%0A            trajectory_fp = open(filename, 'rb')%0A%0A        try
 :%0A  
@@ -734,16 +734,27 @@
 jectory(
+trajectory_
 fp))%0A   
@@ -750,24 +750,72 @@
 ry_fp))%0A    
+    finally:%0A            trajectory_fp.close()%0A%0A
         for 
@@ -860,28 +860,24 @@
             
-    
 for key in p
@@ -897,20 +897,16 @@
 keys():%0A
-    
         

Update transition investigation script to handle bz2