/*============================================================================

  Program:     DTI ToolKit (DTI-TK)
  Module:      $RCSfile: Affine3D.cpp,v $
  Language:    C++
  Date:        $Date: 2012/10/02 18:20:15 $
  Version:     $Revision: 1.2 $

  Copyright (c) Gary Hui Zhang (garyhuizhang@gmail.com).
  All rights reserverd.

  DTI-TK is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  DTI-TK is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with DTI-TK.  If not, see <http://www.gnu.org/licenses/>.
============================================================================*/

#include "Affine3D.h"
#include "Rotation3D.h"
#include "SymTensor3D.h"
#include <fstream>
#include "../io/util.h"
#include "../volume/Vector3DVolume.h"
#include <iomanip>

namespace geometry {

	using volume::Vector3DVolume;
	using io::getNoOfWords;
  using io::parseDTITKFilename;
	
	Affine3D::Affine3D ()
		: vec(), mat(true) {}
	
	Affine3D::Affine3D (const char *filename)
		: vec(), mat(true) {
		load(filename);
	}
	
	Affine3D::Affine3D (const Vector3D& v)
		: vec(v), mat(true) {}
	
	Affine3D::Affine3D (const Matrix3D& m)
		: vec(), mat(m) {}
	
	Affine3D::Affine3D (const Vector3D& v, const Matrix3D& m)
		: vec(v), mat(m) {
		vec *= m;
		vec.negateEqual();
		vec += v;
	}
	
	Affine3D::Affine3D (const Affine3D& rhs) {
		*this = rhs;
	}
	
	Affine3D::~Affine3D () {}

	Affine3D& Affine3D::operator= (const Affine3D& rhs) {
		vec = rhs.vec;
		mat = rhs.mat;
		return *this;
	}
	
	Affine3D& Affine3D::operator= (const Matrix3D& rhs) {
		vec = 0.0;
		mat = rhs;
		return *this;
	}
	
	Affine3D& Affine3D::operator*= (const double rhs) {
		mat *= rhs;
		vec *= rhs;
		return *this;
	}
	
	Affine3D Affine3D::operator* (const double rhs) const {
		Affine3D aff(*this);
		aff *= rhs;
		return aff;
	}
	
	Affine3D operator* (const double lhs, const Affine3D& rhs) {
		return rhs * lhs;
	}
	
	Affine3D& Affine3D::operator+= (const Affine3D& rhs) {
		mat += rhs.mat;
		vec += rhs.vec;
		return *this;
	}

	Affine3D Affine3D::operator+ (const Affine3D& rhs) const {
		Affine3D aff(*this);
		aff += rhs;
		return aff;
	}
	
	Affine3D Affine3D::operator* (const Affine3D& rhs) const {
		Affine3D p(*this);
		p *= rhs;
		return p;
	}
	
	Affine3D& Affine3D::operator*= (const Affine3D& rhs) {
		multiplicationBy(rhs);
		return *this;
	}
	
	Affine3D& Affine3D::operator*= (const Matrix3D& rhs) {
		mat *= rhs;
		return *this;
	}
	
	void Affine3D::setQS (const double para[12]) {
		for (int i = 0; i < 3; ++i) {
			vec[i] = para[i];
		}
		mat.setQS(para + 3);
	}
	
	void Affine3D::setQS (const double para[12], const Vector3D& center) {
		setQS(para);
		moveOriginBy(center);
	}
	
	void Affine3D::getQS (double para[12]) const {
		// translation parameters
		for (int i = 0; i < 3; ++i) {
			para[i] = vec[i];
		}
		mat.getQS(para + 3);
	}
	
	void Affine3D::getQS (double para[12], const Vector3D& center) const {
		getQS(para);
		Vector3D tmp(center);
		tmp *= mat;
		tmp -= center;
		for (int i = 0; i < 3; ++i) {
			para[i] += tmp[i];
		}
	}
	
	void Affine3D::setQR (const double para[12], const bool upperTriangle) {
		for (int i = 0; i < 3; ++i) {
			vec[i] = para[i];
		}
		mat.setQR(para + 3, upperTriangle);
	}
	
	void Affine3D::moveOriginBy (const Vector3D& shift) {
		Vector3D inv(shift);
		inv.negateEqual();
		inv *= mat;
		vec += inv;
		vec += shift;
	}
	
	void Affine3D::setMatrix (const Matrix3D& in) {
		mat = in;
	}
	
	void Affine3D::setVector (const Vector3D& in) {
		vec = in;
	}
	
	void Affine3D::getMatrix (Matrix3D& out) const {
		out = mat;
	}

	void Affine3D::getVector (Vector3D& out) const {
		out = vec;
	}
	
	Affine3D& Affine3D::toIdentity () {
		mat.toIdentity();
		vec = Vector3D::zero;
		return *this;
	}
	
	Affine3D& Affine3D::inverseEqual () {
		mat.inverseEqual();
		vec *= mat;
		vec.negateEqual();
		return *this;
	}
	
	Affine3D Affine3D::inverse () const {
		Affine3D inv(*this);
		inv.inverseEqual();
		return inv;
	}
	
