#!/usr/bin/env python
"""A simple viewer for Stokes patterns based on two far-field pattern files.
(Possibly based on one FF pattern files if it has two requests: one for each
polarization channel.)"""
import os
import argparse
import numpy
import matplotlib.pyplot as plt
from antpat.reps.sphgridfun.tvecfun import TVecFields
from antpat.radfarfield import RadFarField
from antpat.dualpolelem import DualPolElem

FEKOsuffix = 'ffe'
GRASPsuffix = 'swe'
NECsuffix = 'out'


def Jones2Stokes(Jones):
    """Convert Jones matrix to Stokes vector. This assumes dual-pol antenna receiving unpolarized unit
    valued radiation i.e. incoming Stokes = (1,0,0,0)."""
    brightmat = numpy.matmul(Jones, numpy.swapaxes(numpy.conjugate(Jones),-1,-2))
    StokesI = numpy.real(brightmat[...,0,0]+brightmat[...,1,1])
    StokesQ = numpy.real(brightmat[...,0,0]-brightmat[...,1,1])
    StokesU = numpy.real(brightmat[...,0,1]+brightmat[...,1,0])
    StokesV = numpy.imag(brightmat[...,0,1]-brightmat[...,1,0])
    return StokesI, StokesQ, StokesU, StokesV


def plotStokes_fromFEKOfiles(p_chan_file, q_chan_file, freq):
    (tvf_p, tvf_q) = (TVecFields(), TVecFields())
    tvf_p.load_ffe(p_chan_file)
    tvf_q.load_ffe(q_chan_file)
    (ant_p, ant_q) = (RadFarField(tvf_p), RadFarField(tvf_q))
    (p_chan_name, q_chan_name) = (os.path.basename(p_chan_file), os.path.basename(q_chan_file))
    (ant_p.name, ant_q.name) = (p_chan_name, q_chan_name)
    dualpolAnt = DualPolElem(ant_p, ant_q)
    THETA, PHI, Jones = dualpolAnt.getJonesPat(freq)
    (StokesI, StokesQ, StokesU, StokesV) = Jones2Stokes(Jones)

    x = THETA*numpy.cos(PHI)
    y = THETA*numpy.sin(PHI)
    #x= THETA
    #y=PHI
    xyNames = ('theta*cos(phi)','theta*sin(phi)')
    fig = plt.figure()
    ax1 = fig.add_subplot(221)
    plt.pcolormesh(x, y, 10*numpy.log10(StokesI), label="I")
    #plt.pcolormesh(x, y, StokesI, label="I")
    plt.colorbar()
    ax1.set_title('I (dB)')

    ax2 = fig.add_subplot(222)
    plt.pcolormesh(x, y, StokesQ/StokesI, label="Q")
    plt.colorbar()
    ax2.set_title('Q/I')

    ax3 = fig.add_subplot(223)
    plt.pcolormesh(x, y, StokesU/StokesI, label="U")
    plt.colorbar()
    ax3.set_title('U/I')

    ax4 = fig.add_subplot(224)
    plt.pcolormesh(x, y, StokesV/StokesI, label="V")
    plt.colorbar()
    ax4.set_title('V/I')
    fig.suptitle('Stokes (azimuthal-equidistant proj) @ ' +str(freq/1e9)+' GHz')
    plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("p_chan_file",
                 help='Filename of polarization channel p')
    parser.add_argument("q_chan_file",
                 help='Filename of polarization channel p')
    parser.add_argument("freq", nargs='?', type=float,
                        help="Frequency in Hertz")
    args = parser.parse_args()
    
    if args.p_chan_file.endswith(FEKOsuffix):
        plotStokes_fromFEKOfiles(args.p_chan_file, args.q_chan_file, args.freq)
    elif args.p_chan_file.endswith(GRASPsuffix):
        print("Not implemented yet.")
    elif args.p_chan_file.endswith(NECsuffix):
        print("Not implemented yet.")
    else:
        print("Far-field pattern file type not known")
        exit(1)

