import numpy as np
from pixell import utils, mpi, bunch

def read_sim(fname):
	return bunch.Bunch(**np.load(fname, allow_pickle=True).reshape(-1)[0])

def yr2mjd(yr): return utils.ctime2mjd(157766400+(yr-1975)*utils.yr)

def fix_wrap(cov):
	ocov = cov.copy()
	imin = np.argmin(cov[0])
	for i in range(len(cov)-imin):
		ocov[i,imin+i:] = 0
	ocov = ocov.T
	for i in range(len(cov)-imin):
		ocov[i,imin+i:] = 0
	ocov = ocov.T
	return ocov

comm = mpi.COMM_WORLD
# Read in the data points
data = np.loadtxt("haystack_J2134-0153.csv", delimiter=",", skiprows=1).T
data[0] = yr2mjd(data[0])
# Read in the first sim to find the index correspondence with
# the data
sim  = read_sim("sims_haystack/J2134-0153_epoch1_sim_EMM_2.01_0.npy")
dinds, sinds = np.array(utils.crossmatch(data[0,:,None], sim.time[:,None], rmax=1, mode="first")).T
# Override time with higher-precision one from sims file
data[0,dinds] = sim.time[dinds]
# Now read in all the sims and build the covmat
simfiles = sorted(utils.glob("sims_haystack/J2134-0153_epoch1_sim_EMM_2.01_*.npy"))
nsim     = len(simfiles)
a  = 0
aa = 0
nread = 0
for si in range(comm.rank, nsim, comm.size):
	simfile = simfiles[si]
	try:
		sim     = read_sim(simfile)
	except:
		print("Error reading %s. Skipping" % (simfile))
		continue
	flux    = sim.flux[sinds]
	a      += flux
	aa     += flux[:,None]*flux[None,:]
	nread  += 1
	print("%6d/%d %s" % (si+1, nsim, simfile))
nread = comm.allreduce(nread)
a   = utils.allreduce(a,  comm)/nread
aa  = utils.allreduce(aa, comm)/nread
if comm.rank == 0:
	cov = aa - a[:,None]*a[None,:]
	# This messed up the eigenvalues.
	#cov = fix_wrap(cov)
	bunch.write("haystack_J2134-0153.hdf", bunch.Bunch(
		time=data[0,dinds], flux=data[1,dinds], dflux=data[2,dinds], cov=cov))