import logging
import os
import time as ttime
from collections.abc import Iterable

import arrow
import dask.array as da
import numpy as np
import scipy as sp

from ..calibration import Calibration
from ..constants import MARIA_MAX_NU_HZ, MARIA_MIN_NU_HZ
from ..errors import FrequencyOutOfBoundsError, ShapeError
from ..units import Quantity, parse_units

logger = logging.getLogger("maria")

here, this_filename = os.path.split(__file__)

SLICE_DIMS = {
    "stokes": {
        "dtype": str,
        "default": "I",
    },
    "nu": {
        "dtype": float,
        "default": 150e9,
    },
    "t": {
        "dtype": float,
        "default": arrow.now().timestamp(),
    },
}


class Map:
    """
    The base class for maps. Maps have data
    """

    def __init__(
        self,
        data: float,
        weight: float,
        stokes: str,
        nu: Iterable[float],
        t: Iterable[float],
        z: Iterable[float],
        beam: tuple[float, float, float],  # noqa
        map_dims: dict,
        units: str = "K_RJ",
        degrees: bool = True,  # noqa
        dtype: type = np.float32,
    ):
        self.data = da.asarray(data).astype(dtype)
        self.weight = (da.asarray(weight) if weight is not None else da.ones_like(self.data)).astype(dtype)
        self.data, self.weight = np.broadcast_arrays(self.data, self.weight)

        self.units = parse_units(units)["units"]

        if not hasattr(beam, "__len__"):
            beam = (beam, beam, 0)
        self.beam = Quantity(beam, "deg" if degrees else "rad")

        if np.shape(self.beam)[-1] != 3:
            raise ValueError("'beam' must be either a number or a tuple of (major, minor, angle)")

        self.map_dims = map_dims

        self.slice_dims = {}
        if stokes is not None:
            stokes = "".join(list(stokes))
            if stokes not in ["I", "IQU", "IQUV"]:
                raise ValueError(f"Invalid stokes parameter '{stokes}' (must be either 'I', 'IQU', or 'IQUV').")
            self.stokes = stokes.upper()
            self.slice_dims["stokes"] = len(self.stokes)
        else:
            self.stokes = ""

        if nu is not None:
            self.nu = Quantity(nu, "Hz")
            if self.nu.ndim == 0:
                self.nu = Quantity([self.nu.Hz], "Hz")
            if self.nu.ndim > 1:
                raise ShapeError("'nu' can be at most one-dimensional")
            bad_freqs = list(self.nu[(self.nu.Hz < MARIA_MIN_NU_HZ) | (self.nu.Hz > MARIA_MAX_NU_HZ)])
            if bad_freqs:
                raise FrequencyOutOfBoundsError(bad_freqs)
            self.slice_dims["nu"] = len(self.nu)
        else:
            self.nu = Quantity([], "Hz")

        if t is not None:
            self.t = np.atleast_1d(t)
            self.slice_dims["t"] = len(self.t)
        else:
            self.t = np.array([])

        if z is not None:
            self.z = np.atleast_1d(z)
            self.slice_dims["z"] = len(self.z)
        else:
            self.z = np.array([])

        self.data *= np.ones(list(self.dims.values()))
        self.weight *= np.ones(list(self.dims.values()))

        # for i, (dim, n) in enumerate(slice_dims.items()):
        #     if self.data.shape[i] != n:
        #         raise ValueError(
        #             f"'{dim}' has length {n} but map has {dim} length {self.data.shape[i]}.",
        #         )

        implied_shape = tuple(list(self.dims.values()))
        # if len(implied_shape) != self.data.ndim:
        #     raise ValueError(f"Inputs imply rank {len(self.dims)} for map data, but data has rank {self.data.ndim}")

        if implied_shape != self.data.shape:
            raise ValueError(
                f"Inputs imply shape {self.dims_string}={implied_shape} for map data, but data has shape {self.data.shape}"
            )

        if self.u["quantity"] == "spectral_flux_density_per_beam":
            if not np.all(self.beam_area > 0):
                raise ValueError(
                    f"Map is given in units {self.units}, but specified beam(major, minor, angle) = {beam} has zero area"
                )

    @property
    def dims(self):
        return {**self.slice_dims, **self.map_dims}

    @property
    def ndim(self):
        return len(self.dims)

    @property
    def shape(self):
        return tuple(self.dims.values())

    @property
    def dims_string(self):
        return f"({', '.join(self.dims.keys())})"

    @property
    def dims_list(self):
        return list(self.dims.keys())

    @property
    def nu_bins(self):
        return np.array([0, *(self.nu.Hz[1:] + self.nu.Hz[:-1]), 1e6])

    @property
    def nu_side(self):
        return (self.nu_bins[:-1] + self.nu_bins[1:]) / 2

    @property
    def t_bins(self):
        """
        This might break in the year 3000.
        """
        t_min = arrow.get("0001-01-01").timestamp()
        t_max = arrow.get("3000-01-01").timestamp()

        return np.array([t_min, *(self.t[1:] + self.t[:-1]), t_max])

    @property
    def t_side(self):
        return (self.t_bins[:-1] + self.t_bins[1:]) / 2

    @property
    def u(self):
        return parse_units(self.units)

    # @property
    # def weight(self):
    #     return self._weight if self._weight is not None else da.ones(shape=self.data.shape)

    def append(self, map, dim):
        return concatenate([self, map], dim=dim)

    def extend(self, maps, dim):
        return concatenate([self, *maps], dim=dim)

    def squeeze(self, dims=None):
        dims = np.atleast_1d(dims) if dims is not None else [dim for dim, n in self.slice_dims.items() if n == 1]
        package = self.package()
        dim_indices_to_squeeze = []

        for dim_index, name in enumerate(dims):
            n = self.slice_dims.get(name)
            if n is None:
                raise ValueError(f"{self.__class__.__name__} has no dimension '{name}'")
            if n != 1:
                raise ValueError(f"Cannot squeeze dimension '{name}' with length {n} > 1")

            dim_indices_to_squeeze.append(dim_index)
            package.pop(name)

        package["data"] = package["data"].squeeze(tuple(dim_indices_to_squeeze))
        package["weight"] = package["weight"].squeeze(tuple(dim_indices_to_squeeze))

        return type(self)(**package)

    def unsqueeze(self, dim, value=None):
        if dim not in SLICE_DIMS:
            raise ValueError(f"'{dim}' is not a valid map dimension")
        if dim in self.dims:
            raise Exception(f"{self.__class__.__name__} already has dimension '{dim}'")

        new_dim_index = 0
        for d in ["stokes", "nu", "t"]:
            if d == dim:
                break
            if d in self.dims:
                new_dim_index += 1

        package = self.package()
        package["data"] = np.expand_dims(package["data"], new_dim_index)
        package["weight"] = np.expand_dims(package["weight"], new_dim_index)
        if value is None:
            value = SLICE_DIMS[dim]["default"]
        package[dim] = value

        return type(self)(**package)

    @property
    def pixel_area(self):
        if hasattr(self, "resolution"):
            return Quantity(self.resolution.rad**2, "sr")
        else:
            return Quantity(self.x_res.rad * self.y_res.rad, "sr")

    @property
    def beam_area(self):
        """
        Returns the beam area in steradians
        """
        area = (np.pi / 4) * self.beam[..., 0].radians * self.beam[..., 1].radians
        return Quantity(area * np.ones(tuple(self.slice_dims.values())), "sr")

    def to(self, units: str):
        if units == self.units:
            return self

        u = parse_units(units)

        package = self.package().copy()

        # this is just a scaling by some factor
        if u["quantity"] == self.u["quantity"]:
            package["data"] *= self.u["factor"] / u["factor"]

        else:
            if "nu" not in self.dims:
                raise ValueError(
                    f"Cannot convert from quantity {self.u['quantity']} to quantity {u['quantity']} "
                    f"when map has no frequency."
                )

            # data = package["data"].swapaxes(0, self.dims_list.index("nu"))  # put the nu index in front
            for nu_index, nu in enumerate(self.nu.Hz):
                nu_key = tuple([nu_index if dim == "nu" else slice(None, None, None) for dim in self.slice_dims])

                cal = Calibration(
                    f"{self.units} -> {units}",
                    nu=nu,
                    pixel_area=self.pixel_area.sr,
                    beam_area=self.beam_area[nu_key].sr,
                )
                package["data"][nu_key] = cal(package["data"][nu_key])

            # package["data"] = data.swapaxes(0, self.dims_list.index("nu"))  # swap the axes back

        package["units"] = units
        return type(self)(**package)

    def sample_nu(self, nu):
        map_nu_interpolator = sp.interpolate.interp1d(self.nu.Hz, self.data, axis=1, kind="linear")

        nu_maps = []
        for nu in np.atleast_1d(nu):
            if nu < self.nu.Hz[0]:
                nu_maps.append(self.data[:, 0])
            elif not (nu < self.nu.Hz[-1]):  # this will include nan
                nu_maps.append(self.data[:, -1])
            else:
                nu_maps.append(map_nu_interpolator(nu))

        return np.stack(nu_maps, axis=1)

    @property
    def nu_bin_bounds(self):
        nu_boundaries = [0, *(self.nu.Hz[:-1] + self.nu.Hz[1:]) / 2, np.inf]
        return [(Quantity(nu1, "Hz"), Quantity(nu2, "Hz")) for nu1, nu2 in zip(nu_boundaries[:-1], nu_boundaries[1:])]

    def copy(self):
        return type(self)(**self.package().copy())


def concatenate(maps: list[Map], dim: str) -> Map:
    packages = []
    concat_dim_values = []
    dims_list = {}
    for m in maps:
        for d, n in m.dims.items():
            if d not in dims_list:
                dims_list[d] = []
            dims_list[d].append(n)
        for attr in ["center", "width", "height"]:
            # TODO: checks for healpix maps
            if getattr(m, attr, None) != getattr(maps[0], attr, None):
                print("warning")
        concat_dim_values.extend(getattr(m, dim))
        packages.append(m.to(maps[0].units).package())

    for d, ns in dims_list.items():
        if d != dim:
            if len(set(ns)) > 1:
                raise ShapeError(
                    "Map dimensions must be equal except along the concatenation axis "
                    f"(maps have shapes {[m.shape for m in maps]})."
                )

    dim_index = list(dims_list.keys()).index(dim)
    total_package = packages[0].copy()
    total_package["data"] = np.concatenate([p["data"] for p in packages], axis=dim_index)
    total_package["weight"] = np.concatenate([p["weight"] for p in packages], axis=dim_index)
    total_package[dim] = concat_dim_values

    return type(maps[0])(**total_package)
