from __future__ import with_statement
import motmot.ufmf.ufmf as ufmf_mod
import sys, os, tempfile, re, contextlib, warnings
from optparse import OptionParser
import flydra_analysis.a2.auto_discover_ufmfs as auto_discover_ufmfs
import flydra_analysis.a2.auto_discover_movies as auto_discover_movies
import numpy as np
import tables
import flydra_analysis.a2.utils as utils
import flydra_analysis.analysis.result_utils as result_utils
import subprocess, collections
import flydra_analysis.a2.ufmf_tools as ufmf_tools
import flydra_analysis.a2.core_analysis as core_analysis
import flydra_core.kalman.dynamic_models as dynamic_models
from flydra_analysis.a2.orientation_ekf_fitter import compute_ori_quality
import flydra_core.reconstruct as reconstruct
import flydra_core.geom as geom
import cherrypy  # ubuntu: install python-cherrypy3
import benu

from tables_tools import open_file_safe

def get_config_defaults():
    # keep in sync with usage in main() below
    what = {'show_2d_position': False,
            'show_2d_orientation': False,
            'show_3d_MLE_position': False,
            'show_3d_smoothed_position': False,
            'show_3d_raw_orientation': False,
            'show_3d_raw_chosen_orientation': False,
            'show_3d_smoothed_orientation': False,
            'show_3d_obj_position_text': False,
            'minimum_display_orientation_quality':0,
            'zoom_obj':None,
            'zoom_orig_pixels':50,
            'zoom_factor':5,
            'white_background': False,
            'max_resolution': None,
            'border_pixels': 0,
            'obj_labels':False,
            'linewidth':1.0,
            'show_cam_id':False,
            'image_manipulation':'raw',
            }
    default = collections.defaultdict(dict)
    default['what to show']=what
    return default

def montage(fnames, title, target):

    def get_tile(N):
        rows = int(np.ceil(np.sqrt(float(N))))
        cols = rows
        return '%dx%d'%(rows,cols)

    tile = get_tile( len(fnames) )
    imnames = ' '.join(fnames)

    CMD=("gm montage %s -mode Concatenate -tile %s -bordercolor white "
         "-title '%s' "
         "%s"%(imnames, tile, title, target))
    #print CMD
    subprocess.check_call(CMD,shell=True)

def load_3d_raw_data(kalman_filename,require_qual=True,**kwargs):
    with open_file_safe( kalman_filename, mode='r' ) as kh5:
        ca = core_analysis.get_global_CachingAnalyzer()
        all_obj_ids, obj_ids, is_mat_file, data_file, extra = \
                     ca.initial_file_load(kalman_filename)
        allrows = []
        if require_qual:
            allqualrows = []
        this_kw = {'min_ori_quality_required':kwargs['min_ori_quality_required'],
                   'ori_quality_smooth_len':kwargs['ori_quality_smooth_len']}
        for obj_id in obj_ids:
            rows = ca.load_dynamics_free_MLE_position( obj_id, data_file=kh5,
                                                       **this_kw)
            allrows.append( rows )
            if require_qual:
                if np.any( ~np.isnan( rows['hz_line0'] )):
                    qualrows = compute_ori_quality(data_file,
                                                   rows['frame'],
                                                   obj_id,
                                                   smooth_len=0)
                else:
                    qualrows = np.zeros_like(rows['hz_line0'])
                allqualrows.append( qualrows )
    data3d = np.concatenate(allrows)
    if require_qual:
        dataqual3d = np.concatenate(allqualrows)
        return data3d, dataqual3d
    else:
        return data3d

def is_obj_in_frame_range( obj_id, all_obj_ids, frames, start=None, stop=None):
    valid_cond = np.ones( frames.shape, dtype=np.bool)
    if start is not None:
        valid_cond = valid_cond & (frames >= start)
    if stop is not None:
        valid_cond = valid_cond & (frames <= stop )
    valid_obj_ids = all_obj_ids[ valid_cond ]
    return bool(np.sum(valid_obj_ids==obj_id))

def load_3d_data(kalman_filename,start=None,stop=None,require_qual=True,**kwargs):
    with open_file_safe( kalman_filename, mode='r' ) as kh5:
        ca = core_analysis.get_global_CachingAnalyzer()
        all_obj_ids, obj_ids, is_mat_file, data_file, extra = \
                     ca.initial_file_load(kalman_filename)
        fps = extra['frames_per_second']
        dynamic_model_name = None
        if dynamic_model_name is None:
            dynamic_model_name = extra.get('dynamic_model_name',None)
            if dynamic_model_name is None:
                dynamic_model_name = dynamic_models.DEFAULT_MODEL
                warnings.warn('no dynamic model specified, using "%s"'%dynamic_model_name)
            else:
                print 'detected file loaded with dynamic model "%s"'%dynamic_model_name
            if dynamic_model_name.startswith('EKF '):
                dynamic_model_name = dynamic_model_name[4:]
            print '  for smoothing, will use dynamic model "%s"'%dynamic_model_name
        allrows = []
        if require_qual:
            allqualrows = []
        for obj_id in obj_ids:
            if not is_obj_in_frame_range( obj_id, all_obj_ids, extra['frames'], start=start, stop=stop):
                # obj_id not in range of frames that we're analyzing now
                continue
            try:
                rows = ca.load_data( obj_id, kalman_filename,
                                     use_kalman_smoothing=True,
                                     frames_per_second=fps,
                                     dynamic_model_name=dynamic_model_name,
                                     return_smoothed_directions = True,
                                     **kwargs)
            except core_analysis.NotEnoughDataToSmoothError:
                warnings.warn('not enough data to smooth obj_id %d, skipping.'%(obj_id,))
                continue
            allrows.append( rows )
            if require_qual:
                qualrows = compute_ori_quality(data_file,
                                               rows['frame'],
                                               obj_id,
                                               smooth_len=0)
                allqualrows.append( qualrows )
    data3d = np.concatenate(allrows)
    if require_qual:
        dataqual3d = np.concatenate(allqualrows)
        return data3d, dataqual3d
    else:
        return data3d

