#!/usr/bin/python

# 	 Copyright (C) 2015 Francois Schiettekatte, Universite de Montreal
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.

#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.

#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.


# This script downloads a reasonable number (~85) of cross-section datasets from 
# Alex Gurbich's SigmaCalc website for a specific reaction and interpolates in 
# them (using Pylab bicubic interpolation method) in order to write a fine-grained 
# angle-energy cross-section map to be used by Corteo. (The script actually starts 
# by downloading all the datasets, then generates the map.)

# usage: Python importSC reaction element [skip_download]

# where 'reaction' is aa for alpha-alpha, pp for proton-proton, etc.
# 'element' is C-12 for carbon-12, Si-nat for Si with natural abundances of its isotopes, etc.

# reaction names include: pp, aa, pp0, aa0, ap, dd, dp0, da0, da1
# see SigmaCalc website for the reactions and elements available

# If 'skip_download' is specified as 3rd argument, the download stage is skipped  
# and interpolation is carried out based on previously saved data. This is to 
# avoid unnecessary download when making possible ajustements to the interpolation.

# IMPORTANT NOTE: Please understand that if you modify the script in order to download
# more datasets, it might be perceived as a cyber-attack, and may results in server 
# being shut down, etc. Please contact Alex Gurbich directly if you have special needs.

from pylab import *
from scipy import interpolate

import os
import urllib
import sys

# function to download the SigmaCalc-computed cross section of a given reaction at a given angle
# the file is saved in R33 format
# returns the name of the file generated
def downloadSC(angle, reactionName, element):
	'''download a dataset from SigmaCalc website and save it in a directory corresponding to the'''
	'''reaction (e.g. C-12_aa) with name C-12_aa_NNNNNN_sc.r33 where NNNNNN is the angle*1000'''
	'''angle is in degrees, reaction name = "aa" for alpha-alpha, etc., element = "O-16", etc.'''
	
	# first, ask the server to generate the datafile for that reaction at that angle
	address = 'http://sigmacalc.iate.obninsk.ru/cgi-bin/select_run.pl?angle='+ str(angle) \
		+ '&reaction='+reactionName+'&Element='+element \
		+'&r33=r33format&tl=o6a&dt=ebs' # R33 format
	print('\n'+address)
	try:
		url = urllib.urlopen(address)
		st=url.read()
	except Exception:
		print("Error when requesting data for reaction " + element + "(" + reactionName + ") at " +str(angle)+" deg.")
		print("URL requested: "+ address)
		raise
	s = st.decode("utf-8")		

	# in its response, the server tells us the filename of the generated data, lets find it
	findGenFilename = 'onclick=sub6(\"SC\",\"' 
	start = s.find(findGenFilename)+len(findGenFilename)
	stop =  s.find("\"",start+1);
	genFilename = s[start:stop]
	#print(genFilename)

	# now, download the datafile
	adresse = 'http://sigmacalc.iate.obninsk.ru/cgi-bin/download.pl?file='+genFilename
	#print(adresse)
	url = urllib.urlopen(adresse)
	st=url.read()
	s = st.decode("utf-8")

	name = element+'_'+reactionName

	# in the header, find the angle and check if it corresponds to the requested angle
	start = s.find("Theta:")+7;
	stop =  s.find("\n",start);
	# print(s)
	# print(start, stop, s[start:stop])
	try:
		dataangle = float(s[start:stop])
	except ValueError:
		return "Apparent server error at angle = "+str(angle)+", continuing without these data."
		
	if abs(dataangle-angle)>0.001:
		print("Requested and dataset angles not the same for " + reactionName)
		print("Requested angle = "+str(angle)+", dataset angle = " + str(dataangle))
		exit()
	
	# create a string that represents 1000*angle without dot (e.g. 55.2 deg becomes 055200)
	sangle = "{:06.0f}".format(angle*1000) 

	# in the header, check that units are 'rr' (i.e. Rutherford ratio)
	start = s.find("Units:")+7;
	stop =  s.find("\n",start);
	if s[start:stop].lower().strip() != 'rr' :
		print("Data expected to be in 'rr' units (ratio to Rutherford), but here is what is specified: "+s[start-7:stop])
		print("Giving up.")
		exit()
		

	# append the angle and an extension to the reaction name	
	name=name+"/"+name+'_'+sangle+'_sc.r33'
	
	# write what we downloaded into a file
	try:
		with open(name, "w") as outfile:
			outfile.write(s)
	except Exception:
		print("Error attempting to save file " + name)
		raise

	return name


