# Licensed under a 3-clause BSD style license - see LICENSE.rst
from numpy.testing import assert_array_equal

from astropy import modeling
from astropy.io.misc.asdf.tags.transform.basic import TransformType

from . import _parameter_to_value

__all__ = ["AffineType", "Rotate2DType", "Rotate3DType", "RotationSequenceType"]


class AffineType(TransformType):
    name = "transform/affine"
    version = "1.3.0"
    types = ["astropy.modeling.projections.AffineTransformation2D"]

    @classmethod
    def from_tree_transform(cls, node, ctx):
        matrix = node["matrix"]
        translation = node["translation"]
        if matrix.shape != (2, 2):
            raise NotImplementedError(
                "asdf currently only supports 2x2 (2D) rotation transformation matrices"
            )
        if translation.shape != (2,):
            raise NotImplementedError(
                "asdf currently only supports 2D translation transformations."
            )

        return modeling.projections.AffineTransformation2D(
            matrix=matrix, translation=translation
        )

    @classmethod
    def to_tree_transform(cls, model, ctx):
        return {
            "matrix": _parameter_to_value(model.matrix),
            "translation": _parameter_to_value(model.translation),
        }

    @classmethod
    def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert a.__class__ == b.__class__
        assert_array_equal(a.matrix, b.matrix)
        assert_array_equal(a.translation, b.translation)


class Rotate2DType(TransformType):
    name = "transform/rotate2d"
    version = "1.3.0"
    types = ["astropy.modeling.rotations.Rotation2D"]

    @classmethod
    def from_tree_transform(cls, node, ctx):
        return modeling.rotations.Rotation2D(node["angle"])

    @classmethod
    def to_tree_transform(cls, model, ctx):
        return {"angle": _parameter_to_value(model.angle)}

    @classmethod
    def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert isinstance(a, modeling.rotations.Rotation2D) and isinstance(
            b, modeling.rotations.Rotation2D
        )
        assert_array_equal(a.angle, b.angle)


class Rotate3DType(TransformType):
    name = "transform/rotate3d"
    version = "1.3.0"
    types = [
        "astropy.modeling.rotations.RotateNative2Celestial",
        "astropy.modeling.rotations.RotateCelestial2Native",
        "astropy.modeling.rotations.EulerAngleRotation",
    ]

    @classmethod
    def from_tree_transform(cls, node, ctx):
        if node["direction"] == "native2celestial":
            return modeling.rotations.RotateNative2Celestial(
                node["phi"], node["theta"], node["psi"]
            )
        elif node["direction"] == "celestial2native":
            return modeling.rotations.RotateCelestial2Native(
                node["phi"], node["theta"], node["psi"]
            )
        else:
            return modeling.rotations.EulerAngleRotation(
                node["phi"], node["theta"], node["psi"], axes_order=node["direction"]
            )

    @classmethod
    def to_tree_transform(cls, model, ctx):
        if isinstance(model, modeling.rotations.RotateNative2Celestial):
            try:
                node = {
                    "phi": _parameter_to_value(model.lon),
                    "theta": _parameter_to_value(model.lat),
                    "psi": _parameter_to_value(model.lon_pole),
                    "direction": "native2celestial",
                }
            except AttributeError:
                node = {
                    "phi": model.lon,
                    "theta": model.lat,
                    "psi": model.lon_pole,
                    "direction": "native2celestial",
                }
        elif isinstance(model, modeling.rotations.RotateCelestial2Native):
            try:
                node = {
                    "phi": _parameter_to_value(model.lon),
                    "theta": _parameter_to_value(model.lat),
                    "psi": _parameter_to_value(model.lon_pole),
                    "direction": "celestial2native",
                }
            except AttributeError:
                node = {
                    "phi": model.lon,
                    "theta": model.lat,
                    "psi": model.lon_pole,
                    "direction": "celestial2native",
                }
        else:
            node = {
                "phi": _parameter_to_value(model.phi),
                "theta": _parameter_to_value(model.theta),
                "psi": _parameter_to_value(model.psi),
                "direction": model.axes_order,
            }

        return node

    @classmethod
    def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert a.__class__ == b.__class__
        if a.__class__.__name__ == "EulerAngleRotation":
            assert_array_equal(a.phi, b.phi)
            assert_array_equal(a.psi, b.psi)
            assert_array_equal(a.theta, b.theta)
        else:
            assert_array_equal(a.lon, b.lon)
            assert_array_equal(a.lat, b.lat)
            assert_array_equal(a.lon_pole, b.lon_pole)


