#include <mitsuba/core/track.h>
#include <mitsuba/core/aabb.h>
#include <Eigen/SVD>

MTS_NAMESPACE_BEGIN

AnimatedTransform::AnimatedTransform(const AnimatedTransform *trafo)
		: m_transform(trafo->m_transform) {
	m_tracks.reserve(trafo->getTrackCount());
	for (size_t i=0; i<trafo->getTrackCount(); ++i) {
		AbstractAnimationTrack *track = trafo->getTrack(i)->clone();
		m_tracks.push_back(track);
		track->incRef();
	}
}

AnimatedTransform::AnimatedTransform(Stream *stream) {
	size_t nTracks = stream->readSize();
	if (nTracks == 0) {
		m_transform = Transform(stream);
	} else {
		for (size_t i=0; i<nTracks; ++i) {
			AbstractAnimationTrack::EType type =
				(AbstractAnimationTrack::EType) stream->readUInt();
			AbstractAnimationTrack *track = NULL;
			switch (type) {
				case AbstractAnimationTrack::ETranslationX:
				case AbstractAnimationTrack::ETranslationY:
				case AbstractAnimationTrack::ETranslationZ:
				case AbstractAnimationTrack::EScaleX:
				case AbstractAnimationTrack::EScaleY:
				case AbstractAnimationTrack::EScaleZ:
				case AbstractAnimationTrack::ERotationX:
				case AbstractAnimationTrack::ERotationY:
				case AbstractAnimationTrack::ERotationZ:
					track = new FloatTrack(type, stream);
					break;
				case AbstractAnimationTrack::ETranslationXYZ:
				case AbstractAnimationTrack::EScaleXYZ:
					track = new VectorTrack(type, stream);
					break;
				case AbstractAnimationTrack::ERotationQuat:
					track = new QuatTrack(type, stream);
					break;
				default:
					Log(EError, "Encountered an unknown animation track type (%i)!", type);
			}

			track->incRef();
			m_tracks.push_back(track);
		}
	}
}

void AnimatedTransform::addTrack(AbstractAnimationTrack *track) {
	track->incRef();
	m_tracks.push_back(track);
}

AABB1 AnimatedTransform::getTimeBounds() const {
	if (m_tracks.size() == 0)
		return AABB1(0.0f, 0.0f);

	Float min =  std::numeric_limits<Float>::infinity();
	Float max = -std::numeric_limits<Float>::infinity();

	for (size_t i=0; i<m_tracks.size(); ++i) {
		const AbstractAnimationTrack *track = m_tracks[i];
		size_t size = track->getSize();
		SAssert(size > 0);
		min = std::min(min, track->getTime(0));
		max = std::max(max, track->getTime(size-1));
	}

	return AABB1(min, max);
}

AABB AnimatedTransform::getTranslationBounds() const {
	if (m_tracks.size() == 0) {
		Point p = m_transform(Point(0.0f));
		return AABB(p, p);
	}

	AABB aabb;

	for (size_t i=0; i<m_tracks.size(); ++i) {
		const AbstractAnimationTrack *absTrack = m_tracks[i];
		switch (absTrack->getType()) {
			case AbstractAnimationTrack::ETranslationX:
			case AbstractAnimationTrack::ETranslationY:
			case AbstractAnimationTrack::ETranslationZ: {
					int idx  = absTrack->getType() - AbstractAnimationTrack::ETranslationX;
					const FloatTrack *track =
						static_cast<const FloatTrack *>(absTrack);
					for (size_t j=0; j<track->getSize(); ++j) {
						Float value = track->getValue(j);
						aabb.max[idx] = std::max(aabb.max[idx], value);
						aabb.min[idx] = std::min(aabb.min[idx], value);
					}
				}
				break;

			case AbstractAnimationTrack::ETranslationXYZ: {
					const VectorTrack *track =
						static_cast<const VectorTrack *>(absTrack);
					for (size_t j=0; j<track->getSize(); ++j)
						aabb.expandBy(Point(track->getValue(j)));
				}
				break;
			default:
				break;
		}
	}
	for (int i=0; i<3; ++i) {
		if (aabb.min[i] > aabb.max[i])
			aabb.min[i] = aabb.max[i] = 0.0f;
	}

	return aabb;
}

AABB AnimatedTransform::getSpatialBounds(const AABB &aabb) const {
	AABB result;

	if (m_tracks.size() == 0) {
		for (int j=0; j<8; ++j)
			result.expandBy(m_transform(aabb.getCorner(j)));
	} else {
		/* Compute approximate bounds */
		int nSteps = 100;
		AABB1 timeBounds = getTimeBounds();
		Float step = timeBounds.getExtents().x / (nSteps-1);

		for (int i=0; i<nSteps; ++i) {
			const Transform &trafo = eval(timeBounds.min.x + step * i);
			for (int j=0; j<8; ++j)
				result.expandBy(trafo(aabb.getCorner(j)));
		}
	}

	return result;
}

