import w7x
import numpy as np
import warnings
import transcoding as tc
import tfields
import subprocess
import logging
from past.builtins import basestring


class mgrid_file:
    ''' A set of routines to read/write mgrid files for VMEC/EXTENDER'''

    def __init__(self, parameter):
        if isinstance(parameter, basestring):
            self.read(parameter)
        else:
            self.fromWebserviceConfig(parameter)

    def toWebserviceConfig_flt(self):
        '''
        Store the mgrid file contents in an equivalent field line tracer compatible
        magnetic conf datatype.
        '''
        tracer = w7x.get_server(w7x.Server.addr_field_line_server)

        # Figure out the coil groups (same coilgroup and current)
        coils = []
        currents = []
        for coil in self.coils:
            # Prepare the vertices
            pf1 = tracer.types.PolygonFilament()
            pf1.vertices = tracer.types.Points3D()
            pf1.name = coil['name']
            pf1.vertices.x1 = coil['x'].tolist()
            pf1.vertices.x2 = coil['y'].tolist()
            pf1.vertices.x3 = coil['z'].tolist()

            coils.append(pf1)
            currents.append(coil['I'][0])

        config = tracer.types.MagneticConfig()
        config.coils = coils
        config.coilsCurrents = currents

        if self.filename is None:
            config.name = 'Config from mgrid_coilfile.py'
        else:
            config.name = self.filename

        # Add a grid for caching the Biot-Savart law stage.
        gridSymmetry = 5
        grid = tracer.types.Grid()
        grid.fieldSymmetry = gridSymmetry

        cyl = tracer.types.CylindricalGrid()
        cyl.numR, cyl.numZ, cyl.numPhi = 181, 181, 481
        cyl.RMin, cyl.RMax, cyl.ZMin, cyl.ZMax = 4.05, 6.75, -1.35, 1.35

        grid.cylindrical = cyl
        config.grid = grid

        return config

    def toWebserviceConfig_vmec(self):
        extender = w7x.get_server(w7x.Server.addr_extender_server)

        # Figure out the coil groups (same coilgroup and current)
        coilGroups = []
        for coil in self.coils:
            coilGroups.append([coil['coilgroup'], abs(coil['I'][0])])

        # Pick unique rows:
        coilGroups = np.array(coilGroups)
        coilGroups = np.vstack({tuple(row) for row in coilGroups})

        print('There will be {} cirquits (coil groups).'.format(len(coilGroups)))
        circuits = []  # coilgroup
        coilCurrents = []
        for i in range(0, len(coilGroups)):

            sc = extender.types.SerialCircuit()
            sc.currentCarrier = []

            for coil in self.coils:
                # Pick only coils belonging to the current coil group
                if coil['coilgroup'] != int(coilGroups[i][0]) or np.abs(np.abs(coil['I'][0]) - np.abs(coilGroups[i][1])) > 1e-3:
                    continue
                
                sc.name = 'coil group'
                sc.current = np.abs(coil['I'][0])  # The current direction goes to the winding direction

                # Prepare the vertexes
                pf1 = extender.types.PolygonFilament()
                pf1.vertices = extender.types.Points3D()
                pf1.vertices.x1 = coil['x']
                pf1.vertices.x2 = coil['y']
                pf1.vertices.x3 = coil['z']

                # Set up the coil parameters
                coil1 = extender.types.Coil()
                coil1.name = coil['name']
                if coil['I'][0] > 0:
                    coil1.windingDirection = 1.
                else:
                    coil1.windingDirection = -1.
                coil1.numWindings = 1

                # Include the vertexes in the coil parameters
                coil1.currentCarrierPrimitive = [pf1]

                # include the coil paramaters in the coil group
                sc.currentCarrier.append(coil1)

            # include the coil group in the config
            circuits.append(sc)
            coilCurrents.append(sc.current)

        config = extender.types.MagneticConfiguration()
        config.circuit = circuits
        if self.filename is not None:
            config.name = self.filename
        else:
            config.name = 'Config from mgrid_coilfile.py'

        return config, coilCurrents

    def fromWebserviceConfig(self, config):
        import numpy as np
        coils = []
        # for coilgroup in range(0,len(config))
        for coilgroup, circuit in enumerate(config.circuit):
            if np.abs(circuit.current) == 0.0:
                print('No current in circuit/coilgroup {}/{} => '
                      'skipping.'.format(coilgroup, coilgroup + 1))
                continue
            for coil in circuit.currentCarrier:
                name = coil.name
                nWinds = coil.numWindings
                direction = coil.windingDirection
                x = np.array(coil.currentCarrierPrimitive[0].vertices.x1).T
                y = np.array(coil.currentCarrierPrimitive[0].vertices.x2).T
                z = np.array(coil.currentCarrierPrimitive[0].vertices.x3).T
                I = circuit.current * direction * nWinds * np.ones_like(x)
                I[-1] = 0.0
                coils.append({'x': x, 'y': y, 'z': z, 'I': I,
                              'coilgroup': coilgroup + 1,
                              'name': name})
        self.coils = coils
        self.filename = None
        self.nPeriods = 5

    def write(self, filename):
        with open(filename, 'w') as f:
            # Header
            f.write('periods {}\n'.format(self.nPeriods))
            f.write('begin filament\n')
            f.write('mirror NULL')  # The new line comes from next print

            # Each coil
            for c in self.coils:
                for i in range(0, len(c['x'])):
                    f.write('\n{: 12.8f} {: 12.8f} {: 12.8f} {: 15.8e}'
                            .format(c['x'][i], c['y'][i], c['z'][i], c['I'][i]))
                f.write(' {} {}'.format(c['coilgroup'], c['name']))
            f.write('\nend\n')

    def read(self, filename):
        import numpy as np

        with open(filename, 'r') as f:

            # Read number of periods.
            str = f.readline()
            if str[0:8] != 'periods ':
                raise ValueError('Missing "periods " in the beginning of file.')
            nPeriods = int(str[8:])

            # Read begin filament
            str = f.readline()
            if str[0:14] != 'begin filament':
                raise ValueError('Missing "begin filament" in the beginning of file.')

            # Read/skip "mirror NULL"
            f.readline()

            # Read coils
            coils = []
            while True:
                # Read segments
                x = np.array([], float)
                y = np.array([], float)
                z = np.array([], float)
                I = np.array([], float)
                while True:
                    strs = f.readline().split()
                    v = [float(strs[i]) for i in range(0, 4)]
                    x = np.append(x, v[0])
                    y = np.append(y, v[1])
                    z = np.append(z, v[2])
                    I = np.append(I, v[3])

                    # Each coil is marked to end with a 0 Amps segment
                    if v[3] == 0.0:
                        coilgroup = int(strs[4])
                        name = strs[5:]
                        coils.append({'x': x, 'y': y, 'z': z, 'I': I,
                                      'coilgroup': coilgroup,
                                      'name': name})
                        break
                # Check if we are at the end of file
                last_pos = f.tell()
                str = f.readline()
                if str[0:3].lower() == 'end':
                    break
                f.seek(last_pos)  # Go back where we were.
        self.filename = filename
        self.coils = coils
        self.nPeriods = nPeriods

    def scaleCurrents(self, scale):
        for i in range(0, len(self.coils)):
            self.coils[i]['I'] = scale * self.coils[i]['I']
        if not hasattr(self, 'currentScalingHistory'):
            self.currentScalingHistory = []
        self.currentScalingHistory.append(scale)


