import os
os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE'
import h5py
import hdf5plugin  # REQUIRED to have the compression plugins available
import numpy as np
import pandas as pd
import glob
import os
import matplotlib.pyplot as plt
from hera_cal import io, utils
from hera_qm import ant_class
from uvtools.plot import plot_antpos, plot_antclass
%matplotlib inline
from IPython.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))
_ = np.seterr(all='ignore')  # get rid of red warnings
%config InlineBackend.figure_format = 'retina'

# Parse settings from environment
ANT_CLASS_FOLDER = os.environ.get("ANT_CLASS_FOLDER", "./")
SUM_FILE = os.environ.get("SUM_FILE", None)
# ANT_CLASS_FOLDER = "/mnt/sn1/2460330"
# SUM_FILE = "/mnt/sn1/2460330/zen.2460330.25463.sum.uvh5"
OC_SKIP_OUTRIGGERS = os.environ.get("OC_SKIP_OUTRIGGERS", "TRUE").upper() == "TRUE"

for param in ['ANT_CLASS_FOLDER', 'SUM_FILE', 'OC_SKIP_OUTRIGGERS']:
    print(f"{param} = '{eval(param)}'")

ANT_CLASS_FOLDER = '/mnt/sn1/data1/2460743'
SUM_FILE = '/mnt/sn1/data1/2460743/zen.2460743.45941.sum.uvh5'
OC_SKIP_OUTRIGGERS = 'True'

if SUM_FILE is not None:
    from astropy.time import Time, TimeDelta
    utc = Time(float(SUM_FILE.split('zen.')[-1].split('.sum.uvh5')[0]), format='jd').datetime
    print(f'Date: {utc.month}-{utc.day}-{utc.year}')

Date: 3-8-2025

# set thresholds for fraction of the day
overall_thresh = .1
all_zero_thresh = .1
eo_zeros_thresh = .1
xengine_diff_thresh = .1
cross_pol_thresh = .5
bad_fem_thresh = .1
high_power_thresh = .1
low_power_thresh = .1
low_corr_thresh = .1
bad_shape_thresh = .5
excess_rfi_thresh = .1
chisq_thresh = .25

# Load csvs
csv_files = sorted(glob.glob(os.path.join(ANT_CLASS_FOLDER, '*.ant_class.csv')))
jds = [float(f.split('/')[-1].split('zen.')[-1].split('.sum')[0]) for f in csv_files]
tables = [pd.read_csv(f).dropna(axis=0, how='all') for f in csv_files]
table_cols = tables[0].columns[1::2]
class_cols = tables[0].columns[2::2]
print(f'Found {len(csv_files)} csv files starting with {csv_files[0]}')

Found 1842 csv files starting with /mnt/sn1/data1/2460743/zen.2460743.25249.sum.ant_class.csv

# parse ant_strings
ap_strs = np.array(tables[0]['Antenna'])
ants = sorted(set(int(a[:-1]) for a in ap_strs))
translator = ''.maketrans('e', 'n') | ''.maketrans('n', 'e')

# get node numbers
node_dict = {ant: 'Unknown' for ant in ants}
try:
    from hera_mc import cm_hookup
    hookup = cm_hookup.get_hookup('default')
    for ant_name in hookup:
        ant = int("".join(filter(str.isdigit, ant_name)))
        if ant in node_dict:
            if hookup[ant_name].get_part_from_type('node')['E<ground'] is not None:
                node_dict[ant] = int(hookup[ant_name].get_part_from_type('node')['E<ground'][1:])
except:
    pass
nodes = sorted(set(node_dict.values()))

def classification_array(col):
    class_array = np.vstack([t[col] for t in tables])
    class_array[class_array == 'good'] = 1.7
    class_array[class_array == 'suspect'] = 1
    class_array[class_array == 'bad'] = 0
    return class_array.astype(float)

