"""
@package pmsco.igor
data exchange with wavemetrics igor pro.

this module provides functions for loading/saving pmsco data in igor pro.

@author Matthias Muntwiler

@copyright (c) 2019 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from pmsco.compat import open


def _escape_igor_string(s):
    s = s.replace('\\', '\\\\')
    s = s.replace('"', '\\"')
    return s


def namefix_double(name):
    """
    fix 1-character wave name by doubling

    replaces length-1 string by a doubled version.

    @param name: (str) proposed wave name

    @return: corrected name
    """
    return name*2 if len(name) == 1 else name


def namefix_etpais(name):
    """
    fix 1-character wave name according to ETPAIS scheme

    replaces 'e' by 'en' etc.

    @param name: (str) proposed wave name

    @return: corrected name
    """
    name_map = {'e': 'en', 't': 'th', 'p': 'ph', 'i': 'in', 'm': 'mo', 's': 'si'}
    try:
        return name_map[name]
    except KeyError:
        return name


class IgorExport(object):
    """
    class exports pmsco data to an Igor text (ITX) file.

    usage:
    1) create an object instance.
    2) set @ref data.
    3) set optional attributes: @ref prefix and @ref namefix.
    4) call @ref export.
    """

    def __init__(self):
        super(IgorExport, self).__init__()
        self.data = None
        self.prefix = ""
        self.namefix = namefix_double

    def set_data(self, data):
        """
        set the data array to export.

        this must (currently) be a one-dimensional structured array.
        the column names will become wave names.

        @param data: numpy.ndarray
        @return:
        """
        self.data = data

    def export(self, filename):
        """
        write to igor file.
        """
        with open(filename, 'w') as f:
            self._write_header(f)
            self._write_data(f)

    def _fix_name(self, name):
        """
        fix a wave name.

        this function first applies @ref namefix and @ref prefix to the proposed wave name.

        @param name: (str) proposed wave name

        @return: corrected name
        """
        if self.namefix is not None:
            name = self.namefix(name)
        return self.prefix + name

    def _write_header(self, f):
        """
        write the header of the igor text file

        @param f: open file or stream

        @return: None
        """
        f.write('IGOR' + '\n')
        f.write('X // pmsco data export\n')

    def _write_data(self, f):
        """
        write a data section to the igor text file.

        @param f: open file or stream

        @return: None
        """
        assert isinstance(self.data, np.ndarray)
        assert len(self.data.shape) == 1
        assert len(self.data.dtype.names[0]) >= 1

        arr = self.data
        shape = ",".join(map(str, arr.shape))
        names = (self._fix_name(name) for name in arr.dtype.names)
        names = ", ".join(names)

        f.write('Waves/O/D/N=({shape}) {names}\n'.format(shape=shape, names=names))
        f.write('BEGIN\n')
        np.savetxt(f, arr, fmt='%g')
        f.write('END\n')