#!/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)

[((755, 810), 'numpy.real', 'numpy.real', (['(brightmat[..., 0, 0] + brightmat[..., 1, 1])'], {}), '(brightmat[..., 0, 0] + brightmat[..., 1, 1])\n', (765, 810), False, 'import numpy\n'), ((819, 874), 'numpy.real', 'numpy.real', (['(brightmat[..., 0, 0] - brightmat[..., 1, 1])'], {}), '(brightmat[..., 0, 0] - brightmat[..., 1, 1])\n', (829, 874), False, 'import numpy\n'), ((883, 938), 'numpy.real', 'numpy.real', (['(brightmat[..., 0, 1] + brightmat[..., 1, 0])'], {}), '(brightmat[..., 0, 1] + brightmat[..., 1, 0])\n', (893, 938), False, 'import numpy\n'), ((947, 1002), 'numpy.imag', 'numpy.imag', (['(brightmat[..., 0, 1] - brightmat[..., 1, 0])'], {}), '(brightmat[..., 0, 1] - brightmat[..., 1, 0])\n', (957, 1002), False, 'import numpy\n'), ((1454, 1479), 'antpat.dualpolelem.DualPolElem', 'DualPolElem', (['ant_p', 'ant_q'], {}), '(ant_p, ant_q)\n', (1465, 1479), False, 'from antpat.dualpolelem import DualPolElem\n'), ((1740, 1752), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (1750, 1752), True, 'import matplotlib.pyplot as plt\n'), ((1895, 1909), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (1907, 1909), True, 'import matplotlib.pyplot as plt\n'), ((1974, 2024), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (['x', 'y', '(StokesQ / StokesI)'], {'label': '"""Q"""'}), "(x, y, StokesQ / StokesI, label='Q')\n", (1988, 2024), True, 'import matplotlib.pyplot as plt\n'), ((2027, 2041), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (2039, 2041), True, 'import matplotlib.pyplot as plt\n'), ((2103, 2153), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (['x', 'y', '(StokesU / StokesI)'], {'label': '"""U"""'}), "(x, y, StokesU / StokesI, label='U')\n", (2117, 2153), True, 'import matplotlib.pyplot as plt\n'), ((2156, 2170), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (2168, 2170), True, 'import matplotlib.pyplot as plt\n'), ((2232, 2282), 'matplotlib.pyplot.pcolormesh', 'plt.pcolormesh', (['x', 'y', '(StokesV / StokesI)'], {'label': '"""V"""'}), "(x, y, StokesV / StokesI, label='V')\n", (2246, 2282), True, 'import matplotlib.pyplot as plt\n'), ((2285, 2299), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ([], {}), '()\n', (2297, 2299), True, 'import matplotlib.pyplot as plt\n'), ((2410, 2420), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (2418, 2420), True, 'import matplotlib.pyplot as plt\n'), ((2463, 2488), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (2486, 2488), False, 'import argparse\n'), ((1129, 1141), 'antpat.reps.sphgridfun.tvecfun.TVecFields', 'TVecFields', ([], {}), '()\n', (1139, 1141), False, 'from antpat.reps.sphgridfun.tvecfun import TVecFields\n'), ((1143, 1155), 'antpat.reps.sphgridfun.tvecfun.TVecFields', 'TVecFields', ([], {}), '()\n', (1153, 1155), False, 'from antpat.reps.sphgridfun.tvecfun import TVecFields\n'), ((1243, 1261), 'antpat.radfarfield.RadFarField', 'RadFarField', (['tvf_p'], {}), '(tvf_p)\n', (1254, 1261), False, 'from antpat.radfarfield import RadFarField\n'), ((1263, 1281), 'antpat.radfarfield.RadFarField', 'RadFarField', (['tvf_q'], {}), '(tvf_q)\n', (1274, 1281), False, 'from antpat.radfarfield import RadFarField\n'), ((1317, 1346), 'os.path.basename', 'os.path.basename', (['p_chan_file'], {}), '(p_chan_file)\n', (1333, 1346), False, 'import os\n'), ((1348, 1377), 'os.path.basename', 'os.path.basename', (['q_chan_file'], {}), '(q_chan_file)\n', (1364, 1377), False, 'import os\n'), ((1611, 1625), 'numpy.cos', 'numpy.cos', (['PHI'], {}), '(PHI)\n', (1620, 1625), False, 'import numpy\n'), ((1640, 1654), 'numpy.sin', 'numpy.sin', (['PHI'], {}), '(PHI)\n', (1649, 1654), False, 'import numpy\n'), ((710, 732), 'numpy.conjugate', 'numpy.conjugate', (['Jones'], {}), '(Jones)\n', (725, 732), False, 'import numpy\n'), ((1812, 1832), 'numpy.log10', 'numpy.log10', (['StokesI'], {}), '(StokesI)\n', (1823, 1832), False, 'import numpy\n')]