if SUM_FILE is not None:
    hd = io.HERADataFastReader(SUM_FILE)
    ap_tuples = [(int(ap[:-1]), {'e': 'Jee', 'n': 'Jnn'}[ap[-1]]) for ap in ap_strs]
    bad_bools = np.mean(classification_array('Antenna Class') == 0, axis=0) > overall_thresh
    bad_aps = [ap_tuples[i] for i in np.arange(len(ap_tuples))[bad_bools]]
    suspect_bools = np.mean(classification_array('Antenna Class') == 1, axis=0) > overall_thresh
    suspect_aps = [ap_tuples[i] for i in np.arange(len(ap_tuples))[suspect_bools] if ap_tuples[i] not in bad_aps]
    good_aps = [ap for ap in ap_tuples if ap not in bad_aps and ap not in suspect_aps]
    overall_class = ant_class.AntennaClassification(bad=bad_aps, suspect=suspect_aps, good=good_aps)
    autos, _, _ = hd.read(bls=[bl for bl in hd.bls if utils.split_bl(bl)[0] == utils.split_bl(bl)[1]], read_flags=False, read_nsamples=False)
    avg_unflagged_auto = {}
    for pol in ['ee', 'nn']:
        unflagged_autos = [autos[bl] for bl in autos if bl[2] == pol and overall_class[utils.split_bl(bl)[0]] != 'bad']
        if len(unflagged_autos) > 0:
            avg_unflagged_auto[pol] = np.mean(unflagged_autos, axis=(0, 1))
        else:
            avg_unflagged_auto[pol] = np.zeros(len(hd.freqs), dtype=complex)

def print_issue_summary(bad_ant_strs, title, notes='', plot=False):
    '''Print report for list of bad antenna polarizations strings'''
    unique_bad_antnums = [int(ap[:-1]) for ap in bad_ant_strs]
    display(HTML(f'<h2>{title}: ({len(bad_ant_strs)} antpols across {len(set([ba[:-1] for ba in bad_ant_strs]))} antennas)</h2>'))
    if len(notes) > 0:
        display(HTML(f'<h4>{notes}</h4>'))
    if len(bad_ant_strs) > 0:
        print(f'All Bad Antpols: {", ".join(bad_ant_strs)}\n')
    for node in nodes:
        if np.any([node == node_dict[a] for a in unique_bad_antnums]):
            aps = [ap for ap in bad_ant_strs if node_dict[int(ap[:-1])] == node]
            whole_ants = [str(wa) for wa in set([int(ap[:-1]) for ap in aps if ap.translate(translator) in bad_ant_strs])]
            single_pols =  [ap for ap in aps if ap.translate(translator) not in bad_ant_strs]
            print(f'Node {node}:')
            print(f'\tAntpols ({len(aps)} total): {", ".join(aps)}')
            print(f'\tWhole Ants ({len(whole_ants)} total): {", ".join(whole_ants)}')
            print(f'\tSingle Pols ({len(single_pols)} total): {", ".join(single_pols)}')
            if plot and SUM_FILE is not None:
                fig, axes = plt.subplots(1, 2, figsize=(12,4), dpi=70, sharey=True, gridspec_kw={'wspace': 0})
                for ax, pol in zip(axes, ['ee', 'nn']):                    
                    ax.semilogy(autos.freqs / 1e6, avg_unflagged_auto[pol], 'k--', label='Average\nUnflagged\nAuto')
                    for ap in aps:
                        ant = int(ap[:-1]), utils.comply_pol(ap[-1])
                        auto_bl = utils.join_bl(ant, ant)
                        if auto_bl[2] == pol:
                            ax.semilogy(autos.freqs / 1e6, np.mean(autos[auto_bl], axis=0), label=ap)
                    ax.legend()
                    ax.set_xlim([40, 299])
                    ax.set_title(f'{title} on Node {node} ({pol}-antennas)')
                    ax.set_xlabel('Frequency (MHz)')
                axes[0].set_ylabel('Single File Raw Autocorrelation')
                plt.tight_layout()
                plt.show()

