#!/usr/bin/env python

import aipy as ap
import numpy as np
import commands, os, time, math, ephem, optparse, sys, warnings
import omnical.calibration_omni as omni
import cPickle as pickle
import scipy.signal as ss
import scipy.ndimage.filters as sfil
from scipy import interpolate
from scipy.stats import nanmedian

FILENAME = "omnical4.py"
print "#Omnical Version %s#"%omni.__version__
PI = np.pi
TPI = 2 * np.pi
TIMER = omni.Timer()
######################################################################
##############Config parameters###################################
######################################################################
o = optparse.OptionParser()

#ap.scripting.add_standard_options(o, cal=True, pol=True)
o.add_option('-t', '--tag', action = 'store', default = 'PSA128', help = 'tag name of this calibration')
o.add_option('-C', '--cal', action = 'store', default = None, help = 'calfile for processing uv files')
o.add_option('-d', '--datatag', action = 'store', default = None, help = 'tag name of this data set')
o.add_option('-i', '--infopath', action = 'store', default = '/data2/home/hz2ug/omnical/doc/redundantinfo_PSA128_17ba.bin', help = 'redundantinfo file to read')
o.add_option('-r', '--rawcalpath', action = 'store', default = 'NORAWCAL', help = 'raw calibration parameter file to read. The file should be a pickle file generated by first_cal.py')
o.add_option('--add', action = 'store_true', help = 'whether to enable crosstalk removal')
o.add_option('--nadd', action = 'store', type = 'int', default = -1, help = 'time steps w to remove additive term with. for running average its 2w + 1 sliding window.')
o.add_option('--flagsigma', action = 'store', type = 'float', default = 5, help = 'Number of sigmas to flag on chi^2 distribution. 4 sigma by default. For full cadence data, 3 is recommended.')
o.add_option('--flagt', action = 'store', type = 'int', default = 999999, help = 'Number of time slices to run the minimum filter when flagging. 999999 by default.')
o.add_option('--flagf', action = 'store', type = 'int', default = 3, help = 'Number of frequency slices to run the minimum filter when flagging. 4 by default.')
o.add_option('--datapath', action = 'store', default = 'NOBINDATA', help = 'binary data file folder to save/load binary data converted from uv file. Omit this option if you dont want to save binary data.')
o.add_option('--treasure', action = 'store', default = None, help = 'trasure folder to update.')
o.add_option('--healthbar', action = 'store', default = '2', help = 'health threshold (0-100) over which an antenna is marked bad.')
o.add_option('-o', '--outputpath', action = 'store', default = ".", help = 'output folder')
o.add_option('-k', '--skip', action = 'store_true', help = 'whether to skip data importing from uv')
o.add_option('-u', '--newuv', action = 'store_true', help = 'whether to create new uv files with calibration applied')
o.add_option('--flag', action = 'store_true', help = 'whether to create new flagging')
o.add_option('-f', '--overwrite', action = 'store_true', help = 'whether to overwrite if the new uv files already exists')
o.add_option('-s', '--singlethread', action = 'store_true', help = 'whether to disable multiprocessing for calibration and use only one thread. May need this option for things like grid engine.')
o.add_option('--chemo', action = 'store', type='float', default = 1, help = 'whether to apply chemotherapy when flagging. >= 1 enables. The larger the number, the stricter the chemo is. A 10 means all will be flagged if more than 1/10 of the data is flagged')
o.add_option('--plot', action = 'store_true', help = 'whether to make plots in the end')
o.add_option('--skip_sun', action = 'store_true', help = 'whether to calibrate data set with sun up.')
o.add_option('--mem', action = 'store', type = 'float', default = 4e9, help = 'Amount of initial memory to reserve when parsing uv files in number of bytes.')
o.add_option('--chisq_leniency', action = 'store', type = 'float', default = 1.5, help = 'Factor to multiply noise-model by to serve as one of the flagging thresholds.')
o.add_option('--chemof', action = 'store', type = 'float', default = 4, help = 'Fraction threshold of time in a frequency bin, above which that frequency will be flagged. If set to 4, more than 1/4 of time being bad will flag all time in a frequency bin.')
o.add_option('--chemot', action = 'store', type = 'float', default = 4, help = 'Fraction threshold of frequency in a time stamp, above which that time will be flagged. If set to 4, more than 1/4 of frequency being bad will flag all frequency in a time stamp.')
o.add_option('--model_noise', action = 'store', default = None, help = 'A model .omnichisq file that contains the noise model (sigma^2) with the first two columns being lst in range [0,2pi). Separate by , the same order as -p. Need to be the same unit with data or model_treasure.')
o.add_option('--model_treasure', action = 'store', default = None, help = 'A treasure file that contains good foreground visibilities. ')
o.add_option('--delay_compression', action = 'store', type = 'int', default = None, help = 'Compress to the specified number of frequency bins.')

opts,args = o.parse_args(sys.argv[1:])
uvfiles = [os.path.expanduser(arg) for arg in args]
skip = opts.skip
create_new_uvs = opts.newuv
need_new_flag = opts.flag
calibrate_sun = not opts.skip_sun
overwrite_uvs = opts.overwrite
make_plots = opts.plot
if make_plots:
    import matplotlib.pyplot as plt
ano = opts.tag##This is the file name difference for final calibration parameter result file. Result will be saved in miriadextract_xx_ano.omnical
dataano = opts.datatag#ano for existing data and lst.dat
sourcepath = os.path.expanduser(opts.datapath)
oppath = os.path.expanduser(opts.outputpath)
chemo = opts.chemo
chemof = opts.chemof
chemot = opts.chemot

if opts.treasure is not None:
    if not os.path.isdir(os.path.expanduser(opts.treasure)):
        raise IOError("Treasure path %s does not exist."%opts.treasure)
    treasurePath = os.path.expanduser(opts.treasure)
else:
    treasurePath = None




wantpols = {}
for p in ['xx', 'yy']:
    wantpols[p] = ap.miriad.str2pol[p]

for uvf in uvfiles:
    if not os.path.isdir(uvf):
        uvfiles.remove(uvf)
        print "WARNING: uv file path %s does not exist!"%uvf

if '.odf' in uvfiles[0]:
    input_type = 'odf'
else:
    input_type = 'uv'

if len(uvfiles) == 0:
    raise IOError("ERROR: No valid uv files detected in input. Exiting!")
