# coding=utf-8
import geometry
from abc import ABCMeta, abstractmethod
from typing import NewType

import numpy as np
from duckietown_serialization_ds1 import Serializable

import geometry as geo
from contracts import contract, new_contract
from duckietown_world.seqs.tsequence import Timestamp
from duckietown_world.seqs import GenericSequence


__all__ = ["TransformSequence", "Transform", "SE2Transform", "Scale2D", "Matrix2D"]

SE2value = NewType("SE2value", np.ndarray)


class Transform(metaclass=ABCMeta):
    @abstractmethod
    def asmatrix2d(self) -> "Matrix2D":
        """ """

    def as_SE2(self):
        m = self.asmatrix2d().m
        try:
            geo.SE2.belongs(m)
        except:
            msg = "Cannot convert %s" % type(self).__name__
            msg += "m = %s" % m
            raise Exception(msg)
        return m


new_contract("Transform", Transform)


class TransformSequence(Transform):

    # @contract(transforms='list[>=1](Transform)')
    def __init__(self, transforms):
        self.transforms = list(transforms)

    def asmatrix2d(self):
        ms = [_.asmatrix2d() for _ in self.transforms]
        # print('ms: %s' % ms)
        result = ms[0].m

        for mi in ms[1:]:
            result = np.dot(result, mi.m)
        return Matrix2D(result)

    def __repr__(self):
        return "TransformSequence(%s)" % self.transforms


class VariableTransformSequence(TransformSequence, GenericSequence[Transform]):  # XXX
    def at(self, t: Timestamp):
        res = []
        for transform in self.transforms:
            if isinstance(transform, GenericSequence):
                transform = transform.at(t)
            res.append(transform)
        return TransformSequence(res)

    def get_end(self):
        return None

    def get_start(self):
        return None

    def get_sampling_points(self):
        pass


class SE2Transform(Transform, Serializable):
    @contract(p="seq[2](float|int)")
    def __init__(self, p, theta):
        self.p = np.array(p, dtype="float64")
        self.theta = float(theta)

    def __repr__(self):
        d = np.rad2deg(self.theta)
        return "SE2Transform(%s,%.1f)" % (self.p.tolist(), d)

    @classmethod
    def identity(cls) -> "SE2Transform":
        return SE2Transform([0.0, 0.0], 0.0)

    @classmethod
    def from_SE2(cls, q: SE2value) -> "SE2Transform":
        """ From a matrix """
        translation, angle = geo.translation_angle_from_SE2(q)
        return SE2Transform(translation, angle)

    def params_to_json_dict(self):
        res = {}
        if np.allclose(self.theta, 0):
            pass
        else:
            for a in [-270, -180, -90, -45, +45, 90, 180, 270]:
                if np.allclose(a, np.rad2deg(self.theta)):
                    res["theta_deg"] = a
                    break
            else:
                res["theta"] = self.theta
        if np.allclose(np.linalg.norm(self.p), 0):
            pass
        else:
            res["p"] = self.p
        return res

    @classmethod
    def params_from_json_dict(cls, d):
        if d is None:
            d = {}
        p = d.pop("p", [0.0, 0.0])

        if "theta" in d:
            theta = d.pop("theta")

        elif "theta_deg" in d:
            theta_deg = d.pop("theta_deg")
            theta = np.deg2rad(theta_deg)
        else:
            theta = 0.0

        return dict(p=p, theta=theta)

    def as_SE2(self) -> SE2value:

        M = geometry.SE2_from_translation_angle(self.p, self.theta)
        return M

    def asmatrix2d(self):
        M = self.as_SE2()
        return Matrix2D(M)


new_contract("SE2Transform", SE2Transform)


class Scale2D(Transform, Serializable):
    def __init__(self, scale):
        self.scale = scale

    def asmatrix2d(self):
        S = self.scale
        M = np.array([[S, 0, 0], [0, S, 0], [0, 0, 1]])
        return Matrix2D(M)


class Matrix2D(Transform, Serializable):
    def __init__(self, m):
        self.m = np.array(m, "float32")
        assert self.m.shape == (3, 3)

    def asmatrix2d(self):
        return self