# precompute various helpful quantities
all_slopes = np.vstack([t['Autocorr Slope'] for t in tables])
median_slope = np.median(all_slopes)
bad_slopes = np.vstack([t['Autocorr Slope Class'] for t in tables]) == 'bad'
suspect_slopes = np.vstack([t['Autocorr Slope Class'] for t in tables]) == 'suspect'
bad_shapes = np.vstack([t['Autocorr Shape Class'] for t in tables]) == 'bad'
suspect_shapes = np.vstack([t['Autocorr Shape Class'] for t in tables]) == 'suspect'
all_powers = np.vstack([t['Autocorr Power'] for t in tables])
median_power = np.median(all_powers)
bad_powers = np.vstack([t['Autocorr Power Class'] for t in tables]) == 'bad'
suspect_powers = np.vstack([t['Autocorr Power Class'] for t in tables]) == 'suspect'
bad_rfi = np.vstack([t['Auto RFI RMS Class'] for t in tables]) == 'bad'
suspect_rfi = np.vstack([t['Auto RFI RMS Class'] for t in tables]) == 'suspect'

# find all zeros
all_zeros_strs = ap_strs[np.mean(np.vstack([t['Dead? Class'] for t in tables]) == 'bad', axis=0) > all_zero_thresh]

# find even/odd zeros
eo_zeros_strs = ap_strs[np.mean(np.vstack([t['Even/Odd Zeros Class'] for t in tables]) == 'bad', axis=0) > eo_zeros_thresh]
eo_zeros_strs = [ap for ap in eo_zeros_strs if ap not in all_zeros_strs]

# find cross-polarized antennas
cross_pol_strs = ap_strs[np.mean(np.vstack([t['Cross-Polarized Class'] for t in tables]) == 'bad', axis=0) > cross_pol_thresh]
cross_pol_strs = [ap for ap in cross_pol_strs if ap not in all_zeros_strs]

# find FEM power issues: must be low power, high slope, and bad or suspect in power, slope, rfi, and shape
fem_off_prod = (bad_powers + .5 * suspect_powers) * (bad_slopes + .5 * suspect_slopes)
fem_off_prod *= (bad_rfi + .5 * suspect_rfi) * (bad_shapes + .5 * suspect_shapes)
fem_off_strs = ap_strs[np.mean(fem_off_prod * (all_powers < median_power) * (all_slopes > median_slope), axis=0) > .1]

# find high power issues
high_power_strs = ap_strs[np.mean(bad_powers & (all_powers > median_power), axis=0) > high_power_thresh]

# find other low power issues
low_power_strs = ap_strs[np.mean(bad_powers & (all_powers < median_power), axis=0) > low_power_thresh]
low_power_strs = [ap for ap in low_power_strs if ap not in all_zeros_strs and ap not in fem_off_strs]

# find low correlation (but not low power)
low_corr_strs = ap_strs[np.mean(np.vstack([t['Low Correlation Class'] for t in tables]) == 'bad', axis=0) > low_corr_thresh]
low_corr_strs = [ap for ap in low_corr_strs if ap not in (set(low_power_strs) | set(all_zeros_strs) | set(fem_off_strs))]

# find bad bandpasses
bad_bandpass_strs = ap_strs[np.mean(bad_shapes, axis=0) > bad_shape_thresh]
bad_bandpass_strs = [ap for ap in bad_bandpass_strs if ap not in (set(low_power_strs) | set(all_zeros_strs) | set(high_power_strs) | set(fem_off_strs))]

# find antennas with excess RFI
excess_rfi_strs = ap_strs[np.mean(np.vstack([t['Auto RFI RMS Class'] for t in tables]) == 'bad', axis=0) > excess_rfi_thresh]
excess_rfi_strs = [ap for ap in excess_rfi_strs if ap not in (set(low_power_strs) | set(all_zeros_strs) |  set(fem_off_strs) |
                                                              set(bad_bandpass_strs) | set(high_power_strs))]

# find bad x-engine diffs
xengine_diff_strs = ap_strs[np.mean(np.vstack([t['Bad Diff X-Engines Class'] for t in tables]) == 'bad', axis=0) > xengine_diff_thresh]
xengine_diff_strs = [ap for ap in xengine_diff_strs if ap not in (set(bad_bandpass_strs) | set(low_power_strs) | set(excess_rfi_strs) | set(low_corr_strs) |
                                                                  set(all_zeros_strs) | set(high_power_strs) | set(fem_off_strs) | set(eo_zeros_strs))]

