#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
write_config.py

Copyright (C) 2009 Nikolaus Rath <Nikolaus@rath.org>

This program can be distributed under the terms of the GNU LGPL.
'''

# NEED MIN DISTANCE: 22.22cm

from __future__ import division, print_function, absolute_import

import sys
from optparse import OptionParser
import logging
import subprocess
import textwrap
import tempfile
import os 
import base64
import matplotlib.pyplot as plt
import multiprocessing
import numpy as np
import glob
import h5py
import shutil
import hashlib
import math
import re

log = logging.getLogger()

def parse_args(args):
    '''Parse command line'''

    parser = OptionParser(
        usage="%prog [options]\n"
              "%prog --help",
        description=textwrap.dedent('''\
        Runs TokaMac.
        '''))

    parser.add_option("--debug", action="store_true", default=False,
                      help="Activate debugging output")
    parser.add_option("--quiet", action="store_true", default=False,
                      help="Be really quiet")
    parser.add_option("-f", dest="configfile", metavar="<file>",
                      default='TokaMac.rc', action='store', type='string',
                      help="Read configuration from <file>. Default: %default")

    (options, pps) = parser.parse_args(args)

    # Verify parameters
    if len(pps) != 0:
        parser.error("Incorrect number of arguments.")

    return options

class Config(object):
    '''Holds TokaMac configuration data'''
    
    pass 

def init_logging(options):
    root_logger = logging.getLogger()
    formatter = logging.Formatter('%(message)s') 
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    if options.quiet:
        root_logger.setLevel(logging.WARN)
    elif options.debug:
        root_logger.setLevel(logging.DEBUG)
    else:
        root_logger.setLevel(logging.INFO)
    root_logger.addHandler(handler)
    
def main(args=None):

    if args is None:
        args = sys.argv[1:]

    options = parse_args(args)
    init_logging(options)
    
    currents = np.linspace(500, 5000, num=20)
    angles = np.linspace(4, 18, 14)
    coil_r = 0.28
    base_angle = 162
       
    pool = multiprocessing.Pool()     
    sep_dist = np.zeros((len(currents), len(angles))) 
    sep_dist_set = np.zeros(sep_dist.shape, dtype='bool') 
    for (i, current) in enumerate(currents):
        results = list()
        
        for angle in angles:
            results.append(pool.apply_async(run_with, 
                                            (options.configfile, angle, current,
                                             base_angle, coil_r)))
        
        for (j, r) in enumerate(results):
            for d in r.get():
                if not sep_dist_set[i,j]:
                    sep_dist[i, j] = d
                    sep_dist_set[i, j] = True
                elif d < sep_dist[i, j]:
                    sep_dist[i, j] = d
  
        #fig = plt.figure()
        #ax = fig.add_subplot(111)
        #for i in range(len(dW[0])):
        #    xdata = [ x[i] for x in dW ]
        #    ydata = np.arange(1,16)
        #    ax.plot(ydata, xdata)
        #fig.savefig('dW-%03dA.png' % i, format='png')
            
    np.savez('sep_axis', sep_dist=sep_dist,
             sep_dist_set=sep_dist_set)
    sep_dist[~ sep_dist_set] = np.max(sep_dist[sep_dist_set])
    
    fig = plt.figure()
    fig.suptitle(u'Coil distance: %.2f, base angle: %d°' % (coil_r, base_angle))
    ax = fig.add_subplot(111)
    cs = ax.contour(angles, currents, np.ma.masked_array(sep_dist, ~sep_dist_set))
    ax.clabel(cs, inline=1, fontsize=10)
    ax.grid(True)
    ax.set_title('Separatrix Distance from Center')
    ax.set_xlabel('Coil separation (degrees)')
    ax.set_ylabel('Coil current (A)')
    fig.savefig('contours-%.2f,%d°.png' % (coil_r, base_angle), format='png')

        

def run_with(configfile, angle, current, base_angle, coil_r):
    '''Run DCON and TokaMac with sligthly changed configuration
    
    This function:
    - Works in a temporary directory
    - Reads TokaMac config from `configfile`
    - Adjusts shaping coil angular position to `angle` degrees
    - Sets shaping coil current to `current`
    - Writes TokaMac and DCON configuration
    - Runs TokaMac
    '''
    
    config = Config()
    execfile(configfile, config.__dict__)
      
    workdir = tempfile.mkdtemp(dir=os.getcwd())
    try:
        R = 0.9733 # 38.320 in according to Drawing in Mike's Notebook
        r = coil_r
        base = base_angle
        assert config.measurements[-2]['Name'] == 'SH_Coil_Current'
        config.measurements[-2]['Value'] = (current, 0.25e3)
        config.coils['SH'] = [  
                (R + r * math.cos((base+angle)*math.pi/180),
                 r * math.sin((base+angle)*math.pi/180), -3),
                (R + r * math.cos(base*math.pi/180), 
                 r* math.sin(base*math.pi/180), 6),
                (R + r * math.cos((base-angle)*math.pi/180),
                 r * math.sin((base-angle)*math.pi/180), -3),
                 ]        
        write_tokamac_config(config, workdir)
        run_tokamac(workdir)
        sep = get_separatrix(workdir)
        fig = plot_surfaces(workdir, [ x[2] for x in sep ])
        fig.suptitle(u'%d A, %d°' % (current, angle))
        fig.savefig('surfaces-%03dA-%02d°.png' % (current, angle), 
                    format='png')
        if not sep:
            log.warn('No separatrix for %d°, %d A (%s)', angle, current, workdir)
        sep = [ math.sqrt((x-R)**2 + y**2) for (x,y, _) in sep ]

    except subprocess.CalledProcessError:
        log.warn('Run failed for %d°, %d A in %s', 
                      angle, current, workdir)
        return []
    else:
        shutil.rmtree(workdir)
         
    return sep
    
def plot_pressure(workdir):
    '''Plot TokaMac pressure profile in `workdir`
    
    Returns a `Figure` instance.
    '''
    f = h5py.File(os.path.join(workdir, 'TokOut_Plasma.h5'), 'r')
    
    Pressure = f['Data-Set-6']
    Y = f[Pressure.attrs['DIMENSION_NAMELIST'][0]]
    X = f[Pressure.attrs['DIMENSION_NAMELIST'][1]]  
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.contour(X, Y, Pressure)
    ax.set_title('Pressure Profile')
    
    return fig
    
def plot_current(workdir):
    '''Plot TokaMac current profile in `workdir`
    
    Returns a `Figure` instance.
    '''
     
    f = h5py.File(os.path.join(workdir, 'TokOut_PsiGrid.h5'), 'r')
    
    Current = f['Data-Set-2']
    Y = f[Current.attrs['DIMENSION_NAMELIST'][0]]
    X = f[Current.attrs['DIMENSION_NAMELIST'][1]]  
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.contour(X, Y, Current)
    ax.set_title('Current Profile')
    
    return fig
        
def plot_surfaces(workdir, extra_contours=None):
    '''Plot TokaMac magnetic surfaces in `workdir`
    
    Return a `Figure` instance
    '''
    
    f = h5py.File(os.path.join(workdir, 'TokOut_PsiGrid.h5'), 'r')
    
    Psi = f['Data-Set-3']
    Z = f[Psi.attrs['DIMENSION_NAMELIST'][0]]
    X = f[Psi.attrs['DIMENSION_NAMELIST'][1]]  
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.contour(X, Z, Psi, 15, colors='black')
    if extra_contours:
        ax.contour(X, Z, Psi, extra_contours, colors='red')
    #ax.contourf(X, Y, Current)
    ax.set_title('Flux Surfaces')
    
    return fig
            
def run_dcon(workdir):
    '''Run DCON in `workdir`
    
    Returns the three lowest energy eigenvalues.
    '''
    
    log.debug('Running dcon')
    
    # Check Cache
    md5 = file_md5(*[os.path.join(workdir, x) for x in 
                    ('dcon.in', 'equil.in', 'vac.in', 'TokOut_DCON.dat')])
    cachedir = os.path.join(CACHEDIR, md5)
    if os.path.exists(os.path.join(cachedir, 'dcon.out')):
        rcopy(cachedir, workdir)
    
    else: 
        with open('/dev/null', 'wb') as fh:
            subprocess.check_call(['dcon'], stdout=fh, stderr=fh, cwd=workdir)
    
        # Add to cache
        if not os.path.exists(cachedir):
            os.makedirs(cachedir)
        for name in [ 'dcon.in', 'equil.in', 'vac.in', 'TokOut_DCON.dat',
                      'dcon.out' ]:
            shutil.copy2(os.path.join(workdir, name), cachedir)
                    
    with open(os.path.join(workdir, 'dcon.out'), 'r') as fh:
        
        for line in fh:
            if line.strip() == 'Total Energy Eigenvalues:':
                break
        else:
            raise RuntimeError()
            
        for _ in range(3):
            fh.next()
            
        dW = list()
        for i in range(3):
            line = fh.next()
            fields = line.split()
            assert int(fields[0]) == i+1
            dW.append(float(fields[3]))
        
    return dW
        
    
def file_md5(*a):
    '''Calculate one MD5 sum for all files in *a'''
    
    m = hashlib.md5()
    
    for name in a:
        with open(name, 'rb') as fh:
            while True:
                buf = fh.read(4096)
                if not buf:
                    break
                m.update(buf)
    
    return base64.b64encode(m.digest(), '+-')
              
def rcopy(src, dest):
    '''Copy all files in `src` to `dest`'''
    
    for name in os.listdir(src):
        shutil.copy2(os.path.join(src, name), dest)
                  
def run_tokamac(workdir):
    '''Run TokaMac in `workdir`'''
    log.debug('Running tokamac')
        
    # Check Cache
    cachedir = os.path.join(CACHEDIR, file_md5(os.path.join(workdir, 'TokIn.dat')))
    
    if os.path.exists(os.path.join(cachedir, 'TokLog.out')):
        rcopy(cachedir, workdir)
    else:   
        with open('/dev/null', 'wb') as fh:
            subprocess.check_call(['TokaMac'], stdout=fh, stderr=fh,
                                  cwd=workdir)
            
        log.debug('Running h4toh5')
        for i in ('TokOut_Plasma.h5', 'TokOut_PsiGrid.h5'):
            if os.path.exists(os.path.join(workdir, i)):
                os.unlink(os.path.join(workdir, i))
        #subprocess.check_call(['h4toh5', 'TokOut_Plasma.HDF'],
        #                      cwd=workdir)
        #os.unlink(os.path.join(workdir, 'TokOut_Plasma.HDF'))
        subprocess.check_call(['h4toh5', 'TokOut_PsiGrid.HDF'],
                              cwd=workdir)
        os.unlink(os.path.join(workdir, 'TokOut_PsiGrid.HDF'))
                    
        # Add to cache
        if not os.path.exists(cachedir):
            os.makedirs(cachedir)
        for name in glob.glob(os.path.join(workdir, 'TokOut_*')):
            shutil.copy2(name, cachedir)
        for name in [ 'TokIn.dat', 'TokLog.out' ]:
            shutil.copy2(os.path.join(workdir, name), cachedir)
                                       
    with open(os.path.join(workdir, 'TokOut_SVDFit.out'), 'r') as fh:
        for _ in range(3):
            fh.next()
            
        line = fh.next().strip()
        match = re.match(r'^Chisq1 = ([0-9.]+), Chisq2 = ([0-9.]+)$', line)
        
    if not match:
        raise RuntimeError('Cannot parse SVDFit.out')
        
    chi1 = float(match.group(1))
    chi2 = float(match.group(2))
    
    log.info('Chi^2 = %.3f, %.3f', chi1, chi2)


def get_separatrix(workdir):
    '''Get separatrix coordinates from TokaMac run in `workdir`
    
    Returns list if (X, Z, Psi) coordinates. If there is no separatrix,
    or if the plasma is limited, return empty list.
    '''
    
    fh = open(os.path.join(workdir, 'TokLog.out'))
    
    while True:
        line = fh.readline()
        if not line:
            break
        if line.strip().startswith('[Free iteration'):
            pos = fh.tell()
    fh.seek(pos)
    
    sep = list()
    while True:
        line = fh.readline().strip()
        if not line:
            break
        match = re.match(r'\[Sep  found at \(X = (-?[0-9.]+), Z = (-?[0-9.]+)\) Psi = (-?[0-9.]+)\]',
                    line)
        if match:
            x = float(match.group(1))
            z = float(match.group(2))
            psi = float(match.group(3))
            sep.append((x, z, psi))
        
        if line.startswith('[Limited plasma,'):
            return []
    
    return sep
    
def write_dcon_config(workdir, n):
    '''Write DCON configuration with torodoidal mode number `n` into `workdir`'''
    
    log.debug('Writing DCON configuration')
   
    with open(os.path.join(workdir, 'dcon.in'), 'w') as fh:
        fh.write(textwrap.dedent('''
            &dcon_control
                    bal_flag=t
                    mat_flag=t
                    ode_flag=t
                    vac_flag=t
                    nn=%d
                    delta_mlow=4
                    delta_mhigh=8
                    delta_mband=0
                    sing_start=0
                    mthvac=960
                    thmax0=1
                    tol_nr=1e-4
                    tol_r=1e-4
                    crossover=1e-2
                    singfac_min=1e-5
                    ucrit=1e4
            /
            &dcon_output
                    crit_break=t
                    ahb_flag=t
                    msol_ahb=1
                    mthsurf0=1
            /
            ''' % n))
        
    with open(os.path.join(workdir, 'equil.in'), 'w') as fh:
        fh.write(textwrap.dedent('''
              &equil_control
                eq_type="tokamac"
                eq_filename="TokOut_DCON.dat"
                jac_type="hamada"
                power_b=0
                power_r=0
                grid_type="ldp"
                psilow=1e-4
                psihigh=0.99
                mpsi=128
                mtheta=128
                newq0=0
                input_only=f
            /
            
            &equil_output
                gse_flag=f
                out_eq_1d=f
                bin_eq_1d=f
                out_eq_2d=f
                bin_eq_2d=f
                out_2d=f
                bin_2d=f
            /
            '''))

    with open(os.path.join(workdir, 'vac.in'), 'w') as fh:
        fh.write(VACIN)
        
        
def write_tokamac_config(config, workdir):
    log.debug('Writing TokaMac configuration')
    
    PsiGrid = config.PsiGrid
    limiters = config.limiters
    measurements = config.measurements
    coils = config.coils
    Plasma = config.Plasma
    separatrix_areas = config.separatrix_areas
    
    CodeControl = {
        'MaxIterFixed': 0,
        'MaxIterFree': 22,
        'NumShells': 0,
        'Oname': 'TokOut',
        'RSname': 'Restart.bin',
        'MGname': 'MGreen.bin',
        'LHname': 'HGreens.bin',
        'SMname': 'SInduct.bin',
        'SGname': 'SGreen.bin',
        'RestartStatus': 0,
        'LHGreenStatus': 0,
        'MGreenStatus': 0,
        'SGreenStatus': 0,
        'SInductStatus': 0,
        }
    
    PsiGrid['Symmetric'] = 0
    CodeControl['NumLimiters'] = len(limiters)
    CodeControl['NumMeasures'] = len(measurements)
    CodeControl['NumCoils'] = len(coils)
    CodeControl['NumSeps'] = len(separatrix_areas)  
    
    fh = open(os.path.join(workdir, 'TokIn.dat'), 'w')
    
    fh.write('K_PsiGrid\n')
    for (key, val) in PsiGrid.iteritems():
        fh.write('    %s = %s\n' % (key, val))
        
    fh.write('K_CodeControl\n')
    for (key, val) in CodeControl.iteritems():
        fh.write('    %s = %s\n' % (key, val))
        
    fh.write('K_Plasma\n')
    for (key, val) in Plasma.iteritems():
        fh.write('    %s = %s\n' % (key, val))
        
    for tuple_ in limiters:
        fh.write('K_Limiter\n')
        fh.write('    Name = %s\n'
                 '    X1 = %.3g\n'
                 '    Z1 = %.3g\n'
                 '    X2 = %.3g\n'
                 '    Z2 = %.3g\n'
                 '    Enabled = 1\n' % tuple_)
        
    for (i, tuple_) in enumerate(separatrix_areas):
        fh.write('K_Separatrix\n')
        fh.write('    Name = Separatrix_%d\n'
                 '    X1 = %.3g\n'
                 '    Z1 = %.3g\n'
                 '    X2 = %.3g\n'
                 '    Z2 = %3.g\n'
                 '    Enabled = 1\n' % ((i,) + tuple_))
                
    coilCurrents = dict()
    for mes in measurements:
        if mes['mType'] == 8 and 'Coil' in mes:
            coilCurrents[mes['Coil']] = mes['Value'][0]
    
    coilNo = dict()
    for (no, (name, subcoils)) in enumerate(coils.iteritems()):
        coilNo[name] = no+1
        
        fh.write('K_Coil\n')
        fh.write('    Name = %s\n' % name)
        fh.write('    Enabled = 1\n')
        fh.write('    InitialCurrent = %s\n' % coilCurrents[name])
        
        for (no, tuple_) in enumerate(subcoils):
            fh.write('K_SubCoil\n')
            fh.write('    Name = %s_%s\n' % (name, no)) 
            fh.write('    X = %.3g\n'
                     '    Z = %.3g\n'
                     '    CurrentFraction = %.3g\n' % tuple_)
                    
    for mes in measurements:
        fh.write('K_Measure\n')
        for (key, val) in mes.iteritems():
            if key == 'Value':
                fh.write('    Value = %.3g\n'
                         '    StdDev = %.3g\n' % val)
                continue
            
            if key == 'Position':
                fh.write('    X = %.3g\n'
                         '    Z = %.3g\n' % val)
                continue
                        
            if key == 'Coil':
                key = 'CoilNum'
                val = coilNo[val]
                
            fh.write('    %s = %s\n' % (key, val))


'''This will be written into vac.in by write_dcon_config'''
VACIN = '''
x    **** Running DCON's input ***

