/*
   For more information, please see: http://software.sci.utah.edu

   The MIT License

   Copyright (c) 2009 Scientific Computing and Imaging Institute,
   University of Utah.

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included
   in all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
   OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE.
*/

#include <Core/Math/MiscMath.h>
#include <Core/Algorithms/BrainStimulator/ModelGenericCoilAlgorithm.h>
#include <Core/Algorithms/BrainStimulator/BiotSavartSolverAlgorithm.h>
#include <Core/Datatypes/Legacy/Field/Field.h>
#include <Core/Datatypes/Legacy/Field/VField.h>
#include <Core/Datatypes/Legacy/Field/VMesh.h>
#include <Core/Datatypes/Legacy/Field/FieldInformation.h>
#include <Core/GeometryPrimitives/Vector.h>
#include <Core/GeometryPrimitives/Transform.h>
#include <Core/GeometryPrimitives/Point.h>
#include <Core/Algorithms/Base/AlgorithmPreconditions.h>

//#include <cassert>
//#include <exception>

#include <boost/lexical_cast.hpp>

using namespace SCIRun::Core::Algorithms;
using namespace SCIRun::Core::Algorithms::BrainStimulator;
using namespace SCIRun::Core::Datatypes;
using namespace SCIRun::Core::Geometry;
using namespace SCIRun::Core::Utility;
using namespace SCIRun;