# find antennas with high redcal chi^2
chisq_strs = ap_strs[np.mean(np.vstack([t['Redcal chi^2 Class'] for t in tables]) == 'bad', axis=0) > chisq_thresh]
chisq_strs = [ap for ap in chisq_strs if ap not in (set(bad_bandpass_strs) | set(low_power_strs) | set(excess_rfi_strs) | set(low_corr_strs) |
                                                    set(all_zeros_strs) | set(high_power_strs) | set(fem_off_strs) | set(eo_zeros_strs) | set(xengine_diff_strs))]
if OC_SKIP_OUTRIGGERS:
    chisq_strs = [ap for ap in chisq_strs if int(ap[:-1]) < 320]

# collect all results
to_print = [(all_zeros_strs, 'All-Zeros', 'These antennas have visibilities that are more than half zeros.'),
            (eo_zeros_strs, 'Excess Zeros in Either Even or Odd Spectra', 
             'These antennas are showing evidence of packet loss or X-engine failure.', True),
            (xengine_diff_strs, 'Excess Power in X-Engine Diffs', 
             'These antennas are showing evidence of mis-written packets in either the evens or the odds.', True),            
            (cross_pol_strs, 'Cross-Polarized', 'These antennas have their east and north cables swapped.'),
            (fem_off_strs, 'Likely FEM Power Issue', 'These antennas have low power and anomolously high slopes.', True),
            (high_power_strs, 'High Power', 'These antennas have high median power.', True),
            (low_power_strs, 'Other Low Power Issues', 'These antennas have low power, but are not all-zeros and not FEM off.', True),
            (low_corr_strs, 'Low Correlation, But Not Low Power', 'These antennas are low correlation, but their autocorrelation power levels look OK.'),
            (bad_bandpass_strs, 'Bad Bandpass Shapes, But Not Bad Power', 
             'These antennas have unusual bandpass shapes, but are not all-zeros, high power, low power, or FEM off.', True),
            (excess_rfi_strs, 'Excess RFI', 'These antennas have excess RMS after DPSS filtering (likely RFI), but not low or high power or a bad bandpass.', True),
            (chisq_strs, 'Redcal chi^2', 'These antennas have been idenfied as not redundantly calibrating well, even after passing the above checks.')]

def print_high_level_summary():
    for tp in sorted(to_print, key=lambda x: len(x[0]), reverse=True):
        print(f'{len(tp[0])} antpols (on {len(set([ap[:-1] for ap in tp[0]]))} antennas) frequently flagged for {tp[1]}.')
        
def print_all_issue_summaries():
    for tp in to_print:
        print_issue_summary(*tp)

print_high_level_summary()

370 antpols (on 230 antennas) frequently flagged for Excess Power in X-Engine Diffs.
56 antpols (on 49 antennas) frequently flagged for Excess RFI.
45 antpols (on 41 antennas) frequently flagged for Redcal chi^2.
20 antpols (on 18 antennas) frequently flagged for Likely FEM Power Issue.
11 antpols (on 11 antennas) frequently flagged for Bad Bandpass Shapes, But Not Bad Power.
8 antpols (on 7 antennas) frequently flagged for Other Low Power Issues.
6 antpols (on 3 antennas) frequently flagged for All-Zeros.
5 antpols (on 4 antennas) frequently flagged for High Power.
5 antpols (on 5 antennas) frequently flagged for Low Correlation, But Not Low Power.
0 antpols (on 0 antennas) frequently flagged for Excess Zeros in Either Even or Odd Spectra.
0 antpols (on 0 antennas) frequently flagged for Cross-Polarized.

print_all_issue_summaries()

All Bad Antpols: 218e, 218n, 233e, 233n, 234e, 234n

Node 17:
	Antpols (6 total): 218e, 218n, 233e, 233n, 234e, 234n
	Whole Ants (3 total): 233, 218, 234
	Single Pols (0 total):