	Affine3D& Affine3D::sqrtEqual () {
		mat.sqrtEqual();
		Matrix3D tmp(mat);
		tmp += Matrix3D(true);
		tmp.inverseEqual();
		vec *= tmp;
		return *this;
	}
	
	Affine3D Affine3D::sqrt () const {
		Affine3D tmp(*this);
		tmp.sqrtEqual();
		return tmp;
	}
	
	void Affine3D::multiplicationBy (const Affine3D& rhs) {
		const Vector3D& rVec = rhs.vec;
		const Matrix3D& rMat = rhs.mat;
		vec += mat * rVec;
		mat *= rMat;
	}
	
	Affine3D& Affine3D::buildAffine (const Vector3D in[4], const Vector3D out[4]) {
		Matrix3D inMat;
		for (int i = 0; i < 3; ++i) {
			// ith column
			for (int j = 0; j < 3; ++j) {
				// jth row
				mat.setElement(j, i, out[i+1][j] - out[0][j]);
				inMat.setElement(j, i, in[i+1][j] - in[0][j]);
			}
		}
		
		// invert inMat
		inMat.inverseEqual();
		
		// compute the mat
		mat *= inMat;
		
		// compute the vec
		vec = out[0];
		vec -= mat * in[0];
		
		return *this;
	}
	
	Vector3D operator* (const Affine3D& lhs, const Vector3D& rhs) {
		Vector3D v(rhs);
		v *= lhs;
		return v;
	}
	
	Vector3D& operator *= (Vector3D& lhs, const Affine3D& rhs) {
		const Matrix3D& mat = rhs.mat;
		const Vector3D& tr = rhs.vec;
		lhs *= mat;
		lhs += tr;
		return lhs;
	}
	
	ostream& operator<< (ostream& out, const Affine3D& rhs) {
		out << rhs.vec << endl;
		out << rhs.mat;
		return out;
	}
	
	Affine3D Affine3D::computeShapeAverage (const char *in) {
		vector<string> files;
		io::parseFileList(in, files);
		Affine3D tmp;
		Matrix3D average;
		for (unsigned int i = 0; i < files.size(); ++i) {
			if (tmp.load(files[i].c_str())) {
				// TODO: still need to be determined
				// if QS or SQ should be used
				tmp.mat.toLogSPD();
				average += tmp.mat;
			}
		}
		average *= 1.0/files.size();
		SymTensor3D sym(average);
		sym.exp();
		average = sym;
		return Affine3D(average);
	}
	
  bool Affine3D::load (const char *filename) {
    string prefix;
    string suffix;
    parseDTITKFilename(filename, prefix, suffix);
    if (suffix == ".aff") {
      return loadDTITKFormat(filename);
    } else if (suffix == ".txt") {
      return loadITKFormat(filename);
	 } else if (suffix == ".mat") {
      return loadFSLFormat(filename);
    } else {
      cerr << "Unrecognized suffix for affine transform file" << endl;
      return false;
    }
  }
	
  bool Affine3D::loadDTITKFormat (const char *filename) {
		// local storage before read complete
		double matrix[3][3];
		double vector[3];
		
		// read
		ifstream in(filename);
		if (! in) {
			cerr << filename << " doesn't exist or can't be opened" << endl;
			return false;
		}
		
		const int bufSize = 256;
		char buf[bufSize];
		int counter = 0;
		
		++counter;
		in.getline(buf, bufSize);
		const char *line = "MATRIX";
		if (strcmp(buf, line) != 0) {
			cerr << "Invalid DTI-TK Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
		}
		
		for (int i = 0; i < 3; ++i) {
			++counter;
			in.getline(buf, bufSize);
			if (getNoOfWords(buf) != 3) {
				cerr << "Invalid Affine Transform File Format." << endl;
				cerr << "Line " << counter << endl;
				cerr << buf << endl;
				return false;
			}
			sscanf(buf, "%lf %lf %lf", matrix[i], matrix[i] + 1, matrix[i] + 2);
		}
		
		++counter;
		in.getline(buf, bufSize);
		line = "VECTOR";
		if (strcmp(buf, line) != 0) {
			cerr << "Invalid Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
		}
		
		++counter;
		in.getline(buf, bufSize);
		if (getNoOfWords(buf) != 3) {
			cerr << "Invalid Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
		}
		sscanf(buf, "%lf %lf %lf", vector, vector + 1, vector + 2);
		
		// copy
		for (int i = 0; i < 3; ++i) {
			vec[i] = vector[i];
			for (int j = 0; j < 3; ++j) {
				mat.setElement(i, j, matrix[i][j]);
			}
		}
		
		return true;
	}
	
