# -*- coding: utf-8 -*-
"""
plot_fchar.py
"""

import numpy as np
import matplotlib.pyplot as plt
import sol.farfield

EPS = 1e-12

# plot F-char
def plot(Post, Ne, Nfreq, Freq, Nfeed, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin):

    if (Post['smith'] > 0):
        _smith(Post, Nfeed, Nfreq, Freq, Title, Z0, Zin)

    if (Post['zin'][0] > 0):
        _fchar(1, Post, Ne, Nfeed, Nfreq, Freq, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin)

    if (Post['yin'][0] > 0):
        _fchar(2, Post, Ne, Nfeed, Nfreq, Freq, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin)

    if (Post['ref'][0] > 0):
        _fchar(3, Post, Ne, Nfeed, Nfreq, Freq, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin)

    if (Post['f0d'][2] > 0):
        _fchar(4, Post, Ne, Nfeed, Nfreq, Freq, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin)

# input impedance/input admittance/reflection/far point (private)
def _fchar(itype, Post, Ne, Nfeed, Nfreq, Freq, E_posc, E_lng, E_tan, E_ifeed, E_feed, Iground, Title, Z0, Xc, Zin):

    if Nfreq < 2:
        print('*** single frequency.')
        return
    if (itype < 1) or (itype > 4):
        return

    name = ['input impedance', 'input admittance', 'reflection', 'far field']
    comp = [['Rin', 'Xin'], ['Gin', 'Bin'], ['', ''], ['', '']]
    unit = ['[ohm]', '[mS]', '[dB]', '']

    # open far0d.log
    if itype == 4:
        fname = 'far0d.log'
        fp = open(fname, 'wt', encoding='utf-8')
    
    # pages
    npage = Nfeed if (itype == 1) or (itype == 2) or (type == 3) else 1

    # X data : frequency
    x = Freq * Post['fscale']

    for ifeed in range(npage):
        # Rin/Gin/Ref/F0d (real)
        if   itype == 1:
            y = Zin[ifeed].real
            color = 'r'
        elif itype == 2:
            y = (1e3 / Zin[ifeed]).real
            color = 'r'
        elif itype == 3:
            gamma = (Zin[ifeed] - Z0) / (Zin[ifeed] + Z0)
            y = 20 * np.log10(np.maximum(np.abs(gamma), EPS))
            color = 'k'
        elif itype == 4:
            # log (header)
            _log_f0d_1(fp, Post)

            y = np.zeros(Nfreq, float)
            for ifreq, freq in enumerate(Freq):
                ffctr = sol.farfield.factor(Ne, Nfeed, E_ifeed, E_feed, Zin, Z0, ifreq, Post['mloss'])
                #_, pfar = sol.farfield.field(Freq, E_posc, E_lng, E_tan, Iground, Xc, ifreq, Post['f0d'][0], Post['f0d'][1], fctr)
                _, pfar = sol.farfield.field(ifreq, freq, Post['f0d'][0], Post['f0d'][1], ffctr, Ne, E_posc, E_lng, E_tan, Iground, Xc)
                y[ifreq] = 10 * np.log10(max(pfar[0], EPS))

                # log (body)
                _log_f0d_2(fp, ifreq, freq, pfar)

            color = 'k'

        # figure
        strfig = 'OpenMOM - %s (%d/%d)' % (name[itype - 1], ifeed + 1, Nfeed)
        fig = plt.figure(strfig, figsize=(Post['w2d'][0], Post['w2d'][1]))
        ax = fig.add_subplot()

        # plot
        ax.plot(x, y, color=color, label=comp[itype - 1][0])
        ax.grid()

        # Xin/Bin
        if (itype == 1) or (itype == 2):
            if   (itype == 1):
                y = Zin[ifeed].imag
            elif (itype == 2):
                y = (1e3 / Zin[ifeed]).imag
            ax.plot(x, y, color='b', label=comp[itype - 1][1])

        # X-axis
        ax.set_xlim(x[0], x[-1])
        if Nfreq > 1:
            fdiv = np.linspace(x[0], x[-1], Post['freqdiv'] + 1)
        else:
            fdiv = x[0]
        ax.set_xticks(fdiv)

        # X-label
        ax.set_xlabel('frequency' + Post['funit'])

        # Y-label
        if (itype == 1) or (itype == 2) or (itype == 3):
            strunit = unit[itype - 1]
        elif itype == 4:
            name[itype - 1] = Post['farname']
            strunit = Post['f0dunit']
        ax.set_ylabel(name[itype - 1] + ' ' + strunit)

        # title
        if   (itype == 1) or (itype == 2):
            strtitle = '%s\n(feed# = %d)' % (Title, ifeed + 1)
        elif (itype == 3):
            strtitle = '%s\nfeed# = %d, Z0 = %g[ohm], min = %.3f[dB]' % (Title, ifeed + 1, Z0, np.min(y))
        elif (itype == 4):
            strtitle = '%s\n(theta, phi) = (%g, %g)[deg], max = %.3f%s' % (Title, Post['f0d'][0], Post['f0d'][1], np.max(y), strunit)
        ax.set_title(strtitle)

        # user scale
        if   (itype == 1) and (Post['zin'][0] > 1):
            ax.set_ylim(Post['zin'][1], Post['zin'][2])
            ydiv = np.linspace(Post['zin'][1], Post['zin'][2], Post['zin'][3] + 1)
            ax.set_yticks(ydiv)
        elif (itype == 2) and (Post['yin'][0] > 1):
            ax.set_ylim(Post['yin'][1], Post['yin'][2])
            ydiv = np.linspace(Post['yin'][1], Post['yin'][2], Post['yin'][3] + 1)
            ax.set_yticks(ydiv)
        elif (itype == 3) and (Post['ref'][0] > 1):
            ax.set_ylim(Post['ref'][1], Post['ref'][2])
            ydiv = np.linspace(Post['ref'][1], Post['ref'][2], Post['ref'][3] + 1)
            ax.set_yticks(ydiv)
        elif (itype == 4) and (Post['f0d'][2] > 1):
            ax.set_ylim(Post['f0d'][3], Post['f0d'][4])
            ydiv = np.linspace(Post['f0d'][3], Post['f0d'][4], Post['f0d'][5] + 1)
            ax.set_yticks(ydiv)

        # get limits
        xlim = ax.get_xlim()#; print(xlim)
        ylim = ax.get_ylim()#; print(ylim)

        # Y=0 line
        if (ylim[0] < 0) and (ylim[1] > 0):
            ax.plot([x[0], x[-1]], [0, 0], color='gray')

        # legend
        if (itype == 1) or (itype == 2):
            ax.legend(loc='best')

        # VSWR (right side)
        if itype == 3:
            vs = np.array([1.1, 1.2, 1.5, 2.0, 3.0, 5.0, 10.0])
            vsdb = 20 * np.log10((vs - 1) / (vs + 1))
            lx = xlim[1] - xlim[0]
            ly = ylim[1] - ylim[0]
            ax.text(xlim[1], ylim[1], 'VSWR', color='r')
            for ivs in range(len(vs)):
                if (vsdb[ivs] >= ylim[0]) and (vsdb[ivs] <= ylim[1]):
                    xvs = [xlim[1], xlim[1] - 0.015 * lx]
                    yvs = [vsdb[ivs], vsdb[ivs]]
                    ax.plot(xvs, yvs, color='r')
                    ax.text(xlim[1] + 0.005 * lx, vsdb[ivs] - 0.01 * ly, str(vs[ivs]), color='r')

