import argparse
parser = argparse.ArgumentParser()
parser.add_argument("ifiles", nargs="+")
parser.add_argument("ofile")
parser.add_argument("--nper", type=int, default=30000)
parser.add_argument("--nmul", type=int, default=2)
args = parser.parse_args()
import numpy as np, emcee
from pixell import utils, bunch

def fit_lin(dsets, P, amps, t0=0):
	# Fit A*cos+B*sin+C to the data sets assuming the period and relative amplitudes
	# are already known
	n   = 2+len(dsets)
	rhs = np.zeros(n)
	div = np.zeros((n,n))
	for di,dset in enumerate(dsets):
		phi  = 2*np.pi*(dset.time-t0)/P
		B    = np.array([amps[di]*np.sin(phi),amps[di]*np.cos(phi),phi*0+1])
		Nd   = dset.icov.dot(dset.flux)
		inds = [0,1,2+di]
		rhs[inds] += B.dot(Nd)
		ldiv = B.dot(dset.icov.dot(B.T))
		# There should be a better way to do this.
		# This is basically just div[inds,inds] = ldiv
		for i,ii in enumerate(inds):
			div[ii,inds] += ldiv[i]
	A = np.linalg.inv(div)
	a = A.dot(rhs)
	fit = bunch.Bunch(a=a, A=A, iA=div, P=P, amps=amps, t0=t0)
	return fit

# -2log lik = 2 log|2piN"| + (d-Pa)N"(d-Pa)
#  = ... + d'N"d + (a-â)'A"(a-â) - â'A"â
# Marginalize over a:
#  -2log lik = 2 log |2piN"| - 2 log |2piA"| + d'N"d - âA"â
# Jeffrey's prior = ignore |2piA"| part, I think. And we can ignore
# constants. We we are left with
#  -2log post = -âA"â + const

def fit_model(dsets, Prange, Arange, t0=0, nmul=2, nper=10000, nburn=10000):
	na      = 2+len(dsets)
	dummy   = np.zeros(na+na*na)
	def zip(P, amps): return np.concatenate([[P],amps[:-1]])
	def unzip(x): return x[0], np.concatenate([x[1:],x[:1]*0+1],0)
	def calc_logP(x):
		P, amps = unzip(x)
		if P < Prange[0] or P > Prange[1]: return -np.inf, dummy
		for amp in amps:
			if amp < Arange[0] or amp > Arange[1]: return -np.inf, dummy
		try: lin = fit_lin(dsets, P, amps, t0=t0)
		except np.linalg.LinAlgError: return -np.inf, dummy
		logP  = 0.5*lin.a.dot(lin.iA.dot(lin.a))
		if not np.isfinite(logP): return -np.inf, dummy
		#print(P, amps, logP)
		return logP, np.concatenate([lin.a, lin.A.reshape(-1)])
	ndim    = 1+len(dsets)-1
	nwalker = ndim*nmul
	# Initialize the walkers
	x0      = zip(np.mean(Prange), np.ones(len(dsets))) + np.random.randn(nwalker,ndim)*0.01
	sampler = emcee.EnsembleSampler(nwalker, ndim, calc_logP)
	sampler.run_mcmc(x0, nper+nburn, progress=True)
	samples = sampler.get_chain(discard=nburn, flat=True)
	blobs   = sampler.get_blobs(discard=nburn, flat=True)
	# Result
	res     = bunch.Bunch()
	res.P, res.amps = unzip(samples.T)
	res.a   = blobs.T[:na]
	res.A   = blobs.T[na:].reshape(na,na,-1)
	res.cos, res.sin = res.a[:2]
	res.offs= res.a[2:]
	# Recover the phase and overall amplitude from a
	hA      = np.moveaxis(np.linalg.cholesky(np.moveaxis(res.A,2,0)),0,2)
	avals   = np.einsum("abi,bi->ai", hA, np.random.randn(*res.a.shape))+res.a
	# sin(a-b) = sin(a)cos(b) - cos(a)sin(b), so avals[0] = cos(b), avals[1] = -sin(b)
	# b = atan2(-avals[1],avals[0])
	res.phase = utils.rewind(np.arctan2(-avals[1], avals[0]),ref="auto")
	# Want phase in fraction of period
	res.phase /= 2*np.pi
	res.amps *= np.sum(avals[:2]**2,0)**0.5
	return res

dsets = [bunch.read(ifile) for ifile in args.ifiles]
for dset in dsets:
	dset.icov = np.linalg.inv(dset.cov)

t0 = 51000
samps = fit_model(dsets, Prange=[1000,3000], Arange=[0.1,4],
		nmul=args.nmul, nper=args.nper, nburn=args.nper, t0=t0)
bunch.write(args.ofile, samps)

# Fit summary

#    ------- both ------  ---- ovro only ----
# P  1736.328 ±   14.594  1759.214 ±   64.428
# φ    -0.471 ±    0.020    -0.478 ±    0.025 # old phi def
# Ao    0.453 ±    0.052     0.464 ±    0.053
# Ah    0.835 ±    0.310