#include <iostream>

#include "LSHDERIVS.h"
#include "LOCALSUPERHELIX.h"

MATRIX3 asCross(const VEC3& v)
{
    MATRIX3 cr;
    cr << 0, -v[2], v[1],
          v[2], 0, -v[0],
          -v[1], v[0], 0;
    return cr;
}

MATRIX3 DdarbEqnDom(const VEC3& globalDarb, const VEC3& base, const VEC3& t0, const double& s)
{
    MATRIX3 derivFinal;
    const double om = globalDarb.norm();
    const VEC3 nDarb = globalDarb/om;
    const MATRIX3 nDarbOuter = nDarb * nDarb.transpose();

    const double om2 = om*om;
    const double sinOmS = sin(om * s);
    const double cosOmS = cos(om * s);
    const double tau = nDarb.dot(t0) ;


    // The linear term (mult rule)
    derivFinal = nDarb / om * s * t0.transpose();
    derivFinal += tau * s * 1/om * (MATRIX3::Identity() - 2 * nDarbOuter);

    // std::cout << "derivFinal after linear term: \n" << derivFinal << std::endl;

    // The sin term (mult rule two termed)
    derivFinal += (t0 - tau * nDarb) * 1/om * (s * cosOmS - sinOmS/om) * nDarb.transpose();
    derivFinal += sinOmS/om * -1/om * (nDarb * t0.transpose() + tau * (MATRIX3::Identity() - 2 * nDarbOuter));

    // std::cout << "derivFinal after sin term: \n" << derivFinal << std::endl;

    // The cos term (mult rule two termed)
    derivFinal += globalDarb.cross(t0) * 1/om2 * (s * sinOmS - (1 - cosOmS) * 2/om) * nDarb.transpose();
    derivFinal += (1 - cosOmS)/om2 * -asCross(t0);

    // std::cout << "derivFinal after cos term: \n" << derivFinal << std::endl;
    
    return derivFinal;
}
MATRIX3 Dlsh(const VEC3& lDarb, const VEC3& base, const MATRIX3& F0, const double& s)
{
    return DdarbEqnDom(F0 * lDarb, base, F0.col(0), s) * F0;
}
MATRIX3 DdarbEqnDt(const VEC3& globalDarb, const VEC3& t0, const double& s)
{
    MATRIX3 derivFinal;
    const double om = globalDarb.norm();
    const VEC3 nDarb = globalDarb/om;
    const MATRIX3 nDarbOuter = nDarb * nDarb.transpose();

    const double om2 = om*om;
    const double sinOmS = sin(om * s);
    const double cosOmS = cos(om * s);

    // linear term
    derivFinal = nDarb * nDarb.transpose() * s;

    // sin term
    derivFinal += sinOmS / om * (MATRIX3::Identity() - nDarb * nDarb.transpose());

    // cos term
    derivFinal += (1 - cosOmS) / om * asCross(nDarb);

    return derivFinal;
}

