# -*- coding: utf-8 -*-
"""
Created on Thu May 24 16:52:15 2018

@author: rstreet
"""

import numpy as np
from astropy.table import Table

# Tables of synthesized u'g'r'i'z' colors for the Gunn-Stryker Stellar Atlas,
# from Lens, D.D et al. (1998), Apj, 119, 121, Table 2.
# "Photometric Separation of Stellar Properties Using SDSS Filters"

def load_lenz_table2():
    """Function to return the full table 2 from 
    Lens, D.D et al. (1998), Apj, 119, 121, except those entries with
    no given spectral type."""
    
    star_colours = np.array([ [  -0.4807 ,  -0.5639 ,  -0.3989 ,  -0.2855   ],
                            [  -0.4369 ,  -0.5392 ,  -0.4216 ,  -0.3269   ],
                            [  -0.4019 ,  -0.5466 ,  -0.4072 ,  -0.3431   ],
                            [ -0.2565  , -0.4751  , -0.3620  , -0.2623    ],
                            [ -0.2052  , -0.5083  , -0.3858  , -0.2985    ],
                            [ -0.1145  , -0.4791  , -0.3933  , -0.2963    ],
                            [ 0.0447   , -0.4195  , -0.3255  , -0.2223    ],
                            [ 0.0714   , -0.4341  , -0.3423  , -0.2447    ],
                            [ 0.0887   , -0.4461  , -0.3498  , -0.2515    ],
                            [ 0.2729   , -0.3745  , -0.3096  , -0.2088    ],
                            [ 0.4969   , -0.3304  , -0.2853  , -0.1840   ],
                            [ 0.5026   , -0.3289  , -0.2643  , -0.1613   ],
                            [ 0.5617   , -0.3068  , -0.2600  , -0.1280   ],
                            [ 0.9790   , -0.2533  , -0.2225  , -0.1205   ],
                            [ 0.8832   , -0.2590  , -0.2161  , -0.1021   ],
                            [ 1.0182   , -0.2463  , -0.2304  , -0.1583   ],
                            [ 0.9620   , -0.2488  , -0.2151  , -0.1243   ],
                            [ 1.0155   , -0.2413  , -0.2225  , -0.1270   ],
                            [ 1.0510   , -0.2024  , -0.1974  , -0.1136   ],
                            [ 1.0547   , -0.1742  , -0.1844  , -0.0739   ],
                            [ 1.0741   , -0.1945  , -0.1848  , -0.1115    ],
                            [ 1.1581   , -0.1584  , -0.1904  , -0.1318    ],
                            [ 1.2430   , -0.1163  , -0.1286  , -0.0540    ],
                            [ 1.1306   , -0.1032  , -0.1390  , -0.0953    ],
                            [ 1.1502   , 0.0013   , -0.0701  , -0.0590    ],
                            [ 1.1785   , 0.1160   , -0.0198  , -0.0618    ],
                            [ 1.1026   , 0.2064   , 0.0404   , -0.0488    ],
                            [ 1.1193   , 0.3048   , 0.0735   , 0.0231     ],
                            [ 0.9831   , 0.3406   , 0.1116   , 0.0298     ],
                            [ 1.0683   , 0.3282   , 0.1230   , 0.0166     ],
                            [ 1.2358   , 0.3467   , 0.0966   , -0.0261    ],
                            [ 1.1709   , 0.3540   , 0.0944   , -0.0061    ],
                            [ 1.2734   , 0.3865   , 0.1285   , 0.0231     ],
                            [ 1.2919   , 0.4361   , 0.1381   , 0.0214     ],
                            [ 1.5195   , 0.4926   , 0.1711   , 0.0265     ],
                            [ 1.5028   , 0.4912   , 0.1658   , 0.0510     ],
                            [ 1.6281   , 0.6197   , 0.2160   , 0.0759     ],
                            [ 2.0313   , 0.6999   , 0.2245   , 0.1008     ],
                            [ 2.2264   , 0.8254   , 0.2901   , 0.1451     ],
                            [ 2.4454   , 0.9289   , 0.3221   , 0.1348     ],
                            [ 2.6485   , 1.1869   , 0.5150   , 0.2592     ],
                            [ 2.7246   , 1.2219   , 0.5650   , 0.2531     ],
                            [ 2.7529   , 1.2146   , 0.5816   , 0.2864     ],
                            [ 2.7564   , 1.2529   , 0.6159   , 0.2540     ],
                            [ 2.7902   , 1.4056   , 0.9429   , 0.4277     ],
                            [ 2.7670   , 1.3563   , 0.9599   , 0.4434     ],
                            [ 2.8687   , 1.3709   , 1.2489   , 0.5682     ],
                            [ 3.1937   , 1.5201   , 1.5139   , 0.6697     ],
                            [ 3.1577   , 1.5271   , 1.7723   , 0.8039     ],
                            [ 2.9529   , 1.5728   , 2.2560   , 1.0712     ],
                            [ -0.1459  , -0.5016  , -0.3974  , -0.2952    ],
                            [ -0.1307  , -0.4757  , -0.3640  , -0.2731    ],
                            [ 0.2144   , -0.4008  , -0.3076  , -0.2245    ],
                            [ 0.1761   , -0.4177  , -0.3396  , -0.2365    ],
                            [ 0.3502   , -0.3619  , -0.3079  ,  -0.1603   ],
                            [ 0.4977   , -0.3511  , -0.2708  ,  -0.1957   ],
                            [ 1.0347   , -0.2961  , -0.2450  ,  -0.1252   ],
                            [  0.8936  ,  -0.2989 ,  -0.2556 ,  -0.1600   ],
                            [  0.8609  ,  -0.2653 ,  -0.2690 ,  -0.1160   ],
                            [  1.1148  ,  -0.1660 ,  -0.1719 ,  -0.1077   ],
                            [ 1.1348  ,  -0.1279 ,  -0.1449 ,  -0.0822   ],
                            [ 1.1474  ,  -0.0796 ,  -0.1440 ,  -0.0797   ],
                            [ 1.2581  ,  -0.0135 ,  -0.1217 ,  -0.0705   ],
                            [ 1.2022  ,  0.0061  ,  -0.0631 ,  -0.0569   ],
                            [ 1.1289  ,  0.0881  ,  -0.0753 ,  -0.0548   ],
                            [ 1.2355  ,  0.0960  ,  -0.0488 ,  -0.0300   ],
                            [ 1.2694  ,  0.0756  ,  -0.0432 ,  -0.0425   ],
                            [ 1.1052  ,  0.1341  ,  0.0011  ,  -0.0187   ],
                            [ 1.1260  ,  0.2172  ,  0.0163  ,  -0.0248   ],
                            [ 1.1007  ,  0.2289  ,  0.0204  ,  -0.0127   ],
                            [ 1.2200  ,  0.3839  ,  0.0734  ,  -0.0039   ],
                            [ 1.4744  ,  0.4287  ,  0.1354  ,  0.0401    ],
                            [ 1.4779  ,  0.4165  ,  0.1351  ,  0.0462    ],
                            [ 1.4871  ,  0.4271  ,  0.1081  ,  0.0155    ],
                            [ 1.6009  ,  0.4546  ,  0.1095  ,  0.0176    ],
                            [ 1.6414  ,  0.4989  ,  0.1533  ,  0.0393    ],
                            [ 1.6267  ,  0.5282  ,  0.1433  ,  0.0001    ],
                            [ 1.7526  ,  0.6150  ,  0.1859  ,  0.0735    ],
                            [ 1.6891  ,  0.6785  ,  0.2440  ,  0.1455    ],
                            [ 1.8016  ,  0.5780  ,  0.2349  ,  0.1113    ],
                            [ 1.8309  ,  0.6400  ,  0.2365  ,  0.1399    ],
                            [ 1.9539  ,  0.6396  ,  0.2201  ,  0.1274    ],
                            [ 2.1162  ,  0.6701  ,  0.2299  ,  0.1202    ],
                            [ 2.0322  ,  0.6821  ,  0.2604  ,  0.1180    ],
                            [ 2.1908  ,  0.6934  ,  0.2115  ,  0.1852    ],
                            [ 2.1373  ,  0.7436  ,  0.2656  ,  0.1572    ],
                            [ 2.1523  ,  0.6971  ,  0.2516  ,  0.1080    ],
                            [ 2.2028  ,  0.7236  ,  0.3084  ,  0.1465    ],
                            [ 2.3106  ,  0.7425  ,  0.2537  ,  0.0915    ],
                            [ 2.1322  ,  0.8062  ,  0.3064  ,  0.1651    ],
                            [ 2.1672  ,  0.7995  ,  0.2808  ,  0.1523    ],
                            [ 2.3185  ,  0.7760  ,  0.2779  ,  0.1558    ],
                            [ 2.4179  ,  0.7835  ,  0.2734  ,  0.1226    ],
                            [ 2.5586  ,  0.7847  ,  0.2193  ,  0.1510    ],
                            [ 2.4713  ,  0.7687  ,  0.2917  ,  0.1561    ],
                            [ 2.4261  ,  0.7988  ,  0.2966  ,  0.1747    ],
                            [ 2.5670  ,  0.8043  ,  0.2606  ,  0.1681    ],
                            [ 2.4812  ,  0.7942  ,  0.2733  ,  0.1435    ],
                            [ 2.5223  ,  0.8352  ,  0.3035  ,  0.1528    ],
                            [ 2.5817  ,  0.8421  ,  0.2994  ,  0.1769    ],
                            [ 2.6338  ,  0.8084  ,  0.2040  ,  0.1904    ],
                            [ 2.7931  ,  0.8609  ,  0.2528  ,  0.2141    ],
                            [ 2.7177  ,  0.8876  ,  0.2815  ,  0.1631    ],
                            [ 2.7224  ,  0.8986  ,  0.3519  ,  0.1214    ],
                            [ 2.7784  ,  0.9671  ,  0.3012  ,  0.1828    ],
                            [ 3.0990  ,  0.9164  ,  0.2583  ,  0.2162    ],
                            [ 2.7153  ,  0.8428  ,  0.3036  ,  0.1735    ],
                            [ 2.9109  ,  0.9317  ,  0.3039  ,  0.1915    ],
                            [ 2.9833  ,  0.9190  ,  0.3478  ,  0.2205    ],
                            [ 3.0511  ,  1.0383  ,  0.4008  ,  0.2023    ],
                            [ 3.1050  ,  1.0559  ,  0.4158  ,  0.2104    ],
                            [ 2.9778  ,  1.0276  ,  0.3966  ,  0.2387    ],
                            [ 2.9427  ,  1.0103  ,  0.4325  ,  0.2494    ],
                            [ 3.4217  ,  1.1062  ,  0.3697  ,  0.3125    ],
                            [ 3.3603  ,  1.1000  ,  0.5092  ,  0.3341    ],
                            [ 3.3588  ,  1.1731  ,  0.5445  ,  0.3211    ],
                            [ 3.4138  ,  1.1913  ,  0.5982  ,  0.3455    ],
                            [ 3.5737  ,  1.2502  ,  0.6240  ,  0.3713    ],
                            [ 3.5501  ,  1.2207  ,  0.6964  ,  0.3556    ],
                            [ 3.6958  ,  1.3473  ,  0.9798  ,  0.4808    ],
                            [ 3.5277  ,  1.2846  ,  1.0005  ,  0.4866    ],
                            [ 3.9552  ,  1.4460  ,  0.8862  ,  0.6270    ],
                            [ 3.8238  ,  1.4933  ,  1.2612  ,  0.6047    ],
                            [ 3.9797  ,  1.6081  ,  1.2657  ,  0.6798    ],
                            [ 3.6524  ,  1.4707  ,  1.4525  ,  0.6833    ],
                            [ 3.6774  ,  1.4914  ,  1.5627  ,  0.7506    ],
                            [ 3.4818  ,  1.4567  ,  1.6882  ,  0.7932    ],
                            [ 3.2648  ,  1.4693  ,  1.8615  ,  0.9387    ],
                            [ 3.1680  ,  1.5534  ,  2.0137  ,  0.9746    ],
                            [ 2.0864  ,  1.5538  ,  2.6493  ,  1.2675    ],
                            [ 2.7613  ,  1.8513  ,  2.3176  ,  1.3506    ],
                            [ 2.4980  ,  2.1526  ,  2.6300  ,  1.4449    ],
                            [ 1.3762  ,  2.1467  ,  2.8639  ,  1.7287    ],
                            [ -0.3239 ,  -0.4611 ,  -0.4118 , -0.2051    ],
                            [ 0.1527  ,  -0.4138 ,  -0.3565 , -0.2217   ],
                            [ 0.2536  ,  -0.1537 ,  -0.1980 , -0.1098   ] ])
    
    star_labels = np.array([ 'O5   ',
                     'O8F  ',
                     'O6   ',
                     'B1V  ',
                     'B1V  ',
                     'B2V  ',
                     'B3V  ',
                     'B3V  ',
                     'B4V  ',
                     'A1V  ',
                     'B6V  ',
                     'B7V  ',
                     'B9V  ', 
                     'A0V  ',
                     'B9V  ',
                     'A2V  ', 
                     'A1V  ',
                     'A2V  ',
                     'A3V  ',
                     'A0V  ',
                     'B9V  ',
                     'A7V  ',
                     'A2V  ',
                     'A5V  ',
                     'A5V  ',
                     'F4V  ',
                     'F6V  ',
                     'F8V  ',
                     'F9V  ',
                     'F8V  ',
                     'F8V  ',
                     'F8V  ',
                     'G2V  ',
                     'K0V  ',
                     'G8V  ',
                     'G5V  ',
                     'G8V  ',
                     'K3V  ',
                     'K4V  ',
                     'K8V  ',
                     'M0V  ',
                     'K7V  ',
                     'K7V  ',
                     'M0V  ',
                     'M0V  ',
                     'M2V  ',
                     'M4V  ',
                     'M6V  ',
                     'M8V  ',
                     'M5V  ',
                     'B1IV ',
                     'B2III',
                     'B3IV ',
                     'B3III',
                     'B7IV ',
                     'B7III',
                     'A3III',
                     'A0IV ',
                     'B9IV ',
                     'A3IV ',
                     'A4IV ',
                     'A5III' ,
                     'A5IV ',
                     'F0IV ',
                     'F0IV ',
                      'A9IV ',
                     'F0IV ',
                     'F2IV ',
                     'F7IV ',
                     'F5IV ',
                     'G2IV ',
                     'G8IV ',
                     'G4IV ',
                     'G6IV ',
                     'G6IV ',
                     'G3IV ',
                     'G8IV ',
                     'G5IV ',
                     'G5IV ',
                     'G5IV ',
                     'G7IV ',
                     'G5IV ',
                     'G8III',
                     'G5IV ',
                     'K0III',
                     'G8IV ',
                     'K0III',
                     'K0III',
                     'K0III',
                     'G8IV ',
                     'K0IV ',
                     'K0III',
                     'K0III',
                     'K3III',
                     'K0IV ',
                     'K0III',
                     'K2III',
                     'K2III',
                     'K1III',
                     'K1III',
                     'K1III',
                     'K3III',
                     'K3III',
                     'K2IIIP',
                     'K3III',
                     'K3III',
                     'K2III',
                     'K2III',
                     'K2IIIP',
                     'K1IV ',
                     'K3III',
                     'K3III',
                     'K3III',
                     'K4III',
                     'K4III',
                     'K5III',
                     'M0III',
                     'M0III',
                     'K5III',
                     'M0III',
                     'M2III',
                     'M1III',
                     'M3III',
                     'M2III',
                     'M2III',
                      'M2III',
                     'M8III',
                     'M5III',
                     'M5III',
                     'M6III',
                     'M8   ',
                     'M7III',
                     'M6III',
                     'B0IB ',
                     'B5IB ',
                     'B8IA '])
    
    return star_colours, star_labels
    
