import os
import warnings
import copy
import logging

import numpy as np
import pandas as pd
import sqlite3
import healpy as hp
import matplotlib.pyplot as plt
import colorcet as cc
import skyproj
from IPython.display import display, HTML

from astropy.time import Time, TimeDelta
from astropy.coordinates import SkyCoord
import astropy.units as u

from rubin_scheduler.site_models import Almanac
from rubin_scheduler.scheduler.features import Conditions
from rubin_scheduler.utils import ddf_locations, angular_separation, approx_ra_dec2_alt_az, Site, SURVEY_START_MJD

import rubin_sim.maf as maf

from rubin_nights import connections
import rubin_nights.lfa_data as rn_lfa
import rubin_nights.dayobs_utils as rn_dayobs
import rubin_nights.plot_utils as rn_plots
import rubin_nights.augment_visits as augment_visits
import rubin_nights.rubin_scheduler_addons as rn_sch
import rubin_nights.rubin_sim_addons as rn_sim

import importlib

from lsst_survey_sim import lsst_support, simulate_lsst, plot

band_colors = rn_plots.PlotStyles.band_colors
logging.getLogger('lsst_survey_sim').setLevel(logging.INFO)

#%load_ext memory_profiler

day_obs = 20251020
survey_start = SURVEY_START_MJD
CONFIG_SCRIPT_PATH = "/Users/lynnej/lsst_repos/ts_config_scheduler/Scheduler/feature_scheduler/maintel/fbs_config_lsst_survey.py"
tokenfile = "/Users/lynnej/.lsst/usdf_rsp"
site = "usdf"

sunset, sunrise = rn_dayobs.day_obs_sunset_sunrise(day_obs, sun_alt=-12)
print(sunset.iso, sunrise.iso)

2025-10-20 23:51:29.829 2025-10-21 09:04:37.825

# Retrieve previous visits to start scheduler 

refresh_visits = True

if refresh_visits:
    initial_opsim = simulate_lsst.fetch_previous_visits(day_obs, tokenfile=tokenfile, site=site)
    if initial_opsim is not None:
        initial_opsim.to_hdf("opsim.h5", key='visits')
else:
    initial_opsim = pd.read_hdf("opsim.h5")
    
if initial_opsim is None:
    print("No visits found")
else:
    print(f"Found {len(initial_opsim)} visits")

No visits found

# Configure scheduler and add initial visits

starting_scheduler, initial_opsim, nside = simulate_lsst.setup_scheduler(config_script_path=CONFIG_SCRIPT_PATH, day_obs=day_obs, initial_opsim=initial_opsim, tokenfile=tokenfile, site=site)