&modes
    ldcon=1
    mth=480
    xiin(1:9)=7*0 1 0
    lsymz=t
    leqarcw=1

    lzio=0
    lgato=0
    lrgato=0
/
    
&debugs
    checkd=f
    check1=f
    check2=f
    checke=f
    checks=f
    wall=f
    lkplt=0
/
    
&vacdat
    ishape=6
    aw=.05
    bw=1.5
    cw=0
    dw=0.5
    tw=.05
    nsing=500
    epsq=1e-5
    noutv=37
    idgt=6
    idot=0
    idsk=0
    delg=15.01
    delfac=1e-3
    cn0=1
/
    
&shape
    ipshp=0
    xpl=100
    apl=1
    a=20
    b=170
    bpl=1
    dpl=0
    r=1
    abulg=0.932
    bbulg=17.0
    tbulg=.02
    qain=2.5
/
    
&diagns
    lkdis=f
    ieig=0
    iloop=0
    lpsub=1
    
    nloop=16
    nloopr=5
    nphil=3
    nphse=1
    xofsl=-0
    ntloop=60
    aloop=.17
    bloop=1.6
    dloop=0.5
    rloop=1.0
    deloop=.001
    mx=21
    mz=21
    nph=0
    xloop=3*2.17 1.45 1.25 0.96 0.82 1.05 0.82 0.96 1.25 1.45 1.05 2.17 1.05 2.17
    zloop=-0.26 0.0 0.26 3*.67 0.45 -0.025 -0.45 3*-0.67 -0.05 0.0 -0.05 0.0
/
    
&sprk
    nminus=0
    nplus=0
    mphi=16
    lwrt11=0
    civ=0.0
    sp2sgn1=1
    sp2sgn2=1
    sp2sgn3=1
    sp2sgn4=1
    sp2sgn5=1
    sp3sgn1=-1
    sp3sgn2=-1
    sp3sgn3=-1
    sp3sgn4=1
    sp3sgn5=1
    lff=0
    ff=1.6
    fv=5*1.6 1.0 3*1.6 1.6 25*1.6
/
'''

CACHEDIR = os.path.join(os.getcwd(), 'cache') 
if __name__ == '__main__':
    main(sys.argv[1:])
#    options = parse_args(sys.argv[1:])
#    init_logging(options)
#    config = Config()
#    execfile(options.configfile, config.__dict__)
#    write_tokamac_config(config, '.')