import matplotlib.ticker
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
from adjustText import adjust_text
from matplotlib import font_manager

##This code creates plots of TNO dynamics and characteristics from csv data files derived from Johnston's Archive (size data) and JPL's Small Body Dynamics database (orbital data).
#A lot of things are customizable if you change the values below.

##Plotting options and parameters
#Name of the files
johnstonfile='tnos.csv'
jplfile='sbdb_query_results(1).csv'

#Font and background layout
matplotlib.rcParams['font.sans-serif'] = "Lato"
plt.rcParams['legend.title_fontsize'] = 'large'
plt.style.use('dark_background')

#Figure layout and size
fig = plt.figure(figsize=(14, 10))
ax = fig.add_axes([0.05,0.055,0.935,0.85])
colorlegendloc=[0,1]
sizelegendloc=[1,0]
titleypos=1.06

#Labels
fileout='kbos-testai-cc'
maintitle= 'The Kuiper Belt'
xlabel = 'Semi-major axis (AU)'
secxlabel = 'Period (years)'
ylabel = 'Inclination (degrees)'
grouptitle = 'Dynamical groups'
sizetitle = 'Diameter (km)'
restitle = 'Resonances'
sizecircles=[50,100,200,500,1000,2000] #sizes of the size circles
secax = True

#Params and axis limits. Reasonable values for ylim:
#perihelion: [23.5,82], minor ticks every 5.
#eccentricity: [0,1], [0,0.4] (KBOs), minor ticks every 0.05
#inclination: [-4,115], [-2,50] (KBOs), minor ticks every 5 or 10
paramx='a_y'
paramy='i_y'
axis_doublelog=False #make a log(log(x)) axis for semi-major
xlim=[37.5,50]
ylim=[-1.5,50]
xticks=[40,50,60,70,80,90,100,200,300,400,500,600,700,800,900,1000,2000,3000]
xticklabels=['40','50','60','','80','','100','200','','400','','600','','','','1,000','2,000','3,000']

secxticks=[200,300,400,500,600,700,800,900,1000,2000,3000,4000,5000,6000,7000,8000,9000,10000,20000,30000,40000,50000,60000,70000,80000,90000,100000]
secxticklabels=['200','300','400','500','600','','800','','1,000','2,000','','4,000','','6,000','','','','10,000','20,000','','','50,000','','','','','100,000']
minorxtickspace=0.5
minorytickspace=5
minorsecaxtickspace=5

#circle parameters
sizescale=np.sqrt(10)/150 #size of circles
circlewidth=1/600 #width of circle edges
maxcirclewidth=2 #maximum width
labelsize=1.8 #text size scale of labels

#Colors for the groups. If you want to label Haumeids, add a separate class_dict entry for Haumea, and change color_dict to match.
class_dict = {'Resonant TNO':'red', 'Plutino':'orange', 'Cold cubewano':'blue', 'Hot cubewano':'deepskyblue', 'Haumeid': 'mediumslateblue', 'Other TNO':'#8b2be2', 'Scattered disk':'lightgrey', 'Centaur':'lawngreen'}

#Defining the main resonances to be marked. These won't be shown if they lie outside the axis bounds.
neptuneorbit=30.07 #AU
mainres=[1,5/4, 4/3, 7/4, 5/3, 7/3, 2, 5/2, 7/2, 9/2, 3, 4, 5, 6]
minorres=[7,8,9,10,11,12,11/2,8/3,10/3,11/3,9/4,11/4,13/4,7/5,8/5,9/5,11/5,12/5,11/6,10/7,11/7,12/7,11/9]
mainresstr=['1:1','5:4','4:3','7:4','5:3','7:3','2:1','5:2','7:2','9:2','3:1','4:1','5:1','6:1']
minorresstr=['7:1','8:1','9:1','10:1','11:1','12:1','11:2','8:3','10:3','11:3','9:4','11:4','13:4','7:5','8:5','9:5','11:5','12:5','11:6','10:7','11:7','12:7','11:9']
toplabels=True

#Resonance tick/label params
restick=[-1.5,-0.75]
reslabelloc=-0.5
restitleloc=[27.5,0]

#manual label adjustments go here. This is going to be different for each plot, good luck!
xoffset=[0,-1.5,0.5,-30,-0, -2,0,0,-1,-1.5, 0,-60,-31,-0.5,-1.5, 24,1.5,-2,1]
yoffset=[-1,-21.5,-3,0,-11, 0,-0.5,-36,-3,-24, -7,4,4.5,0,-1.5, -23,-4,-21,0]