ALGORITHM_PARAMETER_DEF(BrainStimulator, Type);
ALGORITHM_PARAMETER_DEF(BrainStimulator, FigureOf8CoilShape);
ALGORITHM_PARAMETER_DEF(BrainStimulator, Rings);
ALGORITHM_PARAMETER_DEF(BrainStimulator, WingsAngle);
ALGORITHM_PARAMETER_DEF(BrainStimulator, Current);
ALGORITHM_PARAMETER_DEF(BrainStimulator, Radius);
ALGORITHM_PARAMETER_DEF(BrainStimulator, InnerRadius);
ALGORITHM_PARAMETER_DEF(BrainStimulator, OuterRadius);
ALGORITHM_PARAMETER_DEF(BrainStimulator, Distance);
ALGORITHM_PARAMETER_DEF(BrainStimulator, Layers);
ALGORITHM_PARAMETER_DEF(BrainStimulator, LayerStepSize);
ALGORITHM_PARAMETER_DEF(BrainStimulator, LevelOfDetail);

		class BaseSegments
		{
			public:
				BaseSegments(
					std::vector<Vector>& p,
					std::vector<size_t>& i,
					std::vector<double>& v):
					points(p),
					indices(i),
					values(v),
					pc(0),
					start_idx(0),
					end_idx(0)

				{
				}
				virtual ~BaseSegments()
				{
					this->Terminate();
				}

				void AddPoint(Vector point, double value)
				{
					points.push_back(point);

					size_t psize = points.size();

					if( pc > 0)
					{

						indices.push_back(psize-2);
						indices.push_back(psize-1);
						values.push_back(value);
					}
					else
					{
						start_idx = psize - 1;
					}

					++pc;
				}

				void Transform(Transform& t)
				{
					for(size_t i = start_idx; i < start_idx + pc; ++i)
					{
						points[i] = t * points[i];
					}
				}

				virtual void Terminate()
				{
					pc = 0;
				}

				std::vector<Vector>& points;
				std::vector<size_t>& indices;
				std::vector<double>& values;
				size_t pc;
				size_t start_idx;
				size_t end_idx;

			private:

				//! Prevent copying
    			BaseSegments & operator = (const BaseSegments & other);
    			BaseSegments(const BaseSegments & other);
		};

		class ClosedSegments : public BaseSegments
		{
			public:
				ClosedSegments(
					std::vector<Vector>& p,
					std::vector<size_t>& i,
					std::vector<double>& v)
						:BaseSegments(p,i,v)
				{

				}
				virtual ~ClosedSegments()
				{
					if(pc)
					{
						this->Terminate();
					}

					//std::cout << "~ClosedSegments() Points Indices Values: " <<  points.size() << " " << indices.size() << " " << values.size() << std::endl;

					//assert(points.size() > 0);
					//assert(points.size() == values.size());
					//assert(points.size()*2 == indices.size());
				}

				void Terminate()
				{
					size_t psize = points.size();

					//close the segments and make a circle
					indices.push_back(psize-1);
					indices.push_back(psize-pc);

					values.push_back(values[values.size()-1]);

					//std::cout << "ClosedSegment-Teminate-> psize:" << psize << " pc:" << pc << "    " << std::endl;

					pc = 0;
				}
		};

		class OpenSegments : public BaseSegments
		{
			public:
				OpenSegments(
					std::vector<Vector>& p,
					std::vector<size_t>& i,
					std::vector<double>& v)
						:BaseSegments(p,i,v)
				{}
				~OpenSegments()
				{
					//assert(points.size() > 0);
					//assert(indices.size() == values.size()*2);
				}
		};

		//--------------------------------------------------------------

		class BaseCoilgen
		{
			public:
				BaseCoilgen(const AlgorithmBase* algo, ModelTMSCoilAlgorithm::Args& args) :
				  ref_cnt(0),
				  algo(algo),

				  coilLOD(args.coilLevelDetails),
				  coilType(args.type),
				  coilLayers(args.coilLayers),
				  rings(args.rings),

				  coilLayersStep(args.coilLayersStep),
				  wingsAngle(args.wingsAngle),
				  current(args.current),
				  innerR(args.coilRadiusInner),
				  outerR(args.coilRadiusOuter),
				  outerD(args.coilDistance)

				{

				}

				virtual ~BaseCoilgen()
				{
				}

				void Execute(FieldHandle& meshFieldHandle) const
				{

					std::vector<Vector> coilPoints;
					std::vector<size_t> coilIndices;

					this->Generate(meshFieldHandle);
				}

				//! Global reference counting
				int ref_cnt;

			protected:

				//! ref to the executing algorithm context
				const AlgorithmBase* algo;
				size_t coilLOD;
				size_t coilType;
				size_t coilLayers;
				const size_t rings;
				double coilLayersStep;
				double wingsAngle;
				double current;
				const double innerR;
				const double outerR;
				const double outerD;

				//! Local entry function, must be implemented by each specific kernel
				virtual void Generate(FieldHandle& meshHandle) const = 0;

				void GenPointsCircular(
					BaseSegments& segments,
					Vector origin,
					double radius,
					double value,
					double fromPI,
					double toPI,
					double extLOD = 0.0) const
				{

					double dPI = fabs(toPI - fromPI);

					double minPI = M_PI /  ( 8.0 * coilLOD + coilLOD * extLOD );

					assert(dPI > minPI);

					size_t nsegments = 2;
					double iPI = dPI / nsegments;

					while(iPI > minPI)
					{
						nsegments++;
						iPI = dPI / nsegments;
					}

					//algo->remark("#Segments(LOD):  " +  boost::lexical_cast<std::string>(nsegments) );

					dPI = toPI - fromPI;

					iPI = dPI / nsegments;

					for(size_t i = 0; i < nsegments; i++)
					{
						Vector point(origin.x() + radius * cos(fromPI + iPI*i), origin.y() + radius * sin(fromPI + iPI*i), origin.z());
						segments.AddPoint(point,value);
					}
				}

				void BuildScirunMesh(
						const std::vector<Vector>& points,
						const std::vector<size_t>& indices,
						const std::vector<double>& values,
						FieldHandle& meshHandle) const
				{

					VMesh* mesh = meshHandle->vmesh();

					//! add nodes to the new mesh
					for(size_t i = 0; i < points.size(); i++)
					{
						const Point p(points[i]);
						mesh->add_point(p);
					}

					//! add edges to mesh
					VMesh::Node::array_type edge;

					for(size_t i = 0; i < indices.size(); i++)
					{
					  VMesh::Node::index_type p = indices[i];
					  edge.push_back(p);

					  if(edge.size() == 2)
					  {
						mesh->add_elem(edge);
						edge.clear();
					  }
					}

					//! add data to mesh

					VField* field = meshHandle->vfield();

					field->resize_values();
					field->set_values(values);
				}

		};

		//! piece-wise wire discretization
		class CircularWireCoilgen : public BaseCoilgen
		{

			public:

				CircularWireCoilgen(
					const AlgorithmBase* algo,
					ModelTMSCoilAlgorithm::Args& args )
					:
					BaseCoilgen( algo, args )
				{
					coilLayers = coilLayers == 0 ? 1 : coilLayers;
				}

				~CircularWireCoilgen()
				{
				}

				virtual void Generate(FieldHandle& meshHandle) const
				{
					std::vector<Vector> coilPoints;
					std::vector<size_t> coilIndices;
					std::vector<double> coilValues;

					Vector step(0,0,coilLayersStep);

					double dr = 0.0;
					if(rings > 1)
						dr = (outerR - innerR) / (rings - 1);

					if(coilType == 1)
					{
						//Vector origin(0, 0, -0.5*(1.0/coilLayers));
						Vector origin(0, 0, -coilLayersStep*(coilLayers/2) );

						for(size_t l = 0; l < coilLayers; l++)
						{
							for (size_t i = 0; i < rings; i++)
							{
								ClosedSegments segments(coilPoints,coilIndices,coilValues);
								GenPointsCircular(segments, origin, innerR + i*dr, current, 0.0 , 2*M_PI );
							}

							origin += step;
						}

					}
					else if(coilType == 2)
					{
						Vector originLeft ( -outerR - (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );
						Vector originRight(  outerR + (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );

						auto transLeft = Transform::Identity();
						transLeft.post_rotate(M_PI*(wingsAngle/180),{0,1,0});

						auto transRight = Transform::Identity();
						transRight.post_rotate(M_PI*(-wingsAngle/180),{0,1,0});

						for(size_t l = 0; l < coilLayers; l++)
						{

							for (size_t i = 0; i < rings; i++)
							{
								ClosedSegments segments(coilPoints,coilIndices,coilValues);
								GenPointsCircular(segments, originLeft, innerR + i*dr, current, 0.0 , 2*M_PI );
								segments.Transform(transLeft);
							}

							originLeft += step;
						}

						for(size_t l = coilLayers; l < 2*coilLayers; l++)
						{

							for (size_t i = 0; i < rings; i++)
							{
								ClosedSegments segments(coilPoints,coilIndices,coilValues);
								GenPointsCircular(segments, originRight, innerR + i*dr, -current, 0.0 , 2*M_PI);
								segments.Transform(transRight);
							}

							originRight += step;
						}

					}
					else
					{
						algo->error("coil type value expeced: 1/2 (0-shape/8-shape)");
						return;
					}

					//SCIrun API creating a new mesh
					//0 data on elements; 1 data on nodes
					FieldInformation fi("CurveMesh",0,"double");
					fi.make_curvemesh();
					fi.make_constantdata();
					fi.make_scalar();

					meshHandle = CreateField(fi);

					BuildScirunMesh(coilPoints,coilIndices,coilValues,meshHandle);
				}

		};

		//! piece-wise wire discretization
		class MultiloopsCoilgen : public BaseCoilgen
		{
			public:

				MultiloopsCoilgen(
					const AlgorithmBase* algo,
					ModelTMSCoilAlgorithm::Args args )
					:
					BaseCoilgen( algo, args )
				{
					coilLayers = coilLayers == 0 ? 1 : coilLayers;

					//no auto current adjustment for each layer
					//hidden functionality (rather be explicit up front)
					//leave it up to users judgment

				}

				~MultiloopsCoilgen()
				{
				}

				virtual void Generate(FieldHandle& meshHandle) const
				{
					std::vector<Vector> coilPoints;
					std::vector<size_t> coilIndices;
					std::vector<double> coilValues;

					Vector step(0,0,coilLayersStep);

					if(coilType == 1)
					{
						//Vector origin(0, 0, -0.5*(1.0/coilLayers));
						Vector origin(0, 0, -coilLayersStep*(coilLayers/2) );

						for(size_t l = 0; l < coilLayers; l++)
						{
							OpenSegments segments( coilPoints, coilIndices, coilValues );

							///SINGLE coil
							GenPointsSpiralLeft(segments, origin);
							origin += step;
						}

					}
					else if(coilType == 2)
					{
						Vector originLeft ( -outerR - (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );
						Vector originRight(  outerR + (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );

						for(size_t l = 0; l < coilLayers; l++)
						{
							OpenSegments segments( coilPoints, coilIndices, coilValues );

							//LEFT coil
							GenPointsSpiralLeft(segments, originLeft);

							originLeft += step;
						}

						for(size_t l = coilLayers; l < 2*coilLayers; l++)
						{
							OpenSegments segments( coilPoints, coilIndices, coilValues );

							//RIGHT coil
							GenPointsSpiralRight(segments, originRight);

							originRight += step;
						}

					}
					else
					{
						algo->error("coil type value expeced: 1/2 (0-shape/8-shape)");
						return;
					}

					//SCIrun API creating a new mesh
					//0 data on elements; 1 data on nodes
					FieldInformation fi("CurveMesh",0,"double");
					fi.make_curvemesh();
					fi.make_constantdata();
					fi.make_scalar();

					meshHandle = CreateField(fi);

					BuildScirunMesh(coilPoints,coilIndices,coilValues,meshHandle);
				}

			protected:

				void GenPointsSpiralLeft(OpenSegments& segments, Vector center) const

				{
					double dr = (outerR - innerR) / rings;

					Vector center_offset (center.x() + dr/2, center.y(), center.z() );

					for (size_t i = 0; i < rings; i++)
					{
						GenPointsCircular(segments, center, innerR + i*dr, current, 0   , M_PI, i );

						GenPointsCircular(segments, center_offset, innerR + i*dr + dr/2, current, M_PI, 2*M_PI, i );
					}

					//TODO refactor to avoid this
					Vector endp(center.x() + outerR * cos(2*M_PI), center.y() + outerR * sin(2*M_PI), center.z());
					segments.AddPoint(endp, current);
				}

				void GenPointsSpiralRight(OpenSegments& segments, Vector center) const
				{
					double dr = (outerR - innerR) / rings;

					Vector center_offset( center.x() + dr/2, center.y(), center.z() );

					for (size_t i = rings; i > 0; i--)
					{
						GenPointsCircular(segments, center, innerR + i*dr, -current, M_PI, 2*M_PI, i );

						GenPointsCircular(segments, center_offset, innerR + i*dr - dr/2, -current, 0, M_PI, i );
					}

					//TODO refactor to avoid this
					Vector endp(center.x() + innerR * cos(M_PI), center.y() + innerR * sin(M_PI), center.z());
					segments.AddPoint(endp, -current);
				}

				/// this is tricky but doable, the idea is to distribute the current along each coil winding
				/// so that the top surface flux is not linear but curved (bell shaped like)
				/// this is required since in the AC profile there is inter-winding coupling increasing the resistivity of the net
				/// please see: https://en.wikipedia.org/wiki/Proximity_effect_%28electromagnetism%29
				void AdjustForProximityEffect()
				{

				}

		};

		//! dipoles domain discretization
		// (replicating paper doi:10.1006/nimg.2002.1282)
		class DipolesCoilgen : public BaseCoilgen
		{
			public:

				DipolesCoilgen(
					const AlgorithmBase* algo,
					ModelTMSCoilAlgorithm::Args args )
					:
					BaseCoilgen( algo, args )

				{
					coilLayers = coilLayers == 0 ? 1 : coilLayers;

				}

				~DipolesCoilgen()
				{
				}

				virtual void Generate(FieldHandle& meshHandle) const
				{
					std::vector<Vector> dipolePoints;
					std::vector<Vector> dipoleValues;
					//std::vector<size_t> coilIndices;

					//std::vector<double> radiiInner = preRadiiInner();
					//std::vector<double> radiiOuter = preRadiiOuter();
					//std::vector<double> numElements = preNumElem(radiiInner);
					//std::vector<double> numCoupling = preNumAdjElem(radiiInner);

					//print_vector(radiiInner);
					//print_vector(radiiOuter);

					//print_vector(numElements);
					//print_vector(numCoupling);

					//assert(radiiInner.size() == radiiOuter.size());

					//algo->remark("#Rings:  " +  boost::lexical_cast<std::string>(radiiOuter.size()) + " ring-step:" + boost::lexical_cast<std::string>(lod_step_m));

					Vector step(0,0,coilLayersStep);

					double dr = 0.0;
					if(rings > 1)
						dr = (outerR - innerR) / (rings - 1);

					if(coilType == 1)
					{
						Vector origin(0, 0, -coilLayersStep*(coilLayers/2) );

						for(size_t l = 0; l < coilLayers; l++)
						{
							double dipoleMoment = 0;
							for (size_t i = 0; i < rings; i++)
							{
								double ringArea = M_PI * (innerR + i*dr) * (innerR + i*dr);
								dipoleMoment += current * ringArea;
							}
							Vector dipoleNormL(0,0,1.0*dipoleMoment);
							dipolePoints.push_back(origin);
							dipoleValues.push_back(dipoleNormL);

							origin += step;
						}
					}
					else if(coilType == 2)
					{
						Vector originLeft ( -outerR - (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );
						Vector originRight(  outerR + (outerD/2), 0.0, -coilLayersStep*(coilLayers/2) );

						auto transLeft = Transform::Identity();
						transLeft.post_rotate(M_PI*(wingsAngle/180),{0,1,0});

						auto transRight = Transform::Identity();
						transRight.post_rotate(M_PI*(-wingsAngle/180),{0,1,0});

						for(size_t l = 0; l < coilLayers; l++)
						{
							double dipoleMoment = 0;
							for (size_t i = 0; i < rings; i++)
							{
								double ringArea = M_PI * (innerR + i*dr) * (innerR + i*dr);
								dipoleMoment += current * ringArea;
							}
							Vector dipoleNormL(0,0,1.0*dipoleMoment);
							dipoleNormL = transLeft * dipoleNormL;
							Vector originLeftTrans = transLeft * originLeft;
							dipolePoints.push_back(originLeftTrans);
							dipoleValues.push_back(dipoleNormL);

							originLeft += step;
						}

						for(size_t l = 0; l < coilLayers; l++)
						{
							double dipoleMoment = 0;
							for (size_t i = 0; i < rings; i++)
							{
								double ringArea = M_PI * (innerR + i*dr) * (innerR + i*dr);
								dipoleMoment += current * ringArea;
							}
							Vector dipoleNormR(0,0,-1.0*dipoleMoment);
							dipoleNormR = transRight * dipoleNormR;
							Vector originRightTrans = transRight * originRight;
							dipolePoints.push_back(originRightTrans);
							dipoleValues.push_back(dipoleNormR);

							originRight += step;
						}

					}
					else
					{
						algo->error("coil type value expeced: 1/2 (0-shape/8-shape)");
						return;
					}

					///basic topoly assumptions needs to be correct
					assert(dipolePoints.size() > 0);
					assert(dipolePoints.size() == dipoleValues.size());

					///SCIrun API creating a new mesh
					///0 data on elements; 1 data on nodes
					FieldInformation fi("PointCloudMesh",1,"vector");
					fi.make_pointcloudmesh();
					fi.make_lineardata();
					fi.make_vector();

					meshHandle = CreateField(fi);

					BuildScirunMesh(dipolePoints,dipoleValues,meshHandle);
				}

		protected:

				void print_vector(const std::vector<double>& v) const
				{
					std::cout << std::endl;
					for(int i=0;i<v.size();++i)
					{
						std::cout << v[i] << " ";
					}
					std::cout << std::endl;
				}

				const std::vector<double> preRadiiInner() const
				{
					std::vector<double> preRadii;

					double step = (outerR - innerR) / rings;

					double d = innerR;

					//add first element
					//preRadii.push_back(0.00d);

					while( d < outerR)
					{
						preRadii.push_back(d);
						d += step;
					}

					//const double vals[16] = {0.00d, 0.003d, 0.007d, 0.011d, 0.015d, 0.019d, 0.023d, 0.026d, 0.028d, 0.030d, 0.032d, 0.034d, 0.036d, 0.038d, 0.040d, 0.042d};
					//std::vector<double> preRadii(vals,vals+16);
					return preRadii;
				}

				const std::vector<double> preRadiiOuter() const
				{
					std::vector<double> preRadii;

					double step = (outerR - innerR) / rings;

					double d = innerR;

					//add first element
					//preRadii.push_back(d);

					while( d < outerR)
					{
						d += step;
						preRadii.push_back(d);
					}

					//add last
					//preRadii.push_back(outerR);

					//override last to fill to exacct outer radius
					preRadii[preRadii.size()-1u] = outerR;

					//const double vals[16] = {0.003d, 0.007d, 0.011d, 0.015d, 0.019d, 0.023d, 0.026d, 0.028d, 0.030d, 0.032d, 0.034d, 0.036d, 0.038d, 0.040d, 0.042d, 0.044d};
					//std::vector<double> preRadii(vals,vals+16);
					return preRadii;
				}

				const std::vector<double> preNumElem(std::vector<double>& radii) const
				{
					std::vector<double> preNumElem;

					for(size_t i = 1; i <= radii.size(); ++i)
					{

						size_t n = M_PI_2 / (1.0 / ( radii[i]* coilLOD ) );

						//size_t n = segments +  ( pow(i,1.5) / coilLOD );
						preNumElem.push_back(n);
					}

					//const double vals[16] = {3.0d, 9.0d, 12.0d, 16.0d, 20.0d, 24.0d, 28.0d, 30.0d, 32.0d, 34.0d, 36.0d, 38.0d, 40.0d, 42.0d, 44.0d, 44.0d};
					//std::vector<double> preNumElem(vals,vals+16);
					return preNumElem;

				}

				const std::vector<double> preNumAdjElem(std::vector<double>& radii) const
				{
					std::vector<double> preNumAdjElem;

					for(size_t i = 1; i <= radii.size(); ++i)
					{
						preNumAdjElem.push_back(1.0);
					}

					//const double vals[16] = {9.0d, 9.0d, 9.0d, 9.0d, 9.0d, 9.0d, 9.0d, 9.0d, 8.0d, 7.0d, 6.0d, 5.0d, 4.0d, 3.0d, 2.0d, 1.0d};
					//std::vector<double> preNumAdjElem(vals,vals+16);
					return preNumAdjElem;
				}

				void GenSegmentValues(const std::vector<Vector>& points, std::vector<Vector>& values, Vector val) const
				{
					assert(points.size() > 0);

					for(size_t i = values.size(); i < points.size(); i++)
					{
						values.push_back(val);
					}
				}

				size_t GenPointsCircular2(
					std::vector<Vector>& points,
					Vector origin,
					double radius,
					double fromPI,
					double toPI,
					double extLOD) const
				{

					double dPI = fabs(toPI - fromPI);

					double minPI = M_PI /  ( 8.0 * coilLOD + coilLOD * extLOD );

					assert(dPI > minPI);

					size_t nsegments = 2;
					double iPI = dPI / nsegments;

					while(iPI > minPI)
					{
						nsegments++;
						iPI = dPI / nsegments;
					}

					//algo->remark("#Segments(LOD):  " +  boost::lexical_cast<std::string>(nsegments) );

					dPI = toPI - fromPI;

					iPI = dPI / nsegments;

					for(size_t i = 0; i < nsegments; i++)
					{
						Vector point(origin.x() + radius * cos(fromPI + iPI*i), origin.y() + radius * sin(fromPI + iPI*i), origin.z());
						//segments.AddPoint(point,value);
						points.push_back(point);
					}

					return nsegments;
				}

				void BuildScirunMesh(const std::vector<Vector>& points,
						const std::vector<Vector>& values,
						FieldHandle& meshHandle) const
				{

					VMesh* mesh = meshHandle->vmesh();

					//! add nodes to the new mesh
					for(size_t i = 0; i < points.size(); i++)
					{
						const Point p(points[i]);
						mesh->add_point(p);
					}

					//! add data to mesh
					VField* field = meshHandle->vfield();
					field->resize_values();
					field->set_values(values);
				}
		};

int ResolveAlgoName(std::string& name)
{
	if(name == "Thin Circular Wires")
	{
		return 0;
	}

	if(name == "Thin Spiral Wires")
	{
		return 1;
	}

	if(name == "Magnetic Dipoles")
	{
		return 2;
	}

	return 0;

}

std::unique_ptr<BaseCoilgen> AlgoSelector(int idx,const AlgorithmBase* scirunAlgoBase, ModelTMSCoilAlgorithm::Args args)
{
	switch(idx)
	{
		case 0: return std::unique_ptr<BaseCoilgen>(new CircularWireCoilgen(scirunAlgoBase,args));
		case 1: return std::unique_ptr<BaseCoilgen>(new MultiloopsCoilgen(scirunAlgoBase,args));
		case 2: return std::unique_ptr<BaseCoilgen>(new DipolesCoilgen(scirunAlgoBase,args));
		default: return std::unique_ptr<BaseCoilgen>(new CircularWireCoilgen(scirunAlgoBase,args));
	}

}

AlgorithmOutput ModelTMSCoilAlgorithm::run(const AlgorithmInput& input) const
{
	AlgorithmOutput output;
	FieldHandle ofield;

	std::string model_type = static_cast<std::string>(get(Parameters::Type).toString());
	//int model_type = static_cast<int>(get(Parameters::Type).toInt());
	//remark("model_type:  " +  boost::lexical_cast<std::string>(model_type) );

	ModelTMSCoilAlgorithm::Args algoArgs;
	algoArgs.current = static_cast<double>(get(Parameters::Current).toDouble())*1e6; ///MD: here is the unit scaling, wire current module input is in [megaA/s]
	algoArgs.rings = static_cast<size_t>(get(Parameters::Rings).toInt());
	algoArgs.wingsAngle = static_cast<double>(get(Parameters::WingsAngle).toDouble());
	algoArgs.coilRadiusInner = static_cast<double>(get(Parameters::InnerRadius).toDouble());
	algoArgs.coilRadiusOuter = static_cast<double>(get(Parameters::OuterRadius).toDouble());
	algoArgs.coilLayers = static_cast<size_t>(get(Parameters::Layers).toInt());
	algoArgs.coilLevelDetails = static_cast<size_t>(get(Parameters::LevelOfDetail).toInt());
	algoArgs.coilLayersStep = static_cast<double>(get(Parameters::LayerStepSize).toDouble());
	algoArgs.coilDistance = static_cast<double>(get(Parameters::Distance).toDouble());
	algoArgs.type = get(Parameters::FigureOf8CoilShape).toInt() ? 2 : 1;

	auto algo = AlgoSelector(ResolveAlgoName(model_type),this,algoArgs);
	//auto algo = AlgoSelector(model_type,this,algoArgs);

	algo->Execute(ofield);

	output[Parameters::Mesh] = ofield;
	return output;
}


/**************************************************************************\
|
|    Copyright (C) 2009 Marc Stevens
|
|    This program 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.
|
|    This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
|
\**************************************************************************/

#include <vector>
#include <algorithm>
#include <stdexcept>
#include <map>
#include <utility>
#include <algorithm>
#include <string>
#include <iostream>
#include <time.h>

#include <boost/lexical_cast.hpp>

#include <hashclash/saveload_gz.hpp>
#include <hashclash/md5detail.hpp>
#include <hashclash/rng.hpp>
#include <hashclash/differentialpath.hpp>
#include <hashclash/progress_display.hpp>

#include "main.hpp"

using namespace hashclash;
using namespace std;

void random_permutation(vector<differentialpath>& paths)
{
	// use a pseudo-random permutation fixed by the number of paths
	seed(paths.size());
	for (unsigned i = 0; i < paths.size(); ++i)
	{
		unsigned k = xrng64() % paths.size();
		paths[i].swap(paths[k]);
	}
	addseed(time(NULL));
}

inline std::string pathsstring(const std::string& basepath, unsigned modi, unsigned modn)
{
	return workdir +  "/"  + basepath
		+ "_" + boost::lexical_cast<std::string>(modi)
		+ "of" + boost::lexical_cast<std::string>(modn);
}

// map: (dQ1, dF1) => (count, Q1_set0, Q1_set1)
typedef map< pair<sdr,sdr>, pair<unsigned, pair<uint32, uint32> > > step1diffs_type;
void generate_step1diffs(step1diffs_type& step1diffs, const path_container_autobalance& container)
{
	const unsigned m = container.modn;
	const unsigned i = container.modi;
	const uint32* m_diff = container.m_diff;

	uint32 Q1[4], Q2[4];
	Q1[0] = container.IHV1[0]; Q1[1] = container.IHV1[3]; Q1[2] = container.IHV1[2]; Q1[3] = container.IHV1[1];
	Q2[0] = container.IHV2[0]; Q2[1] = container.IHV2[3]; Q2[2] = container.IHV2[2]; Q2[3] = container.IHV2[1];

	step1diffs_type::iterator step1diffs_it, step1diffs_it2;
	uint32 Ta = md5_ff(Q1[3], Q1[2], Q1[1]) + Q1[0] + md5_ac[0];
	uint32 Tb = md5_ff(Q2[3], Q2[2], Q2[1]) + Q2[0] + md5_ac[0] + m_diff[0];

	uint64 endcount = uint64(1)<<32;
	uint64 count = 0;
	unsigned k = 0;
	progress_display show_progress(endcount, true, cout, "t=0:  ", "      ", "      ");
	uint32 Q1a, Fa, Q1b, Fb;
	sdr dQ1, dF;
	pair<sdr,sdr> index;
	for (uint64 count = 0; count < endcount; ++count,++show_progress,++Ta,++Tb)
	{
		++k;
		if (k == m) k = 0;
		if (k != i) continue;

		Q1a = rotate_left(Ta, 7) + Q1[Qoff+0];
		Fa = md5_ff(Q1a, Q1[Qoff+0], Q1[Qoff-1]);
		Q1b = rotate_left(Tb, 7) + Q2[Qoff+0];
		Fb = md5_ff(Q1b, Q2[Qoff+0], Q2[Qoff-1]);
		dQ1.set(Q1a, Q1b);
		dF.set(Fa, Fb);

		index.first = dQ1; index.second = dF;
		step1diffs_it = step1diffs.find(index);
		if (step1diffs_it == step1diffs.end()) {
			pair<unsigned, pair<uint32, uint32> > tmp;
			tmp.first = 1;
			tmp.second.first = Q1a;
			tmp.second.second = Q1a;
			step1diffs[index] = tmp;
		} else {
			++step1diffs_it->second.first;
			step1diffs_it->second.second.first |= Q1a;
			step1diffs_it->second.second.second &= Q1a;
		}
	}
}

void combine_step1diffs(step1diffs_type& global_step1diffs, const step1diffs_type& step1diffs)
{
	step1diffs_type::const_iterator step1diffs_it;
	step1diffs_type::iterator step1diffs_it2;
	for (step1diffs_it = step1diffs.begin(); step1diffs_it != step1diffs.end(); ++step1diffs_it)
	{
		step1diffs_it2 = global_step1diffs.find(step1diffs_it->first);
		if (step1diffs_it2 == global_step1diffs.end())
			global_step1diffs.insert(*step1diffs_it);
		else {
			step1diffs_it2->second.first += step1diffs_it->second.first;
			step1diffs_it2->second.second.first |= step1diffs_it->second.second.first;
			step1diffs_it2->second.second.second &= step1diffs_it->second.second.second;
		}
	}
}

void dostep0(const path_container_autobalance& container)
{
	// map: (dQ1, dF1) => (count, Q1_set0, Q1_set1)
	map< pair<sdr,sdr>, pair<unsigned, pair<uint32, uint32> > >			step1diffs;
	map< pair<sdr,sdr>, pair<unsigned, pair<uint32, uint32> > >::iterator	step1diffs_it, step1diffs_it2;

	cout << "Searching first step differentials: " << endl;
	generate_step1diffs(step1diffs, container);
	hashclash::save_gz(step1diffs, pathsstring("paths0", container.modi, container.modn), binary_archive);
	cout << "Saved " << step1diffs.size() << " first step differentials." << endl;
}

void dostep1(const path_container_autobalance& container)
{
	const unsigned modn = container.modn;
	const unsigned modi = container.modi;

	map< pair<sdr,sdr>, pair<unsigned, pair<uint32, uint32> > >	step1diffs, step1diffs2;
	map< pair<sdr,sdr>, pair<unsigned, pair<uint32, uint32> > >::iterator	step1diffs_it, step1diffs_it2;

	if (modi != 0) return;
	for (unsigned j = 0; j < modn; ++j)
	{
		cout << "Loading " << pathsstring("paths0", j, modn) << "..." << flush;
		hashclash::load_gz(step1diffs2, pathsstring("paths0", j, modn), binary_archive);
		cout << "done." << endl;
		combine_step1diffs(step1diffs, step1diffs2);
	}

	unsigned count_max = 0;
	unsigned count_good = 0;

	step1diffs_it2 = step1diffs.begin();
	vector<unsigned> countvec;
	for (step1diffs_it = step1diffs.begin(); step1diffs_it != step1diffs.end(); ++step1diffs_it)
	{
		countvec.push_back(step1diffs_it->second.first);
		if (step1diffs_it->second.first > count_max)
			count_max = step1diffs_it->second.first;
	}
	std::sort(countvec.begin(), countvec.end());
	int countindex = int(countvec.size()) - int(65536);
	if (countindex < 0) countindex = 0;
	unsigned count_min = countvec[countindex];

	uint32 Q1[4], Q2[4];
	Q1[0] = container.IHV1[0]; Q1[1] = container.IHV1[3]; Q1[2] = container.IHV1[2]; Q1[3] = container.IHV1[1];
	Q2[0] = container.IHV2[0]; Q2[1] = container.IHV2[3]; Q2[2] = container.IHV2[2]; Q2[3] = container.IHV2[1];
	differentialpath tmp;
	for (int i = -3; i <= 0; ++i)
		tmp[i].set(Q2[3+i]-Q1[3+i], Q1[3+i], Q1[3+i]);

	vector< differentialpath > goodpaths;
	for (step1diffs_it = step1diffs.begin(); step1diffs_it != step1diffs.end(); ++step1diffs_it)
	{
		if (step1diffs_it->second.first >= count_min)
		{
			++count_good;
			uint32 dQ1 = step1diffs_it->first.first.adddiff();
			uint32 dF1 = step1diffs_it->first.second.adddiff();
			uint32 dT1 = dF1 + (Q2[Qoff-2]-Q1[Qoff-2]) + container.m_diff[1];
			uint32 dR1 = best_rotated_difference(dT1, 12);
			uint32 dQ2 = dQ1 + dR1;

			tmp[1].set( dQ1, step1diffs_it->second.second.first, step1diffs_it->second.second.second );
			tmp[2] = naf(dQ2);
			goodpaths.push_back(tmp);
		}
	}
	cout << "Saving " << goodpaths.size() << " paths..." << flush;
	save_gz(goodpaths, pathsstring("paths1", modi, modn), binary_archive);
	cout << "done." << endl;
	if (goodpaths.size() > 0)
		show_path(goodpaths[0], container.m_diff);
}

progress_display* dostep_progress = 0;
unsigned dostep_index = 0;
struct dostep_thread {
	dostep_thread(vector<differentialpath>& in, path_container_autobalance& out)
		: pathsin(in), container(out)
	{}
	vector<differentialpath>& pathsin;
	path_container_autobalance& container;
	md5_forward_thread worker;
	void operator()() {
		try {
			while (true) {
				mut.lock();
				unsigned i = dostep_index;
				unsigned iend = i + (unsigned(pathsin.size() - dostep_index)>>4);
				if (iend > i+4) iend = i+4;
				if (iend == i && i < pathsin.size())
					iend = i+1;
				if (iend != i)
					(*dostep_progress) += iend-i;
				dostep_index = iend;
				mut.unlock();
				if (i >= pathsin.size())
					break;
				for (; i < iend; ++i)
					worker.md5_forward_differential_step(pathsin[i], container);
			}
		} catch (std::exception & e) { cerr << "Worker thread: caught exception:" << endl << e.what() << endl; } catch (...) {}
	}
};
void dostep_threaded(vector<differentialpath>& in, path_container_autobalance& out)
{
	dostep_index = 0;
	std::string tstring = "t=" + boost::lexical_cast<std::string>(out.t) + ": ";
	if (tstring.size() == 5) tstring += " ";
	if (out.estimatefactor)
		dostep_progress = new progress_display(in.size(), true, cout, tstring, "      ", "e     ");
	else
		dostep_progress = new progress_display(in.size(), true, cout, tstring, "      ", "      ");
	boost::thread_group mythreads;
	for (unsigned i = 0; i < out.threads; ++i)
		mythreads.create_thread(dostep_thread(in,out));
	mythreads.join_all();
	if (dostep_progress->expected_count() != dostep_progress->count())
		*dostep_progress += dostep_progress->expected_count() - dostep_progress->count();
	delete dostep_progress;
}

vector<differentialpath> pathscache;
void dostep(path_container_autobalance& container, bool savetocache)
{
	const unsigned t = container.t;
	const unsigned modn = container.modn;
	const unsigned modi = container.modi;

	if (t == 0 && !container.normalt01) {
		dostep0(container);
		return;
	}
	if (t == 1 && !container.normalt01) {
		dostep1(container);
		return;
	}

	vector< differentialpath > pathsin, pathstmp, pathsout;
	if (pathscache.size() != 0) {
		pathsin.swap(pathscache);
		random_permutation(pathsin);
	} else if (container.inputfile.size() == 0) {
		for (unsigned k = 0; k < modn; ++k)
		{
			try {
				std::string filename = pathsstring("paths" + boost::lexical_cast<std::string>(t-1), k, modn);
				cout << "Loading " << filename << "..." << flush;
				load_gz(pathstmp, filename, binary_archive);
				random_permutation(pathstmp);
				for (unsigned j = modi; j < pathstmp.size(); j += modn)
					pathsin.push_back(pathstmp[j]);
				cout << "done: " << pathstmp.size() << " (work:" << pathsin.size() << ")." << endl;
			} catch(...) {
				cout << "failed." << endl;
			}
		}
	} else {
		bool failed = false;
		try {
			cout << "Loading " << container.inputfile << "..." << flush;
			load_gz(pathstmp, binary_archive, container.inputfile);
			random_permutation(pathstmp);
			for (unsigned j = modi; j < pathstmp.size(); j += modn)
				pathsin.push_back(pathstmp[j]);
			cout << "done: " << pathsin.size() << "." << endl;
		} catch(...) {
			failed = true;
			cout << "failed." << endl;
		}
		if (failed) {
			try {
				cout << "Loading (text) " << container.inputfile << "..." << flush;
				load_gz(pathstmp, text_archive, container.inputfile);
				random_permutation(pathstmp);
				for (unsigned j = modi; j < pathstmp.size(); j += modn)
					pathsin.push_back(pathstmp[j]);
				cout << "done: " << pathsin.size() << "." << endl;
			} catch(...) {
				cout << "failed." << endl;
			}
		}
	}
	if (container.showinputpaths)
		for (unsigned r = 0; r < pathsin.size(); ++r)
			show_path(pathsin[r], container.m_diff);

	std::string tstring = "t=" + boost::lexical_cast<std::string>(t) + ": ";
	if (tstring.size() == 5) tstring += " ";

	if (container.estimatefactor != 0) {
		cout << "Estimating maxcond for upper bound " << unsigned(double(container.ubound)*container.estimatefactor)
			<< " (=" << container.ubound << " * " << container.estimatefactor << ")..." << endl;
		dostep_threaded(pathsin, container);
//		progress_display show_progress(pathsin.size(), true, cout, tstring, "      ", "e     ");
//		for (unsigned k = 0; k < pathsin.size(); ++k,++show_progress)
//			md5_forward_differential_step(pathsin[k], container);
		container.finish_estimate();
		cout << "Found maxcond = " << container.maxcond << endl;
	}

	dostep_threaded(pathsin, container);
//	progress_display show_progress(pathsin.size(), true, cout, tstring, "      ", "      ");
//	for (unsigned k = 0; k < pathsin.size(); ++k,++show_progress)
//		md5_forward_differential_step(pathsin[k], container);

	pathstmp.swap(pathsout);
	container.export_results(pathsout);

	if (pathsout.size() > 0)
		show_path(pathsout[0], container.m_diff);
	else
		throw std::runtime_error("No valid differential paths found!");
	std::string filenameout = pathsstring("paths" + boost::lexical_cast<std::string>(t), modi, modn);
	cout << "Saving " << pathsout.size() << " paths..." << flush;
	if (savetocache)
		pathscache.swap(pathsout);
	else
		save_gz(pathsout, filenameout, binary_archive);
	cout << "done." << endl;
}


// Written by Angjoo Kanazawa & Abhishek Sharma 2013
// Implementation of Transformation Invariant Convolution Layer
// The transformation parameter are set in the transformation variable in layer
// parameters.

#include <vector>
#include <boost/lexical_cast.hpp>
#include <cstdio>

#include "caffe/layer.hpp"
#include "caffe/vision_layers.hpp"
#include "caffe/util/im2col.hpp"
#include "caffe/filler.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/imshow.hpp"

#include <opencv2/highgui/highgui.hpp>
namespace caffe {

template <typename Dtype>
void
TiedConvolutionLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype> *> &bottom,
                                        vector<Blob<Dtype> *> *top) {
  ConvolutionParameter conv_param = this->layer_param_.convolution_param();
  CHECK(!conv_param.has_kernel_size() !=
        !(conv_param.has_kernel_h() && conv_param.has_kernel_w()))
      << "Filter size is kernel_size OR kernel_h and kernel_w; not both";
  CHECK(conv_param.has_kernel_size() ||
        (conv_param.has_kernel_h() && conv_param.has_kernel_w()))
      << "For non-square filters both kernel_h and kernel_w are required.";
  CHECK((!conv_param.has_pad() && conv_param.has_pad_h() &&
         conv_param.has_pad_w()) ||
        (!conv_param.has_pad_h() && !conv_param.has_pad_w()))
      << "pad is pad OR pad_h and pad_w are required.";
  CHECK((!conv_param.has_stride() && conv_param.has_stride_h() &&
         conv_param.has_stride_w()) ||
        (!conv_param.has_stride_h() && !conv_param.has_stride_w()))
      << "Stride is stride OR stride_h and stride_w are required.";
  if (conv_param.has_kernel_size()) {
    kernel_h_ = kernel_w_ = conv_param.kernel_size();
  } else {
    kernel_h_ = conv_param.kernel_h();
    kernel_w_ = conv_param.kernel_w();
  }
  CHECK_GT(kernel_h_, 0) << "Filter dimensions cannot be zero.";
  CHECK_GT(kernel_w_, 0) << "Filter dimensions cannot be zero.";
  if (!conv_param.has_pad_h()) {
    pad_h_ = pad_w_ = conv_param.pad();
  } else {
    pad_h_ = conv_param.pad_h();
    pad_w_ = conv_param.pad_w();
  }
  if (!conv_param.has_stride_h()) {
    stride_h_ = stride_w_ = conv_param.stride();
  } else {
    stride_h_ = conv_param.stride_h();
    stride_w_ = conv_param.stride_w();
  }
  // Configure output channels and groups.
  channels_ = bottom[0]->channels();
  num_output_ = conv_param.num_output();
  CHECK_GT(num_output_, 0);
  group_ = conv_param.group();
  CHECK_EQ(channels_ % group_, 0);
  CHECK_EQ(num_output_ % group_, 0)
      << "Number of output should be multiples of group.";

  // Handle the parameters: weights and biases.
  // - blobs_[0] holds the filter weights
  // - blobs_[1] holds the biases (optional)
  bias_term_ = conv_param.bias_term();
  if (this->blobs_.size() > 0) {
    LOG(INFO) << "Skipping parameter initialization";
  } else {
    if (bias_term_) {
      this->blobs_.resize(2);
    } else {
      this->blobs_.resize(1);
    }
    // Intialize the weight
    // output channels x input channels per-group x kernel height x kernel width
    this->blobs_[0].reset(
        new Blob<Dtype>(num_output_, channels_ / group_, kernel_h_, kernel_w_));
    // fill the weights
    shared_ptr<Filler<Dtype> > weight_filler(
        GetFiller<Dtype>(conv_param.weight_filler()));
    weight_filler->Fill(this->blobs_[0].get());
    // If necessary, intiialize and fill the biases:
    // 1 x 1 x 1 x output channels.
    if (bias_term_) {
      this->blobs_[1].reset(new Blob<Dtype>(1, 1, 1, num_output_));
      shared_ptr<Filler<Dtype> > bias_filler(GetFiller<Dtype>(
          conv_param.bias_filler()));
      bias_filler->Fill(this->blobs_[1].get());
    }
  }
  // Propagate gradients to the parameters (as directed by backward pass).
  this->param_propagate_down_.resize(this->blobs_.size(), true);
};

template <typename Dtype>
void TiedConvolutionLayer<Dtype>::Reshape(const vector<Blob<Dtype> *> &bottom,
                                          vector<Blob<Dtype> *> *top) {
  num_in_ = bottom.size();           // total number of data to convolve
  num_ = bottom[0]->num();           // batch
  CHECK_EQ(bottom[0]->channels(), channels_) << "Input size incompatible with"
    " convolution kernel.";
  height_.resize(num_in_);
  width_.resize(num_in_);
  N_.resize(num_in_);

  for (int i = 0; i < num_in_; ++i) {
    CHECK_EQ(channels_, bottom[i]->channels())
        << "channels has to be the same for all bottoms";
    CHECK_EQ(num_, bottom[i]->num())
        << "batch size has to be the same for all bottoms";
    height_[i] = bottom[i]->height();
    width_[i] = bottom[i]->width();
  }
  // Prepare the matrix multiplication computation.
  M_ = num_output_ / group_;
  K_ = channels_ * kernel_h_ * kernel_w_ / group_;
  col_buffers_.resize(num_in_);
  if (bias_term_) {
    bias_multipliers_.resize(num_in_);
  }
  for (int i = 0; i < num_in_; ++i) {
    // The im2col result buffer would only hold one image at a time to avoid
    // overly large memory usage.
    int height_out = (height_[i] + 2 * pad_h_ - kernel_h_) / stride_h_ + 1;
    int width_out = (width_[i] + 2 * pad_w_ - kernel_w_) / stride_w_ + 1;
    // AJ: Conserve on memory by not making diff on col_buffer.
    this->col_buffers_[i].reset(new Blob<Dtype>(
        1, channels_ * kernel_h_ * kernel_w_, height_out, width_out, false));
    // Figure out the dimensions for individual gemms.
    N_[i] = height_out * width_out;
    (*top)[i]->Reshape(num_, num_output_, height_out, width_out);

    if (bias_term_) {
      // Set up the all ones "bias multiplier" for adding biases by BLAS
      // bias_multipliers_[i].reset(new Blob<Dtype>(1, 1, 1, N_[i], false));
      // caffe_set(N_[i], Dtype(1), bias_multipliers_[i].mutable_cpu_data());
      bias_multipliers_[i].reset(new SyncedMemory(N_[i] * sizeof(Dtype)));
      Dtype *bias_multiplier_data =
          reinterpret_cast<Dtype *>(bias_multipliers_[i]->mutable_cpu_data());
      for (int j = 0; j < N_[i]; ++j) {
        bias_multiplier_data[j] = 1.;
      }
    }
  }
}

template <typename Dtype>
void
TiedConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype> *> &bottom,
                                         vector<Blob<Dtype> *> *top) {
  const Dtype *weight = this->blobs_[0]->cpu_data();
  const int weight_offset = M_ * K_; // number of filter parameters in a group
  for (int i = 0; i < num_in_; ++i) {
    //-----Same concept as Forward_cpu of convolutionlayer-----
    const Dtype *bottom_data = bottom[i]->cpu_data();
    const int col_offset = K_ * N_[i];
    const int top_offset = M_ * N_[i];
    Dtype *top_data = (*top)[i]->mutable_cpu_data();
    Dtype *col_data = this->col_buffers_[i]->mutable_cpu_data();
    for (int n = 0; n < num_; ++n) {
      // im2col transformation: unroll input regions for filtering
      // into column matrix for multplication.
      im2col_cpu(bottom_data + bottom[i]->offset(n), channels_, height_[i],
                 width_[i], kernel_h_, kernel_w_, pad_h_, pad_w_, stride_h_,
                 stride_w_, col_data);
      // Take innerproduct for groups.
      for (int g = 0; g < group_; ++g) {
        caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, M_, N_[i], K_,
                              (Dtype)1., weight + weight_offset * g,
                              col_data + col_offset * g, (Dtype)0.,
                              top_data + (*top)[i]->offset(n) + top_offset * g);
      }
      // Add bias.
      if (bias_term_) {
        caffe_cpu_gemm<Dtype>(
            CblasNoTrans, CblasNoTrans, num_output_, N_[i], 1, (Dtype)1.,
            this->blobs_[1]->cpu_data(),
            reinterpret_cast<const Dtype *>(bias_multipliers_[i]->cpu_data()),
            (Dtype)1., top_data + (*top)[i]->offset(n));
      }
    }
    //---------------------------------------------------------
  }
}

template <typename Dtype>
void
TiedConvolutionLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype> *> &bottom,
                                         vector<Blob<Dtype> *> *top) {

  const Dtype *weight = this->blobs_[0]->gpu_data();
  const int weight_offset = M_ * K_;
  for (int i = 0; i < num_in_; ++i) {
    //-----Same concept as Forward_gpu of convolutionlayer-----
    const Dtype *bottom_data = bottom[i]->gpu_data();
    const int col_offset = K_ * N_[i];
    const int top_offset = M_ * N_[i];
    Dtype *top_data = (*top)[i]->mutable_gpu_data();
    Dtype *col_data = this->col_buffers_[i]->mutable_gpu_data();
    for (int n = 0; n < num_; ++n) {
      // First, im2col
      im2col_gpu(bottom_data + bottom[i]->offset(n), channels_, height_[i],
                 width_[i], kernel_h_, kernel_w_, pad_h_, pad_w_, stride_h_,
                 stride_w_, col_data);
      // Second, innerproduct with groups.
      for (int g = 0; g < group_; ++g) {
        caffe_gpu_gemm<Dtype>(CblasNoTrans, CblasNoTrans, M_, N_[i], K_,
                              (Dtype)1., weight + weight_offset * g,
                              col_data + col_offset * g, (Dtype)0.,
                              top_data + (*top)[i]->offset(n) + top_offset * g);
      }
      // third, add bias
      if (bias_term_) {
        caffe_gpu_gemm<Dtype>(
            CblasNoTrans, CblasNoTrans, num_output_, N_[i], 1, (Dtype)1.,
            this->blobs_[1]->gpu_data(),
            reinterpret_cast<const Dtype *>(bias_multipliers_[i]->gpu_data()),
            (Dtype)1., top_data + (*top)[i]->offset(n));
      }
    }
    //---------------------------------------------------------
  }
  // montage(this->blobs_[0].get(), "tconv" +
  // boost::lexical_cast<std::string>(M_));
}

template <typename Dtype>
void
TiedConvolutionLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype> *> &top,
                                          const vector<bool> &propagate_down,
                                          vector<Blob<Dtype> *> *bottom) {
  //-----Same concept as Backward_cpu of convolutionlayer-----
  // but multiple times for each bottom-top pair, and accumulating dW
  const Dtype *weight = NULL;
  Dtype *weight_diff = NULL;
  if (this->param_propagate_down_[0]) {
    weight = this->blobs_[0]->cpu_data();
    weight_diff = this->blobs_[0]->mutable_cpu_diff();
    // Init weight diff to all 0s.
    caffe_set(this->blobs_[0]->count(), Dtype(0), weight_diff);
  }
  // bias gradient if necessary
  Dtype *bias_diff = NULL;
  if (bias_term_ && this->param_propagate_down_[1]) {
    bias_diff = this->blobs_[1]->mutable_cpu_diff();
    caffe_set(this->blobs_[1]->count(), Dtype(0), bias_diff);
  }

  const int weight_offset = M_ * K_;
  for (int i = 0; i < num_in_; ++i) {
    const Dtype *top_diff = NULL;
    // Bias gradient if necessary.
    if (bias_term_ && this->param_propagate_down_[1]) {
      top_diff = top[i]->cpu_diff();
      for (int n = 0; n < num_; ++n) {
        caffe_cpu_gemv<Dtype>(
            CblasNoTrans, num_output_, N_[i], 1., top_diff + top[i]->offset(n),
            reinterpret_cast<const Dtype *>(bias_multipliers_[i]->cpu_data()),
            1., bias_diff);
      }
    }
    if (this->param_propagate_down_[0] || propagate_down[i]) {
      if (!top_diff) {
        top_diff = top[i]->cpu_diff();
      }
      Dtype* col_data = this->col_buffers_[i]->mutable_cpu_data();
      const Dtype* bottom_data = (*bottom)[i]->cpu_data();
      Dtype* bottom_diff = (*bottom)[i]->mutable_cpu_diff();

      const int col_offset = K_ * N_[i];
      const int top_offset = M_ * N_[i];
      for (int n = 0; n < num_; ++n) {
	// Since we saved memory in the forward pass by not storing all col
	// data, we will need to recompute them.
	im2col_cpu(bottom_data + (*bottom)[i]->offset(n), channels_, height_[i],
		   width_[i], kernel_h_, kernel_w_, pad_h_, pad_w_, stride_h_,
		   stride_w_, col_data);
	// gradient w.r.t. weight. Note that we will accumulate diffs.
	// AJ: propagate error Delta W_ij = error from above * this_activation^T
        if (this->param_propagate_down_[0]) {
	  for (int g = 0; g < group_; ++g) {
	    caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasTrans, M_, K_, N_[i],
				  (Dtype)1.,
				  top_diff + top[i]->offset(n) + top_offset * g,
				  col_data + col_offset * g, (Dtype)1.,
				  weight_diff + weight_offset * g);
	  }
	}
	// gradient w.r.t. bottom data, if necessary
	// AJ: error here = W*error from above
	if (propagate_down[i]) {
          if (weight == NULL) {
            weight = this->blobs_[0]->cpu_data();
          }
	  for (int g = 0; g < group_; ++g) {
	    caffe_cpu_gemm<Dtype>(CblasTrans, CblasNoTrans, K_, N_[i], M_,
				  (Dtype)1., weight + weight_offset * g,
				  top_diff + top[i]->offset(n) + top_offset * g,
				  (Dtype)0., col_data + col_offset * g);
	  }
	  // col2im back to the data
	  col2im_cpu(col_data, channels_, height_[i], width_[i], kernel_h_,
		     kernel_w_, pad_h_, pad_w_, stride_h_, stride_w_,
		     bottom_diff + (*bottom)[i]->offset(n));
	}
      }
    }
    //---------------------------------------------------------
  }
}