def make_montage( h5_filename,
                  cfg_filename=None,
                  ufmf_dir=None,
                  dest_dir = None,
                  save_ogv_movie = False,
                  no_remove = False,
                  max_n_frames = None,
                  start = None,
                  stop = None,
                  movie_fnames = None,
                  movie_cam_ids = None,
                  caminfo_h5_filename = None,
                  colormap = None,
                  kalman_filename = None,
                  candidate_index = 0,
                  nth_frame=1,
                  verbose = False,
                  reconstructor=None,
                  **kwargs):
    config = get_config_defaults()
    if cfg_filename is not None:
        loaded_cfg = cherrypy.lib.reprconf.as_dict( cfg_filename )
        for section in loaded_cfg:
            config[section].update( loaded_cfg.get(section,{}) )
    else:
        warnings.warn('no configuration file specified -- using defaults')

    orientation_3d_line_length = 0.1

    if (config['what to show']['show_3d_smoothed_position'] or
        config['what to show']['show_3d_MLE_position'] or
        config['what to show']['show_3d_raw_orientation'] or
        config['what to show']['show_3d_raw_chosen_orientation'] or
        config['what to show']['show_3d_smoothed_orientation'] or
        config['what to show']['show_3d_obj_position_text']
        ):
        if kalman_filename is None:
            raise ValueError('need kalman filename to show requested 3D data')

    if config['what to show']['obj_labels']:
        if kalman_filename is None:
            raise ValueError('need kalman filename to show object labels')

    if kalman_filename is not None:
        if (config['what to show']['show_3d_smoothed_orientation'] or
            config['what to show']['show_3d_raw_orientation'] or
            config['what to show']['show_3d_raw_chosen_orientation']):
            need_quality_data = True
        else:
            need_quality_data = False

        if need_quality_data:
            # need data about quality of tracking
            data3d, dataqual_3d = load_3d_data(kalman_filename,start=start,stop=stop,
                                               require_qual=True, **kwargs)
        else:
            data3d = load_3d_data(kalman_filename,start=start,stop=stop,
                                  require_qual=False, **kwargs)
            dataqual_3d = None

        if (config['what to show']['show_3d_MLE_position'] or
            config['what to show']['show_3d_raw_orientation']):
            if need_quality_data:
                data_raw_3d, dataqual_raw_3d = load_3d_raw_data(kalman_filename,**kwargs)
            else:
                data_raw_3d = load_3d_raw_data(kalman_filename,require_qual=False,**kwargs)
                dataqual_raw_3d = None
        else:
            data_raw_3d, dataqual_raw_3d = None, None
        if reconstructor is None:
            R = reconstruct.Reconstructor(kalman_filename)
        else:
            R = reconstruct.Reconstructor(reconstructor)
    else:
        data3d = R = data_raw_3d = None
        dataqual_raw_3d = None
        dataqual_3d = None

    min_ori_qual = config['what to show']['minimum_display_orientation_quality']

    if movie_fnames is None:
        # This works based on UUIDs
        movie_fnames = auto_discover_movies.find_movies( h5_filename,
                                                         ufmf_dir=ufmf_dir,
                                                         candidate_index=candidate_index,
                                                         verbose=verbose)
        if verbose:
            print 'autodiscovery: found movie_fnames: %r'%(movie_fnames,)
    else:
        if verbose:
            print 'autodiscovery: movie_fnames specified, not finding movies'

    if len(movie_fnames)==0:
        if verbose:
            print 'autodiscovery: no FMF files found, looking for ufmfs'
        movie_fnames = auto_discover_ufmfs.find_ufmfs( h5_filename,
                                                       ufmf_dir=ufmf_dir,
                                                       careful=True,
                                                       verbose=verbose,
                                                       )
    else:
        if verbose:
            print 'autodiscovery: prefixing directory'
        if ufmf_dir is not None:
            if verbose:
                print 'autodiscovery: prefixing movie names with directory %r'%(ufmf_dir,)
            movie_fnames = [ os.path.join(ufmf_dir,f) for f in movie_fnames]

    if len(movie_fnames)==0:
        raise ValueError('no input movies -- nothing to do')
    elif verbose:
        print 'movie_fnames:',movie_fnames

    if dest_dir is None:
        dest_dir = os.curdir
    else:
        if not os.path.exists( dest_dir ):
            os.makedirs(dest_dir)

    # get name of data

    datetime_str = os.path.splitext(os.path.split(h5_filename)[-1])[0]
    if datetime_str.startswith('DATA'):
        datetime_str = datetime_str[4:19]

    workaround_ffmpeg2theora_bug = True

    if caminfo_h5_filename is None:
        caminfo_h5_filename = h5_filename

    if caminfo_h5_filename is not None:
        with open_file_safe( caminfo_h5_filename, mode='r' ) as h5:
            camn2cam_id, tmp = result_utils.get_caminfo_dicts(h5)
            del tmp
    else:
        camn2cam_id = None

    blank_images = {}

    all_frame_montages = []
    for frame_enum,(frame_dict,frame) in enumerate(ufmf_tools.iterate_frames(
        h5_filename, movie_fnames,
        movie_cam_ids = movie_cam_ids,
        white_background=config['what to show']['white_background'],
        max_n_frames = max_n_frames,
        start = start,
        stop = stop,
        rgb8_if_color = True,
        camn2cam_id = camn2cam_id,
        )):

        if frame_enum%nth_frame!=0:
            continue

        tracker_data = frame_dict['tracker_data']
        global_data = frame_dict['global_data']

        if data3d is not None:
            this_frame_3d_data = data3d[data3d['frame']==frame]
            if dataqual_3d is None:
                this_frame_dataqual = None
            else:
                this_frame_dataqual = dataqual_3d[data3d['frame']==frame]
        else:
            this_frame_3d_data = None
            this_frame_dataqual = None

        if data_raw_3d is not None:
            this_frame_raw_3d_data = data_raw_3d[data_raw_3d['frame']==frame]
            if dataqual_raw_3d is None:
                this_frame_raw_dataqual = None
            else:
                this_frame_raw_dataqual = dataqual_raw_3d[data_raw_3d['frame']==frame]
        else:
            this_frame_raw_3d_data = None
            this_frame_raw_dataqual = None

        if config['what to show']['zoom_obj']:
            zoom_cond_3d = this_frame_3d_data['obj_id']==config['what to show']['zoom_obj']
            if np.sum( zoom_cond_3d ) == 0:
                # object not in this frame
                this_frame_this_obj_3d_data = None
            else:
                this_frame_this_obj_3d_data = this_frame_3d_data[ zoom_cond_3d ]

        if (frame_enum%100)==0:
            print '%s: frame %d'%(datetime_str,frame)

        saved_fnames = []
        for movie_idx,ufmf_fname in enumerate(movie_fnames):
            try:
                frame_data = frame_dict[ufmf_fname]
            except KeyError:
                # no data saved (frame skip on Prosilica camera?)
                if movie_cam_ids is not None:
                    cam_id = movie_cam_ids[movie_idx]
                else:
                    cam_id = ufmf_tools.get_cam_id_from_ufmf_fname(ufmf_fname)
                camn = None
                if cam_id not in blank_images:
                    im_w, im_h = global_data['width_heights'][cam_id]
                    image = np.empty((im_h,im_w),dtype=np.uint8); image.fill(255)
                    blank_images[cam_id] = image
                image = blank_images[cam_id]
                mean_image = None
            else:
                cam_id = frame_data['cam_id']
                camn = frame_data['camn']
                image = frame_data['image']
                if config['what to show']['image_manipulation'] == 'absdiff':
                    mean_image = frame_data['mean']
                del frame_data
            save_fname = 'tmp_frame%07d_%s.png'%(frame,cam_id)
            save_fname_path = os.path.join(dest_dir, save_fname)

            pixel_aspect = config[cam_id].get('pixel_aspect',1)
            transform = config[cam_id].get('transform','orig')

            border_pixels = config['what to show']['border_pixels']

            if config['what to show']['max_resolution'] is not None:
                b2 = border_pixels*2
                fix_w, fix_h = config['what to show']['max_resolution']
                fix_aspect = (fix_w-b2)/float(fix_h-b2)
                desire_aspect = image.shape[1]/float(image.shape[0]*pixel_aspect)
                if desire_aspect >= fix_aspect:
                    # image is wider than resolution given
                    device_w = fix_w-b2
                    device_h = (fix_w-b2)/desire_aspect
                    device_x = border_pixels
                    device_y = (fix_h-device_h+border_pixels)/2.0
                else:
                    # image is taller than resolution given
                    device_h = fix_h-b2
                    device_w = (fix_h-b2)*desire_aspect
                    device_y = border_pixels
                    device_x = (fix_w-device_w+border_pixels)/2.0
                user_rect = (0,0,image.shape[1],image.shape[0])
            elif config['what to show']['zoom_obj']:
                if border_pixels != 0:
                    raise NotImplementedError()
                device_x = 0
                device_y = 0
                device_w = config['what to show']['zoom_orig_pixels']*config['what to show']['zoom_factor']
                device_h = device_w
                fix_w = device_w
                fix_h = device_h

                if this_frame_this_obj_3d_data is not None:
                    X = np.array([this_frame_this_obj_3d_data['x'],
                                  this_frame_this_obj_3d_data['y'],
                                  this_frame_this_obj_3d_data['z'],
                                  np.ones_like(this_frame_this_obj_3d_data['x'])]).T
                    xarr,yarr = R.find2d( cam_id, X, distorted = True )
                    assert len(xarr)==1
                    x=xarr[0]
                    y=yarr[0]
                    r = config['what to show']['zoom_orig_pixels']*0.5
                    user_rect = (x-r,y-r,r*2,r*2)
                else:
                    # we're not tracking object -- don't draw anything
                    user_rect = (-1000,-1000,10,10)
            else:
                device_x = border_pixels
                device_y = border_pixels
                device_w = image.shape[1]
                device_h = int(image.shape[0]*pixel_aspect) # compensate for pixel_aspect
                fix_w = device_w+2*border_pixels
                fix_h = device_h+2*border_pixels
                user_rect = (0,0,image.shape[1],image.shape[0])

            canv=benu.Canvas(save_fname_path,fix_w,fix_h)
            device_rect = (device_x,device_y,device_w,device_h)
            with canv.set_user_coords(device_rect, user_rect,
                                      transform=transform):
                if config['what to show']['image_manipulation'] == 'raw':
                    canv.imshow(image,0,0,cmap=colormap)
                if config['what to show']['image_manipulation'] == 'absdiff':
                    if mean_image is not None:
                        adsdiff_image = abs(image.astype(np.int16)-mean_image.astype(np.int16))
                        scaled_show = np.clip((5*adsdiff_image)+127, 0, 255).astype(np.uint8)
                        canv.imshow(scaled_show,0,0,cmap=colormap)
                if config['what to show']['show_2d_position'] and camn is not None:
                    cond = tracker_data['camn']==camn
                    this_cam_data = tracker_data[cond]
                    xarr = np.atleast_1d(this_cam_data['x'])
                    yarr = np.atleast_1d(this_cam_data['y'])
                    canv.scatter(xarr, yarr,
                                 color_rgba=(0,0,0,1),
                                 radius=10,
                                 markeredgewidth=config['what to show']['linewidth'],
                                 )
                    # draw shadow
                    canv.scatter(xarr+config['what to show']['linewidth'], yarr+config['what to show']['linewidth'],
                                 color_rgba=(1,1,1,1),
                                 radius=10,
                                 markeredgewidth=config['what to show']['linewidth'],
                                 )
                if config['what to show']['show_2d_orientation'] and camn is not None:
                    cond = tracker_data['camn']==camn
                    this_cam_data = tracker_data[cond]
                    xarr = np.atleast_1d(this_cam_data['x'])
                    yarr = np.atleast_1d(this_cam_data['y'])
                    slope = np.atleast_1d(this_cam_data['slope'])
                    thetaarr = np.arctan(slope)
                    line_len = 30.0
                    xinc = np.cos(thetaarr)*line_len
                    yinc = np.sin(thetaarr)*line_len/float(pixel_aspect)
                    for x,y,xi,yi in zip(xarr,yarr,xinc,yinc):
                        xarr = np.array([x-xi,x+xi])
                        yarr = np.array([y-yi,y+yi])
                        if np.any(np.isnan( xarr )) or np.any(np.isnan( yarr )):
                            continue
                        canv.plot( xarr, yarr,
                                  color_rgba=(0,1,0,0.4),
                                  linewidth=config['what to show']['linewidth'],
                                  )
                if config['what to show']['show_3d_smoothed_position'] and camn is not None:
                    if len(this_frame_3d_data):
                        X = np.array([this_frame_3d_data['x'], this_frame_3d_data['y'], this_frame_3d_data['z'], np.ones_like(this_frame_3d_data['x'])]).T
                        xarr,yarr = R.find2d( cam_id, X, distorted = True )
                        canv.scatter(xarr, yarr,
                                     color_rgba=(0,1,1,1),
                                     radius=10,
                                     markeredgewidth=config['what to show']['linewidth'],
                                     )

                if config['what to show']['show_3d_MLE_position'] and camn is not None:
                    if len(this_frame_raw_3d_data):
                        X = np.array([this_frame_raw_3d_data['x'], this_frame_raw_3d_data['y'], this_frame_raw_3d_data['z'], np.ones_like(this_frame_raw_3d_data['x'])]).T
                        xarr,yarr = R.find2d( cam_id, X, distorted = True )
                        canv.scatter(xarr, yarr,
                                     color_rgba=(0.2,0.2,0.5,1),
                                     radius=8,
                                     markeredgewidth=config['what to show']['linewidth'],
                                     )
                        # draw shadow
                        canv.scatter(xarr+config['what to show']['linewidth'], yarr+config['what to show']['linewidth'],
                                     color_rgba=(0.7,0.7,1,1), # blue
                                     radius=8,
                                     markeredgewidth=config['what to show']['linewidth'],
                                     )

                if config['what to show']['show_3d_raw_orientation'] and camn is not None:
                    if len(this_frame_raw_3d_data):
                        hzs = np.array([this_frame_raw_3d_data['hz_line0'],
                                        this_frame_raw_3d_data['hz_line1'],
                                        this_frame_raw_3d_data['hz_line2'],
                                        this_frame_raw_3d_data['hz_line3'],
                                        this_frame_raw_3d_data['hz_line4'],
                                        this_frame_raw_3d_data['hz_line5']]).T
                        Xs = np.array([this_frame_raw_3d_data['x'],
                                       this_frame_raw_3d_data['y'],
                                       this_frame_raw_3d_data['z']]).T
                        cam_center = R.get_camera_center( cam_id )[:,0]
                        for (X,hz,this_dataqual) in zip(Xs,hzs,this_frame_raw_dataqual):
                            if this_dataqual < min_ori_qual:
                                continue
                            cam_ray = geom.line_from_points( geom.ThreeTuple(cam_center), geom.ThreeTuple(X) )
                            raw_ori_line = geom.line_from_HZline(hz)
                            X_ = raw_ori_line.get_my_point_closest_to_line(cam_ray)

                            ld = raw_ori_line.direction()
                            dmag = abs(ld)
                            du = ld*(1./dmag) # unit length direction (normalize)

                            length = 0.5 # arbitrary, 0.5 meters
                            N = 100 # n segments (to deal with distortion)

                            X0 = X_.vals + du.vals*-length/2.0
                            X = X0[:,np.newaxis] + np.linspace(0,length,N)[np.newaxis,:]*du.vals[:,np.newaxis]
                            Xh = np.vstack( (X, np.ones_like( X[0,np.newaxis,:] ) ) ).T
                            xarr,yarr = R.find2d( cam_id, Xh, distorted = True )
                            canv.plot( xarr, yarr,
                                       color_rgba=(0,0,1,1), # blue
                                       linewidth=config['what to show']['linewidth'],
                                       )

                if config['what to show']['show_3d_smoothed_orientation'] and camn is not None:
                    if len(this_frame_3d_data):
                        for (row,ori_qual) in zip(this_frame_3d_data,this_frame_dataqual):
                            if ori_qual < min_ori_qual:
                                continue
                            X0 = np.array([row['x'], row['y'], row['z'], np.ones_like(row['x'])]).T
                            dx = np.array([row['dir_x'], row['dir_y'], row['dir_z'], np.zeros_like(row['x'])]).T
                            X1 = X0 + dx*orientation_3d_line_length
                            if np.any( np.isnan(X1) ):
                                continue
                            pts = np.vstack( [X0, X1] )
                            xarr,yarr = R.find2d( cam_id, pts, distorted = True )
                            canv.plot(xarr, yarr,
                                      color_rgba=(1,0,0,1), # red
                                      linewidth=config['what to show']['linewidth'],
                                      )

                if config['what to show']['show_3d_raw_chosen_orientation'] and camn is not None:
                    if len(this_frame_3d_data):
                        for (row,ori_qual) in zip(this_frame_3d_data,this_frame_dataqual):
                            if ori_qual < min_ori_qual:
                                continue
                            X0 = np.array([row['x'], row['y'], row['z'], np.ones_like(row['x'])]).T
                            dx = np.array([row['rawdir_x'], row['rawdir_y'], row['rawdir_z'], np.zeros_like(row['x'])]).T
                            X1 = X0 + dx*orientation_3d_line_length
                            if np.any( np.isnan(X1) ):
                                continue
                            pts = np.vstack( [X0, X1] )
                            xarr,yarr = R.find2d( cam_id, pts, distorted = True )
                            canv.plot(xarr, yarr,
                                      color_rgba=(1,159./255,0,1), # orange
                                      linewidth=config['what to show']['linewidth'],
                                      )

                if config['what to show']['obj_labels'] and camn is not None:
                    if len(this_frame_3d_data):
                        X = np.array([this_frame_3d_data['x'], this_frame_3d_data['y'], this_frame_3d_data['z'], np.ones_like(this_frame_3d_data['x'])]).T
                        xarr,yarr = R.find2d( cam_id, X, distorted = True )
                        for i in range(len(xarr)):
                            obj_id = this_frame_3d_data['obj_id'][i]
                            canv.text( '%d'%obj_id, xarr[i], yarr[i],
                                       font_size=14,
                                       color_rgba=(1,0,0,1) )

                if config['what to show']['show_3d_obj_position_text'] and camn is not None:
                    if len(this_frame_3d_data):
                        X = np.array([this_frame_3d_data['x'], this_frame_3d_data['y'], this_frame_3d_data['z'], np.ones_like(this_frame_3d_data['x'])]).T
                        xarr,yarr = R.find2d( cam_id, X, distorted = True )
                        for i in range(len(xarr)):
                            canv.text( '(%.1f, %.1f, %.1f) mm'%( X[i,0]*1000.0, X[i,1]*1000.0, X[i,2]*1000.0 ), xarr[i]+10, yarr[i],
                                       font_size=14,
                                       color_rgba=(0,1,1,1) )

                if config['what to show']['show_cam_id']:
                    canv.text( '%s'%cam_id, 0,20,
                               font_size=14,
                               color_rgba=(1,0,0,1) )

                if workaround_ffmpeg2theora_bug:
                    # first frame should get a colored pixel so that
                    # ffmpeg doesn't interpret the whole move as grayscale
                    canv.plot([0,1],[0,1],
                              color_rgba=(1,0,0,0.1),
                              )
                    workaround_ffmpeg2theora_bug = False # Now we already did it.

            canv.save()
            saved_fnames.append( save_fname_path )

        target = os.path.join(dest_dir, 'movie%s_frame%07d.png'%(
            datetime_str,frame_enum+1 ))
        # All cameras saved for this frame, make montage
        title = '%s frame %d'%( datetime_str, frame )
        montage(saved_fnames, title, target)
        all_frame_montages.append( target )
        if not no_remove:
            for fname in saved_fnames:
                os.unlink(fname)
    print '%s: %d frames montaged'%(datetime_str,len(all_frame_montages),)

    if save_ogv_movie:
        orig_dir = os.path.abspath(os.curdir)
        os.chdir(dest_dir)
        try:
            CMD = 'ffmpeg2theora -v 10 movie%s_frame%%07d.png -o movie%s.ogv'%(
                datetime_str,datetime_str)
            subprocess.check_call(CMD,shell=True)
        finally:
            os.chdir(orig_dir)

        if not no_remove:
            for fname in all_frame_montages:
                os.unlink(fname)

