#!/usr/bin/python

from math import *

import input_wrapper
InputFile = input_wrapper.InputFile
rw_data = list(input_wrapper.rw_data)
rw_data.sort()
ro_data = list(input_wrapper.rw_data)
ro_data.sort()

import sys, os
import tempfile

editor = os.getenv("EDITOR",default="vi")

## Data cleaning
##################

class Cleaner(object):
  def __init__(self,inp,name):
    self.dir = inp.name
    self.name = name.lower()
    self.inp = inp
  def action(self):
    raise TypeError
  def clean(self):
    print "Removing "+self.name
    self.action()

class GridDataCleaner(Cleaner):
  def action(self):
    os.system("rm -f %s/grid_data/%s.gz"%(self.dir,self.name) )

class AllCleaner(Cleaner):
  def action(self):
    os.system("rm -rf %s/grid"%(self.dir) )
    os.system("rm -rf %s/grid_data"%(self.dir) )

def clear(x,inp):
  d = { "all": AllCleaner,
      }
  if x in d.keys():
    type = d[x]
  else:
    type = GridDataCleaner

  type(inp,x.replace("_"," ")).clean()

## Data editing
###############

class Editor(object):
  def __init__(self,inp,name):
    self.dir = inp.name
    self.name = name
    self.inp = inp
  def action(self):
    raise TypeError
  def edit(self):
    print "Editing "+self.name
    self.action()
    
class GeometryEditor(Editor):
  def action(self):
    self.charge = self.inp.nucl_charge
    self.coord = self.inp.nucl_coord
    edit_temp_file(self.inp,self.write_geom_file,self.read_geom_file)

  def write_geom_file(self,file,inp):
    print >>file, "###########################"
    print >>file, "# Geometry : %s"%(inp.name)
    print >>file, "###########################"
    print >>file, ""
    fields = "Charge X Y Z"
    format = "%10.1f "+"%10.6f "*3
    fields = fields.split()
    print >>file, " ".join(["#"]+map(lambda x: x.center(10),fields))
    print >>file, "# "+"-"*11*len(fields)+"\n"
    charge = self.charge
    coord = self.coord
    for i in xrange(len(charge)):
      buffer = [ charge[i], coord[0][i], coord[1][i], coord[2][i] ]
      print >>file, "  "+format%tuple(buffer)

  def read_geom_file(self,file,inp):
    lines = file.readlines()
    lines = filter(lambda x: len(x.strip())>0,lines)
    lines = filter(lambda x: not x.startswith("#"),lines)
    lines = map(lambda x: x.split(),lines)
    coord = [ [], [], [] ]
    charge = []
    for line in lines:
      charge.append(float(line[-4]))
      coord[0].append(float(line[-3]))
      coord[1].append(float(line[-2]))
      coord[2].append(float(line[-1]))
    inp.coord = coord
      
class GridEditor(Editor):
  def action(self):
    self.origin= self.inp.origin
    self.opposite = self.inp.opposite
    self.step_size = self.inp.step_size
    self.point_num = self.inp.point_num
#   if self.origin is None:
#     self.origin = [ 0., 0., 0. ]
#     for l in xrange(3): 
#       self.origin[l] = self.inp.nucl_coord[l][1]
#       for i in self.inp.nucl_coord[l]:
#         self.origin[l] = min(self.origin[l],i)
#       self.origin[l] -= 4.
#   if self.opposite is None:
#     self.opposite = [ 0., 0., 0. ]
#     for l in xrange(3): 
#       self.opposite[l] = self.inp.nucl_coord[l][1]
#       for i in self.inp.nucl_coord[l]:
#         self.opposite[l] = max(self.opposite[l],i)
#       self.opposite[l] += 4.
#   if self.step_size is None:
#     self.step_size = [ 0., 0., 0. ]
#     for l in xrange(3): 
#       self.step_size[l] = (self.opposite[l] - self.origin[l])/float(self.point_num[l])