#arrows: play around with shrinkB until the arrow terminates on the *outside* of the circle for readability.
arrowprops=[None for i in range(23)]
arrowprops[7]={'arrowstyle':'-', 'color':'blue', 'lw': 0.5, 'relpos': (0.34,0.6), 'shrinkB':3.5}
arrowprops[17]={'arrowstyle':'-', 'color':'blue', 'lw': 0.5, 'relpos': (0.33,0.6), 'shrinkB':3.4}
arrowprops[15]={'arrowstyle':'-', 'color':'orange', 'lw': 0.5, 'relpos': (0.1,0.5), 'shrinkB':9}


##Main Code

pd.options.display.max_rows=100

#read in data, filter by high condition code, make sure it's in both johnston's and JPL
johnston=pd.read_csv(johnstonfile, delimiter=";",skipinitialspace=True)
johnston['pdes']=johnston['number'].str.strip(to_strip="()").fillna(johnston['provisional'])
jpl=pd.read_csv(jplfile)
tnos=pd.merge(johnston, jpl, how='inner', on='pdes')
tnofil=tnos[tnos['condition_code']<=7]

#Filter the notable TNOs that should be labeled. This is all custom.
largetnos=tnofil[(tnofil['diameter_x']>200)]
largetnos=largetnos[((largetnos['diameter_x']>700) | (~tnofil['name_x'].isna()))]
largetnos=largetnos[~largetnos['name_x'].isin(['Mbabamwanawaresa'])]
largetnos=largetnos[(largetnos['a_y']>=37.5) & (largetnos['a_y']<=50) & (largetnos['e_y']<0.4)]
#largetnos=tnofil[(tnofil['diameter_x']>210) | ((tnofil['name_x'].isin(['Alicanto'])) | tnofil['dynamics'].isin(['Sednoid']))]
#largetnos=tnofil[(tnofil['diameter']>150)]
#largetnos=largetnos[(~largetnos['dynamics'].isin(['other TNO'])) | ((largetnos['diameter_x']>499) & (largetnos['q_y']>50)) | (largetnos['diameter_x']>1000)]
#largetnos=largetnos[(~largetnos['dynamics'].isin(('cubewano-hot', 'cubewano-cold', 'Haumea', 'cubewano', 'plutino', 'SDO', 'twotino', 'Centaur', 'Nep Trj L4', 'Nep Trj L5')) & ~largetnos['dynamics'].str.contains('res')) | (largetnos['diameter_x']>600)|((~tnofil['name_x'].isna()))]
#largetnos=largetnos[(~largetnos['dynamics'].isin(['SDO'])) | (largetnos['diameter_x']>700) | ((~tnofil['name_x'].isna())) | (largetnos['a_y']>80)]
#largetnos=largetnos[(~largetnos['dynamics'].isin(('cubewano-hot', 'cubewano-cold', 'Haumea', 'cubewano', 'plutino'))|(tnofil['diameter_x']>1500))]
#largetnos=largetnos[~((largetnos['dynamics']=='SDO') & (largetnos['name_x'].isna()) & (largetnos['a_y']<85))]
largetnos['name_x']=largetnos['name_x'].fillna(largetnos['provisional'])  #add provisional desig to name to those with no names yet
print(largetnos[['name_x','provisional','diameter_x']])

#more subgroups for analysis
namedtnos=tnos[~tnos['name_x'].isna()]
p9tnos=tnofil[(tnofil['q_x']>=35) & (tnofil['a_x']>=200)]
large=tnos.sort_values(by= 'diameter_x', ascending=False)

largetnos=largetnos.reset_index()
namedtnos=namedtnos.reset_index()
#print(tnos.sort_values(by='t_jup').head(100)[['number','name_x','t_jup']])

