import numpy as np
import matplotlib.pyplot as plt

from astroquery.vizier import Vizier
from astroquery.simbad import Simbad
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.coordinates import AltAz
from astropy.coordinates import EarthLocation
from astropy.time import Time
from astropy.utils.iers import conf
conf.auto_max_age = None

%matplotlib inline



def unix_timestamp_from_iso_date(datetime):
    return Time(datetime, format = 'isot').unix

def iso_date_from_unix_timestamp(datetime):
    return Time(datetime, format = 'unix').isot

def radec_to_altaz(lat, lon, datetime, ra, dec):
    """
    Convert equatorial coordinates to local altitude and azimuth. The algorithm is approximate, not accounting
    for the atmosphere, precession, nutation and other higher order effects. For XX and XXI century dates, the
    expected error is 0.3 deg in alt and az * cos(alt). Maximum errors of 1.4 deg are possible for early XX or
    late XXI century dates

    args:
        lat       :    Latitude of the observatory in degrees (-90 <= lat <= 90)
        lon       :    Longitude of the observatory in degrees
        datetime  :    ISO date and time of observation (e.g. 2000-01-01T12:00:00Z for J2000)
        ra        :    Right Ascension of the target in degrees, J2000
        dec       :    Declination of the target in degrees (-90 <= dec <= 90), J2000

    returns:
        alt       :    Altitude of the target in degrees
        az        :    Azimuth of the target in degrees
    """
    # Days since J2000
    unix_timestamp_requested = unix_timestamp_from_iso_date(datetime)
    unix_timestamp_J2000 = unix_timestamp_from_iso_date('2000-01-01T12:00:00Z')
    J2000 = (unix_timestamp_requested - unix_timestamp_J2000) / (60 * 60 * 24)

    # Check that the longitude complies with the format required by this approximation
    lon = lon % 360
    if lon > 180.0:
        lon = lon - 360
    # Siderial time and hour angle
    LST = 100.46 + 0.985647 * J2000 + lon + 15 * (J2000 - 0.5 - np.floor(J2000)) * 24
    LST = LST % 360
    HA = LST - ra

    # Hour angle
    dec, lat, HA = np.radians((dec, lat, HA))
    alt = np.arcsin(np.sin(dec) * np.sin(lat) + np.cos(dec) * np.cos(lat) * np.cos(HA))
    az = np.arccos((np.sin(dec) - np.sin(alt) * np.sin(lat)) / (np.cos(alt) * np.cos(lat)))
    alt, az = np.degrees((alt, az))
    if np.sin(HA) > 0:
        az = 360 - az
    return alt, az