elif len(uvfiles) != 4 and len(uvfiles) != 1 and input_type == 'uv':
    raise IOError("ERROR: %i uvfiles are inputed and assumed to be corresponding to different polarizations, but we expect %s polarizations. Exiting!"%(len(uvfiles), 4))
elif len(uvfiles) != 1 and input_type == 'odf':
    raise IOError("ERROR: %i odfs are inputed and only one at a time is supported."%len(uvfiles))
else:
    if input_type == 'uv' and len(uvfiles) == 4:
        uvfiles_dic = {}
        for i, p in enumerate(['xx', 'xy', 'yx', 'yy']):
            uvfiles_dic[p] = uvfiles[i]


flag_thresh = opts.flagsigma
flagt = opts.flagt
flagf = opts.flagf
if opts.singlethread:
    nthread = 1
else:
    nthread = None

keep_binary_data = False
if os.path.isdir(sourcepath):
    keep_binary_data = True
elif opts.skip:
    raise IOError("Direct binary data import requested by -k or --skip option, but the --datapth %s doesn't exist."%sourcepath)

keep_binary_calpar = False

#print opts.healthbar, opts.healthbar.split(), len(opts.healthbar.split())
if len(opts.healthbar.split(',')) == 1:
    healthbar = float(opts.healthbar)
    ubl_healthbar = 100
elif len(opts.healthbar.split(',')) == 2:
    healthbar = float(opts.healthbar.split(',')[0])
    ubl_healthbar = float(opts.healthbar.split(',')[1])
else:
    raise Exception("User input healthbar option (--healthbar %s) is not recognized."%opts.healthbar)

init_mem = opts.mem



if dataano is None:
    dataano = ''
    for i, uvf in enumerate(uvfiles):
        if i!= 0:
            dataano = dataano + '_'
        while os.path.basename(uvf) == '' and len(uvf) > 0:
            uvf = uvf[:-1]
        dataano = dataano + os.path.basename(uvf)



have_model_noises = False
if opts.model_noise is not None:
    if not need_new_flag:
        raise IOError("--noise_model supplied without --flag. Noise model is only useful when doing new flagging.")
    if len(opts.model_noise.split(',')) != len(opts.pol.split(',')):
        raise ValueError("--model_noise got argument %s that does not have the same number of polarizations as -p argument %s."%(opts.model_noise, opts.pol))
    model_noises = {}
    for p, pol in enumerate(opts.pol.split(',')):
        model_noise_file = os.path.expanduser(opts.model_noise.split(',')[p])
        if not os.path.isfile(model_noise_file):
            raise IOError("model noise file %s does not exist"%model_noise_file)
        model_noises[pol] = omni.load_omnichisq(model_noise_file)
        if np.max(omni.get_omnitime(model_noises[pol])) >= TPI or np.min(omni.get_omnitime(model_noises[pol])) < 0:
            raise ValueError("Times stored in noise model %s is outside the range [0, 2pi)."%model_noise_file)
    have_model_noises = True
    chisq_leniency = opts.chisq_leniency

have_model_treasure = False
if opts.model_treasure is not None:
    if os.path.isdir(os.path.expanduser(opts.model_treasure)):
        model_treasure = omni.Treasure(os.path.expanduser(opts.model_treasure))
        for pol in opts.pol.split(','):
            if pol not in model_treasure.ubls.keys():
                raise ValueError("Polarization %s not found in the model treasure file %s."%(pol, os.path.expanduser(opts.model_treasure)))
    else:
        raise IOError("Model treasure folder not found: %s."%os.path.expanduser(opts.model_treasure))
    have_model_treasure = True

delay_compression = opts.delay_compression
if delay_compression is not None:
    if delay_compression%2 != 1 or delay_compression <= 0:
        raise ValueError('Delay compression bin number must be odd and positive. Input: %i.'%delay_compression)

if input_type == 'uv':
    print "Reading calfile %s"%opts.cal,
    sys.stdout.flush()
    aa = ap.cal.get_aa(opts.cal, np.array([.15]))
    print "Done"
    sys.stdout.flush()

infopaths = {}
for pol in wantpols.keys():
    infopaths[pol]= os.path.expanduser(opts.infopath)


removedegen = 5
if opts.add and opts.nadd > 0:
    removeadditive = True
    removeadditiveperiod = opts.nadd
else:
    removeadditive = False
    removeadditiveperiod = -1

crudecalpath = os.path.expanduser(opts.rawcalpath)
needrawcal = False
if os.path.isfile(crudecalpath):
    needrawcal = True
    with open(crudecalpath, 'rb') as crude_calpar_file:
        crude_calpar = pickle.load(crude_calpar_file)
elif crudecalpath != 'NORAWCAL':
    raise IOError("Input rawcalpath %s doesn't exist on disk."%crudecalpath)


converge_percent = 0.001
max_iter = 20
step_size = .3

######################################################################
######################################################################
######################################################################

########Massage user parameters###################################
sourcepath += '/'
oppath += '/'
utcPath = sourcepath + 'miriadextract_' + dataano + "_localtime.dat"
lstPath = sourcepath + 'miriadextract_' + dataano + "_lsthour.dat"

####get some info from the first uvfile   ################
print "Getting some basic info from %s"%uvfiles[0],
sys.stdout.flush()

sa = ephem.Observer()
sa.pressure = 0

if input_type == 'odf':
    with open(uvfiles[0]+'/header.txt') as f:
        for l in f.readlines():
            if l.split()[0] == 'nChannels':
                nfreq = int(l.split()[1])
            elif l.split()[0] == 'nAntennas':
                nant = int(l.split()[1])
            elif l.split()[0] == 'longitude':#odf header has bug that switched lat and lon
                sa.lat = float(l.split()[1]) * PI/180.
            elif l.split()[0] == 'latitude':
                sa.lon = float(l.split()[1]) * PI/180.
            elif l.split()[0] == 'startFreq':
                startfreq = float(l.split()[1]) / 1000.
            elif l.split()[0] == 'endFreq':
                endfreq = float(l.split()[1]) / 1000.
        dfreq = (endfreq - startfreq) / nfreq

    header_antloc = np.loadtxt(uvfiles[0].replace('.odf','.odfa') + '/antlocx5.dat')
else:
    uv=ap.miriad.UV(uvfiles[0])
    nfreq = uv.nchan;
    nant = uv['nants']
    sa.lon = aa.lon
    sa.lat = aa.lat
    startfreq = uv['sfreq']#GHz
    dfreq = uv['sdf']#GHz
    del(uv)
