#include "gpklib/assert.h"
#include "nltrans.h"


nltransform::nltransform(C tensor222& n, C matrix22&w, C vector2& s,
							C vector2& c)
: nl(n), warp(w), shift(s), center(c)
{}

nltransform::nltransform(C vector2& init_point, C vector2& final_point,
				C matrix22& w, C tensor222& n)
: nl(n), warp(w), shift(final_point), center(init_point)
{ }

nltransform nltransform::identity(C vector2& ctr)
{
 return nltransform(tensor222(), matrix22::identity(), ctr, ctr);
}

matrix22 nltransform::Jacobian(C vector2& x) C
{
 C vector2 X = x-center;
 assert(vector_space::vsok(X.space()));
 C matrix22 nltmp(2*nl(0,0,0)*X(0)+(nl(0,0,1)+nl(0,1,0))*X(1),
 			    	2*nl(0,1,1)*X(1)+(nl(0,0,1)+nl(0,1,0))*X(0),
 					2*nl(1,0,0)*X(0)+(nl(1,0,1)+nl(1,1,0))*X(1),
 			    	2*nl(1,1,1)*X(1)+(nl(1,0,1)+nl(1,1,0))*X(0));
 				 
 return nltmp + warp;
}

vector2 nltransform::operator()(C vector2& x) C
{
 C vector2 X = x-center;
 return (nl*X + warp)*X + shift;
}

nltransform nltransform::operator()(C nltransform& a) C
{
 assert(vector_space::vseQ(center.space(), a.Shift().space()));
 assert(vector_space::vseQ(center.space(), warp.in));
 assert(vector_space::vseQ(shift.space(), warp.out));
 C vector2 p = a.Shift() - center;
 C tensor222 nn = (nl*p)*a.Nl() + mid(nl, a.Warp())*a.Warp()
 						+ mid(nl, p)*a.Nl() + warp*a.Nl();
 C matrix22 ww = (nl*p)*a.Warp() + mid(nl, p)*a.Warp() + warp*a.Warp();
 C vector2 ss = (nl*p)*p + warp*p + shift;

 return nltransform(nn, ww, ss, a.Center()
					); /* ... plus higher order terms. */
}
#if 0
nltransform nltransform::operator()(C nltransform& a) C
{
 C vector2 p = a.Shift() - center;
 C vector2 A1 = a.Nl()*a.Center()*a.Center() - a.Warp()*a.Center() + p;
 C matrix22 B1 = a.Warp() - a.Nl()*a.Center() - mid(a.Nl(), a.Center());
 C tensor222 C1 = a.Nl();

 C tensor222 nn = (nl*A1)*C1 + mid(nl, B1)*B1 + mid(nl, A1)*C1 + warp*C1;
 C matrix22 ww = (nl*A1)*B1 + mid(nl, A1)*B1 + warp*B1;
 C vector2 ss = (nl*A1)*A1 + warp*A1 + shift;

 C vector2 nc = inverse(a.linearize())(0.5*center + 0.5*a.Shift());
 /* cout << "nc = "<< nc << "\n"; */

 C matrix22 www = ww + nn*nc + mid(nn, nc);

 return nltransform(nn, www,
					ss - (nn*nc)*nc + www*nc,
					nc
					); /* ... plus higher order terms, and debris from HOT. */
}
#endif

nltransform nltransform::operator()(C ltransform& aa) C
{
 assert(vector_space::vseQ(center.space(), aa.Shift().space()));
 C vector2 p = aa.Shift() - center;
 assert(vector_space::vsok(p.space()));
 C vector2 c0 = ::inverse(aa)(center);
 C tensor222 nn = mid(nl*aa.Warp(), aa.Warp());
 C matrix22 ww = warp*aa.Warp() + mid(nl, p)*aa.Warp() + mid(nl, aa.Warp())*p
					 	+ nn*c0 + mid(nn,c0);
 C vector2 ss = nl*p*p + warp*p + shift - nn*c0*c0 + ww*c0;

 return nltransform(nn, ww, ss, c0);
}

