# AUTOGENERATED! DO NOT EDIT! File to edit: 03_ringutils.ipynb (unless otherwise specified).

__all__ = ['get_order', 'get_resonances', 'get_prime_resonances', 'get_janus_epimetheus_resonances',
           'get_prime_jan_epi', 'get_all_prime_resonances', 'which_epi_janus_resonance']

# Cell
import pandas as pd
import pkg_resources as pr
from astropy.time import Time

# Cell
def get_order(name):
    ratio = name.split()[1]
    a, b = ratio.split(":")
    return int(a) - int(b)


def get_resonances():
    with pr.resource_stream("pyciss", "data/ring_resonances.csv") as f:
        resonances = pd.read_csv(f)
    resonances.columns = ["name", "radius", "a_moon", "n", "kappa"]
    resonances = resonances.sort_values(by="radius", ascending=True)
    resonances["order"] = resonances.name.map(get_order)
    return resonances

# Cell
def get_prime_resonances():
    resonances = get_resonances()
    prime_resonances = resonances[resonances.order == 1].drop("order", axis=1)
    # filter out Janus and Epimetheus as we have a more precise file for that.
    prime_resonances = prime_resonances.loc[
        ~prime_resonances.name.str.startswith("Janus")
    ]
    prime_resonances = prime_resonances.loc[
        ~prime_resonances.name.str.startswith("Epimetheus")
    ]
    return prime_resonances

# Cell
def get_janus_epimetheus_resonances():
    "Filter resonances for only Janus and Epithemeus resonances."
    w = [len("              Janus1"), len(" reson"), len("  Resonance radius R")]

    def get_janos_epi_order(reso):
        a, b = reso.split(":")
        return int(a) - int(b)

    fname = pr.resource_filename("pyciss", "data/ring_janus_epimetheus_resonances.txt")
    with open(fname) as f:
        jan_epi_resonances = pd.read_fwf(
            f, skiprows=15, header=0, widths=w, skipfooter=1
        )

    # replace column names
    jan_epi_resonances.columns = ["moon", "reson", "radius"]

    # calculate order from resonance name
    jan_epi_resonances["order"] = jan_epi_resonances.reson.map(get_janos_epi_order)

    def func(x):
        "Remove space from resonce string"
        return ":".join(i.strip() for i in x.split(":"))

    jan_epi_resonances.reson = jan_epi_resonances.reson.map(func)

    # calculate name for axes display
    jan_epi_resonances["name"] = (
        jan_epi_resonances.moon + " " + jan_epi_resonances.reson
    )
    return jan_epi_resonances

# Cell
def get_prime_jan_epi():
    jan_epi_resonances = get_janus_epimetheus_resonances()
    # remove orders > 1 and drop unrequired columns
    prime_jan_epis = jan_epi_resonances[jan_epi_resonances.order == 1]
    to_drop = ["order", "moon"]
    prime_jan_epis = prime_jan_epis.drop(to_drop, axis=1)
    return prime_jan_epis

# Cell
def get_all_prime_resonances():
    prime_resonances = get_prime_resonances()
    prime_jan_epis = get_prime_jan_epi()
    all_resonances = pd.concat(
        [prime_resonances, prime_jan_epis], ignore_index=True, sort=False
    )
    all_resonances.sort_values(by="radius", inplace=True)
    all_resonances["moon"] = all_resonances.name.map(lambda x: x.split()[0].lower())
    return all_resonances

# Cell
def which_epi_janus_resonance(name, time):
    """Find which swap situtation we are in by time.

    Starting from 2006-01-21 where a Janus-Epimetheus swap occured, and
    defining the next 4 years until the next swap as `scenario1, and the 4
    years after that `scenario2`.
    Calculate in units of 4 years, in which scenario the given time falls.

    Parameters
    ----------
    time : timestring, datetime
        Time of the image. The astropy Time object can deal with both formats.

    Returns
    -------
    str
        The given name string (either `janus` or `epimetheus`) and attach
        a 1 or 2, as appropriate.
    """
    t1 = Time("2002-01-21").to_datetime()
    delta = Time(time).to_datetime() - t1
    yearfraction = delta.days / 365
    if int(yearfraction / 4) % 2 == 0:
        return name + "2"
    else:
        return name + "1"