#!/hiskp2/jost/code/Enthought/System/bin/python
#!/hiskp2/werner/libraries/Python-2.7.12/python
import numpy as np

import infile_handler
import raw_data
import utils
import wick
import subduction
import setup_gevp
import plot

# TODO: pull out irrep as outer loop and do not use it in any of the files
# TODO: refactor towards a more objectoriented design
def main():

  ################################################################################
  # Parameters ###################################################################

  args, config = infile_handler.get_parameters()

  verbose = args.verbose
  continuum_basis = args.basis

  flag_pion        = config.getboolean('parameters', 'Read pion') 
  flag_read        = config.getboolean('parameters', 'Read rho') 
  flag_subduction  = config.getboolean('parameters', 'Subduce')
  flag_contraction = config.getboolean('parameters', 'Contract')
  flag_gevp        = config.getboolean('parameters', 'Create Gevp')
  flag_plot        = config.getboolean('parameters', 'Plot data')
  flag_old         = config.getboolean('parameters', 'Old data')
  flag_ana         = config.getboolean('parameters', 'Rho analysis')
  
  if verbose:
    print '#################################################################'\
                                                               '###############'
    print 'Reading infile'
  
    print flag_pion
    print flag_read
    print flag_subduction
    print flag_contraction
    print flag_gevp
    print flag_plot
    print flag_old
    print flag_ana
  
  sta_cnfg = config.getint('gauge configuration numbers', 'First configuration')
  end_cnfg = config.getint('gauge configuration numbers', 'Last configuration')
  del_cnfg = config.getint('gauge configuration numbers', \
                                                         'Configuration stepping')
  missing_configs = config.get('gauge configuration numbers', \
                                                         'Missing configurations')
  # turns missing configs into list of integers
  if(missing_configs == ''):
    missing_configs = []
  else:
    missing_configs = [int(m) for m in missing_configs.split(',')]
  
  if verbose:
    print sta_cnfg
    print end_cnfg
    print del_cnfg
    print missing_configs
  
  ensemble = config.get('ensemble parameters', 'Ensemble Name')
  T = config.getint('ensemble parameters', 'T')
  
  if verbose:
    print ensemble
    print T
  
  # gamma structure wich shall be averaged. Last entry of gamma must contain all
  # names in LaTeX compatible notation for plot labels
  gamma_dictionary = {
  'gamma_i' :   [1, 2, 3, \
                ['\gamma_1', '\gamma_2', '\gamma_3', '\gamma_i']],
  'gamma_0i' :  [10, 11, 12, \
                ['\gamma_0\gamma_1', '\gamma_0\gamma_2', '\gamma_0\gamma_3', \
                 '\gamma_0\gamma_i']],
  'gamma_50i' : [13, 14, 15, \
                ['\gamma_5\gamma_0\gamma_1', '\gamma_5\gamma_0\gamma_2', \
                 '\gamma_5\gamma_0\gamma_3', '\gamma_5\gamma_0\gamma_i']],
  'gamma_5' :   [5, ['\gamma_5']]
  }
  
  p_max = config.getint('gevp parameters', 'p_max')
  p = config.get('gevp parameters', 'p_cm')
  p = [int(k) for k in p.split(',')]
  gamma_input = config.get('gevp parameters', 'Dirac structure')
  # translates list of names for gamma structures to indices used in contraction 
  # code
  gamma_input = gamma_input.replace(" ", "").split(',')
  gammas = [gamma_dictionary[g] for g in gamma_input]
  
  if verbose:
    print p_max
    print p
    print gamma_input
  
  diagrams = config.get('contraction details', 'Diagram')
  diagrams = diagrams.replace(" ", "").split(',')
  directories = config.get('contraction details', 'Input Path')
  directories = directories.replace(" ", "").replace("\n", "")
  directories = directories.split(',')
  # use the same directory for all diagrams if only one is given
  if(len(directories) == 1):
    directories = directories*len(diagrams)
  
  if verbose:
    print diagrams
    print directories 
  
  outpath = config.get('other parameters', 'Output Path')
  
  if verbose:
    print outpath
  
  sep_rows_sep_mom = config.getboolean('plot details', 'Plot sep_rows_sep_mom') 
  sep_rows_sum_mom = config.getboolean('plot details', 'Plot sep_rows_sum_mom') 
  avg_rows_sep_mom = config.getboolean('plot details', 'Plot avg_rows_sep_mom')
  avg_rows_sum_mom = config.getboolean('plot details', 'Plot avg_rows_sum_mom')
  
  logscale = config.getboolean('plot details', 'Logscale')
  
  bootstrapsize = config.getint('plot details', 'Number of bootstrap samples')
  
  if verbose:
    print sep_rows_sum_mom
    print avg_rows_sep_mom
    print avg_rows_sum_mom
    print logscale
    print bootstrapsize
 
  ############################################################################## 
  # Main
  for p_cm in p:
    if verbose:
      print '################################################################'\
                                                              '################'
      print 'p_cm = ', p_cm

    ############################################################################ 
    # read data for pion
    if flag_pion and 'C2+' in diagrams:
      diagram = 'C2+'
      directory = directories[0]

      lookup_cnfg = raw_data.set_lookup_cnfg(sta_cnfg, end_cnfg, del_cnfg, \
                                                       missing_configs, verbose)
      pion_qn = raw_data.set_lookup_qn(diagram, p_cm, p_max, gammas, skip=flag_ana,
          verbose=verbose)
    
      if flag_old:
        pion_data = raw_data.read_old(lookup_cnfg, pion_qn, diagram, T, 
                                                             directory, verbose)
      else:
        pion_data = raw_data.read(lookup_cnfg, pion_qn, diagram, T, 
                                                             directory, verbose)
      if verbose:
        print 'Pion mass for p_cm = %1d' % (p_cm)
        print pion_qn
        print pion_data.mean(axis=1).apply(np.real)

      # write data
      path = '%s/%s/0_raw-data/' % (outpath, ensemble)
      filename = '%s_p%1i.h5' % ('pion', p_cm)
      utils.write_hdf5_correlators(path, filename, pion_data, 'data', verbose)
      filename = '%s_p%1i_qn.h5' % ('pion', p_cm)
      utils.write_hdf5_correlators(path, filename, pion_qn, 'qn', verbose=False)

      path = '%s/%s/3_gevp-data/' % (outpath, ensemble)
      filename = 'Pion_p%1i.dat' % (p_cm)
      utils.write_ascii_correlators(path, filename, 
                                 pion_data.mean(axis=1).apply(np.real), verbose)

    elif 'C2+' in diagrams:
      # helper function to read all raw data from disk
      path = '%s/%s/0_raw-data/' % (outpath, ensemble)
  
      filename = '%s_p%1i.h5' % ('C2+', p_cm)
      pion_data = utils.read_hdf5_correlators(path+filename, 'data')
      filename = '%s_p%1i_qn.h5' % ('C2+', p_cm)
      pion_qn = utils.read_hdf5_correlators(path+filename, 'qn')

    ############################################################################ 
    # read diagrams for correlators contributing to rho
    if flag_read:
      data = {}
      lookup_qn = {}
      for diagram, directory in zip(diagrams, directories):
  
        if verbose:
          print '\treading data for %s' % (diagram)
        lookup_cnfg = raw_data.set_lookup_cnfg(sta_cnfg, end_cnfg, del_cnfg, \
                                                         missing_configs, verbose)

        # for moving frames, sum of individual component's absolute value (i.e.
        # total kindetic energy) might be larger than center of mass absolute 
        # value. Modify the cutoff accordingly.
  #      p_cm_max = np.asarray([4,5,6,7,4], dtype=int)[p_cm]
        # TODO: that needs to be refactored when going to a larger operator basis
        lookup_qn[diagram] = raw_data.set_lookup_qn(diagram, p_cm, p_max, 
            gammas, skip=flag_ana, verbose=verbose)
    
        if flag_old:
          data[diagram] = raw_data.read_old(lookup_cnfg, lookup_qn[diagram], diagram, T, 
                                                               directory, verbose)
        else:
          data[diagram] = raw_data.read(lookup_cnfg, lookup_qn[diagram], diagram, T, 
                                                               directory, verbose)
        # write data
        # TODO: Writing into same file only works in append mode
        path = '%s/%s/0_raw-data/' % (outpath, ensemble)
        filename = '%s_p%1i.h5' % (diagram, p_cm)
        utils.write_hdf5_correlators(path, filename, data[diagram], 'data', verbose)
        filename = '%s_p%1i_qn.h5' % (diagram, p_cm)
        utils.write_hdf5_correlators(path, filename, lookup_qn[diagram], 'qn', verbose=False)
 
    else:
      # helper function to read all raw data from disk
      path = '%s/%s/0_raw-data/' % (outpath, ensemble)
  
      data = {}
      lookup_qn = {}
      for diagram in diagrams:
        filename = '%s_p%1i.h5' % (diagram, p_cm)
        data[(diagram)] = utils.read_hdf5_correlators(path+filename, 'data')
        filename = '%s_p%1i_qn.h5' % (diagram, p_cm)
        lookup_qn[(diagram)] = utils.read_hdf5_correlators(path+filename, 'qn')
    #print("diagram %s, shape %r, levels %d, axis %r" % (diagram, data[diagram].shape,
    #    data[diagram].index.nlevels, data[diagram].axes))

    ############################################################################
    # Subduction

    if flag_subduction:
      j_ana = 1 if flag_ana else 0
      basis = subduction.get_lattice_basis(p_cm, verbose, j=j_ana)
  
      # List with the names of all contributing irreducible representations and 
      # their multiplicity
      lookup_irreps = basis['Irrep'].unique()
  
      subduced_data = {}
      lookup_qn_irrep = {}
      for diagram in diagrams:
  
        print '\tsubducing data for %s' % diagram
  #      for irrep, multiplicity in lookup_irreps:
        for irrep in lookup_irreps:
  
          print '\t  subducing into %s' % irrep
          # read coefficients for correlation function, little group and irreducible 
          # representation given by diagram, p_cm irrep and mult
          coefficients_irrep = subduction.get_coefficients(diagram, gammas, p_cm, 
              irrep, basis, continuum_basis, verbose)
          lookup_qn_irrep[(diagram, irrep)] = subduction.set_lookup_qn_irrep( \
                                                coefficients_irrep, \
                                                lookup_qn[diagram], verbose)
          subduced_data[(diagram, irrep)] = subduction.ensembles(data[diagram], \
                                                 lookup_qn_irrep[(diagram,irrep)])
          # write data to disc
          path = '%s/%s/1_subduced-data/' % (outpath, ensemble)
          filename = '/%s_p%1i_%s.h5' % (diagram, p_cm, irrep)
          utils.write_hdf5_correlators(path, filename, 
                                  subduced_data[(diagram,irrep)], 'data', verbose)
    else:
      # helper function to read all subduced data from disk
      j_ana = 1 if flag_ana else 0
      basis = subduction.get_lattice_basis(p_cm, verbose, j_ana)

      # List with the names of all contributing irreducible representations and 
      # their multiplicity
      lookup_irreps = basis['Irrep'].unique()

      path = '%s/%s/1_subduced-data/' % (outpath, ensemble)
  
      subduced_data = {}
      for diagram in diagrams:
        for irrep in lookup_irreps:
          filename = '/%s_p%1i_%s.h5' % (diagram, p_cm, irrep)
          subduced_data[(diagram,irrep)] = \
                                utils.read_hdf5_correlators(path+filename, 'data')

    ############################################################################ 
    # Wick contraction

    if flag_contraction:
      correlators = wick.set_lookup_correlators(diagrams)

      contracted_data = {}
      contracted_data_avg = {}

      for correlator in correlators:

        print '\tcontracting data for %s' % correlator 
        for irrep in lookup_irreps:

          if flag_ana:
              # rho analysis
              contracted_data[(correlator,irrep)] = wick.rho(subduced_data, \
                                                           correlator, irrep, verbose)
          else:
              # pipi I=2 analysis
              contracted_data[(correlator,irrep)] = wick.pipi(subduced_data, \
                  correlator, irrep, verbose)

          # write data to disc
          path = '%s/%s/2_contracted-data/' % (outpath, ensemble)
          filename = '/%s_p%1i_%s.h5' % (correlator, p_cm, irrep)
          utils.write_hdf5_correlators(path, filename, \
                             contracted_data[(correlator,irrep)], 'data', verbose)

          # sum over gamma structures. 
          # Only real part is physically relevant at that point
          contracted_data_avg[(correlator,irrep)] = \
              contracted_data[(correlator,irrep)].apply(np.real).\
                                                            sum(level=[0,1,2,3,4,6])
          # sum over equivalent momenta
          contracted_data_avg[(correlator,irrep)] = \
                        contracted_data_avg[(correlator,irrep)].sum(level=[0,1,2,3])
          # average over rows
          contracted_data_avg[(correlator,irrep)] = \
                         contracted_data_avg[(correlator,irrep)].mean(level=[0,1,2])

          # write data to disc
          path = '%s/%s/2_contracted-data/' % (outpath, ensemble)
          filename = '/%s_p%1i_%s_avg.h5' % (correlator, p_cm, irrep)
          utils.write_hdf5_correlators(path, filename, \
                         contracted_data_avg[(correlator,irrep)], 'data', verbose)

    ############################################################################ 
    # Gevp construction

    if flag_gevp:
      print '\tcreating gevp'
      if flag_ana:
        for irrep in lookup_irreps:
          gevp_data = setup_gevp.build_gevp(contracted_data_avg, irrep, verbose)
  
          path = '%s/%s/3_gevp-data/' % (outpath, ensemble)
          filename = 'Gevp_p%1i_%s.h5' % (p_cm, irrep)
          utils.write_hdf5_correlators(path, filename, gevp_data, 'data', verbose)
  
          path = '%s/%s/3_gevp-data/p%1i/%s/' % (outpath, ensemble, p_cm, irrep)
          utils.write_ascii_gevp(path, gevp_data, p_cm, irrep, verbose)
      else:
        for irrep in lookup_irreps:
          gevp_data = contracted_data_avg[("C4", irrep)]
          # NOTE: Delete all rows and colums with only NaN's
          gevp_data = gevp_data.dropna(axis=0,how='all').dropna(axis=1,how='all')

          # loop over target irreps
          for tirr, select in gevp_data.groupby(level=[0]):
            #print(select)
            path = '%s/%s/3_gevp-data/' % (outpath, ensemble)
            filename = 'Gevp_p%1i_%s_%s.h5' % (p_cm, irrep, tirr)
            utils.write_hdf5_correlators(path, filename, select, 'data', verbose)
  
            path = '%s/%s/3_gevp-data/p%1i/%s/%s/' % (outpath, ensemble, p_cm,
                irrep, tirr)
            utils.write_ascii_gevp(path, select, p_cm, irrep, verbose)

    ############################################################################ 
    # Plotting 

    if flag_plot:
      for irrep in lookup_irreps:
  
        path = '%s/%s/4_plots/p%1i/%s/' % (outpath, ensemble, p_cm, irrep)
        for correlator in correlators:
  
          if sep_rows_sum_mom:
            filename = '/%s_sep_rows_sum_mom_p%1i_%s_%s.pdf' % (correlator, p_cm, \
                                                                            irrep, continuum_basis)
            pdfplot = utils.create_pdfplot(path, filename)
            plot.sep_rows_sum_mom(contracted_data[(correlator,irrep)], \
                            correlator, bootstrapsize, pdfplot, logscale, verbose)
            pdfplot.close()
  
          if avg_rows_sep_mom:
            filename = '%s_avg_rows_sep_mom_p%1i_%s.pdf' % (correlator, p_cm, \
                                                                            irrep)
            pdfplot = utils.create_pdfplot(path, filename)
            plot.avg_rows_sep_mom(contracted_data[(correlator,irrep)], \
                            correlator, bootstrapsize, pdfplot, logscale, verbose)
            pdfplot.close()
  
          if sep_rows_sep_mom:
            filename = '%s_sep_rows_sep_mom_real_p%1i_%s.pdf' % (correlator, p_cm, \
                                                                            irrep)
            pdfplot = utils.create_pdfplot(path, filename)
            plot.sep_rows_sep_mom(contracted_data[(correlator,irrep)].apply(np.real), \
                            correlator, bootstrapsize, pdfplot, logscale, verbose)
            pdfplot.close()

            filename = '%s_sep_rows_sep_mom_imag_p%1i_%s.pdf' % (correlator, p_cm, \
                                                                            irrep)
            pdfplot = utils.create_pdfplot(path, filename)
            plot.sep_rows_sep_mom(contracted_data[(correlator,irrep)].apply(np.imag), \
                            correlator, bootstrapsize, pdfplot, logscale, verbose)
            pdfplot.close()
  
        if flag_gevp:
          if avg_rows_sum_mom:
            #gevp_data = setup_gevp.build_gevp(contracted_data_avg, irrep, verbose)
            gevp_data = contracted_data_avg[("C4", irrep)].dropna(axis=0,how='all').dropna(axis=1,how='all')
            filename = 'Gevp_p%1i_%s.pdf' % (p_cm, irrep)
            pdfplot = utils.create_pdfplot(path, filename)
            plot.avg_row_sum_mom(gevp_data, bootstrapsize, pdfplot)
            pdfplot.close()
        else:
          print 'Warning: skipped avg_rows_sum_mom because gevp is incomplete'


################################################################################
if __name__ == '__main__':
  try:
    main()
  except KeyboardInterrupt:
    pass


