import warnings
from math import floor
from datetime import UTC, datetime

import astropy.units as u
import numpy as np
import remove_starfield
from astropy.io.fits import getheader
from astropy.nddata import StdDevUncertainty
from astropy.wcs import WCS
from dateutil.parser import parse as parse_datetime_str
from ndcube import NDCollection, NDCube
from prefect import get_run_logger
from remove_starfield import ImageHolder, ImageProcessor, Starfield
from remove_starfield.reducers import GaussianReducer
from solpolpy import resolve
from solpolpy.util import solnorth_from_wcs

from punchbowl.data import NormalizedMetadata, load_ndcube_from_fits, write_ndcube_to_fits
from punchbowl.data.wcs import (
    calculate_celestial_wcs_from_helio,
    calculate_helio_wcs_from_celestial,
    celestial_north_from_wcs,
)
from punchbowl.prefect import punch_flow, punch_task
from punchbowl.util import average_datetime

warnings.filterwarnings("ignore")


def polarize_solar_to_celestial(input_data: NDCube) -> NDCube:
    """
    Convert polarization from mzpsolar to Celestial frame.

    All images need their polarization converted to Celestial frame
    to generate the background starfield model.
    """
    # Create a data collection for M, Z, P components
    mzp_angles = [-60, 0, 60]*u.degree

    ncols, nrows = input_data.data[0].shape
    full_header = input_data.meta.to_fits_header(wcs=input_data.wcs,
                                                write_celestial_wcs=not False)
    wcs1 = WCS(full_header).dropaxis(2)
    wcs2 = WCS(full_header, key="A").dropaxis(2)

    # Converting polarization w.r.t. Celestial North
    angle_solar_north = solnorth_from_wcs(wcs1, (nrows, ncols))
    angle_celest_north = celestial_north_from_wcs(wcs2, (nrows, ncols))

    zoff = (angle_celest_north.value - angle_solar_north.value) * u.degree
    new_angles = np.stack([zoff - 60 * u.deg, zoff, zoff + 60 * u.deg])

    collection_contents = [
        (label,
         NDCube(data=input_data[i].data,
                wcs=wcs1,
                meta={"POLAR": angle}))
        for label, i, angle in zip(["M", "Z", "P"], [0, 1, 2], mzp_angles, strict=False)
    ]
    data_collection = NDCollection(collection_contents, aligned_axes="all")

    # Resolve data to celestial frame
    celestial_data_collection = resolve(data_collection, "npol", out_angles=new_angles)

    valid_keys = [key for key in celestial_data_collection if key != "alpha"]
    new_data = [celestial_data_collection[key].data for key in valid_keys]
    new_wcs = input_data.wcs.copy()

    output_meta = NormalizedMetadata.load_template("PTM", "3")
    output_meta["DATE-OBS"] = input_data.meta["DATE-OBS"].value

    output = NDCube(data=new_data, wcs=new_wcs, meta=output_meta)
    output.meta.history.add_now("LEVEL3-convert2celestial", "Convert mzpsolar to Celestial")

    return output


def polarize_celestial_to_solar(input_data: NDCube) -> NDCube:
    """
    Convert polarization from Celestial frame to mzpsolar.

    All images need their polarization converted back to Solar frame
    after removing the stellar polarization.
    """
    # Compute new angles for celestial frame
    ncols, nrows = input_data.data[0].shape
    full_header = input_data.meta.to_fits_header(wcs=input_data.wcs,
                                                write_celestial_wcs=not False)
    wcs1 = WCS(full_header).dropaxis(2)
    wcs2 = WCS(full_header, key="A").dropaxis(2)

    # Converting polarization w.r.t. Celestial North
    angle_solar_north = solnorth_from_wcs(wcs1, (nrows, ncols))
    angle_celest_north = celestial_north_from_wcs(wcs2, (nrows, ncols))

    zoff = (angle_celest_north.value - angle_solar_north.value) * u.degree
    new_angles = np.stack([zoff - 60 * u.deg, zoff, zoff + 60 * u.deg])

    collection_contents = [
        (f"{np.round(new_angles[i, nrows//2, ncols//2].value)} deg",
         NDCube(data=input_data[i].data,
                wcs=wcs1,
                meta={"POLAR": angle}))
        for i, angle in enumerate(new_angles)
    ]
    data_collection = NDCollection(collection_contents, aligned_axes="all")

    # Resolve data to mzpsolar frame
    solar_data_collection = resolve(data_collection, "mzpsolar", in_angles=new_angles)

    valid_keys = [key for key in solar_data_collection if key != "alpha"]
    new_data = [solar_data_collection[key].data for key in valid_keys]
    new_wcs = input_data.wcs.copy()

    output_meta = NormalizedMetadata.load_template("PTM", "3")
    output_meta["DATE-OBS"] = input_data.meta["DATE-OBS"].value

    output = NDCube(data=new_data, wcs=new_wcs, meta=output_meta, uncertainty=input_data.uncertainty)
    output.meta.history.add_now("LEVEL3-convert2mzpsolar", "Convert Celestial to mzpsolar")

    return output


