import argparse
parser = argparse.ArgumentParser()
parser.add_argument("ifile")
parser.add_argument("ofile")
args = parser.parse_args()
import numpy as np
from pixell import utils, bunch

# We want to set correlations beyond some distance
# to zero while keeping the matrix healthy. Simply
# zeroing out the corners led to a non-positive-definite
# matrix, so let's do this carefully.

def last(a, default=0):
	v = np.where(a)[0]
	return v[-1] if len(v) > 0 else default
def first(a, default=0):
	v = np.where(a)[0]
	return v[0] if len(v) > 0 else default

def fix_1d(data):
	iref   = len(data.flux)//2
	dt     = data.time-data.time[iref]
	covfun = data.cov[iref]
	# Zero out negative parts of covmat, and everything beyond that
	i1     = last (covfun[:iref]<0,0)+1
	i2     = first(covfun[iref:]<0,len(covfun)-iref)+iref
	covfun[:i1] = 0
	covfun[i2:] = 0
	print(covfun)
	# Reconstruct full covmat
	dts = data.time[:,None]-data.time
	cov = np.interp(dts, dt, covfun, left=0, right=0)
	cov = 0.5*(cov+cov.T)
	return cov

def fix_1d2(data):
	n      = len(data.time)
	iref   = n//2
	dt     = data.time-data.time[iref]
	covfun = data.cov[iref]
	# Zero out negative parts of covmat, and everything beyond that
	i1     = last (covfun[:iref]<0,0)+1
	i2     = first(covfun[iref:]<0,n-iref)+iref
	covfun[:i1] = 0
	covfun[i2:] = 0
	cov = np.zeros((n,n))
	for i in range(n):
		# Index of covfun[0] and covfun[1] into this row of the matrix, before bounds
		a1 = i-iref
		a2 = a1+n
		# Offset of start and end to keep us in bounds
		o1 = max( -a1,0)
		o2 = min(n-a2,0)
		cov[i,a1+o1:a2+o2] = covfun[o1:n+o2]
	cov = 0.5*(cov+cov.T)
	#from pixell import enmap
	#enmap.write_map("test4.fits", cov)
	#E,V = np.linalg.eigh(cov)
	#print(E)
	#1/0
	return cov

data = bunch.read(args.ifile)
data.cov = fix_1d2(data)
bunch.write(args.ofile, data)