# Python modules
import ctypes

# 3rd party modules
import numpy as np
import multiprocessing as mp
import xml.etree.cElementTree as ElementTree

# Our modules
import vespa.analysis.constants as constants
import vespa.analysis.mrs_user_prior as mrs_user_prior
import vespa.analysis.block_fit_voigt as block_fit_voigt

import vespa.analysis.algos.wavelet_filter as wavelet_filter

import vespa.common.mrs_prior             as mrs_prior
import vespa.common.util.ppm              as util_ppm
import vespa.common.util.xml_             as util_xml
import vespa.common.util.math_            as util_math
import vespa.common.util.generic_spectral as util_spectral

from vespa.common.constants import Deflate
from vespa.analysis.constants import FitLineshapeModel

from vespa.common.constants import DEGREES_TO_RADIANS as DTOR
from vespa.analysis.algos.constrained_levenberg_marquardt import constrained_levenberg_marquardt as clm

# Sad but true, we need these globals to be able to minimize the amount
# of data and fitting parameters that are transferred to the processes

data = None
chain = None

class ProcessChain(block_fit_voigt._Settings):
    """
    A subclass of block_fit_voigt._Settings that will let us serialize
    fitting information that can be sent to each process.  We add a dash
    of values from mrs_dataset.Dataset to be able to use the util_ppm
    methods.

    """
    # This is the version of this object's XML output format.
    XML_VERSION = "1.0.0"

    def __init__(self, attributes=None):

        block_fit_voigt._Settings.__init__(self)

        # Spectral data information
        self.frequency              = 124.0
        self.sw                     = 1000.0
        self.raw_dims               = []
        self.raw_hpp                = 1.0
        self.resppm                 = 4.7
        self.echopeak               = 0.0
        self.spectral_dims          = []
        self.spectral_hpp           = 1.0
        self.phase_1_pivot          = 2.01
        self.zero_fill_multiplier   = 1

        self.user_prior = mrs_user_prior.UserPrior()
        self.metinfo    = self.user_prior.metinfo

        # These are run-time attributes

        self.peakpts = []
        self.basis_mets = None
        self.init_results = None
        self.fit_results = None

        # Now if we are loading in data we can re-set a few attributes

        if attributes is not None:
            self.inflate(attributes)

            # Here are attributes that can be calculated if we are loading
            # a ProcessChain object back into this instance.

            if self.prior.names != []:
                self.prior_list = self.prior.names
                #FIXME-bjs need new code to account for GISO fit method here! if needed.
                self.prior.calculate_full_basis_set(self.prior_ppm_start,
                                                    self.prior_ppm_end,
                                                    self)
                basis_mets = []
                ppms       = []
                peaks      = []
                alist      = self.prior.basis_set_names
                for name in sorted(alist):
                    item = self.prior.basis_set[name]
                    basis_mets.append(item.fid.copy())
                    ppms += item.all_ppms
                    peaks.append(item.peak_ppm)

                self.peakpts        = self.ppm2pts(np.array(ppms))
                self.prior_peak_ppm = np.array(peaks)  # ppm is OK here
                self.basis_mets     = np.array(basis_mets)

                self.fit_results    = np.zeros((self.nmet,), 'float')
                self.init_results   = np.zeros((self.nmet,), 'float')

        self.current_lw     = 5.0
        self.data           = None
        self.limits         = None
        self.weight_array   = None
        self.weight_array   = None
        self.fit_baseline   = None
        self.fit_stats      = None

    @property
    def nmet(self):
        if self.prior is not None:
            return len(self.prior.names)
        else:
            return 0

    def deflate(self, flavor=Deflate.ETREE):
        if flavor == Deflate.ETREE:
            # Make my base class do its deflate work
            e = block_fit_voigt._Settings.deflate(self, flavor)

            # Alter the tag name & XML version info
            e.tag = "process_chain"
            e.set("version", self.XML_VERSION)

            # Add custom attributes here.

            # These attributes are all lists.
            for attribute in ("raw_dims",
                              "spectral_dims", ):
                for value in getattr(self, attribute):
                    util_xml.TextSubElement(e, attribute, value)

            # These atttributes are all scalars and map directly to
            # XML elements of the same name.
            for attribute in (  "frequency",
                                "sw",
                                "raw_hpp",
                                "resppm",
                                "echopeak",
                                "spectral_hpp",
                                "phase_1_pivot",
                                "zero_fill_multiplier",):
                util_xml.TextSubElement(e, attribute, getattr(self, attribute))

            return e

    def inflate(self, source):
        if hasattr(source, "makeelement"):
            # Quacks like an ElementTree.Element

            block_fit_voigt._Settings.inflate(self, source)

            # We inflate in attributes grouped by type since there's so
            # doggone many attrs on this class.

#             # Booleans
#             for attribute in ("xxx", "xxx", ):
#                 item = source.findtext(attribute)
#                 if item is not None:
#                     setattr(self, attribute, util_xml.BOOLEANS[item])
            # floats
            for attribute in ("frequency",
                              "sw",
                              "raw_hpp",
                              "resppm",
                              "echopeak",
                              "spectral_hpp",
                              "phase_1_pivot", ):
                item = source.findtext(attribute)
                if item is not None:
                    setattr(self, attribute, float(item))

            # ints
            for attribute in ("zero_fill_multiplier", ):
                item = source.findtext(attribute)
                if item is not None:
                    setattr(self, attribute, int(float(item)))

#             # No translation required for these text attrs
#             for attribute in ("xxx", "xxx", ):
#                 item = source.findtext(attribute)
#                 if item is not None:
#                     setattr(self, attribute, item)

            # lists
            self.raw_dims       = [int(val.text) for val in source.getiterator("raw_dims")]
            self.spectral_dims  = [int(val.text) for val in source.getiterator("spectral_dims")]

    def fit_function(self, a,   pderflg=True,
                                nobase=False,
                                indiv=False,
                                finalwflg=False):

        # Setup constants and flags

#         import os
#         print "fit_function: process id = {0}  got here 100 data={1})".format(os.getpid(), self.data.shape )

        nmet    = self.nmet
        npts    = self.data.shape[0]
        zfmult  = self.zero_fill_multiplier
        nptszf  = round(npts * zfmult)
        sw      = 1.0 * self.sw
        td      = 1.0/sw
        piv     = self.ppm2pts(self.phase_1_pivot, acq=True)
        t2fix   = self.prior_fix_t2

        arr1    = np.zeros(npts,float) + 1.0
        f       = np.zeros((nmet,nptszf),complex)
        mf      = np.zeros((nptszf,),complex)

        t = (np.arange(nmet * npts) % npts) * td
        t.shape = nmet, npts
        mt = np.arange(npts) * td

        # get prior max peak ppm vals for metabs which are flagged ON
        peaks   = self.prior_peak_ppm

        # setup Lineshape
        if self.lineshape_model != FitLineshapeModel.GAUSS:
            # voigt and lorentzian models
            expo     = t/a[nmet*2] + (t/a[nmet*2+1])**2
            lshape   = util_math.safe_exp(-expo)
        else:
            # Note. in the case of the Gaussian lineshape, we now allow the user to
            # set a fixed T2 value for each metabolite. In the model call (fitt_funct.pro)
            # we now create a lineshape array that takes each fixed value into account.
            # BUT! we are still passing in a Tb parameter here, and though that parameter
            # will be tightly constrained, it should still be in the range of the fixed
            # physiologic params choosen by the user. At the moment, we have choosen to
            # just set this parameter to 0.250 sec, which should be a reasonable average
            # of 1H metabolite T2 values in the brain. It is then allowed to bop plus or
            # minus 0.001 as set further below.  In the fitting function, we adjust each
            # fixed value by the delta from 0.250 so that the pder will fluctuate as the
            # parameter changes and not confuse the poor optimization routine. In reality,
            # the 0.250 is never used, just the delta amount of the actual a[nmet*2]
            # parameter from that value.

            ma     = (self.fix_t2_center - a[nmet*2]) + t2fix    # delta for Ta param that is set at 0.25 sec
            ma     =  t/np.outer(ma, arr1)
            mb     = (t / a[nmet*2+1])**2
            expo   = ma+mb
            lshape = util_math.safe_exp(-expo)

        if finalwflg:
            finalw = lshape[:,0]
            finalw = util_spectral.full_width_half_max(np.fft.fft(util_spectral.chop(finalw))/len(finalw)) * self.spectral_hpp
            return finalw

        # if FID, then for correct area, first point must be divided by 2

        tmp  = self.basis_mets.copy()
        fre  = a[nmet:nmet*2] - self.ppm2hz(peaks)*2.0*np.pi    # in Radians here
        fre  = np.exp( 1j * (np.outer(fre, arr1)) * t ) # outer is matrix multiplication
        amp  = np.outer(a[0:nmet], arr1)
        ph0  = np.outer(np.exp(1j * (np.zeros(nmet) + a[nmet*2+2])), arr1)
        tmp *= amp * fre * ph0 * lshape
        f[:,0:npts] = tmp

        f[:,0] = f[:,0] / 2.0

        # Calc Phase1
        phase1 = np.exp(1j * (a[nmet*2+3]*DTOR*(np.arange(nptszf,dtype=float)-piv)/nptszf))

        # Calculate Partial Derivatives
        pder = None
        if pderflg:
            pder = np.zeros((len(a),nptszf), complex)

            pall = np.sum(f,axis=0)   # all lines added

            pind = f
            tt         = np.zeros(nptszf,float)
            tt[0:npts] = np.arange(npts,dtype=float) * td

            for i in range(nmet):   # Calc the Ampl and Freq pders
                pder[i,:]      = (np.fft.fft(pind[i,:] / a[i]   )/nptszf) * phase1
                pder[i+nmet,:] = (np.fft.fft(tt * 1j * pind[i,:])/nptszf) * phase1
            pder[nmet*2+0,:]  = (np.fft.fft(     tt     * pall/(a[nmet*2+0]**2))/nptszf) * phase1
            pder[nmet*2+1,:]  = (np.fft.fft(2.0*(tt**2) * pall/(a[nmet*2+1]**3))/nptszf) * phase1

            pder[nmet*2+2,:]  = (np.fft.fft(1j*pall)/nptszf) * phase1 * nptszf
            pder[nmet*2+3,:]  = (np.fft.fft(   pall)/nptszf) * (1j*DTOR*(np.arange(nptszf,dtype=float)-piv)/nptszf) * phase1

        # Do the FFT
        if indiv:   # return individual lines
            if nmet != 1:
                for i in range(nmet):
                    f[i,:] = (np.fft.fft(f[i,:])/nptszf) * phase1
            else:
                f = (np.fft.fft(f[0,:])/nptszf) * phase1
        else:  # return summed spectrum
            if (nmet) != 1:
                f = np.sum(f,axis=0)
                f = (np.fft.fft(f)/nptszf) * phase1
            else:
                f = (np.fft.fft(f[0,:])/nptszf) * phase1

        # Add in baseline unless nobase is True ---
        if not nobase:
            if f.ndim > 1:
                for i in range(len(f)): f[i,:] = f[i,:] + self.fit_baseline
            else:
                f = f + self.fit_baseline

        return f, pder

    def set_weight_array(self, lwidth=None, wtmult=None, wtmax=None):
        """
        Creates a weight array to be used in the optimization based on setting in
        the chain structure.  *(self.wtarr)) is set as the output.

        chain: ptr to optimization control structure

        """
        prior     = self.prior
        metinfo   = self.metinfo

        abbr = [metinfo.get_abbreviation(item) for item in self.prior_list]
        dim0 = self.spectral_dims[0]

        if not lwidth:
            lwidth = self.initial_linewidth_value
        else:
            lwidth = float(lwidth) if lwidth > 0.1 else 0.1

        if not wtmult:
            wtmult = self.optimize_weights_width_factor
        else:
            wtmult = float(wtmult) if wtmult>0.0001 else 0.001

        if not wtmax:
            wtmax  = dim0-1
        else:
            wtmax  = float(wtmax)

        wtarr = np.zeros(dim0, float)

        if self.optimize_weights_method == constants.FitOptimizeWeightsMethod.EVEN_WEIGHTING:
            wtarr = wtarr + 1.0

        elif self.optimize_weights_method == constants.FitOptimizeWeightsMethod.LOCAL_WEIGHTING:

            lw = lwidth / self.spectral_hpp   # in points
            centers = self.peakpts

            wid = lw * wtmult
            wid = wid if lw<wtmax else wtmax

            for ctr in self.peakpts:
                cs  = int(np.where(round(ctr-wid)>0, round(ctr-wid), 0))
                cs  = int(np.where(cs<dim0, cs, dim0))
                ce  = int(np.where(round(ctr+wid)>0, round(ctr+wid), 0))
                ce  = int(np.where(ce<dim0, ce, dim0))
                wtarr[cs:ce] = 1.0

            # set small pk weight scale higher if needed len(chain.peakpts)
            if self.optimize_weights_small_peak_factor != 1.0:

                ws = np.clip(int(np.round(self.ppm2pts(14.0))),0,dim0)
                we = np.clip(int(np.round(self.ppm2pts(1.25))),0,dim0)
                wtarr[ws:we] = wtarr[ws:we] * self.optimize_weights_small_peak_factor

                if 'lac' in abbr:
                    ws = np.clip(int(np.round(self.ppm2pts(1.45))),0,dim0)
                    we = np.clip(int(np.round(self.ppm2pts(1.25))),0,dim0)
                    wtarr[ws:we] = 1.0

                if 'naa' in abbr:
                    ws = np.clip(int(np.round(self.ppm2pts(2.12))),0,dim0)
                    we = np.clip(int(np.round(self.ppm2pts(1.85))),0,dim0)
                    wtarr[ws:we] = 1.0

                if 'cr' in abbr or 'cho' in abbr:
                    ws = np.clip(int(np.round(self.ppm2pts(3.30))),0,dim0)
                    we = np.clip(int(np.round(self.ppm2pts(2.85))),0,dim0)
                    wtarr[ws:we] = 1.0

            # Set and filter the weights
            indx0 = np.where(wtarr == 0.0)[0]
            if np.size(indx0) != 0:
                wtarr[indx0] = 1.0 / self.optimize_weights_scale_factor

            # set pks in water suppression low
            if self.optimize_weights_water_flag:
                ws = np.clip(int(np.round(self.ppm2pts(self.optimize_weights_water_end))),0,dim0)
                we = np.clip(int(np.round(self.ppm2pts(self.optimize_weights_water_start))),0,dim0)
                wtarr[ws:we] = 1.0 / self.optimize_weights_scale_factor

            # set pks in lipid area low
            if self.optimize_weights_lipid_flag == 1:
                ws = np.clip(int(np.round(self.ppm2pts(self.optimize_weights_lipid_end))),0,dim0)
                we = np.clip(int(np.round(self.ppm2pts(self.optimize_weights_lipid_start))),0,dim0)
                wtarr[ws:we] = 1.0 / self.optimize_weights_scale_factor

            wtarr = wtarr / max(wtarr)

        return wtarr

    def set_initial_values(self, indx):

        # indx will be used in future

        inival = [  1.7,6.0,10.0, 3.21,3.01,2.02, 0.06, 0.06,   0.1,  0.1]
        limits = [[ 1.4,4.8, 8.0, 3.31,2.85,1.80, 0.03, 0.06, -20.0,-50.1],
                  [ 1.9,7.2,12.0, 3.14,3.10,2.22, 100.0,100.0, 20.0,50.1]]

        self.current_lw    = 5.0
        self.init_results  = np.array(inival)
        self.fit_results   = np.array(inival)

        self.fit_results[0:3] *= 0.8
        self.fit_results[4:6] += 0.015
        self.fit_results[9]    = 12.0

        self.limits = np.array(limits)