All Bad Antpols: 3e, 3n, 4n, 5e, 5n, 7n, 9e, 9n, 10e, 15e, 15n, 16e, 16n, 17e, 19e, 19n, 20e, 21n, 22e, 28n, 30e, 30n, 31e, 32e, 33e, 34n, 35e, 36e, 36n, 38e, 38n, 40e, 41e, 41n, 43e, 43n, 44n, 45e, 46e, 46n, 47n, 48e, 48n, 49e, 49n, 50e, 50n, 51n, 52e, 52n, 53e, 53n, 54e, 54n, 55n, 56e, 56n, 57e, 57n, 58e, 58n, 59e, 60e, 60n, 61e, 61n, 63e, 64e, 64n, 65n, 66e, 67e, 68n, 69e, 69n, 70e, 71e, 72e, 73e, 79e, 79n, 80e, 80n, 81e, 82e, 83n, 84e, 84n, 85e, 86e, 86n, 87e, 87n, 88e, 88n, 89e, 89n, 91e, 91n, 92e, 93e, 93n, 94e, 94n, 95e, 96e, 96n, 97e, 98e, 99e, 100n, 101e, 101n, 102e, 103n, 105n, 106n, 108n, 109e, 110e, 110n, 111e, 111n, 112n, 113e, 113n, 114e, 114n, 115e, 115n, 116e, 116n, 118e, 118n, 119e, 119n, 120n, 122e, 122n, 123e, 123n, 124e, 124n, 126e, 126n, 127e, 127n, 128e, 128n, 129e, 129n, 130e, 130n, 131e, 132e, 132n, 133e, 133n, 134e, 134n, 135n, 136e, 137n, 138e, 138n, 139n, 140e, 140n, 141e, 141n, 142e, 142n, 143e, 144e, 144n, 145e, 145n, 146e, 147e, 147n, 148e, 149e, 149n, 150e, 150n, 151n, 152e, 152n, 153e, 153n, 154e, 154n, 155e, 155n, 156e, 156n, 157e, 157n, 158e, 158n, 159e, 159n, 160e, 161e, 162e, 162n, 163e, 163n, 164e, 164n, 165e, 165n, 166e, 166n, 167e, 167n, 168e, 168n, 169e, 169n, 170n, 171e, 172e, 172n, 173e, 173n, 174e, 175e, 175n, 176e, 176n, 177e, 177n, 178e, 179e, 179n, 181e, 181n, 182e, 182n, 183e, 183n, 184e, 184n, 185e, 185n, 186e, 186n, 187e, 187n, 188e, 190e, 190n, 191e, 191n, 192e, 192n, 193e, 193n, 194e, 194n, 195e, 195n, 196e, 196n, 197e, 197n, 198n, 201e, 201n, 203e, 203n, 204e, 204n, 205e, 205n, 206e, 206n, 207e, 207n, 208n, 209e, 210e, 210n, 211n, 212e, 213n, 214e, 214n, 217n, 219e, 219n, 220e, 220n, 221e, 221n, 222e, 222n, 223e, 223n, 224n, 225e, 225n, 226e, 227e, 227n, 228e, 228n, 229e, 229n, 232n, 235e, 235n, 237e, 237n, 239n, 240e, 240n, 241e, 241n, 243n, 244e, 244n, 245e, 245n, 246n, 250n, 251n, 252e, 252n, 253e, 255e, 256e, 256n, 261e, 261n, 266e, 266n, 267e, 267n, 268e, 269e, 269n, 270e, 272n, 281e, 281n, 282e, 283e, 283n, 285e, 285n, 295e, 295n, 320n, 321e, 321n, 322n, 323e, 323n, 325e, 325n, 327e, 327n, 328n, 331n, 333n, 336e, 336n, 340n

Node 1:
	Antpols (13 total): 3e, 3n, 4n, 5e, 5n, 15e, 15n, 16e, 16n, 17e, 28n, 30e, 30n
	Whole Ants (5 total): 3, 5, 15, 16, 30
	Single Pols (3 total): 4n, 17e, 28n

Casting complex values to real discards the imaginary part
Casting complex values to real discards the imaginary part