print "Done."
sys.stdout.flush()

need_compression = (delay_compression > 0 and delay_compression < nfreq)

#######import data file###########################
if input_type == 'odf':
    nTimes = []
    for uvfile in uvfiles:
        with open(uvfile+'/header.txt') as f:
            for l in f.readlines():
                if l.split()[0] == 'nIntegrations':
                    nTimes = nTimes + [int(l.split()[1])]
    pol_select = []
    for key in wantpols.keys():
        if 'xx' == key:
            pol_select = pol_select + [0]
        elif 'xy' == key:
            pol_select = pol_select + [1]
        elif 'yx' == key:
            pol_select = pol_select + [2]
        elif 'yy' == key:
            pol_select = pol_select + [3]
    data = np.zeros((len(pol_select), np.sum(nTimes), nfreq, nant * (nant + 1) / 2), dtype='complex64')
    rawflag = np.zeros(data.shape, dtype='bool')
    timing = []
    lst = []
    for i, uvfile in enumerate(uvfiles):
        data[:, len(timing):len(timing)+nTimes[i]] = np.fromfile(uvfile+'/visibilities', dtype='complex64').reshape((4, nTimes[i], nfreq, nant * (nant + 1) / 2))[pol_select]
        with open(uvfile.replace('.odf', '.odfa') + '/timing.dat') as f:
            for l in f.readlines():
                sa.date = l.replace('\n','')
                sa.date += ephem.second * 3600 * 4
                lst += [sa.sidereal_time()]
                timing += [sa.date.__str__()]
else:
    ###start reading miriads################
    if skip:
        if need_compression:
            raise Exception('skip option -k and compression --delay_compress are currently not supported at the same time.')
        print FILENAME + " MSG: SKIPPED reading uvfiles. Reading binary data files directly...",
        sys.stdout.flush()
        with open(utcPath) as f:
            timing = [t.replace('\n','') for t in f.readlines()]
        with open(lstPath) as f:
            lst = [float(t) for t in f.readlines()]
        print (len(timing), nfreq, len(aa) * (len(aa) + 1) / 2), "...",
        data = np.array([np.fromfile(sourcepath + 'data_' + dataano + '_' + pol, dtype = 'complex64').reshape((len(timing), nfreq, len(aa) * (len(aa) + 1) / 2)) for pol in wantpols.keys()])
        rawflag = np.array([np.fromfile(sourcepath + 'flag_' + dataano + '_' + pol, dtype = 'bool').reshape((len(timing), nfreq, len(aa) * (len(aa) + 1) / 2)) for pol in wantpols.keys()])
        print "Done."
        sys.stdout.flush()

    else:
        print FILENAME + " MSG:",  len(uvfiles), "%s files to be processed for "%input_type + ano
        sys.stdout.flush()
        if len(uvfiles) == 1:
            data, jd, timing, lst, rawflag = omni.importuvs(uvfiles, wantpols, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)#, nTotalAntenna = len(aa))
        else:

            for p, pol in enumerate(wantpols.keys()):
                if p == 0:
                    data, jd, timing, lst, rawflag = omni.importuvs([uvfiles_dic[pol]], {pol:wantpols[pol]}, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)
                else:
                    tmpdata, jd, timing, lst, tmpflag = omni.importuvs([uvfiles_dic[pol]], {pol:wantpols[pol]}, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)
                    data = np.concatenate((data, tmpdata))
                    rawflag = np.concatenate((rawflag, tmpflag))

        print FILENAME + " MSG:",  len(jd), "slices read. data shape: ", data.shape
        sys.stdout.flush()

        if keep_binary_data:
            print FILENAME + " MSG: saving binary data to disk...",
            sys.stdout.flush()
            f = open(utcPath,'w')
            for qaz in timing:
                f.write("%s\n"%qaz)
            f.close()
            f = open(lstPath,'w')
            for l in lst:
                f.write("%s\n"%l)
            f.close()
            for p,pol in zip(range(len(wantpols)), wantpols.keys()):
                data[p].tofile(sourcepath + 'data_' + dataano + '_' + pol)
                rawflag[p].tofile(sourcepath + 'flag_' + dataano + '_' + pol)
            print "Done."
            sys.stdout.flush()



#####print some astronomical info
sun = ephem.Sun()
sunpos  = np.zeros((len(timing), 2))
cenA = ephem.FixedBody()
cenA._ra = 3.5146
cenA._dec = -.75077
cenApos = np.zeros((len(timing), 2))
for nt,tm in zip(range(len(timing)),timing):
    sa.date = tm

    sun.compute(sa)
    sunpos[nt] = sun.alt, sun.az
    cenA.compute(sa)
    cenApos[nt] = cenA.alt, cenA.az
print FILENAME + " MSG: data time range UTC: %s to %s, sun altaz from (%f,%f) to (%f,%f)"%(timing[0], timing[-1], sunpos[0,0], sunpos[0,1], sunpos[-1,0], sunpos[-1,1])#, "CentaurusA altaz from (%f,%f) to (%f,%f)"%(cenApos[0,0], cenApos[0,1], cenApos[-1,0], cenApos[-1,1])
sys.stdout.flush()

if (not calibrate_sun) and (not (sunpos[:,0] < -.1).all()):
    print "Sun up", sunpos[:,0]
    sys.exit(0)

###########initialize stuff
calibrators = {}
omnigains = {}
adds = {}
flags = {}
uvflags = {}