# smith chart (private)
def _smith(Post, Nfeed, Nfreq, Freq, Title, Z0, Zin):

    if (Nfeed < 1) or (Nfreq < 1):
        return

    # figure
    strfig = 'OpenMOM - SmithChart'
    fig = plt.figure(strfig, figsize=(Post['w2d'][0], Post['w2d'][1]))
    ax = fig.add_subplot()

    # Smith Chart
    _smithchart(ax, Z0)

    
    # Zin
    for ifeed in range(Nfeed):
        gamma = (Zin[ifeed] - Z0) / (Zin[ifeed] + Z0)
        ax.plot(gamma.real, gamma.imag, color='k', marker='o', markersize=3)
        # #n : start frequency
        ax.text(gamma[0].real, gamma[0].imag, '#' + str(ifeed + 1), color='r')
    
    # title
    ax.set_title('%s\nf%s = %g - %g, Z0 = %g[ohm], No. of feeds = %d' %
        (Title, Post['funit'], Freq[0] * Post['fscale'], Freq[-1] * Post['fscale'], Z0, Nfeed))

    # layout
    ax.set_aspect('equal')
    ax.axis('off')
    
# Smith Chart (private)
def _smithchart(ax, Z0):
    dh = 0.06 # string shift
    
    # R=0
    div = 50
    angle = np.linspace(0, 2 * np.pi, div + 1)
    ax.plot(np.cos(angle), np.sin(angle), color='gray')

    # X=0
    ax.plot([-1, +1], [0, 0], color='gray', linestyle='--')
    
    # DB
    const = [0.2, 0.5, 1, 2, 5]
    zmin  = [0,    1.1, 2.2, 3.5,  5,  20,  200]
    zmax  = [1,    2,   3,   4,   10, 100, 1000]
    zdiv  = [20,   9,   4,   1,    5,   8,    4]
    
    # R-const
    for ir in range(len(const)):
        xlist = np.array([])
        for ix in range(len(zmin)):
            xlist = np.hstack((xlist, np.linspace(zmin[ix], zmax[ix], zdiv[ix] + 1)))

        z = const[ir] + (1j * xlist)
        g = (z - 1) / (z + 1)
        ax.plot(g.real, +g.imag, color='gray', linestyle='--')
        ax.plot(g.real, -g.imag, color='gray', linestyle='--')
        strr = '%d' % (Z0 * const[ir])
        ax.text(g[0].real - dh, 0, strr)

    # X-const
    for ix in range(len(const)):
        rlist = np.array([])
        for ir in range(len(zmin)):
            rlist = np.hstack((rlist, np.linspace(zmin[ir], zmax[ir], zdiv[ir] + 1)))

        z = rlist + (1j * const[ix])
        g = (z - 1) / (z + 1)
        ax.plot(g.real, +g.imag, color='gray', linestyle='--')
        ax.plot(g.real, -g.imag, color='gray', linestyle='--')
        strx = '%dj' % (Z0 * const[ix])
        ax.text(g[0].real - dh, +g[0].imag, '+' + strx)
        ax.text(g[0].real - dh, -g[0].imag, '-' + strx)

# far0d.log (header) (private)
def _log_f0d_1(fp, Post):
    fp.write('theta=%.3f[deg] phi=%.3f[deg]\n' % (Post['f0d'][0], Post['f0d'][1]))
    fp.write('  No. frequency[Hz]    E-abs[dB]  E-theta[dB]    E-phi[dB]  E-major[dB]  E-minor[dB]   E-RHCP[dB]   E-LHCP[dB] AxialRatio[dB]\n')

# far0d.log (body) (private)
def _log_f0d_2(fp, ifreq, freq, pfar):
    fmt = '%4d%15.5e%13.4f%13.4f%13.4f%13.4f%13.4f%13.4f%13.4f%13.4f\n'
    fdb = 10 * np.log10(np.maximum(pfar, EPS))
    fp.write(fmt % (ifreq, freq, fdb[0], fdb[1], fdb[2], fdb[3], fdb[4], fdb[5], fdb[6], fdb[3] - fdb[4]))
