import time
tstart = time.time()

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 copy
import matplotlib
import matplotlib.pyplot as plt
from pyuvdata import UVFlag, UVData, UVCal
from hera_cal import io, utils
from hera_qm.time_series_metrics import true_stretches, impose_max_flag_gap, metric_convolution_flagging
from hera_qm.metrics_io import read_a_priori_int_flags
from uvtools.plot import plot_antpos, plot_antclass
from scipy.ndimage import convolve
%matplotlib inline
from IPython.display import display, HTML

# get files
SUM_FILE = os.environ.get("SUM_FILE", None)
# SUM_FILE = '/users/jsdillon/lustre/H6C/abscal/2459853/zen.2459853.25518.sum.uvh5'
SUM_SUFFIX = os.environ.get("SUM_SUFFIX", 'sum.uvh5')
CAL_SUFFIX = os.environ.get("CAL_SUFFIX", 'sum.omni.calfits')
ANT_CLASS_SUFFIX = os.environ.get("ANT_CLASS_SUFFIX", 'ant_class.csv')
OUT_FLAG_SUFFIX = os.environ.get("OUT_FLAG_SUFFIX", 'sum.antenna_flags.h5')
APRIORI_YAML_PATH = os.environ.get("APRIORI_YAML_PATH", None)

# Parameters for harmonizing partially-flagged antennas
SMOOTHING_SCALE_NFILES = float(os.environ.get("SMOOTHING_SCALE_NFILES", 30))
MAX_FLAG_GAP_NFILES = int(os.environ.get("MAX_FLAG_GAP_NFILES", 30))

# Max flag fractions (before just flagging the whole antenna)
POWER_MAX_FLAG_FRAC = float(os.environ.get("POWER_MAX_FLAG_FRAC", .5))
AUTO_POWER_MAX_FLAG_FRAC = float(os.environ.get("AUTO_POWER_MAX_FLAG_FRAC", .5))
AUTO_SHAPE_MAX_FLAG_FRAC = float(os.environ.get("AUTO_SHAPE_MAX_FLAG_FRAC", .25))
AUTO_SLOPE_MAX_FLAG_FRAC = float(os.environ.get("AUTO_SLOPE_MAX_FLAG_FRAC", .25))
AUTO_RFI_MAX_FLAG_FRAC = float(os.environ.get("AUTO_RFI_MAX_FLAG_FRAC", .25))
CHISQ_MAX_FLAG_FRAC = float(os.environ.get("CHISQ_MAX_FLAG_FRAC", .5))
OVERALL_MAX_FLAG_FRAC = float(os.environ.get("OVERALL_MAX_FLAG_FRAC", .5))

for setting in ['SUM_FILE', 'SUM_SUFFIX', 'CAL_SUFFIX', 'ANT_CLASS_SUFFIX', 'OUT_FLAG_SUFFIX', 'APRIORI_YAML_PATH',
                'SMOOTHING_SCALE_NFILES', 'MAX_FLAG_GAP_NFILES', 'POWER_MAX_FLAG_FRAC', 'AUTO_POWER_MAX_FLAG_FRAC', 
                'AUTO_SHAPE_MAX_FLAG_FRAC', 'AUTO_SLOPE_MAX_FLAG_FRAC', 'AUTO_RFI_MAX_FLAG_FRAC', 
                'CHISQ_MAX_FLAG_FRAC', 'OVERALL_MAX_FLAG_FRAC']:
        print(f'{setting} = {eval(setting)}')

SUM_FILE = /mnt/sn1/data2/2460548/zen.2460548.16914.sum.uvh5
SUM_SUFFIX = sum.uvh5
CAL_SUFFIX = sum.omni.calfits
ANT_CLASS_SUFFIX = sum.ant_class.csv
OUT_FLAG_SUFFIX = sum.antenna_flags.h5
APRIORI_YAML_PATH = /mnt/sn1/data2/2460548/2460548_apriori_flags.yaml
SMOOTHING_SCALE_NFILES = 30.0
MAX_FLAG_GAP_NFILES = 30
POWER_MAX_FLAG_FRAC = 30.0
AUTO_POWER_MAX_FLAG_FRAC = 0.5
AUTO_SHAPE_MAX_FLAG_FRAC = 0.25
AUTO_SLOPE_MAX_FLAG_FRAC = 0.25
AUTO_RFI_MAX_FLAG_FRAC = 0.25
CHISQ_MAX_FLAG_FRAC = 0.75
OVERALL_MAX_FLAG_FRAC = 0.75

