"""
Copyright 2025 - David Minton.

This file is part of Swiftest.
Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License
as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with Swiftest.
If not, see: https://www.gnu.org/licenses.
"""

from __future__ import annotations

import datetime
import warnings
from typing import Any, Union

import astropy.units as u
import numpy as np
from astropy.coordinates import SkyCoord
from astroquery.jplhorizons import Horizons, HorizonsClass

import swiftest

from . import constants


def get_solar_system_body_mass_rotation(
    id: str, jpl: HorizonsClass = None, ephemerides_start_date: str = constants.MINTON_BCL, verbose: bool = True, **kwargs: Any
) -> dict:
    """
    Parses the raw output from JPL Horizons in order to extract physical properties of a body if they exist.

    Parameters
    ----------
    id : string or list of strings
        A string identifying which body is requested from JPL/Horizons (or a list of strings if multiple ids are possible, such as an altid list)
    jpl : HorizonsClass
        An astroquery.jplhorizons HorizonsClass object. If None, a new query will be made.
    ephemerides_start_date : string
        Date to use when obtaining the ephemerides in the format YYYY-MM-DD. Default is constants.MINTON_BCL
    verbose : bool
        Indicates whether to print messages about the query or not. Default is False
    **kwargs: Any
        Additional keyword arguments to pass to the query method (see https://astroquery.readthedocs.io/en/latest/jplhorizons/jplhorizons.html)

    Returns
    -------
    A dictionary containing the following elements

    Gmass : float
        G*Mass of the body in m^3/s^2
    radius : float
        The radius of the body in m
    rot : (3) float vector
        The rotation rate vector oriented toward the north pole in deg/s
    """

    def get_Gmass(raw_response):
        GM = [s for s in raw_response.split("\n") if "GM" in s and "Rel" not in s and "GMT" not in s and "ANGMOM" not in s]
        if len(GM) == 0:
            # Try an alternative name for the Mass found in some satellite queries
            M = [s for s in raw_response.split("\n") if "Mass" in s]
            if len(M) > 0:
                M = M[0].split("Mass")[-1].strip()
                if "kg" in M:
                    try:
                        unit_conv_str = M.split("kg")[0].strip()
                        unit_conv_str = unit_conv_str.split("^")[1].strip()
                        unit_conv = 10 ** int(unit_conv_str)
                        M = M.split("=")[1].strip()
                        if "^" in M:
                            mult = M.split("^")[1].split(")")[0].strip()
                            mult = 10 ** int(mult)
                        else:
                            mult = 1.0
                        M = M.split()[0].strip()
                        M = float(M) * mult * unit_conv
                        return M * swiftest.GC * 1e-9  # Return units of km**3 / s**2 for consistency
                    except Exception:
                        return None
            return None
        GM = GM[0]
        if len(GM) > 1:
            GM = GM.split("GM")[1]
            if len(GM) > 1:
                GM = GM.split("=")
                if len(GM) > 1:
                    GM = GM[1].strip().split(" ")[0].split("+")[0]
        try:
            GM = float(GM)
        except Exception:
            GM = None
        return GM

    def get_radius(raw_response):
        radius = [s for s in raw_response.split("\n") if "mean radius" in s.lower() or "RAD" in s or "Radius (km" in s]
        if len(radius) == 0:
            return None
        radius = radius[0]
        if "Radius (km" in radius or "RAD" in radius:
            if "Radius (km" in radius:
                radius = radius.split("Radius (km")[1].strip().strip("=").split("+")[0].strip()
            elif "RAD" in radius:
                radius = radius.split("RAD")[1].strip().strip("=").split("+")[0].strip().split()[0].strip()
            if "=" in radius:
                radius = radius.split("=")[1].strip()
            if "x" in radius:  # Triaxial ellipsoid bodies like Haumea may have multiple dimensions which need to be averaged
                radius = radius.split("x")
                try:
                    for i, v in enumerate(radius):
                        radius[i] = float(v.split()[0].strip())
                    radius = np.average(radius)
                except Exception:
                    radius = None
            else:
                radius = radius.split()[0].strip()
        elif "R_eff" in radius:  # Some small bodies list an effective radius
            radius = radius.split("R_eff")[1].strip(" =").split("+")[0].split()[0].strip()
        else:  # Handles most major bodies
            radius = radius.split("=")[1].strip().split("+")[0].split()[0].strip()
        try:
            radius = float(radius)
        except Exception:
            radius = None
        return radius

    def get_rotrate(raw_response):
        if "rot. rat" in raw_response.lower():  # or 'ROTPER' in raw_response.upper() or 'Orbital period' in raw_response:
            rotrate = [s for s in raw_response.split("\n") if "rot. rat" in s.lower()]
            rotrate = rotrate[0].lower().split("rot. rat")[1].split("=")[1].strip().split("  ")[0].strip()
            try:
                rotrate = float(rotrate)
            except Exception:
                rotrate = None
        elif "ROTPER" in raw_response.upper():
            rotrate = [s for s in raw_response.split("\n") if "ROTPER" in s.upper()]  # Try the small body version
            rotrate = rotrate[0].split("ROTPER=")[1].strip()
            try:
                rotrate = 2 * np.pi / (float(rotrate) * 3600)
            except Exception:
                rotrate = None
        elif "Synchronous" in raw_response:  # Satellites have this:
            rotrate = [s for s in raw_response.split("\n") if "Orbital period" in s][0]
            rotrate = rotrate.split("Orbital period")[1].replace("~", " ").replace("d", " ").replace("=", " ").strip()
            rotrate = 2 * np.pi / (float(rotrate) * swiftest.JD2S)
        elif "Orbital period" in raw_response:
            rotrate = [s for s in raw_response.split("\n") if "Orbital period" in s][0]
            rotrate = rotrate.split("Orbital period")[1].split("=")[1].strip().split()[0].strip()
            rotrate = 2 * np.pi / (float(rotrate) * swiftest.JD2S)
        else:
            rotrate = None
        return rotrate

    def get_rotpole(jpl):
        RA = jpl.ephemerides()["NPole_RA"][0]
        DEC = jpl.ephemerides()["NPole_DEC"][0]

        if np.ma.is_masked(RA) or np.ma.is_masked(DEC):
            return np.array([0.0, 0.0, 1.0])

        rotpole = SkyCoord(ra=RA * u.degree, dec=DEC * u.degree, frame="icrs").transform_to("barycentricmeanecliptic").cartesian

        return np.array([rotpole.x.value, rotpole.y.value, rotpole.z.value])

    if ephemerides_start_date is None:
        ephemerides_start_date = constants.MINTON_BCL

    if jpl is None:
        if not isinstance(id, list):
            id = [id]
        jpl, altid, namelist = horizons_query(id=id[0], ephemerides_start_date=ephemerides_start_date, verbose=verbose, **kwargs)
    else:
        if not isinstance(id, list):
            altid = [id]
        else:
            altid = id

    for i in altid:
        if jpl is None:
            jpl, _, namelist = horizons_query(id=i, ephemerides_start_date=ephemerides_start_date, verbose=verbose, **kwargs)
        namelist = [jpl.table["targetname"][0]]
        raw_response = jpl.vectors_async().text
        Rpl = get_radius(raw_response)
        if Rpl is not None:
            Rpl *= 1e3
            rotpole = get_rotpole(jpl)
            break
        else:
            jpl = None
    Gmass = get_Gmass(raw_response)
    if Rpl is None or Gmass is None:
        mass = None
    else:
        Gmass *= 1e9  # JPL returns GM in units of km**3 / s**2, so convert to SI
        mass = Gmass / swiftest.GC
    if Rpl is not None:
        rotrate = get_rotrate(raw_response)
        if rotrate is None:
            rotrate = 0.0
        else:
            rotrate = np.rad2deg(rotrate)

        rot = rotpole * rotrate
    else:
        rot = np.full(3, np.nan)
    if verbose and (Gmass is not None or Rpl is not None):
        print(f"Physical properties found for {namelist[0]}")

    return {"Gmass": Gmass, "mass": mass, "radius": Rpl, "rot": rot}