#   self.inp.point_num = self.point_num
#   self.inp.origin = self.origin
#   self.inp.opposite = self.opposite
#   self.inp.step_size = self.step_size
    edit_temp_file(self.inp,self.write_grid_file,self.read_grid_file)

  def write_grid_file(self,file,inp):
    print >>file, "###########################"
    print >>file, "# Grid : %s"%(inp.name)
    print >>file, "###########################"
    print >>file, ""
    print >>file, "# Number of points along x, y, z"
    if self.point_num is None:
      print >>file, "Default"
    else:
      print >>file, "%d %d %d"%tuple(self.point_num)
    print >>file, ""
    print >>file, "# Coordinates of the origin (x,y,z)"
    if self.origin is None:
      print >>file, "Default"
    else:
      print >>file, "%f %f %f"%tuple(self.origin)
    print >>file, ""
    print >>file, "# Coordinates of the opposite point (x,y,z)"
    if self.opposite is None:
      print >>file, "Default"
    else:
      print >>file, "%f %f %f"%tuple(self.opposite)
    print >>file, ""
    print >>file, "# Step sizes (x,y,z)"
    if self.step_size is None:
      print >>file, "Default"
    else:
      print >>file, "%f %f %f"%tuple(self.step_size)

  def read_grid_file(self,file,inp):
    lines = file.readlines()
    lines = filter(lambda x: len(x.strip())>0,lines)
    lines = filter(lambda x: not x.startswith("#"),lines)

    buffer = lines.pop(0).lower().split()
    if buffer[0] != "default":
      self.inp.point_num = map(int,buffer)

    buffer = lines.pop(0).lower().split()
    if buffer[0] != "default":
      self.inp.origin = map(float,buffer)

    buffer = lines.pop(0).lower().split()
    if buffer[0] != "default":
      self.inp.opposite = map(float,buffer)

    buffer = lines.pop(0).lower().split()
    if buffer[0] != "default":
      self.inp.step_size = self.step_size

      

def edit(x,inp):
  d = { "geometry": GeometryEditor,
        "grid_parameters": GridEditor,
      }
  d[x](inp,x.replace("_"," ")).edit()

## Write temporary input file
#############################

def write_main_file(file,inp):
  print >>file, "####################"
  print >>file, "# %s"%(inp.name)
  print >>file, "####################"
  print >>file, ""
  print >>file, "# Edit"
  print >>file, "# --------------------\n"
  print >>file, '# edit(geometry)'
  print >>file, '# edit(grid_parameters)'
  print >>file, ""
  compute = filter(lambda x: x.startswith("compute"),rw_data)
  print >>file, "# Clear"
  print >>file, "# --------------------\n"
  print >>file, '# clear(all)'
  for p in compute:
    print >>file, '# clear(%s)'%(p[8:])
  print >>file, ""
  print >>file, "# Computation"
  print >>file, "# --------------------\n"
  compute = filter(lambda x: not x.endswith("_grad"),compute)
  compute = filter(lambda x: not x.endswith("_lapl"),compute)
  for p in compute:
    x = ' '
    exec "if inp.%s: x = 'X'"%(p)
    print >>file, "(%s) %s"%(x,p[8:])


## Execute temporary input file
###############################

def read_main_file(file,inp):

  lines = file.readlines()
  for line in lines:
    line = line.lstrip()
    if line == "":
      pass
    elif line[0] == "(":
      line = line.replace("( ) ","=False inp.compute_")
      line = line.replace("(X) ","=True inp.compute_")
      line = line.replace("(x) ","=True inp.compute_")
      buffer = line.split()
      line = ' '.join([buffer[1].lower(),buffer[0]])
    elif line.startswith("edit") \
      or line.startswith("clear"):
      line = line.replace("(","('").replace(")","',inp)").lower()
    exec line


def edit_temp_file(input_file,write_file,read_file,saved_file=None):
  if saved_file is None:
    file,filename = tempfile.mkstemp()
    file = open(filename,"w")
    write_file(file,input_file)
    file.close()
    os.system("%s %s"%(editor,filename))
  else:
    file = open(saved_file,'r')
    buffer = file.read()
    file.close()

    file,filename = tempfile.mkstemp()
    file = open(filename,"w")
    file.write(buffer)    
    file.close()
  file = open(filename,"r")
  read_file(file,input_file)
  file.close()

  os.remove(filename)


def build_script(inp):
  run_script = "run_"+inp.name
  file = open(run_script,'w')
  buffer = """
  #!/bin/bash
source ${HOME}/.eplfrc
$RUNCOMMAND
""".lstrip()
  command = "${EPLF_PATH}/bin/eplf %s"%(inp.name)
  if inp.has_mpi:
    command = "${MPIRUN} -np %d "%(inp.nproc)+command
  buffer = buffer.replace("$RUNCOMMAND",command)
  print >>file, buffer
  file.close()
  os.chmod(run_script,0744)


def main():
  if len(sys.argv) not in [2,3]:
    print "Syntax : %s <EZFIO_File> [input_file]"%(sys.argv[0])
    sys.exit(1)
  inp = InputFile(sys.argv[1])
  if len(sys.argv) == 3:
    saved_file = sys.argv[2]
  else:
    saved_file = None
  edit_temp_file(inp,write_main_file,read_main_file,saved_file)
  build_script(inp)

main()