# ant_metrics bounds for low correlation / dead antennas
am_corr_bad = (0, float(os.environ.get("AM_CORR_BAD", 0.3)))
am_corr_suspect = (float(os.environ.get("AM_CORR_BAD", 0.3)), float(os.environ.get("AM_CORR_SUSPECT", 0.5)))

# ant_metrics bounds for cross-polarized antennas
am_xpol_bad = (-1, float(os.environ.get("AM_XPOL_BAD", -0.1)))
am_xpol_suspect = (float(os.environ.get("AM_XPOL_BAD", -0.1)), float(os.environ.get("AM_XPOL_SUSPECT", 0)))

# bounds on solar altitude (in degrees)
good_solar_altitude = (-90, float(os.environ.get("SUSPECT_SOLAR_ALTITUDE", 0)))
suspect_solar_altitude = (float(os.environ.get("SUSPECT_SOLAR_ALTITUDE", 0)), 90)

# bounds on zeros in spectra
good_zeros_per_eo_spectrum = (0, int(os.environ.get("MAX_ZEROS_PER_EO_SPEC_GOOD", 2)))
suspect_zeros_per_eo_spectrum = (0, int(os.environ.get("MAX_ZEROS_PER_EO_SPEC_SUSPECT", 8)))

# bounds on autocorrelation power
auto_power_good = (float(os.environ.get("AUTO_POWER_GOOD_LOW", 5)), float(os.environ.get("AUTO_POWER_GOOD_HIGH", 30)))
auto_power_suspect = (float(os.environ.get("AUTO_POWER_SUSPECT_LOW", 1)), float(os.environ.get("AUTO_POWER_SUSPECT_HIGH", 60)))

# bounds on autocorrelation slope
auto_slope_good = (float(os.environ.get("AUTO_SLOPE_GOOD_LOW", -0.4)), float(os.environ.get("AUTO_SLOPE_GOOD_HIGH", 0.4)))
auto_slope_suspect = (float(os.environ.get("AUTO_SLOPE_SUSPECT_LOW", -0.6)), float(os.environ.get("AUTO_SLOPE_SUSPECT_HIGH", 0.6)))

# bounds on autocorrelation RFI
auto_rfi_good = (0, float(os.environ.get("AUTO_RFI_GOOD", 1.5)))
auto_rfi_suspect = (0, float(os.environ.get("AUTO_RFI_SUSPECT", 2.0)))

# bounds on autocorrelation shape
auto_shape_good = (0, float(os.environ.get("AUTO_SHAPE_GOOD", 0.1)))
auto_shape_suspect = (0, float(os.environ.get("AUTO_SHAPE_SUSPECT", 0.2)))

# bounds on chi^2 per antenna in omnical
oc_cspa_good = (0, float(os.environ.get("OC_CSPA_GOOD", 2)))
oc_cspa_suspect = (0, float(os.environ.get("OC_CSPA_SUSPECT", 3)))

OC_SKIP_OUTRIGGERS = os.environ.get("OC_SKIP_OUTRIGGERS", "TRUE").upper() == "TRUE"

for bound in ['am_corr_bad', 'am_corr_suspect', 'am_xpol_bad', 'am_xpol_suspect', 
              'good_solar_altitude', 'suspect_solar_altitude',
              'good_zeros_per_eo_spectrum', 'suspect_zeros_per_eo_spectrum',
              'auto_power_good', 'auto_power_suspect', 'auto_slope_good', 'auto_slope_suspect',
              'auto_rfi_good', 'auto_rfi_suspect', 'auto_shape_good', 'auto_shape_suspect',
              'oc_cspa_good', 'oc_cspa_suspect', 'OC_SKIP_OUTRIGGERS']:
    print(f'{bound} = {eval(bound)}')

