import warnings
import functools

import numpy as np

import astropy.units as u
from astropy.coordinates import SkyCoord

import sunpy.map

import sunkit_magex.pfss.coords

# Default colourmap for magnetic field maps
_MAG_CMAP = 'RdBu'


class Output:
    '''
    Output of PFSS modelling.

    Parameters
    ----------
    alr :
        Vector potential * grid spacing in radial direction.
    als :
        Vector potential * grid spacing in elevation direction.
    alp :
        Vector potential * grid spacing in azimuth direction.
    grid : Grid
        Grid that the output was calculated on.
    input_map : sunpy.map.GenericMap
        The input map.

    Notes
    -----
    Instances of this class are intended to be created by `sunkit_magex.pfss.pfss`, and
    not by users.
    '''
    def __init__(self, alr, als, alp, grid, input_map=None):
        self._alr = alr
        self._als = als
        self._alp = alp
        self.grid = grid
        self.input_map = input_map

        # Cache attributes
        self._common_b_cache = None
        self._rgi = None

    def _wcs_header(self):
        """
        Construct a world coordinate system describing the sunkit_magex.pfss solution.
        """
        return self.input_map.wcs

    @property
    def _lon0(self):
        """Longitude offset of the map."""
        return self.input_map.reference_coordinate.lon

    @property
    def coordinate_frame(self):
        """
        Coordinate frame of the input map.

        Notes
        -----
        This is either a `~sunpy.coordinates.frames.HeliographicCarrington` or
        `~sunpy.coordinates.frames.HeliographicStonyhurst` frame, depending on
        the input map.
        """
        return self.input_map.coordinate_frame

    @property
    def dtime(self):
        """
        Date and time of the input map.
        """
        return self.input_map.date

    @property
    def bunit(self):
        """
        `~astropy.units.Unit` of the input map data.

        If no metadata is available, returns dimensionless units.
        """
        unit = self.input_map.unit
        return u.dimensionless_unscaled if unit is None else unit

    @property
    def source_surface_br(self):
        """
        Radial magnetic field component on the source surface.

        Returns
        -------
        :class:`sunpy.map.GenericMap`
        """
        # Get radial component at the top
        br = self.bc[0][:, :, -1].to_value(self.bunit)
        # Remove extra ghost cells off the edge of the grid
        m = sunpy.map.Map((br.T, self._wcs_header()))
        if 'bunit' in self.input_map.meta:
            m.meta['bunit'] = self.input_map.meta['bunit']
        vlim = np.max(np.abs(br))
        m.plot_settings['cmap'] = _MAG_CMAP
        m.plot_settings['vmin'] = -vlim
        m.plot_settings['vmax'] = vlim
        return m

    @property
    def source_surface_pils(self):
        """
        Coordinates of the polarity inversion lines on the source surface.

        Notes
        -----
        This is always returned as a list of `~astropy.coordinates.SkyCoord`,
        as in general there may be more than one polarity inversion line.
        """
        from skimage import measure
        m = self.source_surface_br
        contours = measure.find_contours(m.data, 0)
        contours = [m.wcs.pixel_to_world(c[:, 1], c[:, 0]) for c in contours]
        return contours

    @property
    def _brgi(self):
        """
        Regular grid interpolator for B.
        """
        from sunkit_magex.pfss.interpolator import RegularGridInterpolator as rgi
        if self._rgi is not None:
            return self._rgi

        f32 = np.float32
        # - (rho,s,phi) coordinates:
        rho = self.grid.rg.astype(f32)
        s = self.grid.sg.astype(f32)
        phi = self.grid.pg.astype(f32)
        br, bth, bph = self.bg[..., 2], self.bg[..., 1], self.bg[..., 0]

        # Because we need the cartesian grid to stretch just beyond r=rss,
        # add an extra dummy layer of magnetic field pointing radially outwards
        rho = np.append(rho, rho[-1] + 0.01)
        extras = np.ones(br.shape[0:2] + (1, )) * self.bunit
        br = np.concatenate((br, extras), axis=2).astype(f32)
        bth = np.concatenate((bth, 0 * extras), axis=2).astype(f32)
        bph = np.concatenate((bph, 0 * extras), axis=2).astype(f32)

        # - convert to Cartesian components and make interpolator on
        # (rho,s,phi) grid:
        ph3, s3, rh3 = np.meshgrid(phi, s, rho, indexing='ij')
        sin_th = np.sqrt(1 - s3**2)
        cos_th = s3
        sin_ph = np.sin(ph3)
        cos_ph = np.cos(ph3)

        # Directly stack the expressions below, to save a bit of memory
        # bx = (sin_th * cos_ph * br) + (cos_th * cos_ph * bth) - (sin_ph * bph)
        # by = (sin_th * sin_ph * br) + (cos_th * sin_ph * bth) + (cos_ph * bph)
        # bz = (cos_th * br) - (sin_th * bth)
        bstack = np.stack(((sin_th * cos_ph * br) + (cos_th * cos_ph * bth) - (sin_ph * bph),
                           (sin_th * sin_ph * br) + (cos_th * sin_ph * bth) + (cos_ph * bph),
                           (cos_th * br) - (sin_th * bth)),
                          axis=-1)

        self._rgi = rgi((phi, s, rho), bstack)
        return self._rgi

    def _bTrace(self, t, coord, direction):
        """
        Return B/|B| for use by the field line tracer.
        """
        x, y, z = coord
        # (ph, s, rh) coordinates of current point:
        rho, s, phi = sunkit_magex.pfss.coords.cart2strum(x, y, z)

        # Check if position vector is outside the data limits
        if rho < 0 or rho > np.log(self.grid.rss):
            return np.array([0, 0, 0])
            # raise _OutOfBoundsError

        b1 = self._brgi(np.stack((phi, s, rho)))[0]
        return direction * b1 / np.linalg.norm(b1)

    def trace(self, tracer, seeds):
        """
        Trace magnetic field lines.

        Parameters
        ----------
        tracer : tracing.Tracer
            Field line tracer.
        seeds : astropy.coordinates.SkyCoord
            Starting coordinates.
        """
        return tracer.trace(seeds, self)

    def _integrate_one_way(self, dt, start_point, rtol, atol):
        import scipy.integrate

        direction = np.sign(dt)

        def finish_integration(t, coord):
            r = np.linalg.norm(coord)
            ret = (r - 1) * (r - self.grid.rss)
            return ret

        finish_integration.terminal = True
        # The integration domain is deliberately huge, because the
        # the integration automatically stops when an out of bounds error
        # is thrown
        t_span = (0, 1e4)

        def fun(t, y):
            return self._bTrace(t, y, direction)

        res = scipy.integrate.solve_ivp(
            fun, t_span, start_point, method='RK23',
            rtol=rtol, atol=atol, events=finish_integration)

        xout = res.y
        return xout

    @property
    def _al(self):
        """
        Vector potential times cell edge lengths.

        Returns ar*Lr, as*Ls, ap*Lp on cell edges.
        """
        return self._alr, self._als, self._alp

    @property
    def bc(self):
        """
        B on the centres of the cell faces.

        Returns
        -------
        br : astropy.units.Quantity
            A ``nphi, ns, nrho + 1`` shaped Quantity.
        btheta : astropy.units.Quantity
            A ``nphi, ns + 1, nrho`` shaped Quantity.
        bphi : astropy.units.Quantity
            A ``nphi + 1, ns, nrho`` shaped Quantity.

        Notes
        -----
        The three components are not co-located at the same locations.
        Component ``n`` is located on cell faces of constant ``n``, e.g.
        the ``r`` component is located on the cell faces at constant ``r``
        values.
        """
        br, bs, bp, Sbr, Sbs, Sbp = self._common_b()
        # Remove area factors:
        br = br.copy()
        bs = bs.copy()
        bp = bp.copy()
        for i in range(self.grid.nphi + 2):
            br[i, :, :] = br[i, :, :] / Sbr
            bs[i, :, :] = bs[i, :, :] / Sbs
        for i in range(self.grid.nphi + 1):
            bp[i, :, :] = bp[i, :, :] / Sbp

        # Slice to remove ghost cells
        return (br[1:-1, 1:-1, :] * self.bunit,
                -bs[1:-1, :, 1:-1] * self.bunit,
                bp[:, 1:-1, 1:-1] * self.bunit)

    @property
    @functools.lru_cache(maxsize=1)
    def bg(self):
        """
        B as a (weighted) averaged on grid points.

        Returns
        -------
        astropy.units.Quantity
            A ``(nphi + 1, ns + 1, nrho + 1, 3)`` shaped Quantity.
            The last index gives the corodinate axis, 0 for Bphi, 1 for Bs, 2
            for Brho. Because the phi dimension is periodic,
            ``bg[0, :, :] == bg[-1, :, :]``.
        """
        br, bs, bp, Sbr, Sbs, Sbp = self._common_b()
        # Weighted average to grid points:
        brg = br[:-1, :-1, :] + br[1:, :-1, :] + br[1:, 1:, :] + br[:-1, 1:, :]
        bsg = bs[:-1, :, :-1] + bs[1:, :, :-1] + bs[1:, :, 1:] + bs[:-1, :, 1:]
        bpg = bp[:, :-1, :-1] + bp[:, 1:, :-1] + bp[:, 1:, 1:] + bp[:, :-1, 1:]
        for i in range(self.grid.nphi + 1):
            brg[i, :, :] /= 2 * (Sbr[:-1, :] + Sbr[1:, :])
            bsg[i, :, :] /= 2 * (Sbs[:, :-1] + Sbs[:, 1:])
        for i in range(self.grid.nphi + 1):
            bpg[i, :, :] /= (Sbp[:-1, :-1] + Sbp[1:, :-1] +
                             Sbp[1:, 1:] + Sbp[:-1, 1:])
        bsg *= -1
        out = np.stack((bpg, bsg, brg), axis=-1)
        out.flags.writeable = False
        return out * self.bunit

    @property
    @functools.lru_cache(maxsize=1)
    def _modbg(self):
        return np.linalg.norm(self.bg, axis=-1)

    def _common_b(self):
        """
        Common code needed to calculate magnetic field from vector potential.
        """
        if self._common_b_cache is not None:
            return self._common_b_cache

        dr = self.grid.dr
        ds = self.grid.ds
        dp = self.grid.dp

        nr = self.grid.nr
        ns = self.grid.ns
        nphi = self.grid.nphi

        rss = self.grid.rss

        rc = self.grid.rc
        sc = self.grid.sc

        rg = self.grid.rg
        sg = self.grid.sg

        alr, als, alp = self._al

        # Centre of cells in rho (including ghost cells)
        rc = np.linspace(-0.5 * dr, np.log(rss) + 0.5 * dr, nr + 2)
        rrc = np.exp(rc)
        thc = np.zeros(ns + 2) - 1
        thc[1:-1] = np.arccos(sc)
        # Centre of cells in phi (including ghost cells)
        np.linspace(-0.5 * dp, 2 * np.pi + 0.5 * dp, nphi + 2)

        # Required face normals:
        dnp = np.zeros((ns + 2, 2))
        dns = np.zeros((ns + 1, 2))
        dnr = np.zeros(ns + 2)
        for k in range(2):
            for j in range(1, ns + 1):
                dnp[j, k] = rrc[k] * np.sqrt(1 - sc[j - 1]**2) * dp
            dnp[0, k] = dnp[1, k]
            dnp[-1, k] = dnp[-2, k]
            for j in range(1, ns):
                dns[j, k] = rrc[k] * (np.arcsin(sc[j]) - np.arcsin(sc[j - 1]))
            dns[0, k] = dns[1, k]
            dns[-1, k] = dns[-2, k]
        for j in range(ns + 2):
            dnr[j] = rrc[0] * (np.exp(dr) - 1)
        dnr[0] = -dnr[0]
        dnr[-1] = -dnr[-1]

        # Required area factors:
        Sbr = np.zeros((ns + 2, nr + 1))
        for k in range(nr + 1):
            Sbr[1:-1, k] = np.exp(2 * rg[k]) * ds * dp
            Sbr[0, k] = Sbr[1, k]
            Sbr[-1, k] = Sbr[-2, k]
        Sbs = np.zeros((ns + 1, nr + 2))
        for k in range(nr + 2):
            for j in range(1, ns):
                Sbs[j, k] = 0.5 * np.exp(2 * rc[k] - dr) * dp * (np.exp(2 * dr) - 1) * np.sqrt(1 - sg[j]**2)
            Sbs[0, k] = Sbs[1, k]
            Sbs[-1, k] = Sbs[-2, k]
        Sbp = np.zeros((ns + 2, nr + 2))
        for k in range(nr + 2):
            for j in range(1, ns + 1):
                Sbp[j, k] = 0.5 * np.exp(2 * rc[k] - dr) * (np.exp(2 * dr) - 1) * (np.arcsin(sg[j]) - np.arcsin(sg[j - 1]))
            Sbp[0, k] = Sbp[1, k]
            Sbp[-1, k] = Sbp[-2, k]

        # Compute br*Sbr, bs*Sbs, bp*Sbp at cell centres by Stokes theorem:
        br = np.zeros((nphi + 2, ns + 2, nr + 1))
        bs = np.zeros((nphi + 2, ns + 1, nr + 2))
        bp = np.zeros((nphi + 1, ns + 2, nr + 2))
        br[1:-1, 1:-1, :] = als[1:, :, :] - als[:-1, :, :] + alp[:, :-1, :] - alp[:, 1:, :]
        bs[1:-1, :, 1:-1] = alp[:, :, 1:] - alp[:, :, :-1]
        bp[:, 1:-1, 1:-1] = als[:, :, :-1] - als[:, :, 1:]

        # Fill ghost values with boundary conditions:
        # - zero-gradient at outer boundary:
        bs[1:-1, :, -1] = 2 * bs[1:-1, :, -2] - bs[1:-1, :, -3]
        bp[:, 1:-1, -1] = 2 * bp[:, 1:-1, -2] - bp[:, 1:-1, -3]
        # - periodic in phi:
        bs[0, :, :] = bs[-2, :, :]
        bs[-1, :, :] = bs[1, :, :]
        br[0, :, :] = br[-2, :, :]
        br[-1, :, :] = br[1, :, :]
        # js = jp = 0 at photosphere:
        for i in range(nphi + 1):
            bp[i, :, 0] = Sbp[:, 0] / dnp[:, 0] * (bp[i, :, 1] * dnp[:, 1] / Sbp[:, 1] + br[i, :, 0] * dnr[:] / Sbr[:, 0] - br[i + 1, :, 0] * dnr[:] / Sbr[:, 0])
        for i in range(nphi + 2):
            bs[i, :, 0] = Sbs[:, 0] / dns[:, 0] * (bs[i, :, 1] * dns[:, 1] / Sbs[:, 1] + br[i, :-1, 0] * dnr[:-1] / Sbr[:-1, 0] - br[i, 1:, 0] * dnr[1:] / Sbr[1:, 0])
        # - polar boundaries as in dumfric:
        for i in range(nphi + 2):
            i1 = (i + nphi // 2) % nphi
            br[i, -1, :] = br[i1, -2, :]
            br[i, 0, :] = br[i1, 1, :]
            bs[i, -1, :] = 0.5 * (bs[i, -2, :] - bs[i1, -2, :])
            bs[i, 0, :] = 0.5 * (bs[i, 1, :] - bs[i1, 1, :])
        for i in range(nphi + 1):
            i1 = (i + nphi // 2) % nphi
            bp[i, -1, :] = -bp[i1, -2, :]
            bp[i, 0, :] = -bp[i1, 1, :]

        self._common_b_cache = br, bs, bp, Sbr, Sbs, Sbp
        return self._common_b_cache

    def get_bvec(self, coords, out_type="spherical"):
        """
        Interpolate magnetic vectors at arbitrary coordinates.

        Parameters
        ----------
        coords : `astropy.coordinates.SkyCoord`
            An arbitrary point or set of points (length N >= 1) in the PFSS
            model domain (1Rs < r < Rss).

        out_type : str
            Takes values 'spherical' (default) or 'cartesian' and specifies
            whether the output vector is in spherical coordinates
            (B_r, B_theta, B_phi) or cartesian (B_x, B_y, B_z).

        Returns
        -------
        bvec : astropy.units.Quantity
            (N, 3) shaped Quantity of magnetic field vectors.

        Notes
        -----
        The output coordinate system is defined by the input
        magnetogram with x-z plane equivalent to the plane
        containing the Carrington meridian (0 deg longitude)

        The spherical coordinates follow the physics convention:
        https://upload.wikimedia.org/wikipedia/commons/thumb/4/4f/3D_Spherical.svg/240px-3D_Spherical.svg.png)
        Therefore the polar angle (theta) is the co-latitude, rather
        than the latitude, with range 0 (north pole) to 180 degrees
        (south pole)

        The conversion which relates the spherical and cartesian
        coordinates is as follows:

        .. math:: B_R = sin\\theta cos\\phi B_x + sin\\theta sin\\phi B_y + cos\\theta B_z
        .. math:: B_\\theta = cos\\theta cos\\phi B_x + cos\\theta sin\\phi B_y - sin\\theta B_z
        .. math:: B_\\phi = -sin\\phi B_x + cos\\phi B_y

        The above equations may be written as a (3x3) matrix and
        inverted to retrieve the inverse transformation (cartesian from spherical)
        """

        # Assert skycoord is type astropy.coordinates.SkyCoord
        if not isinstance(coords, SkyCoord):
            raise ValueError("coords must be of type "
                             "astropy.coordinates.SkyCoord")

        # Ensure representation type is spherical for input to interpolator
        coords.representation_type = "spherical"

        # Check coord_type is cartesian or spherical
        if out_type not in ["cartesian", "spherical"]:
            raise ValueError("out_type must be 'cartesian' or 'spherical' "
                             f"(got {out_type})")

        # Raise warning if input skycoord obstime does not match
        # self.coordinate_frame.obstime
        if np.any(self.coordinate_frame.obstime.to_datetime() !=
                  coords.obstime.to_datetime()):
            warnings.warn("The obstime of one of more input coordinates "
                          "do not match the pfss model obstime.")

        # Convert SkyCoord to sunkit_magex.pfss.Output coordinate frame
        coords.transform_to(self.coordinate_frame)

        # Do interpolation (returns cartesian vector)
        bvecs = self._brgi(np.array([coords.lon.to("rad").value,
                                     np.sin(coords.lat).value,
                                     np.log(coords.radius.to("R_sun").value)]).T
                           )

        # Convert to spherical if requested
        if out_type == "spherical":
            # Generate vector of 3x3 rotation matrices
            M = np.array([
                [np.cos(coords.lat).value * np.cos(coords.lon).value,
                 np.cos(coords.lat).value * np.sin(coords.lon).value,
                 np.sin(coords.lat).value],
                [np.sin(coords.lat).value * np.cos(coords.lon).value,
                 np.sin(coords.lat).value * np.sin(coords.lon).value,
                 -np.cos(coords.lat).value],
                [-np.sin(coords.lon).value,
                 np.cos(coords.lon).value,
                 np.zeros(len(coords))],
            ])
            bvecs = np.array([np.dot(M_.T, v) for M_, v in zip(M.T, bvecs)])
        return bvecs * self.bunit