MATRIX3 DFdarbEqnDom(const VEC3& globalDarb, const VEC3& t0, const double& s)
{
    MATRIX3 derivFinal;
    const double om = globalDarb.norm();
    const VEC3 nDarb = globalDarb/om;
    const MATRIX3 nDarbOuter = nDarb * nDarb.transpose();

    const double sinOmS = sin(om * s);
    const double cosOmS = cos(om * s);
    const double tau = nDarb.dot(t0) ;


    // The linear term (mult rule)
    derivFinal = nDarb / om * t0.transpose();
    derivFinal += tau/om * (MATRIX3::Identity() - 2 * nDarbOuter);

    // std::cout << "derivFinal after linear term: \n" << derivFinal << std::endl;

    // The sin term (mult rule two termed)
    derivFinal += (t0 - tau * nDarb) * (-s * sinOmS) * nDarb.transpose();
    derivFinal += cosOmS * -1/om * (nDarb * t0.transpose() + tau * (MATRIX3::Identity() - 2 * nDarbOuter));

    // std::cout << "derivFinal after sin term: \n" << derivFinal << std::endl;

    // The cos term (mult rule two termed)
    derivFinal += nDarb.cross(t0) / om * (s * om * cosOmS - sinOmS) * nDarb.transpose();
    derivFinal += sinOmS / om * -asCross(t0);

    // std::cout << "derivFinal after cos term: \n" << derivFinal << std::endl;
    
    return derivFinal;
}
MATRIX3 DFlsh(const VEC3& lDarb, const MATRIX3& F0, const double& s, const int col)
{
    return DFdarbEqnDom(F0 * lDarb, F0.col(col), s) * F0;
}
MATRIX3 DFdarbEqnDt(const VEC3& globalDarb, const VEC3& t0, const double& s)
{
    MATRIX3 derivFinal;
    const double om = globalDarb.norm();
    const VEC3 nDarb = globalDarb/om;
    const MATRIX3 nDarbOuter = nDarb * nDarb.transpose();

    const double om2 = om*om;
    const double sinOmS = sin(om * s);
    const double cosOmS = cos(om * s);

    // linear term
    derivFinal = nDarb * nDarb.transpose();

    // sin term
    derivFinal += cosOmS * (MATRIX3::Identity() - nDarb * nDarb.transpose());

    // cos term
    derivFinal += sinOmS * asCross(nDarb);

    return derivFinal;
}

void fillFrames(const std::vector<superHelix>& lshs, std::vector<MATRIX3>& frames, const MATRIX3& F0)
{
    frames.resize(lshs.size());
    frames[0] = F0;
    for(unsigned int i = 1; i < lshs.size(); i++)
    {
        const superHelix lsh = lshs[i - 1];
        frames[i] = FsampleLocalSupHelix(lsh, frames[i-1], lsh.sQ);
    }
}
void fillRootsAndFrames(const std::vector<superHelix>& lshs, std::vector<VEC3>& roots, std::vector<MATRIX3>& frames, const VEC3& p0, const MATRIX3& F0)
{
    frames.resize(lshs.size());
    roots.resize(lshs.size());

    roots[0] = p0;
    frames[0] = F0;

    for(unsigned int i = 1; i < lshs.size(); i++)
    {
        const superHelix lsh = lshs[i - 1];
        roots[i] = sampleLocalSupHelix(lsh, roots[i-1], frames[i-1], lsh.sQ);
        frames[i] = FsampleLocalSupHelix(lsh, frames[i-1], lsh.sQ);
    }
}

MATRIX3 D2aliDsDqj(const std::vector<superHelix>& lshs, const MATRIX3& F0, const double s, const int i, const int j, const int col)
{
    const int nLSHS = lshs.size();
    assert(i <= nLSHS);
    assert(j >= 0);
    assert(col >=0 && col <=2);

    // base case: j > i: variations on a future superhelix does not impact past superhelix final tangent 
    if(j > i) return MATRIX3::Zero();

    std::vector<MATRIX3> frames;
    fillFrames(lshs, frames, F0);

    // base case: j == i: 
    if(j == i) return DFlsh(lshs[i].darboux, frames[i], s, col);

    // recursion case: 
    // speedup: cache these results (they get overcalled a BUNCH, eh)
    std::vector<MATRIX3> tuvDerivs(3);
    tuvDerivs[0] = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 0);
    tuvDerivs[1] = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 1);
    tuvDerivs[2] = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 2);

    //( Dframe recursion GOES HERE);
    MATRIX3 D2aiDsDqj = DFdarbEqnDom(frames[i] * lshs[i].darboux, frames[i].col(col), s) * 
    (tuvDerivs[0] * lshs[i].darboux[0] + tuvDerivs[1] * lshs[i].darboux[1] + tuvDerivs[2] * lshs[i].darboux[2]);

    //( Dtangent recursion)
    D2aiDsDqj += DFdarbEqnDt(frames[i] * lshs[i].darboux, frames[i].col(col), s) * tuvDerivs[col];

    return D2aiDsDqj;
}

