# hack to get the path right
import sys
sys.path.append('..')

from ztf_sim.SkyBrightness import SkyBrightness
from ztf_sim.magnitudes import limiting_mag
import astropy.units as u
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline

sns.set_style('ticks')
sns.set_context('talk')

# load the training data
df = pd.read_csv('../data/ptf-iptf_diq.csv.gz')

df.head()

"""
Explanation: plot_sky_brightness_model
Visualizing the machine-learned (PTF) sky brightness model
End of explanation
"""


for col in df.columns:
    print('{}: {}-{}'.format(col, df[col].min(), df[col].max()))

for filterkey in [1,2,4]:
    wbright = (df.filterkey==filterkey) & (df.moonillf > 0.95) & (df.moonalt > 70)
    print(filterkey, np.sum(wbright))
    plt.hist(df.loc[wbright,'limiting_mag'],normed=True,label=f'{filterkey}')
plt.legend()

for filterkey in [1,2,4]:
    w = (df.filterkey==filterkey)
    print(filterkey, np.percentile(df.loc[w,'limiting_mag'],[5,95]))

"""
Explanation: Let's look at the range of the data:
End of explanation
"""


plt.hist(np.abs(df.moonillf))
plt.xlabel('Moon Illumination Fraction')
plt.ylabel('Number of Exposures')
sns.despine()

plt.hist(df.moonalt)
plt.xlabel('Moon Altitude (degrees)')
plt.ylabel('Number of Exposures')
sns.despine()

plt.hist(df.moon_dist)
plt.xlabel('Moon Distance (degrees)')
plt.ylabel('Number of Exposures')
sns.despine()

plt.hexbin(df['moonalt'],df['moon_dist'])
plt.xlabel('Moon Altitude (deg)')
plt.ylabel('Moon Distance (deg)')

plt.hist(df.azimuth)
plt.xlabel('Telescope Azimuth (degrees)')
plt.ylabel('Number of Exposures')
plt.xlim(0,360)
sns.despine()

plt.hist(df.altitude)
plt.xlabel('Telescope Altitude (degrees)')
plt.ylabel('Number of Exposures')
plt.xlim(0,90)
sns.despine()

plt.hist(df.sunalt)
plt.xlabel('Sun Altitude (degrees)')
plt.ylabel('Number of Exposures')
plt.xlim(-90,-10)
sns.despine()

w = (df['filterkey'] == 2)
wdark = (df['filterkey'] == 2) & (np.abs(df['moonillf']) < .1) & (df['moonalt'] < -10.)
wbright = (df['filterkey'] == 2) & (np.abs(df['moonillf']) > .9) & (df['moonalt'] > 70.)

print(df[w]['limiting_mag'].median())
normed=True
plt.hist(df[w]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)
plt.hist(df[wdark]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)
plt.hist(df[wbright]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)

plt.xlabel('Limiting Magnitude (R-band mag arcsec$^{-2}$)')
plt.ylabel('Number of Exposures')
sns.despine()

for filterkey, filtername in {1:'g',2:'R',4:'i'}.items():
    plt.figure()
    w = (df['filterkey'] == filterkey)
    wdark = (df['filterkey'] == filterkey) & (np.abs(df['moonillf']) < .1) & (df['moonalt'] < -10.)
    wbright = (df['filterkey'] == filterkey) & (np.abs(df['moonillf']) > .9) & (df['moonalt'] > 70.)

    normed=True
    plt.hist(df[w]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)
    plt.hist(df[wdark]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)
    plt.hist(df[wbright]['limiting_mag'],bins=np.linspace(18,23,100),normed=normed,alpha=0.5)
    print(filterkey,df.loc[wdark,'limiting_mag'].median(),df.loc[wbright,'limiting_mag'].median())

    plt.xlabel(f'Limiting Magnitude ({filtername}-band mag arcsec$^{-2}$)')
    plt.ylabel('Number of Exposures')
    sns.despine()

