import os
import numpy as np
import matplotlib.pyplot as plt

import astropy.units as u
from astropy.table import Table, Column
from astropy.coordinates import SkyCoord

import fitsio
from pydl.pydlutils.spheregroup import spheregroup

%matplotlib inline

LSLGAdir = os.getenv('LSLGA_DIR')

mindiameter = 0.25 # [arcmin]
linking_length = 2.5 # [arcmin]

"""
Explanation: Group Galaxy Catalog for the DR5 Gallery
The purpose of this notebook is to build a group catalog (using a simple friends-of-friends algorithm) from a diameter-limited (D25>5 arcsec) parent sample of galaxies defined and documented as part of the Legacy Survey Large Galaxy Atlas. 
Preliminaries
Import the libraries we need, define the I/O path, and specify the desired linking length (in arcminutes) and the minimum D(25) of the galaxy sample.
End of explanation
"""


suffix = '0.05'

ledafile = os.path.join(LSLGAdir, 'sample', 'leda-logd25-{}.fits'.format(suffix))
leda = Table.read(ledafile)

keep = (np.char.strip(leda['OBJTYPE']) != 'g') * (leda['D25'] / 60 > mindiameter)
leda = leda[keep]

keep = ['SDSS' not in gg and '2MAS' not in gg for gg in leda['GALAXY']]
#keep = np.logical_and( (np.char.strip(leda['OBJTYPE']) != 'g'), ~((leda['D25'] / 60 > 2.5) * (leda['BMAG'] > 16)) )
leda = leda[keep]
leda

fig, ax = plt.subplots()
ax.scatter(leda['RA'], leda['DEC'], s=1, alpha=0.5)

fig, ax = plt.subplots()
ax.hexbin(leda['BMAG'], leda['D25'] / 60, extent=(5, 20, 0, 20),
          mincnt=1)
ax.set_xlabel('B mag')
ax.set_ylabel('D(25) (arcmin)')

if False:
    these = (leda['RA'] > 200) * (leda['RA'] < 210) * (leda['DEC'] > 5) * (leda['DEC'] < 10.0)
    leda = leda[these]
    print(np.sum(these))

"""
Explanation: Read the parent HyperLeda catalog.
We immediately throw out objects with objtype='g' in Hyperleda, which are "probably extended" and many (most? all?) have incorrect D(25) diameters.  We also toss out objects with D(25)>2.5 arcmin and B>16, which are also probably incorrect.
End of explanation
"""


%time grp, mult, frst, nxt = spheregroup(leda['RA'], leda['DEC'], linking_length / 60.0)

npergrp, _ = np.histogram(grp, bins=len(grp), range=(0, len(grp)))
nbiggrp = np.sum(npergrp > 1).astype('int')
nsmallgrp = np.sum(npergrp == 1).astype('int')
ngrp = nbiggrp + nsmallgrp

print('Found {} total groups, including:'.format(ngrp))
print('  {} groups with 1 member'.format(nsmallgrp))
print('  {} groups with 2-5 members'.format(np.sum( (npergrp > 1)*(npergrp <= 5) ).astype('int')))
print('  {} groups with 5-10 members'.format(np.sum( (npergrp > 5)*(npergrp <= 10) ).astype('int')))
print('  {} groups with >10 members'.format(np.sum( (npergrp > 10) ).astype('int')))

"""
Explanation: Run FoF with spheregroup
Identify groups using a simple angular linking length.  Then construct a catalog of group properties.
End of explanation
"""


groupcat = Table()
groupcat.add_column(Column(name='GROUPID', dtype='i4', length=ngrp, data=np.arange(ngrp))) # unique ID number
groupcat.add_column(Column(name='GALAXY', dtype='S1000', length=ngrp))
groupcat.add_column(Column(name='NMEMBERS', dtype='i4', length=ngrp))
groupcat.add_column(Column(name='RA', dtype='f8', length=ngrp))  # average RA
groupcat.add_column(Column(name='DEC', dtype='f8', length=ngrp)) # average Dec
groupcat.add_column(Column(name='DIAMETER', dtype='f4', length=ngrp))
groupcat.add_column(Column(name='D25MAX', dtype='f4', length=ngrp))
groupcat.add_column(Column(name='D25MIN', dtype='f4', length=ngrp))

