import os 
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import healpy as hp
from rubin_scheduler.scheduler.utils import get_current_footprint

nside = 64

# get_current_footprint returns the 'target map' weightings over the sky per filter,
# and the 'labels' indicating the primary survey purpose in that area
# note that DDF and other microsurveys get added in the FBS, so are not included in these maps

maps, labels = get_current_footprint(nside=nside)

# what labels do we have?
label_names = [l for l in np.unique(labels) if l != '']
label_names

# let's show them on the sky

label_dict = {}
for i, name in enumerate(label_names):
    label_dict[name] = i + 1
label_dict

ft = np.zeros(len(maps['r'])) 
for name in label_names:
    match = np.where(labels == name)
    ft[match] += label_dict[name]
ft = np.where(ft == 0, hp.UNSEEN, ft)

hp.mollview(ft)

# Area on-sky in each of these regions?
pixarea = hp.nside2pixarea(nside, degrees=True)
# fudge factor to cover X amount of area 
# overlap between pointings in the tesselation is 10%,
# but edge effects of the footprint increases the needed
# number of visits (to fully cover the area, as some visits 'hang over' the edge)
# NOTE THIS IS APPROXIMATE and will vary with footprint details
pixarea_to_pointings = (1 / 9.6) * 1.3

skarea = {}
skapprox_pointings = {}
for name in label_names:
    npix = len(np.where(labels == name)[0])
    skarea[name] = npix * pixarea
    skapprox_pointings[name] = np.ceil(skarea[name] * pixarea_to_pointings)
    # this would be 1.1 for a 10% overlap between pointings, but 
# allfootprint
npix = len(np.where(labels != '')[0])
skarea['All'] = npix * pixarea
skapprox_pointings['All'] = np.ceil(skarea['All'] * pixarea_to_pointings)
    
print("Area per SKY_AREA section")
pd.DataFrame([skarea, skapprox_pointings], index=['Area (sq deg)', 'Approx NPointing']).round(1).T

# how many visits in the current baseline fall into each of these regions?
from rubin_sim.data import get_baseline
import sqlite3
from rubin_sim.maf import WFDlabelStacker
import rubin_sim.maf as maf

opsdb = get_baseline()
print(opsdb)
conn = sqlite3.connect(opsdb)
opsim = os.path.split(opsdb)[-1].replace('.db.', '')
print(opsim, ":", opsdb)
visits = pd.read_sql('select observationStartMJD, filter, fieldRA, fieldDec, rotSkyPos, visitExposureTime, target_name, scheduler_note from observations', conn)

visits.head()

twilight_neo = visits.query('scheduler_note.str.contains("twilight_near_sun")')
ddf = visits.query('scheduler_note.str.contains("DD") and not scheduler_note.str.contains("RGES")')
rges = visits.query('scheduler_note.str.contains("RGES")')
triplets = visits.query('scheduler_note == "long"')
nvis = {}
nvis['all visits'] = len(visits)
nvis['TwilightNEO'] = len(twilight_neo)
nvis['DD'] =  len(ddf)
nvis['Triplet (third)'] = len(triplets)
nvis['RGES'] = len(rges)
nvis

area = {}
approx_pointings = {}
med_nvis = {}
mean_nvis = {}

name = 'all visits'
print(f"Plotting for {name}")
vis = visits.to_records()
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'DD'
print(f"Plotting for {name}")
vis = ddf.to_records()
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'RGES'
print(f"Plotting for {name}")
vis = rges.to_records()
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'TwilightNEO'
print(f"Plotting for {name}")
vis = twilight_neo.to_records()
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'Triplet (third)'
print(f"Plotting for {name}")
vis = triplets.to_records()
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

other_visits = visits.query('~scheduler_note.str.contains("DD") and ~scheduler_note.str.contains("twilight_near_sun")')
ovisitsnp = other_visits.to_records()

# We need to run this stacker separately for different labels,
# as visits can count for more than one area
# Already discounting microsurveys, although triplet 'long' should remain in consideration for WFD
name = 'NES'
footpt = np.where(labels == 'nes', 1, 0)
stacker = WFDlabelStacker(footprint=footpt, area_id_name=name)
ovisitsnp = stacker.run(sim_data=ovisitsnp)
nvis[name] = len(np.where(ovisitsnp['area_id'] == name)[0])

print(f"Plotting for {name}")
vis = ovisitsnp[np.where(ovisitsnp['area_id'] == name)]
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'SCP'
footpt = np.where(labels == 'scp', 1, 0)
stacker = WFDlabelStacker(footprint=footpt, area_id_name=name)
ovisitsnp = stacker.run(sim_data=ovisitsnp)
nvis[name] = len(np.where(ovisitsnp['area_id'] == name)[0])

print(f"Plotting for {name}")
vis = ovisitsnp[np.where(ovisitsnp['area_id'] == name)]
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'MW'
footpt = np.where((labels == 'bulgy') | (labels == 'dusty_plane') | (labels == 'LMC_SMC'), 1, 0)
stacker = WFDlabelStacker(footprint=footpt, area_id_name=name)
ovisitsnp = stacker.run(sim_data=ovisitsnp)
nvis[name] = len(np.where(ovisitsnp['area_id'] == name)[0])

print(f"Plotting for {name}")
vis = ovisitsnp[np.where(ovisitsnp['area_id'] == name)]
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'ExGal'
footpt = np.where((labels == 'lowdust') | (labels == 'euclid_overlap'), 1, 0)
stacker = WFDlabelStacker(footprint=footpt, area_id_name=name)
ovisitsnp = stacker.run(sim_data=ovisitsnp)
nvis[name] = len(np.where(ovisitsnp['area_id'] == name)[0])

print(f"Plotting for {name}")
vis = ovisitsnp[np.where(ovisitsnp['area_id'] == name)]
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0])  * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

name = 'WFD'
footpt = np.where((labels == 'lowdust') | (labels == 'euclid_overlap') | (labels == "bulgy") | (labels == "virgo") | (labels=="LMC_SMC"), 1, 0)
stacker = WFDlabelStacker(footprint=footpt, area_id_name=name)
ovisitsnp = stacker.run(sim_data=ovisitsnp)
nvis[name] = len(np.where(ovisitsnp['area_id'] == name)[0])

print(f"Plotting for {name}")
vis = ovisitsnp[np.where(ovisitsnp['area_id'] == name)]
metric = maf.CountMetric(col='observationStartMJD')
slicer = maf.HealpixSlicer(nside=nside)
summaries = [maf.MeanMetric(), maf.MedianMetric(), maf.RmsMetric()]
b = maf.MetricBundle(metric, slicer, constraint = '', 
                     summary_metrics=summaries, run_name=opsim, info_label=name)
g = maf.MetricBundleGroup({0: b}, None)
g.run_current(sim_data=vis, constraint='')

area[name] = len(np.where(b.metric_values > 0)[0]) * pixarea
approx_pointings[name] = np.ceil(area[name] * pixarea_to_pointings)
med_nvis[name] = b.summary_values['Median']
mean_nvis[name] = b.summary_values['Mean']

print(b.summary_values)
figs = b.plot()

fraction = dict(zip(list(nvis.keys()), np.array(list(nvis.values()))/nvis['all visits'] * 100))
fraction['TwilightNEO'] /= 2  # visits are half as long
pd.DataFrame([nvis, fraction, area, approx_pointings, med_nvis, mean_nvis], 
             index=['Nvisits', 'Percent', 'Area', 'Approx NPointings', 'Median Nvis', 'Mean Nvis']).round(1).T