am_corr_bad = (0, 0.2)
am_corr_suspect = (0.2, 0.4)
am_xpol_bad = (-1, -0.1)
am_xpol_suspect = (-0.1, 0.0)
good_solar_altitude = (-90, 0.0)
suspect_solar_altitude = (0.0, 90)
good_zeros_per_eo_spectrum = (0, 2)
suspect_zeros_per_eo_spectrum = (0, 8)
auto_power_good = (5.0, 30.0)
auto_power_suspect = (1.0, 60.0)
auto_slope_good = (-0.4, 0.4)
auto_slope_suspect = (-0.6, 0.6)
auto_rfi_good = (0, 1.5)
auto_rfi_suspect = (0, 2.0)
auto_shape_good = (0, 0.1)
auto_shape_suspect = (0, 0.2)
oc_cspa_good = (0, 2.0)
oc_cspa_suspect = (0, 3.0)
OC_SKIP_OUTRIGGERS = True

hd = io.HERAData(SUM_FILE)

sum_glob = '.'.join(SUM_FILE.split('.')[:-3]) + '.*.' + SUM_SUFFIX
cal_files_glob = sum_glob.replace(SUM_SUFFIX, CAL_SUFFIX)
cal_files = sorted(glob.glob(cal_files_glob))
print(f'Found {len(cal_files)} *.{CAL_SUFFIX} files starting with {cal_files[0]}.')

Found 1571 *.sum.omni.calfits files starting with /mnt/sn1/data2/2460548/zen.2460548.16914.sum.omni.calfits.

ant_class_csvs_glob = sum_glob.replace(SUM_SUFFIX, ANT_CLASS_SUFFIX)
ant_class_csvs = sorted(glob.glob(ant_class_csvs_glob))
jds = [float(f.split('/')[-1].split('zen.')[-1].split('.sum')[0]) for f in ant_class_csvs]
frac_jds = jds - np.floor(jds[0])
Ncsvs = len(ant_class_csvs)
print(f'Found {Ncsvs} *.{ANT_CLASS_SUFFIX} files starting with {ant_class_csvs[0]}.')

Found 1571 *.sum.ant_class.csv files starting with /mnt/sn1/data2/2460548/zen.2460548.16914.sum.ant_class.csv.

# Load ant_class csvs
tables = [pd.read_csv(f).dropna(axis=0, how='all') for f in ant_class_csvs]
table_cols = tables[0].columns[1::2]
class_cols = tables[0].columns[2::2]
ap_strs = np.array(tables[0]['Antenna'])
ants = sorted(set(int(a[:-1]) for a in ap_strs))

# build up dictionaries for full night metrics and classifications
replace = {'-': np.nan, 'INF': np.inf, 'No': False, 'Yes': True}

metric_data = {tc: {} for tc in table_cols}
class_data = {tc: {} for tc in table_cols}
for tc, cc in zip(table_cols, class_cols):
    class_array = np.vstack([t[cc] for t in tables]).T
    metric_array = np.vstack([t[tc] for t in tables]).T
    for ca, ma, ap_str in zip(class_array, metric_array, ap_strs):
        class_data[tc][ap_str] = ca
        if tc == 'Antenna':
            metric_data[tc][ap_str] = ma
        else:
            metric_data[tc][ap_str] = np.array([replace[val] if val in replace else float(val) for val in ma])

original_flags = {ap_str: np.array(class_data['Antenna'][ap_str] == 'bad') for ap_str in ap_strs}
print(f'{np.mean(list(original_flags.values())):.2%} of antenna-files flagged by RTP.')

38.75% of antenna-files flagged by RTP.

class FlagHistory:
    '''Helps keep strack of why flags get applied'''
    
    def __init__(self, ap_strs, Nfiles):
        self.final_flags = {ap_str: np.zeros(Nfiles, dtype=bool) for ap_str in ap_strs}
        self.history = {}
    
    def update(self, ap_str, rtp_flags, updated_flags, rationale):
        self.history[(ap_str, rationale)] = (np.array(self.final_flags[ap_str]), np.array(rtp_flags), np.array(updated_flags))
        self.final_flags[ap_str] |= updated_flags
        
    def summarize(self, description):
        print(f'{np.mean(list(self.final_flags.values())):.2%} of antenna-files flagged after {description}.')