for p, pol in zip(range(len(data)), wantpols.keys()):
    ####create redundant calibrators################
    if input_type == 'uv':
        calibrators[pol] = omni.RedundantCalibrator_PAPER(aa)
    elif input_type == 'odf':
        calibrators[pol] = omni.RedundantCalibrator_X5(header_antloc)
    calibrators[pol].read_redundantinfo(infopaths[pol], verbose=False)
    info = calibrators[pol].Info.get_info()
    calibrators[pol].nTime = len(timing)
    calibrators[pol].nFrequency = nfreq

    ####consolidate 3D flags from uv files into per time/freq 2D flags
    uvflags[pol] = np.any(rawflag[p,:,:,calibrators[pol].subsetbl[calibrators[pol].crossindex]], axis = 0)# aweird transpose happens when slicing

    ###apply, if needed, raw calibration################
    if needrawcal:
        original_data = np.copy(data[p])
        data[p] = omni.apply_calpar(data[p], crude_calpar[pol], calibrators[pol].totalVisibilityId)

    ####calibrate################
    calibrators[pol].removeDegeneracy = removedegen
    calibrators[pol].convergePercent = converge_percent
    calibrators[pol].maxIteration = max_iter
    calibrators[pol].stepSize = step_size

    ################first round of calibration  #########################
    print FILENAME + " MSG: starting calibration on %s %s. nTime = %i, nFrequency = %i ..."%(dataano, pol, calibrators[pol].nTime, calibrators[pol].nFrequency),
    sys.stdout.flush()
    timer = time.time()
    additivein = np.zeros_like(data[p])

    #######################Logcal###############################
    calibrators[pol].logcal(data[p], additivein, nthread = nthread, verbose=True)

    if needrawcal:#restore original data's amplitude after logcal. dont restore phase because it may cause problem in degeneracy removal. Also phase wouldnt affect accuracy in lincal
        calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] + np.log10(np.abs(crude_calpar[pol][:, calibrators[pol].subsetant]))
        #calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] + np.angle(crude_calpar[pol][:, calibrators[pol].subsetant])
        #data[p] = np.copy(original_data)
        data[p] = omni.apply_calpar(data[p], 1/np.abs(crude_calpar[pol]), calibrators[pol].totalVisibilityId)
        #del original_data

    #######################Lincal###############################
    additiveout = calibrators[pol].lincal(data[p], additivein, nthread = nthread, verbose=True)

    #######################remove additive###############################
    if removeadditive:
        nadditiveloop = 1
        for i in range(nadditiveloop):
            #subtimer = omni.Timer()
            additivein[:,:,calibrators[pol].Info.subsetbl] = additivein[:,:,calibrators[pol].Info.subsetbl] + additiveout
            weight = ss.convolve(np.ones(additivein.shape[0]), np.ones(removeadditiveperiod * 2 + 1), mode='same')
            #for f in range(additivein.shape[1]):#doing for loop to save memory usage at the expense of negligible time
                #additivein[:,f] = ss.convolve(additivein[:,f], np.ones(removeadditiveperiod * 2 + 1)[:, None], mode='same')/weight[:, None]
            additivein = ((sfil.convolve1d(np.real(additivein), np.ones(removeadditiveperiod * 2 + 1), mode='constant') + 1j * sfil.convolve1d(np.imag(additivein), np.ones(removeadditiveperiod * 2 + 1), mode='constant'))/weight[:, None, None]).astype('complex64')
            calibrators[pol].computeUBLFit = False
            additiveout = calibrators[pol].lincal(data[p], additivein, nthread = nthread, verbose=True)

    #if needrawcal:#restore original data's phase. Since we are "changing data" here, ublfits in rawCalpar dont have to change accordingly
        ##calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] + np.log10(np.abs(crude_calpar[pol][:, calibrators[pol].subsetant]))
        #calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] + np.angle(crude_calpar[pol][:, calibrators[pol].subsetant])
        ##data[p] = np.copy(original_data)
        #data[p] = omni.apply_calpar(data[p], np.exp(-1j * np.angle(crude_calpar[pol])), calibrators[pol].totalVisibilityId)
        #additivein = omni.apply_calpar(additivein, np.exp(-1j * np.angle(crude_calpar[pol])), calibrators[pol].totalVisibilityId)
        ##del original_data

    #####################flag bad data part 1#########################
    #####################need to flag before absolute calibration, bc absolutecal supresses high chisq caused by rfi
    if need_new_flag:
        print "%s flagged fraction in uv: "%(pol), float(np.sum(uvflags[pol]))/uvflags[pol].shape[0]/uvflags[pol].shape[1]
        flags[pol] = calibrators[pol].flag(nsigma = flag_thresh, twindow=flagt, fwindow=flagf)#True if bad and flagged
        if input_type == 'odf' and .137 > startfreq and .137 < endfreq:
            flags[pol][:, int((.137-startfreq)/dfreq):int((.1385-startfreq)/dfreq)] = True
        if input_type == 'odf' and .15 > startfreq and .15 < endfreq:
            flags[pol][:, int((.149-startfreq)/dfreq):int((.151-startfreq)/dfreq)] = True
        if input_type == 'odf' and .174 < startfreq:
            flags[pol][:, int((.174-startfreq)/dfreq):] = True
        if input_type == 'odf' and .126 > startfreq:
            flags[pol][:, :int((.126-startfreq)/dfreq)] = True
        flags[pol] = flags[pol] | uvflags[pol]
    else:
        flags[pol] = uvflags[pol]
    print "%s flagged fraction pre-chemo: "%(pol), float(np.sum(flags[pol]))/uvflags[pol].shape[0]/uvflags[pol].shape[1]

    #omni.omniview([data[0,0,110], calibrators['xx'].get_calibrated_data(data[0])[0,110]], info = calibrators['xx'].get_info())
    if have_model_treasure and (sunpos[:,0] < -.1).all():
        raise Exception("Model treasure temporarily disabled.")
        #print np.nanmean(np.linalg.norm((data[0]-calibrators[pol].get_modeled_data())[..., calibrators[pol].subsetbl[calibrators[pol].crossindex]], axis = -1)**2 / calibrators[pol].rawCalpar[...,2])
        amp_cal, phs_cal = calibrators[pol].absolutecal_w_treasure(model_treasure, pol, np.array(lst)*TPI/24., tolerance = 2.)
        #print np.nanmean(np.linalg.norm((data[0]-calibrators[pol].get_modeled_data())[..., calibrators[pol].subsetbl[calibrators[pol].crossindex]], axis = -1)**2 / calibrators[pol].rawCalpar[...,2])
    #omni.omniview([data[0,0,110], calibrators['xx'].get_calibrated_data(data[0])[0,110]/5e4], info = calibrators['xx'].get_info())


    #####################flag bad data part 2#########################
    if need_new_flag:

        if have_model_noises and (sunpos[:,0] < -.1).all():
            if model_noises[pol][0,2] != calibrators[pol].nFrequency:
                raise ValueError('Model noise on %pol has nFrequency %i that differs from calibrators %i.'%(pol, model_noises[pol][0,2], calibrators[pol].nFrequency))
            interp_model = interpolate.interp1d(np.append(omni.get_omnitime(model_noises[pol]), [TPI]), np.append(model_noises[pol][..., 3:], [model_noises[pol][0, ..., 3:]], axis=0), axis = 0)
            model_chisq = chisq_leniency * interp_model(np.array(lst)*TPI/24.) * (len(calibrators[pol].crossindex) - calibrators[pol].nAntenna - calibrators[pol].nUBL + 2) * 10**(4 * np.median(calibrators[pol].rawCalpar[..., 3:3+calibrators[pol].nAntenna], axis = -1))#leniency factor * model * DOF * gain correction
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore",category=RuntimeWarning)
                flags[pol] = flags[pol]| (calibrators[pol].rawCalpar[..., 2] > model_chisq)

    print "Done. %fmin"%(float(time.time()-timer)/60.)
    sys.stdout.flush()