  bool Affine3D::loadITKFormat (const char *filename) {
		// local storage before read complete
		double matrix[3][3];
		double vector[3];
		
		ifstream in(filename);
		if (! in) {
			cerr << filename << " doesn't exist or can't be opened" << endl;
			return false;
		}
		
		const int bufSize = 256;
		char buf[bufSize];
		int counter = 0;
		
		++counter;
		in.getline(buf, bufSize);
    string line = "#Insight Transform File V1.0";
    if (strcmp(buf, line.c_str()) != 0) {
			cerr << "Invalid ITK Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
    }
		
    ++counter;
		in.getline(buf, bufSize);
    line = "# Transform 0";
    string line_newStyle = "#Transform 0";
    if (strcmp(buf, line.c_str()) != 0 && strcmp(buf, line_newStyle.c_str()) != 0 ) {
			cerr << "Invalid ITK Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
    }
    
    ++counter;
		in.getline(buf, bufSize);
    line = "Transform: AffineTransform_double_3_3";
    if (strcmp(buf, line.c_str()) != 0) {
			cerr << "Invalid ITK Affine Transform File Format." << endl;
			cerr << "Line " << counter << endl;
			cerr << buf << endl;
			return false;
    }
    
    ++counter;
    string next;
    in >> next;
    if (next != "Parameters:") {
      cerr << "Invalid ITK Affine Transform File Format." << endl;
      cerr << "Line " << counter << endl;
      cerr << next << endl;
      return false;
    }
    
    for (int i = 0; i < 3; ++i) {
      for (int j = 0; j < 3; ++j) {
        in >> matrix[i][j];
      }
    }
    
    for (int i = 0; i < 3; ++i) {
      in >> vector[i];
    }
     
		// copy
		for (int i = 0; i < 3; ++i) {
			vec[i] = vector[i];
			for (int j = 0; j < 3; ++j) {
				mat.setElement(i, j, matrix[i][j]);
			}
		}
		
    return true;
  }
  
  bool Affine3D::loadFSLFormat (const char *filename) {
		// local storage before read complete
		double matrix[4][4];
		
		ifstream in(filename);
		if (! in) {
			cerr << filename << " doesn't exist or can't be opened" << endl;
			return false;
		}
		
		// read
		for (int i = 0; i < 4; ++i) {
			for (int j = 0; j < 4; ++j) {
				in >> matrix[i][j];
			}
		}
		
		// copy
		for (int i = 0; i < 3; ++i) {
			vec[i] = matrix[i][3];
			for (int j = 0; j < 3; ++j) {
				mat.setElement(i, j, matrix[i][j]);
			}
		}
		
		return true;
  }
  
  void Affine3D::saveAs (const char *filename) const {
    string prefix;
    string suffix;
    parseDTITKFilename(filename, prefix, suffix);
    if (suffix == ".aff") {
      saveAsDTITKFormat(filename);
    } else if (suffix == ".txt") {
      saveAsITKFormat(filename);
    } else if (suffix == ".mat") {
      saveAsFSLFormat(filename);
	 }
  }
	
  void Affine3D::saveAsDTITKFormat (const char *filename) const {
		ofstream out(filename);
		out << "MATRIX" << endl;
		out.precision(16);
		for (int i = 0; i < 3; ++i) {
			for (int j = 0; j < 2; ++j) {
				out << setw(24) << mat.getElement(i, j) << " ";
			}
			out << setw(24) << mat.getElement(i, 2) << endl;
		}
		out << "VECTOR" << endl;
		for (int i = 0; i < 2; ++i) {
			out << setw(24) << vec[i] << " ";
		}
		out << setw(24) << vec[2] << endl;
	}
	
  void Affine3D::saveAsITKFormat (const char *filename) const {
    ofstream out(filename);
    out << "#Insight Transform File V1.0" << endl;
    out << "# Transform 0" << endl;
    out << "Transform: AffineTransform_double_3_3" << endl;
    out << "Parameters: ";
    for (int i = 0; i < 3; ++i) {
      for (int j = 0; j < 3; ++j) {
        out << mat.getElement(i, j) << " ";
      }
    }
    out << vec[0] << " " << vec[1] << " " << vec[2] << endl;
    out << "FixedParameters: 0 0 0" << endl;
  }
  
  void Affine3D::saveAsFSLFormat (const char *filename) const {
    ofstream out(filename);
    for (int i = 0; i < 3; ++i) {
      for (int j = 0; j < 3; ++j) {
        out << mat.getElement(i, j) << " ";
      }
		out << vec[i] << endl;
    }
    out << 0 << " " << 0 << " " << 0 << " " << 1 << endl;
  }
	
  void Affine3D::writeAsDispField (const char *filename, const int size[3], const double vsize[3]) const {
		const double origin[3] = {0.0, 0.0, 0.0};
		Vector3DVolume vv(size);
		vv.setVSize(vsize);
		vv.setOrigin(origin);
		
		Vector3D in;
		Vector3D out;
		for (int k = 0; k < size[2]; ++k) {
			in[2] = k * vsize[2];
			for (int j = 0; j < size[1]; ++j) {
				in[1] = j * vsize[1];
				for (int i = 0; i < size[0]; ++i) {
					in[0] = i * vsize[0];
					out = in;
					out *= *this;
					out -= in;
					vv.voxel[i][j][k] = out;
				}
			}
		}
		vv.writeVolAs(filename);
	}
	
}