fh = FlagHistory(ap_strs, Ncsvs)

# STEP 1: FLAG SUN UP DATA AND OTHER A PRIORI FLAGGED TIMES
for ap_str in ap_strs:
    sun_up = metric_data['Solar Alt'][ap_str] > good_solar_altitude[1]
    fh.update(ap_str, sun_up, sun_up, 'Solar Alt')
solar_flags = np.all([metric_data['Solar Alt'][ap_str] > good_solar_altitude[1] for ap_str in ap_strs], axis=0)
fh.summarize('flagging sun-up data')
if APRIORI_YAML_PATH is not None:
    apriori_flags = np.zeros(len(jds), dtype=bool)
    apriori_flags[read_a_priori_int_flags(APRIORI_YAML_PATH, times=np.array(jds)).astype(int)] = True
    for ap_str in ap_strs:
        fh.update(ap_str, apriori_flags, apriori_flags, 'A Priori')
    fh.summarize('flagging a priori flagged times')
    apriori_flags |= solar_flags
else:
    apriori_flags = solar_flags

# STEP 2: FLAG TOTALLY DEAD ANTENNAS
for ap_str in ap_strs:
    fh.update(ap_str, metric_data['Dead?'][ap_str], metric_data['Dead?'][ap_str], 'Dead?')
fh.summarize('removing dead antennas')    

# STEP 3: FLAG OUTRIGGERS
if OC_SKIP_OUTRIGGERS:
    for ap_str in ap_strs:
        if int(ap_str[:-1]) >= 320:
            fh.update(ap_str, np.ones(Ncsvs, dtype=bool), np.ones(Ncsvs, dtype=bool), 'Outrigger')
    fh.summarize('flagging outriggers')

# STEP 4: FLAG CROSS-POLARIZED ANTENNAS
for ap_str in ap_strs:
    if np.mean(metric_data['Cross-Polarized'][ap_str]) < am_xpol_bad[1]:
        fh.update(ap_str, class_data['Cross-Polarized'][ap_str] == 'bad', np.ones(Ncsvs, dtype=bool), 'Cross-Polarized')
fh.summarize('removing cross-polarized antennas')
    
# STEP 5: FLAG POORLY-CORRELATING ANTENNAS
for ap_str in ap_strs:
    if np.mean(metric_data['Low Correlation'][ap_str]) < am_corr_bad[1]:
        fh.update(ap_str, class_data['Low Correlation'][ap_str] == 'bad', np.ones(Ncsvs, dtype=bool), 'Cross-Polarized')
fh.summarize('removing non-correlating antennas')

# STEP 6: FLAG ON AUTOCORRELATIONS
for category, ok_range, max_flag_frac in zip(['Autocorr Power', 'Autocorr Shape', 'Autocorr Slope', 'Auto RFI RMS'],
                                              [auto_power_suspect, auto_shape_good, auto_slope_good, auto_rfi_good],
                                              [AUTO_POWER_MAX_FLAG_FRAC, AUTO_SHAPE_MAX_FLAG_FRAC, AUTO_SLOPE_MAX_FLAG_FRAC, AUTO_RFI_MAX_FLAG_FRAC]):
    for ap_str in ap_strs:
        # apply RTP flags for these categories
        rtp_flags = (class_data[category][ap_str] == 'bad')
        # if not completely flagged or completely unflagged, grow flags via convolution
        if np.any(rtp_flags) and not np.all(fh.final_flags[ap_str] | rtp_flags):
            metric = metric_data[category][ap_str]    
            new_flags = metric_convolution_flagging(metric, rtp_flags, ok_range, sigma=SMOOTHING_SCALE_NFILES, max_flag_gap=MAX_FLAG_GAP_NFILES)
            # if too many times are flagged for this category, flag the whole antenna (excluding sun-up times)
            if np.mean(new_flags[~apriori_flags]) > max_flag_frac:
                new_flags[:] = True
        else:
            new_flags = rtp_flags
        fh.update(ap_str, rtp_flags, new_flags, category)
    fh.summarize(f'flagging antennas for {category}')
    