####################################################################################
#####################cross-pol calibration############################
####################################################################################
###start reading miriads################
print FILENAME + " MSG: Processing cross pol:"
sys.stdout.flush()

crosspols = {}
for p in ['xy', 'yx']:
    crosspols[p] = ap.miriad.str2pol[p]

if input_type == 'odf':
    nTimes = []
    for uvfile in uvfiles:
        with open(uvfile+'/header.txt') as f:
            for l in f.readlines():
                if l.split()[0] == 'nIntegrations':
                    nTimes = nTimes + [int(l.split()[1])]
    pol_select = []
    for key in crosspols.keys():
        if 'xx' == key:
            pol_select = pol_select + [0]
        elif 'xy' == key:
            pol_select = pol_select + [1]
        elif 'yx' == key:
            pol_select = pol_select + [2]
        elif 'yy' == key:
            pol_select = pol_select + [3]
    rawdata = np.zeros((len(pol_select), np.sum(nTimes), nfreq, nant * (nant + 1) / 2), dtype='complex64')
    timing = []
    lst = []

    for i, uvfile in enumerate(uvfiles):
        rawdata[:, len(timing):len(timing)+nTimes[i]] = np.fromfile(uvfile+'/visibilities', dtype='complex64').reshape((4, nTimes[i], nfreq, nant * (nant + 1) / 2))[pol_select]
        with open(uvfile.replace('.odf', '.odfa') + '/timing.dat') as f:
            for l in f.readlines():
                sa.date = l.replace('\n','')
                sa.date += ephem.second * 3600 * 4
                lst += [sa.sidereal_time()]
                timing += [sa.date.__str__()]
else:
    if len(uvfiles) == 1:
        rawdata, t, timing, lst, rawflag = omni.importuvs(uvfiles, crosspols, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)#, nTotalAntenna = len(aa))
    else:

        for p, pol in enumerate(crosspols.keys()):
            if p == 0:
                rawdata, t, timing, lst, rawflag = omni.importuvs([uvfiles_dic[pol]], {pol:crosspols[pol]}, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)
            else:
                tmpdata, t, timing, lst, tmpflag = omni.importuvs([uvfiles_dic[pol]], {pol:crosspols[pol]}, totalVisibilityId = np.concatenate([[[i,j] for i in range(j + 1)] for j in range(len(aa))]), timingTolerance=100)
                rawdata = np.concatenate((rawdata, tmpdata))
    print FILENAME + " MSG:",  len(t), "slices read."
    sys.stdout.flush()


c = calibrators['xx']
calpar = {}
for pol in ['xx','yy']:
    calpar[pol[0]] = np.ones((c.nTime, c.nFrequency, c.nTotalAnt), dtype='complex64')
    calpar[pol[0]][:, :, c.subsetant] = 10**calibrators[pol].rawCalpar[...,3:3+c.nAntenna] * np.exp(1.j * calibrators[pol].rawCalpar[...,3+c.nAntenna:3+2*c.nAntenna])

for p, pol in enumerate(crosspols.keys()):
    rawdata[p] = omni.apply_calpar2(rawdata[p], calpar[pol[0]], calpar[pol[1]], c.totalVisibilityId)

#construct A and b for cross-pol
A = []
bindex = []
ublcount_sort = np.argsort(c.ublcount)
for u in range(c.nUBL):
    if c.ublcount[u] >= c.ublcount[ublcount_sort[-min(c.nUBL,10)]]:
        cbl = c.ublindex[u][:,2].astype(int)
        mask = c.reversed[cbl] == 1
        if mask.any():
            A = A + (c.antloc[c.ublindex[u][mask, 0].astype(int)] - c.antloc[int(c.ublindex[u][mask, 0][0])]).tolist()
            bindex = bindex + [[a, int(c.ublindex[u][mask, 2][0])] for a in c.ublindex[u][mask, 2].astype(int)]#inside crossindex
        if not mask.all():
            A = A + (c.antloc[c.ublindex[u][~mask, 0].astype(int)] - c.antloc[int(c.ublindex[u][~mask, 0][0])]).tolist()
            bindex = bindex + [[a, int(c.ublindex[u][~mask, 2][0])] for a in c.ublindex[u][~mask, 2].astype(int)]#inside crossindex
A = np.array(A)[:, :2]
A = np.concatenate((A, -A), axis=0)
bindex = c.subsetbl[c.crossindex[np.array(bindex).astype(int)]]#inside all bl

b = np.empty((len(A), rawdata.shape[1], rawdata.shape[2]), dtype='float32')
for p, pol in enumerate(crosspols.keys()):
    b[p*len(bindex):(p+1)*len(bindex)] = (np.angle(rawdata[p][..., bindex[:,0]]) - np.angle(rawdata[p][..., bindex[:,1]])).transpose((2,0,1))

crosstimer = omni.Timer()
print "Solving cross calibration....",
sys.stdout.flush()
sol = omni.solve_slope(A, b, 1)
#sol[:,flags['xx']] = np.nan
#sol[:,flags['yy']] = np.nan
print "Done.",
sys.stdout.flush()
crosstimer.tick()
if make_plots:
    for i in range(len(sol)):
        plt.subplot('21%i'%(i+1))
        min_period = TPI / min(abs(A[:,i][abs(A[:,i]) > 1]))
        plt.imshow((sol[i]+min_period/2.)%min_period-min_period/2.);plt.colorbar()
    plt.show()
#apply to data
crosscalpar = {}
for pol in ['xx','yy']:
    crosscalpar[pol[0]] = np.ones((rawdata.shape[1], rawdata.shape[2], c.nTotalAnt), dtype='complex64')
crosscalpar[crosspols.keys()[0][1]][..., c.subsetant] = np.exp(1.j * c.antloc[:,:2].dot(sol.transpose(1,0,2))).transpose(1,2,0)

