#!/usr/bin/env python3

help =\
'''
Plot trajectory, or frames if solution has two spatial dimensions, generated by
running a PETSc TS program.  Reads output from
   -ts_monitor binary:TDATA -ts_monitor_solution binary:UDATA
Requires copies or sym-links to $PETSC_DIR/lib/petsc/bin/PetscBinaryIO.py and
$PETSC_DIR/lib/petsc/bin/petsc_conf.py.
'''

import PetscBinaryIO

from sys import exit, stdout
from time import sleep
from argparse import ArgumentParser, RawTextHelpFormatter
import numpy as np
import matplotlib.pyplot as plt

parser = ArgumentParser(description=help,
                        formatter_class=RawTextHelpFormatter)
parser.add_argument('tfile',metavar='TDATA',
                    help='from -ts_monitor binary:TDATA')
parser.add_argument('ufile',metavar='UDATA',
                    help='from -ts_monitor_solution binary:UDATA')
parser.add_argument('-mx',metavar='MX', type=int, default=-1,
                    help='spatial grid with MX points in x direction')
parser.add_argument('-my',metavar='MY', type=int, default=-1,
                    help='spatial grid with MY points in y direction; =MX by default')
parser.add_argument('-dof',metavar='DOF', type=int, default=1,
                    help='degrees of freedom of solution; =1 by default')
parser.add_argument('-c',metavar='C', type=int, default=0,
                    help='component; =0,1,..,dof-1; ignored if dof=1)')
parser.add_argument('-o',metavar='FILE',dest='filename',
                    help='image file FILE (trajectory case)')
parser.add_argument('-oroot',metavar='ROOT',dest='rootname',
                    help='frame files ROOT000.png,ROOT001.png,... (movie case)')
args = parser.parse_args()

if args.mx > 0 and args.my < 1:
    args.my = args.mx
frames = (args.mx > 0)

io = PetscBinaryIO.PetscBinaryIO()
t = np.array(io.readBinaryFile(args.tfile)).flatten()
U = np.array(io.readBinaryFile(args.ufile)).transpose()
dims = np.shape(U)

if len(t) != dims[1]:
    print('time dimension mismatch: %d != %d' % (len(t),dims[1]))
    exit(1)
if frames:
    if args.dof == 1:
        if dims[0] != args.mx * args.my:
            print('spatial dimension mismatch: %d != %d * %d (and dof=1)' % \
                  (dims[0],args.mx,args.my))
            exit(2)
        U = np.reshape(U,(args.my,args.mx,len(t)))
        dims = np.shape(U)
        print('solution U is shape=(%d,%d,%d)' % tuple(dims))
    else:
        if dims[0] != args.mx * args.my * args.dof:
            print('spatial dimension mismatch: %d != %d * %d * %d' % \
                  (dims[0],args.mx,args.my,args.dof))
            exit(3)
        U = np.reshape(U,(args.my,args.mx,args.dof,len(t)))
        dims = np.shape(U)
        print('solution U is shape=(%d,%d,%d,%d)' % tuple(dims))
    print('time t has length=%d, with mx x my = %d x %d frames' % \
          (dims[-1],dims[1],dims[0]))
else:
    print('time t has length=%d, solution Y is shape=(%d,%d)' % \
          (len(t),dims[0],dims[1]))

framescmap = 'jet'  # close to the PETSc X windows default
#framescmap = 'inferno'
#framescmap = 'gray'

if frames:
    print('generating files %s000.png .. %s%03d.png:' % \
          (args.rootname,args.rootname,len(t)-1))
    if args.dof == 1:
        plt.imshow(U[:,:,0],cmap=framescmap)
    else:
        plt.imshow(U[:,:,args.c,0],cmap=framescmap)
    plt.title('t = %g' % t[0])
    if args.rootname:
        plt.savefig(args.rootname + "%03d.png" % 0)
    else:
        plt.ion()
        plt.show()
    for k in range(len(t)-1):
        print('.', end =' ')
        stdout.flush()
        if args.dof == 1:
            plt.imshow(U[:,:,k+1],cmap=framescmap)
        else:
            plt.imshow(U[:,:,args.c,k+1],cmap=framescmap)
        plt.title('t = %g' % t[k+1])
        if args.rootname:
            plt.savefig(args.rootname + "%03d.png" % (k+1))
        else:
            plt.pause(0.1)
    print('.')
else:
    if args.dof > 1:
        print('plotting only component %d ...' % args.c)
        plt.plot(t,U[args.c],label='y[%d]' % args.c)
    else:
        for k in range(dims[0]):
            plt.plot(t,U[k],label='y[%d]' % k)
    plt.xlabel('t')
    plt.legend()
    if args.filename:
        print('writing file %s' % args.filename)
        plt.savefig(args.filename)
    else:
        plt.show()