# STEP 7: FLAG FOR HIGH CHI^2
for ap_str in ap_strs:
    flagged_for_cspa = (~fh.final_flags[ap_str]) & (class_data['Even/Odd Zeros'][ap_str] != 'bad') & \
                       (class_data['Bad Diff X-Engines'][ap_str] != 'bad') & (class_data['Redcal chi^2'][ap_str] == 'bad')
    if np.any(flagged_for_cspa) and not np.all(fh.final_flags[ap_str] | flagged_for_cspa):
        cspa = metric_data['Redcal chi^2'][ap_str]
        new_flags = metric_convolution_flagging(cspa, flagged_for_cspa, oc_cspa_suspect, sigma=SMOOTHING_SCALE_NFILES, max_flag_gap=MAX_FLAG_GAP_NFILES)
        if np.mean(new_flags[~apriori_flags]) > CHISQ_MAX_FLAG_FRAC:
            new_flags[:] = True
    else:
        new_flags = flagged_for_cspa
    fh.update(ap_str, flagged_for_cspa, new_flags, 'Redcal chi^2')
fh.summarize('flagging antennas for high redcal chi^2')

# STEP 8: FLAG FOR TOO EVEN/ODD ZEROS AND EXCESS DIFF POWER (USUALLY PACKET ISSUES)
for ap_str in ap_strs:
    rtp_flags = (class_data['Even/Odd Zeros'][ap_str] == 'bad')
    fh.update(ap_str, rtp_flags, rtp_flags, 'Even/Odd Zeros')
fh.summarize('flagging antennas for excess even/odd zeros')
for ap_str in ap_strs:
    rtp_flags = (class_data['Bad Diff X-Engines'][ap_str] == 'bad')
    fh.update(ap_str, rtp_flags, rtp_flags, 'Bad Diff X-Engines')
fh.summarize('flagging antennas for excess power in specific X-engines in the diffs')


# STEP 9: FLAG ANTENNAS THAT ARE ALREADY LARGELY FLAGGED
for ap_str in ap_strs:
    new_flags = np.array(fh.final_flags[ap_str])
    impose_max_flag_gap(new_flags)
    if np.mean(new_flags[~apriori_flags]) > OVERALL_MAX_FLAG_FRAC:
        new_flags[:] = True
    fh.update(ap_str, fh.final_flags[ap_str], new_flags, 'Frequently Flagged')
fh.summarize('flagging frequently-flagged antennas')

1.53% of antenna-files flagged after flagging sun-up data.
1.53% of antenna-files flagged after flagging a priori flagged times.

10.10% of antenna-files flagged after removing dead antennas.
14.79% of antenna-files flagged after flagging outriggers.
15.96% of antenna-files flagged after removing cross-polarized antennas.
18.50% of antenna-files flagged after removing non-correlating antennas.

22.02% of antenna-files flagged after flagging antennas for Autocorr Power.

25.19% of antenna-files flagged after flagging antennas for Autocorr Shape.

26.00% of antenna-files flagged after flagging antennas for Autocorr Slope.

31.15% of antenna-files flagged after flagging antennas for Auto RFI RMS.

54.70% of antenna-files flagged after flagging antennas for high redcal chi^2.
54.70% of antenna-files flagged after flagging antennas for excess even/odd zeros.
54.75% of antenna-files flagged after flagging antennas for excess power in specific X-engines in the diffs.
54.76% of antenna-files flagged after flagging frequently-flagged antennas.

