#!/usr/bin/env python
#-*- coding: utf-8 -*-
"""
Generate nice graphs from mesh-tally 1 mctal files.

At the current stage only simple text processing is done that doesn't work with
multiple tallies or energy dependent meshes, yet.
"""

import gzip
from numpy import *
from pylab import *
from matplotlib.colors import LogNorm

COLORMAP='tab20c'
SCALING=1.56e16/8.7374e15

def read_tally(fname, use_tally=None):
  if fname.endswith('.gz'):
    txt=gzip.open(fname, 'r').read()
  else:
    txt=open(fname, 'r').read()
    
  tally_ranges=[]
  tidx=0
  while 'tally' in txt[tidx+1:]:
    tidx=txt.index('tally ', tidx+1)
    tnr=int(txt[tidx:].split(None,3)[1])
    if len(tally_ranges)>0:
      tally_ranges[-1][2]=tidx-1
    tally_ranges.append([tnr, tidx, -1])
  if use_tally is None:
    txt=txt[tally_ranges[0][1]:tally_ranges[0][2]]
  else:
    tidx=[ti[0] for ti in tally_ranges].index(use_tally)
    txt=txt[tally_ranges[tidx][1]:tally_ranges[tidx][2]]
  header, data=txt.split('vals')
  

  # get dimensions of data
  fstart=header.index('\nf ')
  hi=header[fstart+3:].splitlines()[0]
  N,nE,nX,nY,nZ=map(int, hi.strip().split())
  
  grid=header[fstart+3+len(hi):header.index('\nd ')]
  grid=array(grid.replace('\n', '').strip().split(), dtype=float)
  
  x=grid[:nX+1]
  y=grid[nX+1:nX+nY+2]
  z=grid[nX+nY+2:nX+nY+nZ+3]
  
  data=array(data.replace('\n','').strip().split(), dtype=float)
  I,dI=data.reshape(nZ, nY, nX, 2).transpose(3,2,1,0)
  
  return x,y,z,I,dI

def plot_xy(outfile, res, zidx=5, Imin=1e-6, Imax=1e4, cticks=None):
  x,y,z,I,dI=res
  fig=gcf()
  ax=gca()
  p=ax.pcolormesh(x/100., y/100., I[:,:,zidx].T, norm=LogNorm(Imin, Imax), cmap=COLORMAP)
  ax.set_aspect('equal')
  fig.colorbar(p, orientation='horizontal', shrink=0.9, ticks=cticks, label='Dose rate [Sv/h]')
  xlabel('x [m]')
  ylabel('y [m]')
  title('XY plane at z=%.1fm'%((z[zidx]+z[zidx+1])/200.))
  fig.subplots_adjust()
  #fig.savefig(outfile)
  #pc=p.get_facecolors()
  #pc[:,3]=where(I[:,:,zidx].T.flatten()>0, 0.9-0.25*dI[:,:,zidx].T.flatten(), 0.)
  ax.add_patch(Circle((0,0), radius=5.50, color='black', fill=False, lw=2))
  ax.add_patch(Circle((0,0), radius=11.50, color='black', fill=False, lw=2))
  ax.add_patch(Circle((0,0), radius=15.00, color='black', fill=False, lw=2))
  fig.savefig(outfile, transparent=False, dpi=300)

def plot_xz(outfile, res, yidx=5, Imin=1e-6, Imax=1e4, cticks=None):
  x,y,z,I,dI=res
  fig=gcf()
  ax=gca()
  p=ax.pcolormesh(x/100., z/100., I[:,yidx].T, norm=LogNorm(Imin, Imax), cmap=COLORMAP)
  ax.set_aspect('equal')
  fig.colorbar(p, orientation='horizontal', shrink=0.9, ticks=cticks, label='Dose rate [Sv/h]')
  xlabel('x [m]')
  ylabel('z [m]')
  ypos=(y[yidx]+y[yidx+1])/200.
  title('XZ plane at y=%.1fm'%ypos)
  fig.subplots_adjust()
  #fig.savefig(outfile)
  #pc=p.get_facecolors()
  #pc[:,3]=where(I[:,yidx].T.flatten()>0, 0.9-0.25*dI[:,yidx].T.flatten(), 0.)
  ax.plot([sqrt(5.5**2-ypos**2), sqrt(5.5**2-ypos**2)], [z.min()/100., z.max()/100.], color='black', lw=2)
  ax.plot([sqrt(11.5**2-ypos**2), sqrt(11.5**2-ypos**2)], [z.min()/100., z.max()/100.], color='black', lw=2)
  ax.plot([sqrt(15.0**2-ypos**2), sqrt(15.0**2-ypos**2)], [z.min()/100., z.max()/100.], color='black', lw=2)
  fig.savefig(outfile, transparent=False, dpi=300)

def plot_yz(outfile, res, xidx=5, Imin=1e-6, Imax=1e4, cticks=None):
  x,y,z,I,dI=res
  fig=gcf()
  ax=gca()
  p=ax.pcolormesh(y/100., z/100., I[xidx].T, norm=LogNorm(Imin, Imax), cmap=COLORMAP)
  ax.set_aspect('equal')
  fig.colorbar(p, orientation='horizontal', shrink=0.9, ticks=cticks, label='Dose rate [Sv/h]')
  xlabel('y [m]')
  ylabel('z [m]')
  title('YZ plane at x=%.1fm'%((x[xidx]+x[xidx+1])/200.))
  fig.subplots_adjust()
  #fig.savefig(outfile)
  #pc=p.get_facecolors()
  #pc[:,3]=where(I[xidx].T.flatten()>0, 0.9-0.25*dI[xidx].T.flatten(), 0.)
  fig.savefig(outfile, transparent=False, dpi=300)

if __name__=='__main__':
  import sys, os
  res=read_tally(sys.argv[1], use_tally=1)
  x,y,z,I,dI=res
  I*=SCALING
  fpre=os.path.join('.', 'images', os.path.basename(sys.argv[1])[:-4])
  xi, yi, zi=res[3].shape
  for zidx in range(zi):
    print fpre+'_xy_%02i.png'%zidx
    figure(figsize=(12,8))
    plot_xy(fpre+'_xy_%02i.png'%zidx, res, zidx=zidx, 
            Imin=1e-8, Imax=1e2, cticks=[1e-8, 1e-6, 1e-4, 1e-2, 1, 100])
  for yidx in range(yi):
    print fpre+'_xz_%02i.png'%yidx
    figure(figsize=(12,8))
    plot_xz(fpre+'_xz_%02i.png'%yidx, res, yidx=yidx, 
            Imin=1e-8, Imax=1e2, cticks=[1e-8, 1e-6, 1e-4, 1e-2, 1, 100])
  for xidx in range(xi):
    print fpre+'_yz_%02i.png'%xidx
    figure(figsize=(8,10))
    plot_yz(fpre+'_yz_%02i.png'%xidx, res, xidx=xidx, 
            Imin=1e-8, Imax=1e2, cticks=[1e-8, 1e-6, 1e-4, 1e-2, 1, 100])