#!/usr/bin/env python

import os
import lsd
import lsd.colgroup as colgroup
import lsd.bounds
import numpy as np
import scipy.stats.mstats
from itertools import izip
from collections import defaultdict
from lsd.join_ops import IntoWriter
from scipy.weave import inline
import logging
import cPickle
import surveys.ps1.calib as cal
import lsd.tui

nominal_exptimes = { 'g.0000': 43, 'r.0000': 40, 'i.0000': 45, 'z.0000': 30, 'y.0000': 30 }

def get_zp(mjd):
	self = get_zp
	if getattr(self, 'zps', None) is None:
		# Load the zero points
		fn = os.environ['ZPS_ARRAY']
		self.zps = cPickle.load(file(fn))
	zps = self.zps

	# Find the ZP for each MJD. If it does not exist,
	# return +inf
	idx = np.searchsorted(zps['mjd'], mjd)
	idx[idx == len(zps)] = 0
	found = zps['mjd'][idx] == mjd
#	print "Found:", sum(found)

	# Join the ZPs
	zp = np.zeros(len(idx), dtype='f4')
	zp[:] = zps['ZP'][idx]

	# Correct the ZP from nominal exposures to true exposures
	filterid = zps['filterid'][idx]
	exptime = zps['exptime'][idx]
	delta_mag = np.empty(len(exptime))
	for fid, et in nominal_exptimes.iteritems():
		ii = filterid == fid
		delta_mag[ii] = -2.5 * np.log10( exptime[ii] / et )
	zp -= delta_mag

	# Remove those that are flagged bad
	bad = ~found
	bad |= (zps['time_to_bad'][idx] < 10. / (24*60)) & ~zps['overlaps_sdss'][idx]
	zp[bad] = np.nan

#	print "OK/Not OK:", sum(~bad), sum(bad)

	return zp

def calc_objmag(qresult, db, tabname):
	"""
	Compute object magnitude from detections.

	Compute ndet, median, average, SIQR per detection.
	"""
	# PS1 bug workaround: Ignore the few objects that wound up on the south pole
	if qresult.static_cell & 0xFFFFFFFF00000000 == 0:
		logging.warning("Encountered the cell at the south pole. Dropping it.")
		return

	if True:
		ids = []
		for all_rows in colgroup.partitioned_fromiter(qresult, "_ID", 20*1000*1000, blocks=True):
			ids.append(calc_objmag_aux(all_rows, db, tabname))
			all_rows = None
		ids = np.concatenate(ids)
	else:
		all_rows = colgroup.fromiter(qresult, blocks=True)
		ids = calc_objmag_aux(all_rows, db, tabname)
		all_rows = None

	assert np.all(np.sort(ids) == np.unique(ids))

	import hashlib
	print qresult.static_cell, hashlib.md5(np.sort(ids)).hexdigest(), len(ids)

	yield len(ids)