def ascent_time(lat, lon, datetime_start, datetime_end, ra, dec, alt):
    """
    Determine the times when the target reaches a certain altitude within a window of dates. See the docstring
    of radec_to_altaz() for a note on the accuracy

    args:
        lat             :    Latitude of the observatory in degrees (-90 <= lat <= 90)
        lon             :    Longitude of the observatory in degrees
        datetime_start  :    Start date of the window (ISO)
        datetime_end    :    End date of the window (ISO)
        ra              :    Right Ascension of the target in degrees, J2000
        dec             :    Declination of the target in degrees (-90 <= dec <= 90), J2000
        alt             :    Desired altitude

    returns:
        ascents         :    List of all times (as ISO dates) when the target is at the desired altitude.
                             Empty list if the target never reaches the altitude within the given window
    """
    # Get the hour angle of the target
    alt, dec, lat = np.radians((alt, dec, lat))
    cos_HA = (np.sin(alt) - np.sin(dec) * np.sin(lat)) / (np.cos(dec) * np.cos(lat))
    # The target never rises to the desired altitude
    if np.abs(cos_HA) >= 1.0:
        return []
    # Now, we have the cosine of the HA. Of course, there will be infinitely many values of HA that
    # have that cosine of two families: arccos(cos_HA) + n*360 and -arccos(cos_HA) + n*360

    # Days since J2000
    unix_timestamp_end = unix_timestamp_from_iso_date(datetime_end)
    unix_timestamp_start = unix_timestamp_from_iso_date(datetime_start)
    unix_timestamp_J2000 = unix_timestamp_from_iso_date('2000-01-01T12:00:00Z')
    J2000_end = (unix_timestamp_end - unix_timestamp_J2000) / (60 * 60 * 24)
    J2000_start = (unix_timestamp_start - unix_timestamp_J2000) / (60 * 60 * 24)

    ascents = []     # Store found ascents here

    for phase in [1, -1]: # "phase" refers to the family of HA values considered. We search for ascents in both
        # Check that the longitude complies with the format required by this approximation
        lon = lon % 360
        if lon > 180.0:
            lon = lon - 360
        # Get the HA value of the right family without n*360 for now
        HA = np.degrees(np.arccos(cos_HA)) * phase
        # Get the LST without n*360 as well
        LST = HA + ra
        # Try different values of n until LST matches the window
        while (LST - 100.46 + 180 - lon - 360 * 1) / 0.985647 > J2000_start:
            LST -= 360
        while (LST - 100.46 + 180 - lon) / 0.985647 < J2000_start:
            LST += 360
        # Get the UT time of the ascent on the first day of the window as well as the rate
        # at which it will drift per day
        UT_start = (LST - 100.46 + 180 - lon - 0.985647 * np.floor(J2000_start)) / 360 / (1 + 0.985647 / 360)
        UT_offset = -0.985647 / 360

        # Collect all ascents
        i = 0
        while True:
            if i > 1000:
                raise ValueError()
            ascent = np.floor(J2000_start) + i + UT_start + i * UT_offset
            i += 1
            if ascent < J2000_start:
                continue
            if ascent > J2000_end:
                break
            ascents += [iso_date_from_unix_timestamp(ascent * ((60 * 60 * 24)) + unix_timestamp_J2000)]
    return ascents

def climax(lat, lon, datetime_start, datetime_end, ra, dec):
    """
    Determine the maximum altitude reached by the target and the time when it is reached within a window
    of dates. See the docstring of radec_to_altaz() for a note on the accuracy

    args:
        lat             :    Latitude of the observatory in degrees (-90 <= lat <= 90)
        lon             :    Longitude of the observatory in degrees
        datetime_start  :    Start date of the window (ISO)
        datetime_end    :    End date of the window (ISO)
        ra              :    Right Ascension of the target in degrees, J2000
        dec             :    Declination of the target in degrees (-90 <= dec <= 90), J2000

    returns:
        alt             :    Maximum altitude within the window in degrees
        datetime        :    ISO date of when the altitude is reached. If reached multiple times, only the
                             first instance is returned
    """
    # Find the minimum and maximum altitudes attainable by the target
    extrema = [-lat - dec - 90, dec - lat + 90]
    for i, extremum in enumerate(extrema):
        if extremum > 90 or extremum < -90:
            extrema[i] = (180 - extremum)
        if extremum < -90:
            extrema[i] = extrema[i] - 360

    # Get all times when the object is at its highest
    ascents = ascent_time(lat, lon, datetime_start, datetime_end, ra, dec, np.max(extrema))
    # If the object does reach its highest possible altitude, it will be the desired climax
    if len(ascents) != 0:
        return np.max(extrema), ascents[0]
    else:
    # If not, either the start or the end of the window will be...
        start_alt = radec_to_altaz(lat, lon, datetime_start, ra, dec)[0]
        end_alt = radec_to_altaz(lat, lon, datetime_end, ra, dec)[0]
        if start_alt > end_alt:
            return start_alt, datetime_start
        else:
            return end_alt, datetime_end

