#!/usr/bin/env python

from __future__ import division
import sys
import subprocess
import os
from math import *

# Convert logical into "0" or "1" string for Spectra

def logic_str (logic):
  if logic:
    return '1'
  else:
    return '0'

# Command line arguments

main_input_file = ''

i = 1
while i < len(sys.argv):

  if sys.argv[i][0] == '-':
    print_help()

  else:
   main_input_file = sys.argv[i]

  i += 1

# Read in main input file

if main_input_file == '': 
  sys.exit('No main input file specified. Stopping.')

print ('Main input file: ' + main_input_file)
if main_input_file.endswith('.py'): main_input_file = main_input_file[:-3]
param = __import__(main_input_file)

if param.source_type == 'bend':
  param.x_angle_range = [0, 0, 1]  # Fake

spectra_exe = param.spectra_exe_dir + '/spectra_solver'

# Create data subdirectories if needed

if param.spline_dir[-1] != '/': param.spline_dir = param.spline_dir + '/'
input_dir  = param.spline_dir + 'spectra_input/'
output_dir = param.spline_dir + 'spectra_output/'
spline_dir = param.spline_dir + 'spline/'

if (not os.path.exists(input_dir)):  os.makedirs(input_dir)
if (not os.path.exists(output_dir)): os.makedirs(output_dir)
if (not os.path.exists(spline_dir)): os.makedirs(spline_dir)

#--------------------------------------------------------------------------------
# Create spectra energy probability run file by modifying the base run file.

spectra_energy_in_file = input_dir + 'energy_run.spectra'
spectra_energy_out_file = output_dir + 'energy'  # Full name: + '.dc0'

in_file = open(param.spectra_base_parameter_file)
out_file = open(spectra_energy_in_file, 'w')

section = '[]'
for line in in_file.readlines():
  split_line = line.split('\t')

  if line[0] == '[': section = line.strip()

  if line[:10] == '@ProcessNo': line = '@ProcessNo 0 5 -1 0\n'

  if section == '[OUTPUTDATA]':
    if split_line[0] == 'Name': line = 'Name\t' + spectra_energy_out_file + '\n'

  if section == '[ACCELERATOR]':
    if split_line[0] == 'eGeV': beam_energy = float(split_line[1]) * 1e9

  if section == '[SOURCE]':
    if split_line[0] == 'by':   bend_field  = float(split_line[1])

  if section == '[CALCULATION]':
    if split_line[0] == 'accuracy':  line = 'accuracy\t'  + str(1) + '\n'  # Accuracy level (1 - 10)
    if split_line[0] == 'elogscale': line = 'elogscale\t' + logic_str(param.energy_log_scale) + '\n' 
    if hasattr(param, 'energy_E_range'):
      if split_line[0] == 'epmin':     line = 'epmin\t'     + str(param.energy_E_range[0]) + '\n'
      if split_line[0] == 'epmax':     line = 'epmax\t'     + str(param.energy_E_range[1]) + '\n'
      if split_line[0] == 'meshep':    line = 'meshep\t'    + str(param.energy_E_range[2]) + '\n'

  out_file.write(line)

in_file.close()
out_file.close()

# Run spectra energy probability calc

print ('Running spectra to produce energy probability file')
subprocess.call(spectra_exe + ' ' + spectra_energy_in_file, shell=True)

#--------------------------------------------------------------------------------
# Create angle probability run file by modifying the energy scan file.

spectra_base_angle_in_file = input_dir + 'angle_run.spectra'
spectra_base_angle_out_file = output_dir + 'angle#'  # Full name: + '.dty'

in_file = open(spectra_energy_in_file)
out_file = open(spectra_base_angle_in_file, 'w')

slit_dist = 1

