'''

   3D port array boundary condition


   TE Mode
    Expression are based on Microwave Engineering p122 - p123.
    Note that it consists from two terms
       1)  \int dS W \dot n \times iwH (VectorFETangentIntegrator does this)
       2)  an weighting to evaulate mode amplutude from the E field
           on a boundary
   TEM Mode
       E is parallel to the periodic edge
       Mode number is ignored (of course)

   About sign of port phasing.
       positive phasing means the incoming wave appears later (phase delay)
       at the port

    2020 7/20  generated from portarray
'''
from petram.phys.common.rf_port_geometry import (analyze_geom_norm_ref,
                                                 analyze_rect_geom,
                                                 analyze_coax_geom,
                                                 analyze_circular_geom)
from petram.phys.em3d.em3d_base import EM3D_Bdry, EM3D_Domain
from petram.phys.phys_model import Phys
from petram.model import Bdry, Pair
from petram.phys.vtable import VtableElement, Vtable
from mfem.common.chypre import (LF2PyVec,
                                PyVec2PyMat,
                                Array2PyVec,
                                IdentityPyMat,
                                HStackPyVec)
from petram.mfem_config import use_parallel

from itertools import cycle
import numpy as np

import petram.debug as debug
dprint1, dprint2, dprint3 = debug.init_dprints('EM3D_PortArray')

if use_parallel:
    import mfem.par as mfem
    from mfem.common.mpi_debug import nicePrint
    from petram.helper.mpi_recipes import *

else:
    import mfem.ser as mfem
    nicePrint = dprint1


data = (('inc_amp', VtableElement('inc_amp',
                                  type='array',
                                  guilabel='incoming amp',
                                  default=0.0,
                                  tip="amplitude of incoming wave")),
        ('inc_phase', VtableElement('inc_phase',
                                    type='float',
                                    guilabel='incoming phase (deg)',
                                    default=0.0,
                                    tip="phase of incoming wave")),
        ('inc_dphase', VtableElement('inc_dphase',
                                     type='float',
                                     guilabel='phase progression (deg)',
                                     default=0.0,
                                     tip="phase progression of incoming wave")),
        ('epsilonr', VtableElement('epsilonr',
                                   type='complex',
                                   guilabel='epsilonr',
                                   default=1.0,
                                   tip="relative permittivity")),
        ('mur', VtableElement('mur',
                              type='complex',
                              guilabel='mur',
                              default=1.0,
                              tip="relative permeability")),)


def bdry_constraints():
    return [EM3D_PortArray]