def _create_default_prior():
    # This creates & returns a Prior object populated with some default
    # metabs. It's exists just so that something shows up when users first
    # open this tab.
    metabolites = { "n-acetylaspartate" : { "spins" : 3,
                                            "ppm"   : 2.01,
                                            "area"  : 3.0,
                                            "phase" : 0.0,
                                          },
                    "creatine"          : { "spins" : 3,
                                            "ppm"   : 3.01,
                                            "area"  : 3.0,
                                            "phase" : 0.0,
                                          },
                    "choline"          : { "spins"  : 9,
                                            "ppm"   : 3.21,
                                            "area"  : 9.0,
                                            "phase" : 0.0,
                                          },
                  }

    deflated_metabolites = [ ]
    for name, metabolite in metabolites.items():
        d = { "name" : name,
              "spins" : metabolite["spins"],
              "dims" : [0, 0, 0],
              "group" : [0],
              "ppms" : [metabolite["ppm"]],
              "areas" : [metabolite["area"]],
              "phases" : [metabolite["phase"]],
            }
        deflated_metabolites.append(d)

    d = { "source" : "default",
          "source_id" : "default",
          "comment" : "This is a typical 1H singlet prior basis set.",
          "nucleus" : "1H",
          "seqte"  : 0.07,
          "prior_metabolites" : deflated_metabolites,
        }

    return mrs_prior.Prior(d)

def init_chain(npts=1024, sw=1024.0):
    """
    This is only called in main process on startup to enable us to create
    fake data and a string version of the deflated XML object that can be
    sent to each process and resuscitated into a local copy for fitting data

    """

    # this is local
    chain = ProcessChain()

    chain.prior                     = _create_default_prior()
    chain.prior_list                = chain.prior.names

    chain.frequency                 = 124.0
    chain.sw                        = sw
    chain.raw_dims                  = np.array([npts,])
    chain.raw_hpp                   = chain.sw / float(npts)
    chain.resppm                    = 4.7
    chain.echopeak                  = 0.0
    chain.spectral_dims             = np.array([npts,])
    chain.spectral_hpp              = chain.sw / float(npts)
    chain.phase_1_pivot             = 2.01
    chain.zero_fill_multiplier      = 1

    chain.prior_fix_t2                      = [200.0 for i in range(chain.nmet)]
    chain.initial_linewidth_value           = 6.0
    chain.baseline_wavelet_scale            = 16.0
    chain.baseline_wavelet_min_dyad         = 4
    chain.lineshape_model                   = FitLineshapeModel.VOIGT
    chain.optimize_max_iterations           = 100
    chain.optimize_stop_tolerance           = 0.005
    chain.optimize_global_iterations        = 6
    chain.optimize_weights_width_factor     = 4.0
    chain.optimize_weights_scale_factor     = 3.0
    chain.optimize_weights_method           = constants.FitOptimizeWeightsMethod.LOCAL_WEIGHTING
    chain.optimize_weights_water_flag       = True
    chain.optimize_weights_water_end        = 4.2
    chain.optimize_weights_water_start      = 5.2
    chain.optimize_weights_lipid_flag       = False
    chain.optimize_weights_lipid_end        = 1.7
    chain.optimize_weights_lipid_start      = 0.9
    chain.optimize_weights_small_peak_factor= 1.0

    #FIXME-bjs need new code to account for GISO fit method here! if needed.
    chain.prior.calculate_full_basis_set(chain.prior_ppm_start, chain.prior_ppm_end, chain)

    basis_mets = []
    ppms       = []
    peaks      = []
    alist      = chain.prior.basis_set_names
    for name in sorted(alist):
        item = chain.prior.basis_set[name]
        basis_mets.append(item.fid.copy())
        ppms += item.all_ppms
        peaks.append(item.peak_ppm)

    chain.peakpts        = chain.ppm2pts(np.array(ppms))
    chain.prior_peak_ppm = np.array(peaks)  # ppm is OK here
    chain.basis_mets     = np.array(basis_mets)

    strout = chain.deflate()
    util_xml.indent(strout)
    strout = ElementTree.tostring(strout, "utf-8")

    return strout

def init_model():
    """ Create a single ideal spectrum as basis for creating fake data """

    global chain
    chain.set_initial_values(0)   # need this input parameters here

    # Create fake data model
    a = chain.init_results.copy()
    nmet = chain.nmet
    a[nmet:nmet*2] = chain.ppm2hz(a[nmet:nmet*2], acq=True)*2.0*np.pi
    a[nmet*2+2]    = a[nmet*2+2] * np.pi / 180.0
    if a[nmet*2]   == 0.0: a[nmet*2]   = 0.000001
    if a[nmet*2+1] == 0.0: a[nmet*2+1] = 0.000001

    chain.data = np.zeros((chain.spectral_dims[0],),'complex')
    model, _ = chain.fit_function(a, pderflg=False, nobase=True)

    return model

def do_baseline():

    global chain

#     import os
#     print "do_baseline: process id = {0}  got here 10 )".format(os.getpid(), )

    data = chain.data.copy()

    a = chain.fit_results.copy()
    model, _ = chain.fit_function(a, pderflg=False, nobase=True)

    # Subtract metabolite model from data
    basr = data.real - model.real
    basi = data.imag - model.imag

    # Estimate the baseline from the residual spectral calculated above.
    thresh  = chain.current_lw / chain.spectral_hpp
    scale   = int(chain.baseline_wavelet_scale)
    dyadmin = int(chain.baseline_wavelet_min_dyad)

    baser = wavelet_filter.wavelet_filter(basr, thresh, scale, dyadmin=dyadmin)
    basei = wavelet_filter.wavelet_filter(basi, thresh, scale, dyadmin=dyadmin)

    base = baser + 1j * basei  # Make complex array from real and imaginary parts

    chain.fit_baseline = base

def do_model():

    global chain

    data  = chain.data.copy()
    nmet  = chain.nmet
    a     = chain.fit_results.copy()
    ww    = chain.weight_array
    lim   = chain.limits.copy()
    itmax = chain.optimize_max_iterations
    toler = chain.optimize_stop_tolerance
    funct = chain.fit_function

    # optimize the model
    yfit, a, sig, chis, wchis, badfit = clm(data, ww, a, lim, funct, itmax, toler)

    # copy results into chain object so the do_baseline method can get them
    chain.fit_results = a.copy()
    chain.fit_stats   = np.array([chis, wchis, badfit])
    chain.fit_plot    = yfit

    chain.fitted_lw, _ = util_spectral.voigt_width(a[nmet*2], a[nmet*2+1], chain)

