"""
Abstract parent module for all other instruments
Contains some general functionality, which may be overridden by the children of course
"""

import datetime
import glob
import json
import logging
import os.path
from itertools import product

import numpy as np
import yaml
from astropy.io import fits
from astropy.time import Time
from dateutil import parser
from tqdm import tqdm

from ..clipnflip import clipnflip
from .filters import ChannelFilter, Filter, InstrumentFilter, NightFilter, ObjectFilter
from .models import InstrumentConfig

logger = logging.getLogger(__name__)


def find_first_index(arr, value):
    """find the first element equal to value in the array arr"""
    try:
        return next(i for i, v in enumerate(arr) if v == value)
    except StopIteration as e:
        raise KeyError(f"Value {value} not found") from e


def observation_date_to_night(observation_date):
    """Convert an observation timestamp into the date of the observation night
    Nights start at 12am and end at 12 am the next day

    Parameters
    ----------
    observation_date : datetime
        timestamp of the observation

    Returns
    -------
    night : datetime.date
        night of the observation
    """
    if observation_date == "":
        return None

    observation_date = parser.parse(observation_date)
    oneday = datetime.timedelta(days=1)

    if observation_date.hour < 12:
        observation_date -= oneday
    return observation_date.date()


class getter:
    """Get data from a header/dict, based on the given channel, and applies replacements"""

    def __init__(self, header, info, channel):
        self.header = header
        self.info = info.copy()
        try:
            self.index = find_first_index(info["channels"], channel.upper())
        except KeyError:
            logger.warning("No instrument channels found in instrument info")
            self.index = 0

        # Pick values for the given channel
        for k, v in self.info.items():
            if isinstance(v, list):
                self.info[k] = v[self.index]

    def __call__(self, key, alt=None):
        return self.get(key, alt)

    def get(self, key, alt=None):
        """Get data

        Parameters
        ----------
        key : str
            key of the data in the header
        alt : obj, optional
            alternative value, if key does not exist (default: None)

        Returns
        -------
        value : obj
            value found in header (or alternatively alt)
        """

        value = self.info.get(key, key)
        # if isinstance(value, list):
        #     value = value[self.index]
        if isinstance(value, str):
            value = value.format(**self.info)
            value = self.header.get(value, alt)
        return value


