"""
                      .-.
                 .--.(   ).--.
      <-.  .-.-.(.->          )_  .--.
       `-`(     )-'             `)    )
         (o  o  )                `)`-'
        (      )                ,)
        ( ()  )                 )
         `---"\    ,    ,    ,/`
               `--' `--' `--'
                |  |   |   |
                |  |   |   |
                '  |   '   |     AnneyA
"""
from __future__ import print_function, division
import numpy as np
from astropy import units as u
import matplotlib.pyplot as plt
from matplotlib import patches
from astropy.cosmology import WMAP9 as cosmo
from sklearn.neighbors import BallTree
import pandas as pd
pd.set_option('display.max_rows', 10)
from SciServer import CasJobs
from SciServer import LoginPortal
from io import StringIO
import urllib
plt.rc('font',family='serif')

# get access token
my_access_token = LoginPortal.getToken()

# get satellite data from CasJobs
#query_text = "SELECT * from rasalhague.anneya.satellites0_50__25_36DR11"
query_text = "SELECT * from mydb.satellites10_20__0_10DR11"
response = CasJobs.executeQuery(query_text, "mydb", token = my_access_token)
output = StringIO(response.read())

# put the query output in a pandas dataframe
sats1 = pd.read_csv(output)
sats1

# Since the satellite query was run in patches, get more or them with another call to mydb
#query_text = "SELECT * from rasalhague.anneya.satellites0_50__25_36DR11"
query_text = "SELECT * from mydb.satellites0_10__0_10DR11"
response = CasJobs.executeQuery(query_text, "mydb", token = my_access_token)
output = StringIO(response.read())

# put the query output in a pandas dataframe
sats2 = pd.read_csv(output)
sats = sats1.append(sats2,ignore_index=True)
sats

# Get the unique host pairs too
query_text = "select * from rasalhague.anneya.uniquePairsDR12NoNeighborsv2_0_50__25_36specphoto "+\
             "where ra between 0 and 20 and dec between 0 and 10"
response = CasJobs.executeQuery(query_text, "mydb", token = my_access_token)
output = StringIO(response.read())
pairs = pd.read_csv(output)
pairs

# Plot the pair hosts and the satellites of the primaries
plt.scatter(sats['satRa'], sats['satDec'], s=14, color='0.75', alpha=0.75, label='Satellites')
plt.scatter(sats['ra'], sats['dec'], s=8, color='red', label='Primary Hosts')
plt.scatter(sats['ra2'], sats['dec2'], s=8, color='black',  label='Secondary Hosts')

# Plot a line connecting each pair
for p in pairs.itertuples():
    plt.plot([p.ra,p.ra2],[p.dec,p.dec2],color='black',lw=1)

plt.gca().set_aspect('equal')
plt.xlim(0,20), plt.ylim(0,10)
plt.xlabel('ra'), plt.ylabel('dec')
plt.legend(scatterpoints=1, loc='upper right', fontsize=12, \
          fancybox=True, shadow=True)
plt.gcf().set_size_inches(15,15)
plt.grid()

# Add new coordinates for each satellite, moving the origin to the position of the primary host
sats['shiftRa'] = pd.Series(sats['satRa']-sats['ra'], index=sats.index)
sats['shiftDec'] = pd.Series(sats['satDec']-sats['dec'], index=sats.index)

# Rotate the x-axis to lie on the line connecting the host pair
sats['rotTheta'] = pd.Series(-1.*np.arctan2((sats['dec2']-sats['dec']),(sats['ra2']-sats['ra'])),index=sats.index)
sats['rotRa'] = np.cos(sats['rotTheta'])*sats['shiftRa'] - np.sin(sats['rotTheta'])*sats['shiftDec']
sats['rotDec'] = np.sin(sats['rotTheta'])*sats['shiftRa'] + np.cos(sats['rotTheta'])*sats['shiftDec']

# Calculate angle of each satellite relative to pair
sats['theta'] = pd.Series(np.arctan2(sats['rotDec'],sats['rotRa'])*180./np.pi,index=sats.index)

# Add position of each satellite relative to the host in physical units
sats['xPhys'] = pd.Series(sats['d0']*np.cos(sats['theta']), index=sats.index)
sats['yPhys'] = pd.Series(sats['d0']*np.sin(sats['theta']), index=sats.index)
sats['thetabs'] = pd.Series(np.abs(sats['theta']),index=sats.index)

pd.unique(sats[np.abs(sats['rotTheta']-0.5)<0.1]['objId'])

testPair = sats[sats['objId']==pairs['objId'][1]]
#plt.scatter(testPair['shiftRa'],testPair['shiftDec'],color='blue',alpha=0.3,s=30,\
#            label='Satellite positions before rotation')
plt.scatter(testPair['rotRa'],testPair['rotDec'],color='red',alpha=0.5,s=75,marker='*',\
            label='Satellite positions after rotation')

plt.scatter([0],[0],color='0.4',s=60)
plt.scatter(testPair['ra2']-testPair['ra'],testPair['dec2']-testPair['dec'],color='0.4',s=60)
sep = np.array(np.sqrt((testPair['ra2']-testPair['ra'])**2+(testPair['dec2']-testPair['dec'])**2))
plt.scatter([sep],0.*testPair['ra2'],\
            color='0.4',s=60,label=None)
rotAngle = 1.0*np.array(testPair['rotTheta'])[0]
if rotAngle<0:
    t1 = 0.
    t2 = np.abs(rotAngle)
else:
    t1 = 2*np.pi-rotAngle
    t2 = 0.

print(rotAngle*180/np.pi)

el = patches.Arc((0,0), 2*sep[0],2*sep[0], angle=0.0, theta2=180.*t2/np.pi, theta1=180.*t1/np.pi,lw=3,\
                 color='0.6',alpha=0.66)
plt.grid()
plt.gca().add_patch(el)

ps = []
for p in testPair.itertuples():
    sep = np.sqrt(p.shiftRa**2 + p.shiftDec**2)
    myTheta = np.arctan2(p.shiftDec,p.shiftRa)
    if myTheta<0:
        myTheta = 2*np.pi + myTheta
    if rotAngle<0:
        t2 = myTheta
        t1 = myTheta + rotAngle
    else:
        t1 = myTheta
        t2 = myTheta + rotAngle
    ps.append(patches.Arc((0,0), 2*sep,2*sep, angle=0.0,theta2=180.*t2/np.pi, theta1=180.*t1/np.pi ,\
                          lw=2,color='purple',alpha=0.3))
for p in ps:
    plt.gca().add_patch(p)
plt.gca().set_aspect('equal')
plt.gcf().set_size_inches(15,15)
plt.legend(loc='best',scatterpoints=1)