def horizons_query(
    id: str | int, ephemerides_start_date: str, exclude_spacecraft: bool = True, verbose: bool = False, **kwargs: Any
) -> HorizonsClass | None | (list | None) | (list | None):
    """
    Queries JPL/Horizons for a body matching the id.

    If one is found, a HorizonsClass object is returned for the first object that matches the passed id string. If more than one match is found, a list of alternate ids is also returned. If no object is found then None is returned.

    Parameters
    ----------
    id : string
        A string identifying which body is requested from JPL/Horizons
    ephemerides_start_date : string
        Date to use when obtaining the ephemerides in the format YYYY-MM-DD.
    exclude_spacecraft: bool (optional) - Default True
        Indicate whether spacecraft ids should be exluded from the alternate id list
    verbose: bool (optional) - Default False
        Indicate whether to print messages about the query or not

    Returns
    -------
    jpl : HorizonsClass | None
        An astroquery.jplhorizons HorizonsClass object. Or None if no match was found.
    altid : string list | None
        A list of alternate ids if more than one object matches the list
    """

    def get_altid(errstr, exclude_spacecraft=True):
        """
        Parses the error message returned from a failed Horizons query.

        If the query failed because of an ambiguous id, then it will return a list of id values that could possibly match the query.

        Parameters
        ----------
        errstr : string
            The error message returned from the Horizons query
        exclude_spacecraft: bool (optional) - Default True
            Indicate whether spacecraft ids should be exluded from the alternate id list

        Returns
        -------
        altid: string list | None
            A list of alternate ids if more than one object matches the list
        altname: string list | None
            A list of alternate names if more than one object matches the list

        """
        if "ID" in errstr:
            altid = errstr.split("ID")[1]
            altid = altid.split("\n")[2:-1]
            altname = []
            for n, v in enumerate(altid):
                altid[n] = v.strip().split(" ")[0]
                altname.append(" ".join(v.strip().split(" ")[1:]).strip().split("  ")[0])
            if exclude_spacecraft:
                altname = [altname[i] for i, n in enumerate(altid) if int(n) > 0]
                altid = [n for n in altid if int(n) > 0]

            return altid, altname
        else:
            return None, None

    # Horizons date time internal variables
    tstart = datetime.date.fromisoformat(ephemerides_start_date)
    tstep = datetime.timedelta(days=1)
    tend = tstart + tstep
    ephemerides_end_date = tend.isoformat()
    ephemerides_step = "1d"

    try:
        jpl = Horizons(
            id=id,
            location="@sun",
            epochs={"start": ephemerides_start_date, "stop": ephemerides_end_date, "step": ephemerides_step},
            **kwargs,
        )
        _ = jpl.ephemerides()
        altid = [id]
        altname = [jpl.table["targetname"][0]]
    except Exception as e:
        altid, altname = get_altid(str(e))
        if altid is not None and len(altid) > 0:  # Return the first matching id
            id = altid[0]
            jpl = Horizons(
                id=id,
                location="@sun",
                epochs={"start": ephemerides_start_date, "stop": ephemerides_end_date, "step": ephemerides_step},
            )
            _ = jpl.ephemerides()
        else:
            if verbose:
                warnings.warn(f"Could not find {id} in the JPL/Horizons system", stacklevel=2)
            return None, None, None
    if verbose:
        print(f"Found matching body: {altname[0]} ({altid[0]})")
        if len(altid) > 1:
            print("    Alternate matching bodies:")
            for i, v in enumerate(altid):
                if i > 0:
                    print(f"    {altname[i]} ({v})")

    return jpl, altid, altname


