import pathlib
from typing import TypeAlias

import pandas as pd
import pyarrow as pa
from astropy.io import fits

from swift_comet_pipeline.swift.swift_data import SwiftData
from swift_comet_pipeline.swift.swift_data import (
    swift_observation_id_from_int,
    swift_orbit_id_from_obsid,
)
from swift_comet_pipeline.swift.swift_datamodes import (
    datamode_from_fits_keyword_string,
    float_to_pixel_resolution,
)
from swift_comet_pipeline.swift.swift_filter_to_string import (
    filter_to_obs_string,
    obs_string_to_filter,
)
from swift_comet_pipeline.types.swift_filter import SwiftFilter
from swift_comet_pipeline.types.swift_uvot_image import SwiftUVOTImage


SwiftObservationLog: TypeAlias = pd.DataFrame


def observation_log_schema() -> pa.Schema:
    """
    Returns schema that describes the columns in our observation log dataframe

    Most FITS header keyword documentation can be found at https://archive.stsci.edu/swiftuvot/UVOT_swguide_v2_2.pdf
    """
    schema = pa.schema(
        [
            ### FITS keywords from SWIFT-processed FITS files
            # Observation ID: uniquely identifies an SWIFT observation set, which may contain multiple FITS image extensions.  Represented as a fixed-width string of 11 characters,
            # with the first 8 digits being the 'target id' and the last 3 being 'segment number': '000', '001', ... to denote multiple observations taken on the target consecutively.
            # We store this in the db as an integer, and convert back to SwiftObservationID string representation when we read the file
            pa.field("OBS_ID", pa.int64()),
            # Start time of the observation, as a string with the format 'YYYY-MM-DD HH:MM:SS'
            pa.field("DATE_OBS", pa.string()),
            # End time of the observation, as a string with the format 'YYYY-MM-DD HH:MM:SS'
            pa.field("DATE_END", pa.string()),
            # Mid time of the observation, as a string with the format 'YYYY-MM-DD HH:MM:SS'
            pa.field("MID_TIME", pa.string()),
            # Which SwiftFilter was used for this observation
            pa.field("FILTER", pa.string()),
            # TODO: we probably don't need PA_PNT
            # Position angle, degrees of roll
            pa.field("PA_PNT", pa.float64()),
            # Right ascension, degrees
            pa.field("RA_OBJ", pa.float64()),
            # Declination, degrees
            pa.field("DEC_OBJ", pa.float64()),
            # Total exposure time after all known corrections applied
            pa.field("EXPOSURE", pa.float64()),
            ### Added during observation log building
            # Each image extension gets its own header and entry in the observation log: keep track of it here
            pa.field("EXTENSION", pa.int16()),
            # The filename (not path) of the FITS file this observation came from: together with the extension, we can find an observation's image to read and process later
            pa.field("FITS_FILENAME", pa.string()),
            # heliocentric distance at time of observation, in AU
            pa.field("HELIO", pa.float64()),
            # heliocentric velocity at time of observation, in km/s
            pa.field("HELIO_V", pa.float64()),
            # distance to comet, AU
            pa.field("OBS_DIS", pa.float64()),
            # phase angle, degrees (Sun-Target-Object angle)
            pa.field("PHASE", pa.float64()),
            # coordinates given by Horizons for the center of the comet
            pa.field("RA", pa.float64()),
            pa.field("DEC", pa.float64()),
            # PX, PY: pixel x and y coordinates of comet center after converting RA, DEC using WCS
            pa.field("PX", pa.float64()),
            pa.field("PY", pa.float64()),
            # x and y coordinates of comet center as selected by user during veto step
            pa.field("USER_CENTER_X", pa.float64()),
            pa.field("USER_CENTER_Y", pa.float64()),
            # string as taken from the FITS header keyword DATAMODE
            pa.field("DATAMODE", pa.string()),
            # how many arcseconds per pixel in this data mode
            pa.field("ARCSECS_PER_PIXEL", pa.float64()),
            # conversion from pixels to kilometers, based on the distance of the comet from the observation point
            pa.field("KM_PER_PIX", pa.float64()),
            # String holding the SWIFT software that processed the image, which includes the version.  Currently unused - this is here just in case.
            pa.field("CREATOR", pa.string()),
            # Holds whether or not the user has chosen to veto the image's inclusion in stacking: manual_veto = true means exclude
            pa.field("manual_veto", pa.bool_()),
        ]
    )

    return schema