def main():
    # keep default config file in sync with get_config_defaults() above
    usage = """%prog DATAFILE2D.h5 [options]

The default configuration correspondes to a config file:

[what to show]
show_2d_position = False
show_2d_orientation = False
show_3d_MLE_position = False
show_3d_smoothed_position = False
show_3d_raw_orientation = False
show_3d_raw_chosen_orientation = False
show_3d_smoothed_orientation = False
show_3d_obj_position_text = False
minimum_display_orientation_quality = 0
white_background =  False
max_resolution = None
border_pixels = 0
zoom_obj = None
zoom_orig_pixels = 50
zoom_factor = 5
obj_labels = False
linewidth = 1.0
show_cam_id = False
image_manipulation = 'raw'

Config files may also have sections such as:

[cam7_1]
pixel_aspect=2 # each pixel is twice as wide as tall
transform='rot 180' # rotate the image 180 degrees (See transform
                    # keyword argument of
                    # :meth:`flydra_analysis.a2.benu.Canvas.set_user_coords`
                    # for all possible transforms.)

"""

    parser = OptionParser(usage)

    parser.add_option('-k', "--kalman-file", dest="kalman_filename",
                      type='string',
                      help=".h5 file with 3D kalman data (and reconstructor if not given)")

    parser.add_option('-R', "--reconstructor", dest="reconstructor",
                      type='string',
                      help="reconstructor used for computing 2D coordinates")

    parser.add_option("--dest-dir", type='string',
                      help="destination directory to save resulting files")

    parser.add_option("--ufmf-dir", type='string',
                      help="directory with .ufmf files")

    parser.add_option("--config", type='string',
                      help="configuration file name")

    parser.add_option("--max-n-frames", type='int', default=None,
                      help="maximum number of frames to save")

    parser.add_option("--start", type='int', default=None,
                      help="start frame")

    parser.add_option("--stop", type='int', default=None,
                      help="stop frame")

    parser.add_option("--ogv", action='store_true', default=False,
                      help="export .ogv video")

    parser.add_option('-n', "--no-remove", action='store_true', default=False,
                      help="don't remove intermediate images")

    parser.add_option('--movie-fnames', type='string', default=None,
                      help=("names of movie files (separator is %r; "
                            "don't autodiscover from .h5)"%(os.pathsep,)))

    parser.add_option('--movie-cam-ids', type='string', default=None,
                      help=("cam_ids of movie files (separator is %r; "
                            "don't autodiscover from .h5)"%(os.pathsep,)))

    parser.add_option("--verbose", action='store_true', default=False,
                      help="verbose mode (help understand autodiscovery)")

    parser.add_option('--colormap', type='string', default=None)

    parser.add_option( "--caminfo-h5-filename", type="string",
                       help="path of h5 file from which to load caminfo")

    parser.add_option( "--candidate", type="int", default=0,
                       help="when multiple auto-discovered movie options, the index to take")

    parser.add_option( "--nth-frame", type="int", default=1,
                       help="save every Nth frame")

    core_analysis.add_options_to_parser(parser)
    (options, args) = parser.parse_args()

    if len(args)<1:
        parser.print_help()
        return

    movie_fnames = options.movie_fnames
    if movie_fnames is not None:
        movie_fnames = movie_fnames.split( os.pathsep )

    movie_cam_ids = options.movie_cam_ids
    if movie_cam_ids is not None:
        movie_cam_ids = movie_cam_ids.split( os.pathsep )

    reconstructor = options.reconstructor

    h5_filename = args[0]
    kwargs = core_analysis.get_options_kwargs(options)
    make_montage( h5_filename,
                  kalman_filename = options.kalman_filename,
                  cfg_filename = options.config,
                  ufmf_dir = options.ufmf_dir,
                  dest_dir = options.dest_dir,
                  save_ogv_movie = options.ogv,
                  no_remove = options.no_remove,
                  max_n_frames = options.max_n_frames,
                  start = options.start,
                  stop = options.stop,
                  movie_fnames = movie_fnames,
                  movie_cam_ids = movie_cam_ids,
                  caminfo_h5_filename = options.caminfo_h5_filename,
                  colormap = options.colormap,
                  candidate_index = options.candidate,
                  verbose = options.verbose,
                  nth_frame = options.nth_frame,
                  reconstructor = reconstructor,
                  **kwargs)


from   __future__ import print_function
from   itertools import islice, product
import logging
import MDAnalysis as md
import math
import random
import numpy as np
import pandas as pd
import plotly
import plotly.graph_objs as go
import subprocess
import scipy
import scipy.stats
import string
import time

import settings

class Atom(object):

    def __init__(self, identifier, **kwargs):
        self.id = identifier
        self.type = kwargs.get('type', None)
        self.element = kwargs.get('element', None)
        self.xyz = kwargs.get('xyz', None)
        self.stress = kwargs.get('stress', None)
        self.normal = kwargs.get('normal', False)
        self.distance = None
        self.sin_theta = None
        self.cos_theta = None
        self.sin_phi   = None
        self.cos_phi   = None
        self.spherical_stress = None
        self.voro_volume = 0

    def calc_spherical_stress(self):
        """
        Calculate spherical stress tensor from cartesian one
        ref: http://www.brown.edu/Departments/Engineering/Courses/En221/Notes/Polar_Coords/Polar_Coords.htm
        """
        xx, yy, zz, xy, xz, yz = self.stress
        cart = np.array( [ [xx, xy, xz], [xy, yy, yz], [xz, yz, zz] ] )

        # 1 for theta, the angle between xyz and z axis, 2 for phi,
        # angle between x axis and the projection on xy-plane
        sin1 = self.sin_theta
        cos1 = self.cos_theta
        sin2 = self.sin_phi
        cos2 = self.cos_phi

        conv = np.array( [ [sin1*cos2, cos1*cos2, -sin2],
                           [sin1*sin2, cos1*sin2, -cos2],
                           [cos1,      -sin1,     0], ] )

        sphe = np.dot( conv, cart.dot( np.transpose(conv) ) )

        # Of format [ [rr, rTheta, rPhi], [rTheta, thetaTheta, thetaPhi], [rPhi, thetaPhi, phiPhi] ]
        self.spherical_stress = sphe