class Run(object):
    def __init__(self, vmec_id):
        self._vmec_run = None
        self.vmec_run = w7x.vmec.Run(vmec_id)

    @property
    def vmec_run(self):
        return self._vmec_run

    @vmec_run.setter
    def vmec_run(self, vmec_run):
        self._vmec_run = vmec_run

    def get_coils(self):
        """
        Mostly taken from joachim geiger
        """
        log = logging.getLogger()
        extenderServer = w7x.get_server(w7x.Server.addr_extender_server)
        coilsDbServer = w7x.get_server(w7x.Server.addr_coils_db)
        coilCurrents = self.vmec_run.magnetic_config.coil_currents('A')
        # read the W7-X coils data from CoilsDB...

        # coilCurrents[0:5] = non-planar coils
        # coilCurrents[5:7] = planar coils
        # coilCurrents[7]   = sweep coils half-modules 0
        # coilCurrents[8]   = sweep coils half-modules 1

        if len(coilCurrents) > 7:
            print('Using also sweep coils')
            sweepCoils = True
        else:
            sweepCoils = False

        for i in range(7):
            coilCurrents[i] = -coilCurrents[i]

        log.info('Generate MagneticConfiguration from CoilsDB!')
        log.info('Define Coils')
        # data used in the codes lowendel and w7
        # coils collected as one type, i.e. having the same coil current
        # -------------------------
        type1 = [0, 5, 10, 15, 20, 25, 30, 35, 40, 45]
        type2 = [1, 6, 11, 16, 21, 26, 31, 36, 41, 46]
        type3 = [2, 7, 12, 17, 22, 27, 32, 37, 42, 47]
        type4 = [3, 8, 13, 18, 23, 28, 33, 38, 43, 48]
        type5 = [4, 9, 14, 19, 24, 29, 34, 39, 44, 49]
        typeA = [50, 52, 54, 56, 58, 60, 62, 64, 66, 68]
        typeB = [51, 53, 55, 57, 59, 61, 63, 65, 67, 69]
        # mirrored coils
        # -------------------------
        type1m = [80, 85, 90, 95, 100, 105, 110, 115, 120, 125]
        type2m = [81, 86, 91, 96, 101, 106, 111, 116, 121, 126]
        type3m = [82, 87, 92, 97, 102, 107, 112, 117, 122, 127]
        type4m = [83, 88, 93, 98, 103, 108, 113, 118, 123, 128]
        type5m = [84, 89, 94, 99, 104, 109, 114, 119, 124, 129]
        typeAm = [130, 132, 134, 136, 138, 140, 142, 144, 146, 148]
        typeBm = [131, 133, 135, 137, 139, 141, 143, 145, 147, 149]

        # sweep coils 0-halfmodules and 1-halfmodules
        # --------------------------
        sweep0 = range(150, 160, 2)
        sweep1 = range(151, 160, 2)

        # assure stellarator symmetry by combining coils with mirrored coils
        # -------------------------
        l_sym = True
        # combination of coils and mirrored coils to ensure stellarator symmetry
        # -------------------------
        if l_sym:
            type1 = type1 + type1m
            type2 = type2 + type2m
            type3 = type3 + type3m
            type4 = type4 + type4m
            type5 = type5 + type5m
            typeA = typeA + typeAm
            typeB = typeB + typeBm

        # generate list of the coil types used.
        # -------------------------
        coilIds = [type1, type2, type3, type4, type5, typeA, typeB]
        if sweepCoils:
            coilIds = coilIds + [sweep0, sweep1]
        circuits = []
        # i=0
        # for c in coilIds:
        #     if len(c) == 20 :
        #         coilCurrents[i]=coilCurrents[i]/2.0
        #     i=i+1

        log.info(' Retrieve Coil geometries from CoilsDB ')
        log.info(' and build up MagneticConfiguration ')
        # explanation to current implementation:
        #    the current implementation is not quite as it consistent as it could be
        #    but it is the most convenient one. Each filament is treated a coil and not as one
        #    might like when being a purist as a filament. In the later sense, a coil would consist
        #    of two filaments and the currentCarrierRatios could be used correctly.
        #    To do this the filaments of the typexm coils would have to be extracted and to be added
        #    either to the typex-coils or a new twofilament coil would have to be built up.
        #    The current treatment uses the convenience of adding the two coil setups by
        #    concatenating one after the other. Therefore the coil current has to be divided by 2
        #    in order to get the correct field strength. The currentCarrierRatios are commented
        #    presently (might change with a more consistent treatment in future).

        for i in range(0, len(coilIds)):
            sc = extenderServer.types.SerialCircuit()
            sc.name = 'coil group'
            sc.current = coilCurrents[i]
            sc.currentCarrier = []
            coils = coilsDbServer.service.getCoilData(coilIds[i])
            l = 0
            for c in coils:
                pf1 = extenderServer.types.PolygonFilament()
                pf1.vertices = extenderServer.types.Points3D()
                pf1.vertices.x1 = c.vertices.x1
                pf1.vertices.x2 = c.vertices.x2
                pf1.vertices.x3 = c.vertices.x3
                coil1 = extenderServer.types.Coil()
                coil1.name = c.name
                coil1.windingDirection = 1.
                if i < 5:
                    coil1.numWindings = 108
                    # Mirrored coils have half the windings
                    if len(type1) == 20:
                        coil1.numWindings = 54
                        # sc.current = coilCurrents[i]
                        if l > 9:
                            coil1.windingDirection = -1.
                elif i < 7:
                    coil1.numWindings = 36
                    # Mirrored coils have half the windings
                    if len(typeA) == 20:
                        coil1.numWindings = 18
                        # sc.current = coilCurrents[i]
                        if l > 9:
                            coil1.windingDirection = -1.
                else:
                    # sweep coils
                    coil1.numWindings = 8

                l = l + 1
                coil1.currentCarrierPrimitive = [pf1]
                sc.currentCarrier.append(coil1)

            circuits.append(sc)

        log.info(' Retrieval of Coil geometries from CoilsDB done ')
        # use the coil data for a 'getExtendedField' call ...
        # -------------------------
        log.info(' ... and finalize MagneticConfiguration for extender call ')

        my_config = extenderServer.types.MagneticConfiguration()
        my_config.circuit = circuits
        if sweepCoils:
            my_config.name = 'w7x_sweepCoils'
        else:
            my_config.name = 'w7x'

        log.info('Finalization of MagneticConfiguration done ')

        return my_config