w = (df['filterkey'] == 1)
wdark = (df['filterkey'] == 1) & (np.abs(df['moonillf']) < .1) & (df['moonalt'] < -10.)
print(df[w]['limiting_mag'].median())
plt.hist(df[w]['limiting_mag'],bins=np.linspace(18,23,100))
plt.hist(df[wdark]['limiting_mag'],bins=np.linspace(18,23,100))

plt.xlabel('Limiting Magnitude (g-band mag arcsec$^{-2}$)')
plt.ylabel('Number of Exposures')
sns.despine()

w = (df['filterkey'] == 4)
wdark = (df['filterkey'] == 4) & (np.abs(df['moonillf']) < .1) & (df['moonalt'] < -10.)
print(df[w]['limiting_mag'].median())
plt.hist(df[w]['limiting_mag'],bins=np.linspace(18,23,100))
plt.hist(df[wdark]['limiting_mag'],bins=np.linspace(18,23,100))
plt.xlabel('Limiting Magnitude (i-band mag arcsec$^{-2}$)')
plt.ylabel('Number of Exposures')
sns.despine()

"""
Explanation: So there are clearly outliers.  Let's look at histograms to find good fiducial values:
End of explanation
"""


w = (df['filterkey'] == 2)
w &= (df['altitude'] > 20) & (df['fwhmsex'] < 5)
plt.hexbin(df[w]['altitude'],df[w]['fwhmsex'])
plt.xlabel('Altitude (deg)')
plt.ylabel('FWHM (arcsec)')
alts = np.linspace(20,90,100)
airmasses = 1./np.cos(np.radians(90.-alts))
plt.plot(alts,2*airmasses**(3./5))

"""
Explanation: Let's look at the dependence of FWHM on altitude:
End of explanation
"""


w = (df['filterkey'] == 1)
w &= (df['altitude'] > 20) & (df['fwhmsex'] < 5)
plt.hexbin(df[w]['altitude'],df[w]['fwhmsex'])
plt.xlabel('Altitude (deg)')
plt.ylabel('FWHM (arcsec)')
plt.plot(alts,2.4*airmasses**(3./5))

"""
Explanation: g-band
End of explanation
"""


w = (df['filterkey'] == 4)
w &= (df['altitude'] > 20) & (df['fwhmsex'] < 5)
plt.hexbin(df[w]['altitude'],df[w]['fwhmsex'])
plt.xlabel('Altitude (deg)')
plt.ylabel('FWHM (arcsec)')
plt.plot(alts,2.0*airmasses**(3./5))

# load our model
Sky = SkyBrightness()

def test_conditions(moonillf=0,moonalt=40,moon_dist=60.,azimuth=120,altitude=60.,sunalt=-35,filter_id=2):
    """define convenience function for building inputs to SkyBrightness.
    
    Inputs should be scalars expect for one and only one vector.
    
    Note that manual inputs can yield unphysical moonalt/moon_dist pairs, but the real simulator will 
    calculate these correctly."""
    
    moonillf = np.atleast_1d(moonillf)
    moonalt = np.atleast_1d(moonalt)
    moon_dist = np.atleast_1d(moon_dist)
    azimuth = np.atleast_1d(azimuth)
    altitude = np.atleast_1d(altitude)
    sunalt = np.atleast_1d(sunalt)
    filter_id = np.atleast_1d(filter_id)
    
    maxlen = np.max([len(moonillf), len(moonalt), len(moon_dist), len(azimuth), len(altitude), len(sunalt),
                     len(filter_id)])
    
    def blow_up_array(arr,maxlen=maxlen):
        if (len(arr) == maxlen):
            return arr
        elif (len(arr) == 1):
            return np.ones(maxlen)*arr[0]
        else:
            raise ValueError
    
    moonillf = blow_up_array(moonillf)
    moonalt = blow_up_array(moonalt)
    moon_dist = blow_up_array(moon_dist)
    azimuth = blow_up_array(azimuth)
    altitude = blow_up_array(altitude)
    sunalt = blow_up_array(sunalt)
    filter_id = blow_up_array(filter_id)
    
    pars = {'moonillf':moonillf, 'moonalt':moonalt, 'moon_dist':moon_dist, 'azimuth':azimuth, 'altitude':altitude, 
           'sunalt':sunalt,'filter_id':filter_id}
    
    return pd.DataFrame(pars,index=np.arange(maxlen))