def read_observation_log(
    obs_log_path: pathlib.Path, additional_schema: pa.Schema = None
) -> SwiftObservationLog:
    """
    Reads an observation log generated by build_observation_log and converts column data types where needed
    """
    schema = observation_log_schema()
    if additional_schema is not None:
        schema = pa.unify_schemas([schema, additional_schema])

    obs_log = pd.read_parquet(obs_log_path, schema=schema)

    # post-processing on the columns to get them into a datatype we want: we have to store these as strings
    # but datetime objects are much nicer to work with
    obs_log[["DATE_OBS", "DATE_END", "MID_TIME"]] = obs_log[
        ["DATE_OBS", "DATE_END", "MID_TIME"]
    ].apply(pd.to_datetime)

    # filters are stored as strings as well, convert them to SwiftFilter objects
    obs_log.FILTER = obs_log["FILTER"].astype(str).map(obs_string_to_filter)

    # observation ID is stored as an int: convert to SwiftObservationID and SwiftOrbitID types
    obs_log.OBS_ID = obs_log["OBS_ID"].apply(swift_observation_id_from_int)
    obs_log["ORBIT_ID"] = obs_log["OBS_ID"].apply(swift_orbit_id_from_obsid)

    # DATAMODE is stored as a string: convert to valid SwiftImageDataMode
    obs_log.DATAMODE = obs_log.DATAMODE.apply(datamode_from_fits_keyword_string)

    # ARCSECS_PER_PIXEL is stored as a float: convert to valid SwiftPixelResolution
    obs_log.ARCSECS_PER_PIXEL = obs_log.ARCSECS_PER_PIXEL.apply(
        float_to_pixel_resolution
    )

    return obs_log


def write_observation_log(
    obs_log: SwiftObservationLog,
    obs_log_path: pathlib.Path,
    # additional_schema: pa.lib.Schema = None,
) -> None:
    """
    Copy obs_log and process the columns to data types that fit our schema, then save
    """
    oc = obs_log.copy()

    oc[["DATE_OBS", "DATE_END", "MID_TIME"]] = oc[
        ["DATE_OBS", "DATE_END", "MID_TIME"]
    ].astype(str)

    oc["OBS_ID"] = oc["OBS_ID"].astype(int)
    oc["FILTER"] = oc["FILTER"].map(filter_to_obs_string)

    oc.DATAMODE = oc.DATAMODE.astype(str)
    oc.ARCSECS_PER_PIXEL = oc.ARCSECS_PER_PIXEL.astype(float)

    # the rest of the column conversions should be taken care of automatically through type hints from the schema

    schema = observation_log_schema()
    # if additional_schema is not None:
    #     schema = pa.unify_schemas([schema, additional_schema])

    oc.to_parquet(obs_log_path, schema=schema)


def includes_uvv_and_uw1_filters(
    obs_log: SwiftObservationLog,
) -> bool:
    has_uvv_filter = obs_log[obs_log["FILTER"] == SwiftFilter.uvv]
    has_uvv_set = set(has_uvv_filter["ORBIT_ID"])

    has_uw1_filter = obs_log[obs_log["FILTER"] == SwiftFilter.uw1]
    has_uw1_set = set(has_uw1_filter["ORBIT_ID"])

    has_both = len(has_uvv_set) > 0 and len(has_uw1_set) > 0

    return has_both


def get_image_path_from_obs_log_row(swift_data: SwiftData, obs_log_row) -> pathlib.Path:
    image_path = (
        swift_data.get_uvot_image_directory(obsid=obs_log_row.OBS_ID)
        / obs_log_row.FITS_FILENAME
    )

    return image_path


def get_image_from_obs_log_row(swift_data: SwiftData, obs_log_row) -> SwiftUVOTImage:
    image_path = get_image_path_from_obs_log_row(
        swift_data=swift_data, obs_log_row=obs_log_row
    )
    image_data = fits.getdata(image_path, ext=obs_log_row.EXTENSION)

    return image_data  # type: ignore


def get_header_from_obs_log_row(swift_data: SwiftData, obs_log_row):
    image_path = (
        swift_data.get_uvot_image_directory(obsid=obs_log_row.OBS_ID)
        / obs_log_row.FITS_FILENAME
    )
    header = fits.getheader(image_path)

    return header


# For saving metadata with a parquet file
# schema = pa.schema(
#     [],
#     metadata={
#         "coincidence_correction": str(do_coincidence_correction),
#         "pixel_resolution": str(epoch_pixel_resolution),
#     },
# )


# def print_parquet_metadata():
#
#     m = pa.parquet.read_metadata(parquet_path)
#     print(m.metadata[b'metadata_key'])