class EM3D_PortArray(EM3D_Bdry):
    extra_diagnostic_print = True    
    vt = Vtable(data)

    def __init__(self, mode='TE', mn='0,1', inc_amp='1',
                 inc_phase='0', inc_dphase='90', port_idx=1):
        super(EM3D_PortArray, self).__init__(mode=mode,
                                             mn=mn,
                                             inc_amp=inc_amp,
                                             inc_phase=inc_phase,
                                             inc_dphase=inc_dphase,
                                             port_idx=port_idx)
        Phys.__init__(self)

    def extra_DoF_name(self):
        p1 = int(self.port_idx)
        p2 = len(self._sel_index)
        txt = str(p1) + '_' + str(p2)
        return self.get_root_phys().dep_vars[0] + "_port_" + txt

    def get_probes(self):
        p1 = int(self.port_idx)
        p2 = len(self._sel_index)
        txt = str(p1) + '_' + str(p2)
        return [self.get_root_phys().dep_vars[0] + "_port_" + txt]

    def attribute_set(self, v):
        super(EM3D_PortArray, self).attribute_set(v)
        v['port_idx'] = 1
        v['mode'] = 'TE'
        v['mn'] = [1, 0]
        v['inc_amp_txt'] = '1.0'
        v['inc_phase_txt'] = '0.0'
        v['inc_dphase_txt'] = '90.0'
        v['inc_amp'] = 1.0
        v['inc_phase'] = 0.0
        v['inc_dphase'] = 0.0
        v['epsilonr'] = 1.0
        v['mur'] = 1.0
        v['sel_readonly'] = False
        v['sel_index'] = []
        v['isTimeDependent_RHS'] = True
        v['ref_pt'] = '1'        
        return v

    def panel1_param(self):
        return ([["port group id", str(self.port_idx), 0, {}],
                 ["mode", self.mode, 4, {"readonly": True,
                                         "choices": ["TE", "TEM", "Coax(TEM)"]}],
                 ["m/n", ','.join(str(x) for x in self.mn), 0, {}], 
                 ["ref. pt.", "", 0, {}],] +                
                self.vt.panel_param(self))

    def get_panel1_value(self):
        return ([str(self.port_idx),
                 self.mode, ','.join(str(x) for x in self.mn),
                 self.ref_pt] +
                self.vt.get_panel_value(self))

    def import_panel1_value(self, v):
        self.port_idx = v[0]
        self.mode = v[1]
        self.mn = [int(x) for x in v[2].split(',')]
        self.ref_pt = v[3]        
        self.vt.import_panel_value(self, v[4:])

    def panel4_param(self):
        ll = super(EM3D_PortArray, self).panel4_param()
        ll.append(['Varying (in time/for loop) RHS', False, 3, {"text": ""}])
        return ll

    def panel4_tip(self):
        return None

    def import_panel4_value(self, value):
        super(EM3D_PortArray, self).import_panel4_value(value[:-1])
        self.isTimeDependent_RHS = value[-1]

    def get_panel4_value(self):
        value = super(EM3D_PortArray, self).get_panel4_value()
        value.append(self.isTimeDependent_RHS)
        return value

    def verify_setting(self):
        if self.isTimeDependent_RHS:
            flag = True
        else:
            flag = False
        return flag, 'Varying RHS is not set', 'This potntially causes an error with PortScan. Set it Time/NL Dep. panel '

    def update_param(self):
        self.vt.preprocess_params(self)
        inc_amp, inc_phase, inc_dphase, eps, mur = self.vt.make_value_or_expression(
            self)

        '''
        self.update_inc_amp_phase()

    def update_inc_amp_phase(self):
        try:
            self.inc_amp = self.eval_phys_expr(
                str(self.inc_amp_txt),  'inc_amp')[0]
            self.inc_phase = self.eval_phys_expr(
                str(self.inc_phase_txt), 'inc_phase', chk_float = True)[0]
        except:
            raise ValueError("Cannot evaluate amplitude/phase to float number")
        '''

    def do_preprocess_params(self, engine):
        self.a = 0
        self.b = 0
        self.a_vec = 0
        self.b_vec = 0
        self.c = 0
        self.ctr = 0
        
        # find normal (outward) vector and reference point location.
        mesh = engine.get_emesh(mm=self)
        nbe = mesh.GetNBE()
        ibe = np.array([i for i in range(nbe)
                        if mesh.GetBdrElement(i).GetAttribute() ==
                        self._sel_index[0]])

        norm, rptx = analyze_geom_norm_ref(mesh, ibe, self.ref_pt)
        self.norm = norm

        dprint1("Normal Vector " + str(self.norm))
        dprint1("Ref. Point" + str(rptx))

        # find rectangular shape
        if str(self.mode).upper().strip() in ['TE', 'TM', 'TEM']:
            geom_data, vv = analyze_rect_geom(self, mesh, ibe, norm)
            self.a = geom_data["a"]
            self.a_vec = geom_data["a_vec"]
            self.b = geom_data["b"]
            self.b_vec = geom_data["b_vec"]
            self.c = geom_data["c"]
            dprint1("Long Edge  " + self.a.__repr__())
            dprint1("Long Edge Vec." + list(self.a_vec).__repr__())
            dprint1("Short Edge  " + self.b.__repr__())
            dprint1("Short Edge Vec." + list(self.b_vec).__repr__())
            
        elif self.mode == 'Coax(TEM)':
            geom_data, vv = analyze_coax_geom(mesh, ibe, norm, rptx)
            self.a = geom_data["a"]
            self.b = geom_data["b"]
            self.ctr = geom_data["ctr"]
            self.ax1 = geom_data["ax1"]
            self.ax2 = geom_data["ax2"]
            
        C_Et, C_jwHt = self.get_coeff_cls()

        self.vt.preprocess_params(self)
        inc_amp, inc_phase, inc_dphase, eps, mur = self.vt.make_value_or_expression(
            self)

        dprint1("E field pattern", eps, mur)
        Et = C_Et(3, self, real=True, eps=eps, mur=mur,
                  m=self.mn[0], n=self.mn[1])
        for p in vv:
            dprint1(p.__repr__() + ' : ' + Et.EvalValue(p).__repr__())
        dprint1("H field pattern")
        cnorm = self.get_root_phys().get_coeff_norm()
        Ht = C_jwHt(3, self, real=False, eps=eps, mur=mur, cnorm=cnorm,
                    m=self.mn[0], n=self.mn[1])
        for p in vv:
            dprint1(p.__repr__() + ' : ' + Ht.EvalValue(p).__repr__())

    def read_portparams(self):
        return (self.a, self.b, self.a_vec, self.b_vec, self.c, self.ctr)

    def set_portparams(self, params):
        self.a, self.b, self.a_vec, self.b_vec, self.c, self.ctr = params

    def preprocess_params(self, engine):
        sels = self._sel_index

        port_params = []
        for s in sels:
            self._sel_index = [s]
            self.do_preprocess_params(engine)
            params = self.read_portparams()
            port_params.append(params)

        self._port_params = port_params
        self._sel_index = sels

    def get_coeff_cls(self):
        from petram.phys.common.rf_portmode import get_portmode_coeff_cls

        return get_portmode_coeff_cls(self.mode)

    def has_lf_contribution(self, kfes):
        if kfes != 0:
            return False
        self.vt.preprocess_params(self)
        inc_amp, inc_phase, inc_dphase, eps, mur = self.vt.make_value_or_expression(
            self)

        return np.any(np.array(inc_amp) != 0)

    def _fixup_inc_amp(self, inc_amp):

        num_ports = len(self._sel_index)

        if len(inc_amp) == 1:
            inc_amp = [inc_amp[0]]*num_ports
        elif len(inc_amp) != num_ports:
            assert False, "Number of ports are not the same as number of incoming amplitude"
        else:
            pass  # looks good
        return inc_amp, num_ports

    def add_lf_contribution(self, engine, b, real=True, kfes=0):
        if real:
            dprint1("Add LF contribution(real)" + str(self._sel_index))
        else:
            dprint1("Add LF contribution(imag)" + str(self._sel_index))

        self.vt.preprocess_params(self)
        inc_amp, inc_phase, inc_dphase, eps, mur = self.vt.make_value_or_expression(
            self)

        C_Et, C_jwHt = self.get_coeff_cls()

        sels = self._sel_index

        from itertools import cycle
        dphase = cycle(np.atleast_1d(inc_dphase))

        inc_amp, num_ports = self._fixup_inc_amp(inc_amp)
        ph = inc_phase

        dprint1("Power, Phase: ", inc_amp, inc_phase)

        for kport in range(num_ports):
            sel = self._sel_index[kport]
            params = self._port_params[kport]
            inc_a = inc_amp[kport]

            self.set_portparams(params)

            inc_wave = inc_a * np.exp(1j * ph / 180. * np.pi)

            phase = np.angle(inc_wave) * 180 / np.pi
            amp = np.sqrt(np.abs(inc_wave))
            cnorm = self.get_root_phys().get_coeff_norm()

            Ht = C_jwHt(3, self, real=real, amp=amp, eps=eps, mur=mur, phase=phase, cnorm=cnorm,
                        m=self.mn[0], n=self.mn[1])
            Ht = self.restrict_coeff(Ht, engine, vec=True, idx=[sel])

            intg = mfem.VectorFEBoundaryTangentLFIntegrator(Ht)
            b.AddBoundaryIntegrator(intg)

            ph = ph + next(dphase)

    def has_extra_DoF(self, kfes):
        if kfes != 0:
            return False
        return True

    def get_extra_NDoF(self):
        return len(self._sel_index)
        # (in future) must be number of modes on this port...

    def is_extra_RHSonly(self):
        return True

    def postprocess_extra(self, sol, flag, sol_extra):
        name = self.name() + '_' + str(self.port_idx)
        sol_extra[name] = sol.toarray()

    def do_add_extra_contribution(
            self, engine, inc_amp, inc_phase, eps, mur, sel):
        C_Et, C_jwHt = self.get_coeff_cls()
        cnorm = self.get_root_phys().get_coeff_norm()

        fes = engine.get_fes(self.get_root_phys(), 0)

        lf1 = engine.new_lf(fes)
        Ht1 = C_jwHt(3, self, real=True, eps=eps, mur=mur, cnorm=cnorm,
                     m=self.mn[0], n=self.mn[1])
        Ht2 = self.restrict_coeff(Ht1, engine, vec=True, idx=[sel])
        intg = mfem.VectorFEBoundaryTangentLFIntegrator(Ht2)
        lf1.AddBoundaryIntegrator(intg)
        lf1.Assemble()
        lf1i = engine.new_lf(fes)
        Ht3 = C_jwHt(3, self, real=False, eps=eps, mur=mur, cnorm=cnorm,
                     m=self.mn[0], n=self.mn[1])
        Ht4 = self.restrict_coeff(Ht3, engine, vec=True, idx=[sel])
        intg = mfem.VectorFEBoundaryTangentLFIntegrator(Ht4)
        lf1i.AddBoundaryIntegrator(intg)
        lf1i.Assemble()

        lf2r = engine.new_lf(fes)
        Etr = C_Et(3, self, real=True, eps=eps, mur=mur,
                  m=self.mn[0], n=self.mn[1])
        Etr = self.restrict_coeff(Etr, engine, vec=True, idx=[sel])
        intg = mfem.VectorFEDomainLFIntegrator(Etr)
        lf2r.AddBoundaryIntegrator(intg)
        lf2r.Assemble()
        lf2i = engine.new_lf(fes)
        Eti = C_Et(3, self, real=False, eps=eps, mur=mur,
                   m=self.mn[0], n=self.mn[1])
        Eti = self.restrict_coeff(Eti, engine, vec=True, idx=[sel])
        intg = mfem.VectorFEDomainLFIntegrator(Eti)
        lf2i.AddBoundaryIntegrator(intg)
        lf2i.Assemble()

        arr = self.get_restriction_array(engine)

        xr = engine.new_gf(fes)
        xr.Assign(0.0)
        xr.ProjectBdrCoefficientTangent(Etr, arr)
        xi = engine.new_gf(fes)
        xi.Assign(0.0)
        xi.ProjectBdrCoefficientTangent(Eti, arr)

        et = engine.x2X(xr).GetDataArray() + \
                1j * engine.x2X(xi).GetDataArray()

        vec = engine.b2B(lf2r).GetDataArray() - \
                1j * engine.b2B(lf2i).GetDataArray()  # complex conjugate

        weight = np.sum(et.dot(vec))
        if use_parallel:
            weight = np.sum(allgather(weight))

        # t1
        t1 = LF2PyVec(lf1, lf1i)
        t1 *= 1


        # t2
        lf2i *= -1 # complex conjugate
        t2 = LF2PyVec(lf2r, lf2i, horizontal=True)
        t2 *= -1. / weight
        t2 = t2.transpose()

        # t4
        t4 = - np.array(
            [[np.sqrt(inc_amp) * np.exp(1j * inc_phase / 180. * np.pi)]])
        
        return (t1, t2, t4)

    def add_extra_contribution(self, engine, **kwargs):
        '''
        compute extra (multiplier) contribution

        Format of extar   (t2 is returnd as vertical(transposed) matrix)
        [M,  t1]   [  ]
        [      ] = [  ]
        [t2, t3]   [t4]

        and it returns if Lagurangian will be saved.
        '''
        dprint1("Add Extra contribution" + str(self._sel_index))

        self.vt.preprocess_params(self)
        inc_amp, inc_phase, inc_dphase, eps, mur = self.vt.make_value_or_expression(
            self)
        inc_amp, num_ports = self._fixup_inc_amp(inc_amp)
        ph = inc_phase
        dphase = cycle(np.atleast_1d(inc_dphase))

        dprint1("Power, Phase: ", inc_amp, inc_phase)

        t1_arr = []
        t2_arr = []
        t4 = []

        for kport in range(num_ports):
            sel = self._sel_index[kport]
            params = self._port_params[kport]
            inc_a = inc_amp[kport]

            self.set_portparams(params)
            # dprint1("phasing ", ph, sel)
            t1, t2, t4_1 = self.do_add_extra_contribution(
                engine, inc_a, ph, eps, mur, sel)

            t1_arr.append(t1)
            t2_arr.append(t2)

            t4.append(t4_1)
            ph = ph + next(dphase)

        t3 = IdentityPyMat(len(self._sel_index))
        t4 = Array2PyVec(np.array(t4))

        t1 = HStackPyVec(t1_arr)
        t2 = HStackPyVec(t2_arr)
        t2 = t2.transpose()

        #print(t1.shape, t2.shape, t3.shape, t4.shape)

        return (t1, t2, t3, t4, True)