class Box(object):

    PI = 3.1415926

    def __init__(self, timestep=0, radius=None, use_atomic_volume=True, average_on_atom=False, **kwargs):
        # Current timestep.
        self.timestep = timestep
        # Maximum bubble radius in box.
        self.radius = radius
        self.count  = 0
        # XYZ boundaries.
        self.bx = kwargs.get('bx', None)
        self.by = kwargs.get('by', None)
        self.bz = kwargs.get('bz', None)
        # Bubble center coordinates.
        self._center = kwargs.get('center', None)
        # All atoms.
        self.atoms = []
        # Container of atoms for each element.
        self._elements = {}
        # Container of shell stress for each element.
        self._shell_stress       = {}
        self._shell_stress_r     = {}
        self._shell_stress_theta = {}
        self._shell_stress_phi   = { }
        # Container of shell atom count for each element.
        self.nbins = None
        self._shell_atoms = {}
        self._shell_atom_objs = []
        self._shell_volumes = {}
        # Indicator of stats status.
        self._stats_finished = False
        self._measured       = False
        # Dump atom coordinates to calculate voro tessellation volume
        self.voro_file_name = 'atom_coors'
        self.use_atomic_volume  = use_atomic_volume
        self.average_on_atom    = average_on_atom

    @property
    def measured(self):
        """Returns true if all atoms have a distance (to bubble center)."""
        if all([x.distance for x in self.atoms]):
            self._measured = True
        else:
            self._measured = False
        return self._measured

    @property
    def center(self):
        return self._center

    @center.setter
    def center(self, coor):
        self._center         = coor
        self._measured       = False
        self._stats_finished = False

    def combine_water_atoms(self):
        """
        Combine H and O together into a new particle
        stress = S_h + S_o
        coor   = center of mass
        The sequency of H and O atoms are O H H
        """
        self._old_atoms = self.atoms
        self.atoms = []

        self._old_atoms.sort( key=lambda x: x.id )
        water = []
        for atom in self._old_atoms:
            if atom.element not in ['H', 'O']:
                self.atoms.append( atom )
            else:
                water.append(atom)
                if len( water ) == 3:
                    # need to combine the 3 atoms into 1 now
                    assert [ _ele.element for _ele in water ] == ['O', 'H', 'H']
                    new_stress = [a+b+c for a, b, c in zip(water[0].stress, water[1].stress, water[2].stress)]
                    new_volume = sum( _ele.voro_volume for _ele in water )
                    masses = [ 16 if _ele.element == 'O' else 1 for _ele in water ]
                    xs     = [ _ele.xyz[0] for _ele in water]
                    ys     = [ _ele.xyz[ 1 ] for _ele in water ]
                    zs     = [ _ele.xyz[ 2 ] for _ele in water ]

                    cx = sum( m*x for m,x in zip(masses, xs) ) / sum(masses)
                    cy = sum( m * y for m, y in zip( masses, ys ) ) / sum( masses )
                    cz = sum( m * z for m, z in zip( masses, zs ) ) / sum( masses )
                    new_xyz = (cx, cy, cz)

                    new_id = water[0].id
                    normal = water[0].normal

                    self.atoms.append( Atom(new_id, type=3, element='H', xyz=new_xyz, stress=new_stress, normal=normal) )

                    water = []

    def dump_atoms_for_voro( self, length=None ):
        '''
        Dump atom coordinates so we can calculate Voronoi tessellation using Voro++
        from http://math.lbl.gov/voro++/
        The input file format for voro++ is
        <atom id> <x> <y> <z>
        and output file format is
        <atom id> <x> <y> <z> <tessellation volume>
        '''
        logging.info( 'Dump atom coordinates to {}'.format( self.voro_file_name ) )
        fmt = '{} {} {} {}\n'
        if length:
            xmin, xmax = self.center[0] - length, self.center[0] + length
            ymin, ymax = self.center[1] - length, self.center[1] + length
            zmin, zmax = self.center[2] - length, self.center[2] + length

        with open( self.voro_file_name, 'w' ) as output:
            for atom in self.atoms:
                x, y, z = atom.xyz
                if length:
                    if xmin <= x <= xmax and ymin<= y <= ymax and zmin <= z <= zmax:
                        output.write( fmt.format( atom.id, x, y, z ) )
                else:
                    output.write( fmt.format( atom.id, x, y, z ) )

    def voro_cmd( self, gnuplot=False, length=None ):
        '''
        CMD to run voro++ in bash
        gnuplot=True will also export gnu plot file. Be careful when system is large as
        this file will be extremely large
        default to use -o to preserve the atom order. This has small memory and performance
        impact as the documentation says.
        '''

        # when have length -o will not work
        cmd = 'voro++' if length else 'voro++ -o'
        fmt = cmd + ' {opts} {{xmin}} {{xmax}} {{ymin}} {{ymax}} {{zmin}} {{zmax}} {{infile}}'
        opts = '-g' if gnuplot else ''

        fmt = fmt.format( opts=opts )

        if length:
            xmin, xmax = self.center[0] - length, self.center[0] + length
            ymin, ymax = self.center[1] - length, self.center[1] + length
            zmin, zmax = self.center[2] - length, self.center[2] + length
        else:
            xmin, xmax = self.bx
            ymin, ymax = self.by
            zmin, zmax = self.bz

        return fmt.format( xmin=xmin, xmax=xmax,
                           ymin=ymin, ymax=ymax,
                           zmin=zmin, zmax=zmax,
                           infile=self.voro_file_name)

    def run_voro_cmd( self, gnuplot=False, length=None ):
        logging.info( 'Calculating voro volumes for atoms' )
        cmd = self.voro_cmd( gnuplot=gnuplot, length=length )

        logging.info( "Running: {}".format( cmd ))

        sp = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        out, err = sp.communicate()
        if err:
            raise Exception(err)

        logging.info( "Finished: {}".format( cmd ) )

    def read_voro_volumes( self ):
        voro_out = self.voro_file_name + '.vol'
        logging.info( 'Reading voro volumes from {}'.format( voro_out ) )
        with open( voro_out, 'r' ) as volumes:
            idx = 0
            for line in volumes:
                atom_id, x, y, z, vol = [ float(ele) for ele in line.split() ]
                atom_id = int( atom_id )
                atom = self.atoms[ idx ]
                try:
                    assert( atom.id == atom_id )
                except Exception as e:
                    print( atom.id, atom_id )
                    raise e
                atom.voro_volume = vol
                idx += 1

    def calc_voro_volumes( self, gnuplot=False, length=None ):
        ''' Calculate voro tessellation volume using voro '''
        self.dump_atoms_for_voro( length=length )
        self.run_voro_cmd( gnuplot=gnuplot, length=length )
        if not length:
            self.read_voro_volumes()

    def adjust_water_vol(self, ratio=(0.5, 0.25)):
        """ Adjust volume of H and O in water. For pure water system only """
        satoms = sorted( self.atoms, key= lambda x: x.id)
        assert( len( satoms ) % 3 == 0 )
        assert( ratio[0] + 2 * ratio[1] == 1.0)
        for idx in xrange( len(satoms) / 3):
            o = satoms[ idx * 3 ]
            h1 = satoms[ idx * 3 + 1 ]
            h2 = satoms[ idx * 3 + 2 ]

            vsum = sum( ele.voro_volume for ele in [o, h1, h2])
            vo = ratio[0] * vsum
            vh = ratio[1] * vsum

            o.adj_vol = vo
            h1.adj_vol = vh
            h2.adj_vol = vh

    def set_boundary(self, bx, by, bz):
        """Set bx by bz together."""
        self.bx = bx
        self.by = by
        self.bz = bz

    def add_atom(self, atom):
        self.atoms.append(atom)
        self.count += 1
        # Need to run stats after new atom added.
        self._stats_finished = False
        if atom.element in self._elements:
            self._elements[atom.element].append(atom)
        else:
            self._elements[atom.element] = [atom]

    def measure(self):
        """Measure distance to bubble center for each atom."""
        for atom in self.atoms:
            coor          = np.array(atom.xyz)
            atom.distance = np.linalg.norm(coor - self.center)

            if atom.normal:
                # Calculate sin cos for theta and phi
                dx = coor[0] - self.center[0]
                dy = coor[1] - self.center[1]
                dz = coor[2] - self.center[2]

                xy_square = math.sqrt(dx*dx + dy*dy)

                atom.sin_theta = xy_square / atom.distance
                atom.cos_theta = dz / atom.distance

                atom.sin_phi   = dy / xy_square
                atom.cos_phi   = dx / xy_square

        self.calc_voro_volumes()

    def stats(self, dr, normal):
        """
        System stats.
        Generate data for atom stats and stress stats for each element.
        self._shell_atoms = {}
        self._shell_stress = {}
        """
        if not self.measured:
            raise AtomUnmeasuredError("Some atoms are unmeasuerd")

        self.nbins = int(math.ceil(self.radius / float(dr)))
        self._shell_atom_objs = [ { } for x in range( self.nbins ) ]

        for ele, atoms in self._elements.iteritems():
            # Do stats for each element.
            for atom in atoms:
                if atom.distance < self.radius:
                    shell_idx = int( atom.distance / dr )
                    self._shell_atom_objs[ shell_idx ].setdefault(ele, []).append( atom )

                    if normal:
                        atom.calc_spherical_stress()

        self._stats_finished = True

    def atom_stats(self, element, dr):
        """Atom ratio stats inside bubble."""
        if not self._stats_finished:
            self.stats(dr)
        nbins = len(self._shell_atoms[element])
        bubble_atoms = {}
        # Init bubble atoms by copying shell atoms
        for ele, count in self._shell_atoms.iteritems():
            bubble_atoms[ele] = [x for x in count]
            for i in range(1, nbins):
                bubble_atoms[ele][i] += bubble_atoms[ele][i - 1]
            bubble_atoms[ele] = np.array(bubble_atoms[ele])

        return bubble_atoms[element] / sum(bubble_atoms.values())

    def pressure_stats(self, elements, dr):
        """Average pressure stats inside bubble for species in elements."""
        if not self._stats_finished:
            self.stats(dr)

        nbins = len(self._shell_stress[elements[0]])
        # Calculate stress for all element in elements as whole.
        # Convert numpy.Array to mutable list.
        stress_in  = [x for x in sum([self._shell_stress[ele] for ele in elements])]
        stress_out = [x for x in stress_in]
        for i in range(1, nbins):
            # Cumulative stress.
            stress_in[i] += stress_in[i-1]
            stress_out[nbins - 1 - i] += stress_out[nbins - i]
        for i in range(1, nbins):
            # Stress -> pressure.
            stress_in[i] = 0 - stress_in[i] / self.vol_sphere((i+1)*dr) / 3.0
            stress_out[nbins-1-i] = 0 - stress_out[nbins-1-i] / (self.vol_sphere(self.radius) - self.vol_sphere((nbins-i-1)*dr)) / 3
        # Head and tail.
        stress_in[0] = 0 - stress_in[0] / self.vol_sphere(dr) / 3
        stress_out[nbins - 1] = 0 - stress_out[nbins - 1] / (self.vol_sphere(self.radius) - self.vol_sphere((nbins - 1)*dr)) / 3
        return {'in': stress_in, 'out': stress_out}

    def shell_pressure_stats(self, elements, dr, normal=False):
        """Average pressure of elements inside shell."""
        self.stats(dr, normal=normal)

        print( "NNNNNumber of bins: {}".format(self.nbins) )

        if not normal:
            # atom.stress has 3 elements, xx yy zz components
            if self.use_atomic_volume:
                if self.average_on_atom:
                    # atomic volume is used, pressure is calculated for each atom and then averaged together
                    stress = []
                    for idx, shell_atoms in enumerate(self._shell_atom_objs):
                        pressure_raw = {}
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                # P = -(S_xx + S_yy + S_zz)/3/V
                                pressure_raw[element] = [ - sum(atom.stress)/atom.voro_volume/3.0 for atom in atoms ]
                        # Average pressure = sum(Pressure)/n_atoms
                        n_atoms = sum( len(_ele) for _ele in pressure_raw.values() )
                        if n_atoms != 0:
                            pressure_ave = sum( sum(_ele) for _ele in pressure_raw.values() ) / n_atoms
                        else:
                            pressure_ave = 0
                        stress.append(pressure_ave)
                    return stress
                else:
                    # pressure is calculated as sum(atom stress in a shell) / sum(atom volume in a shell)
                    stress = []
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        stress_all = 0
                        volume_all = 0
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                stress_all += sum( sum(atom.stress[:3]) for atom in atoms )
                                volume_all += sum( atom.voro_volume for atom in atoms )
                        if volume_all != 0:
                            pressure_ave = - stress_all / 3.0 / volume_all
                        else:
                            pressure_ave = 0
                        stress.append( pressure_ave )
                    return stress
            else:
                # use shell volume
                stress = [ ]
                for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                    r_min, r_max = idx * dr, (idx + 1)*dr
                    stress_all = 0
                    volume_all = self.vol_sphere(r_max) - self.vol_sphere(r_min)
                    for element, atoms in shell_atoms.iteritems():
                        if element in elements:
                            stress_all += sum( sum( atom.stress[:3] ) for atom in atoms )
                    pressure_ave = - stress_all / 3.0 / volume_all
                    stress.append( pressure_ave )
                return stress
        else:
            # normal pressure, atom.spherical_stress has 6 items: xx, yy, zz, xy, xz, yz.
            stress_r     = []
            stress_theta = []
            stress_phi   = []

            if self.use_atomic_volume:

                if self.average_on_atom:
                    # Pressure is calculate as average of pressure on each atom
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        pressure_r_raw     = {}
                        pressure_theta_raw = {}
                        pressure_phi_raw   = {}
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                pressure_r_raw[element]     = [ - atom.spherical_stress[0][0] / atom.voro_volume for atom in atoms ]
                                pressure_theta_raw[element] = [ - atom.spherical_stress[1][1] / atom.voro_volume for atom in atoms ]
                                pressure_phi_raw[element]   = [ - atom.spherical_stress[2][2] / atom.voro_volume for atom in atoms ]

                        n_atoms = sum( len( _ele ) for _ele in pressure_r_raw.values() )
                        if n_atoms != 0:
                            pressure_r_ave     = sum( sum(_ele) for _ele in pressure_r_raw.values() ) / n_atoms
                            pressure_theta_ave = sum( sum(_ele) for _ele in pressure_theta_raw.values() ) / n_atoms
                            pressure_phi_ave   = sum( sum(_ele) for _ele in pressure_phi_raw.values() ) / n_atoms
                        else:
                            pressure_r_ave = pressure_theta_ave = pressure_phi_ave = 0

                        stress_r.append( pressure_r_ave )
                        stress_theta.append( pressure_theta_ave )
                        stress_phi.append( pressure_phi_ave )
                    return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }

                else:
                    # Pressure is calculated as sum(stress)/sum(atomic_volume)
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        stress_r_all     = 0
                        stress_theta_all = 0
                        stress_phi_all   = 0
                        volume_all       = 0

                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                stress_r_all     += sum( atom.spherical_stress[0][0] for atom in atoms )
                                stress_theta_all += sum( atom.spherical_stress[1][1] for atom in atoms )
                                stress_phi_all   += sum( atom.spherical_stress[2][2] for atom in atoms )
                                volume_all       += sum( atom.voro_volume for atom in atoms )
                        if volume_all != 0:
                            pressure_r_ave     = - stress_r_all / volume_all
                            pressure_theta_ave = - stress_theta_all / volume_all
                            pressure_phi_ave   = - stress_phi_all / volume_all
                        else:
                            pressure_r_ave = pressure_theta_ave = pressure_phi_ave = 0

                        stress_r.append( pressure_r_ave )
                        stress_theta.append( pressure_theta_ave )
                        stress_phi.append( pressure_phi_ave )
                    return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }
            else:
                # Use shell volume
                for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                    r_min, r_max = idx * dr, (idx+1) * dr
                    stress_r_all = 0
                    stress_theta_all = 0
                    stress_phi_all = 0
                    volume_all = self.vol_sphere(r_max) - self.vol_sphere(r_min)

                    for element, atoms in shell_atoms.iteritems():
                        if element in elements:
                            stress_r_all += sum( atom.spherical_stress[ 0 ][ 0 ] for atom in atoms )
                            stress_theta_all += sum( atom.spherical_stress[ 1 ][ 1 ] for atom in atoms )
                            stress_phi_all += sum( atom.spherical_stress[ 2 ][ 2 ] for atom in atoms )

                    pressure_r_ave = - stress_r_all / volume_all
                    pressure_theta_ave = - stress_theta_all / volume_all
                    pressure_phi_ave = - stress_phi_all / volume_all

                    stress_r.append( pressure_r_ave )
                    stress_theta.append( pressure_theta_ave )
                    stress_phi.append( pressure_phi_ave )
                return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }

    def pressure_between(self, rlow, rhigh):
        """Return the average pressure and number of atoms between rlow
        and rhigh."""
        stress = 0
        count = 0
        for atom in self.atoms:
            if atom.distance > rlow and atom.distance <= rhigh:
                count += 1
                stress += sum(atom.stress)
        volume = self.vol_sphere(rhigh) - self.vol_sphere(rlow)
        return stress / volume / 3, count

    def shell_density(self, elements, mole, dr):
        """Shell density for species inside elements.
        mole unit - g/cm^3
        dr unit - angstrom
        """
        # Usually density_dr is different from stats_dr.
        self.stats(dr)
        # Avogadro constant. Modified by coefficient used to
        # convert angstrom^3 to cm^3.
        NA = 6.022 / 10
        nbins = len(self._shell_atoms[elements[0]])
        # Calculate atom count for all species in elements as whole.
        # Convert numpy.Array to mutable list.
        count = [x for x in sum([self._shell_atoms[ele] for ele in elements])]
        # Calculate density.
        for i in range(nbins):
            r_low = i * dr
            r_high = r_low + dr
            # Volume unit is Angstrom^3.
            volume = self.vol_sphere(r_high) - self.vol_sphere(r_low)
            count[i] = count[i] / NA / volume
        return count

    def bubble_density(self, elements, mole, dr):
        pass

    def xyz_density(self, elements, mole, dx):
        """Density distribution along x, y, and z inside box."""
        # Avogadro constant. Modified by coefficient used to
        # convert angstrom^3 to cm^3.
        NA = 6.022 / 10
        nx = int(math.ceil((self.bx[1] - self.bx[0]) / dx))
        ny = int(math.ceil((self.by[1] - self.by[0]) / dx))
        nz = int(math.ceil((self.bz[1] - self.bz[0]) / dx))
        dist = {}
        dist['x'] = [0 for x in range(nx)]
        dist['y'] = [0 for y in range(ny)]
        dist['z'] = [0 for z in range(nz)]

        for ele in elements:
            # Count atoms.
            for atom in self._elements[ele]:
                dist['x'][int(atom.xyz[0] / dx)] += 1
                dist['y'][int(atom.xyz[1] / dx)] += 1
                dist['z'][int(atom.xyz[2] / dx)] += 1

        volx = (self.by[1] - self.by[0]) * (self.bz[1] - self.bz[0]) * dx
        voly = (self.bx[1] - self.bx[0]) * (self.bz[1] - self.bz[0]) * dx
        volz = (self.by[1] - self.by[0]) * (self.bx[1] - self.bx[0]) * dx

        for i in range(nx):
            # Calculate density.
            dist['x'][i] = dist['x'][i] / NA / volx
            dist['y'][i] = dist['y'][i] / NA / voly
            dist['z'][i] = dist['z'][i] / NA / volz
        return dist

    def vol_sphere(self, r):
        """Volume of sphere with radius r."""
        return 4.0/3 * Box.PI * (r ** 3)

    def volume(self):
        """ Box volume """
        return (self.bx[1] - self.bx[0]) * (self.by[1] - self.by[0]) * (self.bz[1] - self.bz[0])