def check_out(dest, chain):
    """ helper function - converts PPMs to Hz and Phase degrees to radians """
    nmet = chain.nmet
    dest[nmet:nmet*2] = chain.ppm2hz(dest[nmet:nmet*2], acq=True)*2.0*np.pi
    dest[nmet*2+2]    = dest[nmet*2+2] * np.pi / 180.0
    if dest[nmet*2]   == 0.0: dest[nmet*2]   = 0.000001
    if dest[nmet*2+1] == 0.0: dest[nmet*2+1] = 0.000001
    return dest

def check_in(dest, chain):
    """ helper function - converts Hz to PPMs and Phase radians to degrees """
    nmet = chain.nmet
    dest[nmet:nmet*2] = chain.hz2ppm(dest[nmet:nmet*2]/(2.0*np.pi), acq=True)
    dest[nmet*2+2]    = dest[nmet*2+2] * 180.0 / np.pi
    return dest

def tonumpyarray(mp_arr):
    """ helper function to re-cast the shared array to a numpy array """
    return np.frombuffer(mp_arr.get_obj())

def do_loop_kernel(indx):
    """
    This is where each spectrum is fitted over N iterations of baseline
    and spectral model estimation and optimization respectively.

    We return the optimized spectral model parameters and index of the voxel
    so we can sort the results list in the main process after they have all
    been fitted.

    """
    global data
    global chain

#     import os
#     print "do_loop_kernel: proc_id={0} indx={1})".format(os.getpid(), indx)

    # load FID from shared array into chain object
    _data  = tonumpyarray(data)
    _datac = _data.view('complex64')
    str = chain.spectral_dims[0] * indx # compute offset since we flattened it
    end = str + chain.spectral_dims[0]
    chain.data = _datac[str:end]

    # set initial values and limits here - hard set for this example
    chain.set_initial_values(indx)
    chain.init_results = check_out(chain.init_results, chain)
    chain.limits[0,:]  = check_out(chain.limits[0,:].copy(), chain)
    chain.limits[1,:]  = check_out(chain.limits[1,:].copy(), chain)

    # calculate weight array ... likely should be in baseline/model loop
    chain.weight_array = chain.set_weight_array()
    chain.fit_results = chain.init_results.copy()
    chain.fit_results[0:chain.nmet] *= 0.8

    for k in range(chain.optimize_global_iterations):
        do_baseline()
        do_model()

    result = check_in(chain.fit_results, chain)

    return (indx, result)

def do_process_init(xmlstring, shared_arr):

    import os
#     narr = 0 if shared_arr is None else len(shared_arr)
#     print "do_process_init: process id = {0}  xml_len({1}) data_len({2}))".format(os.getpid(),len(xmlstring),narr)

    global chain
    tree = ElementTree.ElementTree(ElementTree.fromstring(xmlstring))
    attributes = tree.getroot()
    chain = ProcessChain(attributes)

    if shared_arr is not None:
        global data
        data = shared_arr

def do_loop(nvox, xmlstring, shared_arr):

    global data
    global chain

    inputs = list(range(nvox))

    pool_size = 10 if mp.cpu_count() < 10 else mp.cpu_count()

    if pool_size > nvox:
        pool_size = nvox

    pool = mp.Pool(processes=pool_size,
                   initializer=do_process_init,
                   initargs=(xmlstring,shared_arr) )

    pool_outputs = pool.map(do_loop_kernel, inputs)
    pool.close() # no more tasks
    pool.join()  # wrap up current tasks

    # Get process results from the output list
    pool_outputs.sort()
    results = np.array([r[1] for r in pool_outputs])

    return results

def run_tests():
    """
    Ran some tests on my T3500 with 6 cores and HT = 12 cores

    There's a major hit at the startup of the processes, on the order of 24
    seconds or so ...

    Here's the breakdown:
     1         3.51 sec
     2         5.40 sec
     10       23.44 sec
     50       29.81 sec
     100      31.84 sec
     800      64.41 sec

    """
#     from pylab import *

    nvox = 100

    # I test this convoluted way of initializing global chain variable here so
    # I know that it will perform well in the processes without having to printf
    # my way through a distributed debug session

    xmlstring  = init_chain()           # deflated object with values for fake data set properly
    do_process_init(xmlstring, None)    # inflate and set equal to the global chain variable

    global chain

    # Create fake data and a shared array that will be accessed by all the
    # distributed processes. Init model uses the inital values set in the
    # chain object. We then add different data to this model to create nvox
    # distinct sets of data.

    model = init_model()
    fake_data = np.zeros((nvox,chain.raw_dims[0]), 'complex64')

    ndbl = nvox * chain.spectral_dims[0] * 2    # for complex data
    shared_arr = mp.Array(ctypes.c_float, ndbl)
    arr = tonumpyarray(shared_arr).view('complex64')

    for i in range(nvox):
        noiser = np.random.randn(chain.raw_dims[0])
        noisei = np.random.randn(chain.raw_dims[0])
        noise  = noiser + 1j*noisei
        fake_data[i,:] = model + noise * 0.02

    tmp = _data.flatten()
    arr[:] = tmp

    # We are ready to start the processes, initialize them and then use them
    # to fit the fake data. The results list stores fitted parameters for
    # each data entry in the order they are in the shared_array

    results = do_loop(nvox, xmlstring, shared_arr)

    # Check the results to see if the first set of fitted parameters
    # and the visulization of the overlaid fit on the data is good.
    # Note. We do not actually have the fitted baseline results where
    # we can plot them in this example ... something for the future.

    print('fit - init =' + str(results[0,:] - chain.init_results))

    a = results[0,:].copy()
    a = check_out(a, chain)
    fit, _ = chain.fit_function(a, pderflg=False, nobase=True)

#     plot(_data[0,:])
#     plot(fit)
#     plot(chain.fit_baseline)
#     plot(chain.weight_array)
#     show()

    bob = 10
    bob = bob + 1

if __name__ == '__main__':

    import cProfile
    cProfile.run('run_tests()')

#    run_tests()


# This file is part of LayerModel_lib
#
#     A tool to compute the transmission behaviour of plane electromagnetic waves
#     through human tissue.
#
# Copyright (C) 2018 Jan-Christoph Brumm
#
# Licensed under MIT license.
#
"""
A small example showing how to generate a layer model for specific coordinates inside the chosen VoxelModel.
Additionally, the transfer function for S_21 and E-, and H-field are compared for surface to in-body communication
and vice versa

"""
import numpy as np
import matplotlib.pyplot as plt

from LayerModel_lib import VoxelModel, LayerModel, Coordinate
from config import phantom_path

VoxelModel.working_directory = phantom_path
all_models = VoxelModel.list_all_voxel_models()
for model in all_models:
    print(model)

# Load a virtual human model
vm = VoxelModel('AustinMan_v2.5_2x2x2')

start = Coordinate([231, 270, 1110])  # some point in the colon
end = Coordinate([330, 277, 1110])  # some point outside the body

# calculate the layer model from two of these points
lm = LayerModel(vm, start, end)

# show info about the model
lm.print_info()

# Calculate the transfer function for S21 (square root of transmitted power) from 0 to 10 GHz with 1024 samples,
# the default direction is 'start->end'
(transfer_function, frequency) = lm.S21(f_start=3.1e9, f_end=4.8e9, n_samples=1024)

# plot the magnitude
hf, ha = plt.subplots(nrows=2)
ha[0].plot(frequency / 1e9, 20 * np.log10(np.abs(transfer_function)))
ha[0].set_xlabel("Frequency in GHz")
ha[0].set_ylabel("Magnitude in dB")
ha[0].set_title("Transfer Function between the two coordinates")

ha[1].plot(frequency / 1e9, np.unwrap(np.angle(transfer_function), axis=0))
ha[1].set_xlabel("Frequency in GHz")
ha[1].set_ylabel("Phase in rad")
plt.tight_layout()
plt.show()


import numpy as np
import subprocess

from numpy import (
    zeros,
    pi,
    floor_divide,
    loadtxt,
    max as np_max,
    min as np_min,
    sign,
    angle as np_angle,
)
from os.path import basename, splitext
from SciDataTool import DataTime, VectorField, Data1D
from os.path import join

from ....Functions.Winding.gen_phase_list import gen_name

from ....Classes.Magnetics import Magnetics
from ....Functions.labels import (
    AIRGAP_LAB,
    ROTOR_LAB_S,
    short_label,
    decode_label,
    get_obj_from_label,
    LAM_LAB_S,
    STATOR_LAB_S,
    HOLEV_LAB_S,
    HOLEM_LAB_S,
    WIND_LAB_S,
    MAG_LAB,
    SHAFT_LAB,
    NO_LAM_LAB,
    SLID_LAB,
    BOT_LAB,
)
from ....Functions.Winding.find_wind_phase_color import get_phase_id
from .... import __version__
from ....Functions.get_path_binary import get_path_binary

from ....Classes.HoleM50 import HoleM50
from ....Classes.HoleM51 import HoleM51
from ....Classes.HoleM52 import HoleM52
from ....Classes.HoleM53 import HoleM53
from ....Classes.MachineSIPMSM import MachineSIPMSM
from ....Classes.MachineIPMSM import MachineIPMSM
from ....Methods import NotImplementedYetError

