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

#include <iostream>

VECTOR lshsToParam(const std::vector<superHelix>& lshs)
{
    VECTOR state(3 * lshs.size());
    for(unsigned int i = 0; i < lshs.size(); i++)
    {
        for(int j = 0; j< 3; j++) state[3*i + j] = lshs[i].darboux[j];
    }
    return state;
}

void updateLSHsState(std::vector<superHelix>& lshs, const VECTOR& newState)
{
    for(unsigned int i = 0; i < lshs.size(); i++)
    {
        for(int j = 0; j < 3; j++) lshs[i].darboux[j] = newState[3*i + j];
    }
}

VECTOR adamOptimizeLSHChainEnd(adamParams aP, const std::vector<superHelix>& lshs, const VEC3& p0, const MATRIX3& F0, const VEC3& targetEnd,
                            bool debug, double kSpring)
{
    const int nHelices = lshs.size();

    // std::cout << "making moment0..." << std::endl;
    VECTOR moment0(3 * nHelices);
    moment0.setZero();
    // std::cout << "making moment1..." << std::endl;
    VECTOR moment1(3 * nHelices);
    moment1.setZero();
    int numSteps = 0;
    bool converged = false;
    // std::cout << "making params..." << std::endl;
    VECTOR params = lshsToParam(lshs);
    const VECTOR params0 = params;

    std::vector<VEC3> roots;
    std::vector<MATRIX3> frames;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);
    const superHelix endLSH = lshs[lshs.size() - 1];
    // std::cout << "making pE..." << std::endl;
    VEC3 pE = sampleLocalSupHelix(endLSH, roots[lshs.size() - 1], frames[lshs.size() - 1], endLSH.sQ);
    double lowestLoss = (pE - targetEnd).squaredNorm() + kSpring * (params - params0).squaredNorm();

    VECTOR optimizedParams = params0;

    // std::cout << "making lshs1..." << std::endl;
    std::vector<superHelix> lshs1 = lshs;
    VECTOR ones(3 * nHelices);
    ones.setOnes();

    int stepsSinceImprovement = 0;
    while(!converged)
    {
        numSteps += 1;

        const VECTOR gradEnergy = 2 * (fullConnectedStateDeriv(lshs1, p0, F0, endLSH.sQ, nHelices - 1).transpose() * (pE - targetEnd) + kSpring * (params - params0));

        const VECTOR gradEnergy2 = gradEnergy.cwiseProduct(gradEnergy);
        moment0 = aP.bet1 * moment0 + (1 - aP.bet1) * gradEnergy;
        moment1 = aP.bet2 * moment1 + (1 - aP.bet2) * gradEnergy2;
        const VECTOR mom0Cor = moment0 / (1 - pow(aP.bet1, numSteps));
        const VECTOR mom1Cor = moment1 / (1 - pow(aP.bet2, numSteps));

        params = params - aP.stepSize * mom0Cor.cwiseQuotient(mom1Cor.cwiseSqrt() + aP.eps * ones);

        updateLSHsState(lshs1, params);
        fillRootsAndFrames(lshs1, roots, frames, p0, F0);
        const superHelix endLSH = lshs1[lshs1.size() - 1];
        pE = sampleLocalSupHelix(endLSH, roots[lshs1.size() - 1], frames[lshs1.size() - 1], endLSH.sQ);
        double loss = (pE - targetEnd).squaredNorm() + kSpring * (params - params0).squaredNorm();

        if(debug) std::cout << "loss: " << loss << "\n";
        if(loss < lowestLoss)
        {
            lowestLoss = loss;
            optimizedParams = params;
            stepsSinceImprovement = 0;
        }
        else
        {
            stepsSinceImprovement +=1;
        }
        if(stepsSinceImprovement > 50 || numSteps > 4000)
        {
            converged = true;
        }
        // std::cout << "stepsSinceImproviement: " << stepsSinceImprovement << std::endl;
    }
    if(debug) std::cout << "optimization took " << numSteps << " steps.\n";
    return optimizedParams;
}