cdata = np.empty_like(rawdata)
for p, pol in enumerate(crosspols.keys()):
    cdata[p] = omni.apply_calpar2(rawdata[p], crosscalpar[pol[0]], crosscalpar[pol[1]], c.totalVisibilityId)
del(rawdata)

#################################################################################
############################solve for overall constant###########################
#################################################################################
crossextract, cross_chisq = omni.extract_crosspol_ubl(cdata, c.Info.get_info())
red_enough = c.ublcount > (max(c.ublcount) / 2) #treat all ubls above half max ublcount equally when computing the x-y angle
xy_candidates = np.empty(list(cdata.shape)[1:3] + [np.sum(red_enough)], dtype='float32')

for i,u in enumerate(np.arange(c.nUBL)[red_enough].astype(int)):
    xy_candidates[..., i] = np.angle(crossextract[0,...,u]/crossextract[1,...,u])

with warnings.catch_warnings():
    warnings.filterwarnings("ignore",category=RuntimeWarning)
    overall_angle = omni.medianAngle(xy_candidates, axis = -1) / 2

    median_overall_angle = omni.medianAngle(xy_candidates.transpose(0,2,1).reshape((xy_candidates.shape[0]*xy_candidates.shape[2],xy_candidates.shape[1])), axis = 0) / 2
    prev_valid = -1
    valids1 = np.zeros(len(median_overall_angle), dtype='bool')
    for i in range(1, len(median_overall_angle)):
        if not (np.isnan(median_overall_angle[i-1]) or (prev_valid >= 0 and abs((median_overall_angle[i-1] - median_overall_angle[prev_valid] + PI/2)%PI-PI/2) > PI/4)):
            prev_valid = i - 1
            valids1[i-1] = True
        if prev_valid >= 0:
            offset = median_overall_angle[prev_valid]-PI/2
        else:
            offset = -PI/2
        median_overall_angle[i] = (median_overall_angle[i] - offset)%(PI) + offset

    offset = median_overall_angle[None]-PI/2
    overall_angle = (overall_angle - offset)%(PI) + offset
if make_plots:
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore",category=RuntimeWarning)
        plt.imshow(overall_angle, vmin = -PI, vmax = PI);plt.colorbar()
    plt.title('overall angle solution between xx and yy')
    plt.show()

crosscalpar2={}
for pol in ['xx','yy']:
    crosscalpar2[pol[0]] = np.ones((cdata.shape[1], cdata.shape[2], c.nTotalAnt), dtype='complex64')
crosscalpar2[crosspols.keys()[0][1]][..., c.subsetant] = np.exp(1.j * overall_angle[..., None])

cdata2 = np.empty_like(cdata)
for p, pol in enumerate(crosspols.keys()):
    cdata2[p] = omni.apply_calpar2(cdata[p], crosscalpar2[pol[0]], crosscalpar2[pol[1]], c.totalVisibilityId)
del(cdata)

####apply cross pol solution to 'xx'############
apply_pol = crosspols.keys()[0][1] + crosspols.keys()[0][1]
#gain
calibrators[apply_pol].rawCalpar[..., 3+c.nAntenna:3+2*c.nAntenna] = (calibrators[apply_pol].rawCalpar[..., 3+c.nAntenna:3+2*c.nAntenna] + c.antloc[:,:2].dot(sol.transpose(1,0,2)).transpose(1,2,0) + overall_angle[..., None] + PI)%TPI -PI

###use gain phase to try to fix pi jump for a second time: doesnt do anything  at the moment
with warnings.catch_warnings():
    warnings.filterwarnings("ignore",category=RuntimeWarning)
    median_calpar_angle = omni.medianAngle(omni.medianAngle(calibrators[apply_pol].rawCalpar[..., 3+c.nAntenna:3+2*c.nAntenna], axis = 0), axis = -1)
    prev_valid = -1
    valids2 = np.zeros(len(median_calpar_angle), dtype='bool')
    for i in range(1, len(median_calpar_angle)):
        if not (np.isnan(median_calpar_angle[i-1])or (prev_valid >= 0 and abs((median_calpar_angle[i-1] - median_calpar_angle[prev_valid] + PI/2)%PI-PI/2) > PI/4)):
            prev_valid = i - 1
            valids2[i-1] = True
        if prev_valid >= 0:
            offset = median_calpar_angle[prev_valid]-PI/2
        else:
            offset = -PI/2
        median_calpar_angle[i] = (median_calpar_angle[i] - offset)%(PI) + offset
    pi_jump = median_calpar_angle - omni.medianAngle(omni.medianAngle(calibrators[apply_pol].rawCalpar[..., 3+c.nAntenna:3+2*c.nAntenna], axis = 0), axis = -1)



#ubl fits
old_ubl_fits = calibrators[apply_pol].rawCalpar[..., 3+2*c.nAntenna::2] + 1.j * calibrators[apply_pol].rawCalpar[..., 3+2*c.nAntenna+1::2]
new_ubl_fits = old_ubl_fits / np.exp(1.j * c.ubl[:,:2].dot(sol.transpose(1,0,2)).transpose(1,2,0))
calibrators[apply_pol].rawCalpar[..., 3+2*c.nAntenna::2] = np.real(new_ubl_fits)
calibrators[apply_pol].rawCalpar[..., 3+2*c.nAntenna+1::2] = np.imag(new_ubl_fits)
#cross ubl fits
crossextract[0] = crossextract[0] / np.exp(1.j * overall_angle[..., None])
crossextract[1] = crossextract[1] * np.exp(1.j * overall_angle[..., None])
#========================End of cross-pol==============================#

if needrawcal:#restore original data's phase. Since we are "changing data" here, ublfits in rawCalpar dont have to change accordingly
    for p, pol in enumerate(calibrators.keys()):
        #calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3:3 + calibrators[pol].nAntenna] + np.log10(np.abs(crude_calpar[pol][:, calibrators[pol].subsetant]))
        calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] = calibrators[pol].rawCalpar[:, :, 3 + calibrators[pol].nAntenna:3 + 2 * calibrators[pol].nAntenna] + np.angle(crude_calpar[pol][:, calibrators[pol].subsetant])
        #data[p] = np.copy(original_data)
        data[p] = omni.apply_calpar(data[p], np.exp(-1j * np.angle(crude_calpar[pol])), calibrators[pol].totalVisibilityId)
        additivein = omni.apply_calpar(additivein, np.exp(-1j * np.angle(crude_calpar[pol])), calibrators[pol].totalVisibilityId)
        #del original_data


