import astropy.units as u
from astropy.time import Time
from astropy.coordinates import SkyCoord, EarthLocation, AltAz, get_sun, get_body
from astropy.utils.iers import conf
from scipy.optimize import root_scalar
import numpy as np
from datetime import datetime, timedelta
from colorama import Fore, Back, Style, init
import csv

init(autoreset=True) # Initializes Colorama and auto-resets color after each print

conf.auto_download = False
conf.auto_max_age = None
#
# Script to determine the amount of time a target is in view during astronomical darkness
#
# Some math concepts for this were taken from https://airmass.org/notes
#
# AstroVis v1.5 by xenospiritual @ https://xenomorpheus.neocities.org/vis.html [sorry, too lazy to make/maintain a changelog at this time]
#

# Hardcoded observer location that is not too specific to avoid goons who would poach my gear while I'm napping during a session >:[
location = EarthLocation(lat=39*u.deg, lon=-77*u.deg, height=200*u.m)

# Hardcoded altitude limit (degrees) since I am limited by a roof and light pollution *ick*
ALT_LIMIT = 40

# Hardcoded UTC offset (hours) for local time. DST status is checked later via prompt and then applied if in effect.
utc_offset = -5

# Define find_root function for numerical solving
def find_root(f, bracket):
    try:
        a, b = bracket
        fa, fb = f(a), f(b)
        if fa * fb >= 0:
            print(f"Debug: f(a)={fa:.2f}, f(b)={fb:.2f} at JD {a:.2f}, {b:.2f}")
            raise ValueError("Function does not change sign in bracket.")
        result = root_scalar(f, bracket=bracket, method='brentq')
        return result.root
    except ValueError as e:
        raise ValueError(f"Root finding failed: {str(e)}. Check input or visibility.")

# Round time to nearest minute
def round_time_to_minute(time):
    dt = time.to_datetime()
    dt_rounded = dt + timedelta(seconds=30)  # Round to nearest minute
    dt_rounded = dt_rounded.replace(second=0, microsecond=0)
    return Time(dt_rounded)

# Format time for clean output
def format_time(time):
    return time.iso.split('.')[0]

# Introduction, instructions, usage?
print(Fore.BLUE + "\nAstrophotography Planning Script v1.5 by xenospiritual")
print(Fore.LIGHTBLUE_EX + "\nThis script is for calculating the amount of time a target")
print(Fore.LIGHTBLUE_EX + f"is above {ALT_LIMIT}° altitude during astronomical darkness.\n")

# User inputs for date and DST
try:
    date_str = input("Enter date (YYYY-MM-DD): ")
    try:
        midnight = Time(date_str + ' 00:00:00')
    except ValueError:
        raise ValueError("Invalid date format. Use YYYY-MM-DD.")
    while True:
        dst_input = input("Is DST in effect on this date? (y/n): ").lower()
        if dst_input in ('y', 'n'):
            dst_time = dst_input == 'y'
            if dst_time:
                utc_offset += 1
            break
        print("Please enter 'y' or 'n'.")
except ValueError as e:
    print(f"Input error: {e}")
    exit(1)

# Load targets from CSV or prompt for single target
targets = []
load_csv = input("Load targets from CSV? (y/n): ").lower() == 'y'
if load_csv:
    try:
        with open('targets.csv', 'r') as f:
            reader = csv.reader(f)
            next(reader)  # Skip header
            for row in reader:
                if len(row) != 3:
                    print(f"Skipping invalid CSV row: {row}")
                    continue
                targets.append({'name': row[0], 'ra': row[1], 'dec': row[2]})
        if not targets:
            print("Error: No valid targets found in targets.csv.")
            exit(1)
    except FileNotFoundError:
        print("Error: targets.csv not found.")
        exit(1)
else:
    try:
        ra_str = input("Enter RA (hours minutes seconds, e.g., '1 33 50.9'): ")
        ra_parts = ra_str.split()
        if len(ra_parts) != 3:
            raise ValueError("RA must have 3 parts (hours, minutes, seconds).")
        ra_h, ra_m, ra_s = map(float, ra_parts)
        if not (0 <= ra_h < 24 and 0 <= ra_m < 60 and 0 <= ra_s < 60):
            raise ValueError("RA values out of range (0-24h, 0-60m, 0-60s).")

        dec_str = input("Enter DEC (degrees arcmin arcsec, e.g., '30 39 36'): ")
        dec_parts = dec_str.split()
        if len(dec_parts) != 3:
            raise ValueError("DEC must have 3 parts (degrees, arcmin, arcsec).")
        dec_d, dec_m, dec_s = map(float, dec_parts)
        if not (-90 <= dec_d <= 90 and 0 <= dec_m < 60 and 0 <= dec_s < 60):
            raise ValueError("DEC values out of range (-90 to 90°, 0-60', 0-60\").")

        name = input("Enter target name (e.g., M33): ")
        if not name.strip():
            name = "Target"
        targets.append({'name': name, 'ra': ra_str, 'dec': dec_str})
    except ValueError as e:
        print(f"Input error: {e}")
        exit(1)