VECTOR adamOptimizeLSHChainEndHotStart(adamParams aP, const std::vector<superHelix>& lshs, const VEC3& p0, const MATRIX3& F0, const VEC3& targetEnd,
                            const VECTOR& hotStart,
                            bool debug, double kSpring)
{
    const int nHelices = lshs.size();

    // std::cout << "making moment0..." << std::endl;
    VECTOR moment0(3 * nHelices);
    moment0.setZero();
    // std::cout << "making moment1..." << std::endl;
    VECTOR moment1(3 * nHelices);
    moment1.setZero();
    int numSteps = 0;
    bool converged = false;
    // std::cout << "making params..." << std::endl;
    VECTOR params = lshsToParam(lshs);
    const VECTOR params0 = params;

    std::vector<VEC3> roots;
    std::vector<MATRIX3> frames;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);
    const superHelix endLSH = lshs[lshs.size() - 1];
    // std::cout << "making pE..." << std::endl;
    VEC3 pE = sampleLocalSupHelix(endLSH, roots[lshs.size() - 1], frames[lshs.size() - 1], endLSH.sQ);
    double lowestLoss = (pE - targetEnd).squaredNorm() + kSpring * (params - params0).squaredNorm();

    VECTOR optimizedParams = hotStart;

    // std::cout << "making lshs1..." << std::endl;
    std::vector<superHelix> lshs1 = lshs;
    VECTOR ones(3 * nHelices);
    ones.setOnes();

    int stepsSinceImprovement = 0;
    while(!converged)
    {
        numSteps += 1;

        const VECTOR gradEnergy = 2 * (fullConnectedStateDeriv(lshs1, p0, F0, endLSH.sQ, nHelices - 1).transpose() * (pE - targetEnd) + kSpring * (params - params0));

        const VECTOR gradEnergy2 = gradEnergy.cwiseProduct(gradEnergy);
        moment0 = aP.bet1 * moment0 + (1 - aP.bet1) * gradEnergy;
        moment1 = aP.bet2 * moment1 + (1 - aP.bet2) * gradEnergy2;
        const VECTOR mom0Cor = moment0 / (1 - pow(aP.bet1, numSteps));
        const VECTOR mom1Cor = moment1 / (1 - pow(aP.bet2, numSteps));

        params = params - aP.stepSize * mom0Cor.cwiseQuotient(mom1Cor.cwiseSqrt() + aP.eps * ones);

        updateLSHsState(lshs1, params);
        fillRootsAndFrames(lshs1, roots, frames, p0, F0);
        const superHelix endLSH = lshs1[lshs1.size() - 1];
        pE = sampleLocalSupHelix(endLSH, roots[lshs1.size() - 1], frames[lshs1.size() - 1], endLSH.sQ);
        double loss = (pE - targetEnd).squaredNorm() + kSpring * (params - params0).squaredNorm();

        if(debug) std::cout << "loss: " << loss << "\n";
        if(loss < lowestLoss)
        {
            lowestLoss = loss;
            optimizedParams = params;
            stepsSinceImprovement = 0;
        }
        else
        {
            stepsSinceImprovement +=1;
        }
        if(stepsSinceImprovement > 50 || numSteps > 4000)
        {
            converged = true;
        }
        // std::cout << "stepsSinceImproviement: " << stepsSinceImprovement << std::endl;
    }
    if(debug) std::cout << "optimization took " << numSteps << " steps.\n";
    return optimizedParams;
}