def flagging_board():
    cmap = matplotlib.colors.ListedColormap(['blue', 'black', 'red', 'orange'])
    to_plot = np.vstack([np.where(~fh.final_flags[ap_str] & ~original_flags[ap_str], 0,
                                  np.where(~fh.final_flags[ap_str] & original_flags[ap_str], -1,
                                           np.where(fh.final_flags[ap_str] & ~original_flags[ap_str], 2, 1))) for ap_str in ap_strs])

    plt.figure(figsize=(14, len(ants) / 10), dpi=100)
    plt.imshow(to_plot, aspect='auto', interpolation='none', cmap=cmap, vmin=-1.5, vmax=2.5,
               extent=[frac_jds[0], frac_jds[-1], len(ants), 0])
    plt.xlabel(f'JD - {int(jds[0])}')
    plt.yticks(ticks=np.arange(.5, len(ants)+.5), labels=[ant for ant in ants], fontsize=6)
    plt.ylabel('Antenna Number (East First, Then North)')
    plt.gca().tick_params(right=True, top=True, labelright=True, labeltop=True)
    cbar = plt.colorbar(location='top', aspect=40)
    cbar.set_ticks([-1, 0, 1, 2])
    cbar.set_ticklabels(['RTP Flag Removed', 'No Flags', 'Flagged by RTP and Here', 'Flagged Here Only'])
    plt.tight_layout()

flagging_board()

def flag_frac_array_plot():
    fig, axes = plt.subplots(1, 2, figsize=(14, 8), dpi=100, gridspec_kw={'width_ratios': [2, 1]})

    def flag_frac_panel(ax, antnums, radius=7, legend=False):

        ang_dict = {'e': (225, 405), 'n': (45, 225)}

        xpos = np.array([hd.antpos[antnum][0] for antnum in ants if antnum in antnums])
        ypos = np.array([hd.antpos[antnum][1] for antnum in ants if antnum in antnums])
        scatter = ax.scatter(xpos, ypos, c='w', s=0)
        for ap_str in ap_strs:
            antnum, pol = int(ap_str[:-1]), ap_str[-1]
            if antnum in antnums:
                ax.add_artist(matplotlib.patches.Wedge(tuple(hd.antpos[antnum][0:2]), radius, *ang_dict[pol], color='grey'))
                flag_frac = np.mean(fh.final_flags[ap_str][~solar_flags])
                if flag_frac > .05:
                    ax.add_artist(matplotlib.patches.Wedge(tuple(hd.antpos[antnum][0:2]), radius * np.sqrt(flag_frac), *ang_dict[pol], color='r'))
                ax.text(hd.antpos[antnum][0], hd.antpos[antnum][1], str(antnum), color='w',  va='center', ha='center', zorder=100)

        ax.axis('equal')
        ax.set_xlim([np.min(xpos) - radius * 2, np.max(xpos) + radius * 2])
        ax.set_ylim([np.min(ypos) - radius * 2, np.max(ypos) + radius * 2])
        ax.set_xlabel("East-West Position (meters)", size=12)
        ax.set_ylabel("North-South Position (meters)", size=12)

        if legend:
            legend_objs = []
            legend_labels = []

            legend_objs.append(matplotlib.lines.Line2D([0], [0], marker='o', color='w', markeredgecolor='grey', markerfacecolor='grey', markersize=15))
            unflagged_nights = lambda pol: np.sum([np.mean(~fh.final_flags[ap_str][~solar_flags]) for ap_str in ap_strs if ap_str[-1] == pol])
            legend_labels.append((' \u2571\n').join([f'{unflagged_nights(pol):.1f} unflagged {pol}-polarized\nantenna-nights.' for pol in ['e', 'n']]))

            legend_objs.append(matplotlib.lines.Line2D([0], [0], marker='o', color='w', markeredgecolor='red', markerfacecolor='red', markersize=15))
            unflagged_nights = lambda pol: np.sum([np.mean(fh.final_flags[ap_str][~solar_flags]) for ap_str in ap_strs if ap_str[-1] == pol])
            legend_labels.append((' \u2571\n').join([f'{unflagged_nights(pol):.1f} flagged {pol}-polarized\nantenna-nights.' for pol in ['e', 'n']]))        
            ax.legend(legend_objs, legend_labels, ncol=1, fontsize=12)

    flag_frac_panel(axes[0], [ant for ant in ants if ant < 320], radius=7)
    flag_frac_panel(axes[1], [ant for ant in ants if ant >= 320], radius=50, legend=True)
    plt.tight_layout()

flag_frac_array_plot()