class PUNCHImageProcessor(ImageProcessor):
    """Special loader for PUNCH data."""

    def __init__(self, layer: int | None = None, apply_mask: bool = True, key: str = " ") -> None:
        """Create PUNCHImageProcessor."""
        self.layer: int | None = layer
        self.apply_mask = apply_mask
        self.key = key

    def load_image(self, filename: str) -> ImageHolder:
        """Load an image."""
        cube = load_ndcube_from_fits(filename, key=self.key)

        if self.layer is None:  # clear data
            data = cube.data
        else:  # it's polarized
            cube = polarize_solar_to_celestial(cube)
            data = cube.data[self.layer]

        if self.apply_mask:
            data[data==0] = np.nan
        return ImageHolder(data, cube.wcs.celestial, cube.meta)


@punch_flow(log_prints=True, timeout_seconds=21_600)
def generate_starfield_background(
        filenames: list[str],
        map_scale: float = 0.01,
        target_mem_usage: float = 1000,
        n_procs: int | None = None,
        reference_time: datetime | None = None,
        is_polarized: bool = False,
        out_file: str | None = None) -> NDCube | None :
    """Create a background starfield map from a series of PUNCH images over a long period of time."""
    logger = get_run_logger()

    if reference_time is None:
        reference_time = datetime.now(UTC)
    elif isinstance(reference_time, str):
        reference_time = parse_datetime_str(reference_time)

    logger.info("construct_starfield_background started")

    # create an empty array to fill with data
    # open the first file in the list to ge the shape of the file
    if len(filenames) == 0:
        msg = "filenames cannot be empty"
        raise ValueError(msg)

    shape = [floor(132 / map_scale), floor(360 / map_scale)]
    starfield_wcs = WCS(naxis=2)
    # n.b. it seems the RA wrap point is chosen so there's 180 degrees
    # included on either side of crpix
    starfield_wcs.wcs.crpix = [shape[1] / 2 + .5, shape[0] / 2 + .5]
    starfield_wcs.wcs.crval = 270, -23.5
    starfield_wcs.wcs.cdelt = map_scale, map_scale
    starfield_wcs.wcs.ctype = "RA---CAR", "DEC--CAR"
    starfield_wcs.wcs.cunit = "deg", "deg"
    starfield_wcs.array_shape = shape

    date_obses = [getheader(f, 1)["DATE-OBS"] for f in filenames]
    times = [datetime.fromisoformat(d) for d in date_obses]

    meta = NormalizedMetadata.load_template("PSM" if is_polarized else "CSM", "3")
    meta["DATE-OBS"] = reference_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
    meta["DATE-BEG"] = min(times).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
    meta["DATE-END"] = max(times).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
    meta["DATE-AVG"] = average_datetime(times).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
    meta["DATE"] = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]

    if is_polarized:
        logger.info("Starting m starfield")
        starfield_m = remove_starfield.build_starfield_estimate(
            filenames,
            attribution=False,
            frame_count=False,
            reducer=GaussianReducer(),
            starfield_wcs=starfield_wcs,
            n_procs=n_procs,
            processor=PUNCHImageProcessor(0, apply_mask=True, key="A"),
            target_mem_usage=target_mem_usage)
        logger.info("Ending m starfield")

        logger.info("Starting z starfield")
        starfield_z = remove_starfield.build_starfield_estimate(
            filenames,
            attribution=False,
            frame_count=False,
            reducer=GaussianReducer(),
            starfield_wcs=starfield_wcs,
            n_procs=n_procs,
            processor=PUNCHImageProcessor(1, apply_mask=True, key="A"),
            target_mem_usage=target_mem_usage)
        logger.info("Ending z starfield")

        logger.info("Starting p starfield")
        starfield_p = remove_starfield.build_starfield_estimate(
            filenames,
            attribution=False,
            frame_count=False,
            reducer=GaussianReducer(),
            starfield_wcs=starfield_wcs,
            n_procs=n_procs,
            processor=PUNCHImageProcessor(2, apply_mask=True, key="A"),
            target_mem_usage=target_mem_usage)
        logger.info("Ending p starfield")

        out_data = np.stack([starfield_m.starfield, starfield_z.starfield, starfield_p.starfield], axis=0)
        out_wcs = calculate_helio_wcs_from_celestial(starfield_m.wcs, meta.astropy_time, starfield_m.starfield.shape)
    else:
        logger.info("Starting clear starfield")
        starfield_clear = remove_starfield.build_starfield_estimate(
            filenames,
            attribution=False,
            frame_count=False,
            reducer=GaussianReducer(),
            starfield_wcs=starfield_wcs,
            n_procs=n_procs,
            processor=PUNCHImageProcessor(None, apply_mask=True, key="A"),
            target_mem_usage=target_mem_usage)
        logger.info("Ending clear starfield")
        out_data = starfield_clear.starfield
        out_wcs = calculate_helio_wcs_from_celestial(starfield_clear.wcs,
                                                        meta.astropy_time,
                                                        starfield_clear.starfield.shape)

    # TODO - Replace uncertainty below with values folded through starfield estimation logic
    output = NDCube(data=out_data, uncertainty=StdDevUncertainty(np.sqrt(out_data)), wcs=out_wcs, meta=meta)
    output.meta.history.add_now("LEVEL3-starfield_background", "constructed starfield_bg model")

    logger.info("construct_starfield_background finished")

    if out_file is not None:
        write_ndcube_to_fits(output, filename=out_file, write_hash=False, overwrite=True)
        return None

    return [output]


