#!/usr/bin/env python
#
# Copyright (C) 2017 Michael Janssen
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This library 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 Library General Public
# License for more details.
#
# You should have received a copy of the GNU Library General Public License
# along with this library; if not, write to the Free Software Foundation,
# Inc., 675 Massachusetts Ave, Cambridge, MA 02139, USA.
#
"""
Script to set up the pipeline.
Can be re-run anytime (e.g., if you have a new CASA version or if you have moved the pipeline to a different directory).
"""
import subprocess
import re
import sys
import os
import socket
import multiprocessing
import platform
import optparse
try:
    import pyfits
    gotpyfits = True
except ImportError:
    gotpyfits = False


def strip_path_at_keyword(_path, keyword):
    """
    Cutoff path after the second to last keyword occurence
    """
    _strip_path = _path.split('/')
    _newpath = []
    doappend = False
    for i,sp in enumerate(_strip_path[::-1]):
        if i>0 and keyword.lower() in sp.lower():
            doappend = True
        if doappend:
            _newpath.append(sp)
    _outpath = ''
    for np in _newpath[::-1]:
        _outpath += np + '/'
    return _outpath


def input23(_msg):
    if sys.version_info >= (3, 0):
        _inp = input(_msg)
    else:
        _inp = raw_input(_msg)
    return _inp


def alter_line(filename, search_pattern, new_line, str_char = '#', illegal_string='--Katzenkratzbaum--', strict=False):
    f     = open(filename, 'r')
    lines = f.readlines()
    for i, l in enumerate(lines):
        if l[0] != str_char:
            if strict:
                if re.search(r'\b' + search_pattern + r'\b', l) and illegal_string not in l:
                    lines[i] = new_line
            else:
                if search_pattern in l and illegal_string not in l:
                    lines[i] = new_line
    f.close()
    f = open(filename, 'w')
    for l in lines:
        f.write(l)
    f.close()


def check_CASA_installation(path, _must_have, _crucial_features={}):
    red    = '\033[91m'
    green  = '\033[92m'
    end    = '\033[0m'
    path0  = path.split('bin/casa', 1)[0]
    outmsg = '\n'
    maxlen = max([len(xx) for xx in _must_have])
    for item in must_have:
        _cmd   = 'find ' + path0 + ' -name ' + item
        _out   = os.popen(_cmd).read()
        xspace = maxlen-len(item)
        xspace = ' ' * xspace
        gotit  = False
        if _out:
            gotit = True
            try:
                filecandidates = _out.split('\n')
                for filecandidate in filecandidates[:-1]:
                    filecontent = open(filecandidate, 'r').read()
                    if 'def {0}'.format(item.split('.')[0]) in filecontent:
                        for cfeature in _crucial_features[item]:
                            if cfeature not in filecontent:
                                gotit = False
                                break
                            else:
                                pass
                    else:
                        pass
            except (KeyError, IOError) as _:
                pass
        if gotit:
            _out  = green + xspace + 'True' + end
        else:
            _out = red + xspace + 'False' + end
        outmsg += '      Has correct ' + item + ': ' + _out + '\n'
    return outmsg


parser = optparse.OptionParser()
parser.add_option("-p", type="str", dest="casapath", metavar="<path>",
                  help="Path to CASA installation. Default: search whole local file system.", default='')
parser.add_option("-a", dest="automatic", metavar="<automatic mode>", action="store_true",
                  help="Automatic build mode. Use it only for Docker!", default=False)
(options, args) = parser.parse_args()


if sys.platform != "linux" and sys.platform != "linux2":
    raise OSError('I have detected ' + str(sys.platform) + ' as your operating system. You will need linux to run the pipeline.')

if not options.automatic:
    print('\n*** This script will link your CASA installation to the pipeline. ***\n'
          '*** It can always be executed again.                              ***\n\n')


array_defaults = {'EHT': [['array_type','array_type = EHT\n'],
                          ['refant','refant = AA, LM, AP, AZ, PV, SP\n'],
                          ['fringe_solint_optimize_search_cal','fringe_solint_optimize_search_cal = 0.0001;30\n'],
                          ['fringe_solint_optimize_search_sci','fringe_solint_optimize_search_sci = 0.0001;30\n'],
                          ['fringe_solint_mb_reiterate', 'fringe_solint_mb_reiterate = 60;120\n']
                         ],
                  'GMVA': [['array_type','array_type = GMVA\n'],
                           ['refant','refant = AA, EF, PV, LA, FD, PT, KP\n'],
                           ['fringe_solint_optimize_search_cal','fringe_solint_optimize_search_cal = 10;180\n'],
                           ['fringe_solint_optimize_search_sci','fringe_solint_optimize_search_sci = 10;180\n'],
                           ['fringe_solint_mb_reiterate', 'fringe_solint_mb_reiterate = 240;300\n']
                          ],
                  'VLBAhi (for high frequencies)': [['array_type','array_type = VLBAhi\n'],
                                                    ['refant','refant = LA, FD, PT, KP\n'],
                                                    ['fringe_solint_optimize_search_cal','fringe_solint_optimize_search_cal = 30;400\n'],
                                                    ['fringe_solint_optimize_search_sci','fringe_solint_optimize_search_sci = 30;400\n']
                                                   ],
                  'VLBAlo (for low frequencies)': [['array_type','array_type = VLBAlo\n'],
                                                   ['refant','refant = LA, FD, PT, KP\n'],
                                                   ['fringe_solint_optimize_search_cal','fringe_solint_optimize_search_cal = 60;400\n'],
                                                   ['fringe_solint_optimize_search_sci','fringe_solint_optimize_search_sci = 60;400\n']
                                                  ],
                  'EVN': [['array_type','array_type = EVN\n'],
                          ['refant','refant = EF, YS, MC, NT\n'],
                          ['fringe_solint_optimize_search_cal','fringe_solint_optimize_search_cal = 60;400\n'],
                          ['fringe_solint_optimize_search_sci','fringe_solint_optimize_search_sci = 60;400\n']
                         ]
                 }