section = '[]'
for line in in_file.readlines():
  split_line = line.split('\t')
  if line[0] == '[': section = line.strip()
  if param.source_type == 'bend':
    if line[:10] == '@ProcessNo': line = '@ProcessNo 0 6 -1 0\n'
  else:
    if line[:10] == '@ProcessNo': line = '@ProcessNo 0 7 -1 0\n'

  if section == '[CALCULATION]':
    if split_line[0] == 'fin_dist':    line = 'fin_dist\t 1'
    if split_line[0] == 'slit_dist':   line = 'slit_dist\t'  + str(slit_dist) + '\n'  
    if split_line[0] == 'minthetax':   line = 'minthetax\t'  + str(1e6*param.x_angle_range[0]) + '\n'  
    if split_line[0] == 'maxthetax':   line = 'maxthetax\t'  + str(1e6*param.x_angle_range[1]) + '\n'  
    if split_line[0] == 'meshx':       line = 'meshx\t'      + str(param.x_angle_range[2]) + '\n'  
    if split_line[0] == 'min_x':       line = 'min_x\t'      + str(1e3*param.x_angle_range[0]*slit_dist) + '\n'  
    if split_line[0] == 'max_x':       line = 'max_x\t'      + str(1e3*param.x_angle_range[1]*slit_dist) + '\n'  
    if split_line[0] == 'minthetay':   line = 'minthetay\t'  + str(1e6*param.y_angle_range[0]) + '\n'  
    if split_line[0] == 'maxthetay':   line = 'maxthetay\t'  + str(1e6*param.y_angle_range[1]) + '\n'  
    if split_line[0] == 'meshy':       line = 'meshy\t'      + str(param.y_angle_range[2]) + '\n'  
    if split_line[0] == 'min_y':       line = 'min_y\t'      + str(1e3*param.y_angle_range[0]*slit_dist) + '\n'  
    if split_line[0] == 'max_y':       line = 'max_y\t'      + str(1e3*param.y_angle_range[1]*slit_dist) + '\n'  

  out_file.write(line)

in_file.close()
out_file.close()

#--------------------------------------------------------------------------------
# Angular dependence: Loop over all energies

for n in xrange(int(param.angle_E_range[2]+0.5)):
  if param.energy_log_scale: 
    energy = param.angle_E_range[0] * exp(log((param.angle_E_range[1] / param.angle_E_range[0])) * \
                                                                          (n / (param.angle_E_range[2]-1)))
  else:
    energy = param.angle_E_range[0] + (param.angle_E_range[1] - param.angle_E_range[0]) * (n / (param.angle_E_range[2]-1))

  #--------------------------------------------------------------------------------
  # Run spectra vertical angle probability and create vertical angle spline file.

  spectra_angle_in_file  = input_dir + 'angle_run' + str(n+1) + '.spectra'

  spectra_angle_out_file = spectra_base_angle_out_file.replace('#', str(n+1))

  in_file = open(spectra_base_angle_in_file)
  out_file = open(spectra_angle_in_file, 'w')

  section = '[]'
  for line in in_file.readlines():
    split_line = line.split('\t')
    if line[0] == '[': section = line.strip()
    if section == '[OUTPUTDATA]':
      if split_line[0] == 'Name': line = 'Name\t' + spectra_angle_out_file + '\n'

    if section == '[CALCULATION]':
      if split_line[0] == 'fixep':  line = 'fixep\t'  + str(energy) + '\n'  

    out_file.write(line)

  in_file.close()
  out_file.close()

  # Run spectra angle probability calc

  print ('Running spectra to produce angle probability file at energy: ' + str(energy))
  run_cmd = spectra_exe + ' ' + spectra_angle_in_file
  print ('     ' + run_cmd)
  subprocess.call(run_cmd, shell = True)


#--------------------------------------------------------------------------------
#  create spline files.

spline_params_file = param.spline_dir + 'spline.params'

param_file = open(spline_params_file, 'w')
param_file.write ('&master_params\n')
param_file.write ('  source_type              = "' + param.source_type + '"\n')
param_file.write ('  spline_space_dimensions  = ' + str(param.spline_space_dimensions) + '\n')
param_file.write ('\n')
param_file.write ('  num_rows_energy          = ' + str(param.energy_E_range[2]) + '\n')
param_file.write ('  dE_spline_max            = ' + str(param.energy_dE_spline_max) + '\n')
param_file.write ('  dP_spline_max            = ' + str(param.energy_dP_spline_max) + '\n')
param_file.write ('\n')
param_file.write ('  num_rows_x_angle         = ' + str(param.x_angle_range[2]) + '\n')
param_file.write ('  num_rows_y_angle         = ' + str(param.y_angle_range[2]) + '\n')
param_file.write ('/\n')
param_file.close()

print ('Running spline program...')
if param.program_bin_dir[-1] != '/': param.program_bin_dir = param.program_bin_dir + '/'
run_cmd = param.program_bin_dir + 'create_splines ' + param.spline_dir
print ('     ' + run_cmd)
subprocess.call(run_cmd, shell=True)