def select_data(df, moonillf=0,moonalt=40,moon_dist=60.,azimuth=120,altitude=60,sunalt=-35,filter_id=2):
    """return a boolean selecting data around fiducial inputs
    
    set parameter to None to skip comparison"""
    
    # start with boolean True array
    w = df['azimuth'] == df['azimuth']
    
    def anding(df,par,value,frac=0.2):
        if value is None:
            return True
        else:
            lima = value*(1-frac)
            limb = value*(1+frac)
            llim = np.atleast_1d(np.where(lima < limb, lima, limb))[0]
            ulim = np.atleast_1d(np.where(lima > limb, lima, limb))[0]
            # figure out which is smaller
            
            return (df[par] >= llim) & (df[par] <= ulim)
    
    w &= anding(df, 'moonillf', moonillf)
    w &= anding(df, 'moonalt', moonalt)
    w &= anding(df, 'moon_dist', moon_dist)
    w &= anding(df, 'azimuth', azimuth)
    w &= anding(df, 'altitude', altitude)
    w &= anding(df, 'sunalt', sunalt)
    
    # dataframe uses filterkey
    w &= (df['filterkey'] == filter_id)
    
    return w

# show that it works!
test_conditions(altitude=np.linspace(0,90,5))

"""
Explanation: i-band
End of explanation
"""


from ztf_sim.magnitudes import limiting_mag

fids = [1,2,3]
dark_condition = test_conditions(moonillf=0,moonalt=0,altitude=90,sunalt=-80,filter_id=fids)
sky_brightness = Sky.predict(dark_condition)
for fid, sky in zip(fids,sky_brightness.values):
    print(fid, sky, limiting_mag(30, 2.0, sky, filter_id = fid, altitude=90, SNR=5.))

bright_condition = test_conditions(moonillf=1.,moonalt=70.,moon_dist=20.,altitude=90,sunalt=-80,filter_id=[1,2,3])
sky_brightness = Sky.predict(bright_condition)
for fid, sky in zip(fids,sky_brightness.values):
    print(fid, sky, limiting_mag(30, 2.0, sky, filter_id = fid, altitude=90, SNR=5.))

"""
Explanation: Calculate some limiting mag values in bright and dark time at zenith:
End of explanation
"""


moonillfs = np.linspace(0,1,20)
dft = test_conditions(moonillf=moonillfs)
skies = Sky.predict(dft)
plt.plot(moonillfs,skies)
plt.gca().invert_yaxis()
plt.xlabel('Moon Illimination Fraction')
plt.ylabel('Sky Brightness (R-band mag arcsec$^{-2}$)')
sns.despine()

w =  select_data(df,moonillf=None)#,moonalt=45, moon_dist=45)
plt.scatter(np.abs(df[w]['moonillf']),df[w]['limiting_mag'],alpha=0.3)
plt.gca().invert_yaxis()
plt.xlabel('Moon Illimination Fraction')
plt.ylabel('Limiting Mag (R-band mag$)')
sns.despine()

moonillfs = np.linspace(0,1,20)
dft = test_conditions(moonillf=moonillfs,filter_id=1)
skies = Sky.predict(dft)
plt.plot(moonillfs,skies)
plt.gca().invert_yaxis()
plt.xlabel('Moon Illimination Fraction')
plt.ylabel('Sky Brightness (g-band mag arcsec$^{-2}$)')
sns.despine()

"""
Explanation: Now let's walk through the various parameters and look at the model behavior.
Vary moon illilumination fraction at 45 degree altitude, R-band:
TOFIX
forgot to calculate limiting mag from the sky brightness through here
End of explanation
"""


moonillfs = np.linspace(0,1,20)
dft = test_conditions(moonillf=moonillfs,filter_id=3)
skies = Sky.predict(dft)
plt.plot(moonillfs,skies)
plt.gca().invert_yaxis()
plt.xlabel('Moon Illimination Fraction')
plt.ylabel('Sky Brightness (i-band mag arcsec$^{-2}$)')
sns.despine()