class Instrument:
    """
    Abstract parent class for all instruments
    Handles the instrument specific information
    """

    def __init__(self):
        #:str: Name of the instrument (lowercase)
        self.name = self.__class__.__name__.lower()
        #:InstrumentConfig: Validated configuration model
        #:dict: Information about the instrument (for backward compatibility)
        self.config, self.info = self.load_info()

        self.filters = {
            "instrument": InstrumentFilter(self.config.instrument, regex=True),
            "night": NightFilter(self.config.date, timeformat=self.config.date_format),
            "target": ObjectFilter(self.config.target, regex=True),
            "bias": Filter(self.config.kw_bias),
            "flat": Filter(self.config.kw_flat),
            "orders": Filter(self.config.kw_orders),
            "curvature": Filter(self.config.kw_curvature),
            "scatter": Filter(self.config.kw_scatter),
            "wave": Filter(self.config.kw_wave),
            "comb": Filter(self.config.kw_comb),
            "spec": Filter(self.config.kw_spec),
        }

        self.night = "night"
        self.science = "science"
        self.shared = ["instrument", "night"]

        # Add channel filter if kw_channel is defined (for instruments with separate files per channel)
        if self.config.kw_channel is not None:
            self.filters["channel"] = ChannelFilter(self.config.kw_channel)
            self.shared.append("channel")
        self.find_closest = [
            "bias",
            "flat",
            "wavecal_master",
            "freq_comb_master",
            "orders",
            "scatter",
            "curvature",
        ]

    def __str__(self):
        return self.name

    @property
    def channels(self) -> list[str] | None:
        """Available instrument channels."""
        return self.config.channels

    @property
    def extension(self) -> int | str | list:
        """FITS extension(s) to read."""
        return self.config.extension

    @property
    def orientation(self) -> int | list[int]:
        """Detector orientation code(s)."""
        return self.config.orientation

    @property
    def id_instrument(self) -> str:
        """Instrument identifier for header matching."""
        return self.config.id_instrument

    def get(self, key, header, channel, alt=None):
        get = getter(header, self.info, channel)
        return get(key, alt=alt)

    def get_extension(self, header, channel):
        channel = channel.upper()
        ext = self.extension  # Use property

        if isinstance(ext, list):
            ichannel = find_first_index(self.channels, channel)
            ext = ext[ichannel]

        return ext

    def load_info(self):
        """
        Load static instrument information
        Either as fits header keywords or static values

        Returns
        ------
        config : InstrumentConfig
            Validated Pydantic model
        info : dict(str:object)
            dictionary of REDUCE names for properties to Header keywords/static values
        """
        # Tips & Tricks:
        # if several channels are supported, use a list for channels
        # if a value changes depending on the channel, use a list with the same order as "channels"
        # you can also use values from this dictionary as placeholders using {name}, just like str.format

        this = os.path.dirname(__file__)

        # Try YAML first, fall back to JSON
        yaml_fname = os.path.join(this, f"{self.name}.yaml")
        json_fname = os.path.join(this, f"{self.name}.json")

        if os.path.exists(yaml_fname):
            with open(yaml_fname) as f:
                info = yaml.safe_load(f)
        elif os.path.exists(json_fname):
            with open(json_fname) as f:
                info = json.load(f)
        else:
            raise FileNotFoundError(
                f"No instrument config found for {self.name} "
                f"(tried {yaml_fname} and {json_fname})"
            )

        # Validate with Pydantic (strict - invalid config is a bug)
        config = InstrumentConfig(**info)

        return config, info

    def load_fits(
        self, fname, channel, extension=None, mask=None, header_only=False, dtype=None
    ):
        """
        load fits file, REDUCE style

        primary and extension header are combined
        channel-specific info is applied to header
        data is clipnflipped
        mask is applied

        Parameters
        ----------
        fname : str
            filename
        instrument : str
            name of the instrument
        channel : str
            instrument channel
        extension : int
            data extension of the FITS file to load
        mask : array, optional
            mask to add to the data
        header_only : bool, optional
            only load the header, not the data
        dtype : str, optional
            numpy datatype to convert the read data to

        Returns
        --------
        data : masked_array
            FITS data, clipped and flipped, and with mask
        header : fits.header
            FITS header (Primary and Extension + channel info)

        ONLY the header is returned if header_only is True
        """

        channel = channel.upper()

        hdu = fits.open(fname)
        h_prime = hdu[0].header
        if extension is None:
            extension = self.get_extension(h_prime, channel)

        header = hdu[extension].header
        if extension != 0:
            header.extend(h_prime, strip=False)
        header = self.add_header_info(header, channel)
        header["e_input"] = (os.path.basename(fname), "Original input filename")

        if header_only:
            hdu.close()
            return header

        data = clipnflip(hdu[extension].data, header)

        if dtype is not None:
            data = data.astype(dtype)

        data = np.ma.masked_array(data, mask=mask)

        hdu.close()
        return data, header

    def add_header_info(self, header, channel, **kwargs):
        """read data from header and add it as REDUCE keyword back to the header

        Parameters
        ----------
        header : fits.header, dict
            header to read/write info from/to
        channel : str
            instrument channel

        Returns
        -------
        header : fits.header, dict
            header with added information
        """

        info = self.info
        get = getter(header, info, channel)

        # Use HIERARCH prefix only for FITS Header objects to avoid warnings
        # For dict objects, HIERARCH is not needed and would break key access
        from astropy.io.fits import Header as FitsHeader

        hierarch = "HIERARCH " if isinstance(header, FitsHeader) else ""

        header[f"{hierarch}e_instrument"] = get("instrument", self.__class__.__name__)
        header[f"{hierarch}e_telescope"] = get("telescope", "")
        header[f"{hierarch}e_exptime"] = get("exposure_time", 0)

        jd = get("date")
        if jd is not None:
            jd = Time(jd, format=self.info.get("date_format", "fits"))
            jd = jd.to_value("mjd")

        header["e_orient"] = get("orientation", 0)
        # As per IDL rotate if orient is 4 or larger and transpose is undefined
        # the image is transposed
        header[f"{hierarch}e_transpose"] = get(
            "transpose", (header["e_orient"] % 8 >= 4)
        )

        naxis_x = get("naxis_x", 0)
        naxis_y = get("naxis_y", 0)

        prescan_x = get("prescan_x", 0)
        overscan_x = get("overscan_x", 0)
        prescan_y = get("prescan_y", 0)
        overscan_y = get("overscan_y", 0)

        header["e_xlo"] = prescan_x
        header["e_xhi"] = naxis_x - overscan_x

        header["e_ylo"] = prescan_y
        header["e_yhi"] = naxis_y - overscan_y

        header["e_gain"] = get("gain", 1)
        header["e_readn"] = get("readnoise", 0)

        header["e_sky"] = get("sky", 0)
        header["e_drk"] = get("dark", 0)
        header["e_backg"] = header["e_gain"] * (header["e_drk"] + header["e_sky"])

        header["e_imtype"] = get("image_type")
        header["e_ctg"] = get("category")

        header["e_ra"] = get("ra", 0)
        header["e_dec"] = get("dec", 0)
        header["e_jd"] = jd

        header["e_obslon"] = get("longitude")
        header["e_obslat"] = get("latitude")
        header["e_obsalt"] = get("altitude")

        if info.get("wavecal_element", None) is not None:
            header["HIERARCH e_wavecal_element"] = get(
                "wavecal_element", info.get("wavecal_element", None)
            )
        return header

    def find_files(self, input_dir):
        """Find fits files in the given folder

        Parameters
        ----------
        input_dir : string
            directory to look for fits and fits.gz files in, may include bash style wildcards

        Returns
        -------
        files: array(string)
            absolute path filenames
        """
        files = glob.glob(input_dir + "/*.fits")
        files += glob.glob(input_dir + "/*.fits.gz")
        files = np.array(files)
        return files

    def get_expected_values(self, target, night, channel=None, **kwargs):
        expectations = {
            "bias": {
                "instrument": self.config.id_instrument,
                "night": night,
                "bias": self.config.id_bias,
            },
            "flat": {
                "instrument": self.config.id_instrument,
                "night": night,
                "flat": self.config.id_flat,
            },
            "orders": {
                "instrument": self.config.id_instrument,
                "night": night,
                "orders": self.config.id_orders,
            },
            "scatter": {
                "instrument": self.config.id_instrument,
                "night": night,
                "scatter": self.config.id_scatter,
            },
            "curvature": {
                "instrument": self.config.id_instrument,
                "night": night,
                "curvature": self.config.id_curvature,
            },
            "wavecal_master": {
                "instrument": self.config.id_instrument,
                "night": night,
                "wave": self.config.id_wave,
            },
            "freq_comb_master": {
                "instrument": self.config.id_instrument,
                "night": night,
                "comb": self.config.id_comb,
            },
            "science": {
                "instrument": self.config.id_instrument,
                "night": night,
                "target": target,
                "spec": self.config.id_spec,
            },
        }

        # Add channel filter if this instrument has separate files per channel
        if channel is not None and self.config.kw_channel is not None:
            id_channel = self.config.id_channel
            channels = self.config.channels
            channel_id = (
                id_channel[channels.index(channel)] if channel in channels else channel
            )
            for key in expectations:
                expectations[key]["channel"] = channel_id

        return expectations

    def populate_filters(self, files):
        """Extract values from the fits headers and store them in self.filters

        Parameters
        ----------
        files : list(str)
            list of fits files

        Returns
        -------
        filters: list(Filter)
            list of populated filters (identical to self.filters)
        """
        # Empty filters
        for _, fil in self.filters.items():
            fil.clear()

        for f in tqdm(files):
            with fits.open(f) as hdu:
                h = hdu[0].header
                for _, fil in self.filters.items():
                    fil.collect(h)

        return self.filters

    def apply_filters(self, files, expected, allow_calibration_only=False):
        """
        Determine the relevant files for a given set of expected values.

        Parameters
        ----------
        files : list(files)
            list if fits files
        expected : dict
            dictionary with expected header values for each reduction step

        Returns
        -------
        files: list((dict, dict))
            list of files. The first element of each tuple is the used setting,
            and the second are the files for each step.
        """

        # Fill the filters with header information
        self.populate_filters(files)

        # Use the header information determined in populate filters
        # to find potential science and calibration files in the list of files
        # result = {step : [ {setting : value}, [files] ] }
        result = {}
        for step, values in expected.items():
            result[step] = []
            data = {}
            for name, value in values.items():
                if isinstance(value, list):
                    for v in value:
                        data[name] = self.filters[name].classify(v)
                        if len(data[name]) > 0:
                            break
                else:
                    data[name] = self.filters[name].classify(value)
            # Get all combinations of possible filter values
            # e.g. if several nights are allowed
            for thingy in product(*data.values()):
                mask = np.copy(thingy[0][1])
                for i in range(1, len(thingy)):
                    mask &= thingy[i][1]
                if np.count_nonzero(mask) == 0:
                    continue
                d = {k: v[0] for k, v in zip(values.keys(), thingy, strict=False)}
                f = files[mask]
                result[step].append((d, f))

        # Filter for only nights that have a science observation
        # files = [{setting: value}, {step: files}]
        files = []
        if allow_calibration_only:
            # Use all unique nights
            settings = {}
            for shared in self.shared:
                keys = [k for k in set(self.filters[shared].data) if k is not None]
                settings[shared] = keys
        else:
            # Or use only science nights
            settings = {}
            for shared in self.shared:
                keys = [key[shared] for key, _ in result[self.science]]
                settings[shared] = keys

        values = [settings[k] for k in self.shared]
        for setting in product(*values):
            setting = dict(zip(self.shared, setting, strict=False))
            night = setting[self.night]
            f = {}
            # For each step look for files with matching settings
            for step, step_data in result.items():
                f[step] = []
                for step_key, step_files in step_data:
                    match = [
                        setting[shared] == step_key[shared]
                        for shared in self.shared
                        if shared in step_key.keys()
                    ]
                    if all(match):
                        f[step] = step_files
                        break
                # If no matching files are found ...
                if len(f[step]) == 0:
                    if step not in self.find_closest:
                        # Show a warning
                        logger.warning(
                            "Could not find any files for step '%s' with settings %s, sharing parameters %s",
                            step,
                            setting,
                            self.shared,
                        )
                    else:
                        # Or find the closest night instead
                        j = None
                        for i, (step_key, _) in enumerate(step_data):
                            match = [
                                setting[shared] == step_key[shared]
                                for shared in self.shared
                                if shared in step_key.keys() and shared != self.night
                            ]
                            if all(match):
                                if j is None:
                                    j = i
                                else:
                                    diff_old = abs(step_data[j][0][self.night] - night)
                                    diff_new = abs(step_data[i][0][self.night] - night)
                                    if diff_new < diff_old:
                                        j = i
                        if j is None:
                            # We still dont find any files
                            logger.warning(
                                "Could not find any files for step '%s' in any night with settings %s, sharing parameters %s",
                                step,
                                setting,
                                self.shared,
                            )
                        else:
                            # We found files in a close night
                            closest_key, closest_files = step_data[j]
                            logger.warning(
                                "Using '%s' files from night %s for observations of night %s",
                                step,
                                night,
                                closest_key["night"],
                            )
                            f[step] = closest_files

            if any(len(a) > 0 for a in f.values()):
                files.append((setting, f))
        if len(files) == 0:
            logger.warning(
                "No %s files found matching the expected values %s",
                self.science,
                expected[self.science],
            )
        return files

    def sort_files(
        self, input_dir, target, night, *args, allow_calibration_only=False, **kwargs
    ):
        """
        Sort a set of fits files into different categories
        types are: bias, flat, wavecal, orderdef, spec

        Parameters
        ----------
        input_dir : str
            input directory containing the files to sort
        target : str
            name of the target as in the fits headers
        night : str
            observation night, possibly with wildcards
        channel : str
            instrument channel
        Returns
        -------
        files_per_night : list[dict{str:dict{str:list[str]}}]
            a list of file sets, one entry per night, where each night consists of a dictionary with one entry per setting,
            each fileset has five lists of filenames: "bias", "flat", "order", "wave", "spec", organised in another dict
        nights_out : list[datetime]
            a list of observation times, same order as files_per_night
        """
        input_dir = input_dir.format(
            **kwargs, target=target, night=night, instrument=self.name
        )
        files = self.find_files(input_dir)
        ev = self.get_expected_values(target, night, *args, **kwargs)
        files = self.apply_filters(
            files, ev, allow_calibration_only=allow_calibration_only
        )
        return files

    def get_wavecal_filename(self, header, channel, **kwargs):
        """Get the filename of the pre-existing wavelength solution for the current setting

        Parameters
        ----------
        header : fits.header, dict
            header of the wavelength calibration file
        channel : str
            instrument channel

        Returns
        -------
        filename : str
            name of the wavelength solution file
        """

        specifier = header.get(self.config.wavecal_specifier or "", "")
        instrument = "wavecal"

        cwd = os.path.dirname(__file__)
        fname = f"{instrument.lower()}_{channel}_{specifier}.npz"
        fname = os.path.join(cwd, "..", "wavecal", fname)
        return fname

    def get_supported_channels(self):
        return self.channels

    def get_mask_filename(self, channel, **kwargs):
        i = self.name.lower()
        c = channel.lower()
        fname = f"mask_{i}_{c}.fits.gz"
        cwd = os.path.dirname(__file__)
        fname = os.path.join(cwd, "..", "masks", fname)
        return fname

    def get_wavelength_range(self, header, channel, **kwargs):
        return self.get("wavelength_range", header, channel)


class COMMON(Instrument):
    pass


def create_custom_instrument(
    name, extension=0, info=None, mask_file=None, wavecal_file=None
):
    class CUSTOM(Instrument):
        def __init__(self):
            super().__init__()
            self.name = name

        def load_info(self):
            if info is None:
                return None, COMMON().info
            try:
                with open(info) as f:
                    data = json.load(f)
                return None, data
            except:
                return None, info

        def get_extension(self, header, channel):
            return extension

        def get_mask_filename(self, channel, **kwargs):
            return mask_file

        def get_wavecal_filename(self, header, channel, **kwargs):
            return wavecal_file

    return CUSTOM()