def writeCoils(vmec_id, coil_file='wout.coil'):
    w7x_config = Run(vmec_id).get_coils()
    G = mgrid_file(w7x_config)
    # invert currents
    print('Scaling the coil currents by -1')
    G.scaleCurrents(-1.0)
    G.write(coil_file)


def write_control_file(control_file='grid_in'):
    gridFile = (
        "nr 173   # number of radial grid points\n"
        "nz 170   # number of vertical grid points\n"
        "nphi 141  # number of toroidal cut planes, never be fooled to believe"
        " you can save too much on nphi\n"
        "rmin 3.810000 # optional param., can be computed autom. from"
        " boundary\n"
        "rmax 6.870000 # optional parameter\n"
        "zmax 1.450000 # optional parameter\n")

    with open(control_file, 'w') as f:
        f.write(gridFile)


def write_parameters_file(vmec_calc_url, vmec_id,
                          param_file='parameters.txt',
                          net_cdf_file="wout.nc"):
    with open(param_file, 'w') as f:
        f.write(vmec_calc_url + '\n')
        f.write(vmec_id + '\n')
        f.write(net_cdf_file + '\n')


def read_cylindrical(hybrid_path):
    hybridFileBox = tc.Block(
        "{header}",
        "cyl.RMin, cyl.RMax = {RMin:f}, {RMax:f}",
        "cyl.ZMin, cyl.ZMax = {ZMin:f}, {ZMax:f}",
        "cyl.numR, cyl.numZ, cyl.numPhi = {numR:d}, {numZ:d}, {numPhi:d}")
    transc = tc.Transcoding(hybridFileBox)
    hybridPars = transc.read(hybrid_path)
    hybridPars.pop('header')
    cyl = w7x.flt.CylindricalGrid()
    cyl.__dict__.update(hybridPars)
    return cyl