must_have        = ['mpi', 'accor.py', 'fringefit.py', 'gaincal.py']
crucial_features = {'fringefit.py': ['globalsolve', 'delaywindow', 'ratewindow', 'niter', 'concatspws'],
                    'gaincal.py': ['solmode', 'rmsthresh'],
                    'accor.py': ['corrdepflags']
                   }
if not options.automatic:
    print ('First we will try to find a suitable CASA installation.\nThe required features for this pipeline are:\n' + \
           str(must_have)
          )
    print ('\nIt is highly recommended to use the exact same CASA version as advertised in the README.md file\n')

#see also https://casa.nrao.edu/installlinux.shtml

pipedir0 = os.path.dirname(os.path.realpath(sys.argv[0])) + '/'
pipedir  = pipedir0 + 'picard/'

if not options.casapath:
    if not options.automatic:
        print ('Looking for CASA in your PATH:')
    cmd       = "which casa"
    proc      = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err  = proc.communicate()
    locsearch = False
    if out:
        if isinstance(out, list):
            out = [iterout.decode().rstrip('\n') for iterout in out]
        else:
            out = [out.decode().rstrip('\n')]
        print ('Printing the versions I found below.\n' \
               ' Please enter the number for the installation you want to use\n' \
               ' or enter <search> to search your local system instead of just your PATH:'
              )
        msg = ''
        for i,x in enumerate(out):
            msg += '\n{0} for {1} {2}\n'.format(str(i), x, check_CASA_installation(x, must_have, crucial_features))
        msg += '\nsearch for searching your whole local system instead of just your PATH\n'
        msg += '\nEnter the number corresponding to the installation you want to use\n' \
               ' or type search to search your local system and press Enter\n' \
               ' after you have installed the correct CASA version from the list printed above.\n >'
        inp = input23(msg)
        if 's' in inp:
            locsearch = True
        else:
            exe = out[int(inp)]
    else:
        msg = '\nDid not find a CASA executable in your path.\n'
        msg +='Press Enter and I will continue to search the whole local filesytem for an executable\n'
        msg += 'after you have installed the correct CASA version from the list printed above.\n >'
        if not options.automatic:
            inp = input23(msg)
        locsearch = True
else:
    locsearch = True

if locsearch:
    if not options.casapath:
        if not options.automatic:
            print ('\nLooking for CASA executables on your local system. This may take up to a minute...\n')
        #find executable(s):
        cmd = "find / /home /homes -xdev -name casa"
    else:
        print ('\nLooking for CASA executables in ' + options.casapath + '\n')
        cmd = "find " + os.path.abspath(options.casapath) + " -name casa"
    #pick the right executable:
    if sys.platform == "linux" or sys.platform == "linux2":
        pattern  = re.compile(b".*/casa[^/]*/bin/casa", re.IGNORECASE)
    else:
        pattern  = re.compile(b".*/casa[^/]*/Contents/MacOS/casa", re.IGNORECASE)
    proc     = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = proc.communicate()
    out      = out.splitlines()
    exe      = []
    for o in out:
        if re.match(pattern, o) and '/lib/casa/bin/casa' not in o.decode('utf-8'):
            exe.append(o.decode('utf-8'))
    if len(exe) > 1:
        print('I have found more than one CASA installation.\nPlease enter the number for the installation you want to use:')
        msg = ''
        for i,x in enumerate(exe):
            msg += '\n{0} for {1} {2}\n'.format(str(i), x, check_CASA_installation(x, must_have, crucial_features))
        msg += '\nEnter the number corresponding to the installation you want to use and press Enter\n >'
        inp = input23(msg)
        exe = exe[int(inp)]
    elif exe:
        exe = exe[0]
    if exe:
        if options.automatic:
            inp = ''
        else:
            inp = input23('\nFound\n ' + exe + '\n as your CASA executable.\n' \
                          ' Checking this CASA version:\n' \
                          + check_CASA_installation(exe, must_have, crucial_features) + '\n' \
                          ' Press Enter and I will use the absolute path to this executable for picard.\n' \
                          ' Write anything else (and then press Enter) to abort.\n >'
                         )
        if inp:
            print ('\nAbort.\n')
            sys.exit()
        else:
            pass
    else:
        raise IOError('\nDid not find a CASA executable.\n' \
                      'Please make sure that CASA is properly installed\n'
                     )