class Trajectory( object ):
    '''Gas molecule trajectory class'''
    def __init__( self, pdbPath, xtcPath ):
        self.universe = md.Universe( pdbPath, xtcPath )
        self.set_density_params()

    @property
    def n_frames( self ):
        return self.universe.trajectory.n_frames

    @property
    def frame( self ):
        return self.universe.trajectory.frame

    def set_density_params(self, low=0.4, high=0.5, length=60 ):
        '''
        Generate grid with length of dnesity_grid_length at x,y,z directions.
        Grids whose density are between low * max_density and high * max_density
        will be used for radius calculation. d
        '''
        self.density_low  = low
        self.density_high = high
        self.density_grid_length = length

    def set_frame( self, frame ):
        self.universe.trajectory[ frame ]

    def radius( self, frame ):
        '''
        Bubble radius at one frame.
        Method:
        1. Load the snapshot at frame
        2. Load x, y, z coordinates
        3. Calculate density grid mesh at grid points
        4. Filter the shell grids with density between low * max density and high * max density
        5. Calculate the average radius
        '''
        start = time.clock()

        self.set_frame( frame )

        # Load x, y, z coordinates
        data = pd.DataFrame( list(self.universe.coord), columns=['x','y','z'])
        x    = data[ 'x' ].values
        y    = data[ 'y' ].values
        z    = data[ 'z' ].values

        # Density grid
        xyz  = scipy.vstack( [ x, y, z ] )
        kde  = scipy.stats.gaussian_kde( xyz )
        xmin, ymin, zmin = x.min(), y.min(), z.min()
        xmax, ymax, zmax = x.max(), y.max(), z.max()
        NI         = complex( imag=self.density_grid_length)
        xi, yi, zi = scipy.mgrid[ xmin:xmax:NI, ymin:ymax:NI, zmin:zmax:NI ]
        coords     = scipy.vstack([item.ravel() for item in [xi, yi, zi]])
        density    = kde(coords).reshape(xi.shape)

        # Filter density grid
        density_max  = density.max()
        density_low  = self.density_low * density_max
        density_high = self.density_high * density_max

        xyzs = []
        N = self.density_grid_length
        for idx, idy, idz in product( xrange(N), xrange(N), xrange(N) ):
            if density_low < density[ idx, idy, idz ] <= density_high:
                xyzs.append( [ xi[ idx, idy, idz ], yi[ idx, idy, idz ], zi[ idx, idy, idz ] ] )
        xyzs = np.array( xyzs )

        # Average radius
        center = xyzs.mean( axis=0 )
        rs = []
        for xyz_ele in xyzs:
            rs.append( np.linalg.norm( center - xyz_ele ) )

        duration = time.clock() - start
        print( "Radius for frame {} calculated in {:.2f} seconds".format( frame, duration ) )

        return center, scipy.mean( rs )

    def radius_for_frames( self, start, end, step=1 ):
        ret = []
        for frame in xrange( start, end, step ):
            center, radius = self.radius( frame )
            ret.append( [ frame, radius ] )
        return ret

    def all_radius( self ):
        return self.radius_for_frames( 0, self.n_frames, 1 )

    def regression( self, radiusList ):
        ''' Input (frame, radius) lists and do linear regression on the data '''
        ts = [ ele[0] for ele in radiusList ]
        rs = [ ele[1] for ele in radiusList ]
        slope, intercept, r_value, p_value, std_err = scipy.stats.linregress( ts, rs )
        return slope, intercept, r_value, p_value, std_err

    def plot_radius( self, rs, notebook=False ):
        ''' plot dots and linear regression results '''

        xs = [ ele[0] for ele in rs ]
        ys = [ ele[1] for ele in rs ]

        x_min = min( xs )
        x_max = max( xs )
        x_min = x_min - ( x_max - x_min ) * 0.05
        x_max = x_max + ( x_max - x_min ) * 0.05

        slope, intercept, r_value, p_value, std_err = self.regression( rs )

        xs_line = [ x_min ] + xs + [ x_max ]
        ys_line = [ ele * slope + intercept for ele in xs_line ]

        # Scatter plot
        scatter = go.Scatter(
            x = [ele[0] for ele in rs],
            y = [ele[1] for ele in rs],
            mode = 'markers',
            name = 'Radius'
            )

        reg_line = go.Scatter(
            x = xs_line, y = ys_line,
            mode='lines', name='y={:.4f}x+{:.4f}, p-value={:.2f}, StdErr={:.3f}'.format(slope, intercept, p_value, std_err)
            )

        data = go.Data([scatter, reg_line])

        plot = plotly.offline.iplot if notebook else plotly.offline.plot

        plot( {
            'data': data,
            'layout': go.Layout( title='Radius vs Frame', xaxis={'title':'Frame'}, yaxis={'title':'Radius'} )
            } )

    def flux_info( self, start, end, step=1 ):
        '''
        Flux info for frames [start:end:step]. Info are, for each step,
        nframe, center, radius, n atoms inside sphere
        '''
        info = []
        for nframe in xrange( start, end, step ):
            center, radius = self.radius( nframe )

            # Selector for AtomGroup in MDAnalysis
            selector = 'point ' + ' '.join( str( ele ) for ele in list( center ) + [ radius ] )

            # Explicitly set frame here
            self.set_frame( nframe )
            atoms = self.universe.select_atoms( selector )
            natoms = atoms.n_atoms
            info.append( (nframe, center, radius, natoms) )
        return info

