import fitsio
import numpy as np
import time
import math
from pydl import pydlutils
import sys
import mksampleutils

def photokey(darr):
  """
  SDSS magic photoid is some combo of these things, so they can be used as a unique key for each object.
  """
  return (darr['RUN'], darr['RERUN'], darr['CAMCOL'], darr['FIELD'], darr['ID'])

def match_sample(target_file,known_file,NSopt,LRopt,spall_file,output_file):
  """
  Same input as match_sample.pro, but added known_file as well rather than hard-coding it.
  Assuming full paths to these files are provided.
  Photo-obj matching is coded up but it doesn't actually work for BOSS, need to spherematch instead.
  """
  dd_targ = fitsio.read(target_file)
  glong, glat = mksampleutils.euler(dd_targ['RA'],dd_targ['DEC'])
  north = np.where(glat > 0)[0]
  south = np.where(glat <= 0)[0]
  print 'NS split',len(dd_targ),len(north),len(south)
  assert len(dd_targ) == len(north) + len(south)
  if NSopt == 0:
    dd_targ = dd_targ[north]
  else:
    dd_targ = dd_targ[south]

  ## LR split.
  # LRopt == 0: do nothing.
  left = np.where(dd_targ['RA'] <= 180.)[0]
  right = np.where(dd_targ['RA'] > 180.)[0]
  print 'LR split',len(dd_targ),len(left),len(right)
  assert len(dd_targ) == len(left) + len(right)

  if LRopt == -1:
    dd_targ = dd_targ[left]
  if LRopt == 1:
    dd_targ = dd_targ[right]

  print 'keeping this many gals:',len(dd_targ)

  #######  set up final data structure for galaxy catalog.
  ddtype = np.dtype([('RA', '>f8'), ('DEC', '>f8'), ('RUN', '>i4'), ('RERUN', '|S12'), ('CAMCOL', '>i4'), ('FIELD', '>i4'), ('ID', '>i4'), ('ICHUNK', '>i4'), ('IPOLY', '>i4'), ('ISECT', '>i4'), ('FRACPSF', '>f4', (5L,)), ('EXPFLUX', '>f4', (5L,)), ('DEVFLUX', '>f4', (5L,)), ('PSFFLUX', '>f4', (5L,)), ('MODELFLUX', '>f4', (5L,)), ('FIBER2FLUX', '>f4', (5L,)), ('R_DEV', '>f4', 5L), ('EXTINCTION', '>f4', (5L,)), ('PSF_FWHM', '>f4', (5L,)), ('AIRMASS', '>f4'), ('SKYFLUX', '>f4', (5L,)), ('EB_MINUS_V', '>f4'), ('IMAGE_DEPTH', '>f4', (5L,)), ('IMATCH', '>i4'), ('Z', '>f4'), ('WEIGHT_FKP', '>f4'), ('WEIGHT_CP', '>f4'), ('WEIGHT_NOZ', '>f4'), ('WEIGHT_STAR', '>f4'), ('WEIGHT_SEEING', '>f4'), ('WEIGHT_SYSTOT', '>f4'), ('NZ', '>f4'), ('COMP', '>f4'), ('CP_COMP', '>f4'), ('PLATE', '>i4'), ('FIBERID', '>i4'), ('MJD', '>i4'), ('NTILES', '>i2'), ('FINALN', '>i4'), ('FINALID', '>i4'), ('NASSIGNED', '>i4'), ('WHICHFINAL', '>i4'), ('TILE', '>i2', (3L,)), ('SPECTILE', '>i4'), ('COLLATEINDX', '>i4'), ('ICOLLIDED', '>i4'), ('COLL_GROUP_FINALID', '>i4', (10L,)), ('COLL_GROUP_BOSSTARGET1', '>i4', (10L,)), ('INGROUP', '>i4'), ('FIRSTGROUP', '>i4'), ('MULTGROUP', '>i4'), ('NEXTGROUP', '>i4')])

  ee = np.recarray(len(dd_targ),dtype=ddtype)
  for ff in ee.dtype.fields:
    if ff in dd_targ.dtype.fields:
      print 'copying over ',ff
      ee[ff] = dd_targ[ff]
  ## fill in blank values for the rest of the new fields.
  # weights.
  ee['WEIGHT_FKP'][:] = 1.0
  ee['WEIGHT_CP'][:] = 1.0
  ee['WEIGHT_NOZ'][:] = 1.0
  ee['WEIGHT_STAR'][:] = 1.0
  ee['WEIGHT_SEEING'][:] = 1.0
  ee['WEIGHT_SYSTOT'][:] = 1.0
  ## specobj info
  ee['PLATE'][:] = -10
  ee['FIBERID'][:] = -10
  ee['MJD'][:] = -10
  ee['CP_COMP'][:] = 1.0

  ee['Z'][:] = -10.0  ## mark things without a redshift.  

  ## then make a lookup dictionary for hte spall file.
  ## actually, I will just make the dictionary for the targets, and loop over spall.
  ## Stephen said this should take a couple minutes.

  specreruntype = '|S3'

  targ_objid = {}
  t1 = time.time()
  for i in range(len(ee)):
    ## native time for RERUN is '|S12', which disagrees with spAll.
    targ_objid[(ee['RUN'][i], ee['RERUN'][i].astype(specreruntype), ee['CAMCOL'][i], ee['FIELD'][i], ee['ID'][i])] = i
  t2 = time.time()
  print 'building RUN/RERUN/CAMCOL/FIELD/ID lookup took ',t2-t1,'seconds.'

  t1=time.time()
  dd_spec_known = fitsio.read(known_file,columns=['PLUG_RA','PLUG_DEC','SPECPRIMARY','ZWARNING','CLASS','Z','PLATE','FIBERID','MJD'])
  print 'read in ',len(dd_spec_known),'elements of known file.'
  dd_spec_known = dd_spec_known[np.where((dd_spec_known['ZWARNING'] == 0) & \
                                (dd_spec_known['SPECPRIMARY'] == 1) & (dd_spec_known['CLASS'] == 'GALAXY'))[0]]
  print 'known reduced to ',len(dd_spec_known),'galaxies'
  cnt = 0
  cntwacky = 0

  sepmax = 2./3600.
  print 'max sep in degrees:',sepmax

  photomatch = 0

  ## nevermind, should be spehrematching here, run/rerun/camcol/field/id still not given for spall_prim
  if photomatch == 1:
    for i in range(len(dd_spec_known)):
      photoid = photokey(dd_spec_known[i])
      if photoid in targ_objid:
        ee['IMATCH'][targ_objid[photoid]] = 2
        ee['Z'][targ_objid[photoid]] = dd_spec_known['Z'][i]
        ee['PLATE'][targ_objid[photoid]] = dd_spec_known['PLATE'][i]
        ee['FIBERID'][targ_objid[photoid]] = dd_spec_known['FIBERID'][i]
        ee['MJD'][targ_objid[photoid]] = dd_spec_known['MJD'][i]
        cnt += 1
        ## compute ang sep!
        sep = mksampleutils.angsepdeg(ee['RA'][targ_objid[photoid]], ee['DEC'][targ_objid[photoid]], dd_spec_known['PLUG_RA'][i], dd_spec_known['PLUG_DEC'][i])
        if sep > sepmax:
          print 'got a match, big separation?',sep, sep/sepmax
          cntwacky += 1
  
    print 'matched',cnt,'known redshifts.'
    print 'had ',cntwacky,'too big separation matches.'

  else:  ## use spherematching.
    ## all  hte inputs assumed in decimal degrees.
    ee_match, dd_spec_known_match, di = pydlutils.spheregroup.spherematch(ee['RA'], ee['DEC'], dd_spec_known['PLUG_RA'], dd_spec_known['PLUG_DEC'],matchlength=sepmax,maxmatch=1) 
    ee['IMATCH'][ee_match] = 2
    ee['Z'][ee_match] = dd_spec_known['Z'][dd_spec_known_match]
    ee['PLATE'][ee_match] = dd_spec_known['PLATE'][dd_spec_known_match]
    ee['FIBERID'][ee_match] = dd_spec_known['FIBERID'][dd_spec_known_match]
    ee['MJD'][ee_match] = dd_spec_known['MJD'][dd_spec_known_match]
    print 'matched',len(ee_match),'known redshifts'

  ## clean up some memory.
  del dd_spec_known
  del dd_spec_known_match

  dd_spec = fitsio.read(spall_file,columns=['RA','DEC','Z_NOQSO','PLATE','FIBERID','MJD','SPECPRIMARY','ZWARNING_NOQSO','CLASS_NOQSO','TILE','RUN','RERUN','CAMCOL','FIELD','ID'],iter_row_buffer=10000)
  t2=time.time()
  print 'spall read took',t2-t1,'seconds for ',len(dd_spec),'elements'
  dd_spec = dd_spec[np.where(dd_spec['SPECPRIMARY'] == 1)[0]]
  print 'downselected specprimary=1',len(dd_spec)

  photomatch = 0
  if photomatch == 1:
    ## this is missing 8500 objects that spherematch finds, why?  They seem to indeed match spatially but have different RUN/RERUN/CAMCOL/FIELD/ID.  bugginess in swap of photo reductions?
    for i in range(len(dd_spec)):
      photoid = photokey(dd_spec[i])
      if photoid in targ_objid:
        # Set the flags, assuming everything is a good match.
        # Below we change imatch to account for stars and redshift failures.
        ee['IMATCH'][targ_objid[photoid]] = 1
        if dd_spec['ZWARNING_NOQSO'][i] > 0:
          ee['IMATCH'][targ_objid[photoid]] = 7
        elif dd_spec['ZWARNING_NOQSO'][i] == 0 and dd_spec['CLASS_NOQSO'][i] == 'STAR  ':
          ee['IMATCH'][targ_objid[photoid]] = 4

          
        ee['Z'][targ_objid[photoid]] = dd_spec['Z_NOQSO'][i]
        ee['SPECTILE'][targ_objid[photoid]] = dd_spec['TILE'][i]
        ee['PLATE'][targ_objid[photoid]] = dd_spec['PLATE'][i]
        ee['FIBERID'][targ_objid[photoid]] = dd_spec['FIBERID'][i]
        ee['MJD'][targ_objid[photoid]] = dd_spec['MJD'][i]
        cnt += 1
        ## compute ang sep!
        sep = mksampleutils.angsepdeg(ee['RA'][targ_objid[photoid]], ee['DEC'][targ_objid[photoid]], dd_spec['RA'][i], dd_spec['DEC'][i])
        if sep > sepmax:
          print 'got a match, big separation?',sep, sep/sepmax
          cntwacky += 1
  
    print 'matched',cnt,'known redshifts.'
    print 'had ',cntwacky,'too big separation matches.'

  else:  ## use spherematching.
    ## all  hte inputs assumed in decimal degrees.
    ee_match, dd_spec_match, di = pydlutils.spheregroup.spherematch(ee['RA'], ee['DEC'], dd_spec['RA'], dd_spec['DEC'],matchlength=sepmax,maxmatch=1) 
    ee['IMATCH'][ee_match] = 1
    ee['SPECTILE'][ee_match] = dd_spec['TILE'][dd_spec_match]
    ee['Z'][ee_match] = dd_spec['Z_NOQSO'][dd_spec_match]
    ee['PLATE'][ee_match] = dd_spec['PLATE'][dd_spec_match]
    ee['FIBERID'][ee_match] = dd_spec['FIBERID'][dd_spec_match]
    ee['MJD'][ee_match] = dd_spec['MJD'][dd_spec_match]
    print 'matched',len(ee_match),'known redshifts'
    ## fix
    xxfail = np.where(dd_spec['ZWARNING_NOQSO'][dd_spec_match] > 0)[0]
    print 'found this many redshift failures:',len(xxfail)
    ee['IMATCH'][ee_match[xxfail]] = 7
    assert (dd_spec['ZWARNING_NOQSO'][dd_spec_match[xxfail]] > 0).all()
    ## CLASS_NOQSO has type '|S6' so i need to pad with zeros (probably a function exists to do this for you, or just approximate match?  anyway this seems to work.

    xxstar = np.where((dd_spec['ZWARNING_NOQSO'][dd_spec_match] == 0) & (dd_spec['CLASS_NOQSO'][dd_spec_match] == 'STAR  '))[0]
    ee['IMATCH'][ee_match[xxstar]] = 4
    assert (dd_spec['ZWARNING_NOQSO'][dd_spec_match[xxstar]] == 0).all()
    assert (dd_spec['CLASS_NOQSO'][dd_spec_match[xxstar]] == 'STAR  ').all()
    
  hi, ii = np.histogram(ee['IMATCH'],bins=np.arange(-0.5,10.5,1.0))
  ii = 0.5*(ii[:-1] + ii[1:])
  print 'histogram of imatch values'
  print hi
  print ii

  ## write the result out.
  fitsio.write(output_file,ee,clobber=True)

if __name__ == '__main__':

  ## these were run using photoid on spall rather than spherematch.
  #match_sample("../data/trimmed-collate-cmasspy-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=0,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="match-cmasspy-dr12v0-N-photoid.fits")
  #match_sample("../data/trimmed-collate-cmasspy-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=1,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="match-cmasspy-dr12v0-S-photoid.fits")


  #match_sample("../data/trimmed-collate-cmasspy-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=0,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="match-cmasspy-dr12v0-N-sphere.fits")
  #match_sample("../data/trimmed-collate-cmasspy-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=1,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="match-cmasspy-dr12v0-S-sphere.fits")

  match_sample("../data/trimmed-collate-cmass-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=0,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="../data/match-cmass-dr12v0-N.fits")
  match_sample("../data/trimmed-collate-cmass-dr12v0.fits","inputfiles/spAll_primary_good.fits",NSopt=1,LRopt=0,spall_file="inputfiles/spAll-v5_6_0.fits",output_file="../data/match-cmass-dr12v0-S.fits")