# Function that returns the cross section from a SigmaCalc file in R33 format
# Unless they are provided, returns the energy and the number of lines to skip to reach the data
# If the energy array E is provided, check if it is identical to that of the loaded data 
def readsc(file, angle, nskip = None, E = None):
	name = file + "{:06.0f}".format(angle*1000) +'_sc.r33'
	#print(name)

	if nskip is None:
		# find number of lines to skip (actually search for "Data:")
		nskip = 0
		with open(name) as f:
			for line in f:
				nskip+=1
				if 'Data:' in line: 
					break
				
	a = loadtxt(name, skiprows=nskip,  comments="E")
	
	if E is not None:
		assert allclose(E,a[:,0], 0.001)
		
	return a[:,2], a[:,0], nskip, name




##### main script starts here #######

# check number of arguments
if len(sys.argv)<3 or len(sys.argv)>4 :
	print("Usage: Python importSC reaction element [skip_download]")
	print("where 'reaction' is aa for alpha-alpha, pp for proton-proton, etc.")
	print("reaction names include: pp, aa, pp0, aa0, ap, dd, dp0, da0, da1")
	print("'element' is C-12 for carbon-12, Si-nat for Si with natural abundances of its isotopes, etc.")
	print("See SigmaCalc website for the reactions and elements available")
	print("If 'skip_download' is specified as 3rd argument, the download stage is skipped and interpolation is carried out based on previously saved data.")
	exit()
	
# check if reaction name seems correct	
reactionName = sys.argv[1]
if reactionName not in ['pp', 'aa', 'pp0', 'aa0', 'ap', 'dd', 'dp0', 'da0', 'da1']:
	print("Warning: reaction name '"+reactionName+"' probably invalid")

# create directory, and warn if it already exists and may be overwritten
element = sys.argv[2]
dir = element+"_"+reactionName

if len(sys.argv)==4 and sys.argv[3]=='skip_download':
	if not os.path.exists(os.getcwd()+"/"+dir):
		print("When skipping download, data have to be in directory '" + dir + "'.")
		exit()
else:
	try:
		os.mkdir(dir)
	except Exception:
		print(os.getcwd()+"/"+dir)
		if os.path.exists(os.getcwd()+"/"+dir):
			print("Directory '" + dir + "' already exists, data may be overwritten.")
			r = raw_input('Do you want to continue? (y/n)')
			if not (r=='y' or r=='Y'):
				exit()
		else:
			raise

	# Download the data
	print("Dowloading data for reaction " + element + "(" + reactionName + ") in directory " + dir) 
	prevname = ''
	for angle in range(15,183,1): # angular range and step size
		name = downloadSC(angle, reactionName, element)
		if(name==prevname):
			print('End of '+name+' ?')
			break;
	
		print(name)
		prevname = name


# ---> At this point, we have downloaded and/or saved data. Whatever is present in the directory with name 
# in format 'reactionName_NNNNNN_sc.r33', where NNNNNN is the angle*1000, we will load them to interpolate.
# (The code is structure like that so one could carry out other interpolations without having to re-download
#  the data, could eliminate datasets, etc. Plus it allows checking the data for a couple of problems.)

# from the directory listing, get the available angles. Will be stored as number in 'angleList'.
reaction = element+'_'+reactionName
angleList = []
for file in os.listdir(reaction):
	if file.startswith(reaction):
		angle = float(file.split("_")[2].split(".")[0])/1000.
		angleList.append(angle)
