"""
This file tests the behavior of subclasses of Representation and Frames
"""
from copy import deepcopy

import pytest

from astropy.coordinates import Longitude, Latitude
from astropy.coordinates.representation import (REPRESENTATION_CLASSES,
                                                SphericalRepresentation,
                                                UnitSphericalRepresentation,
                                                _invalidate_reprdiff_cls_hash)
from astropy.coordinates.baseframe import frame_transform_graph
from astropy.coordinates.transformations import FunctionTransform
from astropy.coordinates import ICRS
from astropy.coordinates.baseframe import RepresentationMapping

import astropy.units as u

import astropy.coordinates

# Classes setup, borrowed from SunPy.

# Here we define the classes *inside* the tests to make sure that we can wipe
# the slate clean when the tests have finished running.


def setup_function(func):
    func.REPRESENTATION_CLASSES_ORIG = deepcopy(REPRESENTATION_CLASSES)


def teardown_function(func):
    REPRESENTATION_CLASSES.clear()
    REPRESENTATION_CLASSES.update(func.REPRESENTATION_CLASSES_ORIG)
    _invalidate_reprdiff_cls_hash()


def test_unit_representation_subclass():

    class Longitude180(Longitude):
        def __new__(cls, angle, unit=None, wrap_angle=180*u.deg, **kwargs):
            self = super().__new__(cls, angle, unit=unit, wrap_angle=wrap_angle,
                                   **kwargs)
            return self

    class UnitSphericalWrap180Representation(UnitSphericalRepresentation):
        attr_classes = {'lon': Longitude180,
                        'lat': Latitude}

    class SphericalWrap180Representation(SphericalRepresentation):
        attr_classes = {'lon': Longitude180,
                        'lat': Latitude,
                        'distance': u.Quantity}

        _unit_representation = UnitSphericalWrap180Representation

    class MyFrame(ICRS):
        default_representation = SphericalWrap180Representation
        frame_specific_representation_info = {
            'spherical': [
                RepresentationMapping('lon', 'ra'),
                RepresentationMapping('lat', 'dec')]
        }
        frame_specific_representation_info['unitsphericalwrap180'] = \
            frame_specific_representation_info['sphericalwrap180'] = \
            frame_specific_representation_info['spherical']

    @frame_transform_graph.transform(FunctionTransform,
                                     MyFrame, astropy.coordinates.ICRS)
    def myframe_to_icrs(myframe_coo, icrs):
        return icrs.realize_frame(myframe_coo._data)

    f = MyFrame(10*u.deg, 10*u.deg)
    assert isinstance(f._data, UnitSphericalWrap180Representation)
    assert isinstance(f.ra, Longitude180)

    g = f.transform_to(astropy.coordinates.ICRS())
    assert isinstance(g, astropy.coordinates.ICRS)
    assert isinstance(g._data, UnitSphericalWrap180Representation)

    frame_transform_graph.remove_transform(MyFrame,
                                           astropy.coordinates.ICRS,
                                           None)