import astropy.units as unit
import numpy as np
from astropy.coordinates import EarthLocation, AltAz
from astropy.coordinates import get_sun
from astropy.time import TimeISO
from input.natural_constants import *


class TimeYearDayTimeCustom(TimeISO):
   """
   day-of-year as "<DOY>".
   The day-of-year (DOY) goes from 001 to 365 (366 in leap years).
   The allowed subformat is:
   - 'doy': day of year
   """
   name = 'doy'  # unique format name
   subfmts = (('doy',
               '%j',
               '{yday:03d}'),
              ('doy',
               '%j',
               '{yday:03d}'),
              ('doy',
               '%j',
               '{yday:03d}'))


def sun_angles_astropy(lat, lon, h, utc):  # get current sun elevation and azimuth through astropy
    loc = EarthLocation(lat=lat*unit.deg, lon=lon*unit.deg, height=h*unit.m)
    ref = AltAz(obstime=utc, location=loc)

    sun_pos = get_sun(utc).transform_to(ref)

    az = sun_pos.az.degree
    elv = sun_pos.alt.degree

    return az, elv


def sun_angles_analytical(lat, lon, utc):  # get current sun elevation and azimuth through several equations (see [xx])
    if np.abs(lat) == 90:                  # handling collapse of longitudes at poles by
        lat = np.sign(lat) * 89.999999     # expanding one point to a very small circle
    else:
        pass

    jd = utc.jd
    jc = (jd - 2451545) / 36525
    gml = (280.46646 + jc * (36000.76983 + jc * 0.0003032)) % 360
    gma = 357.52911 + jc * (35999.05029 - 0.0001537 * jc)
    eeo = 0.016708634 - jc * (0.000042037 + 0.0000001267 * jc)
    sec = np.sin(np.deg2rad(gma)) * (1.914602 - jc * (0.004817 + 0.000014 * jc)) + np.sin(np.deg2rad(2 * gma)) * (
                0.019993 - 0.000101 * jc) + np.sin(np.deg2rad(3 * gma)) * 0.000289
    stl = gml + sec
    sal = stl - 0.00569 - 0.00478 * np.sin(np.deg2rad(125.04 - 1934.136 * jc))
    moe = 23 + (26 + (21.448 - jc * (46.815 + jc * (0.00059 - jc * 0.001813))) / 60) / 60
    oc = moe + 0.00256 * np.cos(np.deg2rad(125.04 - 1934.136 * jc))
    sd = np.rad2deg(np.arcsin(np.sin(np.deg2rad(oc)) * np.sin(np.deg2rad(sal))))  # radian
    var_y = np.tan(np.deg2rad(oc / 2)) ** 2
    eot = 4 * np.rad2deg(
        var_y * np.sin(2 * np.deg2rad(gml)) - 2 * eeo * np.sin(np.deg2rad(gma)) + 4 * eeo * var_y * np.sin(
            np.deg2rad(gma)) * np.cos(2 * np.deg2rad(gml)) - 0.5 * var_y ** 2 * np.sin(
            4 * np.deg2rad(gml)) - 1.25 * eeo ** 2 * np.sin(2 * np.deg2rad(gma)))
    tst = (((jd - 0.5) % 1) * 1440 + eot + 4 * lon) % 1440

    if tst / 4 < 0:
        ha = tst / 4 + 180
    else:
        ha = tst / 4 - 180

    sza = np.rad2deg(np.arccos(
        np.sin(np.deg2rad(lat)) * np.sin(np.deg2rad(sd)) + np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(sd)) * np.cos(
            np.deg2rad(ha))))
    sea = 90 - sza

    if ha > 0:
        saa = (np.rad2deg(np.arccos(((np.sin(np.deg2rad(lat)) * np.cos(np.deg2rad(sza))) - np.sin(np.deg2rad(sd))) / (
                    np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(sza))))) + 180) % 360
    else:
        saa = (540 - np.rad2deg(np.arccos(
            ((np.sin(np.deg2rad(lat)) * np.cos(np.deg2rad(sza))) - np.sin(np.deg2rad(sd))) / (
                        np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(sza)))))) % 360

    return saa, sea  # Azimuth, Elevation


def AirMass(p_air, p_0, ELV, h):      # get atmospheric air mass over balloon
    ELV_rad = np.deg2rad(ELV)         # convert ELV from degree to radian
    Dip = np.arccos(R_E / (R_E + h))  # geometric "dip" in radian

    if ELV_rad >= -Dip and ELV_rad < 0:
        res = p_air/p_0 * (1 + ELV_rad/Dip) - 70 * ELV_rad/Dip
    else:
        res = (p_air/p_0) * ((1229 + (614 * np.sin(ELV_rad)) ** 2) ** (1/2) - 614 * np.sin(ELV_rad))

    return res


def tau(ELV, h, p_air):  # get atmospheric transmissivity as function of balloon altitude and sun elevation
    if ELV >= -(180 / np.pi * np.arccos(R_E / (R_E + h))):
        tau_atm = 0.5 * (
                    np.exp(-0.65 * AirMass(p_air, p_0, ELV, h)) + np.exp(-0.095 * AirMass(p_air, p_0, ELV, h)))
        tau_atmIR = 1.716 - 0.5 * (np.exp(-0.65 * p_air / p_0) + np.exp(-0.095 * p_air / p_0))
    else:
        tau_atm = 0
        tau_atmIR = 0
    return tau_atm, tau_atmIR