#################################################
################# Exceptions ####################
#################################################

class AtomUnmeasuredError(Exception):
    pass

################################################
################## Functions ###################
################################################

def next_n_lines(file_opened, N, strip='right'):
  strip_dic = {
    'right': string.rstrip,
    'left': string.lstrip,
    'both': string.strip
    }
  if strip:
    return [strip_dic[strip](x) for x in islice(file_opened, N)]
  else:
    return list(islice(file_opened, N))

def read_stress(stress_file, N=settings.NLINES, normalPressure=False):
    """
    Read dump file into a list of atoms, which have type / coordinates /
    stresses info stored as Atom properties.
    Dump file data format:
    atom_id atom_type x y z stress_x stress_y stress_z
    """
    atoms = {}
    count = 0
    data = next_n_lines(stress_file, N)[9:]
    while data:
        atoms[count] = []
        for line in data:
            line = line.strip().split()
            identifier = int(line[0])
            atom_type = int(line[1])
            element = settings.ELEMENTS[atom_type]
            xyz = tuple([float(x) for x in line[2:5]])
            if normalPressure:
                # To calculate normal pressure, we need xx, yy, zz, xy, xz, yz
                stress = tuple([float(x) for x in line[5:11]])
            else:
                # To calculate pressure, we need xx, yy, zz
                stress = tuple([float(x) for x in line[5:8]])
            atom = Atom(identifier, type=atom_type, element=element, xyz=xyz, stress=stress, normal=normalPressure)
            atoms[count].append(atom)
        # Process next N lines.
        data = next_n_lines(stress_file, N)[9:]
        count += 1
    return atoms

def read_pdb(filename):
    """
    Read pdb file as a list of atoms
    """
    logging.info( "Reading {}".format(filename) )
    atoms_lines = []
    with open(filename, 'r') as pdbfile:
        for line in pdbfile:
            if line.startswith('CRYST'):
                cryst_line = line
            elif line.startswith('ATOM'):
                atoms_lines.append( line )

    x, y, z = [float(ele) for ele in cryst_line.strip().split()[1:4] ]

    atoms = []
    for line in atoms_lines:
        data = line.strip().split()
        idx = int(data[1])
        element = data[2][:2]
        coor = [ float(ele) for ele in data[5:8] ]
        atoms.append( Atom(identifier=idx, element=element, xyz=coor) )
    return atoms, (x,y,z)

def combine_water(atoms, remove=True):
    """
    Combine water atoms
    """
    combined = []
    ne = [ ele for ele in atoms if ele.element == 'Ne' ]
    wat = [ele for ele in atoms if ele.element != 'Ne' ]
    logging.info("Before:: {} Ne, {} Water atoms".format(len(ne), len(wat)))
    idx_wat = len(ne) + 1
    comb_wat = []
    for idx in range( len( wat ) / 3 ):
        coor1 = np.array( wat[ idx * 3 ].xyz )
        coor2 = np.array( wat[ idx * 3 + 1 ].xyz )
        coor3 = np.array( wat[ idx * 3 + 2 ].xyz )
        coor  = (coor1 + coor2 + coor3) / 3.
        comb_wat.append(Atom(identifier=idx_wat, element='W', xyz=coor))
        idx_wat += 1
    if remove:
        selected = random.sample(comb_wat, len(comb_wat)/4)
    else:
        selected = comb_wat
    n_ne = len(ne)
    for idx in xrange(len(selected)):
        selected[idx].id = idx + 1 + n_ne
    logging.info("After:: {} Ne, {} Water atoms".format(len(ne), len(selected)))
    return ne + selected

def write_lammps_data(atoms, xyz, filename):
    """
    LAMMPS data
    format: atom idx, molecule idx, atom type, x, y, z,
    """
    atom_types = {'Ne':1, 'W':2}
    x, y, z = xyz
    header = "LAMMPS bubble\n\n" \
             "{n_atoms} atoms\n\n" \
             "{n_types} atom types\n" \
             "0 bond types\n" \
             "0 angle types\n\n" \
             "0 {x} xlo xhi\n0 {y} ylo yhi\n0 {z} zlo zhi\n\n"\
             "Atoms\n\n".format(n_atoms=len(atoms), n_types=2,x=x,y=y,z=z)
    print(header)

    fmt = "{idx} {mol} {atype} {charge} {x} {y} {z}\n"

    for idx, atom in enumerate(atoms):
        header += fmt.format(idx=atom.id, mol=atom.id, atype=atom_types[atom.element], charge=0, x=atom.xyz[0], y=atom.xyz[1], z=atom.xyz[2])

    with open(filename, 'w') as output:
        output.write(header)

def average_atom_stress(write=True, step=0, *args):
    """Calculates averaged stress from multiple stress files.
    write determines whether to write output or not.
    step determines which timestep to average."""
    n_files = float(len(args))
    stress_list = []
    for ele in args:
        stress_list.append(read_stress(ele)[step])
        # Sort atoms by id.
        stress_list[-1].sort(key=lambda x: x.id)
    n_atoms = len(stress_list[0])
    atoms = []
    # Average stress for each atom id.
    for i in range(n_atoms):
        sx = sum([x[i].stress[0] for x in stress_list]) / n_files
        sy = sum([x[i].stress[1] for x in stress_list]) / n_files
        sz = sum([x[i].stress[2] for x in stress_list]) / n_files
        atom = stress_list[0][i]
        atoms.append(
            Atom(atom.id, type=atom.type, element=atom.element, xyz=atom.xyz, stress=(sx, sy, sz))
            )
    # Write averaged stress to file.
    if write:
        out_name = '.'.join(args[0].name.split('.')[:-1]) + '_averaged.dat'
        with open(out_name, 'w') as output:
            # Write header lines to be compatitable with LAMMPS dump files.
            output.write('Header line\n' * 9)
            for atom in atoms:
                # Do not write element here to be compatitable with
                # LAMMPS dump files.
                output.write("{} {} {} {} {} {} {} {}\n".format(
                    atom.id, atom.type,
                    atom.xyz[0], atom.xyz[1], atom.xyz[2],
                    atom.stress[0], atom.stress[1], atom.stress[2]))
        print("Average Stress saved to {}.".format(out_name))
    return atoms

def build_box(atoms, timestep, radius, center, use_atomic_volume, average_on_atom, bx, by, bz):
    """Build a box from a list of atoms."""
    box = Box(timestep, radius=radius, center=center, use_atomic_volume=use_atomic_volume, average_on_atom=average_on_atom)
    for atom in atoms:
        box.add_atom(atom)
    box.set_boundary(bx=bx, by=by, bz=bz)
    box.measure()
    return box

def write_density(density, dr, outname, header):
    """Write density (both shell and xyz density) stats to output file.
    One density list at a time.
    """
    with open(outname, 'w') as output:
        output.write(header)
        for i, item in enumerate(density):
            low = i * dr
            high = low + dr
            output.write('{l:.3f}\t{h:.3f}\t{d:.13f}\n'.format(l=low, h=high, d=item))

def write_pressure(pressure, dr, outname, header, bubble=False):
    """Write pressure (both bubble and shell pressure) stats to output file.
    If bubble is True, r_low is always zero.
    """
    logging.info( "Writing output to {}".format(outname) )
    if bubble:
        # Bubble pressure has in pressure and out pressure.
        with open(outname, 'w') as output:
            output.write(header)
            nbins = len(pressure['in'])
            for i in range(nbins):
                low = 0
                high = (i + 1) * dr
                if i < nbins - 1:
                    output.write('{l:.3f}\t{h:.3f}\t{pin:.13f}\t{pout:.13f}\n'.format(
                        l=low, h=high,
                        pin=pressure['in'][i], pout=pressure['out'][i+1]
                        ))
                else:
                    output.write('{l:.3f}\t{h:.3f}\t{pin:.13f}\t{pout:.13f}\n'.format(
                        l=low, h=high,
                        pin=pressure['in'][i], pout=0
                        ))
    else:
        # Shell pressure.
        with open(outname, 'w') as output:
            output.write(header)
            for i, item in enumerate(pressure):
                low = i * dr
                high = low + dr
                output.write('{l:.3f}\t{h:.3f}\t{p:.13f}\n'.format(l=low, h=high, p=item))