color_dict = {'cubewano-cold': class_dict['Cold cubewano'], 'cubewano-hot': class_dict['Hot cubewano'], 'cubewano': class_dict['Hot cubewano'], 
              'plutino': class_dict['Plutino'], 'twotino': class_dict['Resonant TNO'], 'Nep Trj L4': class_dict['Resonant TNO'], 'Nep Trj L5': class_dict['Resonant TNO'],
              'res 1:3': class_dict['Resonant TNO'], 'res 1:4': class_dict['Resonant TNO'], 'res 1:5': class_dict['Resonant TNO'], 'res 1:6': class_dict['Resonant TNO'], 
              'res 1:7': class_dict['Resonant TNO'], 'res 1:8': class_dict['Resonant TNO'], 'res 1:9': class_dict['Resonant TNO'], 'res 1:10': class_dict['Resonant TNO'], 'res 1:11': class_dict['Resonant TNO'],
              'res 2:5': class_dict['Resonant TNO'], 'res 2:7': class_dict['Resonant TNO'], 'res 2:9': class_dict['Resonant TNO'], 'res 2:11': class_dict['Resonant TNO'], 
              'res 3:4': class_dict['Resonant TNO'], 'res 3:5': class_dict['Resonant TNO'], 'res 3:7': class_dict['Resonant TNO'], 'res 3:8': class_dict['Resonant TNO'], 'res 3:10': class_dict['Resonant TNO'], 'res 3:11': class_dict['Resonant TNO'],
              'res 4:5': class_dict['Resonant TNO'], 'res 4:7': class_dict['Resonant TNO'], 'res 4:9': class_dict['Resonant TNO'], 'res 4:11': class_dict['Resonant TNO'], 'res 4:13': class_dict['Resonant TNO'], 
              'res 5:7': class_dict['Resonant TNO'], 'res 5:8': class_dict['Resonant TNO'], 'res 5:9': class_dict['Resonant TNO'], 'res 5:11': class_dict['Resonant TNO'], 'res 5:12': class_dict['Resonant TNO'],
              'res 6:11': class_dict['Resonant TNO'], 'res 7:10': class_dict['Resonant TNO'], 'res 7:11': class_dict['Resonant TNO'], 'res 7:12': class_dict['Resonant TNO'], 'res 9:11': class_dict['Resonant TNO'],
               "SDO": class_dict['Scattered disk'], 'Haumea': class_dict['Haumeid'], "ESDO": class_dict['Scattered disk'], "EDDO": class_dict['Scattered disk'], "Sednoid": class_dict['Scattered disk'],
               'Centaur': class_dict['Centaur'],'Apollo': class_dict['Centaur'], 'Amor': class_dict['Centaur'], 'unusual': class_dict['Centaur'], 'Damocloid': class_dict['Centaur'], 'other TNO': class_dict['Other TNO'], 
               'Ura Trj L4': class_dict['Resonant TNO'], 'comet': class_dict['Centaur']}

colors=tuple(map(color_dict.get, tnofil['dynamics']))
largecolors=tuple(map(color_dict.get,largetnos['dynamics']))
namedcolors=tuple(map(color_dict.get,namedtnos['dynamics']))

##Plotting

#defining fancy log formatters for various stuff (those end up unused)

class myformatter(matplotlib.ticker.LogFormatter):

    def _num_to_string(self, x, vmin, vmax):

        if x > 1000000:
            s = '%1.0e' % x
        elif x < 1 and x >= 0.001:
            s = f'{x:n}'
        elif x < 0.001:
            s = '%1.0e' % x
        else:
            s = f'{x:n}'
        return s
    
def my_locs(self, locs=None):

        b = self._base
        c = np.geomspace(1, b, int(b)//int(locs) + 1)
        self._sublabels = set(np.round(c))

#Title and axis labels.
ax.set_title(maintitle,size='24',y=titleypos)
ax.set_xlabel(xlabel, size=16)
ax.set_ylabel(ylabel, size=16)


#formatter shenanigans for testing, not used in final plot

#formatter = myformatter(labelOnlyBase=False, minor_thresholds=(5, 2.5))
#formatter2 = myformatter(labelOnlyBase=False, minor_thresholds=(5, 2.5))
#fmt = matplotlib.ticker.StrMethodFormatter("{x:g}")
#ax.get_xaxis().set_minor_formatter(formatter)
#ax.get_xaxis().set_major_formatter(fmt)

#Double-log x scale and ticks for semi-major axis. 

if axis_doublelog==True:
    ax.set_xlim(xlim)
    ax.set_xscale("functionlog", functions=(
        lambda x: np.log10(np.log10(x)),
        lambda x: 10**((10**(x)))))
    ax.set_xticks(xticks, labels=xticklabels)
    #set length of minor ticks to same as major to avoid wonkiness
    ax.tick_params(which='minor', length=3.5)

else:
    ax.set_xlim(xlim)
    ax.xaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(minorxtickspace))

#Linear y scale and ticks.
ax.set_ylim(ylim)
ax.yaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(minorytickspace))

#Defining the secondary x scale for orbital period, with ticks and label; this will be a**(3/2).
if secax==True:
    def period(x):
        return x**(3/2)
    def invper(x):
        return x**(2/3)
    secax = ax.secondary_xaxis('top', functions=(period, invper))
    secax.set_xlabel(secxlabel)

    if axis_doublelog==True:
        secax.set_xscale("functionlog", functions=(
            lambda x: np.log10(np.log10(x**(2/3))),
            lambda x: (10**((10**(x))))**(3/2)))
        secax.set_xticks(secxticks, labels=secxticklabels)
        secax.tick_params(which='minor', length=3.5)
    else:
        secax.set_xscale("function", functions=(
            lambda x: x**(2/3),
            lambda x: x**(3/2)))
        secax.xaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(minorsecaxtickspace))

#Uncomment this if you insist on using a linear scale for a. This will create reasonable ticks.
#ax.set_xticks(np.concat((np.arange(30,110,10),np.arange(150,2350,50))))
#secax.set_xticks(np.concat(([160],np.arange(200,1100,100), np.arange(1500,5500,500),np.arange(6000,11000,1000), np.arange(12000,42000,2000),np.arange(40000,115000,5000))))