def solve_FEA(self, output, sym, angle, time, angle_rotor, Is, Ir):
    """
    Solve Elmer model to calculate airgap flux density, torque instantaneous/average/ripple values,
    flux induced in stator windings and flux density, field and permeability maps

    Parameters
    ----------
    self: MagElmer
        A MagElmer object
    output: Output
        An Output object
    sym: int
        Spatial symmetry factor
    time: ndarray
        Time vector for calculation
    angle: ndarray
        Angle vector for calculation
    Is : ndarray
        Stator current matrix (qs,Nt) [A]
    Ir : ndarray
        Stator current matrix (qs,Nt) [A]
    angle_rotor: ndarray
        Rotor angular position vector (Nt,)
    """

    project_name = self.get_path_save_fea(output)
    elmermesh_folder = project_name
    mesh_names_file = join(project_name, "mesh.names")
    boundaries = {}
    bodies = {}
    machine = output.simu.machine
    BHs = output.geo.stator.BH_curve  # Stator B(H) curve
    BHr = output.geo.rotor.BH_curve  # Rotor B(H) curve
    # Is = output.elec.Is  # Stator currents waveforms
    # Ir = output.elec.Ir  # Rotor currents waveforms
    Speed = output.elec.OP.get_N0()
    rotor_mat_file = join(project_name, "rotor_material.pmf")
    stator_mat_file = join(project_name, "stator_material.pmf")

    # TO-DO: Time vector must be greater than one
    timesize_str = np.array2string(
        np.diff(time), separator=" ", formatter={"float_kind": lambda x: "%.2e" % x}
    )
    time = np.append(time, time[1] + time[-1])
    timesize_str = np.array2string(
        np.diff(time), separator=" ", formatter={"float_kind": lambda x: "%.2e" % x}
    )
    timelen = len(time) - 1
    ones_str = np.array2string(
        np.ones(timelen), separator=" ", formatter={"int": lambda x: "%d" % x}
    )
    timeinterval_str = ones_str.replace(".", "")

    with open(mesh_names_file, "rt") as f:
        for line in f:
            fields = line.strip().split()
            if fields[0] == "$":
                field_name = fields[1]
                field_value = fields[3]
                # update dictionary
                # _settings['Geometry'][field_name] = field_value
                if field_name.count("BOUNDARY"):
                    boundaries[field_name] = field_value
                else:
                    bodies[field_name] = {
                        "id": field_value,
                        "mat": 1,  # Air by Default
                        "eq": 1,  # RigidMeshMapper by Default
                        "bf": None,
                        "tg": None,
                    }

    with open(rotor_mat_file, "wt") as ro:
        ro.write("! File Generated by pyleecan v{0}\n".format(__version__))
        ro.write(
            "! Material Name: {0}\n"
            "! B-H Curve Rotor Material\n"
            "Electric Conductivity = 0\n"
            "H-B Curve = Variable coupled iter\n"
            " Real\t\tCubic Monotone\n".format(machine.rotor.mat_type.name)
        )
        for ii in range(BHr.shape[0]):
            ro.write("   {0}\t\t{1}\n".format(BHr[ii][1], BHr[ii][0]))
        ro.write("End\n")

    with open(stator_mat_file, "wt") as ro:
        ro.write("! File Generated by pyleecan v{0}\n".format(__version__))
        ro.write(
            "! Material Name: {0}\n"
            "! B-H Curve Stator Material\n"
            "Electric Conductivity = 0\n"
            "H-B Curve = Variable coupled iter\n"
            " Real\t\tCubic Monotone\n".format(machine.stator.mat_type.name)
        )
        for ii in range(BHs.shape[0]):
            ro.write("   {0}\t\t{1}\n".format(BHs[ii][1], BHs[ii][0]))
        ro.write("End\n")

    elmer_sim_file = join(project_name, "pyleecan_elmer.sif")
    pp = machine.stator.winding.p
    wind_mat = machine.stator.winding.get_connection_mat(machine.stator.slot.Zs)
    surf_wind = machine.stator.slot.comp_surface_active()
    ror = machine.rotor.comp_radius_mec()
    sir = machine.stator.comp_radius_mec()
    with open(elmer_sim_file, "wt") as fo:
        fo.write("! File Generated by pyleecan v{0}\n".format(__version__))
        fo.write(
            "$ WM = 2*pi*{0}/60        ! Mechanical Frequency [rad/s]\n".format(Speed)
        )
        fo.write("$ PP = {0}                ! Pole pairs\n".format(pp))
        fo.write("$ WE = PP*WM              ! Electrical Frequency [Hz]\n")

        if isinstance(machine, MachineSIPMSM):
            # magnet_0 = machine.rotor.slot.magnet[0]
            magnet_0 = machine.rotor.magnet
        elif isinstance(machine, MachineIPMSM):
            magnet_dict = machine.rotor.hole[0].get_magnet_dict()
            magnet_0 = magnet_dict["magnet_0"]
        else:
            self.get_logger().info("ElmerSolver [Error]: Unsupported Machine Geometry")
            return False

        surf_list = machine.build_geometry(sym=sym)
        pm_index = 6
        Mangle = list()
        Ncond_Aplus = 1
        Ncond_Aminus = 1
        Ncond_Bplus = 1
        Ncond_Bminus = 1
        Ncond_Cplus = 1
        Ncond_Cminus = 1
        Ncond_Dplus = 1
        Ncond_Dminus = 1
        Ncond_Eplus = 1
        Ncond_Eminus = 1
        Ncond_Fplus = 1
        Ncond_Fminus = 1
        Npcp = machine.stator.winding.Npcp
        for surf in surf_list:
            label = short_label(surf.label)
            label_dict = decode_label(label)
            point_ref = surf.point_ref
            if HOLEM_LAB_S in label_dict["surf_type"]:  # LamHole
                mag_obj = get_obj_from_label(machine, label_dict=label_dict)
                if mag_obj.type_magnetization == 1:  # Parallel
                    magnetization_type = "parallel"
                    # calculate pole angle and angle of pole middle
                    T_id = label_dict["T_id"]
                    hole = mag_obj.parent
                    Zh = hole.Zh
                    alpha_p = 360 / Zh
                    mag_0 = (
                        floor_divide(np_angle(point_ref, deg=True), alpha_p) + 0.5
                    ) * alpha_p

                    mag_dict = hole.comp_magnetization_dict()
                    mag = mag_0 + mag_dict["magnet_" + str(T_id)] * 180 / pi
                    # modifiy magnetisation of south poles
                    if (label_dict["S_id"] % 2) == 1:
                        mag = mag + 180
                else:
                    raise NotImplementedYetError(
                        "Only parallele magnetization are available for HoleMagnet"
                    )
                if bodies.get(label, None) is not None:
                    Mangle.append(mag)
                    bodies[label]["mat"] = pm_index
                    bodies[label]["eq"] = 1
                    bodies[label]["bf"] = 1
                    bodies[label]["tg"] = 1
                    pm_index = pm_index + 1
            elif MAG_LAB in label_dict["surf_type"]:
                mag_obj = get_obj_from_label(machine, label_dict=label_dict)
                if mag_obj.type_magnetization == 0 and (label_dict["S_id"] % 2) == 0:
                    mag = 0  # North pole magnet
                    magnetization_type = "radial"
                elif mag_obj.type_magnetization == 0:
                    mag = 180  # South pole magnet
                    magnetization_type = "radial"
                elif mag_obj.type_magnetization == 1 and (label_dict["S_id"] % 2) == 0:
                    mag = np_angle(point_ref) * 180 / pi  # North pole magnet
                    magnetization_type = "parallel"
                elif mag_obj.type_magnetization == 1:
                    mag = np_angle(point_ref) * 180 / pi + 180  # South pole magnet
                    magnetization_type = "parallel"
                elif mag_obj.type_magnetization == 2:
                    Zs = mag_obj.parent.slot.Zs
                    mag = str(-(Zs / 2 - 1)) + " * theta + 90 "
                    magnetization_type = "hallback"
                else:
                    continue
                if bodies.get(label, None) is not None:
                    Mangle.append(mag)
                    bodies[label]["mat"] = pm_index
                    bodies[label]["eq"] = 1
                    bodies[label]["bf"] = 1
                    bodies[label]["tg"] = 1
                    pm_index = pm_index + 1
            elif WIND_LAB_S in label_dict["surf_type"]:
                lam_obj = get_obj_from_label(machine, label_dict=label_dict)
                wind_mat = lam_obj.winding.get_connection_mat(lam_obj.get_Zs())
                Nrad_id = label_dict["R_id"]  # zone radial coordinate
                Ntan_id = label_dict["T_id"]  # zone tangential coordinate
                Zs_id = label_dict["S_id"]  # Zone slot number coordinate
                # Get the phase value in the correct slot zone
                q_id = get_phase_id(wind_mat, Nrad_id, Ntan_id, Zs_id)
                Ncond = wind_mat[Nrad_id, Ntan_id, Zs_id, q_id]
                s = sign(Ncond)
                if bodies.get(label, None) is not None:
                    bodies[label]["mat"] = 5
                    bodies[label]["eq"] = 1
                    if q_id == 0 and s == 1:
                        bodies[label]["bf"] = 2
                        Ncond_Aplus = abs(Ncond)
                    elif q_id == 0 and s == -1:
                        bodies[label]["bf"] = 3
                        Ncond_Aminus = abs(Ncond)
                    elif q_id == 1 and s == 1:
                        bodies[label]["bf"] = 4
                        Ncond_Bplus = abs(Ncond)
                    elif q_id == 1 and s == -1:
                        bodies[label]["bf"] = 5
                        Ncond_Bminus = abs(Ncond)
                    elif q_id == 2 and s == 1:
                        bodies[label]["bf"] = 6
                        Ncond_Cplus = abs(Ncond)
                    elif q_id == 2 and s == -1:
                        bodies[label]["bf"] = 7
                        Ncond_Cminus = abs(Ncond)
                    elif q_id == 3 and s == 1:
                        bodies[label]["bf"] = 8
                        Ncond_Dplus = abs(Ncond)
                    elif q_id == 3 and s == -1:
                        bodies[label]["bf"] = 9
                        Ncond_Dminus = abs(Ncond)
                    elif q_id == 4 and s == 1:
                        bodies[label]["bf"] = 10
                        Ncond_Eplus = abs(Ncond)
                    elif q_id == 4 and s == -1:
                        bodies[label]["bf"] = 11
                        Ncond_Eminus = abs(Ncond)
                    elif q_id == 5 and s == 1:
                        bodies[label]["bf"] = 12
                        Ncond_Fplus = abs(Ncond)
                    elif q_id == 5 and s == -1:
                        bodies[label]["bf"] = 13
                        Ncond_Fminus = abs(Ncond)
                    else:
                        pass
            elif (
                LAM_LAB_S in label_dict["surf_type"]
                and ROTOR_LAB_S in label_dict["lam_label"]
                and bodies.get(label, None) is not None
            ):
                bodies[label]["mat"] = 4
                bodies[label]["eq"] = 1
                bodies[label]["bf"] = 1
                bodies[label]["tg"] = 1
            elif (
                LAM_LAB_S in label_dict["surf_type"]
                and STATOR_LAB_S in label_dict["lam_label"]
                and bodies.get(label, None) is not None
            ):
                bodies[label]["mat"] = 3
                bodies[label]["eq"] = 1
            elif (
                SHAFT_LAB in label_dict["surf_type"]
                and bodies.get(label, None) is not None
            ):
                bodies[label]["mat"] = 1
                bodies[label]["eq"] = 1
                bodies[label]["bf"] = 1
                bodies[label]["tg"] = 1
            elif (
                HOLEV_LAB_S in label_dict["surf_type"]
                and bodies.get(label, None) is not None
            ):
                bodies[label]["mat"] = 1
                bodies[label]["eq"] = 1
                bodies[label]["bf"] = 1
                bodies[label]["tg"] = 1
            else:
                pass

        # The following bodies are not in the dictionary
        bodies[ROTOR_LAB_S + "-0_" + AIRGAP_LAB + BOT_LAB]["bf"] = 1
        bodies[NO_LAM_LAB + "_" + SLID_LAB + BOT_LAB]["bf"] = 1  # Sliding band bottom

        No_Magnets = pm_index - 6
        magnet_temp = 20.0  # Magnet Temperature Fixed for now
        Hcm20 = magnet_0.mat_type.mag.Hc
        Brm20 = magnet_0.mat_type.mag.Brm20
        kt = 0.01  # Br Temperature Coefficient fixed for now
        Br = Brm20 * (1 + kt * 0.01 * (magnet_temp - 20.0))
        magnet_permeability = magnet_0.mat_type.mag.mur_lin
        rho20_m = magnet_0.mat_type.elec.rho
        kt_m = 0.01  # Rho Temperature Coefficient fixed for now
        rho_m = rho20_m * (1 + kt_m * (magnet_temp - 20.0))
        conductivity_m = 0.0 * 1.0 / rho_m

        skip_steps = 1  # Fixed for now
        degrees_step = 1  # Fixed for now
        current_angle = 0 - pp * degrees_step * skip_steps
        angle_shift = self.angle_rotor_shift - self.angle_stator_shift
        rotor_init_pos = machine.comp_angle_rotor_initial() + angle_shift
        rotor_d_axis = machine.rotor.comp_angle_d_axis() * 180.0 / pi
        Ncond = 1  # Fixed for Now
        Cp = 1  # Fixed for Now
        qs = len(machine.stator.get_name_phase())

        fo.write(
            "$ H_PM = {0}              ! Magnetization [A/m]\n".format(round(Hcm20, 2))
        )
        fo.write("$ Shift = 2*pi/{0}        ! N-phase machine [rad]\n".format(qs))
        fo.write(
            "$ Gamma = {0}*pi/180      ! Current Angle [rad]\n".format(
                round(current_angle, 2)
            )
        )
        fo.write("$ Ncond = {0}             ! Conductors per coil\n".format(Ncond))
        fo.write("$ Cp = {0}                ! Parallel paths\n".format(Cp))
        fo.write("$ Is = {0}                ! Stator current [A]\n".format(0.0))
        fo.write("$ Aaxis = {0}             ! Axis Coil A [deg]\n".format(0.0))
        fo.write(
            "$ Carea = {0}             ! Coil Side Conductor Area [m2]\n".format(
                surf_wind
            )
        )

        for mm in range(1, No_Magnets + 1):
            fo.write(
                "$ Mangle{0} = {1}     ! Magnetization Angle [deg]\n".format(
                    mm, round(Mangle[mm - 1], 2)
                )
            )

        fo.write("$ Nsteps = {0}            !\n".format(2))
        fo.write("$ StepDegrees = {0}       !\n".format(degrees_step))
        fo.write("$ DegreesPerSec = WM*180.0/pi  !\n")
        fo.write("$ RotorInitPos = {}!\n".format(round(rotor_init_pos * 180.0 / pi, 2)))

        fo.write(
            "\nHeader\n"
            "\tCHECK KEYWORDS Warn\n"
            '\tMesh DB "{0}"\n'
            '\tInclude Path "."\n'
            '\tResults Directory "{1}"\n'
            "End\n".format(elmermesh_folder, elmermesh_folder)
        )

        fo.write("\nConstants\n" "\tPermittivity of Vacuum = 8.8542e-12\n" "End\n")

        fo.write(
            "\nSimulation\n"
            "\tMax Output Level = 4\n"
            "\tCoordinate System = Cartesian 2D\n"
            "\tCoordinate Scaling = {0}\n"
            "\tSimulation Type = Transient\n"
            "\tTimestepping Method = BDF\n"
            "\tBDF Order = 2\n"
            #                 "\tTimestep Sizes = $ (StepDegrees / DegreesPerSec)  ! sampling time\n"
            #                 "\tTimestep Intervals = $ Nsteps              ! steps\n"
            #                 "\tOutput Intervals = 1\n"
            "\tTimestep Sizes({1}) = {2}\n"
            "\tTimestep Intervals({1}) = {3}\n"
            "\tUse Mesh Names = Logical True\n"
            "End\n".format(1.0, timelen, timesize_str[1:-1], timeinterval_str[1:-1])
        )

        fo.write("\n!--- MATERIALS ---\n")
        fo.write(
            "Material 1\n"
            '\tName = "Air"\n'
            "\tRelative Permeability = 1\n"
            "\tElectric Conductivity = 0\n"
            "End\n"
        )

        fo.write(
            "\nMaterial 2\n"
            '\tName = "Insulation"\n'
            "\tRelative Permeability = 1\n"
            "\tElectric Conductivity = 0\n"
            "End\n"
        )

        fo.write(
            "\nMaterial 3\n"
            '\tName = "StatorMaterial"\n'
            '\tInclude "{0}"\n'
            "End\n".format(stator_mat_file)
        )

        fo.write(
            "\nMaterial 4\n"
            '\tName = "RotorMaterial"\n'
            '\tInclude "{0}"\n'
            "End\n".format(rotor_mat_file)
        )

        winding_temp = 20.0  # Fixed for Now
        rho20 = machine.stator.winding.conductor.cond_mat.elec.rho
        kt = 0.01  # Br Temperature Coefficient fixed for now
        rho = rho20 * (1 + kt * (winding_temp - 20.0))
        conductivity = 0.0 * 1.0 / rho

        fo.write(
            "\nMaterial 5\n"
            '\tName = "Copper"\n'
            "\tRelative Permeability = 1\n"
            "\tElectric Conductivity = {0}\n"
            "End\n".format(round(conductivity, 2))
        )

        magnets_per_pole = No_Magnets  # TO-DO: Assumes only one pole drawn
        for m in range(1, magnets_per_pole + 1):
            mat_number = 5 + m
            if magnetization_type == "parallel":
                fo.write(
                    "\nMaterial {0}\n"
                    '\tName = "PM_{1}"\n'
                    "\tRelative Permeability = {2}\n"
                    "\tMagnetization 1 = Variable time, timestep size\n"
                    '\t\tReal MATC  "H_PM*cos(WM*(tx(0)-tx(1)) + {3}*pi/PP + {3}*pi + (RotorInitPos + Mangle{1})*pi/180)"\n'
                    "\tMagnetization 2 = Variable time, timestep size\n"
                    '\t\tReal MATC "H_PM*sin(WM*(tx(0)-tx(1)) + {3}*pi/PP + {3}*pi + (RotorInitPos + Mangle{1})*pi/180)"\n'
                    "\tElectric Conductivity = {4}\n"
                    "End\n".format(
                        mat_number,
                        m,
                        magnet_permeability,
                        int((m - 1) / magnets_per_pole),
                        round(conductivity_m, 2),
                    )
                )
            elif magnetization_type == "radial":
                fo.write(
                    "\nMaterial {0}\n"
                    '\tName = "PM_{1}"\n'
                    "\tRelative Permeability = {2}\n"
                    "\tMagnetization 1 = Variable Coordinate\n"
                    '\t\tReal MATC  "H_PM*cos(atan2(tx(1),tx(0)) + {3}*pi + Mangle{1}*pi/180)"\n'
                    "\tMagnetization 2 = Variable Coordinate\n"
                    '\t\tReal MATC "H_PM*sin(atan2(tx(1),tx(0)) + {3}*pi + Mangle{1}*pi/180)"\n'
                    "\tElectric Conductivity = {4}\n"
                    "End\n".format(
                        mat_number,
                        m,
                        magnet_permeability,
                        m - 1,
                        round(conductivity_m, 2),
                    )
                )
            elif magnetization_type == "perpendicular":
                fo.write(
                    "\nMaterial {0}\n"
                    '\tName = "PM_{1}"\n'
                    "\tRelative Permeability = {2}\n"
                    "\tMagnetization 1 = Variable time, timestep size\n"
                    '\t\tReal MATC  "H_PM*cos(WM*(tx(0)-tx(1)) + {3}*pi/PP + {3}*pi + Aaxis*pi/180 + (Mangle{1}*pi/180))"\n'
                    "\tMagnetization 2 = Variable time, timestep size\n"
                    '\t\tReal MATC "H_PM*sin(WM*(tx(0)-tx(1)) + {3}*pi/PP + {3}*pi + Aaxis*pi/180 + (Mangle{1}*pi/180))"\n'
                    "\tElectric Conductivity = {4}\n"
                    "End\n".format(
                        mat_number,
                        m,
                        magnet_permeability,
                        int((m - 1) / magnets_per_pole),
                        round(conductivity_m, 2),
                    )
                )
            else:
                fo.write(
                    "\nMaterial {0}\n"
                    '\tName = "PM_{1}"\n'
                    "\tRelative Permeability = {2}\n"
                    "\tElectric Conductivity = {4}\n"
                    "End\n".format(
                        mat_number, m, magnet_permeability, round(conductivity_m, 2)
                    )
                )

        fo.write("\n!--- BODY FORCES ---\n")

        # fo.write("Body Force 1\n"
        #          "\tName = \"BodyForce_Rotation\"\n"
        #          "\t$omega = (180/pi)*WM\n"
        #          "\tMesh Rotate 3 = Variable time, timestep size\n"
        #          "\t\tReal MATC \"omega*(tx(0)-tx(1)) + RotorInitPos\"\n"
        #          "End\n")
        fo.write(
            "Body Force 1\n"
            '\tName = "BodyForce_Rotation"\n'
            "\tMesh Rotate 3 = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], angle_rotor[tt - 1] * 180.0 / pi
                )
            )
        fo.write("\tEnd\n" "End\n")

        # fo.write("Body Force 2\n"
        #          "\tName = \"J_A_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) + Gamma)\"\n"
        #          "End\n".format(Ncond_Aplus))

        # fo.write("Body Force 3\n"
        #          "\tName = \"J_A_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) + Gamma)\"\n"
        #          "End\n".format(Ncond_Aminus))

        # fo.write("Body Force 4\n"
        #          "\tName = \"J_B_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Bplus))
        #
        # fo.write("Body Force 5\n"
        #          "\tName = \"J_B_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Bminus))
        #
        # fo.write("Body Force 6\n"
        #          "\tName = \"J_C_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 2*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Cplus))
        #
        # fo.write("Body Force 7\n"
        #          "\tName = \"J_C_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 2*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Cminus))
        #
        # fo.write("Body Force 8\n"
        #          "\tName = \"J_D_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 3*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Dplus))
        #
        # fo.write("Body Force 9\n"
        #          "\tName = \"J_D_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 3*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Dminus))
        #
        # fo.write("Body Force 10\n"
        #          "\tName = \"J_E_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 4*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Eplus))
        #
        # fo.write("Body Force 11\n"
        #          "\tName = \"J_E_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 4*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Eminus))
        #
        # fo.write("Body Force 12\n"
        #          "\tName = \"J_F_PLUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 5*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Fplus))
        #
        # fo.write("Body Force 13\n"
        #          "\tName = \"J_F_MINUS\"\n"
        #          "\tCurrent Density = Variable time, timestep size\n"
        #          "\t\tReal MATC \"-(Is/Carea) * ({0}/Cp) * sin(WE * (tx(0)-tx(1)) - 5*Shift + Gamma)\"\n"
        #          "End\n".format(Ncond_Fminus))

        fo.write(
            "Body Force 2\n"
            '\tName = "J_A_PLUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], Ncond_Aplus * Is[0, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write(
            "Body Force 3\n"
            '\tName = "J_A_MINUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], -Ncond_Aminus * Is[0, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write(
            "Body Force 4\n"
            '\tName = "J_B_PLUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], Ncond_Bplus * Is[1, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write(
            "Body Force 5\n"
            '\tName = "J_B_MINUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], -Ncond_Bminus * Is[1, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write(
            "Body Force 6\n"
            '\tName = "J_C_PLUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], Ncond_Cplus * Is[2, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write(
            "Body Force 7\n"
            '\tName = "J_C_MINUS"\n'
            "\tCurrent Density = Variable time\n"
            "\t\tReal\n"
            "\t\t0.0\t\t0.0\n"
        )
        for tt in range(1, timelen + 1):
            fo.write(
                "\t\t{:.2e}\t\t{:.3f}\n".format(
                    time[tt], -Ncond_Cminus * Is[2, tt - 1] / surf_wind
                )
            )
        fo.write("\tEnd\n" "End\n")

        fo.write("\n!--- BODIES ---\n")
        for k, v in bodies.items():
            bid = bodies[k]["id"]
            beq = bodies[k]["eq"]
            bmat = bodies[k]["mat"]
            bf = bodies[k]["bf"]
            btg = bodies[k]["tg"]
            fo.write(
                "Body {0}\n"
                "\tName = {1}\n"
                "\tEquation = {2}\n"
                "\tMaterial = {3}\n".format(bid, k, beq, bmat)
            )
            if bf is not None:
                fo.write("\tBody Force = {0}\n".format(bf))
            if btg is not None:
                fo.write("\tTorque Groups = Integer {0}\n".format(btg))
            if k == "SB_INT":
                fo.write(
                    "\tR Inner = Real {0}\n" "\tR Outer = Real {1}\n".format(ror, sir)
                )
            fo.write("End\n\n")

        fo.write(
            "Equation 1\n"
            '\tName = "Model_Domain"\n'
            "\tActive Solvers(6) = 1 2 3 4 5 6\n"
            "End\n"
        )

        fo.write("\n!--- SOLVERS ---\n")
        fo.write(
            "Solver 1\n"
            "\tExec Solver = Before Timestep\n"
            "\tEquation = MeshDeform\n"
            '\tProcedure = "RigidMeshMapper" "RigidMeshMapper"\n'
            "End\n"
        )

        fo.write(
            "\nSolver 2\n"
            "\tEquation = MgDyn2D\n"
            '\tProcedure = "MagnetoDynamics2D" "MagnetoDynamics2D"\n'
            "\tExec Solver = Always\n"
            "\tVariable = A\n"
        )
        fo.write("\tNonlinear System Convergence Tolerance = {0}\n".format(1e-6))
        fo.write("\tNonlinear System Max Iterations = {0}\n".format(100))
        fo.write("\tNonlinear System Min Iterations = {0}\n".format(1))
        fo.write("\tNonlinear System Newton After Iterations = {0}\n".format(5))
        fo.write("\tNonlinear System Relaxation Factor = {0}\n".format(0.9))
        fo.write(
            "\tNonlinear System Convergence Without Constraints = {0}\n".format(
                "Logical True"
            )
        )
        fo.write("\tExport Lagrange Multiplier = {0}\n".format("Logical True"))
        fo.write("\tLinear System Abort Not Converged = {0}\n".format("Logical False"))
        fo.write("\tLinear System Solver = {0}\n".format("Direct"))
        fo.write("\tLinear System Direct Method = {0}\n".format("umfpack"))
        fo.write("\tOptimize Bandwidth = {0}\n".format("Logical True"))
        fo.write("\tLinear System Preconditioning =  {0}\n".format("ILU2"))
        fo.write("\tLinear System Max Iterations =  {0}\n".format(5000))
        fo.write("\tLinear System Residual Output =  {0}\n".format(20))
        fo.write("\tLinear System Convergence Tolerance =  {0}\n".format(1e-7))
        fo.write("\tMortar BCs Additive =  {0}\n".format("Logical True"))
        fo.write("End\n")

        fo.write(
            "\nSolver 3\n"
            "\tExec Solver = Always\n"
            "\tEquation = CalcFields\n"
            '\tPotential Variable = "A"\n'
            '\tProcedure = "MagnetoDynamics" "MagnetoDynamicsCalcFields"\n'
            "\tCalculate Nodal Forces = Logical True\n"
            "\tCalculate Magnetic Vector Potential = Logical True\n"
            "\tCalculate Winding Voltage = Logical True\n"
            "\tCalculate Current Density = Logical True\n"
            "\tCalculate Maxwell Stress = Logical True\n"
            "\tCalculate JxB = Logical True\n"
            "\tCalculate Magnetic Field Strength = Logical True\n"
            "End\n"
        )

        fo.write(
            "\nSolver 4\n"
            "\tExec Solver = After Timestep\n"
            '\tProcedure = "ResultOutputSolve" "ResultOutputSolver"\n'
            '\tOutput File Name = "{0}"\n'
            "\tVtu Format = True\n"
            "\tBinary Output = True\n"
            "\tSingle Precision = True\n"
            "\tSave Geometry Ids = True\n"
            "\tShow Variables = True\n"
            "End\n".format("step")
        )

        fo.write(
            "\nSolver 5\n"
            "\tExec Solver = After Timestep\n"
            "\tEquation = SaveLine\n"
            '\tFilename = "{0}"\n'
            '\tProcedure = "SaveData" "SaveLine"\n'
            "\tVariable 1 = Magnetic Flux Density 1\n"
            "\tVariable 2 = Magnetic Flux Density 2\n"
            "\tVariable 3 = Magnetic Flux Density 3\n"
            "\tVariable 4 = Magnetic Flux Density e 1\n"
            "\tVariable 5 = Magnetic Flux Density e 2\n"
            "\tVariable 6 = Magnetic Flux Density e 3\n"
            "End\n".format("lines.dat")
        )

        fo.write(
            "\nSolver 6\n"
            "\tExec Solver = After Timestep\n"
            '\tFilename = "{0}"\n'
            '\tProcedure = "SaveData" "SaveScalars"\n'
            "\tShow Norm Index = 1\n"
            "End\n".format("scalars.dat")
        )

        fo.write("\n!--- BOUNDARIES ---\n")
        for k, v in boundaries.items():
            if k == "VP0_BOUNDARY":
                fo.write(
                    "Boundary Condition {0}\n"
                    "\tName = {1}\n"
                    "\tA = Real 0\n"
                    "End\n\n".format(v, k)
                )
            elif k == "MASTER_STATOR_BOUNDARY":
                for k1, v1 in boundaries.items():
                    if k1 == "SLAVE_STATOR_BOUNDARY":
                        slave = v1
                        break
                if not self.is_periodicity_a:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tMortar BC Static = Logical True\n"
                        "\tRadial Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
                else:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tMortar BC Static = Logical True\n"
                        "\tAnti Radial Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
            elif k == "MASTER_ROTOR_BOUNDARY":
                for k1, v1 in boundaries.items():
                    if k1 == "SLAVE_ROTOR_BOUNDARY":
                        slave = v1
                        break
                if not self.is_periodicity_a:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tMortar BC Static = Logical True\n"
                        "\tRadial Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
                else:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tMortar BC Static = Logical True\n"
                        "\tAnti Radial Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
            elif k == "SB_STATOR_BOUNDARY":
                for k1, v1 in boundaries.items():
                    if k1 == "SB_ROTOR_BOUNDARY":
                        slave = v1
                        break
                if not self.is_periodicity_a:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tRotational Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
                else:
                    fo.write(
                        "Boundary Condition {0}\n"
                        "\tName = {1}\n"
                        "\tMortar BC = Integer {2}\n"
                        "\tAnti Rotational Projector = Logical True\n"
                        "\tGalerkin Projector = Logical True\n"
                        "End\n\n".format(v, k, slave)
                    )
            elif k == "AIRGAP_ARC_BOUNDARY":
                fo.write(
                    "Boundary Condition {0}\n"
                    "\tName = {1}\n"
                    "\tSave Line = True\n"
                    "End\n\n".format(v, k)
                )
            else:
                fo.write(
                    "Boundary Condition {0}\n" "\tName = {1}\n" "End\n\n".format(v, k)
                )

    # setup Elmer solver
    # ElmerSolver v8.4 must be installed and in the PATH

    elmer_settings = join(project_name, "pyleecan_elmer.sif")
    ElmerSolver_binary = get_path_binary("ElmerSolver")
    cmd_elmersolver = [
        ElmerSolver_binary,
        elmer_settings,
    ]
    self.get_logger().info(
        "Calling ElmerSolver: " + " ".join(map(str, cmd_elmersolver))
    )
    elmersolver = subprocess.Popen(
        cmd_elmersolver, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    (stdout, stderr) = elmersolver.communicate()
    elmersolver.wait()
    self.get_logger().info(stdout.decode("UTF-8"))
    if elmersolver.returncode != 0:
        self.get_logger().info("ElmerSolver [Error]: " + stderr.decode("UTF-8"))
        return False
    elmersolver.terminate()
    self.get_logger().info("ElmerSolver call complete!")

    self.get_meshsolution(output)

    Na = angle.size
    Nt = time.size - 1

    # Loading parameters for readibility
    L1 = output.simu.machine.stator.comp_length()
    save_path = self.get_path_save(output)

    scalars_file = join(elmermesh_folder, "scalars.dat")
    ecp, mfe, agt, iv, im, tq = loadtxt(
        scalars_file, unpack=True, usecols=(0, 1, 2, 3, 4, 5)
    )
    # ecp: eddy current power
    # mfe: magnetic field energy
    # agt: air gap torque
    # iv: inertial volume
    # im: inertial moment
    # tq: group 1 torque

    # TODO Load Air gap flux density

    # FEM_dict = output.mag.FEM_dict
    #
    if (
        hasattr(output.simu.machine.stator, "winding")
        and output.simu.machine.stator.winding is not None
    ):
        qs = output.simu.machine.stator.winding.qs  # Winding phase number
        Phi_wind_stator = zeros((Nt, qs))
    else:
        Phi_wind_stator = None

    # Initialize results matrix
    Br = zeros((Nt, Na))
    Bt = zeros((Nt, Na))
    Bz = zeros((Nt, Na))
    Tem = tq * sym * L1

    # Phi_wind_stator = zeros((Nt, qs))

    # compute the data for each time step
    # TODO Other than FEMM, in Elmer I think it's possible to compute
    #      all time steps at once
    self.get_logger().debug("Solving Simulation")

    # run the computation
    if self.nb_worker > 1:
        # TODO run solver in parallel
        pass
    else:
        # TODO run solver 'normal'
        pass

    # get the air gap flux result
    # TODO add function (or method)
    # ii -> Time, jj -> Angle
    # Br[ii, jj], Bt[ii, jj] = get_airgap_flux()

    # get the torque
    # TODO add function (or method)
    # Tem[ii] = comp_Elmer_torque(FEM_dict, sym=sym)

    # flux linkage computation
    # if Phi_wind_stator is not None:
    #     # TODO
    #     # Phi_wind[ii, :] = comp_Elmer_Phi_wind()
    #     pass

    return Br, Bt, Bz, Tem, Phi_wind_stator


"""
This is a module containing functions and classes for imaging propagation with HCIPy, for now LUVOIR A.
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from astropy.io import fits
import hcipy as hc
from hcipy.optics.segmented_mirror import SegmentedMirror

from config import CONFIG_INI

class SegmentedTelescopeAPLC:
    """ A segmented telescope with an APLC and actuated segments.

    Parameters:
    ----------
    aper : Field
        Telescope aperture.
    indexed_aperture : Field
        The *indexed* segmented aperture of the mirror, all pixels each segment being filled with its number for
        segment identification. Segment gaps must be strictly zero.
    seg_pos : CartesianGrid(UnstructuredCoords)
        Segment positions of the aperture.
    apod : Field
        Apodizer
    lyots : Field
        Lyot stop
    fpm : fpm
        Focal plane mask
    focal_grid :
        Focal plane grid to put final image on
    params : dict
        wavelength, diameter, image size in lambda/D, FPM radius
    """

    def __init__(self, aper, indexed_aperture, seg_pos, apod, lyotst, fpm, focal_grid, params):
        self.sm = SegmentedMirror(indexed_aperture=indexed_aperture, seg_pos=seg_pos)
        self.aper = aper
        self.apodizer = apod
        self.lyotstop = lyotst
        self.fpm = fpm   #TODO: this is not actually used inside this class
        self.wvln = params['wavelength']
        self.diam = params['diameter']
        self.imlamD = params['imlamD']
        self.fpm_rad = params['fpm_rad']
        self.lamDrad = self.wvln / self.diam
        self.coro = hc.LyotCoronagraph(indexed_aperture.grid, fpm, lyotst)
        self.prop = hc.FraunhoferPropagator(indexed_aperture.grid, focal_grid)
        self.coro_no_ls = hc.LyotCoronagraph(indexed_aperture.grid, fpm)
        self.wf_aper = hc.Wavefront(aper, wavelength=self.wvln)
        self.focal_det = focal_grid

    def calc_psf(self, ref=False, display_intermediate=False,  return_intermediate=None):
        """Calculate the PSF of the segmented telescope, normalized to contrast units.

        Parameters:
        ----------
        ref : bool
            Keyword for additionally returning the refrence PSF without the FPM.
        display_intermediate : bool
            Keyword for display of all planes.
        return_intermediate : string
            Either 'intensity', return the intensity in all planes; except phase on the SM (first plane)
            or 'efield', return the E-fields in all planes. Default none.
        Returns:
        --------
        wf_im_coro.intensity : Field
            Coronagraphic image, normalized to contrast units by max of reference image (even when ref
            not returned).
        wf_im_ref.intensity : Field, optional
            Reference image without FPM.
        intermediates : dict of Fields, optional
            Intermediate plane intensity images; except for full wavefront on segmented mirror.
        wf_im_coro : Wavefront
            Wavefront in last focal plane.
        wf_im_ref : Wavefront, optional
            Wavefront of reference image without FPM.
        intermediates : dict of Wavefronts, optional
            Intermediate plane E-fields; except intensity in focal plane after FPM.
        """

        # Create fake FPM for plotting
        fpm_plot = 1 - hc.circular_aperture(2 * self.fpm_rad * self.lamDrad)(self.focal_det)

        # Create apodozer as hc.Apodizer() object to be able to propagate through it
        apod_prop = hc.Apodizer(self.apodizer)

        # Calculate all wavefronts of the full propagation
        wf_sm = self.sm(self.wf_aper)
        wf_apod = apod_prop(wf_sm)
        wf_lyot = self.coro(wf_apod)
        wf_im_coro = self.prop(wf_lyot)

        # Wavefronts in extra planes
        wf_before_fpm = self.prop(wf_apod)
        int_after_fpm = np.log10(wf_before_fpm.intensity / wf_before_fpm.intensity.max()) * fpm_plot  # this is the intensity straight
        wf_before_lyot = self.coro_no_ls(wf_apod)

        # Wavefronts of the reference propagation
        wf_ref_pup = hc.Wavefront(self.aper * self.apodizer * self.lyotstop, wavelength=self.wvln)
        wf_im_ref = self.prop(wf_ref_pup)

        # Display intermediate planes
        if display_intermediate:

            plt.figure(figsize=(15, 15))

            plt.subplot(331)
            hc.imshow_field(wf_sm.phase, mask=self.aper, cmap='RdBu')
            plt.title('Seg aperture phase')

            plt.subplot(332)
            hc.imshow_field(wf_apod.intensity, cmap='inferno')
            plt.title('Apodizer')

            plt.subplot(333)
            hc.imshow_field(wf_before_fpm.intensity / wf_before_fpm.intensity.max(), norm=LogNorm(), cmap='inferno')
            plt.title('Before FPM')

            plt.subplot(334)
            hc.imshow_field(int_after_fpm / wf_before_fpm.intensity.max(), cmap='inferno')
            plt.title('After FPM')

            plt.subplot(335)
            hc.imshow_field(wf_before_lyot.intensity / wf_before_lyot.intensity.max(), norm=LogNorm(vmin=1e-3, vmax=1),
                            cmap='inferno')
            plt.title('Before Lyot stop')

            plt.subplot(336)
            hc.imshow_field(wf_lyot.intensity / wf_lyot.intensity.max(), norm=LogNorm(vmin=1e-3, vmax=1),
                            cmap='inferno', mask=self.lyotstop)
            plt.title('After Lyot stop')

            plt.subplot(337)
            hc.imshow_field(wf_im_coro.intensity / wf_im_ref.intensity.max(), norm=LogNorm(vmin=1e-10, vmax=1e-3),
                            cmap='inferno')
            plt.title('Final image')
            plt.colorbar()

        if return_intermediate == 'intensity':

            # Return the intensity in all planes; except phase on the SM (first plane)
            intermediates = {'seg_mirror': wf_sm.phase,
                             'apod': wf_apod.intensity,
                             'before_fpm': wf_before_fpm.intensity / wf_before_fpm.intensity.max(),
                             'after_fpm': int_after_fpm / wf_before_fpm.intensity.max(),
                             'before_lyot': wf_before_lyot.intensity / wf_before_lyot.intensity.max(),
                             'after_lyot': wf_lyot.intensity / wf_lyot.intensity.max()}

            if ref:
                return wf_im_coro.intensity, wf_im_ref.intensity, intermediates
            else:
                return wf_im_coro.intensity, intermediates

        if return_intermediate == 'efield':

            # Return the E-fields in all planes; except intensity in focal plane after FPM
            intermediates = {'seg_mirror': wf_sm,
                             'apod': wf_apod,
                             'before_fpm': wf_before_fpm,
                             'after_fpm': int_after_fpm,
                             'before_lyot': wf_before_lyot,
                             'after_lyot': wf_lyot}

            if ref:
                return wf_im_coro, wf_im_ref, intermediates
            else:
                return wf_im_coro, intermediates

        if ref:
            return wf_im_coro.intensity, wf_im_ref.intensity

        return wf_im_coro.intensity

    def flatten(self):
        self.sm.flatten()

    def set_segment(self, segid, piston, tip, tilt):
        self.sm.set_segment(segid, piston, tip, tilt)

    def apply_aberrations(self, aber_array):
        for vals in aber_array:
            self.sm.set_segment(vals[0], vals[1], vals[2], vals[3])

    def forward(self, wavefront):
        raise NotImplementedError()

    def backward(self, wavefront):
        raise NotImplementedError()

class LuvoirAPLC(SegmentedTelescopeAPLC):
    """ Simple E2E simulator for LUVOIR A (with APLC).

    Parameters:
    ----------
    input dir : string
        Path to input files: apodizer, aperture, indexed aperture, Lyot stop.
    apod_design : string
        Choice of apodizer design from May 2019 delivery. "small", "medium" or "large".
    """
    def __init__(self, input_dir, apod_design, samp):
        self.nseg = 120   # FIXME: this should not be hard-coded
        self.wvln = CONFIG_INI.getfloat('LUVOIR', 'lambda') * 1e-9    # m
        self.diam = 15.  # m   # FIXME: this should not be hard-coded
        self.sampling = samp
        self.lam_over_d = self.wvln / self.diam
        self.apod_dict = {'small': {'pxsize': 1000, 'fpm_rad': 3.5, 'fpm_px': 150, 'iwa': 3.4, 'owa': 12.,
                                    'fname': '0_LUVOIR_N1000_FPM350M0150_IWA0340_OWA01200_C10_BW10_Nlam5_LS_IDD0120_OD0982_no_ls_struts.fits'},
                          'medium': {'pxsize': 1000, 'fpm_rad': 6.82, 'fpm_px': 250, 'iwa': 6.72, 'owa': 23.72,
                                     'fname': '0_LUVOIR_N1000_FPM682M0250_IWA0672_OWA02372_C10_BW10_Nlam5_LS_IDD0120_OD0982_no_ls_struts.fits'},
                          'large': {'pxsize': 1000, 'fpm_rad': 13.38, 'fpm_px': 400, 'iwa': 13.28, 'owa': 46.88,
                                    'fname': '0_LUVOIR_N1000_FPM1338M0400_IWA1328_OWA04688_C10_BW10_Nlam5_LS_IDD0120_OD0982_no_ls_struts.fits'}}
        self.imlamD = 1.2*self.apod_dict[apod_design]['owa']

        # Pupil plane optics
        aper_path = 'inputs/TelAp_LUVOIR_gap_pad01_bw_ovsamp04_N1000.fits'
        aper_ind_path = 'inputs/TelAp_LUVOIR_gap_pad01_bw_ovsamp04_N1000_indexed.fits'
        apod_path = os.path.join(input_dir, 'luvoir_stdt_baseline_bw10', apod_design + '_fpm', 'solutions',
                                 self.apod_dict[apod_design]['fname'])
        ls_fname = 'inputs/LS_LUVOIR_ID0120_OD0982_no_struts_gy_ovsamp4_N1000.fits'

        pup_read = hc.read_fits(os.path.join(input_dir, aper_path))
        aper_ind_read = hc.read_fits(os.path.join(input_dir, aper_ind_path))
        apod_read = hc.read_fits(os.path.join(input_dir, apod_path))
        ls_read = hc.read_fits(os.path.join(input_dir, ls_fname))

        pupil_grid = hc.make_pupil_grid(dims=self.apod_dict[apod_design]['pxsize'], diameter=self.diam)

        self.aperture = hc.Field(pup_read.ravel(), pupil_grid)
        self.aper_ind = hc.Field(aper_ind_read.ravel(), pupil_grid)
        self.apod = hc.Field(apod_read.ravel(), pupil_grid)
        self.ls = hc.Field(ls_read.ravel(), pupil_grid)

        # Load segment positions from fits header
        hdr = fits.getheader(os.path.join(input_dir, aper_ind_path))

        poslist = []
        for i in range(self.nseg):
            segname = 'SEG' + str(i + 1)
            xin = hdr[segname + '_X']
            yin = hdr[segname + '_Y']
            poslist.append((xin, yin))

        poslist = np.transpose(np.array(poslist))
        self.seg_pos = hc.CartesianGrid(poslist)

        # Focal plane mask
        samp_foc = self.apod_dict[apod_design]['fpm_px'] / (self.apod_dict[apod_design]['fpm_rad'] * 2)
        focal_grid_fpm = hc.make_focal_grid(pupil_grid=pupil_grid, q=samp_foc,
                                            num_airy=self.apod_dict[apod_design]['fpm_rad'], wavelength=self.wvln)
        self.fpm = 1 - hc.circular_aperture(2*self.apod_dict[apod_design]['fpm_rad']*self.lam_over_d)(focal_grid_fpm)

        # Final focal plane grid (detector)
        self.focal_det = hc.make_focal_grid(pupil_grid=pupil_grid, q=self.sampling, num_airy=self.imlamD, wavelength=self.wvln)

        luvoir_params = {'wavelength': self.wvln, 'diameter': self.diam, 'imlamD': self.imlamD,
                         'fpm_rad': self.apod_dict[apod_design]['fpm_rad']}

        # Initialize the general segmented telescope with APLC class, includes the SM
        super().__init__(aper=self.aperture, indexed_aperture=self.aper_ind, seg_pos=self.seg_pos, apod=self.apod,
                         lyotst=self.ls, fpm=self.fpm, focal_grid=self.focal_det, params=luvoir_params)

        # Make dark hole mask
        dh_outer = hc.circular_aperture(2 * self.apod_dict[apod_design]['owa'] * self.lam_over_d)(
            self.focal_det)
        dh_inner = hc.circular_aperture(2 * self.apod_dict[apod_design]['iwa'] * self.lam_over_d)(
            self.focal_det)
        self.dh_mask = (dh_outer - dh_inner).astype('bool')

        # Propagators
        self.coro = hc.LyotCoronagraph(pupil_grid, self.fpm, self.ls)
        self.prop = hc.FraunhoferPropagator(pupil_grid, self.focal_det)
        self.coro_no_ls = hc.LyotCoronagraph(pupil_grid, self.fpm)
        #TODO: these three propagators should actually happen in the super init
        # -> how are self.aper_ind and pupil_grid connected?


""" Class for correlating galaxies and IGM within a given field on the sky
"""
from __future__ import print_function, absolute_import, division, unicode_literals

import numpy as np
import warnings
import pdb

import astropy
from astropy import units as u
from astropy.units import Quantity
from astropy.coordinates import SkyCoord
from astropy import constants as const
from astropy.cosmology import Planck15

import pyigm.field.utils as pfu
from pyigm.utils import calc_rho

#from xastropy.xutils import xdebug as xdb

class IgmGalaxyField(object):
    """ Class for a field associating galaxies to the IGM/CGM

    Parameters
    ----------
    radec : tuple or SkyCoord
        (RA,DEC) in deg or astropy.coordinate
    name : str; optional
        Default is set from sexagesimal coordinates
    zem : float; optional
        Redshift of background source (e.g., QSO)
    cosmo : Cosmology; optional
        Default is astropy.cosmology.Planck15

    """

    # Initialize
    def __init__(self, radec, name=None, zem=None, cosmo=None, verbose=False):
        # coord
        if isinstance(radec, (tuple)):
            self.coord = SkyCoord(ra=radec[0], dec=radec[1])
        elif isinstance(radec, SkyCoord):
            self.coord = radec

        # Field
        if name is None:
            self.name = 'IGMFIELD_J{:s}{:s}'.format(
                    self.coord.ra.to_string(unit=u.hour, pad=True),
                    self.coord.dec.to_string(pad=True, alwayssign=True))
        self.name = name

        # Cosmology
        if cosmo is None:
            cosmo = Planck15
            if verbose is True:
                print('IgmGalxyField: Using Planck15 cosmology')
        self.cosmo = cosmo

        # Init
        self.zem = zem
        self.igm = None
        self.targets = None
        self.galaxies = None
        self.observing = None
        self.selection = None

    def calc_rhoimpact(self, obj, comoving=True, los_coord=None):
        """Calculate impact parameter from LOS RA/DEC for a set of objects

        Parameters
        ----------
        obj : Table or dict
          Anything that takes 'Z', and 'RA','DEC' in degrees
          Sources for calculation
        comoving : bool, optional
           If True then comoving, else physical
        los_radec : SkyCoord, optional
          Defaults to field RA/DEC

        Returns
        -------
        rho : Quantity (array usually)
          Impact parameter(s) in kpc
        """
        if los_coord is None:
            los_coord = self.coord
        # Coord
        ora = obj['RA']
        odec = obj['DEC']
        #if ((isinstance(obj['RA'], Quantity)) |
        #        (isinstance(obj['RA'],astropy.table.column.MaskedColumn)) |
        #        (isinstance(obj['RA'], astropy.table.column.Column))):
        #    ora = obj['RA']
        #    odec = obj['DEC']
        #else:
        #    ora = obj['RA']*u.deg
        #    odec = obj['DEC']*u.deg
        o_coord = SkyCoord(ra=ora, dec=odec, unit='deg')
        # Calculate
        rho, ang_sep = calc_rho(los_coord, o_coord, obj['Z'], self.cosmo, comoving=comoving)
        # Return
        return rho

    def get_associated_galaxies(self, z, los_coord=None, R=300*u.kpc, dv_tol=500*u.km/u.s):
        """Return a Table of associated galaxies for a given redshift and separation

        Parameters
        ----------
        z : float
          Redshift for association (usually IGM/CGM absorber redshift)
        R : Quantity
          Radius of impact parameter for association [300kpc]
        dv_tol : Quantity
          Velocity window for association [500km/s]
        los_coord : SkyCoord, optional
          Line-of-sight coordinates

        Returns
        -------
        assoc_gal : Table
          Table of associated galaxies (if any)
        rho : Quantity
          Impact parameters [kpc]
        """
        # los_radec
        if los_coord is None:
            los_coord = self.coord
        # Cut on z first
        dv_gal = const.c.to('km/s') * (self.galaxies['Z']-z)/(1+z)  # Approximate
        gdz = np.where(np.abs(dv_gal)<dv_tol)[0]
        if len(gdz) == 0:
            return None
        #
        gdz_gal = self.galaxies[gdz]
        rho = self.calc_rhoimpact(gdz_gal, los_coord)  # Could add this to Table
        #
        gd_rho = np.where(rho < R)[0]
        if len(gd_rho) == 0:
            return None
        # Return
        return gdz_gal[gd_rho], rho[gd_rho]

    def get_observed(self, theta=None, subtab=None):
        """Generate a Table of observed targets

        Optionally to an angular distance from field center

        Parameters
        ----------
        theta : Quantity or Angle, optional
          Angular radius
        subtab : Table, optional
          User may input a table for processing
          theta is ignored

        Returns
        -------
        obs_targ : Table
          Sub-table of targets that have been observed within theta (if given)
          and/or within subtab (if given)
        obs_dates : List
          List of observing dates [eventually might add to Table]
        indices : array
          Indices from the target table
        """
        if (self.targets is None) or (self.observing is None):
            raise ValueError('IgmGalaxyField: Need to fill the target and/or observing table first!')
        if subtab is None:
            if theta is None:
                subtab = self.targets
            else:
                # Trim on angular cut first
                targ_coord = SkyCoord(ra=self.targets['TARG_RA']*u.deg,
                    dec=self.targets['TARG_DEC']*u.deg)
                sep = self.coord.separation(targ_coord)
                gdsep = np.where(sep < theta)[0]
                if len(gdsep) == 0:
                    return None
                subtab = self.targets[gdsep]
        else:
            gdsep = np.arange(len(subtab)) # For indexing below
        # Generate mask (set all to False; True is masked in numpy)
        tmsk = np.array([False]*len(subtab))
        # Grab those with a MASK_NAME
        have_mask = np.where(~subtab['MASK_NAME'].mask)[0]
        if len(have_mask) == 0:
            warnings.warn("No sources with a MASK_NAME")
            pdb.set_trace()
            return None
        # Get unique mask values
        all_masks = subtab['MASK_NAME'][have_mask]
        uni_masks = np.unique(all_masks)
        obs_dict = {}
        # Loop on these
        for mask in uni_masks:
            obs_dates = self.get_mask_obsdate(mask)
            if len(obs_dates) > 0:
                mt2 = np.where(subtab['MASK_NAME'][have_mask]==mask)
                tmsk[have_mask[mt2]] = True
                obs_dict[mask] = obs_dates
        # Finish
        return subtab[tmsk], obs_dict, gdsep[tmsk]

    def get_unobserved(self, theta=None):
        """Generate a Table of unobserved targets within an angular distance

        Parameters
        ----------
        theta : Quantity, optional
          Angular distance

        Returns
        -------
        unobs_targ : Table
          Sub-table of targets that have been not been observed within theta (if given)
        """
        if (self.targets is None) or (self.observing is None):
            raise ValueError('IgmGalaxyField: Need to fill the target and/or observing table first!')
        # Trim on angular cut first
        if theta is None:
            targ_coord = SkyCoord(ra=self.targets['TARG_RA']*u.deg,
                dec=self.targets['TARG_DEC']*u.deg)
            sep = self.coord.separation(targ_coord)
            gdsep = np.where(sep < theta)[0]
            if len(gdsep) == 0:
                return None
            # Set all to False to start
            subtab = self.targets[gdsep]
        else:
            subtab = self.targets
        tmsk = np.array([True]*len(subtab))
        # Grab observed (short cut!)
        obs_tab, odict, _ = self.get_observed(theta, subtab=subtab)
        # Remove those
        for kk, row in enumerate(subtab):
            if row['TARG_RA'] in obs_tab['TARG_RA']: # Could use DEC too
                tmsk[kk] = False
        # Return
        return subtab[tmsk]

    def get_mask_obsdate(self, mask_name):
        """Given a mask name, find the observing dates

        Parameters
        ----------
        mask_name : str
          Name of the mask

        Returns
        -------
        obs_dates : List
          List of the observing dates (can be empty)
        """
        if self.observing is None:
            raise ValueError('Need to fill observing info!')
        #
        mt = np.where(self.observing['MASK_NAME'] == mask_name)[0]
        if self.observing['DATE_OBS'].mask[mt[0]]:
            return []
        obs_dates = [self.observing['DATE_OBS'][imt] for imt in mt]
        # Return
        return obs_dates

    def clean_duplicates(self, table, tol=1*u.arcsec, method='first'):
        """ Clean duplicates in table based on (ra,dec) coordinates

        Parameters
        ----------
        table : Table
            Table to clean duplicates based on (ra, dec)
        tol : Angle, optional
            Angular tolerance for considering duplicates
        method : str, optional
            Method to use. Current options are:
            ``'first'``: if duplicates exist keep only the first one

        Returns
        -------
        cleaned_table : Table
            A version of `table` without duplicates
        """

        # TODO: add more methods for merging/cleaning duplicates

        if method not in ['first']:
            raise RuntimeError('Not ready for this method=`{}`'.format(method))

        isdup, idx, dcoord = pfu.check_dup_table(table, tol=tol)
        dup_inds = np.where(isdup == True)[0]
        keep = []

        if method == 'first':
            for ii in dup_inds:
                mtch = np.where(dcoord[ii].separation(dcoord) < tol)[0]
                keep.append(min(mtch))

            first_dup = np.unique(np.array(keep))

            no_dup = np.arange(len(idx))[~isdup]
            clean_inds = np.append(no_dup, first_dup)
            clean_inds = np.sort(clean_inds)

        #return
        return table[clean_inds]

    #
    def __repr__(self):
        return ('<{:s}: {:s} {:s} {:s}>'.format(
                self.__class__.__name__,
                 self.name,
                 self.coord.ra.to_string(unit=u.hour, sep=':', pad=True),
                 self.coord.dec.to_string(sep=':', pad=True, alwayssign=True)))

import numpy,import astropy,from astropy