moon_dists = np.linspace(0,180,30)
dft = test_conditions(moonillf=0.8,moon_dist=moon_dists)
skies = Sky.predict(dft)
plt.plot(moon_dists,skies)
plt.gca().invert_yaxis()
plt.xlabel('Moon Distance (deg)')
plt.ylabel('Sky Brightness (R-band mag arcsec$^{-2}$)')
sns.despine()

azimuths = np.linspace(0,360,30)
dft = test_conditions(azimuth=azimuths)
skies = Sky.predict(dft)
plt.plot(azimuths,skies)
plt.gca().invert_yaxis()
plt.xlabel('Telescope Azimuth (deg)')
plt.ylabel('Sky brightness (R-band mag arcsec$^{-2}$)')
sns.despine()

altitudes = np.linspace(0,90,30)
dft = test_conditions(altitude=altitudes)
skies = Sky.predict(dft)
plt.plot(altitudes,skies)
plt.gca().invert_yaxis()
plt.xlabel('Telescope Altitude (deg)')
plt.ylabel('Sky Brightness (R-band mag arcsec$^{-2}$)')
sns.despine()

sunalts = np.linspace(-10,-80,30)
dft = test_conditions(sunalt=sunalts)
skies = Sky.predict(dft)
plt.plot(sunalts,skies)
plt.gca().invert_yaxis()
plt.xlabel('Sun Altitude (deg)')
plt.ylabel('Limiting Magnitude (R-band mag arcsec$^{-2}$)')
sns.despine()

"""
Explanation: Not as smooth as we might expect.  Let's look at the input data:
These data are surprisingly sparse!
End of explanation
"""


filter_id = 2
seeing=2.0
alts = np.linspace(25,90,30)
dft = test_conditions(altitude=alts, azimuth=0, filter_id=filter_id)
skies = Sky.predict(dft)
limmag = limiting_mag(30*u.second,  seeing, skies, filter_id*np.ones(len(alts)), alts )
#plt.plot(alts,skies)
plt.plot(alts, limmag)
plt.gca().invert_yaxis()
plt.xlabel('Altitude (deg)')
plt.ylabel('Limiting Magnitude (mag)')
sns.despine()

from ztf_sim.utils import altitude_to_airmass

altitude_to_airmass(25)

def slot_metric(limiting_mag):
    metric = 10.**(0.6 * (limiting_mag - 21))
    # lock out -99 limiting mags even more aggressively
    return metric.where(limiting_mag > 0, -0.99)

metric = slot_metric(limmag)
plt.plot(alts, metric)
plt.xlabel('Altitude (deg)')
plt.ylabel('Metric value')
sns.despine()

seeing=2.0
alts = np.linspace(25,90,30)
for filter_id in [1,2,3]:
    dft = test_conditions(altitude=alts, azimuth=0, filter_id=filter_id)
    skies = Sky.predict(dft)
    limmag = limiting_mag(30*u.second,  seeing, skies, filter_id*np.ones(len(alts)), alts )
    metric = slot_metric(limmag)
    #airmass = altitude_to_airmass(alts)
    plt.plot(alts, metric/metric.iloc[-1], label=filter_id)
plt.plot(alts,-1e-4*(alts-90)**2.+1)
plt.legend()
plt.xlabel('Altitude (deg)')
plt.ylabel('Metric value')
sns.despine()

seeing=2.0
alts = np.linspace(25,90,30)
for filter_id in [1,2,3]:
    dft = test_conditions(altitude=alts, azimuth=0, filter_id=filter_id)
    skies = Sky.predict(dft)
    limmag = limiting_mag(30*u.second,  seeing, skies, filter_id*np.ones(len(alts)), alts )
    metric = slot_metric(limmag)
    #airmass = altitude_to_airmass(alts)
    plt.plot(alts, metric/(-1e-4*(alts-90)**2.+1), label=filter_id)
plt.legend()
plt.xlabel('Altitude (deg)')
plt.ylabel('Metric value')
sns.despine()

"""
Explanation: low dec fields
End of explanation
"""