exe1     = exe
exe      = exe.rstrip('casa')
if not options.automatic:
    print('\nWriting the CASA executable path to a <your_casapath.txt> file, which will be used by picard.\n')
casapath = open(pipedir0+'your_casapath.txt', 'w')
casapath.write(os.path.abspath(exe))
casapath.close()

if not options.automatic:
    print ('\nMaking picard executable.\n')
os.system('chmod +x ' + pipedir + 'picard')

if not os.path.isdir(pipedir+'input'):
    if not options.automatic:
        print ('No input folder found. Creating one from the default input files stored in input_template.')
    os.system('cp -r ' + pipedir0 + 'input_template ' + pipedir + 'input')
    mpifile = open(pipedir+'input/mpi_host_file', 'w')
    mpifile.write('#See instructions on how to set up hostfiles for MPI: https://www.open-mpi.org/faq/?category=running#mpirun-hostfile.\n')
    mpifile.write('#Syntax for single computer: <pc-name> slots=<number of cores that you want to use (+1)>.\n')
    mpifile.write('mjpc slots=4')
    mpifile.close()

pcname   = socket.gethostname()
ncores   = int(max(multiprocessing.cpu_count(), 1))
if not options.automatic:
    print ('\nEditing the input/mpi_host_file using the determined name of this computer (' + pcname + ')\n' \
           'and ' + str(ncores) + ' cores. Change this setup manually if desired.\n'
          )
mpisetup = '{0} slots={1}'.format(pcname, str(ncores+1))
alter_line(pipedir+'input/mpi_host_file', '', mpisetup)

#cmd      = "find " + strip_path_at_keyword(exe, 'casa') + " -name casapy.py"
#proc     = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
#out, err = proc.communicate()
#out      = out.splitlines()
#exepy    = []
#if out:
#    for o in out:
#        exepy.append(o.decode('utf-8'))
#    if len(exepy) > 1:
#        print('I have found more than one casapy launch script.')
#        print('From the selection below, enter the one that matches to your executable\n ' + exe1)
#        print('Please enter the number for the script you want to use:')
#        msg = ''
#        for i,x in enumerate(exepy):
#            msg += '\n{0} for {1}\n'.format(str(i), x)
#        msg  += '\nEnter the number corresponding to the installation you want to use and press Enter\n >'
#        inp   = input23(msg)
#        exepy = exepy[int(inp)]
#    else:
#        exepy = exepy[0]
#else:
#    raise IOError('\nDid not find a CASA launch script.\n' \
#                  ' If CASA is installed, then this script should be in the casa-release-xxx/lib/python2.7/ folder.\n'
#                 )

if not options.automatic:
    print ('I could put some default values for array.inp depending on which array you inted to use.')
msg = ''
for i,x in enumerate(array_defaults.keys()):
    msg += '\n  {0} for {1}\n'.format(str(i), x)
msg += 'Press enter without entering anything else to continue without altering your array.inp file.\n'
msg += 'Else, enter the number corresponding to the array you want to use and press Enter\n >'
if options.automatic:
    inp = ''
else:
    inp = input23(msg)
if inp:
    array_inputs = array_defaults[array_defaults.keys()[int(inp)]]
    for array_inp in array_inputs:
        alter_line(pipedir+'input/array.inp', array_inp[0], array_inp[1], strict=True)
elif not options.automatic:
    print ('Continuing without touching array.inp')

if not options.automatic:
    print ('\n\nThe pipeline should be ready to run now.\n' \
           'If there are issues with mpicasa contact M.Janssen@astro.ru.nl\n' \
           'or look at https://casa.nrao.edu/casadocs/@@search?SearchableText=mpi\n'
          )
if not options.automatic:
    print('\nIf you want to be able to run the pipeline from everywhere,\n' \
          'then you should add the following line to your .bashrc folder:\n' \
          'export PATH=$PATH:'+pipedir+'\n'
         )
#print ('\nMost python installations should be complete enough to run the pipeline.\n' \
#       'The only module that you may need to install is python-pyfits.\n' \
#       'You have to make sure that you install it for the python version used by CASA.\n'
#      )
#pyvers = platform.python_version()
#if gotpyfits:
#    print ('At least for the python version used to call this script: ' + pyvers + ',\n' \
#           'it looks like pyfits is installed.\n'
#          )
#else:
#    print ('At least for the python version used to call this script: ' + pyvers + ',\n' \
#           'it looks like pyfits is not installed.\n'
#          )
if options.automatic:
    print ('\n ** Finished the automatic setup of the CASA VLBI pipeline within your Docker image. ** \n')
print ('\nRemember set some input parameters in the beginning, before running the pipeline.\n' \
       'At least edit input/observation.inp and input/array.inp\n'
      )
print ('\nPlease read documentation/picard_documentation.pdf and follow the Quick Start Guide chapter to get started.')