# Process each target
for target in targets:
    try:
        ra_str = target['ra']
        dec_str = target['dec']
        name = target['name']

        # Parse RA and DEC
        ra_parts = ra_str.split()
        if len(ra_parts) != 3:
            print(f"Skipping target {name}: RA must have 3 parts.")
            continue
        ra_h, ra_m, ra_s = map(float, ra_parts)
        if not (0 <= ra_h < 24 and 0 <= ra_m < 60 and 0 <= ra_s < 60):
            print(f"Skipping target {name}: RA values out of range.")
            continue
        ra_deg = (ra_h + ra_m/60 + ra_s/3600) * 15  # Convert to degrees

        dec_parts = dec_str.split()
        if len(dec_parts) != 3:
            print(f"Skipping target {name}: DEC must have 3 parts.")
            continue
        dec_d, dec_m, dec_s = map(float, dec_parts)
        if not (-90 <= dec_d <= 90 and 0 <= dec_m < 60 and 0 <= dec_s < 60):
            print(f"Skipping target {name}: DEC values out of range.")
            continue
        dec_deg = abs(dec_d) + dec_m/60 + dec_s/3600
        if dec_d < 0:
            dec_deg = -dec_deg
        # print(f"Debug: {name}: RA = {ra_deg:.2f}°, DEC = {dec_deg:.2f}°")  # Debugging to help fix my negative declination handling error

        # Create SkyCoord object
        obj = SkyCoord(ra=ra_deg*u.deg, dec=dec_deg*u.deg)

        # Print result header
        print(f"\n{Fore.BLUE}Astrophotography Planning for {name} (RA {ra_str}, DEC {dec_str}) on {date_str}")
        print(f"{Fore.LIGHTBLUE_EX}Location: Lat {location.lat.deg:.2f}°, Lon {location.lon.deg:.2f}°, Elevation {location.height.value:.0f} m\n")

        # Calculate astronomical twilight using astropy
        def sun_alt(t_jd):
            time = Time(t_jd, format='jd')
            sun = get_sun(time)
            altaz = sun.transform_to(AltAz(obstime=time, location=location))
            return altaz.alt.deg

        # Find sunset and sunrise for reference [workaround for an error/problem: don't ask!]
        approx_sunset = midnight + (20 - utc_offset) / 24 * u.day  # ~20:00
        f_sunset = lambda t: sun_alt(t)
        try:
            sunset_jd = find_root(f_sunset, [approx_sunset.jd - 0.1, approx_sunset.jd + 0.1])
            sunset = Time(sunset_jd, format='jd')
            sunset_local = round_time_to_minute(sunset + utc_offset * u.hour)
        except ValueError as e:
            print(f"Failed to find sunset for {name}: {e}")
            continue

        approx_sunrise = midnight + 1 * u.day + (6 - utc_offset) / 24 * u.day  # ~06:00
        f_sunrise = lambda t: sun_alt(t)
        try:
            sunrise_jd = find_root(f_sunrise, [approx_sunrise.jd - 0.1, approx_sunrise.jd + 0.1])
            sunrise = Time(sunrise_jd, format='jd')
            sunrise_local = round_time_to_minute(sunrise + utc_offset * u.hour)
        except ValueError as e:
            print(f"Failed to find sunrise for {name}: {e}")
            continue

        # Find astronomical twilight [same workaround lol...]
        approx_twilight_start = midnight + (21.75 - utc_offset) / 24 * u.day  # ~21:45
        f_twilight_start = lambda t: sun_alt(t) + 18
        try:
            dark_start_jd = find_root(f_twilight_start, [approx_twilight_start.jd - 0.1, approx_twilight_start.jd + 0.1])
            dark_start = Time(dark_start_jd, format='jd')
            dark_start_local = round_time_to_minute(dark_start + utc_offset * u.hour)
        except ValueError as e:
            print(f"Failed to find twilight start for {name}: {e}")
            continue

        approx_twilight_end = midnight + 1 * u.day + (4.75 - utc_offset) / 24 * u.day  # ~04:45
        f_twilight_end = lambda t: sun_alt(t) + 18
        try:
            dark_end_jd = find_root(f_twilight_end, [approx_twilight_end.jd - 0.1, approx_twilight_end.jd + 0.1])
            dark_end = Time(dark_end_jd, format='jd')
            dark_end_local = round_time_to_minute(dark_end + utc_offset * u.hour)
        except ValueError as e:
            print(f"Failed to find twilight end for {name}: {e}")
            continue

        # Ensure dark_start_jd < dark_end_jd
        if dark_start_jd > dark_end_jd:
            print(f"Warning: Twilight start is after end for {name}; swapping to correct order.")
            dark_start_jd, dark_end_jd = dark_end_jd, dark_start_jd
            dark_start, dark_end = dark_end, dark_start
            dark_start_local, dark_end_local = dark_end_local, dark_start_local

        # Print twilight times
        print(f"Astronomical darkness starts (Sun at -18°): {format_time(dark_start_local)}")
        print(f"Astronomical darkness ends (Sun at -18°): {format_time(dark_end_local)}")

        # Function to get object altitude
        def obj_alt(t_jd):
            time = Time(t_jd, format='jd')
            altaz = obj.transform_to(AltAz(obstime=time, location=location))
            return altaz.alt.deg

        # Sample full day to find max altitude, then filter for astronomical darkness
        day_start_jd = midnight.jd - 0.5  # Start at previous noon
        day_end_jd = midnight.jd + 1.5  # End at next noon
        num_day_samples = int(24 * 60 / 5) + 1  # Sample every ~5 minutes
        day_times_jd = np.linspace(day_start_jd, day_end_jd, num_day_samples)
        alts = np.array([obj_alt(t) for t in day_times_jd])

        # Calculate maximum altitude and validate
        max_alt = max(alts)
        # Rough estimate of expected max altitude: 90 - |latitude - DEC|
        expected_max_alt = 90 - abs(location.lat.deg - dec_deg)
        if abs(max_alt - expected_max_alt) > 5:
            print(f"Warning: Max altitude {max_alt:.1f}° seems off for {name}; expected ~{expected_max_alt:.1f}° for DEC {dec_deg:.1f}°.")

        # Sample night to find when object is above ALT_LIMIT
        dark_duration_h = (dark_end - dark_start).to(u.hour).value
        if dark_duration_h <= 0:
            print(f"Error: Astronomical darkness duration is non-positive for {name}. Check calculations.")
            continue
        num_samples = int(dark_duration_h * 60 / 5) + 1  # Sample every ~5 minutes
        times_jd = np.linspace(dark_start_jd, dark_end_jd, num_samples)
        alts_night = np.array([obj_alt(t) for t in times_jd])

        # Find indices where object is above ALT_LIMIT
        above = np.where(alts_night > ALT_LIMIT)[0]

        if len(above) == 0:
            print(f"{name} never rises above {ALT_LIMIT}° during the night.")
            total_time_h = 0
        else:
            first_idx = above[0]
            last_idx = above[-1]
            if first_idx == 0:
                rise_jd = dark_start_jd
                rise_time = dark_start
            else:
                bracket_rise = [times_jd[first_idx-1], times_jd[first_idx]]
                f_obj = lambda t: obj_alt(t) - ALT_LIMIT
                try:
                    rise_jd = find_root(f_obj, bracket_rise)
                    rise_time = Time(rise_jd, format='jd')
                except ValueError as e:
                    print(f"Failed to find rise time for {name}: {e}")
                    rise_time = None
            if last_idx == num_samples - 1:
                set_jd = dark_end_jd
                set_time = dark_end
            else:
                bracket_set = [times_jd[last_idx], times_jd[last_idx+1]]
                f_obj = lambda t: obj_alt(t) - ALT_LIMIT
                try:
                    set_jd = find_root(f_obj, bracket_set)
                    set_time = Time(set_jd, format='jd')
                except ValueError as e:
                    print(f"Failed to find set time for {name}: {e}")
                    set_time = None
            if rise_time and set_time:
                rise_time_local = round_time_to_minute(rise_time + utc_offset * u.hour)
                set_time_local = round_time_to_minute(set_time + utc_offset * u.hour)
                total_time_h = (set_time - rise_time).to(u.hour).value
                time_color = Fore.RED if total_time_h < 4 else Fore.GREEN if total_time_h > 4 else Fore.YELLOW
                print(f"{name} rises above {ALT_LIMIT}° at: {format_time(rise_time_local)}")
                print(f"{name} sets below {ALT_LIMIT}° at: {format_time(set_time_local)}")
                print(f"Total time above {ALT_LIMIT}° during astronomical darkness: {Style.BRIGHT}{time_color}{total_time_h:.2f} hours{Style.RESET_ALL}")
            else:
                print(f"{name} time above {ALT_LIMIT}° could not be calculated.")
                total_time_h = 0

        # Color-code max altitude (yellow for <30° to flag low targets) since I personally can see the horizon at about 20° altitude to the south, it's right in the city's light dome.
        alt_color = Fore.YELLOW if max_alt < 30 else ''
        print(f"Maximum altitude during night: {alt_color}{max_alt:.1f}°{Style.RESET_ALL}")

        # Moon calculations over extended window
        def moon_alt(t_jd):
            time = Time(t_jd, format='jd')
            moon = get_body('moon', time)
            altaz = moon.transform_to(AltAz(obstime=time, location=location))
            return altaz.alt.deg

        # Search from 12:00 local time on input date to 23:59 local time the following day [Some dates don't have a moonrise, since the moon may rise after midnight!]
        day_start = midnight + (12 - utc_offset) / 24 * u.day  # 12:00
        day_end = midnight + 1 * u.day + (23.99 - utc_offset) / 24 * u.day  # 23:59 next day
        num_day = 288  # Sample every ~5 minutes
        day_times = np.linspace(day_start.jd, day_end.jd, num_day)
        moon_alts = np.array([moon_alt(t) for t in day_times])

        # Find moon rise and set
        sign_changes = np.where(np.sign(moon_alts[:-1]) != np.sign(moon_alts[1:]))[0]

        moon_rise = None
        moon_set = None
        for i in sign_changes:
            bracket = [day_times[i], day_times[i+1]]
            f_moon = lambda t: moon_alt(t)
            try:
                root_jd = find_root(f_moon, bracket)
                root_time = Time(root_jd, format='jd')
                root_time_local = round_time_to_minute(root_time + utc_offset * u.hour)
                if moon_alts[i] < 0:
                    moon_rise = root_time_local
                else:
                    moon_set = root_time_local
            except ValueError as e:
                print(f"Failed to find moon rise/set for {name}: {e}")

        if moon_rise:
            print(f"Moon rise: {format_time(moon_rise)}")
        else:
            print(f"No moon rise during session ({format_time(day_start + utc_offset * u.hour)} to {format_time(day_end + utc_offset * u.hour)}).")
        if moon_set:
            print(f"Moon set: {format_time(moon_set)}")
        else:
            print(f"No moon set during session ({format_time(day_start + utc_offset * u.hour)} to {format_time(day_end + utc_offset * u.hour)}).")

        # Moon phase at midnight UTC
        time = midnight
        sun = get_sun(time)
        moon = get_body('moon', time)
        elong = sun.separation(moon).radian
        illum = 0.5 * (1 - np.cos(elong)) * 100

        # Determine waxing/waning
        time2 = time + 1*u.hour
        sun2 = get_sun(time2)
        moon2 = get_body('moon', time2)
        elong2 = sun2.separation(moon2).radian
        illum2 = 0.5 * (1 - np.cos(elong2)) * 100

        trend = "Waxing" if illum2 > illum else "Waning"

        # Determine phase
        if illum < 1:
            phase = "New Moon"
        elif illum > 99:
            phase = "Full Moon"
        elif abs(illum - 50) < 1:
            phase = "First Quarter" if trend == "Waxing" else "Last Quarter"
        elif illum < 50:
            phase = trend + " Crescent"
        else:
            phase = trend + " Gibbous"

        print(f"Moon phase: {phase}: {round(illum)}% illuminated")

        # Optionally save output to a txt file
        save_output = input(f"\nSave results for {name} to a file? (y/n): ").lower() == 'y'
        if save_output:
            filename = f"astro_planning_{date_str}_{name.replace(' ', '_')}.txt"
            with open(filename, 'w') as f:
                f.write(f"Astrophotography Planning for {name} (RA {ra_str}, DEC {dec_str}) on {date_str}\n")
                f.write(f"Location: Lat {location.lat.deg:.2f}°, Lon {location.lon.deg:.2f}°, Elevation {location.height.value:.0f} m\n")
                f.write(f"DST in effect: {'Yes' if dst_time else 'No'}\n")
                f.write(f"Astronomical darkness starts (Sun at -18°): {format_time(dark_start_local)}\n")
                f.write(f"Astronomical darkness ends (Sun at -18°): {format_time(dark_end_local)}\n")
                if len(above) > 0 and rise_time and set_time:
                    f.write(f"{name} rises above {ALT_LIMIT}° at: {format_time(rise_time_local)}\n")
                    f.write(f"{name} sets below {ALT_LIMIT}° at: {format_time(set_time_local)}\n")
                    f.write(f"Total time above {ALT_LIMIT}° during astronomical darkness: {total_time_h:.2f} hours\n")
                else:
                    f.write(f"{name} never rises above {ALT_LIMIT}° during the night.\n")
                f.write(f"Maximum altitude during night: {max_alt:.1f}°\n")
                f.write(f"Moon rise: {format_time(moon_rise) if moon_rise else 'None during session'}\n")
                f.write(f"Moon set: {format_time(moon_set) if moon_set else 'None during session'}\n")
                f.write(f"Moon phase: {phase}: {round(illum)}% illuminated\n")
            print(f"Results saved to {filename}")
    except Exception as e:
        print(f"Error processing target {name}: {e}")
        continue