def write_ratio(ratio, dr, outname, header, bubble=True):
    """Write atom ratio stats to output file.
    If bubble is True, r_low is always zero.
    """
    with open(outname, 'w') as output:
        output.write(header)
        for i, item in enumerate(ratio):
            low = 0 if bubble else i * dr
            high = (i + 1) * dr
            output.write('{l:.3f}\t{h:.3f}\t{r:.13f}\n'.format(l=low, h=high, r=item))

def bubble_ratio(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble ratio stats and write results to disk."""
    for eles in elements:
        # Ratio stats for each element.
        e = ''.join(eles)
        print('Bubble ratio stats for {e}'.format(e=e))
        # Calculate ratio.
        ratio = box.atom_stats(eles[0], dr)
        # Write to file.
        outname = out_fmt.format(time=time, ele=e)
        write_ratio(ratio, dr, outname, header, bubble=True)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_ratio(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate shell ratio stats and write results to disk."""
    pass

def bubble_pressure(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble pressure and write results to disk."""
    for eles in elements:
        # Bubble pressure stats for each group of specified elements.
        e = ''.join(eles)
        print("Bubble pressure stats for {e}\n".format(e=e))
        # Calculate bubble pressure.
        bubble_pressure = box.pressure_stats(eles, dr)
        # Write bubble pressure.
        outname = out_fmt.format(time=time, ele=e)
        write_pressure(bubble_pressure, dr, outname, header, bubble=True)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_pressure(box, elements, out_fmt, header, dr, time, container, normal=False, debug=False):
    """Calculate shell pressure and write results to disk."""
    for eles in elements:
        # Shell pressure stats for each group of specified elements.
        e = ''.join(eles)
        print('Shell pressure stats for {e}\n'.format(e=e))
        # Shell pressure.
        if not normal:
            shell_pressure = box.shell_pressure_stats(eles, dr, normal=normal)
            # Write to disk.
            outname = out_fmt.format(time=time, ele=e)
            write_pressure(shell_pressure, dr, outname, header, bubble=False)
            if debug:
                # For testing.
                with open(container, 'a') as cc:
                    cc.write(outname + '\n')
        else:
            shell_pressure   = box.shell_pressure_stats(eles, dr, normal=normal)
            shell_r, shell_theta, shell_phi = shell_pressure['r'], shell_pressure['theta'], shell_pressure['phi']
            # Write to disk.
            outname1 = out_fmt.format(time=time, ele=e) + '_r'
            outname2 = out_fmt.format(time=time, ele=e) + '_theta'
            outname3 = out_fmt.format( time=time, ele=e ) + '_phi'
            write_pressure(shell_r, dr, outname1, header, bubble=False)
            write_pressure(shell_theta, dr, outname2, header, bubble=False)
            write_pressure( shell_phi, dr, outname3, header, bubble=False )

            if debug:
                # For testing.
                with open(container, 'a') as cc:
                    cc.write( outname1 + '\n' )
                    cc.write( outname2 + '\n' )
                    cc.write( outname3 + '\n' )

def bubble_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble density stats and write results to disk."""
    for eles in elements:
        # Bubble density stats for each group of specified elements.
        e = ''.join(eles)
        print('Bubble density stats for {e}\n'.format(e=e))
        # Bubble density.
        bubble_density = box.bubble_density(eles, mole, dr)
        # Write to disk.
        outname = out_fmt.format(time=time, ele=e)
        write_density(bubble_density, dr, outname, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate shell density stats and write results to disk."""
    for eles in elements:
        # Shell density stats for each group of specified elements.
        e = ''.join(eles)
        print('Shell density stats for {e}\n'.format(e=e))
        # Shell density.
        shell_density = box.shell_density(eles, mole, dr)
        # Write to disk.
        outname = out_fmt.format(time=time, ele=e)
        write_density(shell_density, dr, outname, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def xyz_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate xyz density stats and write results to disk."""
    for eles in elements:
        # XYZ density stats for each group of specified elements.
        e = ''.join(eles)
        print('XYZ density stats for {e}\n'.format(e=e))
        # XYZ density.
        xyz_density = box.xyz_density(eles, mole, dr)
        # Write to disk.
        xout = out_fmt.format(time=time, ele=e, xyz='x')
        yout = out_fmt.format(time=time, ele=e, xyz='y')
        zout = out_fmt.format(time=time, ele=e, xyz='z')

        write_density(xyz_density['x'], dr, xout, header)
        write_density(xyz_density['y'], dr, yout, header)
        write_density(xyz_density['z'], dr, zout, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                out = '\n'.join([xout, yout, zout, ''])
                cc.write(out)

def get_radius(box, element, dr, n=1, ratio=0.5):
    """Get the radius of a bubble.
    Radius is determined to be r with closest value of n_element / n_atoms
    to ratio, i.e. within radius, n_element / n_atoms should be as close to
    ratio as possible.
    n specifies number of radiuses to return, i.e. n radiuses that have
    n_element / n_atoms values closest to ratio."""
    bubble_ratio = box.atom_stats(element, dr)
    deltas = [abs(x - ratio) for x in bubble_ratio]
    # Use nanmin to ignore NaNs in ratio vector.
    # Do not select radiuses smaller than 10 angstrom.
    min_index = deltas.index(np.nanmin(deltas))
    n = n / 2
    ret = []
    for i in range(-n, n + 1):
        index = min_index + i
        ret.append((dr * (index + 1), bubble_ratio[index]))
    return ret


# -*- coding: utf-8 -*-
import numpy as np
import math
import tensorflow as tf
from model import UnrealModel

class TestUnrealModel(tf.test.TestCase):
  def test_unreal_variable_size(self):
    """ Check total variable size with all options ON """
    use_pixel_change = True
    use_value_replay = True
    use_reward_prediction = True

    # base: conv=4, fc=2, lstm=2, policy_fc=2, value_fc=2
    # pc:   fc=2, deconv_v=2, deconv_a=2
    # rp:   fc=2
    self.check_model_var_size( use_pixel_change,
                               use_value_replay,
                               use_reward_prediction,
                               20 )

  def test_pc_variable_size(self):
    """ Check total variable size with only pixel change ON """
    use_pixel_change = True
    use_value_replay = False
    use_reward_prediction = False

    # base: conv=4, fc=2, lstm=2, policy_fc=2, value_fc=2
    # pc:   fc=2, deconv_v=2, deconv_a=2
    self.check_model_var_size( use_pixel_change,
                               use_value_replay,
                               use_reward_prediction,
                               18 )

  def test_vr_variable_size(self):
    """ Check total variable size with only value funciton replay ON """
    use_pixel_change = False
    use_value_replay = True
    use_reward_prediction = False

    # base: conv=4, fc=2, lstm=2, policy_fc=2, value_fc=2
    self.check_model_var_size( use_pixel_change,
                               use_value_replay,
                               use_reward_prediction,
                               12 )

  def test_rp_variable_size(self):
    """ Check total variable size with only reward prediction ON """
    use_pixel_change = False
    use_value_replay = False
    use_reward_prediction = True

    # base: conv=4, fc=2, lstm=2, policy_fc=2, value_fc=2
    # rp:   fc=2
    self.check_model_var_size( use_pixel_change,
                               use_value_replay,
                               use_reward_prediction,
                               14 )

  def check_model_var_size(self,
                           use_pixel_change,
                           use_value_replay,
                           use_reward_prediction,
                           var_size):
    """ Check variable size of the model """

    model = UnrealModel(1,
                        -1,
                        use_pixel_change,
                        use_value_replay,
                        use_reward_prediction,
                        1.0,
                        1.0,
                        "/cpu:0");
    variables = model.get_vars()
    self.assertEqual( len(variables), var_size )

if __name__ == "__main__":
  tf.test.main()


#!/usr/bin/env python
"""
# > Script for evaluating 2x SESR
#    - Paper: https://arxiv.org/pdf/2002.01155.pdf
# Maintainer: Jahid (email: islam034@umn.edu)
# Interactive Robotics and Vision Lab (http://irvlab.cs.umn.edu/)
"""
import os
import time
import numpy as np
from glob import glob
from ntpath import basename
from os.path import join, exists
from PIL import Image
from keras.models import model_from_json
## local libs
from utils.data_utils import getPaths
from utils.data_utils import preprocess, deprocess
from utils.data_utils import deprocess_uint8, deprocess_mask

# input and output data shape
scale = 2
hr_w, hr_h = 640, 480 # HR
lr_w, lr_h = 320, 240 # LR (1/2x)
lr_res, lr_shape = (lr_w, lr_h), (lr_h, lr_w, 3)
hr_res, hr_shape = (hr_w, hr_h), (hr_h, hr_w, 3)

## for testing arbitrary local data
data_dir = "data/sample_test_ufo/lrd/"
#data_dir = "data/test_mixed/"
test_paths = getPaths(data_dir)
print ("{0} test images are loaded".format(len(test_paths)))

## load specific model
ckpt_name =  "deep_sesr_2x_1d"
model_h5 = join("models/", ckpt_name+".h5")
model_json = join("models/", ckpt_name + ".json")

# load model
assert (exists(model_h5) and exists(model_json))
with open(model_json, "r") as json_file:
    loaded_model_json = json_file.read()
generator = model_from_json(loaded_model_json)
generator.load_weights(model_h5)
print("\nLoaded data and model")

## create dir for output test data
samples_dir = join("data/output/", "keras_out")
if not exists(samples_dir): os.makedirs(samples_dir)

# testing loop
times = [];
for img_path in test_paths:
    # prepare data
    img_name = basename(img_path).split('.')[0]
    img_lrd = np.array(Image.open(img_path).resize(lr_res))
    im = np.expand_dims(preprocess(img_lrd), axis=0)
    # get output
    s = time.time()
    gen_op = generator.predict(im)
    gen_lr, gen_hr, gen_mask = gen_op[0], gen_op[1], gen_op[2]
    tot = time.time()-s
    times.append(tot)
    # process raw outputs
    gen_lr = deprocess_uint8(gen_lr).reshape(lr_shape)
    gen_hr = deprocess_uint8(gen_hr).reshape(hr_shape)
    gen_mask = deprocess_mask(gen_mask).reshape(lr_h, lr_w)
    # save generated images
    Image.fromarray(img_lrd).save(join(samples_dir, img_name+'.png'))
    Image.fromarray(gen_lr).save(join(samples_dir, img_name+'_En.png'))
    Image.fromarray(gen_mask).save(join(samples_dir, img_name+'_Sal.png'))
    Image.fromarray(gen_hr).save(join(samples_dir, img_name+'_SESR.png'))
    print ("tested: {0}".format(img_path))

# some statistics
num_test = len(test_paths)
if (num_test==0):
    print ("\nFound no images for test")
else:
    print ("\nTotal images: {0}".format(num_test))
    # accumulate frame processing times (without bootstrap)
    Ttime, Mtime = np.sum(times[1:]), np.mean(times[1:])
    print ("Time taken: {0} sec at {1} fps".format(Ttime, 1./Mtime))
    print("\nSaved generated images in in {0}\n".format(samples_dir))


"""
Molecular visualization of caffeine
===================================

.. currentmodule:: biotite.structure.graphics

*Biotite* provides simple interactive molecular visualization via
:func:`plot_atoms()`.
Although it does not produce publication-suitable images,
this function can be a convenient tool for a quick visual analysis of a
structure.

This example displays the small molecule caffeine.
"""

# Code source: Patrick Kunzmann
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import biotite.structure as struc
import biotite.structure.info as info
import biotite.structure.graphics as graphics

# Get an atom array for caffeine
# Caffeine has the PDB reside name 'CFF'
caffeine = info.residue("CFF")

# For cosmetic purposes align central rings to x-y plane
n1 = caffeine[caffeine.atom_name == "N1"][0]
n3 = caffeine[caffeine.atom_name == "N3"][0]
n7 = caffeine[caffeine.atom_name == "N7"][0]
# Normal vector of ring plane
normal = np.cross(n1.coord - n3.coord, n1.coord - n7.coord)
# Align ring plane normal to z-axis
caffeine = struc.align_vectors(caffeine, normal, np.array([0,0,1]))

# Caffeine should be colored by element
colors = np.zeros((caffeine.array_length(), 3))
colors[caffeine.element == "H"] = (0.8, 0.8, 0.8) # gray
colors[caffeine.element == "C"] = (0.0, 0.8, 0.0) # green
colors[caffeine.element == "N"] = (0.0, 0.0, 0.8) # blue
colors[caffeine.element == "O"] = (0.8, 0.0, 0.0) # red

fig = plt.figure(figsize=(8.0, 8.0))
ax = fig.add_subplot(111, projection="3d")
graphics.plot_atoms(
    ax, caffeine, colors, line_width=5, background_color="white",
    zoom=1.5
)
fig.tight_layout()

# Create an animation that rotates the molecule about the x-axis
def update(angle):
    ax.elev = angle

FPS = 50
DURATION = 4
angles = np.linspace(-180, 180, DURATION * FPS)
# Start at 90 degrees
angles = np.concatenate([
    np.linspace(  90, 180, int(DURATION * FPS * 1/4)),
    np.linspace(-180, 90,  int(DURATION * FPS * 3/4))
])
animation = FuncAnimation(fig, update, angles, interval=int(1000/FPS))
plt.show()


"""
Created on Tue Sep 23 17:15:47 2014
@author: Natural Solutions (Thomas)
"""

from pyramid.view import view_config
from sqlalchemy import desc, select, func, insert, join, Integer, cast, and_, Float, or_
from ecorelevesensor.models import AnimalLocation,ProtocolIndividualEquipment, DBSession, DataGsm, SatTrx, ObjectsCaracValues, Individual,V_Individuals_LatLonDate
from ecorelevesensor.utils.distance import haversine

import pandas as pd
import numpy as np
import re
import datetime,transaction,json

prefix = 'dataGsm/'

@view_config(route_name=prefix + 'unchecked/list', renderer='json')
def data_gsm_unchecked_list(request):
    '''List unchecked GSM data.
    '''
    unchecked = select([
        DataGsm.platform_,
        DataGsm.date
    ]).alias()

    print (unchecked)

    pie = ProtocolIndividualEquipment
    unchecked_with_ind = select([
        pie.ind_id.label('ind_id'),
        unchecked.c.platform_,
        func.count().label('nb')
    ]).select_from(
        unchecked.join(SatTrx, SatTrx.ptt == unchecked.c.platform_)
        .outerjoin(
            pie,
            and_(SatTrx.id == pie.sat_id,
                 unchecked.c['DateTime'] >= pie.begin_date,
                 or_(
                     unchecked.c['DateTime'] < pie.end_date,
                     pie.end_date == None
                )
            )
        )
    ).group_by(unchecked.c.platform_, pie.ind_id).order_by(unchecked.c.platform_)
    # Populate Json array
    data = DBSession.execute(unchecked_with_ind).fetchall()
    return [dict(row) for row in data]

@view_config(route_name=prefix + 'unchecked', renderer='json')
def data_gsm_unchecked(request):
    '''Get the unchecked GSM data.
    '''
    platform = int(request.matchdict['id'])

    if request.GET['format'] == 'geojson':
        # Query
        query = select([
            DataGsm.id.label('id'),
            DataGsm.lat,
            DataGsm.lon,
            DataGsm.date
        ]).where(DataGsm.platform_ == platform).where(DataGsm.checked == False
        ).order_by(desc(DataGsm.date)).limit(1000)
        # Create list of features from query result
        features = [
            {
                'type':'Feature',
                'properties':{'date':str(date)},
                'geometry':{'type':'Point', 'coordinates':[float(lon),float(lat)]},
                'id':id_
            }
        for id_, lat, lon, date in DBSession.execute(query).fetchall()]
        transaction.commit()
        result = {'type':'FeatureCollection', 'features':features}
        return result

    elif request.GET['format'] == 'json':
        # Query
        query = select([
            DataGsm.id.label('id'),
            DataGsm.lat.label('lat'),
            DataGsm.lon.label('lon'),
            DataGsm.ele.label('ele'),
            DataGsm.date.label('date')]
        ).where(DataGsm.platform_ == platform
        ).where(DataGsm.checked == False
        ).order_by(desc(DataGsm.date)).limit(1000)
        data = DBSession.execute(query).fetchall()
        # Load data from the DB then
        # compute the distance between 2 consecutive points.
        df = pd.DataFrame.from_records(data, columns=data[0].keys(), coerce_float=True)
        X1 = df.iloc[:-1][['lat', 'lon']].values
        X2 = df.iloc[1:][['lat', 'lon']].values
        df['dist'] = np.append(haversine(X1, X2), 0).round(3)
        # Compute the speed
        df['speed'] = (df['dist']/((df['date']-df['date'].shift(-1)).fillna(1)/np.timedelta64(1, 'h'))).round(3)
        # Values to import : the first per hour
        ids = df.set_index('date').resample('1H', how='first').dropna().id.values
        df['import'] = df.id.isin(ids)
        df['date'] = df['date'].apply(str)
        # Fill NaN
        df.fillna(value={'ele':-999}, inplace=True)
        return df.to_dict('records')

@view_config(route_name=prefix + 'unchecked/import', renderer='json')
def data_gsm_unchecked_import(request):
    '''Import unchecked GSM data.
    '''
    ptt=request.matchdict['id']
    data = list(request.params)
    print(data)
    data=json.loads(data[0]).get('data')
    id_ind=json.loads(data[0]).get('id_ind')
    select_stmt = DBSession.execute(select([DataGsm],DataGsm.id.in_(data))).fetchall()
    query = insert(AnimalLocation, select_stmt)
    DBSession.execute(query)
    # query = select([
    #     DataGsm.platform_,
    #     func.count(DataGsm.id).label('nb')
    # ]).group_by(DataGsm.platform_)

    # return [dict(row) for row in DBSession.execute(query).fetchall()]

def asInt(s):
    try:
        return int(s)
    except:
        return None

@view_config(route_name=prefix + 'upload', renderer='string')
def data_gsm_all(request):
    '''Import unchecked GSM data.
    '''
    response = 'Success'
    ptt_pattern = re.compile('[0]*(?P<platform>[0-9]+)g')
    try:
        file = request.POST['file'].file
        filename = request.POST['file'].filename
        platform = int(ptt_pattern.search(filename).group('platform'))
        query = select([DataGsm.date]).where(DataGsm.platform_ == platform)
        # Read dates that are already in the database
        df = pd.DataFrame.from_records(DBSession.execute(query).fetchall(), index=DataGsm.date.name, columns=[DataGsm.date.name])
        # Load raw csv data
        csv_data = pd.read_csv(file, sep='\t',
            index_col=0,
            parse_dates=True,
            # Read those values as NaN
            na_values=['No Fix', 'Batt Drain', 'Low Voltage'],
            # Only import the first 8 columns
            usecols=range(9)
        )
        # Remove the lines containing NaN
        csv_data.dropna(inplace=True)
        # Filter data with no elevation by converting the column to numeric type
        csv_data[DataGsm.ele.name] = csv_data[DataGsm.ele.name].convert_objects(convert_numeric=True)
        # Get the data to insert
        data_to_insert = csv_data[~csv_data.index.isin(df.index)]
        # Add the platform to the DataFrame
        data_to_insert[DataGsm.platform_.name] = platform
        # Write into the database
        data_to_insert.to_sql(DataGsm.__table__.name, DBSession.get_bind(), if_exists='append')
    except:
        response = 'An error occured.'
        request.response.status_code = 500
    return response

@view_config(route_name=prefix + 'details', renderer='json')
def indiv_details(request):

    params=int(request.matchdict['id'])
    join_table = join(SatTrx, ObjectsCaracValues, SatTrx.ptt == cast(ObjectsCaracValues.value, Integer)
        ).join(Individual, ObjectsCaracValues.object==Individual.id)

    query=select([ObjectsCaracValues.value.label('id'), Individual.id.label('ind_id'),Individual.survey_type.label('survey_type'), Individual.status.label('status')
        , Individual.monitoring_status.label('monitoring_status'), Individual.birth_date.label('birth_date'), Individual.ptt.label('ptt'),ObjectsCaracValues.begin_date.label('begin_date'),ObjectsCaracValues.end_date.label('end_date')]
        ).select_from(join_table
        ).where(and_(SatTrx.model.like('GSM%'),ObjectsCaracValues.carac_type==19,ObjectsCaracValues.object_type=='Individual')
        ).where(ObjectsCaracValues.value==params).order_by(desc(ObjectsCaracValues.begin_date))

    data=DBSession.execute(query).first()
    transaction.commit()
    if data['end_date'] == None :
        end_date=datetime.datetime.now()
    else :
        end_date=data['end_date']

    result=dict([ (key[0],key[1]) for key in data.items()])
    print(result)
    result['duration']=(end_date.month-data['begin_date'].month)+(end_date.year-data['begin_date'].year)*12

    query = select([V_Individuals_LatLonDate.c.date]
                     ).where(V_Individuals_LatLonDate.c.ind_id == result['ind_id']
                     ).order_by(desc(V_Individuals_LatLonDate.c.date)).limit(1)

    lastObs=DBSession.execute(query).fetchone()
    result['last_observation']=lastObs['date'].strftime('%d/%m/%Y')
    if result['birth_date']!= None:
        result['birth_date']=result['birth_date'].strftime('%d/%m/%Y')
    del result['begin_date'], result['end_date']
    print (result)
    return result