### COPYRIGHT AND LICENSE STATEMENT ###
#
#  This file is part of the MLMCLangevin code.
#  
#  (c) Copyright Eike Mueller, Grigoris Katsiolides and Tony Shardlow,
#      University of Bath
#  
#  Main Developer: Eike Mueller
#  
#  Contributors:
#    Grigoris Katsiolides
#    Tony Shardlow
#  
#  MLMCLangevin is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#  
#  MLMCLangevin is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#  
#  You should have received a copy of the GNU Lesser General Public License
#  along with MLMCLangevin (see files COPYING and COPYING.LESSER). If not,
#  see <http://www.gnu.org/licenses/>.
#
### COPYRIGHT AND LICENSE STATEMENT ###

import re
import sys
import numpy as np
from matplotlib import pyplot as plt

regexfloat = '-*[0-9]+\.[0-9]+[eE][+-][0-9]+'

###############################################################
# One trajectory
###############################################################
class Trajectory(object):
    def __init__(self,label=''):
        self.length = 0
        self.label = label
        self.t = []
        self.x = []
        self.u = []

    # Append a new data point to end of trajectory
    def append(self,t,x,u):
        self.t.append(t)
        self.x.append(x)
        self.u.append(u)

    # Convert to numpy-arrays
    def finalise_append(self):
        self.t = np.array(self.t)
        self.x = np.array(self.x)
        self.u = np.array(self.u)

###############################################################
# Collection of trajectories read from a file
###############################################################
class Trajectories(object):
    def __init__(self,filename):
        self.filename = filename
        self.data = []
        self._read_file()

    # Parse file and read trajectories
    def _read_file(self):
        parse_traj = False
        with open(self.filename,'r') as f:
            for line in f:
                m = re.match('^ *# *trajectory *([0-9]+) *$',line)
                if (m):
                    parse_traj = True
                    trajectory = Trajectory(m.group(1))
                if (parse_traj and re.match('^ *$',line)):
                    parse_traj = False
                    self.data.append(trajectory)
                m = re.match(' *('+regexfloat+') +('+regexfloat+') +('+regexfloat+') *',line)
                if (parse_traj and m):
                    trajectory.append(float(m.group(1)),
                                      float(m.group(2)),
                                      float(m.group(3)))

    # Plot position and velocity separately as functions of time
    def plot_xu(self,outputfilename,ntraj=0,
                min_x=-2.0,max_x=+2.0,
                min_u=-2.0,max_u=+2.0):
        plt.clf()
        if ( (ntraj == 0) or (ntraj > len(self.data)) ):
            ntraj = len(self.data)
        # Position
        plt.subplot(2, 1, 1)
        ax = plt.gca()
        ax.set_ylim(min_x,max_x)
        ax.set_ylabel('Position Q')
        color_func = lambda x: str(1.0-0.8*(float(x)/ntraj))
        i = 0
        for traj in self.data:
            if (i < ntraj):
                plt.plot(traj.t,traj.x,color=color_func(traj.label))
            i += 1
        # Velocity
        plt.subplot(2, 1, 2)
        ax = plt.gca()
        ax.set_xlabel('time')
        ax.set_ylabel('Momentum P')
        ax.set_ylim(min_u,max_u)
        i = 0
        for traj in self.data:
            if (i < ntraj):
                plt.plot(traj.t,traj.u,color=color_func(traj.label))
            i += 1
        plt.savefig(outputfilename,bbox_inches='tight')

    # Plot in phase-space
    def plot_phasespace(self,outputfilename,ntraj=0,
                        min_x=-2.0,max_x=+2.0,
                        min_u=-2.0,max_u=+2.0):
        plt.clf()
        if ( (ntraj == 0) or (ntraj > len(self.data)) ):
            ntraj = len(self.data)
        ax = plt.gca()
        ax.set_xlim(min_x,max_x)
        ax.set_ylim(min_u,max_u)
        ax.set_xlabel('Position Q')
        ax.set_ylabel('Momentum P')
        color_func = lambda x: str(1.0-0.8*(float(x)/ntraj))
        i = 0
        for traj in self.data:
            if (i < ntraj):
                plt.plot(traj.x,traj.u,color=color_func(traj.label))
            i += 1
        plt.savefig(outputfilename,bbox_inches='tight')

###############################################################
# M A I N
###############################################################
if (__name__ == '__main__'):
    if (len(sys.argv) != 3):
        print 'Usage: python '+sys.argv[0]+' <filename> <ntraj>'
        sys.exit(0)
    filename = sys.argv[1]    
    ntraj = int(sys.argv[2])
    trajs = Trajectories(filename)
    # Plot ranges
    min_x = -4.0
    max_x =  4.0
    min_u = -4.0
    max_u =  4.0
    # Plot both Position and Momentum separately and plot
    # in phase space
    trajs.plot_xu('traj_xu.pdf',ntraj=ntraj,
                  min_x=min_x,max_x=max_x,min_u=min_u,max_u=max_u)
    trajs.plot_phasespace('traj_phasespace.pdf',ntraj=ntraj,
                          min_x=min_x,max_x=max_x,min_u=min_u,max_u=max_u)