#!/usr/bin/env python
"""Compare and plot results generated by the simulators.
"""
from __future__ import division
import argparse
import matplotlib.pyplot as plt
import json
import numpy

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('dynamics2_data',
        help='Data file generated by dynamics2')
    parser.add_argument('agent2_data',
        help='Data file generated by agent2')

    args = parser.parse_args()

    with open(args.agent2_data) as f:
        agentdata = json.load(f)

    with open(args.dynamics2_data) as f:
        dyndata = json.load(f)

    def rms(data, axis=0):
        return numpy.sqrt(numpy.mean(data**2, axis))

    for k in agentdata.keys():
        print k, rms(numpy.subtract(agentdata[k][:-1], dyndata[k]))

    plt.figure()
    plt.plot(agentdata['susceptible'], 'b', label='Agent: susceptible')
    plt.plot(agentdata['infected'], 'r', label='Agent: infected')
    plt.plot(agentdata['recovered'], 'g', label='Agent: recovered')

    plt.plot(dyndata['susceptible'], 'c', label='System Dynamics: susceptible')
    plt.plot(dyndata['infected'], 'y', label='System Dynamics: infected')
    plt.plot(dyndata['recovered'], 'k', label='System Dynamics: recovered')

    plt.legend()
    plt.xlabel('Time')
    plt.ylabel('Individual Count')
    plt.title('Correspondence Assessment')
    plt.show()


if __name__ == '__main__':
    main()