class RotationSequenceType(TransformType):
    name = "transform/rotate_sequence_3d"
    types = [
        "astropy.modeling.rotations.RotationSequence3D",
        "astropy.modeling.rotations.SphericalRotationSequence",
    ]
    version = "1.0.0"

    @classmethod
    def from_tree_transform(cls, node, ctx):
        angles = node["angles"]
        axes_order = node["axes_order"]
        rotation_type = node["rotation_type"]
        if rotation_type == "cartesian":
            return modeling.rotations.RotationSequence3D(angles, axes_order=axes_order)
        elif rotation_type == "spherical":
            return modeling.rotations.SphericalRotationSequence(
                angles, axes_order=axes_order
            )
        else:
            raise ValueError(f"Unrecognized rotation_type: {rotation_type}")

    @classmethod
    def to_tree_transform(cls, model, ctx):
        node = {"angles": list(model.angles.value)}
        node["axes_order"] = model.axes_order
        if isinstance(model, modeling.rotations.SphericalRotationSequence):
            node["rotation_type"] = "spherical"
        elif isinstance(model, modeling.rotations.RotationSequence3D):
            node["rotation_type"] = "cartesian"
        else:
            raise ValueError(f"Cannot serialize model of type {type(model)}")
        return node

    @classmethod
    def assert_equal(cls, a, b):
        TransformType.assert_equal(a, b)
        assert a.__class__.__name__ == b.__class__.__name__
        assert_array_equal(a.angles, b.angles)
        assert a.axes_order == b.axes_order


class GenericProjectionType(TransformType):
    @classmethod
    def from_tree_transform(cls, node, ctx):
        args = []
        for param_name, default in cls.params:
            args.append(node.get(param_name, default))

        if node["direction"] == "pix2sky":
            return cls.types[0](*args)
        else:
            return cls.types[1](*args)

    @classmethod
    def to_tree_transform(cls, model, ctx):
        node = {}
        if isinstance(model, cls.types[0]):
            node["direction"] = "pix2sky"
        else:
            node["direction"] = "sky2pix"
        for param_name, default in cls.params:
            val = getattr(model, param_name).value
            if val != default:
                node[param_name] = val
        return node

    @classmethod
    def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert a.__class__ == b.__class__


_generic_projections = {
    "zenithal_perspective": (
        "ZenithalPerspective",
        (("mu", 0.0), ("gamma", 0.0)),
        "1.3.0",
    ),
    "gnomonic": ("Gnomonic", (), None),
    "stereographic": ("Stereographic", (), None),
    "slant_orthographic": ("SlantOrthographic", (("xi", 0.0), ("eta", 0.0)), None),
    "zenithal_equidistant": ("ZenithalEquidistant", (), None),
    "zenithal_equal_area": ("ZenithalEqualArea", (), None),
    "airy": ("Airy", (("theta_b", 90.0),), "1.2.0"),
    "cylindrical_perspective": (
        "CylindricalPerspective",
        (("mu", 0.0), ("lam", 0.0)),
        "1.3.0",
    ),
    "cylindrical_equal_area": ("CylindricalEqualArea", (("lam", 0.0),), "1.3.0"),
    "plate_carree": ("PlateCarree", (), None),
    "mercator": ("Mercator", (), None),
    "sanson_flamsteed": ("SansonFlamsteed", (), None),
    "parabolic": ("Parabolic", (), None),
    "molleweide": ("Molleweide", (), None),
    "hammer_aitoff": ("HammerAitoff", (), None),
    "conic_perspective": (
        "ConicPerspective",
        (("sigma", 0.0), ("delta", 0.0)),
        "1.3.0",
    ),
    "conic_equal_area": ("ConicEqualArea", (("sigma", 0.0), ("delta", 0.0)), "1.3.0"),
    "conic_equidistant": (
        "ConicEquidistant",
        (("sigma", 0.0), ("delta", 0.0)),
        "1.3.0",
    ),
    "conic_orthomorphic": (
        "ConicOrthomorphic",
        (("sigma", 0.0), ("delta", 0.0)),
        "1.3.0",
    ),
    "bonne_equal_area": ("BonneEqualArea", (("theta1", 0.0),), "1.3.0"),
    "polyconic": ("Polyconic", (), None),
    "tangential_spherical_cube": ("TangentialSphericalCube", (), None),
    "cobe_quad_spherical_cube": ("COBEQuadSphericalCube", (), None),
    "quad_spherical_cube": ("QuadSphericalCube", (), None),
    "healpix": ("HEALPix", (("H", 4.0), ("X", 3.0)), None),
    "healpix_polar": ("HEALPixPolar", (), None),
}


def make_projection_types():
    for tag_name, (name, params, version) in _generic_projections.items():
        class_name = f"{name}Type"
        types = [
            f"astropy.modeling.projections.Pix2Sky_{name}",
            f"astropy.modeling.projections.Sky2Pix_{name}",
        ]

        members = {"name": f"transform/{tag_name}", "types": types, "params": params}
        if version:
            members["version"] = version

        globals()[class_name] = type(str(class_name), (GenericProjectionType,), members)

        __all__.append(class_name)


make_projection_types()