#todo: flagging using cross_chisq
print "Performing more flagging and saving the results"
sys.stdout.flush()

for p, pol in enumerate(wantpols.keys()):
    for p2, pol2 in enumerate(wantpols.keys()):
        flags[pol] = flags[pol]|flags[pol2]
for p, pol in enumerate(wantpols.keys()):
    if need_new_flag:
        if chemo >= 1:

            flags[pol] = flags[pol]|(ss.convolve(flags[pol],[[.2,.4,1,.4,.2]],mode='same')>=1)
            flags[pol] = flags[pol]|(ss.convolve(flags[pol],[[.2],[.4],[1],[.4],[.2]],mode='same')>=1)

            bad_freq = (np.sum(flags[pol], axis = 0) > (calibrators[pol].nTime / chemof))
            bad_time = (np.sum(flags[pol][:, ~bad_freq], axis = 1) > (np.sum(~bad_freq) / chemot))
            bad_freq = bad_freq | (ss.convolve(bad_freq, [1,1,1],mode='same')>=1)
            bad_time = bad_time | (ss.convolve(bad_time, [1,1,1],mode='same')>=1)
            flags[pol] = flags[pol]|bad_freq[None,:]|bad_time[:, None]
            print "%s flagged fraction post-chemo: "%(pol), float(np.sum(flags[pol]))/uvflags[pol].shape[0]/uvflags[pol].shape[1]

        if chemo > 1 and np.sum(flags[pol]) > flags[pol].shape[0] * flags[pol].shape[1] / chemo:
            flags[pol] = flags[pol] | True





    #######################save results###############################
    print FILENAME + " MSG: saving calibration results on %s %s."%(dataano, pol),
    sys.stdout.flush()
    calibrators[pol].utctimes = timing
    omnigains[pol[0]] = calibrators[pol].get_omnigain()
    adds[pol] = additivein


    if removeadditive:
        adds[pol].tofile(oppath + '/' + dataano + '_' + ano + "_%s.omniadd"%pol + str(removeadditiveperiod))
    if keep_binary_calpar:
        calibrators[pol].rawCalpar.tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnical"%pol)
    else:
        omnichisq = calibrators[pol].get_omnichisq()
        omnichisq.tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnichisq"%pol)
        omnifit = calibrators[pol].get_omnifit()
        omnifit.tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnifit"%pol)
        omnigains[pol[0]].tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnigain"%pol)
        if have_model_treasure and (sunpos[:,0] < -.1).all():
            #amp_cal, phs_cal
            omnichisq[:, 3:] = amp_cal
            omnichisq.tofile(oppath + '/' + dataano + '_' + ano + "_%s.absa.omnichisq"%pol)
            for i in range(phs_cal.shape[-1]):
                omnichisq[:, 3:] = phs_cal[...,i]
                omnichisq.tofile(oppath + '/' + dataano + '_' + ano + "_%s.absp%i.omnichisq"%(pol,i))


    flags[pol].tofile(oppath + '/' + dataano + '_' + ano + "_%s.omniflag"%pol)
    calibrators[pol].write_redundantinfo(oppath + '/' + dataano + '_' + ano + "_%s.binfo"%pol, overwrite=True)
    diag_txt = calibrators[pol].diagnose(data = data[p], additiveout = additiveout, flag = flags[pol], healthbar = healthbar, ubl_healthbar = ubl_healthbar, ouput_txt = True)
    text_file = open(oppath + '/' + dataano + '_' + ano + "_%s.diagtxt"%pol, "a")
    text_file.write(str(sys.argv)+'\n'+diag_txt)
    text_file.close()

    print "Done."
    sys.stdout.flush()
    if (treasurePath is not None) and (np.sum(flags[pol]) < calibrators[pol].nTime * calibrators[pol].nFrequency) and (sunpos[:,0] < -.1).all():
        print FILENAME + " MSG: updating treasure on %s %s in %s."%(dataano, pol, treasurePath),
        sys.stdout.flush()
        calibrators[pol].update_treasure(treasurePath, np.array(lst)*TPI/24., flags[pol], pol, verbose = True)
        print "Done."
        sys.stdout.flush()


for i in range(len(sol)):
    omnichisq[:, 3:] = sol[i]
    omnichisq.tofile(oppath + '/' + dataano + '_' + ano + ".cp%i.omnichisq"%(i))
omnichisq[:, 3:] = overall_angle
omnichisq.tofile(oppath + '/' + dataano + '_' + ano + ".cp2.omnichisq")

for p, pol in enumerate(crosspols.keys()):
    omnichisq[:, 3:] = cross_chisq[p]
    omnichisq.tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnichisq"%pol)

    omnifit[:, :, 6::2] = np.real(crossextract[p]).transpose((0,2,1))
    omnifit[:, :, 7::2] = np.imag(crossextract[p]).transpose((0,2,1))
    omnifit.tofile(oppath + '/' + dataano + '_' + ano + "_%s.omnifit"%pol)
