from __future__ import print_function

"""
Script to read simulated port patterns of the PAF on an ASKAP antenna from .mat files generated by Stuart Hay, 

It will also manipulate coordinate systems and port numberings to match what we use for measurements on the hardware, 
and store the port patterns in a convenient .hdf5 format for others to investigate PAF beam properties.

See ACES-269 for further info
"""

import emsim.readmat as readmat
import emsim.coord as coord
import os
import glob
import numpy as np
import h5py

ROLL_ANGLE = 0  # roll angle of ASKAP antenna in degrees

# freq_ghz = 0.725

root_path = '/data/chi139/askap-paf-model/Struts2/Struts2'

freq_dirs = glob.glob(os.path.join(root_path, '*GHz'))
freq_list = sorted([float(x.split('/')[-1].split('GHz')[0]) for x in freq_dirs])
n_freq = len(freq_list)

update_keys = ['xyz', 'rtp', 'a_r', 'a_t', 'a_p', 'domega', 'e_field', 'frequency']
points_list = set()

with h5py.File('/data/chi139/askap-paf-model/aces/askap_paf_patterns_Struts2.hdf5', 'w') as h5file:
    # noinspection PyTypeChecker
    h5file.attrs['roll_angle_value'] = np.radians(ROLL_ANGLE)
    h5file.attrs['roll_angle_unit'] = 'radians'
    h5file.attrs['frequency_unit'] = 'hertz'
    h5file.create_dataset('frequency', (n_freq,), np.float64, maxshape=(None,))
    h5file.create_dataset('points_per_region', (n_freq,), np.int, maxshape=(None,))
    h5file.create_dataset('frequency_number_in_group', (n_freq,), np.int, maxshape=(None,))

    for i_freq, freq_ghz in enumerate(freq_list):
        # noinspection PyTypeChecker
        frequency = np.array(freq_ghz*1e9)
        h5file['frequency'][i_freq] = frequency

        print('+++ {:g} GHz ({}/{}) +++'.format(freq_ghz, i_freq + 1, n_freq))

        data_path = os.path.join(root_path, '{:g}GHz/R'.format(freq_ghz))
        regions_fname = glob.glob(os.path.join(data_path, 'SpherePatches_*deg.mat'))[0]

        dthetadeg = regions_fname.split('_')[-1].split('deg.')[0]
        field_fname = os.path.join(data_path, 'SpherePatches_{}deg_Loaded'.format(dthetadeg))

        # get calculation point directions xyz and differential solid angle domega from file
        xyz, domega = readmat.grid_from_file(regions_fname)

        # move the main beam direction from positive z to positive x to comply with common practise for measuring the
        # pattern of a high-gain antenna in the forward direction described in IEEE Std 149-1979
        xyz = coord.move_main_beam_from_z_to_x(xyz)

        # rotate the coordinate system to align with specified roll angle of ASKAP antenna
        # noinspection PyTypeChecker
        xyz = coord.rotate_xyz_about_x(xyz, np.radians(ROLL_ANGLE+45.))

        # now we need the spherical basis for each calculation direction so we can describe polarisation in the
        # appropriate tangent plane
        a_r, a_t, a_p = coord.cartesian_to_spherical_basis(xyz)

        # calculate the spherical coordinates for each calculation direction
        rtp = coord.xyz_to_rtp(xyz)

        # read the electric field from file
        e_field = readmat.field_from_file(field_fname)

        # rotated the coordinates of the e_field vectors to also have the main beam in the positive x direction
        e_field = coord.rotate_field_z_to_x(e_field)

        print (e_field.shape)

        # project the electric field onto the spherical coordinate basis vectors
        e_field = coord.project_field(e_field, a_r, a_t, a_p)

        # are the fields single ended
        is_single = False
        if e_field.shape[0] == 376:
            is_single = True

        if is_single:
            print ('SINGLE ENDED')
            print ('PORTS =', e_field.shape[0])
            # convert single ended fields to differential fields
            e_field = (e_field[188:, :, :] - e_field[:188, :, :])/2.
        else:
            print ('DIFFERENTIAL')
            print ('PORTS =', e_field.shape[0])

        # put the field in the
        e_field = e_field[readmat.portmap, :, :]

        points_per_region = e_field.shape[1]

        h5file['points_per_region'][i_freq] = points_per_region

        if points_per_region not in points_list:
            for key in update_keys:
                data = vars()[key]
                full_key = '/{}/{}'.format(points_per_region, key)
                h5file.create_dataset(full_key, data=np.reshape(data, (1,) + data.shape), maxshape=(None,) + data.shape)
                points_list.add(points_per_region)
                h5file['frequency_number_in_group'][i_freq] = 0
        else:
            # extend and populate
            for key in update_keys:
                full_key = '/{}/{}'.format(points_per_region, key)
                h5file[full_key].resize(h5file[full_key].shape[0] + 1, 0)
                h5file[full_key][-1, ...] = vars()[key]
                h5file['frequency_number_in_group'][i_freq] = h5file[full_key].shape[0] - 1
        del e_field