[((22, 14, 22, 63), 'numpy.real', 'numpy.real', ({(22, 25, 22, 62): 'brightmat[..., 0, 0] + brightmat[..., 1, 1]'}, {}), '(brightmat[..., 0, 0] + brightmat[..., 1, 1])', False, 'import numpy\n'), ((23, 14, 23, 63), 'numpy.real', 'numpy.real', ({(23, 25, 23, 62): 'brightmat[..., 0, 0] - brightmat[..., 1, 1]'}, {}), '(brightmat[..., 0, 0] - brightmat[..., 1, 1])', False, 'import numpy\n'), ((24, 14, 24, 63), 'numpy.real', 'numpy.real', ({(24, 25, 24, 62): 'brightmat[..., 0, 1] + brightmat[..., 1, 0]'}, {}), '(brightmat[..., 0, 1] + brightmat[..., 1, 0])', False, 'import numpy\n'), ((25, 14, 25, 63), 'numpy.imag', 'numpy.imag', ({(25, 25, 25, 62): 'brightmat[..., 0, 1] - brightmat[..., 1, 0]'}, {}), '(brightmat[..., 0, 1] - brightmat[..., 1, 0])', False, 'import numpy\n'), ((36, 17, 36, 42), 'antpat.dualpolelem.DualPolElem', 'DualPolElem', ({(36, 29, 36, 34): 'ant_p', (36, 36, 36, 41): 'ant_q'}, {}), '(ant_p, ant_q)', False, 'from antpat.dualpolelem import DualPolElem\n'), ((45, 10, 45, 22), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((49, 4, 49, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((53, 4, 53, 52), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (), '', True, 'import matplotlib.pyplot as plt\n'), ((54, 4, 54, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((58, 4, 58, 52), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (), '', True, 'import matplotlib.pyplot as plt\n'), ((59, 4, 59, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((63, 4, 63, 52), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (), '', True, 'import matplotlib.pyplot as plt\n'), ((64, 4, 64, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((67, 4, 67, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((71, 13, 71, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((30, 22, 30, 34), 'antpat.reps.sphgridfun.tvecfun.TVecFields', 'TVecFields', ({}, {}), '()', False, 'from antpat.reps.sphgridfun.tvecfun import TVecFields\n'), ((30, 36, 30, 48), 'antpat.reps.sphgridfun.tvecfun.TVecFields', 'TVecFields', ({}, {}), '()', False, 'from antpat.reps.sphgridfun.tvecfun import TVecFields\n'), ((33, 22, 33, 40), 'antpat.radfarfield.RadFarField', 'RadFarField', ({(33, 34, 33, 39): 'tvf_p'}, {}), '(tvf_p)', False, 'from antpat.radfarfield import RadFarField\n'), ((33, 42, 33, 60), 'antpat.radfarfield.RadFarField', 'RadFarField', ({(33, 54, 33, 59): 'tvf_q'}, {}), '(tvf_q)', False, 'from antpat.radfarfield import RadFarField\n'), ((34, 34, 34, 63), 'os.path.basename', 'os.path.basename', ({(34, 51, 34, 62): 'p_chan_file'}, {}), '(p_chan_file)', False, 'import os\n'), ((34, 65, 34, 94), 'os.path.basename', 'os.path.basename', ({(34, 82, 34, 93): 'q_chan_file'}, {}), '(q_chan_file)', False, 'import os\n'), ((40, 14, 40, 28), 'numpy.cos', 'numpy.cos', ({(40, 24, 40, 27): 'PHI'}, {}), '(PHI)', False, 'import numpy\n'), ((41, 14, 41, 28), 'numpy.sin', 'numpy.sin', ({(41, 24, 41, 27): 'PHI'}, {}), '(PHI)', False, 'import numpy\n'), ((21, 51, 21, 73), 'numpy.conjugate', 'numpy.conjugate', ({(21, 67, 21, 72): 'Jones'}, {}), '(Jones)', False, 'import numpy\n'), ((47, 28, 47, 48), 'numpy.log10', 'numpy.log10', ({(47, 40, 47, 47): 'StokesI'}, {}), '(StokesI)', False, 'import numpy\n')]