def calc_objmag_aux(all_rows, db, tabname):
	# Sort
	all_rows.sort(["filterid", "_ID", "mag"])

	# Apply flat field corrections
	offs = cal.flat_offs(all_rows.chip_id, all_rows.x_psf, all_rows.y_psf, all_rows.mjd_obs, all_rows.filterid)
	all_rows.mag += offs

	# Prepare the output array
	objs, idx = np.unique(all_rows['_ID'], return_index=True)
	out = colgroup.ColGroup(
		dtype=[
			('obj_id', 'u8'),
			('ra', 'f8'), ('dec', 'f8'),
			('nmag', '5i2'), ('nmag_ok', '5i2'),
			('mean', '5f4'), ('stdev', '5f4'), ('err', '5f4'),
			('median', '5f4'), ('q25', '5f4'), ('q75', '5f4'),
			# Lightcurve samples
			('mag', '(5,10)f4'), ('magErr', '(5,10)f4'), ('mjd', '(5,10)f8')
		      ],
		      size=len(objs)
	      )
	out['obj_id'][:] = objs
	out['ra'][:]  = all_rows['ra'][idx]
	out['dec'][:] = all_rows['dec'][idx]

	# Pull out the arrays we'll be using
	(id_out, ra, dec, nmag, nmag_ok, mean, stdev, merr, median, q25, q75, lc_mag, lc_magErr, lc_mjd) = out.as_columns()
	id_in, mags, errs, filterid, flags, psf_qf, mjd_obs = ( getattr(all_rows, attr) for attr in ['_ID', 'mag', 'err', 'filterid', 'flags', 'psf_qf', 'mjd_obs'] )

	# Join the zero-point information
	zp = get_zp(all_rows.mjd_obs)

	# Convert filterid to index
	band = np.empty(len(all_rows), dtype='i4')
	for f, i in { 'g.0000': 0, 'r.0000': 1, 'i.0000': 2, 'z.0000': 3, 'y.0000': 4 }.iteritems():
		band[filterid == f] = i

	code = \
	"""
	#line 108 "objdata_weave.py"

	assert(Sid_out[0] == sizeof(*id_out));	// Make sure we've got a contiguous array

	uint32_t bad = PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_POOR | PM_SOURCE_MODE_SATSTAR | 
		PM_SOURCE_MODE_BLEND | PM_SOURCE_MODE_EXTERNAL | PM_SOURCE_MODE_BADPSF |
		PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | PM_SOURCE_MODE_CR_LIMIT 
		// | 
		// PM_SOURCE_MODE_EXT_LIMIT | PM_SOURCE_MODE_MOMENTS_FAILURE | 
		// PM_SOURCE_MODE_SIZE_SKIPPED | PM_SOURCE_MODE_BIG_RADIUS
		;

	// stream through the input arrays
	int size = Nmags[0];
	std::vector<double> mags1, wt, mjds;
	for(int i = 0, j = 0; i != size; i = j)
	{
		j = i;
		mags1.clear(); wt.clear(); mjds.clear();
		while(j != size && id_in[i] == id_in[j] && band[i] == band[j])
		{
			float mag = MAGS1(j);
			float err = ERRS1(j);
			float psf_qf = PSF_QF1(j);
			float zp = ZP1(j);
			double mjd = MJD_OBS1(j);
			uint32_t flag = FLAGS1(j);

			if(std::isfinite(mag) && std::isfinite(zp) && ((flag & bad) == 0) && (psf_qf > 0.85))
			{
				mjds.push_back(mjd);
				mag += zp;
				mags1.push_back(mag);
				// add 0.01 in quadrature to compensate for unrealistically small quoted errors
				double wtt = 1. / (err*err + 0.01*0.01);
				wt.push_back(wtt);
			}
			j++;
		}

		// find where to store
		int row = std::lower_bound(id_out, id_out + Nid_out[0], id_in[i]) - id_out;
		int col = band[i];
		assert(id_out[row] == id_in[i]);
		assert(0 <= row && row < Nid_out[0]);
		assert(0 <= col && col < 5);

		// store number of magnitudes
		   NMAG2(row, col) = j - i;
		NMAG_OK2(row, col) = mags1.size();

		if(!mags1.empty())
		{
			// median statistics
			Q252(row, col)    = gsl_stats_quantile_from_sorted_data(&mags1[0], 1, mags1.size(), 0.25);
			MEDIAN2(row, col) = gsl_stats_quantile_from_sorted_data(&mags1[0], 1, mags1.size(), 0.50);
			Q752(row, col)    = gsl_stats_quantile_from_sorted_data(&mags1[0], 1, mags1.size(), 0.75);

			// mean statistics
			MEAN2(row, col)  = gsl_stats_wmean(&wt[0], 1, &mags1[0], 1, mags1.size());
			STDEV2(row, col) = fabs(gsl_stats_wsd(&wt[0], 1, &mags1[0], 1, mags1.size()));	// I wrap it in fabs because for N=0 it returns a -0 (??)

			// mean error computed as 1./sum(wts)
			double w = 0.;
			for(unsigned i = 0; i != wt.size(); i++) { w += wt[i]; }
			MERR2(row, col) = 1. / sqrt(w);
		}
		
		// store the light curves
		for(unsigned i = 0; i != std::min(mags1.size(), (size_t)10); i++)
		{
			LC_MAG3(row, col, i) = mags1[i];
			LC_MAGERR3(row, col, i) = sqrt(1. / wt[i]);
			LC_MJD3(row, col, i) = mjds[i];
		}
	}
	"""
	inline(code,
		['id_out', 'nmag', 'nmag_ok', 'mean', 'stdev', 'merr', 'median', 'q25', 'q75', 'id_in', 'mags', 'errs', 'band', 'flags', 'zp', 'psf_qf', 'mjd_obs', 'lc_mag', 'lc_magErr', 'lc_mjd'],
		headers=['"pmSourceMasks.h"', '<cmath>', '<iostream>', '<algorithm>', '<vector>', '<gsl/gsl_statistics.h>', '<cassert>', '<algorithm>'],
		libraries=['gsl', 'gslcblas'],
		include_dirs=['.'],
		verbose=0,
		undef_macros=['NDEBUG'])

	# Write out the result
	ids = db.table(tabname).append(out, _update=True)
	assert np.all(ids == out.obj_id)
	return ids

db = lsd.DB("pdb")
bounds = None
#bounds = [ (lsd.bounds.rectangle(40, 40, 250, 80), lsd.bounds.intervalset((-np.inf, np.inf))) ]
#q = db.query("SELECT _ID, cal_psf_mag as mag, filterid FROM ps1_obj, ps1_det where (obj_id==6496442462481940483) & (filterid=='g.0000')")
#q = db.query("SELECT _ID, cal_psf_mag as mag, filterid FROM ps1_obj, ps1_det where (obj_id==6496442462481940483) | (obj_id==6496442462481940484)")

#q = db.query("SELECT _ID, cal_psf_mag as mag, cal_psf_mag_sig as err, filterid, flags FROM ps1_obj, ps1_det")
#writer = IntoWriter(db, "cal_mags WHERE obj_id |= obj_id")
#q = db.query("SELECT _ID, psf_mag as mag, psf_inst_mag_sig as err, filterid, flags, mjd_obs, chip_id, x_psf, y_psf, psf_qf FROM ps1_obj, ps1_det")
#writer = IntoWriter(db, "recalib_mags_ap WHERE obj_id |= obj_id")
q = db.query("SELECT _ID, ra, dec, psf_inst_mag as mag, psf_inst_mag_sig as err, filterid, flags, mjd_obs, chip_id, x_psf, y_psf, psf_qf FROM ps1_obj, ps1_det")

#q = db.query("select _ID, cal_psf_mag, filterid, u, g, r, i, z from ps1_obj, ps1_det, sdss where _ROWNUM==1")
nrows = 0
#kk = 0
with db.transaction():
	for rows_added in q.execute([(calc_objmag, db, 'averages')], group_by_static_cell=True, bounds=bounds):
		nrows += rows_added
		#kk += 1
		#import sys
		#print >>sys.stderr, "At", kk, nrows
		#if kk > 10: break
print "Total objects:", nrows