"""
Linear transformations.
"""

import numpy as np
import csv


from rotation cimport Rotation2D, Rotation3D

cdef class Transformation:
    """Rotation followed by translation.

    Dimension should be at least 2.
    """

    def __init__(self, rotation, translation):
        """Initialise transformation.

        Args:
            rotation: Rotation object representing rotation of object.
            translation: Translation vector as (D, ) array.
        """

        self._rotation = rotation
        self._translation = translation

    @property
    def rotation(self):
        return self._rotation

    @rotation.setter
    def rotation(self, value):
        self._rotation = value

    @property
    def translation(self):
        return self._translation # useful for in-place changes

    def __call__(self, x):
        """Returns transformed vectors.

        First rotates, then translates.

        Args:
            x: (N, D) array of data points to be transformed.
            y: (N, D) array where results will be stored.

        Returns:
            (N, D) array of transformed points.
        """

        return self.rotation(x) + self.translation

    def inverse_transform(self, x):
        """Returns inversely transformed vectors.

        This should be the inverse of self.__call__.
        """

        return self.rotation.inverse_rotate(x - np.array(self.translation))

    cpdef transform(self, double[:,:] x, double[:,:] y):
        """Apply transformation without creating new vectors.

        First rotates, then translates.

        Args:
            x: (N, D) array of data points to be transformed.
            y: (N, D) array where results will be stored.

        Note: when N=1 the vectors should still be 2 dimensional.
        """

        cdef:
            int i, j
            ## double[:] translation = self._translation

        self._rotation.rotate(x, y)

        for i in range(y.shape[0]):
            for j in range(y.shape[1]):
                y[i, j] += self._translation[j]


    @classmethod
    def from_rotation(cls, rotation):
        """Constructs a transformation with given rotation.

        Translation is assumed to be zero.

        Args:
            rotation: Rotation object.

        Returns:
            Translation object.
        """

        translation = np.zeros(rotation.dimension)

        return cls(rotation, translation)

    @classmethod
    def random(cls):
        """Returns a random rotation together with a zero translation."""

        raise NotImplementedError

    def save(self, parent_group, name='transformation'):
        """Save to group in HDF5 database.

        Args:
            parent_group: Group within which to create the
                transformation group.
            name: Name of the new transformation group.
        """

        if name in parent_group:
            del parent_group[name]

        group = parent_group.create_group(name)
        self.rotation.save(group)
        translation_dataset = group.create_dataset(
            'translation',
            (<object>self._translation).shape,
            np.double)
        translation_dataset[...] = self.translation

    def copy(self):
        """Return deep copy."""

        rotation2 = self._rotation.copy()
        translation2 = np.array(self._translation).copy()

        return self.__class__(rotation2, translation2)

    def compose(self, other):
        """Return composition of this transformation with 'other'.

        First the other transformation is applied, then this one.
        """

        # Before 25.3.2014:
        ## r = other.rotation.compose(self.rotation)
        ## t = (np.dot(np.array(other.rotation.matrix),
        ##             np.array(self.translation)) +
        ##      np.array(other.translation))

        r = self.rotation.compose(other.rotation)
        t = (np.dot(np.array(self.rotation.matrix),
                    np.array(other.translation)) +
             np.array(self.translation))

        return Transformation(r, t)

    def invert(self):
        """Return inverse transformation."""

        r = self.rotation.invert()
        t = -np.dot(np.array(r.matrix), np.array(self.translation))

        ## return Transformation(r, t)
        return self.__class__(r, t)



cdef class Transformation2D(Transformation):
    """2D transformation."""

    @classmethod
    def random(cls):
        """Returns a random rotation together with a zero translation."""

        return cls.from_rotation(Rotation2D.random())


cdef class Transformation3D(Transformation):
    """3D transformation."""

    @classmethod
    def random(cls):
        """Returns a random rotation together with a zero translation."""

        return cls.from_rotation(Rotation3D.random())

    @classmethod
    def identity(cls):
        """Identity rotation followed by zero translation."""

        return cls.from_rotation(Rotation3D.identity())

    @classmethod
    def load(cls, parent_group, name='transformation'):
        """Loads Transformation from group in HDF5 database."""

        group = parent_group[name]
        rotation = Rotation3D.load(group)
        translation = group['translation'][:]

        return cls(rotation, translation)


    @classmethod
    def load_chimera(cls, filename):
        """Load transformation from chimera text file.

        Text file assumed to be produced during fit-in-map procedure.
        Use only last transformation in text file.
        """

        fin = open(filename, 'r')
        lines = fin.readlines()
        fin.close()

        if len(lines) < 8:
            print 'File not long enough.'
            return

        t = np.array([[float(x) for x in y.split()] for y in lines[-8:-5]])
        m = t[:, :3]
        v = t[:, -1]

        return Transformation(Rotation3D(m), v)

    @classmethod
    def load_chimera_matrix(cls, filename):
        """Load transformation from chimera text file.

        Text file produced by 'matrixget' or 'matrixset' command.
        Use only last transformation in text file.
        """

        fin = open(filename, 'r')
        lines = fin.readlines()
        fin.close()

        if len(lines) < 8:
            print 'File not long enough.'
            return

        t = np.array([[float(x) for x in y.split()] for y in lines[-3:]])
        m = t[:, :3]
        v = t[:, -1]

        return Transformation(Rotation3D(m), v)


    def __str__(self):

        return str(self._rotation) + '\n' + str(np.array(self._translation))



class Transformations(object):
    """List of Transformation objects."""

    def __init__(self, transformations):
        self.ts = transformations

    def to_list(self):
        return self.ts

    def save_csv(self, filename):
        """Save transformations as csv file."""

        with open(filename, 'wb') as csvfile:
            writer = csv.writer(csvfile, delimiter=',')
            for t in self.ts:
                x = np.array(t.rotation.matrix).flatten()
                x = np.append(x, np.array(t.translation))
                writer.writerow(x)

    @classmethod
    def load3D_csv(cls, filename):
        """Load 3D transformations from csv file.

        """

        ts = []

        with open(filename, 'rb') as csvfile:
            reader = csv.reader(csvfile, delimiter=',')
            for row in reader:
                x = np.array([float(a) for a in row])
                rotation = Rotation3D(x[:9].reshape((3, 3)))
                translation = x[-3:]
                ts.append(Transformation3D(rotation, translation))

        return cls(ts)