leda_groupid = leda.copy()
leda_groupid.add_column(Column(name='GROUPID', dtype='i4', length=len(leda)))
leda_groupid

"""
Explanation: Populate the output group catalog
Also add GROUPID to parent catalog to make it easier to cross-reference the two tables.  D25MAX and D25MIN are the maximum and minimum D(25) diameters of the galaxies in the group.
End of explanation
"""


smallindx = np.arange(nsmallgrp)

ledaindx = np.where(npergrp == 1)[0]
groupcat['RA'][smallindx] = leda['RA'][ledaindx]
groupcat['DEC'][smallindx] = leda['DEC'][ledaindx]
groupcat['NMEMBERS'][smallindx] = 1
groupcat['GALAXY'][smallindx] = np.char.strip(leda['GALAXY'][ledaindx])
groupcat['DIAMETER'][smallindx] = leda['D25'][ledaindx] # [arcsec]
groupcat['D25MAX'][smallindx] = leda['D25'][ledaindx]   # [arcsec]
groupcat['D25MIN'][smallindx] = leda['D25'][ledaindx]   # [arcsec]

leda_groupid['GROUPID'][ledaindx] = groupcat['GROUPID'][smallindx]

"""
Explanation: Groups with one member--
End of explanation
"""


bigindx = np.arange(nbiggrp) + nsmallgrp

coord = SkyCoord(ra=leda['RA']*u.degree, dec=leda['DEC']*u.degree)

def biggroups():
    for grpindx, indx in zip(bigindx, np.where(npergrp > 1)[0]):

        ledaindx = np.where(grp == indx)[0]
        _ra, _dec = np.mean(leda['RA'][ledaindx]), np.mean(leda['DEC'][ledaindx])
        d25min, d25max = np.min(leda['D25'][ledaindx]), np.max(leda['D25'][ledaindx])
    
        groupcat['RA'][grpindx] = _ra
        groupcat['DEC'][grpindx] = _dec
        groupcat['D25MAX'][grpindx] = d25max
        groupcat['D25MIN'][grpindx] = d25min

        groupcat['NMEMBERS'][grpindx] = len(ledaindx)
        groupcat['GALAXY'][grpindx] = ','.join(np.char.strip(leda['GALAXY'][ledaindx]))
        leda_groupid['GROUPID'][ledaindx] = groupcat['GROUPID'][grpindx]

        # Get the distance of each object from the group center.
        cc = SkyCoord(ra=_ra*u.degree, dec=_dec*u.degree)
        diameter = 2 * coord[ledaindx].separation(cc).arcsec.max()

        groupcat['DIAMETER'][grpindx] = np.max( (diameter*1.02, d25max) )

%time biggroups()

leda_groupid

groupcat

ww = np.where(groupcat['NMEMBERS'] >= 2)[0]
fig, ax = plt.subplots()
ax.scatter(groupcat['RA'][ww], groupcat['DEC'][ww], s=1, alpha=0.5)

groupfile = os.path.join(LSLGAdir, 'sample', 'leda-logd25-{}-groupcat.fits'.format(suffix))
print('Writing {}'.format(groupfile))
groupcat.write(groupfile, overwrite=True)

ledafile_groupid = os.path.join(LSLGAdir, 'sample', 'leda-logd25-{}-groupid.fits'.format(suffix))
print('Writing {}'.format(ledafile_groupid))
leda_groupid.write(ledafile_groupid, overwrite=True)

"""
Explanation: Groups with more than one member--
End of explanation
"""