angleList = asarray(angleList)
		
# initialize E and nskip
file = reaction+'/'+reaction+'_'
a, E, nskip, fileForHeader = readsc(file, 75)
fulltable = []

# load data available, checking if data with cross-section = 0 exist at low angle
warningPrinted = False
iangle = 0
for angle in angleList:
	try:
		cs = readsc(file, angle, nskip, E)[0] # readsc checks if E is always the same
	except ValueError:
		print("Server apparently got an error at angle = "+str(angle)+", file '"+file+"', continuing without these data.")
		angleList = delete(angleList, iangle)
		iangle+=1
		continue
		
	if(amin(cs)==0 and angle < 25):
		if not warningPrinted:
			print('Reaction '+reaction+' contains data = 0 at the following angle(s):')
			warningPrinted = True
		print(str(angle)+' '), 
		angleList = delete(angleList, iangle)
	else:
		fulltable.append(cs)
	iangle+=1

if(warningPrinted): 
	print('')
	print('These angles are removed from table to avoid discontinuities resulting in bad interpolations.')

fulltable = asarray(fulltable)		

# check for duplicate energy values
dE = diff(E)
if ~all(dE>0):
	print('Duplicate energy entries removed at' + str(E[dE==0]))
	fulltable = fulltable[:,dE>0]
	E = E[dE>0]
	
# check again!
dE = diff(E)
if ~all(dE>0):
	print('Duplicate energy entries removed at' + str(E[dE==0]))
	fulltable = fulltable[:,dE>0]
	E = E[dE>0]

# create a bicubic interpolation object and generate the fine-grained (0.2 deg) table
# (Corteo carries out a bilinear interpolation)
# print(shape(fulltable), shape(E), shape(angleList))
f = interpolate.interp2d(E, angleList, fulltable, kind='cubic')
angleinterp = arange(amin(angleList), amax(angleList), 0.2)
interptable = f(E, angleinterp)
interptable[interptable<0] = 0


# write the Corteo cross-section (ccs) file
dic = {'pp' : 'H-1_H-1', 'aa': 'He-4_He-4', 'pp0' : 'H-1_H-1b', 'aa0' : 'He-4_He-4', 'ap': 'He-4_H-1', 
  'dd': 'H-2_H-2', 'dp0': 'H-2_H-1', 'da0': 'H-2_He-4', 'da1': 'H-2_He-4b'}
ccsfilename = element+'_'+dic[reactionName]+'_scat.ccs'
ccsfile = open(ccsfilename,"w")

# copy header of first SigmaCalc file
f = open(fileForHeader,"r")
print('Generating '+ccsfilename+'\n')
for line in f.readlines()[0:nskip]:
	line = line.replace('\n', ' ').replace('\r', '')
	if line.startswith('Comment:'):
		# insert our own comment about what is this file
		index = len('Comment:')
		line = line[:index] + ' File generated using importCS.py containing interpolated\n' \
			+ 'cross-section values based on cross-sections downloaded from SigmaCalc.\n' \
			+ 'The rest of the comment is one of SigmaCalc headers.\n' \
			+ 'Please cite this info when using these data.\n' + line[index:]
	if not line.startswith('Theta:'):  # don't write the line indicating the angle as we have all angles
	    ccsfile.write(line + '\n')
	    print(line)
f.close()

# write angles in first row, starting at second positon
ccsfile.write(str(len(E))+' '+str(len(angleinterp))+'\n')

# write energy in first column, then cross-section values for each angle at that energy
ccsfile.write('0 '+ ' '.join(format(a, "10.4e") for a in angleinterp) +'\n')
for i in range(0, len(E)):
	ccsfile.write(str(E[i])+ ' ' + ' '.join(format(cs, "10.4e") for cs in interptable[:, i])+'\n')

ccsfile.close()