def write_run_command(vmec_id, vmec_calc_url,
                      net_cdf_file="wout.nc",
                      script_file='runExtender.cmd', n_proc=32,
                      control_file='grid_in', coil_file='w7x.coil'):
    script = (
        "#!/bin/tcsh\n"
        "\n"
        "# Stdout and Sterr redirection\n"
        "#SBATCH -o ./{vmec_id}.o%j\n"
        "#SBATCH -e ./{vmec_id}.e%j\n"
        "# Initial working directory:\n"
        "#SBATCH -D ./\n"
        "# Job Name :\n"
        "#SBATCH -J extender-{vmec_id}\n"
        "# Queue (Partition):\n"
        "#SBATCH --partition=express\n"
        "# Number of nodes and MPI tasks per node:\n"
        "#SBATCH --nodes=20\n"
        "#SBATCH --ntasks-per-node=64\n"
        "# Enable Hyperthreading:\n"
        "#SBATCH --ntasks-per-core=2\n"
        "#\n"
        "# Request 504 GB of main Memory per node in Units of MB:\n"
        "#-- BATCH --mem=512000\n"
        "#\n"
        "#SBATCH --mail-type=all\n"
        "#SBATCH --mail-user=dboe@rzg.mpg.de\n"
        "# Wall clock limit:\n"
        "#SBATCH --time=0:30:00\n"
        "\n"
        "echo Start...\n"
        "\n"
        "# Run the program:\n"
        "\n"
        "set BIN=/u/mad/bin/EXTENDER_P\n"
        "\n"
        "set RUN={vmec_id}\n"
        "\n"
        "echo Downloading netcdf.\n"
        "wget {netCDFurl}\n"
        "\n"
        "set NETCDFfile={net_cdf_file}\n"
        "set COILfile={coil_file}\n"
        "set OUTPUT_SUFFIX=$RUN\n"
        "set CONTROL_FILE={control_file}\n"
        "\n"
        "set OTHER_OPTS=-full\n"
        "\n"
        "\n"
        "module load intel/16.0\n"
        "module load mkl/11.3\n"
        "module load impi/5.1.3\n"
        "module load ddt/7.0\n"
        "module load nag_flib/intel/mk24\n"
        "module load hdf5-serial/1.8.18\n"
        "module load netcdf-serial/4.4.1.1\n"
        "module load anaconda/2/4.4.0\n"
        "\n"
        "echo Starting extender with srun\n"
        "date\n"
        "\n"
        "echo srun $BIN -vmec_nyquist $NETCDFfile -i $CONTROL_FILE -c $COILfile"
        " -s $OUTPUT_SUFFIX $OTHER_OPTS\n"
        "srun $BIN -vmec_nyquist $NETCDFfile -i $CONTROL_FILE -c $COILfile -s"
        " $OUTPUT_SUFFIX $OTHER_OPTS > extender.out\n"
        "\n"
        "date\n"
        "\n"
        "echo Extender finished\n"
        "\n"
        "echo 'Starting conversion to ASCOT format'\n"
        "../to_ascoth5_onHGW.py > toascot.out\n"
        "\n"
        "echo 'Starting conversion to field line tracer format'\n"
        "../ascoth5_to_fieldLineTracer.py $RUN.h5 $RUN.dat > totracer.out\n"
        "\n"
        .format(jobname=vmec_id,
                netCDFurl=vmec_calc_url,
                **locals()))

    with open(script_file, 'w') as f:
        f.write(script)