MATRIX3 DaliDqj(const std::vector<superHelix>& lshs, const VEC3& p0, const MATRIX3& F0, const double s, const int i, const int j)
{
    const int nLSHS = lshs.size();
    assert(i <= nLSHS);
    assert(j >= 0);

    // base case: j > i: providing variations on a future superhelix does not impact the past superhelix's shape in any way
    if(j > i) return MATRIX3::Zero();

    std::vector<MATRIX3> frames;
    std::vector<VEC3> roots;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);
    // base case: j == i: no need to do a bunch of recursing because all the variables are laid bare.
    if (j == i) return Dlsh(lshs[i].darboux, roots[i], frames[i], s);

    // speedup opportunity: cache these results
    const MATRIX3 tDeriv = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 0);
    const MATRIX3 uDeriv = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 1);
    const MATRIX3 vDeriv = D2aliDsDqj(lshs, F0, lshs[i-1].sQ, i-1, j, 2);

    // recursion case: OH GOSH 
    MATRIX3 DaiDqj = DaliDqj(lshs, p0, F0, lshs[i - 1].sQ, i - 1, j);

    //( Dframe recursion GOES HERE);
    DaiDqj += DdarbEqnDom(frames[i] * lshs[i].darboux, roots[i], frames[i].col(0), s) * 
    (tDeriv * lshs[i].darboux[0] + uDeriv * lshs[i].darboux[1] + vDeriv * lshs[i].darboux[2]);

    //( Dtangent recursion)
    DaiDqj += DdarbEqnDt(frames[i] * lshs[i].darboux, frames[i].col(0), s) * tDeriv;

    return DaiDqj;
}

void MakeD2aliDsDqArrays(const std::vector<superHelix>& lshs, const std::vector<MATRIX3>& frames, const int i,
                         MATRIX& tHolder, MATRIX& uHolder, MATRIX& vHolder)
{
    int rows = i + 1;

    tHolder.resize(3*rows, 3*rows);
    uHolder.resize(3*rows, 3*rows);
    vHolder.resize(3*rows, 3*rows);

    tHolder.setZero();
    uHolder.setZero();
    vHolder.setZero();

    for(int r = 0; r < rows; r++)
    {
        for(int k = 0; k < r + 1; k++)
        {
            if(k == r) 
            {
                tHolder.block<3,3>(3*r, 3*k) = DFlsh(lshs[r].darboux, frames[r], lshs[r].sQ, 0);
                uHolder.block<3,3>(3*r, 3*k) = DFlsh(lshs[r].darboux, frames[r], lshs[r].sQ, 1);
                vHolder.block<3,3>(3*r, 3*k) = DFlsh(lshs[r].darboux, frames[r], lshs[r].sQ, 2);
                continue;
            }
            const MATRIX3 tRefDeriv = tHolder.block<3,3>(3*(r - 1), 3 * k);
            const MATRIX3 uRefDeriv = uHolder.block<3,3>(3*(r - 1), 3 * k);
            const MATRIX3 vRefDeriv = vHolder.block<3,3>(3*(r - 1), 3 * k);

            const MATRIX3 frameDerivContrib = lshs[r].darboux[0] * tRefDeriv + lshs[r].darboux[1] * uRefDeriv + lshs[r].darboux[2] * vRefDeriv;
            const VEC3 globalDarb = frames[r] * lshs[r].darboux;
                
            tHolder.block<3,3>(3*r, 3*k) = DFdarbEqnDom(globalDarb, frames[r].col(0), lshs[r].sQ) * frameDerivContrib;
            tHolder.block<3,3>(3*r, 3*k) += DFdarbEqnDt(globalDarb, frames[r].col(0), lshs[r].sQ) * tRefDeriv;

            uHolder.block<3,3>(3*r, 3*k) = DFdarbEqnDom(globalDarb, frames[r].col(1), lshs[r].sQ) * frameDerivContrib;
            uHolder.block<3,3>(3*r, 3*k) += DFdarbEqnDt(globalDarb, frames[r].col(1), lshs[r].sQ) * uRefDeriv;
            
            vHolder.block<3,3>(3*r, 3*k) = DFdarbEqnDom(globalDarb, frames[r].col(2), lshs[r].sQ) * frameDerivContrib;
            vHolder.block<3,3>(3*r, 3*k) += DFdarbEqnDt(globalDarb, frames[r].col(2), lshs[r].sQ) * vRefDeriv;
        }
    }
}

