import logging
import os
from datetime import timedelta
from os.path import basename

import astropy.units as u
import numpy as np

from kis_tools.gris.GrisFitsFile import GrisFitsFile
from kis_tools.gris.coord_correction_ml import get_coords_ml
from kis_tools.gris.util import telescope_to_hpc
from kis_tools.util.constants import gris_uncertainty_no_center_x, gris_uncertainty_no_center_y
from kis_tools.util.util import date_from_fn, gris_run_number

logger = logging.getLogger(__name__)


class IFUFitsFile(GrisFitsFile):

    def parse_filename(self):
        """Extract metadata from filename"""

        # Extract time from Filename
        name = os.path.splitext(basename(self.filename))[0].split("_")
        self.date = date_from_fn(self.filename)

        # Extract Run-, Map-, and Stepnumber, as well as mode string
        self.runnumber = gris_run_number(self.filename)
        self.mapnumber = int(name[5])
        self.modestring = name[3]

    @property
    def n_steps(self):
        """
        Number of steps per map

        Does not account for aborted maps. Contrary to slit operation,
        IFU maps are typically not too large and the error caused by the
        actual difference isn't as severe.
        Returns:
            n_steps: number of steps per map
        """
        nsteps = self["STEPSH"] * self["STEPSV"]
        return nsteps

    @property
    def _projected_observation_end(self):
        remaining_slits = self.n_steps * self.number_of_maps - (
                self.query(["STEP", "ISTEP"]) + (self.mapnumber - 1) * self.n_steps)
        # remaining_time = remaining_slits * self.exposure_time
        # using placeholder for time between slits, I don't know how to calculate correctly,
        # self.exptime does not account for dead time between exposures.
        remaining_time = remaining_slits * 3
        projected_end = self.obs_time + timedelta(seconds=remaining_time)
        return projected_end

    @property
    def _coords_from_simple_header(self):
        """coordinates for slit, the Keywords 'SLITPOSX' and 'SLITPOSY' are assumed to be the center of the slit
         while the keyword 'SLITORIE' describes the rotation of the slit w.r.t. the helioprojective axes. The
         algorithm further assumes that each pixel represents a square area with side length of 'STEPSIZE'.
         These assumptions are used to calculate the coordinates for the center of each pixel within the data array.

        Returns:
            (X,Y):ndarray(float) x and y coordinates for each pixel of the slit


        WARNING! The code below is not checked for correctness! We do not know what direction SLITORIE references
        off and coordinates provided by GREGOR are very error-prone. If you need precision better than ~50" this
        function is not sufficient.
        """

        x_elements = self["NAXIS1"]
        y_elements = self["NAXIS2"]

        # Slit center in telescope coordinates
        x_slitcenter = self['SLITPOSX'] if self.header['SLITPOSX'] != "" else 0
        y_slitcenter = self['SLITPOSY'] if self.header['SLITPOSY'] != "" else 0

        # Stepsizes
        x_stepsize = self['STEPSIZH']
        y_stepsize = self['STEPSIZV']

        coord_ref_x = x_elements // 2
        coord_ref_y = y_elements // 2

        if "SLITCNTR" in self.header:
            if self['SLITCNTR'].strip() != 'scancntr':
                raise ValueError(
                    f'Cannot determine reference pixel for coordinates! SLITCNTR is {self["SLITCNTR"]} instead of scancntr')

        # Get telescope centering and associated uncertainties. Centering is the driving factor for uncertainties
        delta_x, delta_y, std_delta_x, std_delta_y = self.telescope_centering

        # If uncertainties are too large, use ML model to correct the coordinates
        try:
            p0_angle = self.p0_angle
        except ValueError:
            logger.warning(ValueError)
            p0_angle = None
        if p0_angle is not None:
            from kis_tools.gris.coord_correction_ml import coord_model_stdx, coord_model_stdy
            if std_delta_x > coord_model_stdx or std_delta_y > coord_model_stdy:
                try:
                    x_hpc, y_hpc, std_delta_x, std_delta_y = get_coords_ml(self)
                except ValueError as e:
                    logger.warning(f"Error in finding ML coords for {self.__repr__()}:{e}")
                    # Point to slitpos with 1 solar disk uncertainty if coord retrieval didn't work
                    x_hpc, y_hpc, std_delta_x, std_delta_y = x_slitcenter, y_slitcenter, \
                                                             gris_uncertainty_no_center_x, gris_uncertainty_no_center_y
            else:
                x_hpc, y_hpc = telescope_to_hpc(x_slitcenter, y_slitcenter, p0_angle, delta_x, delta_y)
        else:
            # if p0 is broken, only do centering correction
            x_hpc, y_hpc = x_slitcenter - delta_x, y_slitcenter - delta_y

        # Try determining orientation, use 0 if not possible
        try:
            angle = self.slit_orientation
        except ValueError as e:
            logger.warning(f"Could not determine Slit orientation! \n {e}")
            angle = 0

        # With the center pixel value given, the coordinates need to be expanded to a full array.

        # Setup grid of pixels
        delta_pix_x = [i - coord_ref_x for i in range(x_elements)]
        delta_pix_y = [i - coord_ref_y for i in range(y_elements)]
        delta_pix_x = np.array(delta_pix_x)
        delta_pix_y = np.array(delta_pix_y)

        # generate coordinates for full grid of slit pixels
        X = x_hpc + np.cos(np.radians(angle)) * delta_pix_x * x_stepsize
        Y = y_hpc - np.sin(np.radians(angle)) * delta_pix_y * y_stepsize

        X = (X * np.ones((y_elements, 1))).T
        Y = Y * np.ones((x_elements, 1))

        coord_array = np.array([X, Y]) * u.arcsec

        # track coordinate uncertainty within class
        self.coord_uncertainty = (std_delta_x, std_delta_y)
        return coord_array

    def ___init__(self, filename):
        super(IFUFitsFile, self).___init__(filename)