def extend(args):
    import subprocess

    log = logging.getLogger()
    vmec_calc_url = ("http://svvmec1.ipp-hgw.mpg.de:8080/vmecrest/v1/run/" +
                     args.vmec_id + "/wout.nc")
    coil_file = args.vmec_id + '.coil'
    net_cdf_file = "wout.nc"
    script_file = 'runExtender.cmd'

    work_dir = ("~/Ascott/{args.vmec_id}"
                .format(**locals()))
    tfields.lib.in_out.mkdir(work_dir, isDir=True)
    with tfields.lib.in_out.cd_tmp(work_dir):
        log.info("Writing '{script_file}'.".format(**locals()))
        write_run_command(args.vmec_id,
                          vmec_calc_url,
                          net_cdf_file=net_cdf_file,
                          script_file=script_file,
                          n_proc=32,
                          coil_file=coil_file)
        
        log.info("Making the coil file '{coil_file}'.".format(**locals()))
        writeCoils(args.vmec_id, coil_file=coil_file)

        log.info('Making the grid/control file.')
        write_control_file()

        log.info('Making the parameter file')
        write_parameters_file(vmec_calc_url, args.vmec_id,
                              net_cdf_file=net_cdf_file)
        
        log.info('Submitting it...')
        subprocess.call('sbatch ' + script_file, shell=True)
        log.info('  and now we wait...')