void MakeDaliDqArrays(const std::vector<superHelix>& lshs, const std::vector<VEC3>& roots, const std::vector<MATRIX3>& frames, const int i, 
                      const MATRIX& tHolder, const MATRIX& uHolder, const MATRIX& vHolder, MATRIX& difHolder)
{
    const int rows = i + 1;
    difHolder.resize(3*rows, 3*rows);
    difHolder.setZero();

    for(int r = 0; r < rows; r++)
    {
        for(int k = 0; k <= r; k++)
        {
            if(k == r) 
            {
                difHolder.block<3,3>(3*r, 3*k) = Dlsh(lshs[r].darboux, roots[r], frames[r], lshs[r].sQ);
                continue;
            }
            const MATRIX3 tRefDeriv = tHolder.block<3,3>(3*(r - 1), 3 * k);
            const MATRIX3 uRefDeriv = uHolder.block<3,3>(3*(r - 1), 3 * k);
            const MATRIX3 vRefDeriv = vHolder.block<3,3>(3*(r - 1), 3 * k);

            const MATRIX3 frameDerivContrib = lshs[r].darboux[0] * tRefDeriv + lshs[r].darboux[1] * uRefDeriv + lshs[r].darboux[2] * vRefDeriv;
            const VEC3 globalDarb = frames[r] * lshs[r].darboux;
                
            difHolder.block<3,3>(3*r, 3*k) = difHolder.block<3,3>(3*(r - 1), 3 * k);
            difHolder.block<3,3>(3*r, 3*k) += DdarbEqnDom(globalDarb, roots[r], frames[r].col(0), lshs[r].sQ) * frameDerivContrib;
            difHolder.block<3,3>(3*r, 3*k) += DdarbEqnDt(globalDarb, frames[r].col(0), lshs[r].sQ) * tRefDeriv;
        }
    }
}

MATRIX fullConnectedStateDeriv(const std::vector<superHelix>& lshs, const VEC3& p0, const MATRIX3& F0, const double s, const int i)
{
    const int nLSHs = lshs.size();
    assert(i <= nLSHs);

    std::vector<MATRIX3> frames;
    std::vector<VEC3> roots;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);

    MATRIX fullJacob(3, 3 * nLSHs);
    fullJacob.setZero();

    if(i == 0)
    {
        fullJacob.block<3,3>(0,0) = Dlsh(lshs[i].darboux, roots[i], frames[i], s);
    }

    // get the big templating blocks for t, u, v
    MATRIX tHolder, uHolder, vHolder;
    MakeD2aliDsDqArrays(lshs, frames, i - 1, tHolder, uHolder, vHolder);    
    MATRIX DendpointHolder;
    MakeDaliDqArrays(lshs, roots, frames, i - 1, tHolder, uHolder, vHolder, DendpointHolder);

    for(int j = 0; j < nLSHs; j++)
    {
        if(j > i) continue;
        if(j == i) 
        {
            fullJacob.block<3,3>(0, 3*j) = Dlsh(lshs[i].darboux, roots[i], frames[i], s);
            continue;
        }
        const MATRIX3 tRefDeriv = tHolder.block<3,3>(3*(i - 1), 3 * j);
        const MATRIX3 uRefDeriv = uHolder.block<3,3>(3*(i - 1), 3 * j);
        const MATRIX3 vRefDeriv = vHolder.block<3,3>(3*(i - 1), 3 * j);
        const MATRIX3 frameDerivContrib = lshs[i].darboux[0] * tRefDeriv + lshs[i].darboux[1] * uRefDeriv + lshs[i].darboux[2] * vRefDeriv;
        const VEC3 globalDarb = frames[i] * lshs[i].darboux;
            
        fullJacob.block<3,3>(0, 3*j) = DendpointHolder.block<3,3>(3*(i - 1), 3 * j);
        fullJacob.block<3,3>(0, 3*j) += DdarbEqnDom(globalDarb, roots[i], frames[i].col(0), s) * frameDerivContrib;
        fullJacob.block<3,3>(0, 3*j) += DdarbEqnDt(globalDarb, frames[i].col(0), s) * tRefDeriv;
    }

    return fullJacob;
}