def get_dwarf_star_colours():
    """Synthesized u'g'r'i'z' colors for the dwarf (type V) stars in the
    Gunn-Stryker Stellar Atlas, taken from 
    Lens, D.D et al. (1998), Apj, 119, 121, Table 2."""
    
    (star_colours, star_labels) = load_lenz_table2()
    
    idx = []
    labels = []
    
    for i,t in enumerate(star_labels):
        
        if 'V' in t:
            
            idx.append(i)
            labels.append(t)
            
    dwarf_star_colours = star_colours[idx]
    
    return dwarf_star_colours, labels

def get_giant_star_colours():
    """Synthesized u'g'r'i'z' colors for the giant (type III) stars in the
    Gunn-Stryker Stellar Atlas, taken from 
    Lens, D.D et al. (1998), Apj, 119, 121, Table 2."""
    
    (star_colours, star_labels) = load_lenz_table2()
    
    idx = []
    labels = []
    
    for i,t in enumerate(star_labels):
        
        if 'III' in t:
            
            idx.append(i)
            labels.append(t)
            
    giant_star_colours = star_colours[idx]
    
    return giant_star_colours, labels
    

if __name__ == '__main__':
    
    (dwarf_star_colours, labels) = get_dwarf_star_colours()
    
    print(dwarf_star_colours.shape)
    print(len(labels))
    
    (giant_star_colours, labels) = get_giant_star_colours()
    
    print(giant_star_colours.shape)
    print(len(labels))
    