AnimatedTransform::~AnimatedTransform() {
	for (size_t i=0; i<m_tracks.size(); ++i)
		m_tracks[i]->decRef();
}

void AnimatedTransform::sortAndSimplify() {
	bool isStatic = true;

	for (size_t i=0; i<m_tracks.size(); ++i) {
		AbstractAnimationTrack *track = m_tracks[i];
		bool isNeeded = false;
		switch (track->getType()) {
			case AbstractAnimationTrack::ETranslationX:
			case AbstractAnimationTrack::ETranslationY:
			case AbstractAnimationTrack::ETranslationZ:
			case AbstractAnimationTrack::ERotationX:
			case AbstractAnimationTrack::ERotationY:
			case AbstractAnimationTrack::ERotationZ:
			case AbstractAnimationTrack::EScaleX:
			case AbstractAnimationTrack::EScaleY:
			case AbstractAnimationTrack::EScaleZ:
				isNeeded = static_cast<FloatTrack *>(track)->sortAndSimplify();
				break;
			case AbstractAnimationTrack::ETranslationXYZ:
			case AbstractAnimationTrack::EScaleXYZ:
				isNeeded = static_cast<VectorTrack *>(track)->sortAndSimplify();
				break;
			case AbstractAnimationTrack::ERotationQuat:
				isNeeded = static_cast<QuatTrack *>(track)->sortAndSimplify();
				break;
			default:
				Log(EError, "Encountered an unsupported "
					"animation track type: %i!", track->getType());
		}
		if (isNeeded) {
			isStatic &= track->getSize() == 1;
		} else {
			m_tracks.erase(m_tracks.begin() + i);
			track->decRef();
			--i;
		}
	}

	if (isStatic) {
		Transform temp;
		temp = eval(0);
		m_transform = temp;
		for (size_t i=0; i<m_tracks.size(); ++i)
			m_tracks[i]->decRef();
		m_tracks.clear();
	}
}


const AbstractAnimationTrack *AnimatedTransform::findTrack(AbstractAnimationTrack::EType type) const {
	for (size_t i=0; i<m_tracks.size(); ++i) {
		AbstractAnimationTrack *track = m_tracks[i];
		if (track->getType() == type)
			return track;
	}
	return NULL;
}
AbstractAnimationTrack *AnimatedTransform::findTrack(AbstractAnimationTrack::EType type) {
	for (size_t i=0; i<m_tracks.size(); ++i) {
		AbstractAnimationTrack *track = m_tracks[i];
		if (track->getType() == type)
			return track;
	}
	return NULL;
}

void AnimatedTransform::prependScale(const Vector &scale) {
	FloatTrack *trackX = (FloatTrack *) findTrack(AbstractAnimationTrack::EScaleX);
	FloatTrack *trackY = (FloatTrack *) findTrack(AbstractAnimationTrack::EScaleY);
	FloatTrack *trackZ = (FloatTrack *) findTrack(AbstractAnimationTrack::EScaleZ);
	VectorTrack *trackXYZ = (VectorTrack *) findTrack(AbstractAnimationTrack::EScaleXYZ);

	if (m_tracks.empty()) {
		m_transform = m_transform * Transform::scale(scale);
	} else if (trackXYZ) {
		trackXYZ->prependTransformation(scale);
	} else if (trackX && trackY && trackZ) {
		if (trackX) {
			trackX->prependTransformation(scale.x);
		} else {
			trackX = new FloatTrack(AbstractAnimationTrack::EScaleX);
			trackX->append(0.0f, scale.x); addTrack(trackX);
		}

		if (trackY) {
			trackY->prependTransformation(scale.y);
		} else {
			trackY = new FloatTrack(AbstractAnimationTrack::EScaleY);
			trackY->append(0.0f, scale.y); addTrack(trackY);
		}

		if (trackZ) {
			trackZ->prependTransformation(scale.z);
		} else {
			trackZ = new FloatTrack(AbstractAnimationTrack::EScaleZ);
			trackZ->append(0.0f, scale.z); addTrack(trackZ);
		}
	} else {
		trackXYZ = new VectorTrack(AbstractAnimationTrack::EScaleXYZ);
		trackXYZ->append(0.0f, scale);
		addTrack(trackXYZ);
	}
}

void AnimatedTransform::collectKeyframes(std::set<Float> &result) const {
	for (size_t i=0; i<m_tracks.size(); ++i) {
		const AbstractAnimationTrack *track = m_tracks[i];

		for (size_t j=0; j<track->getSize(); ++j)
			result.insert(track->getTime(j));
	}

	if (result.size() == 0)
		result.insert((Float) 0);
}