@punch_task
def subtract_starfield_background_task(data_object: NDCube,
                                       starfield_background_path: str | None,
                                       is_polarized: bool = False) -> NDCube:
    """
    Subtracts a background starfield from an input data frame.

    checks the dimensions of input data frame and background starfield match and
    subtracts the background starfield from the data frame of interest.

    Parameters
    ----------
    data_object : NDCube
        A NDCube data frame to be background subtracted
    starfield_background_path : str
        path to a NDCube background starfield map
    is_polarized : bool
        whether the data is polarized

    Returns
    -------
    NDCube
        A background starfield subtracted data frame

    """
    logger = get_run_logger()
    logger.info("subtract_starfield_background started")

    if starfield_background_path is None:
        output = data_object
        output.meta.history.add_now("LEVEL3-subtract_starfield_background",
                                           "starfield subtraction skipped since path is empty")
    else:
        star_datacube = load_ndcube_from_fits(starfield_background_path)
        wcs_celestial = calculate_celestial_wcs_from_helio(star_datacube.wcs)
        wcs_celestial.wcs.cdelt[0] = wcs_celestial.wcs.cdelt[0] * -1

        original_mask = data_object.data == 0

        if is_polarized:
            starfield_model_m = Starfield(np.stack((star_datacube.data[0], star_datacube.uncertainty.array[0])),
                                          wcs_celestial[0])
            subtracted_m = starfield_model_m.subtract_from_image(
                NDCube(data=np.stack((data_object.data[0], data_object.uncertainty.array[0])),
                       wcs=data_object.wcs[0],
                       meta=data_object.meta),
                       processor=PUNCHImageProcessor(layer=0, key="A"))
            starfield_model_z = Starfield(np.stack((star_datacube.data[1], star_datacube.uncertainty.array[1])),
                                          wcs_celestial[1])
            subtracted_z = starfield_model_z.subtract_from_image(
                NDCube(data=np.stack((data_object.data[1], data_object.uncertainty.array[1])),
                       wcs=data_object.wcs[1],
                       meta=data_object.meta),
                       processor=PUNCHImageProcessor(layer=1, key="A"))
            starfield_model_p = Starfield(np.stack((star_datacube.data[2], star_datacube.uncertainty.array[2])),
                                          wcs_celestial[2])
            subtracted_p = starfield_model_p.subtract_from_image(
                NDCube(data=np.stack((data_object.data[2], data_object.uncertainty.array[2])),
                       wcs=data_object.wcs[2],
                       meta=data_object.meta),
                       processor=PUNCHImageProcessor(layer=2, key="A"))

            data_object.data[...] = np.stack([subtracted_m.subtracted[0],
                                              subtracted_z.subtracted[0],
                                              subtracted_p.subtracted[0]], axis=0)
            data_object.uncertainty.array[...] = np.sqrt(data_object.uncertainty.array**2 +
                                                         np.stack([subtracted_m.subtracted[1],
                                                                   subtracted_z.subtracted[1],
                                                                   subtracted_p.subtracted[1]], axis=0)**2)
        else:
            starfield_model = Starfield(np.stack((star_datacube.data, star_datacube.uncertainty.array)), wcs_celestial)
            subtracted = starfield_model.subtract_from_image(
                NDCube(data=np.stack((data_object.data, data_object.uncertainty.array)),
                       wcs=data_object.wcs,
                       meta=data_object.meta),
                       processor=PUNCHImageProcessor(key="A"))

            data_object.data[...] = subtracted.subtracted[0]
            data_object.uncertainty.array[...] = np.sqrt(data_object.uncertainty.array**2 +
                                                         subtracted.subtracted[1]**2)

        # Reset the data to be zero in invalid regions
        data_object.data[original_mask] = 0
        data_object.data[~np.isfinite(data_object.data)] = 0

        data_object.meta.history.add_now("LEVEL3-subtract_starfield_background", "subtracted starfield background")
        output = polarize_celestial_to_solar(data_object) if is_polarized else data_object
    logger.info("subtract_starfield_background finished")

    return output


def create_empty_starfield_background(data_object: NDCube) -> np.ndarray:
    """Create an empty starfield background map."""
    return np.zeros_like(data_object.data)