################################compression######################################
if need_compression:
    compr_shape = list(data.shape)
    compr_shape[2] = 4
    compr_shape[2] = delay_compression
    compressed_data = np.zeros(compr_shape, dtype=data.dtype)
    compr_flags = np.ones(compr_shape, dtype='bool')
    compr_error = np.zeros((compr_shape[0], compr_shape[1], compr_shape[-1]), dtype='float32')
    for p, pol in enumerate(wantpols.keys()):
        print FILENAME + " MSG: starting compression on %s %s"%(dataano, pol),
        sys.stdout.flush()
        timer = time.time()

        calpar = np.ones((nfreq,nant), dtype='complex64')
        for bl in calibrators[pol].subsetbl:
            compr_flags[p, ..., bl] = False
        for t in range(calibrators[pol].nTime):
            print '.',
            sys.stdout.flush()
            calpar[:, calibrators[pol].subsetant] = (10**(calibrators[pol].rawCalpar[t, :, 3: (3 + calibrators[pol].nAntenna)])) * np.exp(1.j * calibrators[pol].rawCalpar[t, :, (3 + calibrators[pol].nAntenna): (3 + 2 * calibrators[pol].nAntenna)])
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore",category=RuntimeWarning)

                if not removeadditive:
                    d = omni.apply_calpar(data[p,t], calpar, calibrators[pol].totalVisibilityId)
                else:
                    d = omni.apply_calpar(data[p,t] - adds[pol][t], calpar, calibrators[pol].totalVisibilityId)
        #for t, d in enumerate(calibrators[pol].get_calibrated_data(data[p])):
            d[flags[pol][t]] = 0

            compr_result = omni.deconvolve_spectra2(d, ~flags[pol][t], (delay_compression+1)/2, correction_weight=1e-6)
            compressed_data[p, t] = float(delay_compression)/nfreq * compr_result[0]
            compr_error[p,t] = compr_result[1]
            compr_error_bar = float(delay_compression)/nfreq * np.array([compr_result[3][i,i] for i in range(delay_compression)])**.5
            compr_flags[p,t,compr_error_bar > min(compr_error_bar) * 1.05] = True
        print "Done. %fmin"%(float(time.time()-timer)/60.)
        sys.stdout.flush()
    for p, pol in enumerate(crosspols.keys()):
        print FILENAME + " MSG: starting compression on %s %s"%(dataano, pol),
        sys.stdout.flush()
        timer = time.time()

        for bl in c.subsetbl:
            compr_flags[p, ..., bl] = False
        for t in range(c.nTime):
            print '.',
            sys.stdout.flush()

        #for t, d in enumerate(calibrators[pol].get_calibrated_data(data[p])):
            flag = flags['xx'][t]|flags['yy'][t]|np.isnan(cdata2[p, t])|np.isinf(cdata2[p, t])
            cdata2[p, t][flag] = 0

            compr_result = omni.deconvolve_spectra2(cdata2[p, t], ~flag, (delay_compression+1)/2, correction_weight=1e-6)
            compressed_data[p + 2, t] = float(delay_compression)/nfreq * compr_result[0]
            compr_error[p + 2, t] = compr_result[1]
            compr_error_bar = float(delay_compression)/nfreq * np.array([compr_result[3][i,i] for i in range(delay_compression)])**.5
            compr_flags[p + 2, t, compr_error_bar > min(compr_error_bar) * 1.05] = True
        print "Done. %fmin"%(float(time.time()-timer)/60.)
        sys.stdout.flush()
    #######save uvs
    print FILENAME + " MSG: saving compressed uv files",
    sys.stdout.flush()
    if len(wantpols.keys()) == 1:
        compr_uvfile = oppath + '/' + dataano + "_%s.uvOE"%pol
    else:
        compr_uvfile = oppath + '/' + dataano + ".uvOE"


    uv=ap.miriad.UV(uvfiles[0])
    omni.exportuv(compr_uvfile, compressed_data, compr_flags, list(wantpols.values()) + list(crosspols.values()), jd, uv['inttime'], uv['sfreq'], uv['sdf'], sa.lat, sa.lon, overwrite=False, comment=ano)
    del(uv)
    print "Done."
    sys.stdout.flush()

if create_new_uvs:
    print FILENAME + " MSG: saving new uv files",
    sys.stdout.flush()
    infos = {}
    for pol in wantpols.keys():
        infos[pol[0]] = omni.read_redundantinfo(infopaths[pol])
    for uvfile in uvfiles:
        omni.apply_omnigain_uvs([uvfile], omnigains, calibrators[wantpols.keys()[0]].totalVisibilityId, infos, oppath, ano, adds= adds, verbose = True, comment = '_'.join(sys.argv), flags = flags, overwrite = overwrite_uvs)
    print "Done."
    sys.stdout.flush()

TIMER.tick()

if make_plots:
    for p,pol in zip(range(len(wantpols)), wantpols.keys()):
        #plt.subplot(2, len(wantpols), 2 * p + 1)
        #plot_data = (calibrators[pol].rawCalpar[:,:,2]/(len(calibrators[pol].Info.subsetbl)-calibrators[pol].Info.nAntenna - calibrators[pol].Info.nUBL))**.5
        #plt.imshow(plot_data, vmin = 0, vmax = (np.nanmax(calibrators[wantpols.keys()[0]].rawCalpar[:,:,2][~flags[wantpols.keys()[0]]])/(len(calibrators[pol].Info.subsetbl)-calibrators[pol].Info.nAntenna - calibrators[pol].Info.nUBL))**.5, interpolation='nearest')
        #plt.title('RMS fitting error per baseline')
        #plt.colorbar()

        plt.subplot(3, len(wantpols), 0 * len(wantpols) + p + 1)
        flag_plot_data = (calibrators[pol].rawCalpar[:,:,2]/(len(calibrators[pol].Info.subsetbl)-calibrators[pol].Info.nAntenna - calibrators[pol].Info.nUBL))**.5
        vmax = np.percentile(flag_plot_data.flatten()[~np.isnan(flag_plot_data.flatten())], 95)
        vmin = np.percentile(flag_plot_data, 5)
        flag_plot_data[flags[pol]] = np.nan
        plt.imshow(flag_plot_data, vmin = vmin, vmax = vmax, interpolation='none')
        plt.title('flagged RMS fitting error per baseline %s'%pol)
        plt.colorbar()

        plt.subplot(3, len(wantpols), 1 * len(wantpols) + p + 1)
        shortest_ubli = np.argsort(np.linalg.norm(calibrators[pol].ubl, axis = 1))[0]
        shortest_vis = np.angle(calibrators[pol].rawCalpar[:,:,3+2*calibrators[pol].nAntenna+2*shortest_ubli] + 1.j * (calibrators[pol].rawCalpar[:,:,3+2*calibrators[pol].nAntenna+2*shortest_ubli+1]))
        #shortest_vis[flags[pol]] = np.nan
        plt.imshow(shortest_vis, interpolation='nearest', vmin = -PI, vmax = PI)
        plt.title('phase of visibility fit on [%.1f, %.1f, %.1f] baseline %s'%(calibrators[pol].ubl[shortest_ubli][0], calibrators[pol].ubl[shortest_ubli][1], calibrators[pol].ubl[shortest_ubli][2], pol))
        plt.colorbar()

        plt.subplot(3, len(wantpols), 2 * len(wantpols) + p + 1)
        shortest_vis_flag = np.copy(shortest_vis)
        shortest_vis_flag[flags[pol]] = np.nan
        plt.imshow(shortest_vis_flag, interpolation='nearest', vmin = -PI, vmax = PI)
        plt.title('phase of visibility fit on [%.1f, %.1f, %.1f] baseline %s'%(calibrators[pol].ubl[shortest_ubli][0], calibrators[pol].ubl[shortest_ubli][1], calibrators[pol].ubl[shortest_ubli][2], pol))
        plt.colorbar()
    plt.show()