/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 56 ELAISS1 sequences, but only 55 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 COSMOS sequences, but only 123 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 COSMOS sequences, but only 122 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 COSMOS sequences, but only 124 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 XMM_LSS sequences, but only 131 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 XMM_LSS sequences, but only 135 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 ELAISS1 sequences, but only 151 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 ELAISS1 sequences, but only 149 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 ECDFS sequences, but only 154 nights. Decreasing sequences.
  warnings.warn(
/Users/lynnej/lsst_repos/ts_fbs_utils/python/lsst/ts/fbs/utils/maintel/lsst_ddf_presched.py:441: UserWarning: Asked for 200 EDFS_a sequences, but only 165 nights. Decreasing sequences.
  warnings.warn(

# Configure band scheduler
band_scheduler = simulate_lsst.setup_band_scheduler()

# Check what bands it expects for day_obs? 
conditions = Conditions(mjd=sunset.mjd)
band_scheduler(conditions)

['g', 'r']

starting_observatory, survey_info = simulate_lsst.setup_observatory(day_obs=day_obs, nside=nside, add_downtime=True, add_clouds=True, seeing=None, real_downtime=False, initial_opsim=initial_opsim)

mask = np.where(survey_info['dayobsmjd'] < rn_dayobs.day_obs_to_time(day_obs).mjd + 365)
print(f"Total nighttime {survey_info["hours_in_night"][mask].sum()}, ")
print(f"total downtime {survey_info["downtime_per_night"][mask].sum()}, ")
print(f"available time {survey_info["hours_in_night"][mask].sum() - survey_info["downtime_per_night"][mask].sum()}")
print(f"Average availability mask - {(survey_info["hours_in_night"][mask].sum() - survey_info["downtime_per_night"][mask].sum()) / survey_info["hours_in_night"][mask].sum()}")
print(f"Average availability over survey {survey_info['system_availability']}")

Total nighttime 3671.403386067599, 
total downtime 799.6621342669241, 
available time 2871.741251800675
Average availability mask - 0.7821916988741916
Average availability over survey 0.9289383954009189

for start, end in survey_info['downtimes']:
    x = np.floor(start - 0.5)
    plt.plot((x, x), (start - x, end - x) , color='k') 
x = np.floor(survey_info['sunsets12'] - 0.5)
y = survey_info['sunsets12'] - x
plt.fill_between(x, 0.8, y)
x = np.floor(survey_info['sunrises12'] - 0.5)
y = survey_info['sunrises12'] - x
#plt.axvline(rn_dayobs.day_obs_to_time(day_obs).mjd, color='r', linestyle=':', linewidth=1.8)
#plt.axvline(np.floor(Time("2025-09-06T12:00:00").mjd - 0.5))
plt.fill_between(x, y, 1.6)
plt.ylim(0.9, 1.5)
plt.xlim(survey_info['survey_start'].mjd, rn_dayobs.day_obs_to_time(day_obs).mjd + 400)
plt.xlabel("MJD", fontsize='large')
_ = plt.ylabel("Fraction of MJD", fontsize='large')
#plt.savefig("onsky_downtime.png", bbox_inches='tight')

check_ddfs = True
if check_ddfs:
    ddfs = ddf_locations(skycoords=True)
    lsst_site = Site('LSST')
    times = np.arange(survey_info['survey_start'].mjd, survey_info['survey_end'].mjd, 1/24)
    times = np.arange(np.floor(Time.now().mjd - 0.5),survey_info['survey_end'].mjd, 0.25/24)
    times = np.arange(sunset.mjd - 0.1, sunrise.mjd + 0.1, 0.05/24)
    sunmoon = survey_info['almanac'].get_sun_moon_positions(times)
    moon_ra = sunmoon['moon_RA']
    moon_dec = sunmoon['moon_dec']
    sun_alt = np.degrees(sunmoon['sun_alt'])
    ddf_moon_dist = {}
    ddf_alt = {}
    for ddf in ddfs:
        ddf_moon_dist[ddf] = angular_separation(ddfs[ddf].ra.deg, ddfs[ddf].dec.deg, np.degrees(moon_ra), np.degrees(moon_dec))
        alt, az = approx_ra_dec2_alt_az(ddfs[ddf].ra.deg, ddfs[ddf].dec.deg, lsst_site.latitude, lsst_site.longitude, times, lmst=None)
        ddf_alt[ddf] = alt
    
    plt.figure()
    for ddf in ddfs:
        mask = np.where((sun_alt <= -12) & (ddf_alt[ddf] > 40))
        plt.plot(Time(times[mask], format='mjd').to_datetime(), ddf_moon_dist[ddf][mask], marker='.', linestyle='', label=ddf)
    plt.legend(loc=(1.01, 0.5))
    plt.axvline(sunset.to_datetime(), color='k', linestyle=':')
    plt.axvline(sunrise.to_datetime(), color='k', linestyle=':')
    plt.axhline(30)
    plt.xticks(rotation=90)
    plt.xlabel("MJD")
    plt.ylabel("Distance to moon (deg)")
    
    plt.figure()
    for ddf in ddfs:
        mask = np.where((sun_alt <= -12) & (ddf_alt[ddf] > 40))
        plt.plot(Time(times[mask], format='mjd').to_datetime(), ddf_alt[ddf][mask], marker='.', linestyle='', label=ddf)
    plt.legend(loc=(1.01, 0.5))
    plt.axvline(sunset.to_datetime(), color='k', linestyle=':')
    plt.axvline(sunrise.to_datetime(), color='k', linestyle=':')
    plt.axhline(30)
    plt.xticks(rotation=90)
    plt.xlabel("MJD")
    plt.ylabel("Altitude (deg)")

reset = True
if reset:
    observatory = copy.deepcopy(starting_observatory)
    scheduler = copy.deepcopy(starting_scheduler)

# Run the simulation for some nights
sim_nights = 30
observations, scheduler, observatory, rewards, obs_rewards, survey_info = simulate_lsst.run_sim(scheduler=scheduler, 
                                                                                                    band_scheduler=band_scheduler,
                                                                                                    observatory=observatory,
                                                                                                    survey_info=survey_info,
                                                                                                    day_obs=day_obs,
                                                                                                    sim_nights=sim_nights, 
                                                                                                    keep_rewards=False)

progress = 0.18%

/Users/lynnej/lsst_repos/rubin_scheduler/rubin_scheduler/skybrightness_pre/sky_model_pre.py:345: UserWarning: Sun high, using bright sky approx
  warnings.warn("Sun high, using bright sky approx")

progress = 0.91%

/Users/lynnej/lsst_repos/rubin_scheduler/rubin_scheduler/utils/healpy_utils.py:342: UserWarning: Can not connect to any more pixels.
  warnings.warn("Can not connect to any more pixels.")

progress = 3.58%

/Users/lynnej/lsst_repos/rubin_scheduler/rubin_scheduler/scheduler/features/features.py:509: UserWarning: Time must flow forwards to track the feature in NObservationsCurrentSeason.Not updating season_map.
  warnings.warn(

progress = 102.54%Skipped 0 observations
Flushed 59 observations from queue for being stale
Completed 8503 observations
ran in 4 min = 0.1 hours

filename = "test.db"
obsdf = simulate_lsst.consolidate_and_save(observatory, observations, initial_opsim=None, filename=filename)

def add_dayobs(x):
    mjdfloor = Time(np.floor(x.observationStartMJD - 0.5) + 0.5, format="mjd", scale="tai")
    return rn_dayobs.day_obs_str_to_int(mjdfloor.isot.split("T")[0])

obsdf["day_obs"] = obsdf.apply(add_dayobs, axis=1)

# Need to start at a specific time inside the night?

# sunset, sunrise = rn_dayobs.day_obs_sunset_sunrise(day_obs, sun_alt=-12)
# start = sunset.mjd - 0.1 /24
# end = sunrise.mjd

# with warnings.catch_warnings():
#     warnings.simplefilter("ignore", RuntimeWarning)
#     vals = sim_runner(
#         observatory,
#         scheduler,
#         band_scheduler=fs,
#         sim_start_mjd=start,
#         sim_duration=end-start, 
#         record_rewards=False,
#         verbose=True,
#     )
# observatory = vals[0]
# scheduler = vals[1]
# observations = vals[2]
# if len(vals) == 5:
#     rewards = vals[3]
#     obs_rewards = vals[4]

if sim_nights == 1:

    targets = pd.DataFrame(observations)
    
    from zoneinfo import ZoneInfo
    tz = ZoneInfo("Chile/Continental")
    tz_utc = ZoneInfo("UTC")
    telescope = "Simonyi"

    site = Site('LSST')
    almanac = Almanac()
    night_events = almanac.get_sunset_info(evening_date=rn_dayobs.day_obs_int_to_str(day_obs), longitude=site.longitude_rad)

    def mjd_to_datetime(mjd, scale='utc', timezone=tz):
        return Time(mjd, format='mjd', scale=scale).utc.to_datetime(timezone=timezone)
        
    eps = 1
    fig, ax = plt.subplots(figsize=(13, 8))
    ax_utc = ax.twiny()
    
    ax.set_title(f"{telescope} DAYOBS {day_obs}", pad=20)
    
    # Shade astronomical events
    ax.fill_between([mjd_to_datetime(night_events['sun_n12_setting']), 
                      mjd_to_datetime(night_events['sun_n18_setting'])],
                     2.5, 0.0, color='lightgray', alpha=0.3)
    ax.fill_between([mjd_to_datetime(night_events['sunset']), 
                     mjd_to_datetime(night_events['sun_n12_setting'])], 
                      2.5, 0.0, color='gray', alpha=0.3)
    ax.fill_between([mjd_to_datetime(night_events['sun_n18_rising']), 
                     mjd_to_datetime(night_events['sun_n12_rising'])],
                     2.5, 0.0, color='lightgray', alpha=0.3)
    ax.fill_between([mjd_to_datetime(night_events['sun_n12_rising']), 
                     mjd_to_datetime(night_events['sunrise'])],
                     2.5, 0.0, color='gray', alpha=0.3)
    ax.fill_between([mjd_to_datetime(night_events['sunrise'] - 2/24),
                    mjd_to_datetime(night_events['sun_n18_rising'])],
                    2.5, 0.0, color='pink', alpha=0.1)
    
    if not np.isnan(night_events['moonrise']):
        ax.axvline(mjd_to_datetime(night_events['moonrise']), linestyle='-', color='blue', alpha=0.3)
    if not np.isnan(night_events['moonset']):
        ax.axvline(mjd_to_datetime(night_events['moonset']), linestyle='-', color='red', alpha=0.3)
    
    colors = cc.glasbey_category10
    # Assign distinct target sets with different colors
    marker_colors = {}
    labels = {}
    count = 0
    if len(targets) > 0:
        for sp in targets.observation_reason.unique():
            marker_colors[sp] = colors[count]
            labels[sp] = sp
            count += 1
    
    if len(targets) > 0:
        visit_alpha = 0.7
        for sp in targets.observation_reason.unique():
            qq = targets.query("observation_reason == @sp")
            label = sp
            ax.plot(mjd_to_datetime(qq.mjd, 'tai'), qq.airmass, 
                    marker='o', linestyle='',
                    color=marker_colors[sp], label=label,
                    alpha=visit_alpha, markerfacecolor='none', zorder=3)
    
    ax.legend(loc=(1.01, 0.0), ncol=2)
    
    x0 = night_events['sunset']+30/60/24
    
    ax.set_xlim(mjd_to_datetime(night_events['sunset']+30/60/24), 
             mjd_to_datetime(night_events['sunrise']-30/60/24))
    ax_utc.set_xlim(mjd_to_datetime(night_events['sunset']+30/60/24, 'utc', timezone=tz_utc), 
             mjd_to_datetime(night_events['sunrise']-30/60/24, 'utc', timezone=tz_utc))
    
    ax.set_xlim(mjd_to_datetime(night_events['sunset']+30/60/24), 
             mjd_to_datetime(night_events['sunrise']-30/60/24))
    ax_utc.set_xlim(mjd_to_datetime(night_events['sunset']+30/60/24, 'utc', timezone=tz_utc), 
             mjd_to_datetime(night_events['sunrise']-30/60/24, 'utc', timezone=tz_utc))
    
    # Set ticks relevant sides
    ax.tick_params(axis="x", bottom=True, top=False, labelbottom=True, labeltop=False)
    ax_utc.tick_params(axis="x", bottom=False, top=True, labelbottom=False, labeltop=True)
    
    # Rotate and align bottom ticklabels
    plt.setp([tick.label1 for tick in ax.xaxis.get_major_ticks()], rotation=45,
             ha="right", va="center", rotation_mode="anchor")
    
    # Rotate and align top ticklabels
    plt.setp([tick.label2 for tick in ax_utc.xaxis.get_major_ticks()], rotation=45,
             ha="left", va="center",rotation_mode="anchor")
    
    plt.grid(True, alpha=0.2)
    
    plt.ylim(2.5, 0.9)
    
    ax.set_ylabel("Airmass", fontsize="large")
    ax.set_xlabel(f"Time ({tz})", fontsize="large")
    ax_utc.set_xlabel("Time (UTC)", fontsize='large')
    _ = plt.ylabel("Airmass", fontsize="large")

print(obsdf.target_name.unique())
print(obsdf.science_program.unique())
print(obsdf.observation_reason.unique())

['lowdust' 'ddf_edfs_a, lowdust' 'LMC_SMC' 'lowdust, LMC_SMC'
 'euclid_overlap, LMC_SMC' 'euclid_overlap, scp' 'euclid_overlap' 'scp'
 'scp, LMC_SMC' 'lowdust, ddf_edfs_b' 'scp, lowdust'
 'euclid_overlap, lowdust' 'lowdust, nes' 'nes' 'lowdust, ddf_elaiss1'
 'ddf_xmm_lss, lowdust' 'euclid_overlap, scp, LMC_SMC'
 'scp, lowdust, LMC_SMC' 'lowdust, ddf_ecdfs' 'bulgy, lowdust'
 'ddf_edfs_a, lowdust, ddf_edfs_b' 'dusty_plane, lowdust'
 'ddf_ecdfs, lowdust' 'euclid_overlap, lowdust, LMC_SMC'
 'ddf_elaiss1, lowdust' 'dusty_plane']
['BLOCK-407']
['singles_r' 'template_blob_r_33.0' 'template_blob_g_33.0' 'pairs_gr_33.0'
 'pairs_iz_15.0' 'template_blob_u_33.0' 'template_blob_i_33.0'
 'template_blob_z_33.0' 'pairs_yy_15.0' 'template_blob_y_33.0'
 'pairs_yy_33.0' 'ddf_ecdfs' 'ddf_elaiss1' 'pairs_iz_33.0' 'ddf_edfs_a'
 'ddf_xmm_lss' 'singles_i' 'pairs_ri_15.0' 'pairs_ri_33.0']

ddf_visits = obsdf.query("observation_reason.str.contains('DD') or observation_reason.str.contains('ddf')").copy()
print("n visits per band")
ddf_visits.loc[:, 'observation_reason'] = ddf_visits.observation_reason.str.lower()
ss = ddf_visits.groupby(["observation_reason", "band"]).agg({'observationStartMJD': 'count'})
ss = ss.reset_index('band').pivot(columns=["band"]).droplevel(0, axis=1)
ss['all'] = ss.sum(axis=1)
display(ss.query("observation_reason.str.contains('dd')")[['u', 'g', 'r', 'i', 'z', 'y', 'all']])

print("n days per band")
ss = ddf_visits.groupby(["observation_reason", "band"]).agg({'day_obs': 'nunique'})
ss = ss.reset_index('band').pivot(columns=["band"]).droplevel(0, axis=1)
ss['all'] = ddf_visits.groupby("observation_reason").agg({'day_obs': 'nunique'})
display(ss.query("observation_reason.str.contains('dd')")[['u', 'g', 'r', 'i', 'z', 'y', 'all']])

n visits per band

n days per band

run_calc = True
if run_calc:
    nvisits = {}
    coadd = {}
    m_nvis = maf.CountMetric(col='observationStartMJD', metric_name = "Nvisits")
    m_coadd = maf.Coaddm5Metric(m5_col='fiveSigmaDepth')
    s = maf.HealpixSlicer(nside=64, lon_col='fieldRA', lat_col='fieldDec', rot_sky_pos_col_name = 'rotSkyPos')
    for b in ['u', 'g', 'r', 'i', 'z', 'y', 'all']:
        constraint = f"{b}"
        if b == 'all':
            opsvis = obsdf.to_records()
        else:
            opsvis = obsdf.query("band == @b").to_records()
        nvisits[b] = maf.MetricBundle(m_nvis, s, constraint)
        coadd[b] = maf.MetricBundle(m_coadd, s, constraint)
        g = maf.MetricBundleGroup({f'nvisits {b}': nvisits[b], f'coadd {b}': coadd[b]}, None)
        g.run_current(constraint, opsvis)

Healpix slicer using NSIDE=64, approximate resolution 54.967783 arcminutes

background = plot.get_background(nside=64)

fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(16, 10),)
axdict = {"u": ax[0][0], "g": ax[0][1], "r": ax[0][2],
          "i": ax[1][0], "z": ax[1][1], "y": ax[1][2], "all": None}
for b in ["u", "g", "r", "i", "z", "y"]:
    if len(nvisits[b].metric_values.compressed()) > 1:
        vmax = np.percentile(nvisits[b].metric_values.compressed(), 95)
    else:
        vmax = None
    label_dec = False
    if b == 'u' or b == 'i':
        label_dec = True
    fig = plot.make_plot(nvisits[b], background=background, proj='McBryde', vmax=vmax, ax=axdict[b], title=f"LSSTCam band {b}", label_dec=label_dec)
fig.tight_layout()
#fig.savefig(os.path.join(out_dir, f"lsstcam_nvisits_band.png"), bbox_inches='tight')

vmax = np.percentile(nvisits['all'].metric_values.compressed(), 95)
fig = plot.make_plot(nvisits['all'], background=background, proj='mcbryde', vmin=None, vmax=vmax, ax=None, title=f"LSSTCam visits")
#fig.savefig(os.path.join(out_dir, f"lsstcam_nvisits.png"), bbox_inches='tight')

band	u	g	r	i	z	y	all
observation_reason
ddf_ecdfs	15	10	12	10	12	6	65
ddf_edfs_a	16	10	12	10	12	6	66
ddf_elaiss1	18	14	16	14	16	10	88
ddf_xmm_lss	15	12	14	12	14	8	75