def get_sun(datetime):
    """
    Estimate J2000 equatorial coordinates of the Sun at any given time. Use the approximate algorithm
    from stargazing.net/kepler/sun.html

    args:
        datetime        :    Date of interest (ISO)

    returns:
        ra              :    Right ascension of the Sun in degrees
        dec             :    Declination of the Sun in degrees
    """
    # Days since J2000
    unix_timestamp_requested = unix_timestamp_from_iso_date(datetime)
    unix_timestamp_J2000 = unix_timestamp_from_iso_date('2000-01-01T12:00:00Z')
    J2000 = (unix_timestamp_requested - unix_timestamp_J2000) / (60 * 60 * 24)

    # Just following the algorithm from now on...
    L = (280.461 + 0.9856474 * J2000) % 360
    g = np.radians(357.528 + 0.9856003 * J2000)
    lmbda = np.radians(L + 1.915 * np.sin(g) + 0.020 * np.sin(2 * g))
    epsilon = np.radians(23.439 - 0.0000004 * J2000)
    Y = np.cos(epsilon) * np.sin(lmbda)
    X = np.cos(lmbda)

    ra = np.degrees(np.arctan(Y / X))
    if X < 0:
        ra = ra + 180
    elif Y < 0 and X > 0:
        ra = ra + 360
    dec = np.degrees(np.arcsin(np.sin(epsilon) * np.sin(lmbda)))
    return ra, dec

def twilight(lat, lon, datetime, alt, niter = 5):
    """
    Iteratively determine sunrise/sunset/twilight times on a given date. ascent_time() can determine when
    a given pair of equatorial coordinates reaches a given altitude. However, in order to know the coordinates
    of the Sun when its transiting said altitude, one must know the time when it happens and feed it into
    get_sun(). This recursive loop can be solved iteratively.

    Sunset/sunrise/twilight events are searched within 24 hours of the given time. On the 0th iteration, the
    average coordinates of the Sun are assumed for the search period. Once the events are found, the coordinates
    of the Sun are updated and the search is repeated. Starting with the 1st iteration, each event is searched
    separately within 2 hours of the event time known from the previous iteration.

    See the docstring of radec_to_altaz() for a note on the accuracy. Ultimately, a 1 degree error in the RA/Dec
    to AltAz conversion algorithm can result in a 4 minute error in sunset/sunrise times in the worst case
    scenario

    args:
        lat             :    Latitude of the observatory in degrees (-90 <= lat <= 90)
        lon             :    Longitude of the observatory in degrees
        datetime        :    Date of interest (ISO). Events will be searched within 24 hours of this date
                             (12 hours in each direction)
        alt             :    Altitude of the event. "0" for sunset/sunrise, "-6" for civil twlight etc...
        n_iter          :    Number of iterations. Typically, takes only 3 or so to converge to seconds

    returns:
        sunrise         :    Date of the first event (e.g. sunrise) in ISO
        sunset          :    Date of the second event (e.g. sunset) in ISO
    """
    # Get ISO dates defining the initial search window
    datetime_unix = unix_timestamp_from_iso_date(datetime)
    datetime_start = iso_date_from_unix_timestamp(datetime_unix - 12 * 60 * 60)
    datetime_end = iso_date_from_unix_timestamp(datetime_unix + 12 * 60 * 60)

    # Look for events. We must find 2!
    ascents = ascent_time(lat, lon, datetime_start, datetime_end, *get_sun(datetime), alt)
    if len(ascents) != 2:
        raise ValueError('Expected two ascension events within the window, {} found'.format(len(ascents)))

    for i, ascent in enumerate(ascents):
        # Iterate for each event
        for j in range(niter):
            # Define the new search window (narrower on all subsequent iterations)
            ascent_unix = unix_timestamp_from_iso_date(ascent)
            datetime_start = iso_date_from_unix_timestamp(ascent_unix - 60 * 60)
            datetime_end = iso_date_from_unix_timestamp(ascent_unix + 60 * 60)
            # Search again
            updated_ascents = ascent_time(lat, lon, datetime_start, datetime_end, *get_sun(ascent), alt)
            # Due to the narrow window, we must only find one even this time
            if len(updated_ascents) != 1:
                raise ValueError('Ascension events too frequent (more than one in two hours)'.format(len(ascents)))
            ascent = updated_ascents[0]
        ascents[i] = ascent
    return ascents