nltransform sqrt(C nltransform& x)
{
 /* cout << "sqrt of " << x << "\n"; */
 C matrix22 iw = inverse(sqrt(0.5*matrix22::identity() + 0.5*x.Warp().ns()));
 int i;

 vector2 p = 0.5*(x.Shift().ns() - x.Center().ns());
 tensor222 n = 0.5*iw*mid(x.Nl(), iw)*iw;
 matrix22 w = sqrt(x.Warp()).ns();
 for(i=0;i<20;i++)	{
	C tensor222 te = x.Nl() - (mid(n, x.Warp())*x.Warp() + n*p*n + mid(n, p)*n
									+ x.Warp()*n);
	/* cout << "T error= " << te << "\n"; */
	n += 0.5*iw*mid(te, iw)*iw;
	C matrix22 me = x.Warp() - (n*p*w + mid(n,w)*p + w*w);
	/* cout << "M error = " << me << "\n"; */
	w += 0.5*me*inverse(n*p + mid(n,p) + w);
	C vector2 pe = x.Shift().ns() - x.Center().ns()
						- (n*p + w + matrix22::identity())*p;
	p += inverse(w + matrix22::identity())*pe;
 	}

 nltransform rv(n, w, p+x.Center().ns(), x.Center());
 return rv;
}


nltransform inverse(C nltransform& x)
{

 assert(vector_space::vseQ(x.Warp().in, x.Center().space()));

 assert(vector_space::vseQ(x.Warp().out, x.Shift().space()));

 C ltransform llt = inverse(x.linearize());
 C matrix22 iw1 = llt.Warp();

 return nltransform(-iw1*(x.Nl()*iw1)*iw1, llt, x.Shift());
}


nltransform::nltransform(C tensor222& n, C ltransform& t, C vector2& c)
: nl(n), warp(t.Warp()), shift(t(c)), center(c)
{}

ltransform nltransform::linearize() C
{
 assert(vector_space::vsok(vector_space::vseq(Warp().in, Center().space())));
 assert(vector_space::vsok(vector_space::vseq(Warp().out, Shift().space())));
 return ltransform(Center(), Shift(), Warp());
}

#if 0
ostream& operator<<(ostream& s, C nltransform& t)
{
 return s << "{{" << t.shift << "+" << t.warp << "*+"
 			<< t.nl << "**}(x-" << t.center << ")}";
}
#endif

nline nltransform::operator()(C line& l) C
{
 return nline((*this)(l.base()),
				Warp()*l.dir() + (Nl()*l.dir())*(l.base()-Center())
							     + (Nl() * (l.base()-Center()) ) * l.dir(),
				(Nl()*l.dir())*l.dir()
				);
}

nltransform& nltransform::recenter(C vector2& newcenter)
{
 C vector2 chg = newcenter-center;
 warp += nl*chg + mid(nl, chg);
 shift += warp*chg + (nl*center)*center - (nl*newcenter)*newcenter;
 center = newcenter;
 return *this;
}

nltransform::nltransform(C nltransform& a)
: transform(), nl(a.nl), warp(a.warp), shift(a.shift), center(a.center)
{ }

nltransform::nltransform()
: transform(), nl(),  warp(), shift(), center()
{ }

nltransform operator*(C ltransform& t, C nltransform& a)
{
 assert(vector_space::vseQ(t.Center().space(), a.Shift().space()));
 assert(vector_space::vseQ(t.Center().space(), t.Warp().in));
 assert(vector_space::vseQ(t.Shift().space(), t.Warp().out));
 C vector2 p = a.Shift() - t.Center();
 C tensor222 nn = t.Warp()*a.Nl();
 C matrix22 ww = t.Warp()*a.Warp();
 C vector2 ss = t.Warp()*p + t.Shift();

 return nltransform(nn, ww, ss, a.Center());
}

nltransform& nltransform::operator=(C nltransform& i)
		{nl=i.nl; warp=i.warp; shift=i.shift; center=i.center; return *this;}

vector2 nltransform::inverse(C vector2& x) C
{
 int loops=0;
 double d = 0;
 vector2 y = ::inverse(linearize())(x);
 /* cout << "linearized inverse of "<<x<<" is "<<y<<"\n"; */
 do	{
	/* cout << "result is " << (*this)(y) << "\n"; */
 	C vector2 chg = ::inverse(Jacobian(y))*(x-(*this)(y));
	/* cout << "change is " << chg << "for a dist= " << dist(chg) << "\n"; */
 	d = dist(chg);
 	y += chg;
	/* cout << "next y is " << y << "\n"; */
	} while(loops++<10 && d>1e-6);
 assert(loops<=10);
 return y;
}

BOOL nltransform::ok() C
{
 return nl.ok() && warp.ok() && shift.ok() && center.ok()
		&& vector_space::vseQ(center.space(), warp.in)
		&& vector_space::vseQ(shift.space(), warp.out);
}

void nltransform::set_in_space(C vector_space *i)
{ center.set_space(i); warp.in=i;}

void nltransform::set_out_space(C vector_space *o)
{ shift.set_space(o); warp.out=o;}

nltransform::~nltransform() {}
