#!/usr/bin/python
import sys,os


if len(sys.argv) == 1:  
    print ">>> Need, at least, a name of data-file!"
    print '>>> For example, "./data2F90.py AMS02_posifrac.txt AMS02_electron.txt AMS02_positron.txt AMS02_total_elec.txt PAMELA_antip.txt"'
    print '>>> Do you want to use default data-sets?'
    use_default=raw_input('Yes or No?\n')
    if 'Yes' in use_default.strip():
      flist=['dat/ep_ap_bkg_expcon/AMS02_posifrac.txt','dat/ep_ap_bkg_expcon/AMS02_electron.txt',
             'dat/ep_ap_bkg_expcon/AMS02_positron.txt','dat/ep_ap_bkg_expcon/AMS02_total_elec.txt',
             'dat/ep_ap_bkg_expcon/PAMELA_antip.txt']
    else:
      exit()
else:
    flist=sys.argv[1:]



def write_common_block(head_,x):

   content='    '+head_
   
   k=len(x)/4
   l=len(x)-k*4
   for k_ in xrange(k):
      for i in xrange(4):  
         if k_*4+i+1 == len(x):
           content+=x[k_*4+i]+'/)\n'
           return content
         else:
           content+=x[k_*4+i]+','         
      content+=' & \n       ' 

     
   if l==1:
      content+=x[k*4+0]+'/)\n' 
   elif l==2: 
      content+=x[k*4+0]+','+x[k*4+1]+'/)\n' 
   elif l==3: 
      content+=x[k*4+0]+','+x[k*4+1]+','+x[k*4+2]+'/)\n'  

   return content 


def FluxForm(fname):

   print '*****************************************'
   print 'What is the observable in the data-files?'
   print '*****************************************'
   print ' 0. flux (electron or proton)'
   print ' 1. flux (positron or anti-proton)'
   print ' 2. flux (positron + electron)'
   print ' 3. positron/(positron + electron)'
   print ' 4. anti-proton/proton'
   MyOpt=raw_input('Enter the Observable Form with integer number in: '+fname+'\n --->>> ').strip()



   content='\n    do i=1,100,1 \n'

   if int(MyOpt)<3:
     print '*****************************************'
     print 'What is the Flux Form ?'
     print '*****************************************'
     print ' 1. flux'
     print ' 2. E^2*flux'
     print ' 3. E^3*flux'
     print '*****************************************'   
     MyOpt+='_'+raw_input('What is Flux Form with integer number? \n --->>> ').strip()




   if ('1_1' in MyOpt) or ('2_1' in MyOpt) or ('0_1' in MyOpt):
     print '\n Using the default y-Form: flux' 

     if MyOpt.startswith('2_') : content+='       y(i)=pflux(i)+apflux(i) \n'
     if MyOpt.startswith('1_') : content+='       y(i)=apflux(i) \n'       
     if MyOpt.startswith('0_') : content+='       y(i)=pflux(i) \n'         


   elif ('1_2' in MyOpt) or ('2_2' in MyOpt) or ('0_2' in MyOpt):
     print '\n Using the y-Form: E^2*flux' 

     if MyOpt.startswith('2_') : content+='       y(i)=en(i)*en(i)*(pflux(i)+apflux(i)) \n'
     if MyOpt.startswith('1_') : content+='       y(i)=en(i)*en(i)*apflux(i) \n'   
     if MyOpt.startswith('0_') : content+='       y(i)=en(i)*en(i)*pflux(i) \n'   


   elif ('1_3' in MyOpt) or ('2_3' in MyOpt) or ('0_3' in MyOpt):
     print '\n Using the y-Form: E^3*flux' 
     if MyOpt.startswith('2_') : content+='       y(i)=en(i)**3*(pflux(i)+apflux(i)) \n'
     if MyOpt.startswith('1_') : content+='       y(i)=en(i)**3*apflux(i) \n'   
     if MyOpt.startswith('0_') : content+='       y(i)=en(i)**3*pflux(i) \n'   


   elif MyOpt in '3':

     print '\n Using the y-Form: positron/(positron + electron)' 
     content+='       y(i)=apflux(i)/(pflux(i)+apflux(i)) \n'

   elif MyOpt in '4':
     print '\n Using the y-Form: anti-proton/proton' 
     content+='       y(i)=apflux(i)/pflux(i) \n'

   else:
     print 'Not clear your choice, using default option 1 instead.'
     content='\n      y(i)=apflux(i)'       
   
   content+='    enddo\n'

   return content


F90out = open('charge_data.f90', 'w')

contents='''
!############################################################################
!#
!# Note: A module generated by python script "Data2fortran.py" to deal with 
!#       multi-data-sets 
!#
!# Author: Yue-Lin Sming Tsai
!# Email: smingtsai@gmail.com
!# Date: 2015-6-12
!############################################################################


module charge_data
   use MathLib 

   implicit none

   !first col 2 for 2 types particles (e+ and pbar), 2nd col 2 for (En, flux), 100 rows of energy bins
   real*8 :: PreMod__DMpred(2,2,100)   


contains

'''


ii=1
for fname in flist:
   fin = open(fname)
   lines = fin.readlines()
   fin.close()
   print '\n\n\n+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++'
   print str(ii)+') Doing data file:', fname
   print 'What is the column numbers for (energy, f, fmin, fmax)? '
   print '(Energy always in GeV, length in meter, time in sec, and angle in sr.)'
   colform=raw_input("Enter column numbers of 'energy, f, fmin, fmax', e.g., 0 2 3 4 or 0,1,4,5 : \n")
   print '+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++'
   colform=colform.replace(',',' ').split()
   try:
      nx=int(colform[0])
      ny=int(colform[1])
      nymin=int(colform[2])
      nymax=int(colform[3])
   except:
      print '\n Do not detect your inputs in file ',fname
      print '\n Use the default setting: (0,1,4,5). '
      nx,ny,nymin,nymax=0,1,4,5

   # Reading the data in correct column number 
   x=[]
   y=[]
   ymax=[]
   ymin=[]
   for line in lines:
      if not line.startswith("#"):
        info=line.split()
        x.append(info[nx])
        y.append(info[ny])
        ymax.append(info[nymax])
        ymin.append(info[nymin])


   f=fname.split('/')[-1].split('.')[0]
   contents+='  subroutine '+f+"""(en,pflux,apflux,chisq)
    implicit none
    real*8, intent(in) :: en(100),pflux(100),apflux(100)
    real*8, intent(out) :: chisq 
    real*8 y(100)
    integer i
    real*8 sigma,theo
    real*8 xdata("""+str(len(x))+""")
    real*8 ydata("""+str(len(y))+""")
    real*8 yup("""+str(len(ymax))+""")
    real*8 ylo("""+str(len(ymin))+')\n'


   cX=write_common_block('xdata=(/',x)
   cY=write_common_block('ydata=(/',y)
   cYup=write_common_block('yup=(/',ymax)
   cYlo=write_common_block('ylo=(/',ymin)

   yform=FluxForm(fname)

   contents+=cX+cY+cYup+cYlo+yform


   contents+='''
    chisq=0d0
    do i=1,'''+str(len(x))+''',1
      call linInt2p(en,y,100,.true.,xdata(i),theo,.false.)
      sigma=(yup(i)-ylo(i))/2d0
      chisq=chisq+(theo-ydata(i))**2/sigma**2
    enddo
    return 
  end subroutine '''+f+'\n\n\n'

   ii+=1

contents+='''

end module charge_data

'''
#print contents
F90out.write(contents)
F90out.close()