Node 2:
	Antpols (15 total): 7n, 9e, 9n, 10e, 19e, 19n, 20e, 21n, 31e, 32e, 33e, 321e, 321n, 323e, 323n
	Whole Ants (4 total): 321, 9, 323, 19
	Single Pols (7 total): 7n, 10e, 20e, 21n, 31e, 32e, 33e

Node 3:
	Antpols (16 total): 36e, 36n, 38e, 38n, 50e, 50n, 51n, 52e, 52n, 53e, 53n, 65n, 66e, 67e, 68n, 320n
	Whole Ants (5 total): 36, 38, 50, 52, 53
	Single Pols (6 total): 51n, 65n, 66e, 67e, 68n, 320n

Node 4:
	Antpols (15 total): 40e, 41e, 41n, 54e, 54n, 55n, 56e, 56n, 57e, 57n, 69e, 69n, 70e, 71e, 72e
	Whole Ants (5 total): 69, 41, 54, 56, 57
	Single Pols (5 total): 40e, 55n, 70e, 71e, 72e

Node 5:
	Antpols (13 total): 43e, 43n, 44n, 45e, 46e, 46n, 58e, 58n, 59e, 60e, 60n, 73e, 322n
	Whole Ants (4 total): 58, 43, 60, 46
	Single Pols (5 total): 44n, 45e, 59e, 73e, 322n

Node 6:
	Antpols (13 total): 22e, 34n, 35e, 47n, 48e, 48n, 49e, 49n, 61e, 61n, 63e, 64e, 64n
	Whole Ants (4 total): 48, 49, 64, 61
	Single Pols (5 total): 22e, 34n, 35e, 47n, 63e

Node 7:
	Antpols (15 total): 81e, 82e, 83n, 98e, 99e, 100n, 116e, 116n, 118e, 118n, 119e, 119n, 137n, 138e, 138n
	Whole Ants (4 total): 138, 116, 118, 119
	Single Pols (7 total): 81e, 82e, 83n, 98e, 99e, 100n, 137n

Node 8:
	Antpols (16 total): 84e, 84n, 85e, 86e, 86n, 87e, 87n, 101e, 101n, 102e, 103n, 120n, 122e, 122n, 123e, 123n
	Whole Ants (6 total): 101, 84, 86, 87, 122, 123
	Single Pols (4 total): 85e, 102e, 103n, 120n

Antenna Classification Daily Summary¶

• Summary of Per Antenna Issues ¶

• Figure 1: Per File Overall Antenna Classification Summary ¶

• Figure 2: Per Classifier Antenna Flagging Summary ¶

• Figure 3: Array Visualization of Overall Daily Classification ¶

Settings¶

Load classifications and other metadata¶

Figure out and summarize per-antenna issues¶

Summary of Per-Antenna Issues¶

All-Zeros: (6 antpols across 3 antennas)

These antennas have visibilities that are more than half zeros.

Excess Zeros in Either Even or Odd Spectra: (0 antpols across 0 antennas)

These antennas are showing evidence of packet loss or X-engine failure.

Excess Power in X-Engine Diffs: (370 antpols across 230 antennas)

These antennas are showing evidence of mis-written packets in either the evens or the odds.

Antenna Classification Daily Summary¶

• Summary of Per Antenna Issues¶

• Figure 1: Per File Overall Antenna Classification Summary¶

• Figure 2: Per Classifier Antenna Flagging Summary¶

• Figure 3: Array Visualization of Overall Daily Classification¶

Settings¶

Load classifications and other metadata¶

Figure out and summarize per-antenna issues¶

Summary of Per-Antenna Issues¶

All-Zeros: (6 antpols across 3 antennas)

These antennas have visibilities that are more than half zeros.

Excess Zeros in Either Even or Odd Spectra: (0 antpols across 0 antennas)

These antennas are showing evidence of packet loss or X-engine failure.

Excess Power in X-Engine Diffs: (370 antpols across 230 antennas)

These antennas are showing evidence of mis-written packets in either the evens or the odds.

• Summary of Per Antenna Issues ¶

• Figure 1: Per File Overall Antenna Classification Summary ¶

• Figure 2: Per Classifier Antenna Flagging Summary ¶

• Figure 3: Array Visualization of Overall Daily Classification ¶