VECTOR adamOptimizeLSHChainGeneral(adamParams aP, const std::vector<superHelix> & lshs, const VEC3& p0, const MATRIX3& F0, const std::vector<double> sAs, const std::vector<VEC3>& targetEnds, 
                                   bool debug, double kSpring)
{
    const int nHelices = lshs.size();

    // std::cout << "making moment0..." << std::endl;
    VECTOR moment0(3 * nHelices);
    moment0.setZero();
    // std::cout << "making moment1..." << std::endl;
    VECTOR moment1(3 * nHelices);
    moment1.setZero();
    int numSteps = 0;
    bool converged = false;
    // std::cout << "making params..." << std::endl;
    VECTOR params = lshsToParam(lshs);
    const VECTOR params0 = params;

    std::vector<VEC3> roots;
    std::vector<MATRIX3> frames;
    fillRootsAndFrames(lshs, roots, frames, p0, F0);
    const superHelix endLSH = lshs[lshs.size() - 1];
    // std::cout << "making pEs..." << std::endl;

    std::vector<VEC3> pEs = sampleLocalSupHelixes(lshs, p0, F0, sAs); 
    double lowestLoss = kSpring * (params - params0).squaredNorm();
    for(unsigned int i = 0; i < pEs.size(); i++)
    {
        lowestLoss += (pEs[i] - targetEnds[i]).squaredNorm();
    }

    VECTOR optimizedParams = params0;

    // std::cout << "making lshs1..." << std::endl;
    std::vector<superHelix> lshs1 = lshs;
    VECTOR ones(3 * nHelices);
    ones.setOnes();

    std::vector<MATRIX> indivDerivs;

    int stepsSinceImprovement = 0;
    while(!converged)
    {
        numSteps += 1;

        // VECTOR gradEnergy = 2 * (fullConnectedStateDeriv(lshs1, p0, F0, endLSH.sQ, nHelices - 1).transpose() * (pE - targetEnd) + kSpring * (params - params0));
        // first the spring part
        VECTOR gradEnergy = 2 * kSpring * (params - params0);
        for(unsigned int i = 0; i < sAs.size(); i++)
        {
            indivDerivs = fullConnectedStateDerivs(lshs1, p0, F0, sAs);
        }
        for(unsigned int i = 0; i < sAs.size(); i++)
        {
            VECTOR gradComponent = 2 * indivDerivs[i].transpose() * (pEs[i] - targetEnds[i]);
            if(debug) std::cout << "gradComponent.norm(): " << gradComponent.norm() << std::endl;
            gradEnergy += 2 * indivDerivs[i].transpose() * (pEs[i] - targetEnds[i]);
        }

        const VECTOR gradEnergy2 = gradEnergy.cwiseProduct(gradEnergy);
        moment0 = aP.bet1 * moment0 + (1 - aP.bet1) * gradEnergy;
        moment1 = aP.bet2 * moment1 + (1 - aP.bet2) * gradEnergy2;
        const VECTOR mom0Cor = moment0 / (1 - pow(aP.bet1, numSteps));
        const VECTOR mom1Cor = moment1 / (1 - pow(aP.bet2, numSteps));
        
        params = params - aP.stepSize * mom0Cor.cwiseQuotient(mom1Cor.cwiseSqrt() + aP.eps * ones);

        updateLSHsState(lshs1, params);
        fillRootsAndFrames(lshs1, roots, frames, p0, F0);

        pEs = sampleLocalSupHelixes(lshs1, p0, F0, sAs); 
        double loss = kSpring * (params - params0).squaredNorm();
        for(unsigned int i = 0; i < pEs.size(); i++)
        {
            loss += (pEs[i] - targetEnds[i]).squaredNorm();
        }

        if(debug) std::cout << "loss: " << loss << "\n";
        if(loss < lowestLoss)
        {
            lowestLoss = loss;
            optimizedParams = params;
            stepsSinceImprovement = 0;
        }
        else
        {
            stepsSinceImprovement +=1;
        }
        if(stepsSinceImprovement > 50 || numSteps > 4000 || lowestLoss < 1e-6)
        {
            converged = true;
        }
        // std::cout << "stepsSinceImproviement: " << stepsSinceImprovement << std::endl;
    }
    if(debug) std::cout << "optimization took " << numSteps << " steps.\n";
    return optimizedParams;
}