#Plot the circles for each TNO
normalpoints=ax.scatter(tnofil[paramx], tnofil[paramy], s=(tnofil['diameter_x']*sizescale)**2, facecolors='none', edgecolors=colors, linewidths=np.minimum(maxcirclewidth,tnofil['diameter_x']*circlewidth))

#Plot the labels on the notable tnos
textsize=np.floor(labelsize*np.log(largetnos['diameter_x']))
diameters=largetnos['diameter_x']*sizescale*(np.sqrt(2250)/100)

#generate labels
ann=[ax.annotate(largetnos['name_x'][i],(largetnos[paramx][i], largetnos[paramy][i]),xytext=(diameters[i]+xoffset[i], diameters[i]+yoffset[i]), textcoords='offset points',size=str(textsize[i]), color=largecolors[i], arrowprops=arrowprops[i], annotation_clip=False, bbox=dict(pad=-2, facecolor="none", edgecolor="none")) for i in range(largetnos.shape[0])]

#Mark the main resonance positions
#Make resonance lines, orange for plutinos. Note it's currently given in data coordinates, but you can use axis coordinates like so to avoid confusion:
#ax.vlines([1, 2], 0, 0.2, transform=ax.get_xaxis_transform(), colors='r')
ax.vlines(neptuneorbit*(3/2)**(2/3), restick[0], restick[1], colors='orange', label='3:2')
[ax.vlines(neptuneorbit*(mainres[i])**(2/3), restick[0], restick[1], colors='red') for i in range(len(mainres))]
#[ax.vlines(neptuneorbit*(minorres[i])**(2/3), ylim[0], ylim[0]+ylim[1]/60, colors='red', lw=0.5) for i in range(len(minorres))]

#Make resonance labels (again, in data coordinates)
ax.text(neptuneorbit*(3/2)**(2/3), reslabelloc, '3:2', color='orange',horizontalalignment='center')
[ax.text(neptuneorbit*(mainres[i])**(2/3), reslabelloc, mainresstr[i], color='red',horizontalalignment='center') for i in range(len(mainres))]
if toplabels==True:
    ax.vlines(neptuneorbit*(3/2)**(2/3), ylim[1]-(restick[0]-ylim[0]), ylim[1]-(restick[1]-ylim[0]), colors='orange', label='3:2')
    [ax.vlines(neptuneorbit*(mainres[i])**(2/3), ylim[1]-(restick[0]-ylim[0]), ylim[1]-(restick[1]-ylim[0]), colors='red') for i in range(len(mainres))]
    ax.text(neptuneorbit*(3/2)**(2/3), ylim[1]-(reslabelloc-ylim[0]), '3:2', color='orange',horizontalalignment='center', verticalalignment='top')
    [ax.text(neptuneorbit*(mainres[i])**(2/3), ylim[1]-(reslabelloc-ylim[0]), mainresstr[i], color='red',horizontalalignment='center', verticalalignment='top') for i in range(len(mainres))]

#[ax.text(neptuneorbit*(minorres[i])**(2/3), ylim[0]+ylim[1]/50, minorresstr[i], color='red',horizontalalignment='center', size='small') for i in range(len(minorres))]

#Resonance text (again, in data coordinates)
ax.text(restitleloc[0], restitleloc[1], restitle, color='red',size='large')


#creating the legends
circles=[]
circle2=[]

#make some tiny circles for labeling the colors, then put the legend box in a nice place. I used 375 km circles, but you can change this.
for value in class_dict.values():
    circles.append(plt.Line2D([], [], color='None', marker='o', markersize=375*sizescale,  markeredgecolor=value,mew=np.minimum(maxcirclewidth, 375*circlewidth)))

#circles for labeling the sizes, in white
for size in sizecircles:
    circle2.append(plt.Line2D([], [], color='None', marker='o', markersize=size*sizescale,  markeredgecolor='white',mew=np.minimum(maxcirclewidth,size*circlewidth)))

legend1=ax.legend(circles, class_dict.keys(), numpoints=1, bbox_to_anchor=colorlegendloc, loc='upper left', labelcolor=class_dict.values(), title=grouptitle, frameon=False)
legend2=ax.legend(circle2, sizecircles, numpoints=1, bbox_to_anchor=sizelegendloc, loc='lower right', title=sizetitle,labelspacing=3,reverse=True, frameon=False, handletextpad=2)
ax.add_artist(legend1)
ax.add_artist(legend2)

#Save the final figures we have created, as both large png and svg (latter goes to Commons)
fig.savefig(fileout+'.svg', transparent=False)
fig.savefig(fileout+'.png',dpi=750)