void AnimatedTransform::serialize(Stream *stream) const {
	stream->writeSize(m_tracks.size());
	if (m_tracks.size() == 0) {
		m_transform.serialize(stream);
	} else {
		for (size_t i=0; i<m_tracks.size(); ++i)
			m_tracks[i]->serialize(stream);
	}
}

void AnimatedTransform::TransformFunctor::operator()(const Float &t, Transform &trafo) const {
	Vector translation(0.0f);
	Vector scale(1.0f);
	Quaternion rotation;

	for (size_t i=0; i<m_tracks.size(); ++i) {
		AbstractAnimationTrack *track = m_tracks[i];
		switch (track->getType()) {
			case AbstractAnimationTrack::ETranslationX:
				translation.x = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::ETranslationY:
				translation.y = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::ETranslationZ:
				translation.z = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::ETranslationXYZ:
				translation = static_cast<VectorTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::EScaleX:
				scale.x = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::EScaleY:
				scale.y = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::EScaleZ:
				scale.z = static_cast<FloatTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::EScaleXYZ:
				scale = static_cast<VectorTrack *>(track)->eval(t);
				break;
			case AbstractAnimationTrack::ERotationQuat:
				rotation = static_cast<QuatTrack *>(track)->eval(t);
				break;
			default:
				Log(EError, "Encountered an unsupported "
					"animation track type: %i!", track->getType());
		}
	}

	trafo = Transform::translate(translation);

	if (!rotation.isIdentity())
		trafo = trafo * rotation.toTransform();

	if (scale != Vector(0.0f))
		trafo = trafo * Transform::scale(scale);
}

void AnimatedTransform::appendTransform(Float time, const Transform &trafo) {
	/* Compute the polar decomposition and insert into the animated transform;
	   uh oh.. we have to get rid of the two separate matrix libraries at some point :) */
	typedef Eigen::Matrix<Float, 3, 3> EMatrix;

	if (m_tracks.size() == 0) {
		ref<VectorTrack> translation = new VectorTrack(VectorTrack::ETranslationXYZ);
		ref<QuatTrack> rotation = new QuatTrack(VectorTrack::ERotationQuat);
		ref<VectorTrack> scaling = new VectorTrack(VectorTrack::EScaleXYZ);
		translation->reserve(2);
		rotation->reserve(2);
		scaling->reserve(2);
		addTrack(translation);
		addTrack(rotation);
		addTrack(scaling);
	} else if (m_tracks.size() != 3 ||
			m_tracks[0]->getType() != VectorTrack::ETranslationXYZ ||
			m_tracks[1]->getType() != VectorTrack::ERotationQuat ||
			m_tracks[2]->getType() != VectorTrack::EScaleXYZ) {
		Log(EError, "AnimatedTransform::appendTransform(): unsupported internal configuration!");
	}

	const Matrix4x4 m = trafo.getMatrix();
	EMatrix A;

	A << m(0, 0), m(0, 1), m(0, 2),
		 m(1, 0), m(1, 1), m(1, 2),
		 m(2, 0), m(2, 1), m(2, 2);

	Eigen::JacobiSVD<EMatrix> svd(A, Eigen::ComputeFullU | Eigen::ComputeFullV);
	EMatrix U = svd.matrixU(), V = svd.matrixV(), S = svd.singularValues().asDiagonal();

	if (svd.singularValues().prod() < 0) {
		S = -S; U = -U;
	}

	EMatrix Q = U*V.transpose();
	EMatrix P = V*S*V.transpose();

	VectorTrack *translation = (VectorTrack *) m_tracks[0];
	QuatTrack *rotation = (QuatTrack *) m_tracks[1];
	VectorTrack *scaling = (VectorTrack *) m_tracks[2];

	rotation->append(time, Quaternion::fromMatrix(
		Matrix4x4(
			Q(0, 0), Q(0, 1), Q(0, 2), 0.0f,
			Q(1, 0), Q(1, 1), Q(1, 2), 0.0f,
			Q(2, 0), Q(2, 1), Q(2, 2), 0.0f,
			0.0f,    0.0f,    0.0f,    1.0f
		)
	));

	scaling->append(time, Vector(P(0, 0), P(1, 1), P(2, 2)));
	translation->append(time, Vector(m(0, 3), m(1, 3), m(2, 3)));
}

std::string AnimatedTransform::toString() const {
	if (m_tracks.size() == 0) {
		return m_transform.toString();
	} else {
		std::ostringstream oss;
		oss << "AnimatedTransform[tracks=" << m_tracks.size() << "]";
		return oss.str();
	}
}

MTS_IMPLEMENT_CLASS(AbstractAnimationTrack, true, Object)
MTS_IMPLEMENT_CLASS(AnimatedTransform, false, Object)
MTS_NAMESPACE_END