template <typename Dtype>
void
TiedConvolutionLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype> *> &top,
                                          const vector<bool> &propagate_down,
                                          vector<Blob<Dtype> *> *bottom) {
  const Dtype *weight = NULL;
  Dtype *weight_diff = NULL;
  if (this->param_propagate_down_[0]) {
    weight = this->blobs_[0]->gpu_data();
    weight_diff = this->blobs_[0]->mutable_gpu_diff();
    // Init weight diffs to all 0s.
    caffe_gpu_set(this->blobs_[0]->count(), Dtype(0), weight_diff);
  }
  Dtype *bias_diff = NULL;
  if (bias_term_ && this->param_propagate_down_[1]) {
    bias_diff = this->blobs_[1]->mutable_gpu_diff();
    caffe_gpu_set(this->blobs_[1]->count(), Dtype(0), bias_diff);
  }

  const int weight_offset = M_ * K_;
  for (int i = 0; i < num_in_; ++i) {
    //-----Same concept as Backward_cpu of convolutionlayer-----
    const Dtype* top_diff = NULL;
    // Bias gradient if necessary
    if (bias_term_ && this->param_propagate_down_[1]) {
      top_diff = top[i]->gpu_diff();
      for (int n = 0; n < num_; ++n) {
        caffe_gpu_gemv<Dtype>(
            CblasNoTrans, num_output_, N_[i], 1., top_diff + top[i]->offset(n),
            reinterpret_cast<const Dtype *>(bias_multipliers_[i]->gpu_data()),
            1., bias_diff);
      }
    }
    if (this->param_propagate_down_[0] || propagate_down[i]) {
      if (!top_diff) {
        top_diff = top[i]->gpu_diff();
      }
      Dtype* col_data = this->col_buffers_[i]->mutable_gpu_data();
      const Dtype* bottom_data = (*bottom)[i]->gpu_data();
      Dtype* bottom_diff = (*bottom)[i]->mutable_gpu_diff();

      const int col_offset = K_ * N_[i];
      const int top_offset = M_ * N_[i];
      for (int n = 0; n < num_; ++n) {
	// Since we saved memory in the forward pass by not storing all col data,
	// we will need to recompute them.
	im2col_gpu(bottom_data + (*bottom)[i]->offset(n), channels_, height_[i],
                   width_[i], kernel_h_, kernel_w_, pad_h_, pad_w_,
                   stride_h_, stride_w_, col_data);
	// gradient w.r.t. weight. Note that we will accumulate diffs.
        if (this->param_propagate_down_[0]) {
	  for (int g = 0; g < group_; ++g) {
	    caffe_gpu_gemm<Dtype>(CblasNoTrans, CblasTrans, M_, K_, N_[i],
				  (Dtype)1.,
				  top_diff + top[i]->offset(n) + top_offset * g,
				  col_data + col_offset * g, (Dtype)1.,
				  weight_diff + weight_offset * g);
	  }
	}
	// gradient w.r.t. bottom data, if necessary
	if (propagate_down[i]) {
          if (weight == NULL) {
            weight = this->blobs_[0]->gpu_data();
          }
	  for (int g = 0; g < group_; ++g) {
	    caffe_gpu_gemm<Dtype>(CblasTrans, CblasNoTrans, K_, N_[i], M_,
				  (Dtype)1., weight + weight_offset * g,
				  top_diff + top[i]->offset(n) + top_offset * g,
				  (Dtype)0., col_data + col_offset * g);
	  }
	  // col2im back to the data
	  col2im_gpu(col_data, channels_, height_[i], width_[i], kernel_h_,
		     kernel_w_, pad_h_, pad_w_, stride_h_, stride_w_,
		     bottom_diff + (*bottom)[i]->offset(n));
	}
      }
    }
    // montage_channels(this->blobs_[0].get(),
    // boost::lexical_cast<std::string>(M_) + " tconv bprop " +
    // boost::lexical_cast<std::string>(i) , true);
    //// make sure to give back the pointer to gpu after visualization
    // weight_diff = this->blobs_[0]->mutable_gpu_diff();
  } // end for each input
    // montage_channels(this->blobs_[0].get(), "final tconv bprop " +
    // boost::lexical_cast<std::string>(M_), true);
    // cv::waitKey(0);
}

template <typename Dtype>
void TiedConvolutionLayer<Dtype>::Report(const std::string &name) {
  // montage(this->blobs_[0].get(), name + " tconv");
  // montage(this->blobs_[0].get(), name + " tconv bprop", true);
  // // cv::waitKey(0);
  // fprintf(stderr, "%s tied-conv W_diff:\n", name.c_str());
  // Blob<Dtype>* W = this->blobs_[0].get();
  // for (int h = 0; h < W->height(); ++h) {
  //   for (int w = 0; w < W->width(); ++w) {
  //     fprintf(stderr, "%.2g ", (W->diff_at(0, 0, h, w)));
  //   }
  //   fprintf(stderr, "\n");
  // }
}

INSTANTIATE_CLASS(TiedConvolutionLayer);
} // namespace caffe