def get_solar_system_body(
    name: str,
    ephemeris_id: str | None = None,
    ephemerides_start_date: str = constants.MINTON_BCL,
    central_body_name: str = "Sun",
    verbose: bool = True,
    **kwargs: Any,
) -> dict | None:
    """
    Initializes a Swiftest dataset containing the major planets of the Solar System at a particular data from JPL/Horizons.

    Parameters
    ----------
    name  : string
        Name of body to add to Dataset. If `id` is not supplied, this is also what will be searched for in the JPL/Horizon's database.
        The first matching body is found (for major planets, this is the barycenter of a planet-satellite system)
    ephemeris_id : string (optional)
        If passed, this is passed to Horizons instead of `name`. This can be used to find a more precise body than given by `name`.
    ephemerides_start_date : string
        Date to use when obtaining the ephemerides in the format YYYY-MM-DD. Default is constants.MINTON_BCL
    central_body_name : string
        Name of the central body to use when calculating the relative position and velocity vectors. Default is "Sun"
    verbose : bool
        Indicates whether to print messages about the query or not. Default is True
    **kwargs: Any
        Additional keyword arguments to pass to the query method (see https://astroquery.readthedocs.io/en/latest/jplhorizons/jplhorizons.html)

    Returns
    -------
    A dictionary containing the following elements

    name : string
        Name of the body
    rh : (3,) array of np.float64
        Position vector array relative to the central body in m.
    vh : (3,) array of np.float64
        Velocity vector array relative to the central body in m/s.
    Gmass : np.float64
        G*mass values if these are massive bodies in m^3/s^2
    mass : np.float64
        Mass values if these are massive bodies in kg
    radius : np.float64
        Radius values if these are massive bodies in m
    rhill : np.float64
        The Hill's radius values if these are massive bodies in m
    Ip : (3,) array of np.float64
        Principal axes moments of inertia vectors if these are massive bodies.
    rot : (3,) array of np.float
        Rotation rate vectors if these are massive bodies in deg/s
    j2rp2 : np.float64
        J_2R^2 value for the body if known
    j4rp4 : np.float64
        J_4R^4 value for the body if known

    Notes
    -----
    When passing `name` == "Earth" or `name` == "Pluto", it a body is generated that has initial conditions matching the system
    barycenter and mass equal to the sum of Earth+Moon or Pluto+Charon. To obtain initial conditions for either Earth or Pluto alone,
    pass `ephemeris_id` == "399" for Earth or `id` == "999" for Pluto.
    """
    planetIpz = {  # Only the polar moments of inertia are used currently. Where the quantity is unkown, we just use the value of a sphere = 0.4
        "Sun": np.longdouble(0.070),
        "Mercury": np.longdouble(0.346),
        "Venus": np.longdouble(0.4),
        "Earth": np.longdouble(0.3307),
        "Mars": np.longdouble(0.3644),
        "Jupiter": np.longdouble(0.2756),
        "Saturn": np.longdouble(0.22),
        "Uranus": np.longdouble(0.23),
        "Neptune": np.longdouble(0.23),
        "Pluto": np.longdouble(0.4),
    }

    # J2 and J4 for the major planets are from From Murray & Dermott (1999) Table A.4.
    # The Sun is from Mecheri et al. (2004), using Corbard (b) 2002 values (Table II)
    planetJ2 = {
        "Sun": np.longdouble(2.198e-7),
        "Mercury": np.longdouble(60.0 * 1e-6),
        "Venus": np.longdouble(4.0 * 1e-6),
        "Earth": np.longdouble(1083.0 * 1e-6),
        "Mars": np.longdouble(1960.0 * 1e-6),
        "Jupiter": np.longdouble(14736.0 * 1e-6),
        "Saturn": np.longdouble(16298.0 * 1e-6),
        "Uranus": np.longdouble(3343.0 * 1e-6),
        "Neptune": np.longdouble(3411.0 * 1e-6),
    }
    planetJ4 = {
        "Sun": np.longdouble(-4.805e-9),
        "Mercury": np.longdouble(0.0),
        "Venus": np.longdouble(2.0 * 1e-6),
        "Earth": np.longdouble(-2.0 * 1e-6),
        "Mars": np.longdouble(-19.0 * 1e-6),
        "Jupiter": np.longdouble(-587.0 * 1e-6),
        "Saturn": np.longdouble(-915.0 * 1e-6),
        "Uranus": np.longdouble(-29.0 * 1e-6),
        "Neptune": np.longdouble(-35.0 * 1e-6),
    }

    # Unit conversion factors
    DCONV = swiftest.AU2M
    VCONV = swiftest.AU2M / swiftest.JD2S

    rh = np.full(3, np.nan)
    vh = np.full(3, np.nan)
    Ip = np.full(3, np.nan)
    rot = np.full(3, np.nan)
    rhill = None
    Gmass = None
    mass = None
    Rpl = None
    j2rp2 = None
    j4rp4 = None

    if name == "Sun" or ephemeris_id == "0":  # Create central body
        if verbose:
            print("Creating the Sun as a central body")
        # Central body value vectors
        rotpoleSun = SkyCoord(ra=286.13 * u.degree, dec=63.87 * u.degree).cartesian
        rot = (360.0 / 25.05) / constants.JD2S * rotpoleSun
        rot = np.array([rot.x.value, rot.y.value, rot.z.value])
        Gmass = swiftest.GMSun
        mass = Gmass / swiftest.GC
        Rpl = swiftest.RSun
        rh = np.array([0.0, 0.0, 0.0])
        vh = np.array([0.0, 0.0, 0.0])
    else:  # Fetch solar system ephemerides from Horizons
        if ephemeris_id is None:
            ephemeris_id = name

        if verbose:
            print(f"Fetching ephemerides data for {ephemeris_id} from JPL/Horizons")

        jpl, altid, altname = horizons_query(ephemeris_id, ephemerides_start_date, verbose=verbose, **kwargs)
        if jpl is not None:
            if verbose:
                print(f"Found ephemerides data for {altname[0]} ({altid[0]}) from JPL/Horizons")
            if name is None:
                name = altname[0]
        else:
            return None

        if central_body_name != "Sun":
            jplcb, altidcb, _ = horizons_query(central_body_name, ephemerides_start_date, verbose=verbose, **kwargs)
            GMcb = get_solar_system_body_mass_rotation(altidcb, jplcb)["Gmass"]
            cbrx = jplcb.vectors()["x"][0] * DCONV
            cbry = jplcb.vectors()["y"][0] * DCONV
            cbrz = jplcb.vectors()["z"][0] * DCONV
            cbvx = jplcb.vectors()["vx"][0] * VCONV
            cbvy = jplcb.vectors()["vy"][0] * VCONV
            cbvz = jplcb.vectors()["vz"][0] * VCONV
            cbrh = np.array([cbrx, cbry, cbrz])
            cbvh = np.array([cbvx, cbvy, cbvz])
        else:
            GMcb = swiftest.GMSun
            cbrh = np.zeros(3)
            cbvh = np.zeros(3)

        rx = jpl.vectors()["x"][0] * DCONV
        ry = jpl.vectors()["y"][0] * DCONV
        rz = jpl.vectors()["z"][0] * DCONV
        vx = jpl.vectors()["vx"][0] * VCONV
        vy = jpl.vectors()["vy"][0] * VCONV
        vz = jpl.vectors()["vz"][0] * VCONV

        rh = np.array([rx, ry, rz]) - cbrh
        vh = np.array([vx, vy, vz]) - cbvh

        Gmass, _, Rpl, rot = get_solar_system_body_mass_rotation(altid, jpl, verbose=verbose, **kwargs).values()
        # If the user inputs "Earth" or Pluto, then the Earth-Moon or Pluto-Charon barycenter and combined mass is used.
        # To use the Earth or Pluto alone, simply pass "399" or "999", respectively to name
        if name == "Earth":
            if verbose:
                print("Combining mass of Earth and the Moon")
            Gmass_moon = get_solar_system_body_mass_rotation(["301"], verbose=verbose, **kwargs)["Gmass"]
            Gmass += Gmass_moon
        elif name == "Pluto":
            if verbose:
                print("Combining mass of Pluto and Charon")
            Gmass_charon = get_solar_system_body_mass_rotation(["901"], verbose=verbose, **kwargs)["Gmass"]
            Gmass += Gmass_charon

        if Gmass is not None:
            rhill = jpl.elements()["a"][0] * DCONV * (Gmass / (3 * GMcb)) ** (1.0 / 3.0)
            mass = Gmass / swiftest.GC

    if name in planetIpz:
        Ip = np.array([0.0, 0.0, planetIpz[name]])
    else:
        Ip = np.array([0.4, 0.4, 0.4])

    if name in planetJ2:
        j2rp2 = planetJ2[name] * Rpl**2
        j4rp4 = planetJ4[name] * Rpl**4

    return {
        "name": name,
        "rh": rh,
        "vh": vh,
        "Gmass": Gmass,
        "mass": mass,
        "radius": Rpl,
        "rhill": rhill,
        "Ip": Ip,
        "rot": rot,
        "j2rp2": j2rp2,
        "j4rp4": j4rp4,
    }