def flag_summary_vs_jd():
    plt.figure(figsize=(14, 5), dpi=100,)
    plt.plot(frac_jds, np.mean(np.vstack([fh.final_flags[ap_str] for ap_str in ap_strs if ap_str[-1] == 'e']), axis=0), '.', ms=2, label='EW-Polarized')
    plt.plot(frac_jds, np.mean(np.vstack([fh.final_flags[ap_str] for ap_str in ap_strs if ap_str[-1] == 'n']), axis=0), '.', ms=2, label='NS-Polarized')
    plt.plot(frac_jds, np.mean(np.vstack([fh.final_flags[ap_str] for ap_str in ap_strs]), axis=0), 'k.', ms=3, label='All Antpols')
    plt.legend()
    plt.xlabel(f'JD - {int(np.floor(jds[0]))}')
    plt.ylabel('Fraction of Antennas Flagged')
    plt.tight_layout()

flag_summary_vs_jd()

def class_to_int(c):
    return np.where(c == 'bad', 1.7, np.where(c=='suspect', 1, 0))

def per_antenna_flag_harmonization_plots():
    # compute convolution kernel
    kernel = np.exp(-np.arange(-Ncsvs // 2, Ncsvs // 2 + 1)**2 / 2 / SMOOTHING_SCALE_NFILES**2)
    kernel /= np.sum(kernel)
    
    # JD computations
    djd = np.median(np.diff(jds))

    # loop over ant-pols
    for ap_str in ap_strs:
        # if there are new nighttime, not-apriori flags
        if np.sum(fh.final_flags[ap_str][~apriori_flags]) > np.sum(original_flags[ap_str][~apriori_flags]):
            # if the new flags aren't just because of the OVERALL_MAX_FLAG_FRAC
            if np.mean(original_flags[ap_str][~apriori_flags]) <= OVERALL_MAX_FLAG_FRAC:
                for aps, category in fh.history.keys():
                    if category == 'Frequently Flagged':
                        continue
                    if ap_str == aps:
                        previous_flags, rtp_flags, new_flags = fh.history[(aps, category)]
                        # if new flags were added for this reason
                        if np.sum(new_flags[~apriori_flags]) > np.sum(rtp_flags[~apriori_flags]):
                            plt.figure(figsize=(14, 3), dpi=100)
                            
                            # plot 
                            plt.scatter(frac_jds, metric_data[category][ap_str], c=class_to_int(class_data[category][ap_str]), s=3,
                                        vmin=0, vmax=1.7, cmap='RdYlGn_r', label='Metric/Classification')
                            plt.plot(frac_jds, convolve(metric_data[category][ap_str], kernel, mode='reflect'), 'k--', label='Smoothed Metric')
                            plt.ylabel(category)
                            plt.xlabel(f'JD - {int(np.floor(jds[0]))}')
                            plt.xlim([np.min(frac_jds) - 10 * djd, 1.2 * np.max(frac_jds) - .2 * np.min(frac_jds)])
                            
                            # Indicate flagged stretches 
                            for i, bad_stretch in enumerate(true_stretches(rtp_flags)):
                                plt.axvspan(frac_jds[bad_stretch.start] - djd / 2, frac_jds[bad_stretch.stop - 1] + djd / 2, zorder=0, color='red', alpha=.75, lw=0,
                                            label=(f'RTP Flags:\n{np.mean(rtp_flags[~solar_flags]):.2%} of night' if i == 0 else None))
                            for i, bad_stretch in enumerate(true_stretches(new_flags & ~apriori_flags)):
                                plt.axvspan(frac_jds[bad_stretch.start] - djd / 2, frac_jds[bad_stretch.stop - 1] + djd / 2, zorder=0, color='orange', alpha=.75, lw=0,
                                            label=(f'Harmonized Flags:\n{np.mean((new_flags & ~apriori_flags)[~solar_flags]):.2%} of night' if i == 0 else None))
                            for i, bad_stretch in enumerate(true_stretches((fh.final_flags[ap_str] & ~new_flags) | apriori_flags)):
                                plt.axvspan(frac_jds[bad_stretch.start] - djd / 2, frac_jds[bad_stretch.stop - 1] + djd / 2, zorder=0, color='purple', alpha=.75, lw=0,
                                            label=(f'All Final Flags:\n{np.mean(fh.final_flags[ap_str][~solar_flags]):.2%} of night' if i == 0 else None))                            

                            plt.legend(title=f'{ap_str}: {category}', loc='upper right')
                            plt.tight_layout()
                            plt.show()

per_antenna_flag_harmonization_plots()

add_to_history = 'Produced by full_day_antenna_flagging notebook with the following environment:\n' + '=' * 65 + '\n' + os.popen('conda env export').read() + '=' * 65

out_flag_files = [cal.replace(CAL_SUFFIX, OUT_FLAG_SUFFIX) for cal in cal_files]
for i, (cal, out_flag_file) in enumerate(zip(cal_files, out_flag_files)):
    # create UVFlag object based on UVCal
    uvc = UVCal()
    uvc.read_calfits(cal)
    uvf = UVFlag(uvc, mode='flag')
    uvf.use_future_array_shapes()
    
    # fill with flags
    for ant_ind, antnum in enumerate(uvf.ant_array):
        for pol_ind, polnum in enumerate(uvf.polarization_array):
            pol = {'Jee': 'e', 'Jnn': 'n'}[utils.jnum2str(polnum, x_orientation=uvf.x_orientation)]
            uvf.flag_array[ant_ind, :, :, pol_ind] = fh.final_flags[f'{antnum}{pol}'][i]

    # write to disk
    uvf.history += add_to_history        
    uvf.write(out_flag_file, clobber=True)

File /mnt/sn1/data2/2460548/zen.2460548.16914.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.16936.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.16958.sum.antenna_flags.h5 exists; clobbering
File /mnt/sn1/data2/2460548/zen.2460548.16981.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17003.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17025.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17048.sum.antenna_flags.h5 exists; clobbering
File /mnt/sn1/data2/2460548/zen.2460548.17070.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17093.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17115.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17137.sum.antenna_flags.h5 exists; clobbering

File /mnt/sn1/data2/2460548/zen.2460548.17160.sum.antenna_flags.h5 exists; clobbering

for repo in ['hera_cal', 'hera_qm', 'hera_filters', 'hera_notebook_templates', 'pyuvdata']:
    exec(f'from {repo} import __version__')
    print(f'{repo}: {__version__}')

hera_cal: 3.6.dev110+gcc0a13d
hera_qm: 2.1.5.dev6+g23b7cf7
hera_filters: 0.1.5

hera_notebook_templates: 0.1.dev734+g90f16f4
pyuvdata: 2.4.2

print(f'Finished execution in {(time.time() - tstart) / 60:.2f} minutes.')

Finished execution in 12.98 minutes.

Full Day Antenna Flagging¶

• Figure 1: Flag Summary vs. JD ¶

• Figure 2: Array Flag Fraction Summary ¶

• Figure 3: Flag Fraction vs. JD Summary ¶

• Figure 4: Per-Antenna Flag Harmonization Summary ¶

Parse inputs¶

Load data¶

Perform flagging¶

Plotting¶

Figure 1: Flag Summary vs. JD¶

Figure 2: Array Flag Fraction Summary¶

Figure 3: Flag Fraction vs. JD Summary¶

Figure 4: Per-Antenna Flag Harmonization Summary¶

Save results¶

Metadata¶

Full Day Antenna Flagging¶

• Figure 1: Flag Summary vs. JD¶

• Figure 2: Array Flag Fraction Summary¶

• Figure 3: Flag Fraction vs. JD Summary¶

• Figure 4: Per-Antenna Flag Harmonization Summary¶

Parse inputs¶

Load data¶

Perform flagging¶

Plotting¶

Figure 1: Flag Summary vs. JD¶

Figure 2: Array Flag Fraction Summary¶

Figure 3: Flag Fraction vs. JD Summary¶

Figure 4: Per-Antenna Flag Harmonization Summary¶

Save results¶

Metadata¶

• Figure 1: Flag Summary vs. JD ¶

• Figure 2: Array Flag Fraction Summary ¶

• Figure 3: Flag Fraction vs. JD Summary ¶

• Figure 4: Per-Antenna Flag Harmonization Summary ¶