std::vector<MATRIX> fullConnectedStateDerivs(const std::vector<superHelix>& lshs, const VEC3& p0, const MATRIX3& F0, const std::vector<double>& sAs)
{
    std::vector<MATRIX> fullJacobs(sAs.size());
    const int nLSHs = lshs.size();

    std::vector<MATRIX3> frames;
    std::vector<VEC3> roots;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);

    // get the big templating blocks for t, u, v
    // Maximum size for maximum referencing potential
    MATRIX tHolder, uHolder, vHolder;
    MakeD2aliDsDqArrays(lshs, frames, nLSHs - 1 - 1, tHolder, uHolder, vHolder);    
    MATRIX DendpointHolder;
    MakeDaliDqArrays(lshs, roots, frames, nLSHs - 1 - 1, tHolder, uHolder, vHolder, DendpointHolder);

    // precalculating total length and partial length sums of the superhelix
    std::vector<double> accumLen; //accumulated partial lengths
    std::vector<double> partLens; //lengths of each individual part
    double tolLen = 0;
    for(unsigned int i = 0; i < lshs.size(); i++)
    {
        double partLen = lshs[i].sQ;
        partLens.push_back(partLen);
        tolLen += partLen;
        accumLen.push_back(tolLen);
    }

    
    for(unsigned int n = 0; n < sAs.size(); n++)
    {
        MATRIX fullJacob(3, 3 * nLSHs);
        fullJacob.setZero();

        const double sA = sAs[n];
        const double scaledSampleLen = sA * tolLen;
        if(scaledSampleLen < 0 || scaledSampleLen > tolLen) continue;

        int nearestGroundInd = 0;
        for(unsigned int j = 0; j < accumLen.size(); j++)
        {
            if(accumLen[j] < scaledSampleLen) 
            {
                nearestGroundInd += 1;
            }
        }

        // split to proper distance with phiStar proportioning, 
        double leftoverLen = nearestGroundInd > 0 ? scaledSampleLen - accumLen[nearestGroundInd - 1] : scaledSampleLen;

        const int i = nearestGroundInd; // The index of the superhelix that presides over the requested sA
        const double s = leftoverLen; // The effective length along the particular superhelix

        if(i == 0)
        {
            fullJacob.block<3,3>(0,0) = Dlsh(lshs[i].darboux, roots[i], frames[i], s);
        }


        for(int j = 0; j < nLSHs; j++)
        {
            if(j > i) continue;
            if(j == i) 
            {
                fullJacob.block<3,3>(0, 3*j) = Dlsh(lshs[i].darboux, roots[i], frames[i], s);
                continue;
            }
            const MATRIX3 tRefDeriv = tHolder.block<3,3>(3*(i - 1), 3 * j);
            const MATRIX3 uRefDeriv = uHolder.block<3,3>(3*(i - 1), 3 * j);
            const MATRIX3 vRefDeriv = vHolder.block<3,3>(3*(i - 1), 3 * j);
            const MATRIX3 frameDerivContrib = lshs[i].darboux[0] * tRefDeriv + lshs[i].darboux[1] * uRefDeriv + lshs[i].darboux[2] * vRefDeriv;
            const VEC3 globalDarb = frames[i] * lshs[i].darboux;

            fullJacob.block<3,3>(0, 3*j) = DendpointHolder.block<3,3>(3*(i - 1), 3 * j);
            fullJacob.block<3,3>(0, 3*j) += DdarbEqnDom(globalDarb, roots[i], frames[i].col(0), s) * frameDerivContrib;
            fullJacob.block<3,3>(0, 3*j) += DdarbEqnDt(globalDarb, frames[i].col(0), s) * tRefDeriv;
        }

        fullJacobs[n] = fullJacob;
    }
    return fullJacobs;
}